From 28b750517b13cfce3d35419782ca8105616a30a4 Mon Sep 17 00:00:00 2001
From: Trong Huu Nguyen <trong.huu.nguyen@nav.no>
Date: Wed, 29 Sep 2021 10:00:42 +0200
Subject: [PATCH] wip: cookies

---
 pkg/router/cookies.go | 53 +++++++++----------------------------------
 pkg/router/router.go  | 33 ++++++++++++++++-----------
 2 files changed, 31 insertions(+), 55 deletions(-)

diff --git a/pkg/router/cookies.go b/pkg/router/cookies.go
index eabdfc7..9d3b583 100644
--- a/pkg/router/cookies.go
+++ b/pkg/router/cookies.go
@@ -2,6 +2,7 @@ package router
 
 import (
 	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"time"
@@ -13,58 +14,26 @@ type Cookie struct {
 	expiresIn time.Duration
 }
 
-type CallbackCookies struct {
-	State        string
-	Nonce        string
-	CodeVerifier string
-	Referer      string
+type CallbackParams struct {
+	State        string `json:"state"`
+	Nonce        string `json:"nonce"`
+	CodeVerifier string `json:"code_verifier"`
+	Referer      string `json:"referer"`
 }
 
-func NewCookie(name, value string, expiresIn time.Duration) Cookie {
-	return Cookie{
-		name:      name,
-		value:     value,
-		expiresIn: expiresIn,
-	}
-}
-
-func (h *Handler) getCallbackCookies(r *http.Request) (*CallbackCookies, error) {
-	state, err := h.getEncryptedCookie(r, StateCookieName)
+func (h *Handler) getCallbackParams(r *http.Request) (*CallbackParams, error) {
+	callbackCookieString, err := h.getEncryptedCookie(r, CallbackCookieName)
 	if err != nil {
 		return nil, err
 	}
 
-	nonce, err := h.getEncryptedCookie(r, NonceCookieName)
+	var callbackParams CallbackParams
+	err = json.Unmarshal([]byte(callbackCookieString), &callbackParams)
 	if err != nil {
 		return nil, err
 	}
 
-	codeVerifier, err := h.getEncryptedCookie(r, CodeVerifierCookieName)
-	if err != nil {
-		return nil, err
-	}
-
-	referer, err := h.getEncryptedCookie(r, RedirectURLCookieName)
-	if err != nil {
-		return nil, err
-	}
-
-	return &CallbackCookies{
-		State:        state,
-		Nonce:        nonce,
-		CodeVerifier: codeVerifier,
-		Referer:      referer,
-	}, nil
-}
-
-func (h *Handler) setEncryptedCookies(w http.ResponseWriter, cookies ...Cookie) error {
-	for _, cookie := range cookies {
-		err := h.setEncryptedCookie(w, cookie.name, cookie.value, cookie.expiresIn)
-		if err != nil {
-			return err
-		}
-	}
-	return nil
+	return &callbackParams, nil
 }
 
 func (h *Handler) setEncryptedCookie(w http.ResponseWriter, key string, plaintext string, expiresIn time.Duration) error {
diff --git a/pkg/router/router.go b/pkg/router/router.go
index a6d8b46..265d506 100644
--- a/pkg/router/router.go
+++ b/pkg/router/router.go
@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -30,13 +31,10 @@ import (
 )
 
 const (
-	LoginCookieLifetime = 10 * time.Minute
+	SessionCookieName = "io.nais.wonderwall.session"
 
-	SessionCookieName      = "io.nais.wonderwall.session"
-	StateCookieName        = "io.nais.wonderwall.state"
-	NonceCookieName        = "io.nais.wonderwall.nonce"
-	CodeVerifierCookieName = "io.nais.wonderwall.code_verifier"
-	RedirectURLCookieName  = "io.nais.wonderwall.redirect_url"
+	LoginCookieLifetime = 2 * time.Minute
+	CallbackCookieName  = "io.nais.wonderwall.callback"
 
 	RedirectURLParameter           = "redirect"
 	SecurityLevelURLParameter      = "level"
@@ -218,12 +216,21 @@ func (h *Handler) Login(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = h.setEncryptedCookies(w,
-		NewCookie(StateCookieName, params.state, LoginCookieLifetime),
-		NewCookie(NonceCookieName, params.nonce, LoginCookieLifetime),
-		NewCookie(CodeVerifierCookieName, params.codeVerifier, LoginCookieLifetime),
-		NewCookie(RedirectURLCookieName, CanonicalRedirectURL(r), LoginCookieLifetime),
-	)
+	callbackCookies := &CallbackParams{
+		State:        params.state,
+		Nonce:        params.nonce,
+		CodeVerifier: params.codeVerifier,
+		Referer:      CanonicalRedirectURL(r),
+	}
+
+	jsonString, err := json.Marshal(callbackCookies)
+	if err != nil {
+		log.Error(err)
+		w.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	err = h.setEncryptedCookie(w, CallbackCookieName, string(jsonString), LoginCookieLifetime)
 	if err != nil {
 		log.Error(err)
 		w.WriteHeader(http.StatusInternalServerError)
@@ -234,7 +241,7 @@ func (h *Handler) Login(w http.ResponseWriter, r *http.Request) {
 }
 
 func (h *Handler) Callback(w http.ResponseWriter, r *http.Request) {
-	cookies, err := h.getCallbackCookies(r)
+	cookies, err := h.getCallbackParams(r)
 	if err != nil {
 		log.Error(err)
 		w.WriteHeader(http.StatusUnauthorized)