Added enumers for domainMatch and sameSite and testing done

Signed-off-by: Prafulla Mahindrakar <[email protected]>

auth/authzserver/authorize_test.go

@@ -34,7 +34,7 @@ func TestAuthEndpoint(t *testing.T) {
 		authCtx.OnOAuth2Provider().Return(oauth2Provider)
 
 		cookieManager := &mocks.CookieHandler{}
-		cookieManager.OnSetAuthCodeCookie(req.Context(), w, originalURL).Return(nil)
+		cookieManager.OnSetAuthCodeCookie(req.Context(), req, w, originalURL).Return(nil)
 		authCtx.OnCookieManager().Return(cookieManager)
 
 		authEndpoint(authCtx, w, req)
@@ -57,7 +57,7 @@ func TestAuthEndpoint(t *testing.T) {
 		authCtx.OnOAuth2Provider().Return(oauth2Provider)
 
 		cookieManager := &mocks.CookieHandler{}
-		cookieManager.OnSetAuthCodeCookie(req.Context(), w, originalURL).Return(fmt.Errorf("failure injection"))
+		cookieManager.OnSetAuthCodeCookie(req.Context(), req, w, originalURL).Return(fmt.Errorf("failure injection"))
 		authCtx.OnCookieManager().Return(cookieManager)
 
 		authEndpoint(authCtx, w, req)

auth/config/config.go

@@ -1,7 +1,6 @@
 package config
 
 import (
-	"net/http"
 	"net/url"
 	"time"
 
@@ -74,9 +73,9 @@ var (
 					"profile",
 				},
 			},
-			CookieSetting: &CookieSettings{
-				CoverSubdomains: false,
-				SameSite:        http.SameSiteDefaultMode,
+			CookieSetting: CookieSettings{
+				DomainMatchPolicy: DomainMatchExact,
+				SameSitePolicy:    SameSiteDefaultMode,
 			},
 		},
 		AppAuth: OAuth2Options{
@@ -217,14 +216,32 @@ type UserAuthConfig struct {
 	// Possibly add basicAuth & SAML/p support.
 
 	// Secret names, defaults are set in DefaultConfig variable above but are possible to override through configs.
-	CookieHashKeySecretName  string          `json:"cookieHashKeySecretName" pflag:",OPTIONAL: Secret name to use for cookie hash key."`
-	CookieBlockKeySecretName string          `json:"cookieBlockKeySecretName" pflag:",OPTIONAL: Secret name to use for cookie block key."`
-	CookieSetting            *CookieSettings `json:"cookieSetting" pflag:", settings used by cookies created for user auth"`
+	CookieHashKeySecretName  string         `json:"cookieHashKeySecretName" pflag:",OPTIONAL: Secret name to use for cookie hash key."`
+	CookieBlockKeySecretName string         `json:"cookieBlockKeySecretName" pflag:",OPTIONAL: Secret name to use for cookie block key."`
+	CookieSetting            CookieSettings `json:"cookieSetting" pflag:", settings used by cookies created for user auth"`
 }
 
+//go:generate enumer --type=DomainMatch --trimprefix=DomainMatch -json
+type DomainMatch int
+
+const (
+	DomainMatchExact DomainMatch = iota
+	DomainMatchSubdomains
+)
+
+//go:generate enumer --type=SameSite --trimprefix=SameSite -json
+type SameSite int
+
+const (
+	SameSiteDefaultMode SameSite = iota
+	SameSiteLaxMode
+	SameSiteStrictMode
+	SameSiteNoneMode
+)
+
 type CookieSettings struct {
-	SameSite        http.SameSite `json:"sameSite" pflag:",OPTIONAL: Allows you to declare if your cookie should be restricted to a first-party or same-site context."`
-	CoverSubdomains bool          `json:"coverSubDomains" pflag:",OPTIONAL: Allow subdomain access to the created cookies by setting the domain attribute."`
+	SameSitePolicy    SameSite    `json:"sameSitePolicy" pflag:",OPTIONAL: Allows you to declare if your cookie should be restricted to a first-party or same-site context.Wrapper around http.SameSite."`
+	DomainMatchPolicy DomainMatch `json:"domainMatchPolicy" pflag:",OPTIONAL: Allow subdomain access to the created cookies by setting the domain attribute or do an exact match on domain."`
 }
 
 type OpenIDOptions struct {

auth/cookie_manager.go

@@ -5,22 +5,22 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"github.com/flyteorg/flyteadmin/auth/config"
 	"net/http"
 	"time"
 
+	"github.com/flyteorg/flyteadmin/auth/config"
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
-
 	"github.com/flyteorg/flytestdlib/errors"
 	"github.com/flyteorg/flytestdlib/logger"
+
 	"golang.org/x/oauth2"
 )
 
 type CookieManager struct {
-	hashKey         []byte
-	blockKey        []byte
-	coverSubDomains bool
-	sameSite        http.SameSite
+	hashKey           []byte
+	blockKey          []byte
+	domainMatchPolicy config.DomainMatch
+	sameSitePolicy    config.SameSite
 }
 
 const (
@@ -31,7 +31,7 @@ const (
 	ErrNoIDToken errors.ErrorCode = "NO_ID_TOKEN_IN_RESPONSE"
 )
 
-func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded string, cookieSettings *config.CookieSettings) (CookieManager, error) {
+func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded string, cookieSettings config.CookieSettings) (CookieManager, error) {
 	logger.Infof(ctx, "Instantiating cookie manager")
 
 	hashKey, err := base64.RawStdEncoding.DecodeString(hashKeyEncoded)
@@ -45,10 +45,10 @@ func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded strin
 	}
 
 	return CookieManager{
-		hashKey:         hashKey,
-		blockKey:        blockKey,
-		coverSubDomains: cookieSettings.CoverSubdomains,
-		sameSite:        cookieSettings.SameSite,
+		hashKey:           hashKey,
+		blockKey:          blockKey,
+		domainMatchPolicy: cookieSettings.DomainMatchPolicy,
+		sameSitePolicy:    cookieSettings.SameSitePolicy,
 	}, nil
 }
 
@@ -86,7 +86,7 @@ func (c CookieManager) SetUserInfoCookie(ctx context.Context, request *http.Requ
 		return fmt.Errorf("failed to marshal user info to store in a cookie. Error: %w", err)
 	}
 
-	userInfoCookie, err := NewSecureCookie(userInfoCookieName, string(raw), c.hashKey, c.blockKey, c.getCookieDomain(request), c.sameSite)
+	userInfoCookie, err := NewSecureCookie(userInfoCookieName, string(raw), c.hashKey, c.blockKey, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
 	if err != nil {
 		logger.Errorf(ctx, "Error generating encrypted user info cookie %s", err)
 		return err
@@ -124,7 +124,7 @@ func (c CookieManager) RetrieveAuthCodeRequest(ctx context.Context, request *htt
 }
 
 func (c CookieManager) SetAuthCodeCookie(ctx context.Context, request *http.Request, writer http.ResponseWriter, authRequestURL string) error {
-	authCodeCookie, err := NewSecureCookie(authCodeCookieName, authRequestURL, c.hashKey, c.blockKey, c.getCookieDomain(request), c.sameSite)
+	authCodeCookie, err := NewSecureCookie(authCodeCookieName, authRequestURL, c.hashKey, c.blockKey, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
 	if err != nil {
 		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie %s", err)
 		return err
@@ -141,7 +141,7 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, request *http.Reques
 		return errors.Errorf(ErrTokenNil, "Attempting to set cookies with nil token")
 	}
 
-	atCookie, err := NewSecureCookie(accessTokenCookieName, token.AccessToken, c.hashKey, c.blockKey, c.getCookieDomain(request), c.sameSite)
+	atCookie, err := NewSecureCookie(accessTokenCookieName, token.AccessToken, c.hashKey, c.blockKey, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
 	if err != nil {
 		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie %s", err)
 		return err
@@ -150,7 +150,7 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, request *http.Reques
 	http.SetCookie(writer, &atCookie)
 
 	if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted {
-		idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.getCookieDomain(request), c.sameSite)
+		idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
 		if err != nil {
 			logger.Errorf(ctx, "Error generating encrypted id token cookie %s", err)
 			return err
@@ -164,7 +164,7 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, request *http.Reques
 
 	// Set the refresh cookie if there is a refresh token
 	if token.RefreshToken != "" {
-		refreshCookie, err := NewSecureCookie(refreshTokenCookieName, token.RefreshToken, c.hashKey, c.blockKey, cookieDomain, c.sameSite)
+		refreshCookie, err := NewSecureCookie(refreshTokenCookieName, token.RefreshToken, c.hashKey, c.blockKey, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
 		if err != nil {
 			logger.Errorf(ctx, "Error generating encrypted refresh token cookie %s", err)
 			return err
@@ -200,8 +200,23 @@ func (c CookieManager) DeleteCookies(ctx context.Context, writer http.ResponseWr
 	http.SetCookie(writer, getLogoutRefreshCookie())
 }
 
+func (c CookieManager) getHTTPSameSitePolicy() http.SameSite {
+	httpSameSite := http.SameSiteDefaultMode
+	switch c.sameSitePolicy {
+	case config.SameSiteDefaultMode:
+		httpSameSite = http.SameSiteDefaultMode
+	case config.SameSiteLaxMode:
+		httpSameSite = http.SameSiteLaxMode
+	case config.SameSiteStrictMode:
+		httpSameSite = http.SameSiteStrictMode
+	case config.SameSiteNoneMode:
+		httpSameSite = http.SameSiteNoneMode
+	}
+	return httpSameSite
+}
+
 func (c CookieManager) getCookieDomain(request *http.Request) string {
-	if !c.coverSubDomains {
+	if c.domainMatchPolicy == config.DomainMatchExact {
 		return ""
 	}
 	return fmt.Sprintf(".%s", request.URL.Hostname())