Nats integration (#137)

* Add TLS config creation with Ed25519 key validation for NATS

* Enhance concurrency safety with Replace() and Keys()

* Security improvement, more deterministic Constant-Time comparison comparison (only 0 or 1 is possible)

* Add tests to cover concurrency improvements and coverage for Key management logic

Author: namikmesic

rpc/mtls/mtls.go

@@ -45,6 +45,20 @@ func NewTransportSigner(signer crypto.Signer, pubKeys []ed25519.PublicKey) (cred
 	return credentials.NewTLS(c), nil
 }
 
+func NewTLSConfig(privKey ed25519.PrivateKey, pubKeys []ed25519.PublicKey) (*tls.Config, error) {
+	priv, err := ValidPrivateKeyFromEd25519(privKey)
+	if err != nil {
+		return nil, err
+	}
+
+	pubs, err := ValidPublicKeysFromEd25519(pubKeys...)
+	if err != nil {
+		return nil, err
+	}
+
+	return newMutualTLSConfig(priv.key, pubs)
+}
+
 // newMutualTLSConfig uses the private key and public keys to construct a mutual
 // TLS 1.3 config.
 //
@@ -117,6 +131,9 @@ type PublicKeys struct {
 }
 
 func ValidPublicKeysFromEd25519(keys ...ed25519.PublicKey) (*PublicKeys, error) {
+	if len(keys) == 0 {
+		return nil, errors.New("no public keys provided")
+	}
 	for _, key := range keys {
 		if len(key) != ed25519.PublicKeySize {
 			return nil, fmt.Errorf("invalid key length: %d, expected: %d", len(key), ed25519.PublicKeySize)
@@ -129,7 +146,13 @@ func ValidPublicKeysFromEd25519(keys ...ed25519.PublicKey) (*PublicKeys, error)
 }
 
 func (r *PublicKeys) Keys() []ed25519.PublicKey {
-	return r.keys
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	// Return a copy to prevent race conditions
+	keysCopy := make([]ed25519.PublicKey, len(r.keys))
+	copy(keysCopy, r.keys)
+	return keysCopy
 }
 
 // Verifies that the certificate's public key matches with one of the keys in
@@ -160,17 +183,22 @@ func (r *PublicKeys) VerifyPeerCertificate() func(rawCerts [][]byte, verifiedCha
 // Replace replaces the existing keys with new keys. Use this to dynamically
 // update the allowable keys at runtime.
 func (r *PublicKeys) Replace(pubs *PublicKeys) {
+	pubs.mu.RLock()
+	newKeys := make([]ed25519.PublicKey, len(pubs.keys))
+	copy(newKeys, pubs.keys)
+	pubs.mu.RUnlock()
+
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	r.keys = pubs.keys
+	r.keys = newKeys
 }
 
 // isValidPublicKey checks the public key against a list of valid keys.
 func (r *PublicKeys) isValidPublicKey(pub ed25519.PublicKey) bool {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
 	for _, vpub := range r.keys {
-		if subtle.ConstantTimeCompare(pub, vpub) > 0 {
+		if subtle.ConstantTimeCompare(pub, vpub) == 1 {
 			return true
 		}
 	}

rpc/mtls/mtls_test.go

@@ -6,7 +6,9 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"math/big"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -154,3 +156,192 @@ func Test_NewPublicKeys(t *testing.T) {
 		require.Error(t, err)
 	})
 }
+
+func Test_NewTLSConfig(t *testing.T) {
+	t.Run("nil_arguments", func(t *testing.T) {
+		cfg, err := NewTLSConfig(nil, nil)
+		require.Error(t, err)
+		assert.Nil(t, cfg)
+	})
+
+	t.Run("nil_private_key", func(t *testing.T) {
+		pub, _, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		cfg, err := NewTLSConfig(nil, []ed25519.PublicKey{pub})
+		require.Error(t, err)
+		assert.Nil(t, cfg)
+	})
+
+	t.Run("nil_public_keys", func(t *testing.T) {
+		_, priv, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		cfg, err := NewTLSConfig(priv, nil)
+		require.Error(t, err)
+		assert.Nil(t, cfg)
+	})
+
+	t.Run("empty_public_keys", func(t *testing.T) {
+		_, priv, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		cfg, err := NewTLSConfig(priv, []ed25519.PublicKey{})
+		require.Error(t, err)
+		assert.Nil(t, cfg)
+	})
+
+	t.Run("valid_single_key", func(t *testing.T) {
+		pub, priv, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		cfg, err := NewTLSConfig(priv, []ed25519.PublicKey{pub})
+		require.NoError(t, err)
+		assert.NotNil(t, cfg)
+		assert.Equal(t, uint16(tls.VersionTLS13), cfg.MinVersion)
+		assert.Equal(t, uint16(tls.VersionTLS13), cfg.MaxVersion)
+		assert.True(t, cfg.InsecureSkipVerify)
+		assert.Len(t, cfg.Certificates, 1)
+	})
+
+	t.Run("valid_multiple_keys", func(t *testing.T) {
+		pub1, _, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		pub2, _, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		_, priv, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		cfg, err := NewTLSConfig(priv, []ed25519.PublicKey{pub1, pub2})
+		require.NoError(t, err)
+		assert.NotNil(t, cfg)
+		assert.Equal(t, uint16(tls.VersionTLS13), cfg.MinVersion)
+		assert.Equal(t, uint16(tls.VersionTLS13), cfg.MaxVersion)
+		assert.True(t, cfg.InsecureSkipVerify)
+		assert.Len(t, cfg.Certificates, 1)
+	})
+
+	t.Run("invalid_public_key_length", func(t *testing.T) {
+		_, priv, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err)
+
+		// Create an invalid public key with wrong length
+		invalidPub := make([]byte, ed25519.PublicKeySize+1)
+		cfg, err := NewTLSConfig(priv, []ed25519.PublicKey{invalidPub})
+		require.Error(t, err)
+		assert.Nil(t, cfg)
+	})
+}
+
+func Test_PublicKeys_Keys(t *testing.T) {
+	// Create keys
+	pub1, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pub2, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	// Create PublicKeys with both keys
+	pks, err := ValidPublicKeysFromEd25519(pub1, pub2)
+	require.NoError(t, err)
+
+	// Get a copy of the keys
+	keysCopy := pks.Keys()
+	require.Equal(t, 2, len(keysCopy))
+
+	// Verify the keys match
+	assert.ElementsMatch(t, []ed25519.PublicKey{pub1, pub2}, keysCopy)
+
+	// Check original is unaffected
+	keysAfter := pks.Keys()
+	require.Equal(t, 2, len(keysAfter))
+}
+
+func Test_PublicKeys_Replace(t *testing.T) {
+	// Create original keys
+	pub1, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pks1, err := ValidPublicKeysFromEd25519(pub1)
+	require.NoError(t, err)
+
+	// Create replacement keys
+	pub2, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pks2, err := ValidPublicKeysFromEd25519(pub2)
+	require.NoError(t, err)
+
+	// Replace keys
+	pks1.Replace(pks2)
+
+	// Verify the replacement worked
+	assert.False(t, pks1.isValidPublicKey(pub1))
+	assert.True(t, pks1.isValidPublicKey(pub2))
+
+	// Modify the source after replace (shouldn't affect replaced keys)
+	pub3, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pks2.Replace(&PublicKeys{keys: []ed25519.PublicKey{pub3}})
+
+	// Original replacement should be unaffected
+	assert.True(t, pks1.isValidPublicKey(pub2))
+	assert.False(t, pks1.isValidPublicKey(pub3))
+}
+
+func Test_PublicKeys_Concurrency(t *testing.T) {
+	pub1, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pub2, _, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+
+	pks, err := ValidPublicKeysFromEd25519(pub1)
+	require.NoError(t, err)
+
+	// Simulate concurrent reads and writes
+	var wg sync.WaitGroup
+	start := make(chan struct{})
+
+	// Add multiple concurrent readers
+	for i := 0; i < 5; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-start
+			for j := 0; j < 10; j++ {
+				_ = pks.isValidPublicKey(pub1)
+				_ = pks.Keys()
+				time.Sleep(time.Millisecond)
+			}
+		}()
+	}
+
+	// Add concurrent writers
+	for i := 0; i < 3; i++ {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			<-start
+
+			// Alternate between pub1 and pub2
+			for j := 0; j < 5; j++ {
+				if (idx+j)%2 == 0 {
+					pks.Replace(&PublicKeys{keys: []ed25519.PublicKey{pub1}})
+				} else {
+					pks.Replace(&PublicKeys{keys: []ed25519.PublicKey{pub2}})
+				}
+				time.Sleep(time.Millisecond * 2)
+			}
+		}(i)
+	}
+
+	// Start all goroutines
+	close(start)
+	wg.Wait()
+
+	// No assertion needed - if there are no race conditions, the test passes
+}