sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-10-27 06:49:07 +01:00
Code Issues Packages Projects Releases Wiki Activity

extract proxy controller

Browse Source
This commit is contained in:
pascal
2026-02-23 15:28:20 +01:00
parent 3af287ebab
commit 740c726a78
10 changed files with 291 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
management/internals/shared/grpc/onetime_token.go
+15 -9
View File
@@ -7,9 +7,11 @@ import (
"crypto/subtle"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"time"
"github.com/eko/gocache/lib/v4/cache"
"github.com/eko/gocache/lib/v4/store"
log "github.com/sirupsen/logrus"
@@ -26,7 +28,7 @@ type tokenMetadata struct {
// OneTimeTokenStore manages single-use authentication tokens for proxy-to-management RPC.
// Supports both in-memory and Redis storage via NB_IDP_CACHE_REDIS_ADDRESS env var.
type OneTimeTokenStore struct {
store store.StoreInterface
cache *cache.Cache[string]
ctx context.Context
}
@@ -38,7 +40,7 @@ func NewOneTimeTokenStore(ctx context.Context, maxTimeout, cleanupInterval time.
}
return &OneTimeTokenStore{
store: cacheStore,
cache: cache.New[string](cacheStore),
ctx: ctx,
}, nil
}
@@ -64,7 +66,12 @@ func (s *OneTimeTokenStore) GenerateToken(accountID, serviceID string, ttl time.
CreatedAt: time.Now(),
}
if err := s.store.Set(s.ctx, hashedToken, metadata, store.WithExpiration(ttl)); err != nil {
metadataJSON, err := json.Marshal(metadata)
if err != nil {
return "", fmt.Errorf("failed to serialize token metadata: %w", err)
}
if err := s.cache.Set(s.ctx, hashedToken, string(metadataJSON), store.WithExpiration(ttl)); err != nil {
return "", fmt.Errorf("failed to store token: %w", err)
}
@@ -87,20 +94,19 @@ func (s *OneTimeTokenStore) GenerateToken(accountID, serviceID string, ttl time.
func (s *OneTimeTokenStore) ValidateAndConsume(token, accountID, serviceID string) error {
hashedToken := hashToken(token)
value, err := s.store.Get(s.ctx, hashedToken)
metadataJSON, err := s.cache.Get(s.ctx, hashedToken)
if err != nil {
log.Warnf("Token validation failed: token not found (proxy: %s, account: %s)", serviceID, accountID)
return fmt.Errorf("invalid token")
}
metadata, ok := value.(*tokenMetadata)
if !ok {
log.Warnf("Token validation failed: invalid metadata type (proxy: %s, account: %s)", serviceID, accountID)
metadata := &tokenMetadata{}
if err := json.Unmarshal([]byte(metadataJSON), metadata); err != nil {
log.Warnf("Token validation failed: failed to unmarshal metadata (proxy: %s, account: %s): %v", serviceID, accountID, err)
return fmt.Errorf("invalid token metadata")
}
if time.Now().After(metadata.ExpiresAt) {
s.store.Delete(s.ctx, hashedToken)
log.Warnf("Token validation failed: token expired (proxy: %s, account: %s)", serviceID, accountID)
return fmt.Errorf("token expired")
}
@@ -115,7 +121,7 @@ func (s *OneTimeTokenStore) ValidateAndConsume(token, accountID, serviceID strin
return fmt.Errorf("service ID mismatch")
}
if err := s.store.Delete(s.ctx, hashedToken); err != nil {
if err := s.cache.Delete(s.ctx, hashedToken); err != nil {
log.Warnf("Token deletion warning (proxy: %s, account: %s): %v", serviceID, accountID, err)
}
management/internals/shared/grpc/proxy.go
+27 -37
View File
@@ -59,9 +59,6 @@ type ProxyServiceServer struct {
// Map of connected proxies: proxy_id -> proxy connection
connectedProxies sync.Map
// Map of cluster address -> set of proxy IDs
clusterProxies sync.Map
// Channel for broadcasting reverse proxy updates to all proxies
updatesChan chan *proto.ProxyMapping
@@ -71,6 +68,9 @@ type ProxyServiceServer struct {
// Manager for reverse proxy operations
serviceManager rpservice.Manager
// ProxyController for service updates and cluster management
proxyController rpservice.ProxyController
// Manager for proxy connections
proxyManager proxy.Manager
@@ -173,6 +173,10 @@ func (s *ProxyServiceServer) SetProxyManager(manager rpservice.Manager) {
s.serviceManager = manager
}
func (s *ProxyServiceServer) SetProxyController(proxyController rpservice.ProxyController) {
s.proxyController = proxyController
}
// GetMappingUpdate handles the control stream with proxy clients
func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest, stream proto.ProxyService_GetMappingUpdateServer) error {
ctx := stream.Context()
@@ -205,7 +209,9 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
}
s.connectedProxies.Store(proxyID, conn)
s.addToCluster(conn.address, proxyID)
if err := s.proxyController.RegisterProxyToCluster(ctx, conn.address, proxyID); err != nil {
log.WithContext(ctx).Warnf("Failed to register proxy %s in cluster: %v", proxyID, err)
}
// Register proxy in database
if err := s.proxyManager.Connect(ctx, proxyID, proxyAddress, peerInfo); err != nil {
@@ -224,7 +230,9 @@ func (s *ProxyServiceServer) GetMappingUpdate(req *proto.GetMappingUpdateRequest
}
s.connectedProxies.Delete(proxyID)
s.removeFromCluster(conn.address, proxyID)
if err := s.proxyController.UnregisterProxyFromCluster(context.Background(), conn.address, proxyID); err != nil {
log.Warnf("Failed to unregister proxy %s from cluster: %v", proxyID, err)
}
cancel()
log.Infof("Proxy %s disconnected", proxyID)
@@ -446,61 +454,43 @@ func (s *ProxyServiceServer) GetConnectedProxyURLs() []string {
return urls
}
// addToCluster registers a proxy in a cluster.
func (s *ProxyServiceServer) addToCluster(clusterAddr, proxyID string) {
if clusterAddr == "" {
return
}
proxySet, _ := s.clusterProxies.LoadOrStore(clusterAddr, &sync.Map{})
proxySet.(*sync.Map).Store(proxyID, struct{}{})
log.Debugf("Added proxy %s to cluster %s", proxyID, clusterAddr)
}
// removeFromCluster removes a proxy from a cluster.
func (s *ProxyServiceServer) removeFromCluster(clusterAddr, proxyID string) {
if clusterAddr == "" {
return
}
if proxySet, ok := s.clusterProxies.Load(clusterAddr); ok {
proxySet.(*sync.Map).Delete(proxyID)
log.Debugf("Removed proxy %s from cluster %s", proxyID, clusterAddr)
}
}
// 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) {
func (s *ProxyServiceServer) SendServiceUpdateToCluster(ctx context.Context, update *proto.ProxyMapping, clusterAddr string) {
if clusterAddr == "" {
s.SendServiceUpdate(update)
return
}
proxySet, ok := s.clusterProxies.Load(clusterAddr)
if !ok {
log.Debugf("No proxies connected for cluster %s", clusterAddr)
if s.proxyController == nil {
log.WithContext(ctx).Debugf("ProxyController not set, cannot send to cluster %s", clusterAddr)
return
}
proxyIDs := s.proxyController.GetProxiesForCluster(clusterAddr)
if len(proxyIDs) == 0 {
log.WithContext(ctx).Debugf("No proxies connected for cluster %s", clusterAddr)
return
}
log.Debugf("Sending service update to cluster %s", clusterAddr)
proxySet.(*sync.Map).Range(func(key, _ interface{}) bool {
proxyID := key.(string)
for _, proxyID := range proxyIDs {
if connVal, ok := s.connectedProxies.Load(proxyID); ok {
conn := connVal.(*proxyConnection)
msg := s.perProxyMessage(update, proxyID)
if msg == nil {
return true
continue
}
select {
case conn.sendChan <- msg:
log.Debugf("Sent service update with id %s to proxy %s in cluster %s", update.Id, proxyID, clusterAddr)
log.WithContext(ctx).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)
log.WithContext(ctx).Warnf("Failed to send service update to proxy %s in cluster %s (channel full)", proxyID, clusterAddr)
}
}
return true
})
}
}
// perProxyMessage returns a copy of update with a fresh one-time token for
management/internals/shared/grpc/proxy_test.go
+2 -2
View File
@@ -69,7 +69,7 @@ func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) {
},
}
s.SendServiceUpdateToCluster(update, cluster)
s.SendServiceUpdateToCluster(context.Background(), update, cluster)
tokens := make([]string, numProxies)
for i, ch := range channels {
@@ -116,7 +116,7 @@ func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) {
Domain: "test.example.com",
}
s.SendServiceUpdateToCluster(update, cluster)
s.SendServiceUpdateToCluster(context.Background(), update, cluster)
msg1 := drainChannel(ch1)
msg2 := drainChannel(ch2)