Remove domain match policy and instead use domain value for auth cook…

…ie config (flyteorg#446) * Remove domain match policy and instead use domain value for auth cookie config Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com> * lint fixes Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com> * keeping domain attribute as empty by default and only send domain on cookie if this attribute is set Signed-off-by: Prafulla Mahindrakar <prafulla.mahindrakar@gmail.com>
eapolinario · Jun 16, 2022 · abca8b6 · abca8b6
1 parent 3d80392
commit abca8b6
Show file tree

Hide file tree

Showing 13 changed files with 81 additions and 179 deletions.
diff --git a/flyteadmin/auth/authzserver/authorize.go b/flyteadmin/auth/authzserver/authorize.go
@@ -121,7 +121,7 @@ func authEndpoint(authCtx interfaces.AuthenticationContext, rw http.ResponseWrit
 		return
 	}
 
-	err = authCtx.CookieManager().SetAuthCodeCookie(ctx, req, rw, req.URL.String())
+	err = authCtx.CookieManager().SetAuthCodeCookie(ctx, rw, req.URL.String())
 	if err != nil {
 		logger.Infof(ctx, "Error occurred in NewAuthorizeRequest: %+v", err)
 		oauth2Provider.WriteAuthorizeError(rw, ar, err)

diff --git a/flyteadmin/auth/authzserver/authorize_test.go b/flyteadmin/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(), req, w, originalURL).Return(nil)
+		cookieManager.OnSetAuthCodeCookie(req.Context(), 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(), req, w, originalURL).Return(fmt.Errorf("failure injection"))
+		cookieManager.OnSetAuthCodeCookie(req.Context(), w, originalURL).Return(fmt.Errorf("failure injection"))
 		authCtx.OnCookieManager().Return(cookieManager)
 
 		authEndpoint(authCtx, w, req)

diff --git a/flyteadmin/auth/config/config.go b/flyteadmin/auth/config/config.go
@@ -74,8 +74,8 @@ var (
 				},
 			},
 			CookieSetting: CookieSettings{
-				DomainMatchPolicy: DomainMatchExact,
-				SameSitePolicy:    SameSiteDefaultMode,
+				Domain:         "",
+				SameSitePolicy: SameSiteDefaultMode,
 			},
 		},
 		AppAuth: OAuth2Options{
@@ -221,14 +221,6 @@ type UserAuthConfig struct {
 	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
 
@@ -240,8 +232,8 @@ const (
 )
 
 type CookieSettings struct {
-	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."`
+	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."`
+	Domain         string   `json:"domain" pflag:",OPTIONAL: Allows you to set the domain attribute on the auth cookies."`
 }
 
 type OpenIDOptions struct {

diff --git a/flyteadmin/auth/config/config_flags.go b/flyteadmin/auth/config/config_flags.go
diff --git a/flyteadmin/auth/config/config_flags_test.go b/flyteadmin/auth/config/config_flags_test.go
diff --git a/flyteadmin/auth/config/domainmatch_enumer.go b/flyteadmin/auth/config/domainmatch_enumer.go
diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go
@@ -55,10 +55,17 @@ func NewSecureCookie(cookieName, value string, hashKey, blockKey []byte, domain
 	var s = securecookie.New(hashKey, blockKey)
 	encoded, err := s.Encode(cookieName, value)
 	if err == nil {
+		if len(domain) > 0 {
+			return http.Cookie{
+				Name:     cookieName,
+				Value:    encoded,
+				Domain:   domain,
+				SameSite: sameSiteMode,
+			}, nil
+		}
 		return http.Cookie{
 			Name:     cookieName,
 			Value:    encoded,
-			Domain:   domain,
 			SameSite: sameSiteMode,
 		}, nil
 	}

diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go
@@ -17,10 +17,10 @@ import (
 )
 
 type CookieManager struct {
-	hashKey           []byte
-	blockKey          []byte
-	domainMatchPolicy config.DomainMatch
-	sameSitePolicy    config.SameSite
+	hashKey        []byte
+	blockKey       []byte
+	domain         string
+	sameSitePolicy config.SameSite
 }
 
 const (
@@ -45,10 +45,10 @@ func NewCookieManager(ctx context.Context, hashKeyEncoded, blockKeyEncoded strin
 	}
 
 	return CookieManager{
-		hashKey:           hashKey,
-		blockKey:          blockKey,
-		domainMatchPolicy: cookieSettings.DomainMatchPolicy,
-		sameSitePolicy:    cookieSettings.SameSitePolicy,
+		hashKey:        hashKey,
+		blockKey:       blockKey,
+		domain:         cookieSettings.Domain,
+		sameSitePolicy: cookieSettings.SameSitePolicy,
 	}, nil
 }
 
@@ -80,13 +80,13 @@ func (c CookieManager) RetrieveTokenValues(ctx context.Context, request *http.Re
 	return
 }
 
-func (c CookieManager) SetUserInfoCookie(ctx context.Context, request *http.Request, writer http.ResponseWriter, userInfo *service.UserInfoResponse) error {
+func (c CookieManager) SetUserInfoCookie(ctx context.Context, writer http.ResponseWriter, userInfo *service.UserInfoResponse) error {
 	raw, err := json.Marshal(userInfo)
 	if err != nil {
 		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.getHTTPSameSitePolicy())
+	userInfoCookie, err := NewSecureCookie(userInfoCookieName, string(raw), c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
 	if err != nil {
 		logger.Errorf(ctx, "Error generating encrypted user info cookie %s", err)
 		return err
@@ -123,8 +123,8 @@ func (c CookieManager) RetrieveAuthCodeRequest(ctx context.Context, request *htt
 	return authCodeCookie, nil
 }
 
-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.getHTTPSameSitePolicy())
+func (c CookieManager) SetAuthCodeCookie(ctx context.Context, writer http.ResponseWriter, authRequestURL string) error {
+	authCodeCookie, err := NewSecureCookie(authCodeCookieName, authRequestURL, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
 	if err != nil {
 		logger.Errorf(ctx, "Error generating encrypted accesstoken cookie %s", err)
 		return err
@@ -135,13 +135,13 @@ func (c CookieManager) SetAuthCodeCookie(ctx context.Context, request *http.Requ
 	return nil
 }
 
-func (c CookieManager) SetTokenCookies(ctx context.Context, request *http.Request, writer http.ResponseWriter, token *oauth2.Token) error {
+func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) error {
 	if token == nil {
 		logger.Errorf(ctx, "Attempting to set cookies with nil token")
 		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.getHTTPSameSitePolicy())
+	atCookie, err := NewSecureCookie(accessTokenCookieName, token.AccessToken, c.hashKey, c.blockKey, c.domain, 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.getHTTPSameSitePolicy())
+		idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.domain, 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, c.getCookieDomain(request), c.getHTTPSameSitePolicy())
+		refreshCookie, err := NewSecureCookie(refreshTokenCookieName, token.RefreshToken, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())
 		if err != nil {
 			logger.Errorf(ctx, "Error generating encrypted refresh token cookie %s", err)
 			return err
@@ -214,10 +214,3 @@ func (c CookieManager) getHTTPSameSitePolicy() http.SameSite {
 	}
 	return httpSameSite
 }
-
-func (c CookieManager) getCookieDomain(request *http.Request) string {
-	if c.domainMatchPolicy == config.DomainMatchExact {
-		return ""
-	}
-	return fmt.Sprintf(".%s", request.URL.Hostname())
-}
diff --git a/flyteadmin/auth/cookie_manager_test.go b/flyteadmin/auth/cookie_manager_test.go
@@ -20,8 +20,8 @@ func TestCookieManager_SetTokenCookies(t *testing.T) {
 	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
 	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
 	cookieSetting := config.CookieSettings{
-		SameSitePolicy:    config.SameSiteDefaultMode,
-		DomainMatchPolicy: config.DomainMatchSubdomains,
+		SameSitePolicy: config.SameSiteDefaultMode,
+		Domain:         "default",
 	}
 	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
 	assert.NoError(t, err)
@@ -36,9 +36,9 @@ func TestCookieManager_SetTokenCookies(t *testing.T) {
 	})
 
 	w := httptest.NewRecorder()
-	req, err := http.NewRequest("GET", "/api/v1/projects", nil)
+	_, err = http.NewRequest("GET", "/api/v1/projects", nil)
 	assert.NoError(t, err)
-	err = manager.SetTokenCookies(ctx, req, w, token)
+	err = manager.SetTokenCookies(ctx, w, token)
 	assert.NoError(t, err)
 	fmt.Println(w.Header().Get("Set-Cookie"))
 	c := w.Result().Cookies()
@@ -54,8 +54,8 @@ func TestCookieManager_RetrieveTokenValues(t *testing.T) {
 	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
 
 	cookieSetting := config.CookieSettings{
-		SameSitePolicy:    config.SameSiteDefaultMode,
-		DomainMatchPolicy: config.DomainMatchSubdomains,
+		SameSitePolicy: config.SameSiteDefaultMode,
+		Domain:         "default",
 	}
 
 	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
@@ -71,13 +71,13 @@ func TestCookieManager_RetrieveTokenValues(t *testing.T) {
 	})
 
 	w := httptest.NewRecorder()
-	req, err := http.NewRequest("GET", "/api/v1/projects", nil)
+	_, err = http.NewRequest("GET", "/api/v1/projects", nil)
 	assert.NoError(t, err)
-	err = manager.SetTokenCookies(ctx, req, w, token)
+	err = manager.SetTokenCookies(ctx, w, token)
 	assert.NoError(t, err)
 
 	cookies := w.Result().Cookies()
-	req, err = http.NewRequest("GET", "/api/v1/projects", nil)
+	req, err := http.NewRequest("GET", "/api/v1/projects", nil)
 	assert.NoError(t, err)
 	for _, c := range cookies {
 		req.AddCookie(c)
@@ -107,8 +107,8 @@ func TestCookieManager_DeleteCookies(t *testing.T) {
 	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
 	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
 	cookieSetting := config.CookieSettings{
-		SameSitePolicy:    config.SameSiteDefaultMode,
-		DomainMatchPolicy: config.DomainMatchSubdomains,
+		SameSitePolicy: config.SameSiteDefaultMode,
+		Domain:         "default",
 	}
 
 	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
@@ -129,8 +129,8 @@ func TestGetHTTPSameSitePolicy(t *testing.T) {
 	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
 	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
 	cookieSetting := config.CookieSettings{
-		SameSitePolicy:    config.SameSiteDefaultMode,
-		DomainMatchPolicy: config.DomainMatchSubdomains,
+		SameSitePolicy: config.SameSiteDefaultMode,
+		Domain:         "default",
 	}
 
 	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
@@ -146,25 +146,3 @@ func TestGetHTTPSameSitePolicy(t *testing.T) {
 	manager.sameSitePolicy = config.SameSiteNoneMode
 	assert.Equal(t, http.SameSiteNoneMode, manager.getHTTPSameSitePolicy())
 }
-
-func TestGetCookieDomain(t *testing.T) {
-	ctx := context.Background()
-
-	// These were generated for unit testing only.
-	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
-	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
-	cookieSetting := config.CookieSettings{
-		SameSitePolicy:    config.SameSiteDefaultMode,
-		DomainMatchPolicy: config.DomainMatchSubdomains,
-	}
-
-	req, err := http.NewRequest("GET", "http://localhost/api/v1/projects", nil)
-	assert.NoError(t, err)
-
-	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
-	assert.NoError(t, err)
-	assert.Equal(t, ".localhost", manager.getCookieDomain(req))
-
-	manager.domainMatchPolicy = config.DomainMatchExact
-	assert.Equal(t, "", manager.getCookieDomain(req))
-}