From b6df4b9f0abd6e7f69421bc5e65bd6b16da0438b Mon Sep 17 00:00:00 2001
From: Stavros <steveiliop56@gmail.com>
Date: Mon, 11 May 2026 19:10:50 +0300
Subject: [PATCH] fix: coderabbit comments

---
 internal/bootstrap/router_bootstrap.go    | 10 +++++++++-
 internal/controller/user_controller.go    |  2 +-
 internal/middleware/context_middleware.go |  1 +
 internal/model/constants.go               |  2 ++
 internal/service/auth_service.go          |  8 ++++----
 internal/service/tailscale_service.go     |  2 +-
 6 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/internal/bootstrap/router_bootstrap.go b/internal/bootstrap/router_bootstrap.go
index 0bff3b4e..f072134f 100644
--- a/internal/bootstrap/router_bootstrap.go
+++ b/internal/bootstrap/router_bootstrap.go
@@ -1,14 +1,17 @@
 package bootstrap
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
 	"net/http"
 	"os"
+	"time"
 
 	"github.com/tinyauthapp/tinyauth/internal/controller"
 	"github.com/tinyauthapp/tinyauth/internal/middleware"
+	"github.com/tinyauthapp/tinyauth/internal/model"
 
 	"github.com/gin-gonic/gin"
 )
@@ -199,7 +202,12 @@ func (app *BootstrapApp) serveTailscale() error {
 
 func (app *BootstrapApp) serve(listener net.Listener, server *http.Server, name string) error {
 	shutdown := func() {
-		server.Shutdown(app.ctx)
+		ctx, cancel := context.WithTimeout(context.Background(), model.GracefulShutdownTimeout*time.Second)
+		defer cancel()
+		err := server.Shutdown(ctx)
+		if err != nil {
+			app.log.App.Error().Err(err).Msgf("Failed to shutdown %s listener gracefully", name)
+		}
 		listener.Close()
 	}
 
diff --git a/internal/controller/user_controller.go b/internal/controller/user_controller.go
index 72ab6e04..e235e571 100644
--- a/internal/controller/user_controller.go
+++ b/internal/controller/user_controller.go
@@ -424,7 +424,7 @@ func (controller *UserController) tailscaleHandler(c *gin.Context) {
 	cookie, err := controller.auth.CreateSession(c, sessionCookie)
 
 	if err != nil {
-		controller.log.App.Error().Err(err).Str("username", context.GetUsername()).Msg("Failed to create session cookie after successful TOTP verification")
+		controller.log.App.Error().Err(err).Str("username", context.GetUsername()).Msg("Failed to create session cookie after successful Tailscale login")
 		c.JSON(500, gin.H{
 			"status":  500,
 			"message": "Internal Server Error",
diff --git a/internal/middleware/context_middleware.go b/internal/middleware/context_middleware.go
index f8f9ca8e..f5b7fc63 100644
--- a/internal/middleware/context_middleware.go
+++ b/internal/middleware/context_middleware.go
@@ -116,6 +116,7 @@ func (m *ContextMiddleware) Middleware() gin.HandlerFunc {
 			if tailscaleContext != nil {
 				c.Set("context", &model.UserContext{
 					Authenticated: false,
+					Provider:      model.ProviderTailscale,
 					Tailscale:     tailscaleContext,
 				})
 			}
diff --git a/internal/model/constants.go b/internal/model/constants.go
index d9e85e57..d5885dcf 100644
--- a/internal/model/constants.go
+++ b/internal/model/constants.go
@@ -21,3 +21,5 @@ const SessionCookieName = "tinyauth-session"
 const CSRFCookieName = "tinyauth-csrf"
 const RedirectCookieName = "tinyauth-redirect"
 const OAuthSessionCookieName = "tinyauth-oauth"
+
+const GracefulShutdownTimeout = 5 // seconds
diff --git a/internal/service/auth_service.go b/internal/service/auth_service.go
index 4b191c89..944d7307 100644
--- a/internal/service/auth_service.go
+++ b/internal/service/auth_service.go
@@ -292,6 +292,10 @@ func (auth *AuthService) IsEmailWhitelisted(email string) bool {
 }
 
 func (auth *AuthService) CreateSession(ctx context.Context, data repository.Session) (*http.Cookie, error) {
+	if data.Provider == "tailscale" && auth.tailscale == nil {
+		return nil, fmt.Errorf("tailscale service not configured, cannot create session for tailscale user")
+	}
+
 	uuid, err := uuid.NewRandom()
 
 	if err != nil {
@@ -329,10 +333,6 @@ func (auth *AuthService) CreateSession(ctx context.Context, data repository.Sess
 	}
 
 	if data.Provider == "tailscale" {
-		if auth.tailscale == nil {
-			return nil, fmt.Errorf("tailscale service not configured, cannot create session for tailscale user")
-		}
-
 		auth.log.App.Trace().Str("url", fmt.Sprintf("https://%s", auth.tailscale.GetHostname())).Msg("Extracting root domain from Tailscale hostname")
 
 		tsCookieDomain, err := utils.GetCookieDomain(fmt.Sprintf("https://%s", auth.tailscale.GetHostname()))
diff --git a/internal/service/tailscale_service.go b/internal/service/tailscale_service.go
index 0ed74dad..9d341063 100644
--- a/internal/service/tailscale_service.go
+++ b/internal/service/tailscale_service.go
@@ -64,7 +64,7 @@ func NewTailscaleService(log *logger.Logger, config model.Config, ctx context.Co
 		lc:     lc,
 	}
 
-	connectCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
+	connectCtx, cancel := context.WithTimeout(ctx, 2*time.Minute) // large enough timeout to allow for user to manually authenticate with link if needed
 	defer cancel()
 
 	err = service.waitForConn(connectCtx)