From c3da8991903f7235bfaf2a99abde2c3ddcf4b0fe Mon Sep 17 00:00:00 2001
From: Trong Huu Nguyen <trong.huu.nguyen@nav.no>
Date: Tue, 1 Feb 2022 14:31:28 +0100
Subject: [PATCH] feat(loginstatus): ensure that cookie is set in default route

---
 pkg/loginstatus/loginstatus.go      | 10 ++++++++++
 pkg/loginstatus/loginstatus_test.go | 17 +++++++++++++++++
 pkg/router/handler_default.go       | 14 ++++++++++++--
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/pkg/loginstatus/loginstatus.go b/pkg/loginstatus/loginstatus.go
index 81a3e6d..1376a3a 100644
--- a/pkg/loginstatus/loginstatus.go
+++ b/pkg/loginstatus/loginstatus.go
@@ -3,6 +3,7 @@ package loginstatus
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -19,6 +20,7 @@ const (
 type Client interface {
 	ExchangeToken(ctx context.Context, accessToken string) (*TokenResponse, error)
 	SetCookie(w http.ResponseWriter, token *TokenResponse, opts cookie.Options)
+	HasCookie(r *http.Request) bool
 	ClearCookie(w http.ResponseWriter, opts cookie.Options)
 }
 
@@ -75,6 +77,14 @@ func (c client) SetCookie(w http.ResponseWriter, token *TokenResponse, opts cook
 	cookie.Set(w, newCookie)
 }
 
+func (c client) HasCookie(r *http.Request) bool {
+	_, err := r.Cookie(c.config.CookieName)
+	if errors.Is(err, http.ErrNoCookie) {
+		return false
+	}
+	return true
+}
+
 func (c client) ClearCookie(w http.ResponseWriter, opts cookie.Options) {
 	cookieName := c.config.CookieName
 	opts = c.cookieOptions(opts)
diff --git a/pkg/loginstatus/loginstatus_test.go b/pkg/loginstatus/loginstatus_test.go
index 103dee1..f2f9801 100644
--- a/pkg/loginstatus/loginstatus_test.go
+++ b/pkg/loginstatus/loginstatus_test.go
@@ -115,6 +115,23 @@ func TestClient_ClearCookie(t *testing.T) {
 	assert.Equal(t, "/", result.Path)
 }
 
+func TestClient_HasCookie(t *testing.T) {
+	cfg := newCfg("https://some-server")
+	opts := cookie.DefaultOptions()
+
+	c := cookie.Make(cfg.CookieName, "some-value", opts)
+	r := httptest.NewRequest(http.MethodGet, "/", nil)
+	r.AddCookie(c.Cookie)
+
+	client := loginstatus.NewClient(cfg, http.DefaultClient)
+	actual := client.HasCookie(r)
+	assert.True(t, actual)
+
+	r = httptest.NewRequest(http.MethodGet, "/", nil)
+	actual = client.HasCookie(r)
+	assert.False(t, actual)
+}
+
 func newCfg(serverURL string) config.Loginstatus {
 	return config.Loginstatus{
 		Enabled:           true,
diff --git a/pkg/router/handler_default.go b/pkg/router/handler_default.go
index 4042cd9..db68869 100644
--- a/pkg/router/handler_default.go
+++ b/pkg/router/handler_default.go
@@ -12,10 +12,20 @@ func (h *Handler) Default(w http.ResponseWriter, r *http.Request) {
 	isAuthenticated := false
 
 	sessionData, err := h.getSessionFromCookie(w, r)
-	if err == nil && sessionData != nil && len(sessionData.AccessToken) > 0 {
+
+	hasSessionData := err == nil && sessionData != nil
+	hasAccessToken := hasSessionData && len(sessionData.AccessToken) > 0
+	if hasAccessToken {
 		// add authentication if session cookie and token checks out
 		isAuthenticated = true
-	} else if h.Config.AutoLogin {
+
+		// force new authentication if loginstatus is enabled and cookie isn't set
+		if h.Config.Features.Loginstatus.Enabled && !h.Loginstatus.HasCookie(r) {
+			isAuthenticated = false
+		}
+	}
+
+	if !isAuthenticated && h.Config.AutoLogin {
 		r.Header.Add("Referer", r.URL.String())
 		h.Login(w, r)
 		return