Skip to content

Commit

Permalink
Verifiable payer (#544)
Browse files Browse the repository at this point in the history 
Uses xmtp/proto#247

This adds a new field to the payer envelope that prevents a node from
stealing payloads #474

Additionally it gets rid of old AAD validation #463

The bulk of the changeset allows the payer to pick a different node if
the connection to the ideal node has failed. This is obviously not
ideal, because it will keep retrying even if the node is down long term,
but in the short/medium term this is sufficient.

Fixes #474 #463
  • Loading branch information
@mkysel
mkysel authored Feb 24, 2025
1 parent b5dfeff commit 63fe167
Show file tree
Hide file tree
Showing 18 changed files with 320 additions and 117 deletions.
93 changes: 71 additions & 22 deletions pkg/api/message/publish_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,10 @@ func TestPublishEnvelope(t *testing.T) {
api, db, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

payerEnvelope := envelopeTestUtils.CreatePayerEnvelope(t)
payerEnvelope := envelopeTestUtils.CreatePayerEnvelope(
t,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
)

resp, err := api.PublishPayerEnvelopes(
context.Background(),
Expand Down Expand Up @@ -70,7 +73,10 @@ func TestUnmarshalErrorOnPublish(t *testing.T) {
api, _, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

envelope := envelopeTestUtils.CreatePayerEnvelope(t)
envelope := envelopeTestUtils.CreatePayerEnvelope(
t,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
)
envelope.UnsignedClientEnvelope = []byte("invalidbytes")
_, err := api.PublishPayerEnvelopes(
context.Background(),
Expand All @@ -81,19 +87,42 @@ func TestUnmarshalErrorOnPublish(t *testing.T) {
require.ErrorContains(t, err, "invalid wire-format data")
}

func TestMismatchingOriginatorOnPublish(t *testing.T) {
func TestMismatchingAADOriginatorOnPublishNoLongerFails(t *testing.T) {
api, _, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

nid := envelopeTestUtils.DefaultClientEnvelopeNodeId + 100

clientEnv := envelopeTestUtils.CreateClientEnvelope()
nodeId := uint32(2)
// nolint:staticcheck
clientEnv.Aad.TargetOriginator = &nodeId
clientEnv.Aad.TargetOriginator = &nid
_, err := api.PublishPayerEnvelopes(
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{
envelopeTestUtils.CreatePayerEnvelope(
t,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
clientEnv,
),
},
},
)
require.NoError(t, err)
}

func TestMismatchingOriginatorOnPublish(t *testing.T) {
api, _, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

nid := envelopeTestUtils.DefaultClientEnvelopeNodeId + 100

clientEnv := envelopeTestUtils.CreateClientEnvelope()
_, err := api.PublishPayerEnvelopes(
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{
envelopeTestUtils.CreatePayerEnvelope(t, clientEnv),
envelopeTestUtils.CreatePayerEnvelope(t, nid, clientEnv),
},
},
)
Expand All @@ -110,7 +139,11 @@ func TestMissingTopicOnPublish(t *testing.T) {
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{
envelopeTestUtils.CreatePayerEnvelope(t, clientEnv),
envelopeTestUtils.CreatePayerEnvelope(
t,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
clientEnv,
),
},
},
)
Expand All @@ -120,11 +153,12 @@ func TestMissingTopicOnPublish(t *testing.T) {
func TestKeyPackageValidationSuccess(t *testing.T) {
api, _, apiMocks, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()
nodeId := uint32(100)

nid := envelopeTestUtils.DefaultClientEnvelopeNodeId

clientEnv := envelopeTestUtils.CreateClientEnvelope(&envelopes.AuthenticatedData{
TargetTopic: topic.NewTopic(topic.TOPIC_KIND_KEY_PACKAGES_V1, []byte{1, 2, 3}).Bytes(),
TargetOriginator: &nodeId,
TargetOriginator: &nid,
DependsOn: &envelopes.Cursor{},
})
clientEnv.Payload = &envelopes.ClientEnvelope_UploadKeyPackage{
Expand All @@ -150,7 +184,11 @@ func TestKeyPackageValidationSuccess(t *testing.T) {
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{
envelopeTestUtils.CreatePayerEnvelope(t, clientEnv),
envelopeTestUtils.CreatePayerEnvelope(
t,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
clientEnv,
),
},
},
)
Expand All @@ -160,11 +198,12 @@ func TestKeyPackageValidationSuccess(t *testing.T) {
func TestKeyPackageValidationFail(t *testing.T) {
api, _, apiMocks, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()
nodeId := uint32(100)

nid := envelopeTestUtils.DefaultClientEnvelopeNodeId

clientEnv := envelopeTestUtils.CreateClientEnvelope(&envelopes.AuthenticatedData{
TargetTopic: topic.NewTopic(topic.TOPIC_KIND_KEY_PACKAGES_V1, []byte{1, 2, 3}).Bytes(),
TargetOriginator: &nodeId,
TargetOriginator: &nid,
DependsOn: &envelopes.Cursor{},
})
clientEnv.Payload = &envelopes.ClientEnvelope_UploadKeyPackage{
Expand All @@ -190,7 +229,7 @@ func TestKeyPackageValidationFail(t *testing.T) {
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{
envelopeTestUtils.CreatePayerEnvelope(t, clientEnv),
envelopeTestUtils.CreatePayerEnvelope(t, nid, clientEnv),
},
},
)
Expand All @@ -201,11 +240,16 @@ func TestPublishEnvelopeBlockchainCursorAhead(t *testing.T) {
api, _, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

err := publishPayerEnvelopeWithNodeIDAndCursor(t, api, 100, &envelopes.Cursor{
NodeIdToSequenceId: map[uint32]uint64{
1: 105,
err := publishPayerEnvelopeWithNodeIDAndCursor(
t,
api,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
&envelopes.Cursor{
NodeIdToSequenceId: map[uint32]uint64{
1: 105,
},
},
})
)
require.Error(t, err)
require.Contains(t, err.Error(), "DependsOn has not been seen by this node")
}
Expand All @@ -220,7 +264,7 @@ func publishPayerEnvelopeWithNodeIDAndCursor(
context.Background(),
&message_api.PublishPayerEnvelopesRequest{
PayerEnvelopes: []*envelopes.PayerEnvelope{envelopeTestUtils.CreatePayerEnvelope(
t,
t, nodeId,
envelopeTestUtils.CreateClientEnvelope(&envelopes.AuthenticatedData{
TargetOriginator: &nodeId,
TargetTopic: targetTopic,
Expand All @@ -237,11 +281,16 @@ func TestPublishEnvelopeOriginatorUnknown(t *testing.T) {
api, _, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
defer cleanup()

err := publishPayerEnvelopeWithNodeIDAndCursor(t, api, 100, &envelopes.Cursor{
NodeIdToSequenceId: map[uint32]uint64{
1600: 1,
err := publishPayerEnvelopeWithNodeIDAndCursor(
t,
api,
envelopeTestUtils.DefaultClientEnvelopeNodeId,
&envelopes.Cursor{
NodeIdToSequenceId: map[uint32]uint64{
1600: 1,
},
},
})
)
require.Error(t, err)
require.Contains(t, err.Error(), "DependsOn has not been seen by this node")
}
12 changes: 8 additions & 4 deletions pkg/api/message/service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,14 @@ func (s *Service) validatePayerEnvelope(
return nil, err
}

if payerEnv.TargetOriginator != s.registrant.NodeID() {
return nil, status.Errorf(codes.InvalidArgument, "invalid target originator")
}

if _, err = payerEnv.RecoverSigner(); err != nil {
return nil, err
}

if err := s.validateClientInfo(&payerEnv.ClientEnvelope); err != nil {
return nil, err
}
Expand Down Expand Up @@ -444,10 +452,6 @@ func (s *Service) validateKeyPackage(

func (s *Service) validateClientInfo(clientEnv *envelopes.ClientEnvelope) error {
aad := clientEnv.Aad()
// nolint:staticcheck
if aad.GetTargetOriginator() != s.registrant.NodeID() {
return status.Errorf(codes.InvalidArgument, "invalid target originator")
}

if !clientEnv.TopicMatchesPayload() {
return status.Errorf(codes.InvalidArgument, "topic does not match payload")
Expand Down
27 changes: 24 additions & 3 deletions pkg/api/payer/nodeSelector.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import (
)

type NodeSelectorAlgorithm interface {
GetNode(topic topic.Topic) (uint32, error)
GetNode(topic topic.Topic, banlist ...[]uint32) (uint32, error)
}

type StableHashingNodeSelectorAlgorithm struct {
Expand All @@ -30,7 +30,10 @@ func HashKey(topic topic.Topic) uint32 {
}

// GetNode selects a node for a given topic using stable hashing
func (s *StableHashingNodeSelectorAlgorithm) GetNode(topic topic.Topic) (uint32, error) {
func (s *StableHashingNodeSelectorAlgorithm) GetNode(
topic topic.Topic,
banlist ...[]uint32,
) (uint32, error) {
nodes, err := s.reg.GetNodes()
if err != nil {
return 0, err
Expand Down Expand Up @@ -60,5 +63,23 @@ func (s *StableHashingNodeSelectorAlgorithm) GetNode(topic topic.Topic) (uint32,
return topicHash < nodeLocations[i]
})

return nodes[idx%len(nodeLocations)].NodeID, nil
// Flatten banlist
banned := make(map[uint32]struct{})
for _, list := range banlist {
for _, id := range list {
banned[id] = struct{}{}
}
}

// Find the next available node
for i := 0; i < len(nodes); i++ {
candidateIdx := (idx + i) % len(nodeLocations)
candidateNodeID := nodes[candidateIdx].NodeID

if _, exists := banned[candidateNodeID]; !exists {
return candidateNodeID, nil
}
}

return 0, errors.New("no available nodes after considering banlist")
}
37 changes: 37 additions & 0 deletions pkg/api/payer/nodeSelector_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,3 +226,40 @@ func TestGetNode_ConfirmTopicBalance(t *testing.T) {
}

}

func TestGetNode_NodeGetNextIfBanned(t *testing.T) {
mockRegistry := mocks.NewMockNodeRegistry(t)
mockRegistry.On("GetNodes").Return([]registry.Node{
{NodeID: 100},
{NodeID: 200},
{NodeID: 300},
}, nil)
tpc1 := *topic.NewTopic(topic.TOPIC_KIND_IDENTITY_UPDATES_V1, []byte("stable_key"))

selector := payer.NewStableHashingNodeSelectorAlgorithm(mockRegistry)

node1, err := selector.GetNode(tpc1)
require.NoError(t, err)
require.EqualValues(t, 200, node1)

banlist := []uint32{node1}

reselectedNode, err := selector.GetNode(tpc1, banlist)
require.NoError(t, err)
require.NotEqualValues(t, node1, reselectedNode)
require.EqualValues(t, 300, reselectedNode)

banlist = append(banlist, reselectedNode)

reselectedNode, err = selector.GetNode(tpc1, banlist)
require.NoError(t, err)
require.NotEqualValues(t, node1, reselectedNode)
require.EqualValues(t, 100, reselectedNode)

banlist = append(banlist, reselectedNode)

// now we are out of nodes to try
reselectedNode, err = selector.GetNode(tpc1, banlist)
require.Error(t, err)
require.EqualValues(t, 0, reselectedNode)
}
9 changes: 7 additions & 2 deletions pkg/api/payer/publish_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,11 +181,14 @@ func TestPublishToNodes(t *testing.T) {
HttpAddress: formatAddress(originatorServer.Addr().String()),
}, nil)

mockRegistry.On("GetNodes").Return([]registry.Node{
{NodeID: 100},
}, nil)

groupId := testutils.RandomGroupID()
testGroupMessage := envelopesTestUtils.CreateGroupMessageClientEnvelope(
groupId,
[]byte("test message"),
100, // This is the expected originator ID of the test server
)

publishResponse, err := svc.PublishClientEnvelopes(
Expand All @@ -203,6 +206,8 @@ func TestPublishToNodes(t *testing.T) {
require.NoError(t, err)

targetTopic := parsedOriginatorEnvelope.UnsignedOriginatorEnvelope.PayerEnvelope.ClientEnvelope.TargetTopic()

require.Equal(t, targetTopic.Identifier(), groupId[:])

targetOriginator := parsedOriginatorEnvelope.UnsignedOriginatorEnvelope.PayerEnvelope.TargetOriginator
require.EqualValues(t, 100, targetOriginator)
}
Loading

0 comments on commit 63fe167

Please sign in to comment.