Merge pull request #14 from uselagoon/clean-shutdown

Add metrics and ensure clean shutdown for ssh-portal service

cmd/service-api/serve.go

@@ -31,7 +31,7 @@ func (cmd *ServeCmd) Run(log *zap.Logger) error {
 	// will exit immediately if the context is already done.
 	ictx := context.Background()
 	// init metrics
-	m := metrics.NewServer(log)
+	m := metrics.NewServer(log, ":9911")
 	defer m.Shutdown(ictx) //nolint:errcheck
 	// get main process context
 	ctx, cancel := signalctx.GetContext()

cmd/ssh-portal/serve.go

@@ -28,7 +28,7 @@ func (cmd *ServeCmd) Run(log *zap.Logger) error {
 	// will exit immediately if the context is already done.
 	ictx := context.Background()
 	// init metrics
-	m := metrics.NewServer(log)
+	m := metrics.NewServer(log, ":9912")
 	defer m.Shutdown(ictx) //nolint:errcheck
 	// get main process context
 	ctx, cancel := signalctx.GetContext()

internal/metrics/metrics.go

@@ -11,11 +11,11 @@ import (
 // NewServer returns a *http.Server serving prometheus metrics in a new
 // goroutine.
 // Caller should defer Shutdown() for cleanup.
-func NewServer(log *zap.Logger) *http.Server {
+func NewServer(log *zap.Logger, addr string) *http.Server {
 	mux := http.NewServeMux()
 	mux.Handle("/metrics", promhttp.Handler())
 	s := http.Server{
-		Addr:         ":9911",
+		Addr:         addr,
 		Handler:      mux,
 		ReadTimeout:  16 * time.Second,
 		WriteTimeout: 16 * time.Second,

internal/serviceapi/server.go

@@ -7,8 +7,6 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/uselagoon/ssh-portal/internal/lagoondb"
 	"go.uber.org/zap"
 )
@@ -18,13 +16,6 @@ const (
 	pkgName = "github.com/uselagoon/ssh-portal/internal/serviceapi"
 )
 
-var (
-	requestsCounter = promauto.NewCounter(prometheus.CounterOpts{
-		Name: "serviceapi_requests_total",
-		Help: "The total number of requests received",
-	})
-)
-
 // LagoonDBService provides methods for querying the Lagoon API DB.
 type LagoonDBService interface {
 	EnvironmentByNamespaceName(context.Context, string) (*lagoondb.Environment, error)

internal/serviceapi/sshportal.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 
 	"github.com/nats-io/nats.go"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/uselagoon/ssh-portal/internal/lagoondb"
 	"github.com/uselagoon/ssh-portal/internal/permission"
 	"go.opentelemetry.io/otel"
@@ -24,6 +26,13 @@ type SSHAccessQuery struct {
 	EnvironmentID  int
 }
 
+var (
+	requestsCounter = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "serviceapi_requests_total",
+		Help: "The total number of requests received",
+	})
+)
+
 func sshportal(ctx context.Context, log *zap.Logger, c *nats.EncodedConn,
 	l LagoonDBService, k KeycloakService) nats.Handler {
 	return func(_, replySubject string, query *SSHAccessQuery) {

internal/sshserver/authhandler.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/gliderlabs/ssh"
 	"github.com/nats-io/nats.go"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/uselagoon/ssh-portal/internal/k8s"
 	"github.com/uselagoon/ssh-portal/internal/serviceapi"
 	"go.uber.org/zap"
@@ -17,11 +19,23 @@ var (
 	natsTimeout = 8 * time.Second
 )
 
+var (
+	authAttemptsTotal = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "authentication_attempts_total",
+		Help: "The total number of authentication attempts",
+	})
+	authSuccessTotal = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "authentication_success_total",
+		Help: "The total number of successful authentication",
+	})
+)
+
 // pubKeyAuth returns a ssh.PublicKeyHandler which accepts any key, and simply
 // adds the given key to the connection context.
 func pubKeyAuth(log *zap.Logger, nc *nats.Conn,
 	c *k8s.Client) ssh.PublicKeyHandler {
 	return func(ctx ssh.Context, key ssh.PublicKey) bool {
+		authAttemptsTotal.Inc()
 		// parse SSH public key
 		pubKey, err := gossh.ParsePublicKey(key.Marshal())
 		if err != nil {
@@ -33,7 +47,7 @@ func pubKeyAuth(log *zap.Logger, nc *nats.Conn,
 		// get Lagoon labels from namespace if available
 		pid, eid, err := c.NamespaceDetails(ctx.User())
 		if err != nil {
-			log.Info("couldn't get namespace details",
+			log.Debug("couldn't get namespace details",
 				zap.String("session-id", ctx.SessionID()),
 				zap.String("namespace", ctx.User()), zap.Error(err))
 			return false
@@ -62,6 +76,7 @@ func pubKeyAuth(log *zap.Logger, nc *nats.Conn,
 		}
 		// handle response
 		if bytes.Equal(response.Data, []byte("true")) {
+			authSuccessTotal.Inc()
 			log.Debug("authentication successful",
 				zap.String("session-id", ctx.SessionID()),
 				zap.String("fingerprint", fingerprint),

internal/sshserver/serve.go

@@ -2,8 +2,10 @@ package sshserver
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net"
+	"time"
 
 	"github.com/gliderlabs/ssh"
 	"github.com/nats-io/nats.go"
@@ -23,5 +25,18 @@ func Serve(ctx context.Context, log *zap.Logger, nc *nats.Conn,
 			return fmt.Errorf("invalid host key: %v", err)
 		}
 	}
-	return srv.Serve(l)
+	go func() {
+		// As soon as the top level context is cancelled, shut down the server.
+		// Give an 8 second deadline to do this.
+		<-ctx.Done()
+		shutCtx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
+		defer cancel()
+		if err := srv.Shutdown(shutCtx); err != nil {
+			log.Warn("couldn't shutdown cleanly", zap.Error(err))
+		}
+	}()
+	if err := srv.Serve(l); !errors.Is(ssh.ErrServerClosed, err) {
+		return err
+	}
+	return nil
 }

internal/sshserver/sessionhandler.go

@@ -4,12 +4,22 @@ import (
 	"fmt"
 
 	"github.com/gliderlabs/ssh"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/uselagoon/ssh-portal/internal/k8s"
 	"go.uber.org/zap"
 )
 
+var (
+	sessionTotal = promauto.NewCounter(prometheus.CounterOpts{
+		Name: "session_total",
+		Help: "The total number of ssh sessions",
+	})
+)
+
 func sessionHandler(log *zap.Logger, c *k8s.Client) ssh.Handler {
 	return func(s ssh.Session) {
+		sessionTotal.Inc()
 		sid, ok := s.Context().Value(ssh.ContextKeySessionID).(string)
 		if !ok {
 			log.Warn("couldn't get session ID")