JWT Version Checks

cmd/replication/main.go

@@ -124,6 +124,7 @@ func main() {
 			dbInstance,
 			blockchainPublisher,
 			fmt.Sprintf("0.0.0.0:%d", options.API.Port),
+			Commit,
 		)
 		if err != nil {
 			log.Fatal("initializing server", zap.Error(err))

go.mod

@@ -27,6 +27,7 @@ require (
 )
 
 require (
+	github.com/Masterminds/semver/v3 v3.1.1
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 )

pkg/authn/claims.go

@@ -0,0 +1,40 @@
+package authn
+
+import (
+	"fmt"
+	"github.com/Masterminds/semver/v3"
+	"github.com/golang-jwt/jwt/v5"
+)
+
+const (
+	// XMTPD_COMPATIBLE_VERSION_CONSTRAINT major or minor version bumps indicate backwards incompatible changes
+	XMTPD_COMPATIBLE_VERSION_CONSTRAINT = "~ 0.1.3"
+)
+
+type XmtpdClaims struct {
+	Version *string `json:"version,omitempty"`
+	jwt.RegisteredClaims
+}
+
+func ValidateVersionClaimIsCompatible(claims *XmtpdClaims) error {
+	if claims.Version == nil || *claims.Version == "" {
+		return nil
+	}
+
+	c, err := semver.NewConstraint(XMTPD_COMPATIBLE_VERSION_CONSTRAINT)
+	if err != nil {
+		return err
+	}
+
+	v, err := semver.NewVersion(*claims.Version)
+
+	if err != nil {
+		return err
+	}
+
+	if ok := c.Check(v); !ok {
+		return fmt.Errorf("version %s is not compatible", *claims.Version)
+	}
+
+	return nil
+}

pkg/authn/tokenFactory.go

@@ -15,25 +15,33 @@ const (
 type TokenFactory struct {
 	privateKey *ecdsa.PrivateKey
 	nodeID     uint32
+	version    string
 }
 
-func NewTokenFactory(privateKey *ecdsa.PrivateKey, nodeID uint32) *TokenFactory {
+func NewTokenFactory(privateKey *ecdsa.PrivateKey, nodeID uint32, version string) *TokenFactory {
 	return &TokenFactory{
 		privateKey: privateKey,
 		nodeID:     nodeID,
+		version:    version,
 	}
 }
 
 func (f *TokenFactory) CreateToken(forNodeID uint32) (*Token, error) {
 	now := time.Now()
 	expiresAt := now.Add(TOKEN_DURATION)
 
-	token := jwt.NewWithClaims(&SigningMethodSecp256k1{}, &jwt.RegisteredClaims{
-		Subject:   strconv.Itoa(int(f.nodeID)),
-		Audience:  []string{strconv.Itoa(int(forNodeID))},
-		ExpiresAt: jwt.NewNumericDate(expiresAt),
-		IssuedAt:  jwt.NewNumericDate(now),
-	})
+	claims := &XmtpdClaims{
+		Version: &f.version,
+		RegisteredClaims: jwt.RegisteredClaims{
+			Subject:   strconv.Itoa(int(f.nodeID)),
+			Audience:  []string{strconv.Itoa(int(forNodeID))},
+			ExpiresAt: jwt.NewNumericDate(expiresAt),
+			IssuedAt:  jwt.NewNumericDate(now),
+		},
+	}
+
+	// Create a new token with custom claims
+	token := jwt.NewWithClaims(&SigningMethodSecp256k1{}, claims)
 
 	signedString, err := token.SignedString(f.privateKey)
 	if err != nil {

pkg/authn/tokenFactory_test.go

@@ -10,7 +10,7 @@ import (
 
 func TestTokenFactory(t *testing.T) {
 	privateKey := testutils.RandomPrivateKey(t)
-	factory := NewTokenFactory(privateKey, 100)
+	factory := NewTokenFactory(privateKey, 100, "")
 
 	token, err := factory.CreateToken(200)
 	require.NoError(t, err)

pkg/authn/verifier.go

@@ -30,10 +30,14 @@ func NewRegistryVerifier(registry registry.NodeRegistry, myNodeID uint32) *Regis
 func (v *RegistryVerifier) Verify(tokenString string) error {
 	var token *jwt.Token
 	var err error
-	if token, err = jwt.Parse(tokenString, v.getMatchingPublicKey); err != nil {
+
+	if token, err = jwt.ParseWithClaims(
+		tokenString,
+		&XmtpdClaims{},
+		v.getMatchingPublicKey,
+	); err != nil {
 		return err
 	}
-
 	if err = v.validateAudience(token); err != nil {
 		return err
 	}
@@ -42,6 +46,10 @@ func (v *RegistryVerifier) Verify(tokenString string) error {
 		return err
 	}
 
+	if err = v.validateClaims(token); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -85,6 +93,20 @@ func (v *RegistryVerifier) validateAudience(token *jwt.Token) error {
 	return fmt.Errorf("could not find node ID in audience %v", audience)
 }
 
+func (v *RegistryVerifier) validateClaims(token *jwt.Token) error {
+	claims, ok := token.Claims.(*XmtpdClaims)
+	if !ok {
+		return fmt.Errorf("invalid token claims type")
+	}
+
+	// Check if the token is valid
+	if !token.Valid {
+		return fmt.Errorf("invalid token")
+	}
+
+	return ValidateVersionClaimIsCompatible(claims)
+}
+
 // Parse the subject claim of the JWT and return the node ID as a uint32
 func getSubjectNodeId(token *jwt.Token) (uint32, error) {
 	subject, err := token.Claims.GetSubject()

pkg/authn/verifier_test.go

@@ -53,7 +53,7 @@ func buildJwt(
 func TestVerifier(t *testing.T) {
 	signerPrivateKey := testutils.RandomPrivateKey(t)
 
-	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID))
+	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID), "")
 
 	verifier, nodeRegistry := buildVerifier(t, uint32(VERIFIER_NODE_ID))
 	nodeRegistry.EXPECT().GetNode(uint32(SIGNER_NODE_ID)).Return(&registry.Node{
@@ -79,7 +79,7 @@ func TestVerifier(t *testing.T) {
 func TestWrongAudience(t *testing.T) {
 	signerPrivateKey := testutils.RandomPrivateKey(t)
 
-	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID))
+	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID), "")
 
 	verifier, nodeRegistry := buildVerifier(t, uint32(VERIFIER_NODE_ID))
 	nodeRegistry.EXPECT().GetNode(uint32(SIGNER_NODE_ID)).Return(&registry.Node{
@@ -97,7 +97,7 @@ func TestWrongAudience(t *testing.T) {
 func TestUnknownNode(t *testing.T) {
 	signerPrivateKey := testutils.RandomPrivateKey(t)
 
-	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID))
+	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID), "")
 
 	verifier, nodeRegistry := buildVerifier(t, uint32(VERIFIER_NODE_ID))
 	nodeRegistry.EXPECT().GetNode(uint32(SIGNER_NODE_ID)).Return(nil, errors.New("node not found"))
@@ -112,7 +112,7 @@ func TestUnknownNode(t *testing.T) {
 func TestWrongPublicKey(t *testing.T) {
 	signerPrivateKey := testutils.RandomPrivateKey(t)
 
-	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID))
+	tokenFactory := NewTokenFactory(signerPrivateKey, uint32(SIGNER_NODE_ID), "")
 
 	verifier, nodeRegistry := buildVerifier(t, uint32(VERIFIER_NODE_ID))
 

pkg/interceptors/client/auth_test.go

@@ -55,7 +55,7 @@ func TestAuthInterceptor(t *testing.T) {
 	privateKey := testutils.RandomPrivateKey(t)
 	myNodeID := uint32(100)
 	targetNodeID := uint32(200)
-	tokenFactory := authn.NewTokenFactory(privateKey, myNodeID)
+	tokenFactory := authn.NewTokenFactory(privateKey, myNodeID, "")
 	interceptor := NewAuthInterceptor(tokenFactory, targetNodeID)
 	token, err := interceptor.getToken()
 	require.NoError(t, err)

pkg/registrant/registrant.go

@@ -31,6 +31,7 @@ func NewRegistrant(
 	db *queries.Queries,
 	nodeRegistry registry.NodeRegistry,
 	privateKeyString string,
+	version string,
 ) (*Registrant, error) {
 	privateKey, err := utils.ParseEcdsaPrivateKey(privateKeyString)
 	if err != nil {
@@ -47,7 +48,7 @@ func NewRegistrant(
 		return nil, err
 	}
 
-	tokenFactory := authn.NewTokenFactory(privateKey, record.NodeID)
+	tokenFactory := authn.NewTokenFactory(privateKey, record.NodeID, version)
 
 	log.Info(
 		"Registrant identified",

pkg/registrant/registrant_test.go

@@ -29,6 +29,7 @@ type deps struct {
 	privKey1Str string
 	privKey2    *ecdsa.PrivateKey
 	privKey3    *ecdsa.PrivateKey
+	version     string
 }
 
 func setup(t *testing.T) (deps, func()) {
@@ -54,6 +55,7 @@ func setup(t *testing.T) (deps, func()) {
 		privKey1Str: privKey1Str,
 		privKey2:    privKey2,
 		privKey3:    privKey3,
+		version:     "",
 	}, dbCleanup
 }
 
@@ -70,6 +72,7 @@ func setupWithRegistrant(t *testing.T) (deps, *registrant.Registrant, func()) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.NoError(t, err)
 
@@ -86,6 +89,7 @@ func TestNewRegistrantBadPrivateKey(t *testing.T) {
 		deps.db,
 		deps.registry,
 		"badkey",
+		deps.version,
 	)
 	require.ErrorContains(t, err, "parse")
 }
@@ -105,6 +109,7 @@ func TestNewRegistrantNotInRegistry(t *testing.T) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.ErrorContains(t, err, "registry")
 }
@@ -125,6 +130,7 @@ func TestNewRegistrantNewDatabase(t *testing.T) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.NoError(t, err)
 }
@@ -152,6 +158,7 @@ func TestNewRegistrantExistingDatabase(t *testing.T) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.NoError(t, err)
 }
@@ -179,6 +186,7 @@ func TestNewRegistrantMismatchingDatabaseNodeId(t *testing.T) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.ErrorContains(t, err, "does not match")
 }
@@ -206,6 +214,7 @@ func TestNewRegistrantMismatchingDatabasePublicKey(t *testing.T) {
 		deps.db,
 		deps.registry,
 		deps.privKey1Str,
+		deps.version,
 	)
 	require.ErrorContains(t, err, "does not match")
 }
@@ -224,6 +233,7 @@ func TestNewRegistrantPrivateKeyNo0x(t *testing.T) {
 		deps.db,
 		deps.registry,
 		utils.HexEncode(crypto.FromECDSA(deps.privKey1)),
+		deps.version,
 	)
 	require.NoError(t, err)
 }

pkg/server/server.go

@@ -52,6 +52,7 @@ func NewReplicationServer(
 	writerDB *sql.DB,
 	blockchainPublisher blockchain.IBlockchainPublisher,
 	listenAddress string,
+	version string,
 ) (*ReplicationServer, error) {
 	var err error
 
@@ -89,6 +90,7 @@ func NewReplicationServer(
 			queries.New(writerDB),
 			nodeRegistry,
 			options.Signer.PrivateKey,
+			version,
 		)
 		if err != nil {
 			return nil, err

pkg/server/server_test.go

@@ -65,7 +65,7 @@ func NewTestServer(
 		//Payer: config.PayerOptions{
 		//	Enable: true,
 		//},
-	}, registry, db, messagePublisher, fmt.Sprintf("localhost:%d", port))
+	}, registry, db, messagePublisher, fmt.Sprintf("localhost:%d", port), "")
 	require.NoError(t, err)
 
 	return server

pkg/testutils/api/api.go

@@ -74,7 +74,14 @@ func NewTestAPIServer(t *testing.T) (*api.ApiServer, *sql.DB, func()) {
 	mockRegistry.EXPECT().GetNodes().Return([]registry.Node{
 		{NodeID: 100, SigningKey: &privKey.PublicKey},
 	}, nil)
-	registrant, err := registrant.NewRegistrant(ctx, log, queries.New(db), mockRegistry, privKeyStr)
+	registrant, err := registrant.NewRegistrant(
+		ctx,
+		log,
+		queries.New(db),
+		mockRegistry,
+		privKeyStr,
+		"",
+	)
 	require.NoError(t, err)
 	mockMessagePublisher := blockchain.NewMockIBlockchainPublisher(t)