From a7e7b5e85110492b299cbfbe0d3dc86fba0b4263 Mon Sep 17 00:00:00 2001
From: devinchristianson <devin.christianson@maine.edu>
Date: Sun, 3 Mar 2024 07:18:24 -0500
Subject: [PATCH] add cleanup webhook called on disconnect and shutdown

---
 config/appconfig.go                 |  6 ++++++
 internal/backend/handler.go         | 19 +++++++++++++++++++
 internal/backend/handler_factory.go | 13 ++++++++++++-
 message/backend.go                  |  3 +++
 4 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/config/appconfig.go b/config/appconfig.go
index af63a6e0..57c71ef3 100644
--- a/config/appconfig.go
+++ b/config/appconfig.go
@@ -19,6 +19,9 @@ type AppConfig struct {
 	// Auth contains the configuration for user authentication.
 	// swagger:ignore
 	Auth AuthConfig `json:"auth" yaml:"auth"`
+	// CleanupServer contains the settings for fetching the user-specific cleanup webhook.
+	// swagger:ignore
+	CleanupServer HTTPClientConfiguration `json:"cleanupserver" yaml:"cleanupserver"`
 	// Log contains the configuration for the logging level.
 	// swagger:ignore
 	Log LogConfig `json:"log" yaml:"log"`
@@ -119,6 +122,9 @@ func (cfg *AppConfig) Validate(dynamic bool) error {
 	queue.add("geoip", &cfg.GeoIP)
 	queue.add("audit", &cfg.Audit)
 	queue.add("health", &cfg.Health)
+	if cfg.CleanupServer.URL != "" {
+		queue.add("cleanupserver", &cfg.CleanupServer)
+	}
 
 	if cfg.ConfigServer.URL != "" && !dynamic {
 		return queue.Validate()
diff --git a/internal/backend/handler.go b/internal/backend/handler.go
index 20df0218..48587575 100644
--- a/internal/backend/handler.go
+++ b/internal/backend/handler.go
@@ -9,6 +9,7 @@ import (
 
     auth2 "go.containerssh.io/libcontainerssh/auth"
     "go.containerssh.io/libcontainerssh/config"
+    "go.containerssh.io/libcontainerssh/http"
     internalConfig "go.containerssh.io/libcontainerssh/internal/config"
     "go.containerssh.io/libcontainerssh/internal/docker"
     "go.containerssh.io/libcontainerssh/internal/kubernetes"
@@ -27,6 +28,7 @@ type handler struct {
 
 	config                 config.AppConfig
 	configLoader           internalConfig.Loader
+	cleanupClient          http.Client
 	authResponse           sshserver.AuthResponse
 	metricsCollector       metrics.Collector
 	logger                 log.Logger
@@ -219,6 +221,7 @@ func (n *networkHandler) OnDisconnect() {
 	n.lock.Lock()
 	defer n.lock.Unlock()
 	if n.backend != nil {
+		n.cleanupWebhook()
 		n.backend.OnDisconnect()
 		n.backend = nil
 	}
@@ -227,6 +230,7 @@ func (n *networkHandler) OnDisconnect() {
 func (n *networkHandler) OnShutdown(shutdownContext context.Context) {
 	n.lock.Lock()
 	if n.backend != nil {
+		n.cleanupWebhook()
 		backend := n.backend
 		n.lock.Unlock()
 		backend.OnShutdown(shutdownContext)
@@ -234,3 +238,18 @@ func (n *networkHandler) OnShutdown(shutdownContext context.Context) {
 		n.lock.Unlock()
 	}
 }
+
+func (n *networkHandler) cleanupWebhook() {
+	if n.rootHandler.cleanupClient != nil {
+		_, err := n.rootHandler.cleanupClient.Post("", map[string]interface{}{
+			"connectionId": n.connectionID,
+		}, nil)
+		if err != nil {
+			n.logger.Error(message.Wrap(
+				err,
+				message.ECleanupWebhook,
+				"Got unexpected response from cleanup webhook",
+			))
+		}
+	}
+}
diff --git a/internal/backend/handler_factory.go b/internal/backend/handler_factory.go
index 142fd700..bfecf1d3 100644
--- a/internal/backend/handler_factory.go
+++ b/internal/backend/handler_factory.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 
     "go.containerssh.io/libcontainerssh/config"
+    "go.containerssh.io/libcontainerssh/http"
     internalConfig "go.containerssh.io/libcontainerssh/internal/config"
     "go.containerssh.io/libcontainerssh/internal/metrics"
     "go.containerssh.io/libcontainerssh/internal/sshserver"
@@ -26,7 +27,16 @@ func New(
 	if err != nil {
 		return nil, err
 	}
-
+	var cleanupClient http.Client
+	if config.CleanupServer.URL != "" {
+		cleanupClient, err = http.NewClient(
+			config.CleanupServer,
+			logger,
+		)
+		if err != nil {
+			return nil, err
+		}
+	}
 	backendRequestsCounter := metricsCollector.MustCreateCounter(
 		MetricNameBackendRequests,
 		MetricUnitBackendRequests,
@@ -41,6 +51,7 @@ func New(
 	return &handler{
 		config:                 config,
 		configLoader:           loader,
+		cleanupClient:          cleanupClient,
 		authResponse:           defaultAuthResponse,
 		metricsCollector:       metricsCollector,
 		logger:                 logger,
diff --git a/message/backend.go b/message/backend.go
index e20ddc78..5932e2cb 100644
--- a/message/backend.go
+++ b/message/backend.go
@@ -2,3 +2,6 @@ package message
 
 // EBackendConfig indicates that there is an error in the backend configuration.
 const EBackendConfig = "BACKEND_CONFIG_ERROR"
+
+// ECleanupWebhook indicates that in calling the cleanup webhook
+const ECleanupWebhook = "BACKEND_CLEANUP_WEBHOOK_ERROR"