refactor: oauth flow

internal/bootstrap/app_bootstrap.go

@@ -22,16 +22,17 @@ import (
 type BootstrapApp struct {
 	config  config.Config
 	context struct {
-		appUrl              string
-		uuid                string
-		cookieDomain        string
-		sessionCookieName   string
-		csrfCookieName      string
-		redirectCookieName  string
-		users               []config.User
-		oauthProviders      map[string]config.OAuthServiceConfig
-		configuredProviders []controller.Provider
-		oidcClients         []config.OIDCClientConfig
+		appUrl                 string
+		uuid                   string
+		cookieDomain           string
+		sessionCookieName      string
+		csrfCookieName         string
+		redirectCookieName     string
+		oauthSessionCookieName string
+		users                  []config.User
+		oauthProviders         map[string]config.OAuthServiceConfig
+		configuredProviders    []controller.Provider
+		oidcClients            []config.OIDCClientConfig
 	}
 	services Services
 }
@@ -113,6 +114,7 @@ func (app *BootstrapApp) Setup() error {
 	app.context.sessionCookieName = fmt.Sprintf("%s-%s", config.SessionCookieName, cookieId)
 	app.context.csrfCookieName = fmt.Sprintf("%s-%s", config.CSRFCookieName, cookieId)
 	app.context.redirectCookieName = fmt.Sprintf("%s-%s", config.RedirectCookieName, cookieId)
+	app.context.oauthSessionCookieName = fmt.Sprintf("%s-%s", config.OAuthSessionCookieName, cookieId)
 
 	// Dumps
 	tlog.App.Trace().Interface("config", app.config).Msg("Config dump")
@@ -190,12 +192,12 @@ func (app *BootstrapApp) Setup() error {
 
 	// Start db cleanup routine
 	tlog.App.Debug().Msg("Starting database cleanup routine")
-	go app.dbCleanup(queries)
+	go app.dbCleanupRoutine(queries)
 
 	// If analytics are not disabled, start heartbeat
 	if app.config.Analytics.Enabled {
 		tlog.App.Debug().Msg("Starting heartbeat routine")
-		go app.heartbeat()
+		go app.heartbeatRoutine()
 	}
 
 	// If we have an socket path, bind to it
@@ -226,7 +228,7 @@ func (app *BootstrapApp) Setup() error {
 	return nil
 }
 
-func (app *BootstrapApp) heartbeat() {
+func (app *BootstrapApp) heartbeatRoutine() {
 	ticker := time.NewTicker(time.Duration(12) * time.Hour)
 	defer ticker.Stop()
 
@@ -280,7 +282,7 @@ func (app *BootstrapApp) heartbeat() {
 	}
 }
 
-func (app *BootstrapApp) dbCleanup(queries *repository.Queries) {
+func (app *BootstrapApp) dbCleanupRoutine(queries *repository.Queries) {
 	ticker := time.NewTicker(time.Duration(30) * time.Minute)
 	defer ticker.Stop()
 	ctx := context.Background()

internal/bootstrap/router_bootstrap.go

@@ -77,12 +77,13 @@ func (app *BootstrapApp) setupRouter() (*gin.Engine, error) {
 	contextController.SetupRoutes()
 
 	oauthController := controller.NewOAuthController(controller.OAuthControllerConfig{
-		AppURL:             app.config.AppURL,
-		SecureCookie:       app.config.Auth.SecureCookie,
-		CSRFCookieName:     app.context.csrfCookieName,
-		RedirectCookieName: app.context.redirectCookieName,
-		CookieDomain:       app.context.cookieDomain,
-	}, apiRouter, app.services.authService, app.services.oauthBrokerService)
+		AppURL:                 app.config.AppURL,
+		SecureCookie:           app.config.Auth.SecureCookie,
+		CSRFCookieName:         app.context.csrfCookieName,
+		RedirectCookieName:     app.context.redirectCookieName,
+		CookieDomain:           app.context.cookieDomain,
+		OAuthSessionCookieName: app.context.oauthSessionCookieName,
+	}, apiRouter, app.services.authService)
 
 	oauthController.SetupRoutes()
 

internal/bootstrap/service_bootstrap.go

@@ -58,6 +58,16 @@ func (app *BootstrapApp) initServices(queries *repository.Queries) (Services, er
 
 	services.accessControlService = accessControlsService
 
+	oauthBrokerService := service.NewOAuthBrokerService(app.context.oauthProviders)
+
+	err = oauthBrokerService.Init()
+
+	if err != nil {
+		return Services{}, err
+	}
+
+	services.oauthBrokerService = oauthBrokerService
+
 	authService := service.NewAuthService(service.AuthServiceConfig{
 		Users:              app.context.users,
 		OauthWhitelist:     app.config.OAuth.Whitelist,
@@ -70,7 +80,7 @@ func (app *BootstrapApp) initServices(queries *repository.Queries) (Services, er
 		SessionCookieName:  app.context.sessionCookieName,
 		IP:                 app.config.Auth.IP,
 		LDAPGroupsCacheTTL: app.config.Ldap.GroupCacheTTL,
-	}, dockerService, services.ldapService, queries)
+	}, dockerService, services.ldapService, queries, services.oauthBrokerService)
 
 	err = authService.Init()
 
@@ -80,16 +90,6 @@ func (app *BootstrapApp) initServices(queries *repository.Queries) (Services, er
 
 	services.authService = authService
 
-	oauthBrokerService := service.NewOAuthBrokerService(app.context.oauthProviders)
-
-	err = oauthBrokerService.Init()
-
-	if err != nil {
-		return Services{}, err
-	}
-
-	services.oauthBrokerService = oauthBrokerService
-
 	oidcService := service.NewOIDCService(service.OIDCServiceConfig{
 		Clients:        app.config.OIDC.Clients,
 		PrivateKeyPath: app.config.OIDC.PrivateKeyPath,

internal/config/config.go

@@ -73,6 +73,7 @@ var BuildTimestamp = "0000-00-00T00:00:00Z"
 var SessionCookieName = "tinyauth-session"
 var CSRFCookieName = "tinyauth-csrf"
 var RedirectCookieName = "tinyauth-redirect"
+var OAuthSessionCookieName = "tinyauth-oauth"
 
 // Main app config
 

internal/controller/oauth_controller.go

@@ -21,26 +21,25 @@ type OAuthRequest struct {
 }
 
 type OAuthControllerConfig struct {
-	CSRFCookieName     string
-	RedirectCookieName string
-	SecureCookie       bool
-	AppURL             string
-	CookieDomain       string
+	CSRFCookieName         string
+	OAuthSessionCookieName string
+	RedirectCookieName     string
+	SecureCookie           bool
+	AppURL                 string
+	CookieDomain           string
 }
 
 type OAuthController struct {
 	config OAuthControllerConfig
 	router *gin.RouterGroup
 	auth   *service.AuthService
-	broker *service.OAuthBrokerService
 }
 
-func NewOAuthController(config OAuthControllerConfig, router *gin.RouterGroup, auth *service.AuthService, broker *service.OAuthBrokerService) *OAuthController {
+func NewOAuthController(config OAuthControllerConfig, router *gin.RouterGroup, auth *service.AuthService) *OAuthController {
 	return &OAuthController{
 		config: config,
 		router: router,
 		auth:   auth,
-		broker: broker,
 	}
 }
 
@@ -63,21 +62,30 @@ func (controller *OAuthController) oauthURLHandler(c *gin.Context) {
 		return
 	}
 
-	service, exists := controller.broker.GetService(req.Provider)
+	sessionId, session, err := controller.auth.NewOAuthSession(req.Provider)
 
-	if !exists {
-		tlog.App.Warn().Msgf("OAuth provider not found: %s", req.Provider)
-		c.JSON(404, gin.H{
-			"status":  404,
-			"message": "Not Found",
+	if err != nil {
+		tlog.App.Error().Err(err).Msg("Failed to create OAuth session")
+		c.JSON(500, gin.H{
+			"status":  500,
+			"message": "Internal Server Error",
+		})
+		return
+	}
+
+	authUrl, err := controller.auth.GetOAuthURL(sessionId)
+
+	if err != nil {
+		tlog.App.Error().Err(err).Msg("Failed to get OAuth URL")
+		c.JSON(500, gin.H{
+			"status":  500,
+			"message": "Internal Server Error",
 		})
 		return
 	}
 
-	service.GenerateVerifier()
-	state := service.GenerateState()
-	authURL := service.GetAuthURL(state)
-	c.SetCookie(controller.config.CSRFCookieName, state, int(time.Hour.Seconds()), "/", fmt.Sprintf(".%s", controller.config.CookieDomain), controller.config.SecureCookie, true)
+	c.SetCookie(controller.config.OAuthSessionCookieName, sessionId, int(time.Hour.Seconds()), "/", fmt.Sprintf(".%s", controller.config.CookieDomain), controller.config.SecureCookie, true)
+	c.SetCookie(controller.config.CSRFCookieName, session.State, int(time.Hour.Seconds()), "/", fmt.Sprintf(".%s", controller.config.CookieDomain), controller.config.SecureCookie, true)
 
 	redirectURI := c.Query("redirect_uri")
 	isRedirectSafe := utils.IsRedirectSafe(redirectURI, controller.config.CookieDomain)
@@ -95,7 +103,7 @@ func (controller *OAuthController) oauthURLHandler(c *gin.Context) {
 	c.JSON(200, gin.H{
 		"status":  200,
 		"message": "OK",
-		"url":     authURL,
+		"url":     authUrl,
 	})
 }
 
@@ -112,6 +120,17 @@ func (controller *OAuthController) oauthCallbackHandler(c *gin.Context) {
 		return
 	}
 
+	sessionIdCookie, err := c.Cookie(controller.config.OAuthSessionCookieName)
+
+	if err != nil {
+		tlog.App.Warn().Err(err).Msg("OAuth session cookie missing")
+		c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s/error", controller.config.AppURL))
+		return
+	}
+
+	c.SetCookie(controller.config.OAuthSessionCookieName, "", -1, "/", fmt.Sprintf(".%s", controller.config.CookieDomain), controller.config.SecureCookie, true)
+	defer controller.auth.EndOAuthSession(sessionIdCookie)
+
 	state := c.Query("state")
 	csrfCookie, err := c.Cookie(controller.config.CSRFCookieName)
 
@@ -125,28 +144,15 @@ func (controller *OAuthController) oauthCallbackHandler(c *gin.Context) {
 	c.SetCookie(controller.config.CSRFCookieName, "", -1, "/", fmt.Sprintf(".%s", controller.config.CookieDomain), controller.config.SecureCookie, true)
 
 	code := c.Query("code")
-	service, exists := controller.broker.GetService(req.Provider)
+	_, err = controller.auth.GetOAuthToken(sessionIdCookie, code)
 
-	if !exists {
-		tlog.App.Warn().Msgf("OAuth provider not found: %s", req.Provider)
-		c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s/error", controller.config.AppURL))
-		return
-	}
-
-	err = service.VerifyCode(code)
 	if err != nil {
-		tlog.App.Error().Err(err).Msg("Failed to verify OAuth code")
+		tlog.App.Error().Err(err).Msg("Failed to exchange code for token")
 		c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s/error", controller.config.AppURL))
 		return
 	}
 
-	user, err := controller.broker.GetUser(req.Provider)
-
-	if err != nil {
-		tlog.App.Error().Err(err).Msg("Failed to get user from OAuth provider")
-		c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s/error", controller.config.AppURL))
-		return
-	}
+	user, err := controller.auth.GetOAuthUserinfo(sessionIdCookie)
 
 	if user.Email == "" {
 		tlog.App.Error().Msg("OAuth provider did not return an email")
@@ -192,13 +198,21 @@ func (controller *OAuthController) oauthCallbackHandler(c *gin.Context) {
 		username = strings.Replace(user.Email, "@", "_", 1)
 	}
 
+	service, err := controller.auth.GetOAuthService(sessionIdCookie)
+
+	if err != nil {
+		tlog.App.Error().Err(err).Msg("Failed to get OAuth service for session")
+		c.Redirect(http.StatusTemporaryRedirect, fmt.Sprintf("%s/error", controller.config.AppURL))
+		return
+	}
+
 	sessionCookie := repository.Session{
 		Username:    username,
 		Name:        name,
 		Email:       user.Email,
 		Provider:    req.Provider,
 		OAuthGroups: utils.CoalesceToString(user.Groups),
-		OAuthName:   service.GetName(),
+		OAuthName:   service.Name(),
 		OAuthSub:    user.Sub,
 	}
 

internal/controller/proxy_controller_test.go

@@ -85,7 +85,7 @@ func setupProxyController(t *testing.T, middlewares []gin.HandlerFunc) (*gin.Eng
 		LoginTimeout:       300,
 		LoginMaxRetries:    3,
 		SessionCookieName:  "tinyauth-session",
-	}, dockerService, nil, queries)
+	}, dockerService, nil, queries, &service.OAuthBrokerService{})
 
 	// Controller
 	ctrl := controller.NewProxyController(controller.ProxyControllerConfig{

internal/controller/user_controller_test.go

@@ -71,7 +71,7 @@ func setupUserController(t *testing.T, middlewares *[]gin.HandlerFunc) (*gin.Eng
 		LoginTimeout:       300,
 		LoginMaxRetries:    3,
 		SessionCookieName:  "tinyauth-session",
-	}, nil, nil, queries)
+	}, nil, nil, queries, &service.OAuthBrokerService{})
 
 	// Controller
 	ctrl := controller.NewUserController(controller.UserControllerConfig{