fix: use cookies to handle first login

cmd/serve.go

@@ -23,6 +23,7 @@ import (
 	"github.com/canonical/identity-platform-login-ui/internal/monitoring/prometheus"
 	fga "github.com/canonical/identity-platform-login-ui/internal/openfga"
 	"github.com/canonical/identity-platform-login-ui/internal/tracing"
+	"github.com/canonical/identity-platform-login-ui/pkg/kratos"
 	"github.com/canonical/identity-platform-login-ui/pkg/web"
 )
 
@@ -71,6 +72,13 @@ func serve() {
 	kAdminClient := ik.NewClient(specs.KratosAdminURL, specs.Debug)
 	hClient := ih.NewClient(specs.HydraAdminURL, specs.Debug)
 
+	encrypt := kratos.NewEncrypt([]byte(specs.CookiesEncryptionKey), logger, tracer)
+	cookieManager := kratos.NewAuthCookieManager(
+		specs.CookieTTL,
+		encrypt,
+		logger,
+	)
+
 	var authzClient authz.AuthzClientInterface
 	if specs.AuthorizationEnabled {
 		logger.Info("Authorization is enabled")
@@ -85,7 +93,7 @@ func serve() {
 		panic("Invalid authorization model provided")
 	}
 
-	router := web.NewRouter(kClient, kAdminClient, hClient, authorizer, distFS, specs.MFAEnabled, specs.BaseURL, tracer, monitor, logger)
+	router := web.NewRouter(kClient, kAdminClient, hClient, authorizer, cookieManager, distFS, specs.MFAEnabled, specs.BaseURL, tracer, monitor, logger)
 
 	logger.Infof("Starting server on port %v", specs.Port)
 

internal/config/specs.go

@@ -15,6 +15,9 @@ type EnvSpec struct {
 	Port    int    `envconfig:"port" default:"8080"`
 	BaseURL string `envconfig:"base_url" default:""`
 
+	CookiesEncryptionKey string `envconfig:"cookies_encryption_key" required:"true" validate:"required,min=32,max=32"`
+	CookieTTL            int    `envconfig:"cookie_ttl" default:"300"`
+
 	KratosPublicURL string `envconfig:"kratos_public_url"`
 	KratosAdminURL  string `envconfig:"kratos_admin_url"`
 	HydraAdminURL   string `envconfig:"hydra_admin_url"`

pkg/kratos/cookies.go

@@ -0,0 +1,123 @@
+// Copyright 2024 Canonical Ltd.
+// SPDX-License-Identifier: AGPL-3.0
+
+package kratos
+
+import (
+	"encoding/json"
+	"net/http"
+	"time"
+
+	"github.com/canonical/identity-platform-login-ui/internal/logging"
+)
+
+const (
+	defaultCookiePath = "/"
+	stateCookieName   = "login_ui_state"
+)
+
+var (
+	epoch = time.Unix(0, 0).UTC()
+)
+
+type AuthCookieManager struct {
+	cookieTTL time.Duration
+	encrypt   EncryptInterface
+
+	logger logging.LoggerInterface
+}
+
+type FlowStateCookie struct {
+	SessionId string `json:"login_challenge_id,omitempty"`
+	TotpSetup bool   `json:"totp_setup,omitempty"`
+}
+
+func (a *AuthCookieManager) SetStateCookie(w http.ResponseWriter, state FlowStateCookie) error {
+	rawState, err := json.Marshal(state)
+	if err != nil {
+		return err
+	}
+	return a.setCookie(w, stateCookieName, string(rawState), defaultCookiePath, a.cookieTTL, http.SameSiteLaxMode)
+}
+
+func (a *AuthCookieManager) GetStateCookie(r *http.Request) (FlowStateCookie, error) {
+	var ret FlowStateCookie
+	c, err := a.getCookie(r, stateCookieName)
+	if c == "" || err != nil {
+		return FlowStateCookie{}, err
+	}
+	err = json.Unmarshal([]byte(c), &ret)
+	return ret, err
+}
+
+func (a *AuthCookieManager) ClearStateCookie(w http.ResponseWriter) {
+	a.clearCookie(w, stateCookieName, defaultCookiePath)
+}
+
+func (a *AuthCookieManager) setCookie(w http.ResponseWriter, name, value string, path string, ttl time.Duration, sameSitePolicy http.SameSite) error {
+	if value == "" {
+		return nil
+	}
+
+	expires := time.Now().Add(ttl)
+
+	encrypted, err := a.encrypt.Encrypt(value)
+	if err != nil {
+		a.logger.Errorf("can't encrypt cookie value, %v", err)
+		return err
+	}
+
+	http.SetCookie(w, &http.Cookie{
+		Name:     name,
+		Value:    encrypted,
+		Path:     path,
+		Domain:   "",
+		Expires:  expires,
+		MaxAge:   int(ttl.Seconds()),
+		HttpOnly: true,
+		Secure:   true,
+		SameSite: sameSitePolicy,
+	})
+	return nil
+}
+
+func (a *AuthCookieManager) clearCookie(w http.ResponseWriter, name string, path string) {
+	http.SetCookie(w, &http.Cookie{
+		Name:     name,
+		Value:    "",
+		Path:     path,
+		Domain:   "",
+		Expires:  epoch,
+		MaxAge:   -1,
+		Secure:   true,
+		HttpOnly: true,
+	})
+}
+
+func (a *AuthCookieManager) getCookie(r *http.Request, name string) (string, error) {
+	cookie, err := r.Cookie(name)
+	if err != nil {
+		return "", err
+	}
+
+	value, err := a.encrypt.Decrypt(cookie.Value)
+	if err != nil {
+		a.logger.Errorf("can't decrypt cookie value, %v", err)
+		return "", err
+	}
+	return value, nil
+}
+
+func NewAuthCookieManager(
+	cookieTTLSeconds int,
+	encrypt EncryptInterface,
+	logger logging.LoggerInterface,
+) *AuthCookieManager {
+	a := new(AuthCookieManager)
+	a.cookieTTL = time.Duration(cookieTTLSeconds) * time.Second
+	a.encrypt = encrypt
+
+	a.logger = logger
+	return a
+
+}

pkg/kratos/cookies_test.go

@@ -0,0 +1,175 @@
+// Copyright 2024 Canonical Ltd.
+// SPDX-License-Identifier: AGPL-3.0
+
+package kratos
+
+import (
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"go.uber.org/mock/gomock"
+)
+
+//go:generate mockgen -build_flags=--mod=mod -package kratos -destination ./mock_logger.go -source=../../internal/logging/interfaces.go
+//go:generate mockgen -build_flags=--mod=mod -package kratos -destination ./mock_interfaces.go -source=./interfaces.go
+//go:generate mockgen -build_flags=--mod=mod -package kratos -destination ./mock_tracing.go go.opentelemetry.io/otel/trace Tracer
+//go:generate mockgen -build_flags=--mod=mod -package kratos -destination ./mock_monitor.go -source=../../internal/monitoring/interfaces.go
+
+func findCookie(name string, cookies []*http.Cookie) (*http.Cookie, bool) {
+	for _, cookie := range cookies {
+		if name == cookie.Name {
+			return cookie, true
+		}
+	}
+
+	return nil, false
+}
+
+func TestAuthCookieManager_ClearStateCookie(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockEncrypt := NewMockEncryptInterface(ctrl)
+
+	mockRequest := httptest.NewRequest(http.MethodGet, "/", nil)
+	mockRequest.AddCookie(&http.Cookie{Name: "state"})
+
+	mockResponse := httptest.NewRecorder()
+
+	manager := NewAuthCookieManager(5, mockEncrypt, mockLogger)
+	manager.ClearStateCookie(mockResponse)
+
+	c, _ := findCookie("login_ui_state", mockResponse.Result().Cookies())
+
+	if c.Expires != epoch {
+		t.Fatal("did not clear state cookie")
+	}
+}
+
+func TestAuthCookieManager_GetStateCookie(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockEncrypt := NewMockEncryptInterface(ctrl)
+
+	state := FlowStateCookie{}
+	sj, _ := json.Marshal(state)
+
+	mockEncrypt.EXPECT().Decrypt("mock-state").Return(string(sj), nil)
+
+	mockRequest := httptest.NewRequest(http.MethodGet, "/", nil)
+	mockRequest.AddCookie(&http.Cookie{Name: "login_ui_state", Value: "mock-state"})
+
+	manager := NewAuthCookieManager(5, mockEncrypt, mockLogger)
+	cookie, err := manager.GetStateCookie(mockRequest)
+
+	if cookie != state {
+		t.Fatal("state cookie value does not match expected")
+	}
+
+	if err != nil {
+		t.Fatalf("expected error to be nil not  %v", err)
+	}
+}
+
+func TestAuthCookieManager_GetStateCookieFailure(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockRequest := httptest.NewRequest(http.MethodGet, "/", nil)
+
+	manager := NewAuthCookieManager(5, nil, mockLogger)
+	cookie, err := manager.GetStateCookie(mockRequest)
+
+	state := FlowStateCookie{}
+	if cookie != state {
+		t.Fatal("state cookie value does not match expected")
+	}
+
+	if err == nil {
+		t.Fatal("expected error to be not nil")
+	}
+}
+
+func TestAuthCookieManager_GetStateCookieDecryptFailure(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockError := errors.New("mock-error")
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockLogger.EXPECT().Errorf("can't decrypt cookie value, %v", mockError).Times(1)
+
+	mockEncrypt := NewMockEncryptInterface(ctrl)
+	mockEncrypt.EXPECT().Decrypt("mock-state").Return("", mockError)
+
+	mockRequest := httptest.NewRequest(http.MethodGet, "/", nil)
+	mockRequest.AddCookie(&http.Cookie{Name: "login_ui_state", Value: "mock-state"})
+
+	manager := NewAuthCookieManager(5, mockEncrypt, mockLogger)
+	cookie, err := manager.GetStateCookie(mockRequest)
+
+	state := FlowStateCookie{}
+	if cookie != state {
+		t.Fatal("state cookie value does not match expected")
+	}
+
+	if err == nil {
+		t.Fatalf("expected error to be not nil")
+	}
+}
+
+func TestAuthCookieManager_SetStateCookie(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockEncrypt := NewMockEncryptInterface(ctrl)
+
+	state := FlowStateCookie{}
+	js, _ := json.Marshal(state)
+
+	mockEncrypt.EXPECT().Encrypt(string(js)).Return("mock-state", nil)
+
+	mockResponse := httptest.NewRecorder()
+
+	manager := NewAuthCookieManager(5, mockEncrypt, mockLogger)
+	err := manager.SetStateCookie(mockResponse, state)
+
+	c, found := findCookie("login_ui_state", mockResponse.Result().Cookies())
+
+	if !found {
+		t.Fatal("did not set state cookie")
+	}
+
+	if c.Value != "mock-state" {
+		t.Fatal("state cookie value does not match expected")
+	}
+
+	if err != nil {
+		t.Fatalf("expected error to be nil not  %v", err)
+	}
+}
+
+func TestAuthCookieManager_SetStateCookieFailure(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	mockError := errors.New("mock-error")
+	state := FlowStateCookie{}
+	js, _ := json.Marshal(state)
+
+	mockLogger := NewMockLoggerInterface(ctrl)
+	mockLogger.EXPECT().Errorf("can't encrypt cookie value, %v", mockError).Times(1)
+
+	mockEncrypt := NewMockEncryptInterface(ctrl)
+	mockEncrypt.EXPECT().Encrypt(string(js)).Return("", mockError)
+
+	mockResponse := httptest.NewRecorder()
+
+	manager := NewAuthCookieManager(5, mockEncrypt, mockLogger)
+	err := manager.SetStateCookie(mockResponse, state)
+
+	if err == nil {
+		t.Fatalf("expected error to be not nil")
+	}
+}