Add fees to UnsignedOriginatorEnvelope

dev/gen-protos

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 
 rm -rf pkg/proto/**/*.pb.go pkg/proto/**/*.pb.gw.go pkg/proto/**/*.swagger.json
-if ! go tool -modfile=tools/go.mod buf generate https://github.com/xmtp/proto.git#subdir=proto,branch=main; then
+if ! go tool -modfile=tools/go.mod buf generate https://github.com/xmtp/proto.git#subdir=proto,branch=nm/add-fees; then
     echo "Failed to generate protobuf definitions"
     exit 1
 fi

pkg/api/message/publishWorker.go

@@ -6,9 +6,12 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/xmtp/xmtpd/pkg/constants"
+	"github.com/xmtp/xmtpd/pkg/currency"
 	"github.com/xmtp/xmtpd/pkg/db"
 	"github.com/xmtp/xmtpd/pkg/db/queries"
 	"github.com/xmtp/xmtpd/pkg/envelopes"
+	"github.com/xmtp/xmtpd/pkg/fees"
 	"github.com/xmtp/xmtpd/pkg/registrant"
 	"go.uber.org/zap"
 )
@@ -22,13 +25,15 @@ type publishWorker struct {
 	store         *sql.DB
 	subscription  db.DBSubscription[queries.StagedOriginatorEnvelope, int64]
 	lastProcessed atomic.Int64
+	feeCalculator fees.IFeeCalculator
 }
 
 func startPublishWorker(
 	ctx context.Context,
 	log *zap.Logger,
 	reg *registrant.Registrant,
 	store *sql.DB,
+	feeCalculator fees.IFeeCalculator,
 ) (*publishWorker, error) {
 	log = log.Named("publishWorker")
 	q := queries.New(store)
@@ -62,13 +67,14 @@ func startPublishWorker(
 	}
 
 	worker := &publishWorker{
-		ctx:          ctx,
-		log:          log,
-		notifier:     notifier,
-		subscription: *subscription,
-		listener:     listener,
-		registrant:   reg,
-		store:        store,
+		ctx:           ctx,
+		log:           log,
+		notifier:      notifier,
+		subscription:  *subscription,
+		listener:      listener,
+		registrant:    reg,
+		store:         store,
+		feeCalculator: feeCalculator,
 	}
 	go worker.start()
 
@@ -102,7 +108,13 @@ func (p *publishWorker) start() {
 
 func (p *publishWorker) publishStagedEnvelope(stagedEnv queries.StagedOriginatorEnvelope) bool {
 	logger := p.log.With(zap.Int64("sequenceID", stagedEnv.ID))
-	originatorEnv, err := p.registrant.SignStagedEnvelope(stagedEnv)
+	baseFee, congestionFee, err := p.calculateFees(&stagedEnv)
+	if err != nil {
+		logger.Error("Failed to calculate fees", zap.Error(err))
+		return false
+	}
+
+	originatorEnv, err := p.registrant.SignStagedEnvelope(stagedEnv, baseFee, congestionFee)
 	if err != nil {
 		logger.Error(
 			"Failed to sign staged envelope",
@@ -144,6 +156,7 @@ func (p *publishWorker) publishStagedEnvelope(stagedEnv queries.StagedOriginator
 			Topic:                stagedEnv.Topic,
 			OriginatorEnvelope:   originatorBytes,
 			PayerID:              db.NullInt32(payerId),
+			GatewayTime:          stagedEnv.OriginatorTime,
 		},
 	)
 	if p.ctx.Err() != nil {
@@ -171,3 +184,29 @@ func (p *publishWorker) publishStagedEnvelope(stagedEnv queries.StagedOriginator
 
 	return true
 }
+
+func (p *publishWorker) calculateFees(
+	stagedEnv *queries.StagedOriginatorEnvelope,
+) (currency.PicoDollar, currency.PicoDollar, error) {
+	baseFee, err := p.feeCalculator.CalculateBaseFee(
+		stagedEnv.OriginatorTime,
+		int64(len(stagedEnv.PayerEnvelope)),
+		constants.DEFAULT_STORAGE_DURATION_DAYS,
+	)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	// TODO:nm: Set this to the actual congestion fee
+	// For now we are setting congestion to 0
+	congestionFee, err := p.feeCalculator.CalculateCongestionFee(
+		stagedEnv.OriginatorTime,
+		0,
+	)
+
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return baseFee, congestionFee, nil
+}

pkg/api/message/publish_test.go

@@ -7,7 +7,9 @@ import (
 
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"github.com/xmtp/xmtpd/pkg/currency"
 	"github.com/xmtp/xmtpd/pkg/db/queries"
+	envelopeUtils "github.com/xmtp/xmtpd/pkg/envelopes"
 	"github.com/xmtp/xmtpd/pkg/mlsvalidate"
 	apiv1 "github.com/xmtp/xmtpd/pkg/proto/mls/api/v1"
 	"github.com/xmtp/xmtpd/pkg/proto/xmtpv4/envelopes"
@@ -294,3 +296,48 @@ func TestPublishEnvelopeOriginatorUnknown(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "DependsOn has not been seen by this node")
 }
+
+func TestPublishEnvelopeFees(t *testing.T) {
+	api, db, _, cleanup := apiTestUtils.NewTestReplicationAPIClient(t)
+	defer cleanup()
+
+	payerEnvelope := envelopeTestUtils.CreatePayerEnvelope(
+		t,
+		envelopeTestUtils.DefaultClientEnvelopeNodeId,
+	)
+
+	resp, err := api.PublishPayerEnvelopes(
+		context.Background(),
+		&message_api.PublishPayerEnvelopesRequest{
+			PayerEnvelopes: []*envelopes.PayerEnvelope{payerEnvelope},
+		},
+	)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	returnedEnv, err := envelopeUtils.NewOriginatorEnvelope(resp.GetOriginatorEnvelopes()[0])
+	require.NoError(t, err)
+	// BaseFee will always be > 0
+	require.Greater(t, returnedEnv.UnsignedOriginatorEnvelope.BaseFee(), currency.PicoDollar(0))
+	// CongestionFee will be 0 for now.
+	// TODO:nm: Set this to the actual congestion fee
+	require.Equal(t, returnedEnv.UnsignedOriginatorEnvelope.CongestionFee(), currency.PicoDollar(0))
+
+	envs, err := queries.New(db).
+		SelectGatewayEnvelopes(context.Background(), queries.SelectGatewayEnvelopesParams{})
+	require.NoError(t, err)
+	require.Equal(t, len(envs), 1)
+
+	originatorEnv, err := envelopeUtils.NewOriginatorEnvelopeFromBytes(envs[0].OriginatorEnvelope)
+	require.NoError(t, err)
+	require.Equal(
+		t,
+		originatorEnv.UnsignedOriginatorEnvelope.BaseFee(),
+		returnedEnv.UnsignedOriginatorEnvelope.BaseFee(),
+	)
+	require.Equal(
+		t,
+		originatorEnv.UnsignedOriginatorEnvelope.CongestionFee(),
+		returnedEnv.UnsignedOriginatorEnvelope.CongestionFee(),
+	)
+}

pkg/api/message/service.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
-	"github.com/xmtp/xmtpd/pkg/api/metadata"
 	"time"
 
+	"github.com/xmtp/xmtpd/pkg/api/metadata"
+	"github.com/xmtp/xmtpd/pkg/fees"
+
 	"github.com/xmtp/xmtpd/pkg/db"
 	"github.com/xmtp/xmtpd/pkg/db/queries"
 	"github.com/xmtp/xmtpd/pkg/envelopes"
@@ -42,6 +44,7 @@ type Service struct {
 	subscribeWorker   *subscribeWorker
 	validationService mlsvalidate.MLSValidationService
 	cu                metadata.CursorUpdater
+	feeCalculator     fees.IFeeCalculator
 }
 
 func NewReplicationApiService(
@@ -51,8 +54,10 @@ func NewReplicationApiService(
 	store *sql.DB,
 	validationService mlsvalidate.MLSValidationService,
 	updater metadata.CursorUpdater,
+	rateFetcher fees.IRatesFetcher,
 ) (*Service, error) {
-	publishWorker, err := startPublishWorker(ctx, log, registrant, store)
+	feeCalculator := fees.NewFeeCalculator(rateFetcher)
+	publishWorker, err := startPublishWorker(ctx, log, registrant, store, feeCalculator)
 	if err != nil {
 		return nil, err
 	}
@@ -70,6 +75,7 @@ func NewReplicationApiService(
 		subscribeWorker:   subscribeWorker,
 		validationService: validationService,
 		cu:                updater,
+		feeCalculator:     feeCalculator,
 	}, nil
 }
 
@@ -349,7 +355,12 @@ func (s *Service) PublishPayerEnvelopes(
 	}
 	s.publishWorker.notifyStagedPublish()
 
-	originatorEnv, err := s.registrant.SignStagedEnvelope(stagedEnv)
+	baseFee, congestionFee, err := s.publishWorker.calculateFees(&stagedEnv)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "could not calculate fees: %v", err)
+	}
+
+	originatorEnv, err := s.registrant.SignStagedEnvelope(stagedEnv, baseFee, congestionFee)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "could not sign envelope: %v", err)
 	}

pkg/constants/constants.go

@@ -7,6 +7,7 @@ const (
 	ORIGINATOR_DOMAIN_SEPARATION_LABEL = "originator|"
 	NODE_AUTHORIZATION_HEADER_NAME     = "node-authorization"
 	MAX_BLOCKCHAIN_ORIGINATOR_ID       = 100
+	DEFAULT_STORAGE_DURATION_DAYS      = 60
 )
 
 type VerifiedNodeRequestCtxKey struct{}

pkg/db/queries.sql

@@ -13,8 +13,8 @@ WHERE
 	singleton_id = 1;
 
 -- name: InsertGatewayEnvelope :execrows
-INSERT INTO gateway_envelopes(originator_node_id, originator_sequence_id, topic, originator_envelope, payer_id)
-	VALUES (@originator_node_id, @originator_sequence_id, @topic, @originator_envelope, @payer_id)
+INSERT INTO gateway_envelopes(originator_node_id, originator_sequence_id, topic, originator_envelope, payer_id, gateway_time)
+	VALUES (@originator_node_id, @originator_sequence_id, @topic, @originator_envelope, @payer_id, COALESCE(@gateway_time, NOW()))
 ON CONFLICT
 	DO NOTHING;
 

pkg/envelopes/unsignedOriginator.go

@@ -3,6 +3,7 @@ package envelopes
 import (
 	"errors"
 
+	"github.com/xmtp/xmtpd/pkg/currency"
 	envelopesProto "github.com/xmtp/xmtpd/pkg/proto/xmtpv4/envelopes"
 	"github.com/xmtp/xmtpd/pkg/topic"
 	"github.com/xmtp/xmtpd/pkg/utils"
@@ -45,6 +46,16 @@ func (u *UnsignedOriginatorEnvelope) OriginatorNs() int64 {
 	return u.proto.OriginatorNs
 }
 
+func (u *UnsignedOriginatorEnvelope) BaseFee() currency.PicoDollar {
+	// Skip nil check because it is in the constructor
+	return currency.PicoDollar(u.proto.BaseFeePicodollars)
+}
+
+func (u *UnsignedOriginatorEnvelope) CongestionFee() currency.PicoDollar {
+	// Skip nil check because it is in the constructor
+	return currency.PicoDollar(u.proto.CongestionFeePicodollars)
+}
+
 func NewUnsignedOriginatorEnvelopeFromBytes(bytes []byte) (*UnsignedOriginatorEnvelope, error) {
 	message, err := utils.UnmarshalUnsignedEnvelope(bytes)
 	if err != nil {