rewrite peer creation and network map calc [WIP]

management/internals/shared/grpc/proxy.go

@@ -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 {