Implement PublishIdentityUpdate and GetInboxLogs endpoints

pkg/identity/api/v1/service.go → pkg/identity/api/v1/identity_service.go

@@ -75,7 +75,7 @@ Start transaction (SERIALIZABLE isolation level)
 End transaction
 */
 func (s *Service) PublishIdentityUpdate(ctx context.Context, req *api.PublishIdentityUpdateRequest) (*api.PublishIdentityUpdateResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
+	return s.store.PublishIdentityUpdate(ctx, req)
 }
 
 func (s *Service) GetIdentityUpdates(ctx context.Context, req *api.GetIdentityUpdatesRequest) (*api.GetIdentityUpdatesResponse, error) {
@@ -84,7 +84,7 @@ func (s *Service) GetIdentityUpdates(ctx context.Context, req *api.GetIdentityUp
 		1. Query the inbox_log table for the inbox_id, ordering by sequence_id
 		2. Return all of the entries
 	*/
-	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
+	return s.store.GetInboxLogs(ctx, req)
 }
 
 func (s *Service) GetInboxIds(ctx context.Context, req *api.GetInboxIdsRequest) (*api.GetInboxIdsResponse, error) {

pkg/identity/api/v1/identity_service_test.go

@@ -0,0 +1,218 @@
+package api
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/stretchr/testify/require"
+	"github.com/uptrace/bun"
+	mlsstore "github.com/xmtp/xmtp-node-go/pkg/mls/store"
+	identity "github.com/xmtp/xmtp-node-go/pkg/proto/identity/api/v1"
+	associations "github.com/xmtp/xmtp-node-go/pkg/proto/identity/associations"
+	test "github.com/xmtp/xmtp-node-go/pkg/testing"
+)
+
+func newTestService(t *testing.T, ctx context.Context) (*Service, *bun.DB, func()) {
+	log := test.NewLog(t)
+	db, _, mlsDbCleanup := test.NewMLSDB(t)
+	store, err := mlsstore.New(ctx, mlsstore.Config{
+		Log: log,
+		DB:  db,
+	})
+	require.NoError(t, err)
+	natsServer, err := server.NewServer(&server.Options{
+		Port: server.RANDOM_PORT,
+	})
+	require.NoError(t, err)
+	go natsServer.Start()
+	if !natsServer.ReadyForConnections(4 * time.Second) {
+		t.Fail()
+	}
+
+	svc, err := NewService(log, store)
+	require.NoError(t, err)
+
+	return svc, db, func() {
+		svc.Close()
+		mlsDbCleanup()
+	}
+}
+
+func makeCreateInbox(address string) *associations.IdentityAction {
+	return &associations.IdentityAction{
+		Kind: &associations.IdentityAction_CreateInbox{
+			CreateInbox: &associations.CreateInbox{
+				InitialAddress:          address,
+				Nonce:                   0,
+				InitialAddressSignature: &associations.Signature{},
+			},
+		},
+	}
+}
+func makeAddAssociation() *associations.IdentityAction {
+	return &associations.IdentityAction{
+		Kind: &associations.IdentityAction_Add{
+			Add: &associations.AddAssociation{
+				NewMemberIdentifier:     &associations.MemberIdentifier{},
+				ExistingMemberSignature: &associations.Signature{},
+				NewMemberSignature:      &associations.Signature{},
+			},
+		},
+	}
+}
+func makeRevokeAssociation() *associations.IdentityAction {
+	return &associations.IdentityAction{
+		Kind: &associations.IdentityAction_Revoke{
+			Revoke: &associations.RevokeAssociation{
+				MemberToRevoke:           &associations.MemberIdentifier{},
+				RecoveryAddressSignature: &associations.Signature{},
+			},
+		},
+	}
+}
+func makeChangeRecoveryAddress() *associations.IdentityAction {
+	return &associations.IdentityAction{
+		Kind: &associations.IdentityAction_ChangeRecoveryAddress{
+			ChangeRecoveryAddress: &associations.ChangeRecoveryAddress{
+				NewRecoveryAddress:               "",
+				ExistingRecoveryAddressSignature: &associations.Signature{},
+			},
+		},
+	}
+}
+func makeIdentityUpdate(inbox_id string, actions ...*associations.IdentityAction) *associations.IdentityUpdate {
+	return &associations.IdentityUpdate{
+		InboxId:           inbox_id,
+		ClientTimestampNs: 0,
+		Actions:           actions,
+	}
+}
+
+func publishIdentityUpdateRequest(inbox_id string, actions ...*associations.IdentityAction) *identity.PublishIdentityUpdateRequest {
+	return &identity.PublishIdentityUpdateRequest{
+		IdentityUpdate: makeIdentityUpdate(inbox_id, actions...),
+	}
+}
+
+func makeUpdateRequest(inbox_id string, sequence_id uint64) *identity.GetIdentityUpdatesRequest_Request {
+	return &identity.GetIdentityUpdatesRequest_Request{
+		InboxId:    inbox_id,
+		SequenceId: sequence_id,
+	}
+}
+
+func getIdentityUpdatesRequest(requests ...*identity.GetIdentityUpdatesRequest_Request) *identity.GetIdentityUpdatesRequest {
+	return &identity.GetIdentityUpdatesRequest{
+		Requests: requests,
+	}
+}
+
+func TestPublishedUpdatesCanBeRead(t *testing.T) {
+	ctx := context.Background()
+	svc, _, cleanup := newTestService(t, ctx)
+	defer cleanup()
+
+	inbox_id := "test_inbox"
+	address := "test_address"
+
+	_, err := svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeCreateInbox(address)))
+	require.NoError(t, err)
+
+	res, err := svc.GetIdentityUpdates(ctx, getIdentityUpdatesRequest(makeUpdateRequest(inbox_id, 0)))
+	require.NoError(t, err)
+
+	require.Len(t, res.Responses, 1)
+	require.Equal(t, res.Responses[0].InboxId, inbox_id)
+	require.Len(t, res.Responses[0].Updates, 1)
+	require.Len(t, res.Responses[0].Updates[0].Update.Actions, 1)
+	require.Equal(t, res.Responses[0].Updates[0].Update.Actions[0].GetCreateInbox().InitialAddress, address)
+}
+
+func TestPublishedUpdatesAreInOrder(t *testing.T) {
+	ctx := context.Background()
+	svc, _, cleanup := newTestService(t, ctx)
+	defer cleanup()
+
+	inbox_id := "test_inbox"
+	address := "test_address"
+
+	_, err := svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeCreateInbox(address)))
+	require.NoError(t, err)
+	_, err = svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeAddAssociation(), makeChangeRecoveryAddress()))
+	require.NoError(t, err)
+	_, err = svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeRevokeAssociation()))
+	require.NoError(t, err)
+
+	res, err := svc.GetIdentityUpdates(ctx, getIdentityUpdatesRequest(makeUpdateRequest(inbox_id, 0)))
+	require.NoError(t, err)
+
+	require.Len(t, res.Responses, 1)
+	require.Equal(t, res.Responses[0].InboxId, inbox_id)
+	require.Len(t, res.Responses[0].Updates, 3)
+	require.NotNil(t, res.Responses[0].Updates[0].Update.Actions[0].GetCreateInbox())
+	require.NotNil(t, res.Responses[0].Updates[1].Update.Actions[0].GetAdd())
+	require.NotNil(t, res.Responses[0].Updates[1].Update.Actions[1].GetChangeRecoveryAddress())
+	require.NotNil(t, res.Responses[0].Updates[2].Update.Actions[0].GetRevoke())
+
+	res, err = svc.GetIdentityUpdates(ctx, getIdentityUpdatesRequest(makeUpdateRequest(inbox_id, 1)))
+	require.NoError(t, err)
+
+	require.Len(t, res.Responses, 1)
+	require.Equal(t, res.Responses[0].InboxId, inbox_id)
+	require.Len(t, res.Responses[0].Updates, 2)
+	require.NotNil(t, res.Responses[0].Updates[0].Update.Actions[0].GetAdd())
+	require.NotNil(t, res.Responses[0].Updates[0].Update.Actions[1].GetChangeRecoveryAddress())
+	require.NotNil(t, res.Responses[0].Updates[1].Update.Actions[0].GetRevoke())
+}
+
+func TestQueryMultipleInboxes(t *testing.T) {
+	ctx := context.Background()
+	svc, _, cleanup := newTestService(t, ctx)
+	defer cleanup()
+
+	first_inbox_id := "test_inbox"
+	second_inbox_id := "second_inbox"
+	first_address := "test_address"
+	second_address := "test_address"
+
+	_, err := svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(first_inbox_id, makeCreateInbox(first_address)))
+	require.NoError(t, err)
+	_, err = svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(second_inbox_id, makeCreateInbox(second_address)))
+	require.NoError(t, err)
+
+	res, err := svc.GetIdentityUpdates(ctx, getIdentityUpdatesRequest(makeUpdateRequest(first_inbox_id, 0), makeUpdateRequest(second_inbox_id, 0)))
+	require.NoError(t, err)
+
+	require.Len(t, res.Responses, 2)
+	require.Equal(t, res.Responses[0].Updates[0].Update.Actions[0].GetCreateInbox().InitialAddress, first_address)
+	require.Equal(t, res.Responses[1].Updates[0].Update.Actions[0].GetCreateInbox().InitialAddress, second_address)
+}
+
+func TestInboxSizeLimit(t *testing.T) {
+	ctx := context.Background()
+	svc, _, cleanup := newTestService(t, ctx)
+	defer cleanup()
+
+	inbox_id := "test_inbox"
+	address := "test_address"
+
+	_, err := svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeCreateInbox(address)))
+	require.NoError(t, err)
+
+	for i := 0; i < 255; i++ {
+		_, err = svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeAddAssociation()))
+		require.NoError(t, err)
+	}
+
+	_, err = svc.PublishIdentityUpdate(ctx, publishIdentityUpdateRequest(inbox_id, makeAddAssociation()))
+	require.Error(t, err)
+
+	res, err := svc.GetIdentityUpdates(ctx, getIdentityUpdatesRequest(makeUpdateRequest(inbox_id, 0)))
+	require.NoError(t, err)
+
+	require.Len(t, res.Responses, 1)
+	require.Equal(t, res.Responses[0].InboxId, inbox_id)
+	require.Len(t, res.Responses[0].Updates, 256)
+}

pkg/mls/store/models.go

@@ -6,6 +6,15 @@ import (
 	"github.com/uptrace/bun"
 )
 
+type InboxLogEntry struct {
+	bun.BaseModel `bun:"table:inbox_log"`
+
+	SequenceId          uint64 `bun:",autoincrement"`
+	InboxId             string
+	ServerTimestampNs   int64
+	IdentityUpdateProto []byte
+}
+
 type Installation struct {
 	bun.BaseModel `bun:"table:installations"`
 

pkg/mls/store/store.go

@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"crypto/sha256"
+	"database/sql"
 	"errors"
 	"sort"
 	"strings"
@@ -11,8 +12,11 @@ import (
 	"github.com/uptrace/bun"
 	"github.com/uptrace/bun/migrate"
 	migrations "github.com/xmtp/xmtp-node-go/pkg/migrations/mls"
+	identity "github.com/xmtp/xmtp-node-go/pkg/proto/identity/api/v1"
+	"github.com/xmtp/xmtp-node-go/pkg/proto/identity/associations"
 	mlsv1 "github.com/xmtp/xmtp-node-go/pkg/proto/mls/api/v1"
 	"go.uber.org/zap"
+	"google.golang.org/protobuf/proto"
 )
 
 const maxPageSize = 100
@@ -23,7 +27,14 @@ type Store struct {
 	db     *bun.DB
 }
 
+type IdentityStore interface {
+	PublishIdentityUpdate(ctx context.Context, req *identity.PublishIdentityUpdateRequest) (*identity.PublishIdentityUpdateResponse, error)
+	GetInboxLogs(ctx context.Context, req *identity.GetIdentityUpdatesRequest) (*identity.GetIdentityUpdatesResponse, error)
+}
+
 type MlsStore interface {
+	IdentityStore
+
 	CreateInstallation(ctx context.Context, installationId []byte, walletAddress string, credentialIdentity, keyPackage []byte, expiration uint64) error
 	UpdateKeyPackage(ctx context.Context, installationId, keyPackage []byte, expiration uint64) error
 	FetchKeyPackages(ctx context.Context, installationIds [][]byte) ([]*Installation, error)
@@ -53,6 +64,111 @@ func New(ctx context.Context, config Config) (*Store, error) {
 	return s, nil
 }
 
+func (s *Store) PublishIdentityUpdate(ctx context.Context, req *identity.PublishIdentityUpdateRequest) (*identity.PublishIdentityUpdateResponse, error) {
+	new_update := req.GetIdentityUpdate()
+	if new_update == nil {
+		return nil, errors.New("IdentityUpdate is required")
+	}
+
+	// TODO: Implement serializable isolation level once supported
+	if err := s.db.RunInTx(ctx, &sql.TxOptions{ /*Isolation: sql.LevelSerializable*/ }, func(ctx context.Context, tx bun.Tx) error {
+		inbox_log_entries := make([]*InboxLogEntry, 0)
+
+		if err := s.db.NewSelect().
+			Model(&inbox_log_entries).
+			Where("inbox_id = ?", new_update.GetInboxId()).
+			Order("sequence_id ASC").
+			For("UPDATE").
+			Scan(ctx); err != nil {
+			return err
+		}
+
+		if len(inbox_log_entries) >= 256 {
+			return errors.New("inbox log is full")
+		}
+
+		updates := make([]*associations.IdentityUpdate, 0, len(inbox_log_entries)+1)
+		for _, log := range inbox_log_entries {
+			identity_update := &associations.IdentityUpdate{}
+			if err := proto.Unmarshal(log.IdentityUpdateProto, identity_update); err != nil {
+				return err
+			}
+			updates = append(updates, identity_update)
+		}
+		_ = append(updates, new_update)
+
+		// TODO: Validate the updates, and abort transaction if failed
+
+		proto_bytes, err := proto.Marshal(new_update)
+		if err != nil {
+			return err
+		}
+
+		new_entry := InboxLogEntry{
+			InboxId:             new_update.GetInboxId(),
+			ServerTimestampNs:   nowNs(),
+			IdentityUpdateProto: proto_bytes,
+		}
+
+		_, err = s.db.NewInsert().
+			Model(&new_entry).
+			Returning("sequence_id").
+			Exec(ctx)
+		if err != nil {
+			return err
+		}
+
+		// TODO: Insert or update the address_log table using sequence_id
+
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return &identity.PublishIdentityUpdateResponse{}, nil
+}
+
+func (s *Store) GetInboxLogs(ctx context.Context, batched_req *identity.GetIdentityUpdatesRequest) (*identity.GetIdentityUpdatesResponse, error) {
+	reqs := batched_req.GetRequests()
+	resps := make([]*identity.GetIdentityUpdatesResponse_Response, len(reqs))
+
+	for i, req := range reqs {
+		inbox_log_entries := make([]*InboxLogEntry, 0)
+
+		err := s.db.NewSelect().
+			Model(&inbox_log_entries).
+			Where("sequence_id > ?", req.GetSequenceId()).
+			Where("inbox_id = ?", req.GetInboxId()).
+			Order("sequence_id ASC").
+			Scan(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		updates := make([]*identity.GetIdentityUpdatesResponse_IdentityUpdateLog, len(inbox_log_entries))
+		for j, entry := range inbox_log_entries {
+			identity_update := &associations.IdentityUpdate{}
+			if err := proto.Unmarshal(entry.IdentityUpdateProto, identity_update); err != nil {
+				return nil, err
+			}
+			updates[j] = &identity.GetIdentityUpdatesResponse_IdentityUpdateLog{
+				SequenceId:        entry.SequenceId,
+				ServerTimestampNs: uint64(entry.ServerTimestampNs),
+				Update:            identity_update,
+			}
+		}
+
+		resps[i] = &identity.GetIdentityUpdatesResponse_Response{
+			InboxId: req.GetInboxId(),
+			Updates: updates,
+		}
+	}
+
+	return &identity.GetIdentityUpdatesResponse{
+		Responses: resps,
+	}, nil
+}
+
 // Creates the installation and last resort key package
 func (s *Store) CreateInstallation(ctx context.Context, installationId []byte, walletAddress string, credentialIdentity, keyPackage []byte, expiration uint64) error {
 	createdAt := nowNs()