Skip to content

Commit 94e06f8

Browse files
authored
Merge pull request #125 from thand-io/session-handler-overhaul
Session handler overhaul
2 parents 56787c5 + 9816f4e commit 94e06f8

35 files changed

Lines changed: 2035 additions & 618 deletions

.vscode/launch.json

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,17 @@
4848
"env": {},
4949
"showLog": true
5050
},
51+
{
52+
"name": "Run sessions (local)",
53+
"type": "go",
54+
"request": "launch",
55+
"mode": "auto",
56+
"program": "${workspaceFolder}",
57+
"cwd": "${workspaceFolder}",
58+
"args": ["sessions", "--login-server", "http://localhost:9090"],
59+
"env": {},
60+
"showLog": true
61+
},
5162
{
5263
"name": "Run wizzard (remote)",
5364
"type": "go",

cmd/agent/main.go

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
package cli
22

33
import (
4-
"log"
5-
4+
"github.com/sirupsen/logrus"
65
"github.com/spf13/cobra"
76
"github.com/thand-io/agent/internal/agent"
87
"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
2625
// Load configuration
2726
cfg, err := config.Load(configFile)
2827
if err != nil {
29-
log.Fatalf("Failed to load configuration: %v", err)
28+
logrus.Fatalf("Failed to load configuration: %v", err)
3029
}
3130

3231
if _, err := agent.StartWebService(cfg); err != nil {
33-
log.Fatalf("Failed to start web service: %v", err)
32+
logrus.Fatalf("Failed to start web service: %v", err)
3433
}
3534
},
3635
}
@@ -41,6 +40,6 @@ func init() {
4140

4241
func main() {
4342
if err := rootCmd.Execute(); err != nil {
44-
log.Fatalf("Failed to execute command: %v", err)
43+
logrus.Fatalf("Failed to execute command: %v", err)
4544
}
4645
}

cmd/cli/agent.go

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,11 @@ package cli
33
import (
44
"fmt"
55
"os"
6-
"os/signal"
7-
"syscall"
86

97
"github.com/sirupsen/logrus"
108
"github.com/spf13/cobra"
119
"github.com/thand-io/agent/internal/agent"
10+
"github.com/thand-io/agent/internal/common"
1211
"github.com/thand-io/agent/internal/config"
1312
)
1413

@@ -53,8 +52,8 @@ This will run the web service that handles authentication and authorization requ
5352
fmt.Printf("Environment Architecture: %s\n", cfg.Environment.Architecture)
5453

5554
// Set up signal handling for graceful shutdown
56-
sigChan := make(chan os.Signal, 1)
57-
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
55+
sigChan, cleanup := common.NewInterruptChannel()
56+
defer cleanup()
5857

5958
// Start the web service in a goroutine
6059
errChan := make(chan error, 1)

cmd/cli/login.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
package cli
22

33
import (
4+
"context"
45
"fmt"
56
"net/url"
67

78
"github.com/spf13/cobra"
9+
"github.com/thand-io/agent/internal/common"
10+
"github.com/thand-io/agent/internal/models"
811
)
912

1013
var loginCmd = &cobra.Command{
@@ -27,12 +30,26 @@ var loginCmd = &cobra.Command{
2730
}
2831

2932
func runLogin(cmd *cobra.Command, args []string) error {
33+
return authKickStart()
34+
}
35+
36+
func authKickStart() error {
37+
// Set up signal handling for graceful cancellation
38+
ctx, cleanup := common.WithInterrupt(context.Background())
39+
defer cleanup()
3040

3141
hostname := cfg.GetLoginServerHostname()
3242
fmt.Println("Login server hostname:", hostname)
3343

44+
// Prepare callback URL with local server endpoint
45+
46+
if !cfg.GetServices().HasEncryption() {
47+
return fmt.Errorf("encryption service is not configured")
48+
}
49+
3450
callbackUrl := url.Values{
3551
"callback": {cfg.GetLocalServerUrl()},
52+
"code": {createAuthCode()},
3653
}
3754

3855
// Use the configured login server if no override provided
@@ -50,10 +67,15 @@ func runLogin(cmd *cobra.Command, args []string) error {
5067

5168
// Wait for the session to be established (using empty provider for general login)
5269
session := sessionManager.AwaitRefresh(
70+
ctx,
5371
cfg.GetLoginServerHostname(),
5472
)
5573

5674
if session == nil {
75+
// Check if context was cancelled
76+
if ctx.Err() != nil {
77+
return fmt.Errorf("login cancelled")
78+
}
5779
return fmt.Errorf("authentication failed or timed out")
5880
}
5981

@@ -69,3 +91,16 @@ func init() {
6991
// Add the command to the root
7092
rootCmd.AddCommand(loginCmd)
7193
}
94+
95+
func createAuthCode() string {
96+
code := models.EncodingWrapper{
97+
Type: models.ENCODED_SESSION_CODE,
98+
Data: models.NewCodeWrapper(
99+
cfg.GetLoginServerUrl(),
100+
),
101+
}.EncodeAndEncrypt(
102+
cfg.GetServices().GetEncryption(),
103+
)
104+
105+
return code
106+
}

cmd/cli/main.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ import (
44
"errors"
55
"fmt"
66
"os"
7+
"strings"
78

89
"github.com/charmbracelet/huh"
910
"github.com/kardianos/service"
1011
"github.com/sirupsen/logrus"
1112
"github.com/spf13/cobra"
1213
"github.com/thand-io/agent/internal/agent"
14+
"github.com/thand-io/agent/internal/common"
1315
"github.com/thand-io/agent/internal/config"
1416
"github.com/thand-io/agent/internal/sessions"
1517
)
@@ -85,6 +87,15 @@ func preRunConfigE(cmd *cobra.Command, mode config.Mode) error {
8587
}
8688
}
8789

90+
// Generate a global secret if one hasn't been set
91+
if strings.EqualFold(cfg.Secret, common.DefaultServerSecret) {
92+
generatedSecret, err := common.GenerateSecureRandomString(32)
93+
if err != nil {
94+
return fmt.Errorf("failed to generate secret: %w", err)
95+
}
96+
cfg.Secret = generatedSecret
97+
}
98+
8899
// Load users session state before any command runs
89100
sessionManager = loadUserSessionState(cfg.GetLoginServerHostname())
90101

cmd/cli/request.go

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67
"net/http"
@@ -199,6 +200,7 @@ func authenticateUser(request *models.ElevateRequest) error {
199200

200201
callbackUrl := url.Values{
201202
"callback": {cfg.GetLocalServerUrl()},
203+
"code": {createAuthCode()},
202204
}
203205

204206
if len(request.Authenticator) > 0 {
@@ -219,9 +221,12 @@ func authenticateUser(request *models.ElevateRequest) error {
219221
// the auth in the browser
220222
if len(request.Authenticator) > 0 {
221223

222-
if err := sessionManager.AwaitProviderRefresh(
223-
cfg.GetLoginServerHostname(), request.Authenticator); err != nil {
224-
return fmt.Errorf("failed to await provider refresh: %w", err)
224+
if found := sessionManager.AwaitProviderRefresh(
225+
context.Background(),
226+
cfg.GetLoginServerHostname(),
227+
request.Authenticator,
228+
); found == nil {
229+
return fmt.Errorf("failed to await provider refresh. Authentication timed out or failed")
225230
}
226231

227232
session, err := sessionManager.GetSession(
@@ -238,7 +243,7 @@ func authenticateUser(request *models.ElevateRequest) error {
238243
// If no auth provider is specified then we just wait for any
239244
// valid session to be created
240245
sessionHandler := sessionManager.AwaitRefresh(
241-
cfg.GetLoginServerHostname())
246+
context.Background(), cfg.GetLoginServerHostname())
242247

243248
foundProvider, session, err := sessionHandler.GetFirstActiveSession()
244249

cmd/cli/server.go

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,10 @@ package cli
33
import (
44
"fmt"
55
"os"
6-
"os/signal"
7-
"syscall"
86

97
"github.com/spf13/cobra"
108
"github.com/thand-io/agent/internal/agent"
9+
"github.com/thand-io/agent/internal/common"
1110
)
1211

1312
// serverCmd represents the server command
@@ -33,8 +32,8 @@ This will run the web service that handles authentication and authorization requ
3332
fmt.Printf("Environment Architecture: %s\n", cfg.Environment.Architecture)
3433

3534
// Set up signal handling for graceful shutdown
36-
sigChan := make(chan os.Signal, 1)
37-
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
35+
sigChan, cleanup := common.NewInterruptChannel()
36+
defer cleanup()
3837

3938
// Start the web service in a goroutine
4039
errChan := make(chan error, 1)

0 commit comments

Comments
 (0)