Session handler overhaul

.vscode/launch.json

@@ -48,6 +48,17 @@
             "env": {},
             "showLog": true
         },
+        {
+            "name": "Run sessions (local)",
+            "type": "go",
+            "request": "launch",
+            "mode": "auto",
+            "program": "${workspaceFolder}",
+            "cwd": "${workspaceFolder}",
+            "args": ["sessions", "--login-server", "http://localhost:9090"],
+            "env": {},
+            "showLog": true
+        },
         {
             "name": "Run wizzard (remote)",
             "type": "go",

cmd/agent/main.go

@@ -1,8 +1,7 @@
 package cli
 
 import (
-	"log"
-
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/thand-io/agent/internal/agent"
 	"github.com/thand-io/agent/internal/config"
@@ -26,11 +25,11 @@ If no config file is specified, the agent will look for config files in the foll
 		// Load configuration
 		cfg, err := config.Load(configFile)
 		if err != nil {
-			log.Fatalf("Failed to load configuration: %v", err)
+			logrus.Fatalf("Failed to load configuration: %v", err)
 		}
 
 		if _, err := agent.StartWebService(cfg); err != nil {
-			log.Fatalf("Failed to start web service: %v", err)
+			logrus.Fatalf("Failed to start web service: %v", err)
 		}
 	},
 }
@@ -41,6 +40,6 @@ func init() {
 
 func main() {
 	if err := rootCmd.Execute(); err != nil {
-		log.Fatalf("Failed to execute command: %v", err)
+		logrus.Fatalf("Failed to execute command: %v", err)
 	}
 }

cmd/cli/agent.go

@@ -3,12 +3,11 @@ package cli
 import (
 	"fmt"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/thand-io/agent/internal/agent"
+	"github.com/thand-io/agent/internal/common"
 	"github.com/thand-io/agent/internal/config"
 )
 
@@ -53,8 +52,8 @@ This will run the web service that handles authentication and authorization requ
 		fmt.Printf("Environment Architecture: %s\n", cfg.Environment.Architecture)
 
 		// Set up signal handling for graceful shutdown
-		sigChan := make(chan os.Signal, 1)
-		signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+		sigChan, cleanup := common.NewInterruptChannel()
+		defer cleanup()
 
 		// Start the web service in a goroutine
 		errChan := make(chan error, 1)

cmd/cli/login.go

@@ -1,10 +1,13 @@
 package cli
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 
 	"github.com/spf13/cobra"
+	"github.com/thand-io/agent/internal/common"
+	"github.com/thand-io/agent/internal/models"
 )
 
 var loginCmd = &cobra.Command{
@@ -27,12 +30,26 @@ var loginCmd = &cobra.Command{
 }
 
 func runLogin(cmd *cobra.Command, args []string) error {
+	return authKickStart()
+}
+
+func authKickStart() error {
+	// Set up signal handling for graceful cancellation
+	ctx, cleanup := common.WithInterrupt(context.Background())
+	defer cleanup()
 
 	hostname := cfg.GetLoginServerHostname()
 	fmt.Println("Login server hostname:", hostname)
 
+	// Prepare callback URL with local server endpoint
+
+	if !cfg.GetServices().HasEncryption() {
+		return fmt.Errorf("encryption service is not configured")
+	}
+
 	callbackUrl := url.Values{
 		"callback": {cfg.GetLocalServerUrl()},
+		"code":     {createAuthCode()},
 	}
 
 	// Use the configured login server if no override provided
@@ -50,10 +67,15 @@ func runLogin(cmd *cobra.Command, args []string) error {
 
 	// Wait for the session to be established (using empty provider for general login)
 	session := sessionManager.AwaitRefresh(
+		ctx,
 		cfg.GetLoginServerHostname(),
 	)
 
 	if session == nil {
+		// Check if context was cancelled
+		if ctx.Err() != nil {
+			return fmt.Errorf("login cancelled")
+		}
 		return fmt.Errorf("authentication failed or timed out")
 	}
 
@@ -69,3 +91,16 @@ func init() {
 	// Add the command to the root
 	rootCmd.AddCommand(loginCmd)
 }
+
+func createAuthCode() string {
+	code := models.EncodingWrapper{
+		Type: models.ENCODED_SESSION_CODE,
+		Data: models.NewCodeWrapper(
+			cfg.GetLoginServerUrl(),
+		),
+	}.EncodeAndEncrypt(
+		cfg.GetServices().GetEncryption(),
+	)
+
+	return code
+}

cmd/cli/main.go

@@ -4,12 +4,14 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/charmbracelet/huh"
 	"github.com/kardianos/service"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/thand-io/agent/internal/agent"
+	"github.com/thand-io/agent/internal/common"
 	"github.com/thand-io/agent/internal/config"
 	"github.com/thand-io/agent/internal/sessions"
 )
@@ -85,6 +87,15 @@ func preRunConfigE(cmd *cobra.Command, mode config.Mode) error {
 			}
 		}
 
+		// Generate a global secret if one hasn't been set
+		if strings.EqualFold(cfg.Secret, common.DefaultServerSecret) {
+			generatedSecret, err := common.GenerateSecureRandomString(32)
+			if err != nil {
+				return fmt.Errorf("failed to generate secret: %w", err)
+			}
+			cfg.Secret = generatedSecret
+		}
+
 		// Load users session state before any command runs
 		sessionManager = loadUserSessionState(cfg.GetLoginServerHostname())
 

cmd/cli/request.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -199,6 +200,7 @@ func authenticateUser(request *models.ElevateRequest) error {
 
 	callbackUrl := url.Values{
 		"callback": {cfg.GetLocalServerUrl()},
+		"code":     {createAuthCode()},
 	}
 
 	if len(request.Authenticator) > 0 {
@@ -219,9 +221,12 @@ func authenticateUser(request *models.ElevateRequest) error {
 	// the auth in the browser
 	if len(request.Authenticator) > 0 {
 
-		if err := sessionManager.AwaitProviderRefresh(
-			cfg.GetLoginServerHostname(), request.Authenticator); err != nil {
-			return fmt.Errorf("failed to await provider refresh: %w", err)
+		if found := sessionManager.AwaitProviderRefresh(
+			context.Background(),
+			cfg.GetLoginServerHostname(),
+			request.Authenticator,
+		); found == nil {
+			return fmt.Errorf("failed to await provider refresh. Authentication timed out or failed")
 		}
 
 		session, err := sessionManager.GetSession(
@@ -238,7 +243,7 @@ func authenticateUser(request *models.ElevateRequest) error {
 		// If no auth provider is specified then we just wait for any
 		// valid session to be created
 		sessionHandler := sessionManager.AwaitRefresh(
-			cfg.GetLoginServerHostname())
+			context.Background(), cfg.GetLoginServerHostname())
 
 		foundProvider, session, err := sessionHandler.GetFirstActiveSession()
 

cmd/cli/server.go

@@ -3,11 +3,10 @@ package cli
 import (
 	"fmt"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/spf13/cobra"
 	"github.com/thand-io/agent/internal/agent"
+	"github.com/thand-io/agent/internal/common"
 )
 
 // serverCmd represents the server command
@@ -33,8 +32,8 @@ This will run the web service that handles authentication and authorization requ
 		fmt.Printf("Environment Architecture: %s\n", cfg.Environment.Architecture)
 
 		// Set up signal handling for graceful shutdown
-		sigChan := make(chan os.Signal, 1)
-		signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+		sigChan, cleanup := common.NewInterruptChannel()
+		defer cleanup()
 
 		// Start the web service in a goroutine
 		errChan := make(chan error, 1)