From dbc0a47a46324eb28541b5d4ec9a1e2188c1b042 Mon Sep 17 00:00:00 2001
From: Trong Huu Nguyen <trong.huu.nguyen@nav.no>
Date: Thu, 30 Sep 2021 11:44:15 +0200
Subject: [PATCH] refactor: ensure session lifetime does not exceed access
 token lifetime
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: Sindre Rødseth Hansen <sindre.rodseth.hansen@nav.no>
---
 pkg/router/router.go  | 27 +++------------------
 pkg/router/session.go | 55 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 24 deletions(-)

diff --git a/pkg/router/router.go b/pkg/router/router.go
index 87a1746..a6ec58d 100644
--- a/pkg/router/router.go
+++ b/pkg/router/router.go
@@ -80,15 +80,6 @@ func (h *Handler) WithSecureCookie(enabled bool) *Handler {
 	return h
 }
 
-// localSessionID prefixes the given `sid` with the given client ID to prevent key collisions.
-// `sid` is a key that refers to the user's unique SSO session at the Identity Provider, and the same key is present
-// in all tokens acquired by any Relying Party (such as Wonderwall) during that session.
-// Thus, we cannot assume that the value of `sid` to uniquely identify the pair of (user, application session)
-// if using a shared session store.
-func (h *Handler) localSessionID(sid string) string {
-	return fmt.Sprintf("%s-%s", h.Config.ClientID, sid)
-}
-
 type loginParams struct {
 	state        string
 	codeVerifier string
@@ -300,22 +291,10 @@ func (h *Handler) Callback(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusUnauthorized)
 		return
 	}
-	sessionID := h.localSessionID(externalSessionID)
 
-	err = h.setEncryptedCookie(w, h.GetSessionCookieName(), sessionID, h.Config.SessionMaxLifetime)
+	err = h.createSession(w, r, externalSessionID, tokens, idToken)
 	if err != nil {
-		log.Error(err)
-		w.WriteHeader(http.StatusInternalServerError)
-		return
-	}
-
-	err = h.Sessions.Write(r.Context(), sessionID, &session.Data{
-		ExternalSessionID: externalSessionID,
-		OAuth2Token:       tokens,
-		IDTokenSerialized: idToken.Raw,
-	}, h.Config.SessionMaxLifetime)
-	if err != nil {
-		log.Error(err)
+		log.Errorf("creating session: %+v", err)
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -323,7 +302,7 @@ func (h *Handler) Callback(w http.ResponseWriter, r *http.Request) {
 	http.Redirect(w, r, loginCookie.Referer, http.StatusTemporaryRedirect)
 }
 
-// Proxy all requests upstream
+// Default proxies all requests upstream
 func (h *Handler) Default(w http.ResponseWriter, r *http.Request) {
 	ctx, cancel := context.WithCancel(r.Context())
 	defer cancel()
diff --git a/pkg/router/session.go b/pkg/router/session.go
index ca2ed1e..e8a3169 100644
--- a/pkg/router/session.go
+++ b/pkg/router/session.go
@@ -2,10 +2,23 @@ package router
 
 import (
 	"fmt"
+	"github.com/lestrrat-go/jwx/jwt"
 	"github.com/nais/wonderwall/pkg/session"
+	"github.com/nais/wonderwall/pkg/token"
+	"golang.org/x/oauth2"
 	"net/http"
+	"time"
 )
 
+// localSessionID prefixes the given `sid` with the given client ID to prevent key collisions.
+// `sid` is a key that refers to the user's unique SSO session at the Identity Provider, and the same key is present
+// in all tokens acquired by any Relying Party (such as Wonderwall) during that session.
+// Thus, we cannot assume that the value of `sid` to uniquely identify the pair of (user, application session)
+// if using a shared session store.
+func (h *Handler) localSessionID(sid string) string {
+	return fmt.Sprintf("%s-%s", h.Config.ClientID, sid)
+}
+
 func (h *Handler) getSessionFromCookie(r *http.Request) (*session.Data, error) {
 	sessionID, err := h.getEncryptedCookie(r, h.GetSessionCookieName())
 	if err != nil {
@@ -14,3 +27,45 @@ func (h *Handler) getSessionFromCookie(r *http.Request) (*session.Data, error) {
 
 	return h.Sessions.Read(r.Context(), sessionID)
 }
+
+func (h *Handler) getSessionLifetime(accessToken string) (time.Duration, error) {
+	defaultSessionLifetime := h.Config.SessionMaxLifetime
+
+	token, err := jwt.Parse([]byte(accessToken))
+	if err != nil {
+		return 0, err
+	}
+
+	tokenDuration := token.Expiration().Sub(time.Now())
+
+	if tokenDuration <= defaultSessionLifetime {
+		return tokenDuration, nil
+	}
+
+	return defaultSessionLifetime, nil
+}
+
+func (h *Handler) createSession(w http.ResponseWriter, r *http.Request, externalSessionID string, tokens *oauth2.Token, idToken *token.IDToken) error {
+	sessionID := h.localSessionID(externalSessionID)
+
+	sessionLifetime, err := h.getSessionLifetime(tokens.AccessToken)
+	if err != nil {
+		return fmt.Errorf("getting access token lifetime: %w", err)
+	}
+
+	err = h.setEncryptedCookie(w, h.GetSessionCookieName(), sessionID, sessionLifetime)
+	if err != nil {
+		return fmt.Errorf("setting session cookie: %w", err)
+	}
+
+	err = h.Sessions.Write(r.Context(), sessionID, &session.Data{
+		ExternalSessionID: externalSessionID,
+		OAuth2Token:       tokens,
+		IDTokenSerialized: idToken.Raw,
+	}, sessionLifetime)
+	if err != nil {
+		return fmt.Errorf("writing session to store: %w", err)
+	}
+
+	return nil
+}