From 2558489e788d2e1f5551b7461b55590ce220eee8 Mon Sep 17 00:00:00 2001 From: NamanMahor Date: Tue, 5 Aug 2025 17:13:22 +0530 Subject: [PATCH 1/5] refresh cookie on accessing /v1/users/current --- admin/server/auth/middleware.go | 15 +++++++++++++++ admin/server/server.go | 6 +++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/admin/server/auth/middleware.go b/admin/server/auth/middleware.go index 35e777877e5..292ae196024 100644 --- a/admin/server/auth/middleware.go +++ b/admin/server/auth/middleware.go @@ -90,6 +90,21 @@ func (a *Authenticator) HTTPMiddlewareLenient(next http.Handler) http.Handler { return a.httpMiddleware(next, true) } +// HTTPMiddlewareWithCookieRefresh is a middleware that refreshes the auth cookie +func (a *Authenticator) HTTPMiddlewareWithCookieRefresh(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sess := a.cookies.Get(r, cookieName) + if authToken, ok := sess.Values[cookieFieldAccessToken].(string); ok && authToken != "" { + // Re-save the cookie to refresh its expiration + if err := sess.Save(r, w); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + } + next.ServeHTTP(w, r) + }) +} + // httpMiddleware is the actual implementation of HTTPMiddleware and HTTPMiddlewareLenient. func (a *Authenticator) httpMiddleware(next http.Handler, lenient bool) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { diff --git a/admin/server/server.go b/admin/server/server.go index f33697d9623..6001c207ab9 100644 --- a/admin/server/server.go +++ b/admin/server/server.go @@ -91,7 +91,7 @@ func New(logger *zap.Logger, adm *admin.Service, issuer *runtimeauth.Issuer, lim cookieStore := cookies.New(logger, opts.SessionKeyPairs...) // Auth tokens are validated against the DB on each request, so we can set a long MaxAge. - cookieStore.MaxAge(60 * 60 * 24 * 365 * 10) // 10 years + cookieStore.MaxAge(60 * 60 * 24 * 14) // 14 days // Set Secure if the admin service is served over HTTPS (will resolve to true in production and false in local dev environments). cookieStore.Options.Secure = adm.URLs.IsHTTPS() @@ -185,6 +185,10 @@ func (s *Server) HTTPHandler(ctx context.Context) (http.Handler, error) { if err != nil { return nil, fmt.Errorf("failed to create transcoder: %w", err) } + + // Add specific route for GetCurrentUser with cookie refresh + mux.Handle("/v1/users/current", s.authenticator.HTTPMiddlewareWithCookieRefresh(transcoder)) + mux.Handle("/v1/", transcoder) mux.Handle("/rill.admin.v1.AdminService/", transcoder) mux.Handle("/rill.admin.v1.AIService/", transcoder) From 22f5ac77be6001ea40988878cc4fb443f09094ab Mon Sep 17 00:00:00 2001 From: Naman Mahor <45662003+NamanMahor@users.noreply.github.com> Date: Wed, 6 Aug 2025 15:06:37 +0530 Subject: [PATCH 2/5] Update admin/server/auth/middleware.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Benjamin Egelund-Müller --- admin/server/auth/middleware.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/admin/server/auth/middleware.go b/admin/server/auth/middleware.go index 292ae196024..943ee156745 100644 --- a/admin/server/auth/middleware.go +++ b/admin/server/auth/middleware.go @@ -90,8 +90,10 @@ func (a *Authenticator) HTTPMiddlewareLenient(next http.Handler) http.Handler { return a.httpMiddleware(next, true) } -// HTTPMiddlewareWithCookieRefresh is a middleware that refreshes the auth cookie -func (a *Authenticator) HTTPMiddlewareWithCookieRefresh(next http.Handler) http.Handler { +// CookieRefreshMiddleware is a middleware that refreshes the auth cookie. +// This enables us to do rolling cookie refreshes so we can have a relatively short cookie max age. +// Note that it does not update the auth token encrypted inside the cookie. +func (a *Authenticator) CookieRefreshMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { sess := a.cookies.Get(r, cookieName) if authToken, ok := sess.Values[cookieFieldAccessToken].(string); ok && authToken != "" { From e1958fe04b6375a0e00f2b6268b3b829eaa3e5e8 Mon Sep 17 00:00:00 2001 From: Naman Mahor <45662003+NamanMahor@users.noreply.github.com> Date: Wed, 6 Aug 2025 15:06:46 +0530 Subject: [PATCH 3/5] Update admin/server/server.go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Benjamin Egelund-Müller --- admin/server/server.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/admin/server/server.go b/admin/server/server.go index 6001c207ab9..863cf4d3c48 100644 --- a/admin/server/server.go +++ b/admin/server/server.go @@ -186,7 +186,8 @@ func (s *Server) HTTPHandler(ctx context.Context) (http.Handler, error) { return nil, fmt.Errorf("failed to create transcoder: %w", err) } - // Add specific route for GetCurrentUser with cookie refresh + // Add auth cookie refresh specifically for the GetCurrentUser RPC. + // This is sufficient to refresh the cookie on each page load without unnecessarily refreshing cookies in each API call. mux.Handle("/v1/users/current", s.authenticator.HTTPMiddlewareWithCookieRefresh(transcoder)) mux.Handle("/v1/", transcoder) From 6bb4968cb4c355b06fd94bb30771df93e437646c Mon Sep 17 00:00:00 2001 From: NamanMahor Date: Wed, 6 Aug 2025 20:54:51 +0530 Subject: [PATCH 4/5] review --- admin/server/server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/admin/server/server.go b/admin/server/server.go index 863cf4d3c48..56edafc5aee 100644 --- a/admin/server/server.go +++ b/admin/server/server.go @@ -188,7 +188,7 @@ func (s *Server) HTTPHandler(ctx context.Context) (http.Handler, error) { // Add auth cookie refresh specifically for the GetCurrentUser RPC. // This is sufficient to refresh the cookie on each page load without unnecessarily refreshing cookies in each API call. - mux.Handle("/v1/users/current", s.authenticator.HTTPMiddlewareWithCookieRefresh(transcoder)) + mux.Handle("/v1/users/current", s.authenticator.CookieRefreshMiddleware(transcoder)) mux.Handle("/v1/", transcoder) mux.Handle("/rill.admin.v1.AdminService/", transcoder) From a141656c70d9635834257c86a029e36f979627ba Mon Sep 17 00:00:00 2001 From: NamanMahor Date: Fri, 8 Aug 2025 17:08:59 +0530 Subject: [PATCH 5/5] authWithToken now uses nonce_token --- admin/server/auth/handlers.go | 48 ++++++++++++++++++++++++++++------- admin/urls.go | 4 +-- 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/admin/server/auth/handlers.go b/admin/server/auth/handlers.go index a53e87205bd..a885127f079 100644 --- a/admin/server/auth/handlers.go +++ b/admin/server/auth/handlers.go @@ -9,6 +9,7 @@ import ( "net/http" "net/url" "strconv" + "time" "github.com/coreos/go-oidc/v3/oidc" "github.com/rilldata/rill/admin/database" @@ -160,8 +161,8 @@ func (a *Authenticator) authStart(w http.ResponseWriter, r *http.Request, signup // // authLoginCallback doesn't set the auth cookie directly because the final redirect destination may be an org with a custom domain. // So we need to ensure that the auth token is set in a cookie on the custom domain instead of the primary domain. -// So after creating a new token, it redirects to authWithToken on the same domain as the final redirect destination (if any, else it stays on the same domain). -// authWithToken will then set the actual auth cookie and do the final redirect back to the UI. +// So we issue a short-lived nonce token for the browser auth callback, and then redirect to authWithToken on the same domain as the final redirect destination (if any, else it stays on the same domain). +// authWithToken will then validate the nonce token and create a new auth token and set the actual auth cookie and do the final redirect back to the UI. func (a *Authenticator) authLoginCallback(w http.ResponseWriter, r *http.Request) { // Get auth cookie sess := a.cookies.Get(r, cookieName) @@ -261,8 +262,9 @@ func (a *Authenticator) authLoginCallback(w http.ResponseWriter, r *http.Request return } - // Issue a new persistent auth token - authToken, err := a.admin.IssueUserAuthToken(r.Context(), user.ID, database.AuthClientIDRillWeb, "Browser session", nil, nil) + // Issue a short-lived nonce token (2-minute TTL) for browser auth callback + ttl := 2 * time.Minute + authNonceToken, err := a.admin.IssueUserAuthToken(r.Context(), user.ID, database.AuthClientIDRillWeb, "Nonce Token", nil, &ttl) if err != nil { http.Error(w, fmt.Sprintf("failed to issue API token: %s", err), http.StatusInternalServerError) return @@ -274,19 +276,47 @@ func (a *Authenticator) authLoginCallback(w http.ResponseWriter, r *http.Request return } - // Redirect to authWithToken to set the token in a cookie. + // Redirect to authWithToken to validate the nonce token and set the token in a cookie. // We don't set the cookie directly here because the auth flow may have started from an org with a custom domain (see authStart), // in which case the token needs to be set in a cookie on the custom domain instead of the primary domain (where authLoginCallback is called). - tokenStr := authToken.Token().String() - withTokenURL := a.admin.URLs.WithCustomDomainFromRedirectURL(redirect).AuthWithToken(tokenStr, redirect) + nonceTokenStr := authNonceToken.Token().String() + withTokenURL := a.admin.URLs.WithCustomDomainFromRedirectURL(redirect).AuthWithToken(nonceTokenStr, redirect) http.Redirect(w, r, withTokenURL, http.StatusTemporaryRedirect) } // authWithToken extracts an auth token from the query params and sets it in the auth cookie. // It then redirects the user to the frontend. The redirect destination can be overridden by passing a "redirect" query parameter. func (a *Authenticator) authWithToken(w http.ResponseWriter, r *http.Request) { - // Get new auth token - newToken := r.URL.Query().Get("token") + nonceToken := r.URL.Query().Get("nonce_token") + var newToken string + if nonceToken != "" { + validated, err := a.admin.ValidateAuthToken(r.Context(), nonceToken) + if err != nil { + http.Error(w, err.Error(), http.StatusUnauthorized) + return + } + mdl, ok := validated.TokenModel().(*database.UserAuthToken) + if !ok { + http.Error(w, "invalid token model", http.StatusBadRequest) + return + } + newAuthToken, err := a.admin.IssueUserAuthToken(r.Context(), mdl.UserID, database.AuthClientIDRillWeb, "Browser session", nil, nil) + if err != nil { + http.Error(w, fmt.Sprintf("failed to issue API token: %s", err), http.StatusInternalServerError) + return + } + newToken = newAuthToken.Token().String() + + err = a.admin.RevokeAuthToken(r.Context(), nonceToken) + if err != nil { + a.logger.Info("failed to revoke nonce token during auth", zap.Error(err), observability.ZapCtx(r.Context())) + // The nonce token was probably manually revoked. We can still continue. + } + } else { + // Get new auth token + newToken = r.URL.Query().Get("token") + } + if newToken == "" { http.Error(w, "token not provided", http.StatusBadRequest) return diff --git a/admin/urls.go b/admin/urls.go index d0cb4102d28..eaca786fee3 100644 --- a/admin/urls.go +++ b/admin/urls.go @@ -156,11 +156,11 @@ func (u *URLs) AuthLogoutCallback() string { return urlutil.MustJoinURL(u.external, "/auth/logout/callback") // NOTE: Always using the primary external URL. } -// AuthWithToken returns a URL that sets the auth cookie to the provided token. +// AuthWithToken returns a URL that sets the auth cookie to the provided nonce token. // Providing a redirect URL is optional. func (u *URLs) AuthWithToken(tokenStr, redirect string) string { res := urlutil.MustJoinURL(u.External(), "/auth/with-token") // NOTE: Uses custom domain if set. - res = urlutil.MustWithQuery(res, map[string]string{"token": tokenStr}) + res = urlutil.MustWithQuery(res, map[string]string{"nonce_token": tokenStr}) if redirect != "" { res = urlutil.MustWithQuery(res, map[string]string{"redirect": redirect}) }