Working end-to-end

cmd/xmtpd/main.go

@@ -92,6 +92,13 @@ func main() {
 		return
 	}
 
+	if options.CreateMlsMigration != "" && options.MlsStore.DbConnectionString != "" {
+		if err := server.CreateMlsMigration(options.CreateMlsMigration, options.MlsStore.DbConnectionString, options.WaitForDB, options.MlsStore.ReadTimeout, options.MlsStore.WriteTimeout, options.Store.MaxOpenConns); err != nil {
+			log.Fatal("creating authz db migration", zap.Error(err))
+		}
+		return
+	}
+
 	if options.Tracing.Enable {
 		log.Info("starting tracer")
 		tracing.Start(Commit, utils.Logger())

go.mod

@@ -25,12 +25,12 @@ require (
 	github.com/prometheus/client_golang v1.14.0
 	github.com/stretchr/testify v1.8.4
 	github.com/swaggest/swgui v1.6.2
-	github.com/uptrace/bun v1.1.3
-	github.com/uptrace/bun/dialect/pgdialect v1.1.3
-	github.com/uptrace/bun/driver/pgdriver v1.1.3
+	github.com/uptrace/bun v1.1.16
+	github.com/uptrace/bun/dialect/pgdialect v1.1.16
+	github.com/uptrace/bun/driver/pgdriver v1.1.16
 	github.com/waku-org/go-waku v0.8.0
 	github.com/xmtp/go-msgio v0.2.1-0.20220510223757-25a701b79cd3
-	github.com/xmtp/proto/v3 v3.30.0
+	github.com/xmtp/proto/v3 v3.29.1-0.20231023182354-832c8d572ed4
 	github.com/yoheimuta/protolint v0.39.0
 	go.uber.org/zap v1.24.0
 	golang.org/x/sync v0.3.0
@@ -108,7 +108,7 @@ require (
 	github.com/libp2p/go-yamux/v4 v4.0.1 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/marten-seemann/tcp v0.0.0-20210406111302-dfbc87cc63fd // indirect
-	github.com/mattn/go-colorable v0.1.11 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/miekg/dns v1.1.55 // indirect
@@ -148,6 +148,7 @@ require (
 	github.com/shirou/gopsutil v3.21.4-0.20210419000835-c7a38de76ee5+incompatible // indirect
 	github.com/shurcooL/httpgzip v0.0.0-20190720172056-320755c1c1b0 // indirect
 	github.com/spaolacci/murmur3 v1.1.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/syndtr/goleveldb v1.0.1-0.20220614013038-64ee5596c38a // indirect
 	github.com/tinylib/msgp v1.1.2 // indirect
 	github.com/tklauser/go-sysconf v0.3.5 // indirect
@@ -167,12 +168,12 @@ require (
 	go.uber.org/dig v1.17.0 // indirect
 	go.uber.org/fx v1.20.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/crypto v0.12.0 // indirect
+	golang.org/x/crypto v0.13.0 // indirect
 	golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect
 	golang.org/x/mod v0.12.0 // indirect
 	golang.org/x/net v0.14.0 // indirect
-	golang.org/x/sys v0.11.0 // indirect
-	golang.org/x/text v0.12.0 // indirect
+	golang.org/x/sys v0.12.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
 	golang.org/x/time v0.0.0-20220922220347-f3bd1da661af // indirect
 	golang.org/x/tools v0.12.1-0.20230818130535-1517d1a3ba60 // indirect
 	golang.org/x/xerrors v0.0.0-20220517211312-f3a8303e98df // indirect
@@ -181,7 +182,7 @@ require (
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	lukechampine.com/blake3 v1.2.1 // indirect
-	mellium.im/sasl v0.2.1 // indirect
+	mellium.im/sasl v0.3.1 // indirect
 )
 
 // From node-go

pkg/api/config.go

@@ -8,6 +8,7 @@ import (
 	wakunode "github.com/waku-org/go-waku/waku/v2/node"
 	"github.com/xmtp/xmtp-node-go/pkg/authz"
 	"github.com/xmtp/xmtp-node-go/pkg/mlsstore"
+	"github.com/xmtp/xmtp-node-go/pkg/mlsvalidate"
 	"github.com/xmtp/xmtp-node-go/pkg/ratelimiter"
 	"github.com/xmtp/xmtp-node-go/pkg/store"
 	"go.uber.org/zap"
@@ -31,11 +32,12 @@ type Options struct {
 
 type Config struct {
 	Options
-	AllowLister authz.WalletAllowLister
-	Waku        *wakunode.WakuNode
-	Log         *zap.Logger
-	Store       *store.Store
-	MlsStore    mlsstore.MlsStore
+	AllowLister  authz.WalletAllowLister
+	Waku         *wakunode.WakuNode
+	Log          *zap.Logger
+	Store        *store.Store
+	MlsStore     mlsstore.MlsStore
+	MlsValidator mlsvalidate.MlsValidationService
 }
 
 // AuthnOptions bundle command line options associated with the authn package.

pkg/api/message/v3/service.go

@@ -6,6 +6,7 @@ import (
 	wakunode "github.com/waku-org/go-waku/waku/v2/node"
 	proto "github.com/xmtp/proto/v3/go/message_api/v3"
 	"github.com/xmtp/xmtp-node-go/pkg/mlsstore"
+	"github.com/xmtp/xmtp-node-go/pkg/mlsvalidate"
 	"github.com/xmtp/xmtp-node-go/pkg/store"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
@@ -16,21 +17,23 @@ import (
 type Service struct {
 	proto.UnimplementedMlsApiServer
 
-	log          *zap.Logger
-	waku         *wakunode.WakuNode
-	messageStore *store.Store
-	mlsStore     mlsstore.MlsStore
+	log               *zap.Logger
+	waku              *wakunode.WakuNode
+	messageStore      *store.Store
+	mlsStore          mlsstore.MlsStore
+	validationService mlsvalidate.MlsValidationService
 
 	ctx       context.Context
 	ctxCancel func()
 }
 
-func NewService(node *wakunode.WakuNode, logger *zap.Logger, messageStore *store.Store, mlsStore mlsstore.MlsStore) (s *Service, err error) {
+func NewService(node *wakunode.WakuNode, logger *zap.Logger, messageStore *store.Store, mlsStore mlsstore.MlsStore, validationService mlsvalidate.MlsValidationService) (s *Service, err error) {
 	s = &Service{
-		log:          logger.Named("message/v3"),
-		waku:         node,
-		messageStore: messageStore,
-		mlsStore:     mlsStore,
+		log:               logger.Named("message/v3"),
+		waku:              node,
+		messageStore:      messageStore,
+		mlsStore:          mlsStore,
+		validationService: validationService,
 	}
 
 	s.ctx, s.ctxCancel = context.WithCancel(context.Background())
@@ -45,11 +48,50 @@ func (s *Service) Close() {
 }
 
 func (s *Service) RegisterInstallation(ctx context.Context, req *proto.RegisterInstallationRequest) (*proto.RegisterInstallationResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
+	results, err := s.validationService.ValidateKeyPackages(ctx, [][]byte{req.LastResortKeyPackage.KeyPackageTlsSerialized})
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid identity: %s", err)
+	}
+	if len(results) != 1 {
+		return nil, status.Errorf(codes.Internal, "unexpected number of results: %d", len(results))
+	}
+
+	installationId := results[0].InstallationId
+	walletAddress := results[0].WalletAddress
+
+	err = s.mlsStore.CreateInstallation(ctx, installationId, walletAddress, req.LastResortKeyPackage.KeyPackageTlsSerialized)
+	if err != nil {
+		return nil, err
+	}
+
+	return &proto.RegisterInstallationResponse{
+		InstallationId: installationId,
+	}, nil
 }
 
 func (s *Service) ConsumeKeyPackages(ctx context.Context, req *proto.ConsumeKeyPackagesRequest) (*proto.ConsumeKeyPackagesResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
+	ids := req.InstallationIds
+	keyPackages, err := s.mlsStore.ConsumeKeyPackages(ctx, ids)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to consume key packages: %s", err)
+	}
+
+	resPackages := make([]*proto.ConsumeKeyPackagesResponse_KeyPackage, len(keyPackages))
+	for _, keyPackage := range keyPackages {
+		// Return the key packages in the original order
+		targetIndex := indexOf(keyPackage.InstallationId, ids)
+		if targetIndex == -1 {
+			return nil, status.Errorf(codes.Internal, "could not find key package for installation")
+		}
+
+		resPackages[targetIndex] = &proto.ConsumeKeyPackagesResponse_KeyPackage{
+			KeyPackageTlsSerialized: keyPackage.Data,
+		}
+	}
+
+	return &proto.ConsumeKeyPackagesResponse{
+		KeyPackages: resPackages,
+	}, nil
 }
 
 func (s *Service) PublishToGroup(ctx context.Context, req *proto.PublishToGroupRequest) (*emptypb.Empty, error) {
@@ -61,7 +103,27 @@ func (s *Service) PublishWelcomes(ctx context.Context, req *proto.PublishWelcome
 }
 
 func (s *Service) UploadKeyPackages(ctx context.Context, req *proto.UploadKeyPackagesRequest) (*emptypb.Empty, error) {
-	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
+	keyPackageBytes := make([][]byte, len(req.KeyPackages))
+	for i, keyPackage := range req.KeyPackages {
+		keyPackageBytes[i] = keyPackage.KeyPackageTlsSerialized
+	}
+	validationResults, err := s.validationService.ValidateKeyPackages(ctx, keyPackageBytes)
+	if err != nil {
+		// TODO: Differentiate between validation errors and internal errors
+		return nil, status.Errorf(codes.InvalidArgument, "invalid identity: %s", err)
+	}
+
+	keyPackageModels := make([]*mlsstore.KeyPackage, len(validationResults))
+	for i, validationResult := range validationResults {
+		kp := mlsstore.NewKeyPackage(validationResult.InstallationId, keyPackageBytes[i], false)
+		keyPackageModels[i] = &kp
+	}
+	err = s.mlsStore.InsertKeyPackages(ctx, keyPackageModels)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to insert key packages: %s", err)
+	}
+
+	return &emptypb.Empty{}, nil
 }
 
 func (s *Service) RevokeInstallation(ctx context.Context, req *proto.RevokeInstallationRequest) (*emptypb.Empty, error) {
@@ -71,3 +133,13 @@ func (s *Service) RevokeInstallation(ctx context.Context, req *proto.RevokeInsta
 func (s *Service) GetIdentityUpdates(ctx context.Context, req *proto.GetIdentityUpdatesRequest) (*proto.GetIdentityUpdatesResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "unimplemented")
 }
+
+func indexOf(target string, ids []string) int {
+	for i, id := range ids {
+		if id == target {
+			return i
+		}
+	}
+
+	return -1
+}

pkg/api/server.go

@@ -130,8 +130,8 @@ func (s *Server) startGRPC() error {
 	proto.RegisterMessageApiServer(grpcServer, s.messagev1)
 
 	// Enable the MLS server if a store is provided
-	if s.Config.MlsStore != nil && s.Config.EnableMls {
-		s.messagev3, err = messagev3.NewService(s.Waku, s.Log, s.Store, s.Config.MlsStore)
+	if s.Config.MlsStore != nil && s.Config.MlsValidator != nil && s.Config.EnableMls {
+		s.messagev3, err = messagev3.NewService(s.Waku, s.Log, s.Store, s.Config.MlsStore, s.Config.MlsValidator)
 		if err != nil {
 			return errors.Wrap(err, "creating mls service")
 		}

pkg/migrations/mls/20231023050806_init-schema.down.sql

@@ -0,0 +1,8 @@
+SET
+    statement_timeout = 0;
+
+--bun:split
+DROP TABLE IF EXISTS installations;
+
+--bun:split
+DROP TABLE IF EXISTS key_packages;

pkg/migrations/mls/20231023050806_init-schema.up.sql

@@ -0,0 +1,44 @@
+SET
+    statement_timeout = 0;
+
+--bun:split
+CREATE TABLE installations (
+    id TEXT PRIMARY KEY,
+    wallet_address TEXT NOT NULL,
+    created_at BIGINT NOT NULL,
+    revoked_at BIGINT
+);
+
+--bun:split
+CREATE TABLE key_packages (
+    id TEXT PRIMARY KEY,
+    installation_id TEXT NOT NULL,
+    created_at BIGINT NOT NULL,
+    consumed_at BIGINT,
+    is_last_resort BOOLEAN NOT NULL,
+    data BYTEA NOT NULL,
+    -- Add a foreign key constraint to ensure key packages cannot be added for unregistered installations
+    CONSTRAINT fk_installation_id FOREIGN KEY (installation_id) REFERENCES installations (id)
+);
+
+--bun:split
+CREATE INDEX idx_installations_wallet_address ON installations(wallet_address);
+
+--bun:split
+CREATE INDEX idx_installations_created_at ON installations(created_at);
+
+--bun:split
+CREATE INDEX idx_installations_revoked_at ON installations(revoked_at);
+
+--bun:split
+-- Adding indexes for the key_packages table
+CREATE INDEX idx_key_packages_installation_id ON key_packages(installation_id);
+
+--bun:split
+CREATE INDEX idx_key_packages_created_at ON key_packages(created_at);
+
+--bun:split
+CREATE INDEX idx_key_packages_consumed_at ON key_packages(consumed_at);
+
+--bun:split
+CREATE INDEX idx_key_packages_is_last_resort ON key_packages(is_last_resort);

pkg/migrations/mls/migrations.go

@@ -0,0 +1,18 @@
+package mls
+
+import (
+	"embed"
+
+	"github.com/uptrace/bun/migrate"
+)
+
+var Migrations = migrate.NewMigrations()
+
+//go:embed *.sql
+var sqlMigrations embed.FS
+
+func init() {
+	if err := Migrations.Discover(sqlMigrations); err != nil {
+		panic(err)
+	}
+}

pkg/mlsstore/models.go

@@ -0,0 +1,23 @@
+package mlsstore
+
+import "github.com/uptrace/bun"
+
+type Installation struct {
+	bun.BaseModel `bun:"table:installations"`
+
+	ID            string `bun:",pk"`
+	WalletAddress string `bun:"wallet_address,notnull"`
+	CreatedAt     int64  `bun:"created_at,notnull"`
+	RevokedAt     *int64 `bun:"revoked_at"`
+}
+
+type KeyPackage struct {
+	bun.BaseModel `bun:"table:key_packages"`
+
+	ID             string `bun:",pk"` // ID is the hash of the data field
+	InstallationId string `bun:"installation_id,notnull"`
+	CreatedAt      int64  `bun:"created_at,notnull"`
+	ConsumedAt     *int64 `bun:"consumed_at"`
+	IsLastResort   bool   `bun:"is_last_resort,notnull"`
+	Data           []byte `bun:"data,notnull,type:bytea"`
+}