Create unique token per proxy

2026-10-30 00:09:07 +01:00 · 2026-02-12 15:48:35 +08:00
parent 38db42e7d6
commit 381260911b
6 changed files with 226 additions and 27 deletions
@@ -345,14 +345,19 @@ func (s *ProxyServiceServer) SendAccessLog(ctx context.Context, req *proto.SendA
 }

 // SendServiceUpdate broadcasts a service update to all connected proxy servers.
-// Management should call this when services are created/updated/removed
+// Management should call this when services are created/updated/removed.
+// For create/update operations a unique one-time auth token is generated per
+// proxy so that every replica can independently authenticate with management.
 func (s *ProxyServiceServer) SendServiceUpdate(update *proto.ProxyMapping) {
-	// Send it to all connected proxy servers
 	log.Debugf("Broadcasting service update to all connected proxy servers")
 	s.connectedProxies.Range(func(key, value interface{}) bool {
 		conn := value.(*proxyConnection)
+		msg := s.perProxyMessage(update, conn.proxyID)
+		if msg == nil {
+			return true
+		}
 		select {
-		case conn.sendChan <- update:
+		case conn.sendChan <- msg:
 			log.Debugf("Sent service update with id %s to proxy server %s", update.Id, conn.proxyID)
 		default:
 			log.Warnf("Failed to send service update to proxy server %s (channel full)", conn.proxyID)
@@ -420,6 +425,8 @@ func (s *ProxyServiceServer) removeFromCluster(clusterAddr, proxyID string) {

 // SendServiceUpdateToCluster sends a service update to all proxy servers in a specific cluster.
 // If clusterAddr is empty, broadcasts to all connected proxy servers (backward compatibility).
+// For create/update operations a unique one-time auth token is generated per
+// proxy so that every replica can independently authenticate with management.
 func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMapping, clusterAddr string) {
 	if clusterAddr == "" {
 		s.SendServiceUpdate(update)
@@ -437,8 +444,12 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMappi
 		proxyID := key.(string)
 		if connVal, ok := s.connectedProxies.Load(proxyID); ok {
 			conn := connVal.(*proxyConnection)
+			msg := s.perProxyMessage(update, proxyID)
+			if msg == nil {
+				return true
+			}
 			select {
-			case conn.sendChan <- update:
+			case conn.sendChan <- msg:
 				log.Debugf("Sent service update with id %s to proxy %s in cluster %s", update.Id, proxyID, clusterAddr)
 			default:
 				log.Warnf("Failed to send service update to proxy %s in cluster %s (channel full)", proxyID, clusterAddr)
@@ -448,6 +459,42 @@ func (s *ProxyServiceServer) SendServiceUpdateToCluster(update *proto.ProxyMappi
 	})
 }

+// perProxyMessage returns a copy of update with a fresh one-time token for
+// create/update operations. For delete operations the original message is
+// returned unchanged because proxies do not need to authenticate for removal.
+// Returns nil if token generation fails (the proxy should be skipped).
+func (s *ProxyServiceServer) perProxyMessage(update *proto.ProxyMapping, proxyID string) *proto.ProxyMapping {
+	if update.Type == proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED || update.AccountId == "" {
+		return update
+	}
+
+	token, err := s.tokenStore.GenerateToken(update.AccountId, update.Id, 5*time.Minute)
+	if err != nil {
+		log.Warnf("Failed to generate token for proxy %s: %v", proxyID, err)
+		return nil
+	}
+
+	msg := shallowCloneMapping(update)
+	msg.AuthToken = token
+	return msg
+}
+
+// shallowCloneMapping creates a shallow copy of a ProxyMapping, reusing the
+// same slice/pointer fields. Only scalar fields that differ per proxy (AuthToken)
+// should be set on the copy.
+func shallowCloneMapping(m *proto.ProxyMapping) *proto.ProxyMapping {
+	return &proto.ProxyMapping{
+		Type:             m.Type,
+		Id:               m.Id,
+		AccountId:        m.AccountId,
+		Domain:           m.Domain,
+		Path:             m.Path,
+		Auth:             m.Auth,
+		PassHostHeader:   m.PassHostHeader,
+		RewriteRedirects: m.RewriteRedirects,
+	}
+}
+
 // GetAvailableClusters returns information about all connected proxy clusters.
 func (s *ProxyServiceServer) GetAvailableClusters() []ClusterInfo {
 	clusterCounts := make(map[string]int)
@@ -0,0 +1,164 @@
+package grpc
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/shared/management/proto"
+)
+
+// registerFakeProxy adds a fake proxy connection to the server's internal maps
+// and returns the channel where messages will be received.
+func registerFakeProxy(s *ProxyServiceServer, proxyID, clusterAddr string) chan *proto.ProxyMapping {
+	ch := make(chan *proto.ProxyMapping, 10)
+	conn := &proxyConnection{
+		proxyID:  proxyID,
+		address:  clusterAddr,
+		sendChan: ch,
+	}
+	s.connectedProxies.Store(proxyID, conn)
+
+	proxySet, _ := s.clusterProxies.LoadOrStore(clusterAddr, &sync.Map{})
+	proxySet.(*sync.Map).Store(proxyID, struct{}{})
+
+	return ch
+}
+
+func drainChannel(ch chan *proto.ProxyMapping) *proto.ProxyMapping {
+	select {
+	case msg := <-ch:
+		return msg
+	case <-time.After(time.Second):
+		return nil
+	}
+}
+
+func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
+	tokenStore := NewOneTimeTokenStore(time.Hour)
+	defer tokenStore.Close()
+
+	s := &ProxyServiceServer{
+		tokenStore:  tokenStore,
+		updatesChan: make(chan *proto.ProxyMapping, 100),
+	}
+
+	const cluster = "proxy.example.com"
+	const numProxies = 3
+
+	channels := make([]chan *proto.ProxyMapping, numProxies)
+	for i := range numProxies {
+		id := "proxy-" + string(rune('a'+i))
+		channels[i] = registerFakeProxy(s, id, cluster)
+	}
+
+	update := &proto.ProxyMapping{
+		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
+		Id:        "service-1",
+		AccountId: "account-1",
+		Domain:    "test.example.com",
+		Path: []*proto.PathMapping{
+			{Path: "/", Target: "http://10.0.0.1:8080/"},
+		},
+	}
+
+	s.SendServiceUpdateToCluster(update, cluster)
+
+	tokens := make([]string, numProxies)
+	for i, ch := range channels {
+		msg := drainChannel(ch)
+		require.NotNil(t, msg, "proxy %d should receive a message", i)
+		assert.Equal(t, update.Domain, msg.Domain)
+		assert.Equal(t, update.Id, msg.Id)
+		assert.NotEmpty(t, msg.AuthToken, "proxy %d should have a non-empty token", i)
+		tokens[i] = msg.AuthToken
+	}
+
+	// All tokens must be unique
+	tokenSet := make(map[string]struct{})
+	for i, tok := range tokens {
+		_, exists := tokenSet[tok]
+		assert.False(t, exists, "proxy %d got duplicate token", i)
+		tokenSet[tok] = struct{}{}
+	}
+
+	// Each token must be independently consumable
+	for i, tok := range tokens {
+		err := tokenStore.ValidateAndConsume(tok, "account-1", "service-1")
+		assert.NoError(t, err, "proxy %d token should validate successfully", i)
+	}
+}
+
+func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
+	tokenStore := NewOneTimeTokenStore(time.Hour)
+	defer tokenStore.Close()
+
+	s := &ProxyServiceServer{
+		tokenStore:  tokenStore,
+		updatesChan: make(chan *proto.ProxyMapping, 100),
+	}
+
+	const cluster = "proxy.example.com"
+	ch1 := registerFakeProxy(s, "proxy-a", cluster)
+	ch2 := registerFakeProxy(s, "proxy-b", cluster)
+
+	update := &proto.ProxyMapping{
+		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_REMOVED,
+		Id:        "service-1",
+		AccountId: "account-1",
+		Domain:    "test.example.com",
+	}
+
+	s.SendServiceUpdateToCluster(update, cluster)
+
+	msg1 := drainChannel(ch1)
+	msg2 := drainChannel(ch2)
+	require.NotNil(t, msg1)
+	require.NotNil(t, msg2)
+
+	// Delete operations should not generate tokens
+	assert.Empty(t, msg1.AuthToken)
+	assert.Empty(t, msg2.AuthToken)
+
+	// No tokens should have been created
+	assert.Equal(t, 0, tokenStore.GetTokenCount())
+}
+
+func TestSendServiceUpdate_UniqueTokensPerProxy(t *testing.T) {
+	tokenStore := NewOneTimeTokenStore(time.Hour)
+	defer tokenStore.Close()
+
+	s := &ProxyServiceServer{
+		tokenStore:  tokenStore,
+		updatesChan: make(chan *proto.ProxyMapping, 100),
+	}
+
+	// Register proxies in different clusters (SendServiceUpdate broadcasts to all)
+	ch1 := registerFakeProxy(s, "proxy-a", "cluster-a")
+	ch2 := registerFakeProxy(s, "proxy-b", "cluster-b")
+
+	update := &proto.ProxyMapping{
+		Type:      proto.ProxyMappingUpdateType_UPDATE_TYPE_CREATED,
+		Id:        "service-1",
+		AccountId: "account-1",
+		Domain:    "test.example.com",
+	}
+
+	s.SendServiceUpdate(update)
+
+	msg1 := drainChannel(ch1)
+	msg2 := drainChannel(ch2)
+	require.NotNil(t, msg1)
+	require.NotNil(t, msg2)
+
+	assert.NotEmpty(t, msg1.AuthToken)
+	assert.NotEmpty(t, msg2.AuthToken)
+	assert.NotEqual(t, msg1.AuthToken, msg2.AuthToken, "tokens must be unique per proxy")
+
+	// Both tokens should validate
+	assert.NoError(t, tokenStore.ValidateAndConsume(msg1.AuthToken, "account-1", "service-1"))
+	assert.NoError(t, tokenStore.ValidateAndConsume(msg2.AuthToken, "account-1", "service-1"))
+}