rewrite peer creation and network map calc [WIP]

2026-10-20 20:39:08 +02:00 · 2026-02-04 20:01:00 +01:00
parent 8fafde614a
commit 096d4ac529
19 changed files with 974 additions and 284 deletions
@@ -0,0 +1,167 @@
+package grpc
+
+import (
+	"crypto/rand"
+	"crypto/subtle"
+	"encoding/base64"
+	"fmt"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// OneTimeTokenStore manages short-lived, single-use authentication tokens
+// for proxy-to-management RPC authentication. Tokens are generated when
+// a reverse proxy is created and must be used exactly once by the proxy
+// to authenticate a subsequent RPC call.
+type OneTimeTokenStore struct {
+	tokens      map[string]*tokenMetadata
+	mu          sync.RWMutex
+	cleanup     *time.Ticker
+	cleanupDone chan struct{}
+}
+
+// tokenMetadata stores information about a one-time token
+type tokenMetadata struct {
+	ReverseProxyID string
+	AccountID      string
+	ExpiresAt      time.Time
+	CreatedAt      time.Time
+}
+
+// NewOneTimeTokenStore creates a new token store with automatic cleanup
+// of expired tokens. The cleanupInterval determines how often expired
+// tokens are removed from memory.
+func NewOneTimeTokenStore(cleanupInterval time.Duration) *OneTimeTokenStore {
+	store := &OneTimeTokenStore{
+		tokens:      make(map[string]*tokenMetadata),
+		cleanup:     time.NewTicker(cleanupInterval),
+		cleanupDone: make(chan struct{}),
+	}
+
+	// Start background cleanup goroutine
+	go store.cleanupExpired()
+
+	return store
+}
+
+// GenerateToken creates a new cryptographically secure one-time token
+// with the specified TTL. The token is associated with a specific
+// accountID and reverseProxyID for validation purposes.
+//
+// Returns the generated token string or an error if random generation fails.
+func (s *OneTimeTokenStore) GenerateToken(accountID, reverseProxyID string, ttl time.Duration) (string, error) {
+	// Generate 32 bytes (256 bits) of cryptographically secure random data
+	randomBytes := make([]byte, 32)
+	if _, err := rand.Read(randomBytes); err != nil {
+		return "", fmt.Errorf("failed to generate random token: %w", err)
+	}
+
+	// Encode as URL-safe base64 for easy transmission in gRPC
+	token := base64.URLEncoding.EncodeToString(randomBytes)
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.tokens[token] = &tokenMetadata{
+		ReverseProxyID: reverseProxyID,
+		AccountID:      accountID,
+		ExpiresAt:      time.Now().Add(ttl),
+		CreatedAt:      time.Now(),
+	}
+
+	log.Debugf("Generated one-time token for proxy %s in account %s (expires in %s)",
+		reverseProxyID, accountID, ttl)
+
+	return token, nil
+}
+
+// ValidateAndConsume verifies the token against the provided accountID and
+// reverseProxyID, checks expiration, and then deletes it to enforce single-use.
+//
+// This method uses constant-time comparison to prevent timing attacks.
+//
+// Returns nil on success, or an error if:
+// - Token doesn't exist
+// - Token has expired
+// - Account ID doesn't match
+// - Reverse proxy ID doesn't match
+func (s *OneTimeTokenStore) ValidateAndConsume(token, accountID, reverseProxyID string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	metadata, exists := s.tokens[token]
+	if !exists {
+		log.Warnf("Token validation failed: token not found (proxy: %s, account: %s)",
+			reverseProxyID, accountID)
+		return fmt.Errorf("invalid token")
+	}
+
+	// Check expiration
+	if time.Now().After(metadata.ExpiresAt) {
+		delete(s.tokens, token)
+		log.Warnf("Token validation failed: token expired (proxy: %s, account: %s)",
+			reverseProxyID, accountID)
+		return fmt.Errorf("token expired")
+	}
+
+	// Validate account ID using constant-time comparison (prevents timing attacks)
+	if subtle.ConstantTimeCompare([]byte(metadata.AccountID), []byte(accountID)) != 1 {
+		log.Warnf("Token validation failed: account ID mismatch (expected: %s, got: %s)",
+			metadata.AccountID, accountID)
+		return fmt.Errorf("account ID mismatch")
+	}
+
+	// Validate reverse proxy ID using constant-time comparison
+	if subtle.ConstantTimeCompare([]byte(metadata.ReverseProxyID), []byte(reverseProxyID)) != 1 {
+		log.Warnf("Token validation failed: reverse proxy ID mismatch (expected: %s, got: %s)",
+			metadata.ReverseProxyID, reverseProxyID)
+		return fmt.Errorf("reverse proxy ID mismatch")
+	}
+
+	// Delete token immediately to enforce single-use
+	delete(s.tokens, token)
+
+	log.Infof("Token validated and consumed for proxy %s in account %s",
+		reverseProxyID, accountID)
+
+	return nil
+}
+
+// cleanupExpired removes expired tokens in the background to prevent memory leaks
+func (s *OneTimeTokenStore) cleanupExpired() {
+	for {
+		select {
+		case <-s.cleanup.C:
+			s.mu.Lock()
+			now := time.Now()
+			removed := 0
+			for token, metadata := range s.tokens {
+				if now.After(metadata.ExpiresAt) {
+					delete(s.tokens, token)
+					removed++
+				}
+			}
+			if removed > 0 {
+				log.Debugf("Cleaned up %d expired one-time tokens", removed)
+			}
+			s.mu.Unlock()
+		case <-s.cleanupDone:
+			return
+		}
+	}
+}
+
+// Close stops the cleanup goroutine and releases resources
+func (s *OneTimeTokenStore) Close() {
+	s.cleanup.Stop()
+	close(s.cleanupDone)
+}
+
+// GetTokenCount returns the current number of tokens in the store (for debugging/metrics)
+func (s *OneTimeTokenStore) GetTokenCount() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return len(s.tokens)
+}
@@ -21,12 +21,11 @@ import (
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"

+	"github.com/netbirdio/netbird/management/internals/modules/peers"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs"
 	"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey"
-	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/store"
-	"github.com/netbirdio/netbird/management/server/types"
 	proxyauth "github.com/netbirdio/netbird/proxy/auth"
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
@@ -53,11 +52,6 @@ type reverseProxyManager interface {
 	SetStatus(ctx context.Context, accountID, reverseProxyID string, status reverseproxy.ProxyStatus) error
 }

-type keyStore interface {
-	GetGroupByName(ctx context.Context, groupName string, accountID string) (*types.Group, error)
-	CreateSetupKey(ctx context.Context, accountID string, keyName string, keyType types.SetupKeyType, expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool, allowExtraDNSLabels bool) (*types.SetupKey, error)
-}
-
 // ProxyServiceServer implements the ProxyService gRPC server
 type ProxyServiceServer struct {
 	proto.UnimplementedProxyServiceServer
@@ -71,15 +65,18 @@ type ProxyServiceServer struct {
 	// Store of reverse proxies
 	reverseProxyStore reverseProxyStore

-	// Store for client setup keys
-	keyStore keyStore
-
 	// Manager for access logs
 	accessLogManager accesslogs.Manager

 	// Manager for reverse proxy operations
 	reverseProxyManager reverseProxyManager

+	// Manager for peers
+	peersManager peers.Manager
+
+	// Store for one-time authentication tokens
+	tokenStore *OneTimeTokenStore
+
 	// OIDC configuration for proxy authentication
 	oidcConfig ProxyOIDCConfig

@@ -99,13 +96,14 @@ type proxyConnection struct {
 }

 // NewProxyServiceServer creates a new proxy service server
-func NewProxyServiceServer(store reverseProxyStore, keys keyStore, accessLogMgr accesslogs.Manager, oidcConfig ProxyOIDCConfig) *ProxyServiceServer {
+func NewProxyServiceServer(store reverseProxyStore, accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager) *ProxyServiceServer {
 	return &ProxyServiceServer{
 		updatesChan:       make(chan *proto.ProxyMapping, 100),
 		reverseProxyStore: store,
-		keyStore:          keys,
 		accessLogManager:  accessLogMgr,
 		oidcConfig:        oidcConfig,
+		tokenStore:        tokenStore,
+		peersManager:      peersManager,
 	}
 }

@@ -179,33 +177,14 @@ func (s *ProxyServiceServer) sendSnapshot(ctx context.Context, conn *proxyConnec
 			continue
 		}

-		group, err := s.keyStore.GetGroupByName(ctx, rp.Name, rp.AccountID)
+		// Generate one-time authentication token for each proxy in the snapshot
+		// Tokens are not persistent on the proxy, so we need to generate new ones on reconnection
+		token, err := s.tokenStore.GenerateToken(rp.AccountID, rp.ID, 5*time.Minute)
 		if err != nil {
 			log.WithFields(log.Fields{
 				"proxy":   rp.Name,
 				"account": rp.AccountID,
-			}).WithError(err).Error("Failed to get group by name")
-			continue
-		}
-
-		// TODO: should this even be here? We're running in a loop, and on each proxy, this will create a LOT of setup key entries that we currently have no way to remove.
-		key, err := s.keyStore.CreateSetupKey(ctx,
-			rp.AccountID,
-			rp.Name,
-			types.SetupKeyReusable,
-			0,
-			[]string{group.ID},
-			0,
-			activity.SystemInitiator,
-			true,
-			false,
-		)
-		if err != nil {
-			log.WithFields(log.Fields{
-				"proxy":   rp.Name,
-				"account": rp.AccountID,
-				"group":   group.ID,
-			}).WithError(err).Error("Failed to create setup key")
+			}).WithError(err).Error("Failed to generate auth token for snapshot")
 			continue
 		}

@@ -213,7 +192,7 @@ func (s *ProxyServiceServer) sendSnapshot(ctx context.Context, conn *proxyConnec
 			Mapping: []*proto.ProxyMapping{
 				rp.ToProtoMapping(
 					reverseproxy.Create, // Initial snapshot, all records are "new" for the proxy.
-					key.Key,
+					token,
 					s.GetOIDCValidationConfig(),
 				),
 			},
@@ -433,6 +412,59 @@ func protoStatusToInternal(protoStatus proto.ProxyStatus) reverseproxy.ProxyStat
 	}
 }

+// CreateProxyPeer handles proxy peer creation with one-time token authentication
+func (s *ProxyServiceServer) CreateProxyPeer(ctx context.Context, req *proto.CreateProxyPeerRequest) (*proto.CreateProxyPeerResponse, error) {
+	reverseProxyID := req.GetReverseProxyId()
+	accountID := req.GetAccountId()
+	token := req.GetToken()
+	key := req.WireguardPublicKey
+
+	log.WithFields(log.Fields{
+		"reverse_proxy_id": reverseProxyID,
+		"account_id":       accountID,
+	}).Debug("CreateProxyPeer request received")
+
+	if reverseProxyID == "" || accountID == "" || token == "" {
+		log.Warn("CreateProxyPeer: missing required fields")
+		return &proto.CreateProxyPeerResponse{
+			Success:      false,
+			ErrorMessage: strPtr("missing required fields: reverse_proxy_id, account_id, and token are required"),
+		}, nil
+	}
+
+	if err := s.tokenStore.ValidateAndConsume(token, accountID, reverseProxyID); err != nil {
+		log.WithFields(log.Fields{
+			"reverse_proxy_id": reverseProxyID,
+			"account_id":       accountID,
+		}).WithError(err).Warn("CreateProxyPeer: token validation failed")
+		return &proto.CreateProxyPeerResponse{
+			Success:      false,
+			ErrorMessage: strPtr("authentication failed: invalid or expired token"),
+		}, status.Errorf(codes.Unauthenticated, "token validation failed: %v", err)
+	}
+
+	err := s.peersManager.CreateProxyPeer(ctx, accountID, key)
+	if err != nil {
+		log.WithFields(log.Fields{
+			"reverse_proxy_id": reverseProxyID,
+			"account_id":       accountID,
+		}).WithError(err).Error("CreateProxyPeer: failed to create proxy peer")
+		return &proto.CreateProxyPeerResponse{
+			Success:      false,
+			ErrorMessage: strPtr(fmt.Sprintf("failed to create proxy peer: %v", err)),
+		}, status.Errorf(codes.Internal, "failed to create proxy peer: %v", err)
+	}
+
+	return &proto.CreateProxyPeerResponse{
+		Success: true,
+	}, nil
+}
+
+// strPtr is a helper to create a string pointer for optional proto fields
+func strPtr(s string) *string {
+	return &s
+}
+
 func (s *ProxyServiceServer) GetOIDCURL(ctx context.Context, req *proto.GetOIDCURLRequest) (*proto.GetOIDCURLResponse, error) {
 	redirectURL, err := url.Parse(req.GetRedirectUrl())
 	if err != nil {