Add health checks

proxy/server.go

@@ -32,6 +32,7 @@ import (
 	"github.com/netbirdio/netbird/proxy/internal/acme"
 	"github.com/netbirdio/netbird/proxy/internal/auth"
 	"github.com/netbirdio/netbird/proxy/internal/debug"
+	"github.com/netbirdio/netbird/proxy/internal/health"
 	"github.com/netbirdio/netbird/proxy/internal/proxy"
 	"github.com/netbirdio/netbird/proxy/internal/roundtrip"
 	"github.com/netbirdio/netbird/proxy/internal/types"
@@ -41,14 +42,16 @@ import (
 )
 
 type Server struct {
-	mgmtClient proto.ProxyServiceClient
-	proxy      *proxy.ReverseProxy
-	netbird    *roundtrip.NetBird
-	acme       *acme.Manager
-	auth       *auth.Middleware
-	http       *http.Server
-	https      *http.Server
-	debug      *http.Server
+	mgmtClient    proto.ProxyServiceClient
+	proxy         *proxy.ReverseProxy
+	netbird       *roundtrip.NetBird
+	acme          *acme.Manager
+	auth          *auth.Middleware
+	http          *http.Server
+	https         *http.Server
+	debug         *http.Server
+	healthServer  *health.Server
+	healthChecker *health.Checker
 
 	// Mostly used for debugging on management.
 	startTime time.Time
@@ -71,6 +74,8 @@ type Server struct {
 	DebugEndpointEnabled bool
 	// DebugEndpointAddress is the address for the debug HTTP endpoint (default: ":8444").
 	DebugEndpointAddress string
+	// HealthAddress is the address for the health probe endpoint (default: "localhost:8080").
+	HealthAddress string
 }
 
 // NotifyStatus sends a status update to management about tunnel connectivity
@@ -221,7 +226,7 @@ func (s *Server) ListenAndServe(ctx context.Context, addr string) (err error) {
 			Handler: debugHandler,
 		}
 		go func() {
-			s.Logger.WithField("address", debugAddr).Info("starting debug endpoint")
+			s.Logger.Infof("starting debug endpoint on %s", debugAddr)
 			if err := s.debug.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 				s.Logger.Errorf("debug endpoint error: %v", err)
 			}
@@ -233,6 +238,30 @@ func (s *Server) ListenAndServe(ctx context.Context, addr string) (err error) {
 		}()
 	}
 
+	// Start health probe server on separate port for Kubernetes probes.
+	healthAddr := s.HealthAddress
+	if healthAddr == "" {
+		healthAddr = "localhost:8080"
+	}
+	s.healthChecker = health.NewChecker(s.Logger, s.netbird)
+	s.healthServer = health.NewServer(healthAddr, s.healthChecker, s.Logger)
+	healthListener, err := net.Listen("tcp", healthAddr)
+	if err != nil {
+		return fmt.Errorf("health probe server listen on %s: %w", healthAddr, err)
+	}
+	go func() {
+		if err := s.healthServer.Serve(healthListener); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			s.Logger.Errorf("health probe server: %v", err)
+		}
+	}()
+	defer func() {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := s.healthServer.Shutdown(shutdownCtx); err != nil {
+			s.Logger.Debugf("health probe server shutdown: %v", err)
+		}
+	}()
+
 	defer func() {
 		stopCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 		defer cancel()
@@ -252,8 +281,15 @@ func (s *Server) ListenAndServe(ctx context.Context, addr string) (err error) {
 
 func (s *Server) newManagementMappingWorker(ctx context.Context, client proto.ProxyServiceClient) {
 	b := backoff.New(0, 0)
+	initialSyncDone := false
 	for {
 		s.Logger.Debug("Getting mapping updates from management server")
+
+		// Mark management as disconnected while we're attempting to reconnect.
+		if s.healthChecker != nil {
+			s.healthChecker.SetManagementConnected(false)
+		}
+
 		mappingClient, err := client.GetMappingUpdate(ctx, &proto.GetMappingUpdateRequest{
 			ProxyId:   s.ID,
 			Version:   s.Version,
@@ -270,8 +306,14 @@ func (s *Server) newManagementMappingWorker(ctx context.Context, client proto.Pr
 			time.Sleep(backoffDuration)
 			continue
 		}
+
+		// Mark management as connected once stream is established.
+		if s.healthChecker != nil {
+			s.healthChecker.SetManagementConnected(true)
+		}
 		s.Logger.Debug("Got mapping updates client from management server")
-		err = s.handleMappingStream(ctx, mappingClient)
+
+		err = s.handleMappingStream(ctx, mappingClient, &initialSyncDone)
 		backoffDuration := b.Duration()
 		switch {
 		case errors.Is(err, context.Canceled),
@@ -294,7 +336,7 @@ func (s *Server) newManagementMappingWorker(ctx context.Context, client proto.Pr
 	}
 }
 
-func (s *Server) handleMappingStream(ctx context.Context, mappingClient proto.ProxyService_GetMappingUpdateClient) error {
+func (s *Server) handleMappingStream(ctx context.Context, mappingClient proto.ProxyService_GetMappingUpdateClient, initialSyncDone *bool) error {
 	for {
 		// Check for context completion to gracefully shutdown.
 		select {
@@ -338,6 +380,14 @@ func (s *Server) handleMappingStream(ctx context.Context, mappingClient proto.Pr
 				}
 			}
 			s.Logger.Debug("Processing mapping update completed")
+
+			// After the first mapping sync, mark initial sync complete.
+			// Client health is checked directly in the startup probe.
+			if !*initialSyncDone && s.healthChecker != nil {
+				s.healthChecker.SetInitialSyncComplete()
+				*initialSyncDone = true
+				s.Logger.Info("Initial mapping sync complete")
+			}
 		}
 	}
 }
@@ -376,7 +426,7 @@ func (s *Server) updateMapping(ctx context.Context, mapping *proto.ProxyMapping)
 	}
 	if mapping.GetAuth().GetOidc() != nil {
 		oidc := mapping.GetAuth().GetOidc()
-		schemes = append(schemes, auth.NewOIDC(mgmtClient, mapping.GetId(), mapping.GetAccountId(), auth.OIDCConfig{
+		schemes = append(schemes, auth.NewOIDC(s.mgmtClient, mapping.GetId(), mapping.GetAccountId(), auth.OIDCConfig{
 			Issuer:             oidc.GetIssuer(),
 			Audiences:          oidc.GetAudiences(),
 			KeysLocation:       oidc.GetKeysLocation(),