feat: health check port

Signed-off-by: QuentinN42 <quentin@lieumont.fr>

helm/templates/deploy.yaml

@@ -22,6 +22,7 @@ spec:
           command:
             - locations
             - start
+            - "-v"
             - {{ .Values.ESCAPE_PRIVATE_LOCATION | default (include "fullname" .) | quote }}
           imagePullPolicy: Always
           resources:

pkg/locations/health/main.go

@@ -0,0 +1,53 @@
+package health
+
+import (
+	"context"
+	"net/http"
+	"os"
+	"sync/atomic"
+
+	"github.com/Escape-Technologies/cli/pkg/log"
+)
+
+func buildHandler(healthy *atomic.Bool) func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var msg string
+		if healthy.Load() {
+			w.WriteHeader(http.StatusOK)
+			msg = "OK"
+		} else {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			msg = "Not connected"
+		}
+		_, err := w.Write([]byte(msg))
+		if err != nil {
+			log.Debug("Error during health check: %v", err)
+		}
+	}
+}
+
+func Start(ctx context.Context, healthy *atomic.Bool) {
+	if os.Getenv("HEALTH_CHECK_PORT") == "" {
+		log.Trace("HEALTH_CHECK_PORT is not set, skipping health check")
+		return
+	}
+
+	srv := &http.Server{
+		Addr:    ":" + os.Getenv("HEALTH_CHECK_PORT"),
+		Handler: http.HandlerFunc(buildHandler(healthy)),
+	}
+	go func() {
+		<-ctx.Done()
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			log.Error("Error shutting down health check server: %v", err)
+		}
+	}()
+	go func() {
+		err := srv.ListenAndServe()
+		if err != nil {
+			log.Error("Error starting the health check server: %v", err)
+		}
+	}()
+	log.Info("Health check server started on http://0.0.0.0:%s/health", os.Getenv("HEALTH_CHECK_PORT"))
+}

pkg/locations/private/dialssh.go

@@ -4,34 +4,35 @@ import (
 	"context"
 	"crypto/ed25519"
 	"fmt"
+	"sync/atomic"
 
 	"github.com/Escape-Technologies/cli/pkg/log"
 	"golang.org/x/crypto/ssh"
 )
 
-func dialSSH(ctx context.Context, locationId string, sshPrivateKey ed25519.PrivateKey) error {
-	log.Info("Creating signer from private key")
+func dialSSH(ctx context.Context, locationId string, sshPrivateKey ed25519.PrivateKey, healthy *atomic.Bool) error {
+	log.Debug("Creating signer from private key")
 	signer, err := ssh.NewSignerFromKey(sshPrivateKey)
 	if err != nil {
 		return fmt.Errorf("failed to create signer: %w", err)
 	}
-	
+
 	config := &ssh.ClientConfig{
 		User: locationId,
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeys(signer),
 		},
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
-	
+
 	log.Info("Dialing locationID: %s", locationId)
-	client, err := ssh.Dial("tcp", "a814bdc744e1147dd86d66114ed8edcc-2eb18fcf1bd8afa3.elb.eu-west-3.amazonaws.com:2222", config)
+	client, err := ssh.Dial("tcp", "private-location.escape.tech:2222", config)
 	if err != nil {
 		return fmt.Errorf("failed to dial: %w", err)
 	}
 
 	log.Info("Starting listener")
-	err = startListener(ctx, client)
+	err = startListener(ctx, client, healthy)
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}

pkg/locations/private/listener.go

@@ -4,21 +4,23 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync/atomic"
 
 	"github.com/Escape-Technologies/cli/pkg/log"
 	"golang.org/x/crypto/ssh"
 )
 
-func startListener(ctx context.Context, client *ssh.Client) error {
+func startListener(ctx context.Context, client *ssh.Client, healthy *atomic.Bool) error {
 	listener, err := client.Listen("tcp", "0.0.0.0:0")
 	if err != nil {
 		return fmt.Errorf("failed to create reverse tunnel: %w", err)
 	}
 	defer listener.Close()
 
 	log.Info("Established reverse tunnel on remote port %d", listener.Addr().(*net.TCPAddr).Port)
-
-	err = startSocks5Server(ctx, listener)
+
+	err = startSocks5Server(ctx, listener, healthy)
+	healthy.Store(false)
 	if ctx.Err() != nil {
 		return ctx.Err()
 	}

pkg/locations/private/main.go

@@ -3,16 +3,17 @@ package private
 import (
 	"context"
 	"crypto/ed25519"
+	"sync/atomic"
 	"time"
 
 	"github.com/Escape-Technologies/cli/pkg/log"
 )
 
-func StartLocation(ctx context.Context, locationId string, sshPrivateKey ed25519.PrivateKey) error {
+func StartLocation(ctx context.Context, locationId string, sshPrivateKey ed25519.PrivateKey, healthy *atomic.Bool) error {
 	log.Info("Starting location")
 
 	for {
-		err := dialSSH(ctx, locationId, sshPrivateKey)
+		err := dialSSH(ctx, locationId, sshPrivateKey, healthy)
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}

pkg/locations/private/socks5server.go

@@ -4,18 +4,20 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync/atomic"
 
 	"github.com/Escape-Technologies/cli/pkg/log"
 	socks5 "github.com/Escape-Technologies/go-socks5"
 )
 
-func startSocks5Server(ctx context.Context, listener net.Listener) error {
+func startSocks5Server(ctx context.Context, listener net.Listener, healthy *atomic.Bool) error {
 	log.Info("Starting socks5 server")
 	socks5Server, err := socks5.New(&socks5.Config{})
 	if err != nil {
 		return fmt.Errorf("failed to create socks5 server config: %w", err)
 	}
 	log.Info("Socks5 server started")
+	healthy.Store(true)
 
 	errChan := make(chan error)
 	ctx, cancel := context.WithCancel(ctx)

pkg/locations/start.go

@@ -3,15 +3,20 @@ package locations
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"github.com/Escape-Technologies/cli/pkg/api"
+	"github.com/Escape-Technologies/cli/pkg/locations/health"
 	"github.com/Escape-Technologies/cli/pkg/locations/private"
 	"github.com/Escape-Technologies/cli/pkg/log"
 )
 
-
 func Start(ctx context.Context, client *api.ClientWithResponses, name string) error {
+	healthy := &atomic.Bool{}
+	healthy.Store(false)
+	go health.Start(ctx, healthy)
+
 	sshPublicKey, sshPrivateKey := private.GenSSHKeys(name)
 	log.Info("Generated public SSH Key: %s", sshPublicKey)
 
@@ -27,7 +32,7 @@ func Start(ctx context.Context, client *api.ClientWithResponses, name string) er
 	}
 	if location.JSON200 != nil {
 		for {
-			err := private.StartLocation(ctx,location.JSON200.Id.String(), sshPrivateKey)
+			err := private.StartLocation(ctx, location.JSON200.Id.String(), sshPrivateKey, healthy)
 			if err != nil {
 				log.Error("Error starting location: %s", err)
 			} else {