From 19095ccfea8defe8dba40cf71b4eb51fd4f12322 Mon Sep 17 00:00:00 2001
From: Trong Huu Nguyen <trong.huu.nguyen@nav.no>
Date: Sat, 29 Apr 2023 09:09:02 +0200
Subject: [PATCH] feat(openid): store acr in state cookie

---
 pkg/openid/client/login.go | 43 +++++++++++++++++++++-----------------
 pkg/openid/cookies.go      |  1 +
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/pkg/openid/client/login.go b/pkg/openid/client/login.go
index 3d67eaf..c4647a5 100644
--- a/pkg/openid/client/login.go
+++ b/pkg/openid/client/login.go
@@ -40,16 +40,26 @@ var (
 )
 
 func NewLogin(c *Client, r *http.Request) (*Login, error) {
-	params, err := newLoginParameters()
-	if err != nil {
-		return nil, fmt.Errorf("generating parameters: %w", err)
-	}
-
 	callbackURL, err := urlpkg.LoginCallback(r)
 	if err != nil {
 		return nil, fmt.Errorf("generating callback url: %w", err)
 	}
 
+	acr, err := getParameterOrDefault(r, SecurityLevelURLParameter, c.cfg.Client().ACRValues(), c.cfg.Provider().ACRValuesSupported())
+	if err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrInvalidSecurityLevel, err)
+	}
+
+	locale, err := getParameterOrDefault(r, LocaleURLParameter, c.cfg.Client().UILocales(), c.cfg.Provider().UILocalesSupported())
+	if err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrInvalidLocale, err)
+	}
+
+	params, err := newLoginParameters(acr, callbackURL)
+	if err != nil {
+		return nil, fmt.Errorf("generating parameters: %w", err)
+	}
+
 	opts := []oauth2.AuthCodeOption{
 		oauth2.SetAuthURLParam(openid.Nonce, params.Nonce),
 		oauth2.SetAuthURLParam(openid.ResponseMode, ResponseModeQuery),
@@ -63,16 +73,6 @@ func NewLogin(c *Client, r *http.Request) (*Login, error) {
 		opts = append(opts, oauth2.SetAuthURLParam(openid.Resource, resourceIndicator))
 	}
 
-	acr, err := getParameterOrDefault(r, SecurityLevelURLParameter, c.cfg.Client().ACRValues(), c.cfg.Provider().ACRValuesSupported())
-	if err != nil {
-		return nil, fmt.Errorf("%w: %w", ErrInvalidSecurityLevel, err)
-	}
-
-	locale, err := getParameterOrDefault(r, LocaleURLParameter, c.cfg.Client().UILocales(), c.cfg.Provider().UILocalesSupported())
-	if err != nil {
-		return nil, fmt.Errorf("%w: %w", ErrInvalidLocale, err)
-	}
-
 	if len(acr) > 0 {
 		opts = append(opts, oauth2.SetAuthURLParam(LoginParameterMapping[SecurityLevelURLParameter], acr))
 	}
@@ -83,7 +83,7 @@ func NewLogin(c *Client, r *http.Request) (*Login, error) {
 
 	return &Login{
 		authCodeURL: c.oauth2Config.AuthCodeURL(params.State, opts...),
-		cookie:      params.cookie(callbackURL),
+		cookie:      params.cookie(),
 		params:      params,
 	}, nil
 }
@@ -139,13 +139,15 @@ func (l *Login) SetCookie(w http.ResponseWriter, opts cookie.Options, crypter cr
 }
 
 type loginParameters struct {
+	Acr           string
 	CodeVerifier  string
 	CodeChallenge string
 	Nonce         string
+	RedirectURI   string
 	State         string
 }
 
-func newLoginParameters() (*loginParameters, error) {
+func newLoginParameters(acr, redirectUri string) (*loginParameters, error) {
 	codeVerifier, err := strings.GenerateBase64(64)
 	if err != nil {
 		return nil, fmt.Errorf("creating code verifier: %w", err)
@@ -162,19 +164,22 @@ func newLoginParameters() (*loginParameters, error) {
 	}
 
 	return &loginParameters{
+		Acr:           acr,
 		CodeVerifier:  codeVerifier,
 		CodeChallenge: CodeChallenge(codeVerifier),
 		Nonce:         nonce,
+		RedirectURI:   redirectUri,
 		State:         state,
 	}, nil
 }
 
-func (in *loginParameters) cookie(redirectURI string) *openid.LoginCookie {
+func (in *loginParameters) cookie() *openid.LoginCookie {
 	return &openid.LoginCookie{
+		Acr:          in.Acr,
 		State:        in.State,
 		Nonce:        in.Nonce,
 		CodeVerifier: in.CodeVerifier,
-		RedirectURI:  redirectURI,
+		RedirectURI:  in.RedirectURI,
 	}
 }
 
diff --git a/pkg/openid/cookies.go b/pkg/openid/cookies.go
index d645322..6d2f417 100644
--- a/pkg/openid/cookies.go
+++ b/pkg/openid/cookies.go
@@ -11,6 +11,7 @@ import (
 )
 
 type LoginCookie struct {
+	Acr          string `json:"acr"`
 	CodeVerifier string `json:"code_verifier"`
 	Nonce        string `json:"nonce"`
 	RedirectURI  string `json:"redirect_uri"`