Store the IdP access token and expiry into the session (#154)

illumitacit · Jul 10, 2023 · cfd1d5b · cfd1d5b
1 parent 42a0a76
commit cfd1d5b
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 15 deletions.
diff --git a/webstd/chistd/oidc_handlers.go b/webstd/chistd/oidc_handlers.go
@@ -1,7 +1,9 @@
 package chistd
 
 import (
+	"encoding/gob"
 	"net/http"
+	"time"
 
 	"github.com/alexedwards/scs/v2"
 	"github.com/go-chi/chi/v5"
@@ -12,17 +14,24 @@ import (
 	"github.com/fensak-io/gostd/webstd"
 )
 
+func init() {
+	// Register time.Time to gob so it can be stored in the session
+	gob.Register(time.Time{})
+}
+
 const (
 	// URL paths
 	OIDCLoginPath    = "/oidc/login"
 	OIDCLogoutPath   = "/oidc/logout"
 	OIDCCallbackPath = "/oidc/callback"
 
 	// Session keys
-	RefreshTokenSessionKey     = "refresh_token"
-	UserProfileSessionKey      = "profile"
-	ContinueToURLSessionKey    = "continue_to"
-	PKCECodeVerifierSessionKey = "pkce_code_verifier"
+	AccessTokenSessionKey       = "access_token"
+	AccessTokenExpirySessionKey = "access_token_expiry"
+	RefreshTokenSessionKey      = "refresh_token"
+	UserProfileSessionKey       = "profile"
+	ContinueToURLSessionKey     = "continue_to"
+	PKCECodeVerifierSessionKey  = "pkce_code_verifier"
 )
 
 type OIDCHandlerContext[T any] struct {
@@ -159,6 +168,8 @@ func (h OIDCHandlerContext[T]) oidcCallbackHandler(w http.ResponseWriter, r *htt
 		return
 	}
 
+	h.sessMgr.Put(ctx, AccessTokenSessionKey, token.AccessToken)
+	h.sessMgr.Put(ctx, AccessTokenExpirySessionKey, token.Expiry)
 	h.sessMgr.Put(ctx, RefreshTokenSessionKey, token.RefreshToken)
 	h.sessMgr.Put(ctx, UserProfileSessionKey, profile)
 

diff --git a/webstd/idp/aadb2c/aadb2c.go b/webstd/idp/aadb2c/aadb2c.go
@@ -166,12 +166,14 @@ func (a AADB2C) GetLogoutURL(ctx context.Context) (string, error) {
 	}
 
 	refreshToken := a.sessMgr.GetString(ctx, chistd.RefreshTokenSessionKey)
-	rawIDToken, _, newRefreshToken, err := a.auth.RefreshIDToken(ctx, refreshToken)
+	rawIDToken, _, newToken, err := a.auth.RefreshIDToken(ctx, refreshToken)
 	if err != nil {
 		return "", err
 	}
-	if newRefreshToken != "" {
-		a.sessMgr.Put(ctx, chistd.RefreshTokenSessionKey, newRefreshToken)
+	a.sessMgr.Put(ctx, chistd.AccessTokenSessionKey, newToken.AccessToken)
+	a.sessMgr.Put(ctx, chistd.AccessTokenExpirySessionKey, newToken.Expiry)
+	if newToken.RefreshToken != "" {
+		a.sessMgr.Put(ctx, chistd.RefreshTokenSessionKey, newToken.RefreshToken)
 	}
 
 	appURLCopy := a.appURL

diff --git a/webstd/idp/zitadel/zitadel.go b/webstd/idp/zitadel/zitadel.go
@@ -133,12 +133,14 @@ func (z Zitadel) GetLogoutURL(ctx context.Context) (string, error) {
 	}
 
 	refreshToken := z.sessMgr.GetString(ctx, chistd.RefreshTokenSessionKey)
-	rawIDToken, _, newRefreshToken, err := z.auth.RefreshIDToken(ctx, refreshToken)
+	rawIDToken, _, newToken, err := z.auth.RefreshIDToken(ctx, refreshToken)
 	if err != nil {
 		return "", err
 	}
-	if newRefreshToken != "" {
-		z.sessMgr.Put(ctx, chistd.RefreshTokenSessionKey, newRefreshToken)
+	z.sessMgr.Put(ctx, chistd.AccessTokenSessionKey, newToken.AccessToken)
+	z.sessMgr.Put(ctx, chistd.AccessTokenExpirySessionKey, newToken.Expiry)
+	if newToken.RefreshToken != "" {
+		z.sessMgr.Put(ctx, chistd.RefreshTokenSessionKey, newToken.RefreshToken)
 	}
 
 	appURLCopy := z.appURL

diff --git a/webstd/oidc.go b/webstd/oidc.go
@@ -91,22 +91,22 @@ func (a Authenticator) VerifyIDToken(ctx context.Context, token *oauth2.Token) (
 }
 
 // RefreshIDToken obtains a new OIDC ID token using the provided refresh token.
-func (a Authenticator) RefreshIDToken(ctx context.Context, refreshToken string) (string, *oidc.IDToken, string, error) {
+func (a Authenticator) RefreshIDToken(ctx context.Context, refreshToken string) (string, *oidc.IDToken, *oauth2.Token, error) {
 	ts := a.TokenSource(ctx, &oauth2.Token{RefreshToken: refreshToken})
 	token, err := ts.Token()
 	if err != nil {
-		return "", nil, "", err
+		return "", nil, nil, err
 	}
 
 	rawIDToken, ok := token.Extra("id_token").(string)
 	if !ok {
-		return "", nil, "", errors.New("no id_token field in oauth2 token")
+		return "", nil, nil, errors.New("no id_token field in oauth2 token")
 	}
 	idToken, err := a.VerifyIDToken(ctx, token)
 	if err != nil {
-		return "", nil, "", err
+		return "", nil, nil, err
 	}
-	return rawIDToken, idToken, token.RefreshToken, nil
+	return rawIDToken, idToken, token, nil
 }
 
 // VerifyRawToken verifies a given raw JWT token string issued by the OIDC provider. This is useful for verifying tokens