Merge branch 'master' of https://github.com/flyteorg/flyteadmin into …

…k8s-events

Signed-off-by: Andrew Dye <[email protected]>

auth/cookie.go

@@ -9,10 +9,11 @@ import (
 	"net/url"
 	"time"
 
-	"github.com/flyteorg/flyteadmin/auth/interfaces"
 	"github.com/flyteorg/flytestdlib/errors"
 	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/gorilla/securecookie"
+
+	"github.com/flyteorg/flyteadmin/auth/interfaces"
 )
 
 const (
@@ -52,25 +53,18 @@ func HashCsrfState(csrf string) string {
 }
 
 func NewSecureCookie(cookieName, value string, hashKey, blockKey []byte, domain string, sameSiteMode http.SameSite) (http.Cookie, error) {
-	var s = securecookie.New(hashKey, blockKey)
+	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,
-			SameSite: sameSiteMode,
-		}, nil
-	}
-
-	return http.Cookie{}, errors.Wrapf(ErrSecureCookie, err, "Error creating secure cookie")
+	if err != nil {
+		return http.Cookie{}, errors.Wrapf(ErrSecureCookie, err, "Error creating secure cookie")
+	}
+
+	return http.Cookie{
+		Name:     cookieName,
+		Value:    encoded,
+		Domain:   domain,
+		SameSite: sameSiteMode,
+	}, nil
 }
 
 func retrieveSecureCookie(ctx context.Context, request *http.Request, cookieName string, hashKey, blockKey []byte) (string, error) {
@@ -169,10 +163,11 @@ func NewRedirectCookie(ctx context.Context, redirectURL string) *http.Cookie {
 	}
 }
 
+// GetAuthFlowEndRedirect returns the redirect URI according to data in request.
 // At the end of the OAuth flow, the server needs to send the user somewhere. This should have been stored as a cookie
 // during the initial /login call. If that cookie is missing from the request, it will default to the one configured
 // in this package's Config object.
-func getAuthFlowEndRedirect(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request) string {
+func GetAuthFlowEndRedirect(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request) string {
 	queryParams := request.URL.Query()
 	// Use the redirect URL specified in the request if one is available.
 	if redirectURL := queryParams.Get(RedirectURLParameter); len(redirectURL) > 0 {

auth/cookie_manager.go

@@ -8,12 +8,12 @@ import (
 	"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"
+
+	"github.com/flyteorg/flyteadmin/auth/config"
 )
 
 type CookieManager struct {
@@ -175,29 +175,21 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.Response
 	return nil
 }
 
-func getLogoutAccessCookie() *http.Cookie {
-	return &http.Cookie{
-		Name:     accessTokenCookieName,
-		Value:    "",
-		MaxAge:   0,
-		HttpOnly: true,
-		Expires:  time.Now().Add(-1 * time.Hour),
-	}
-}
-
-func getLogoutRefreshCookie() *http.Cookie {
+func (c *CookieManager) getLogoutCookie(name string) *http.Cookie {
 	return &http.Cookie{
-		Name:     refreshTokenCookieName,
+		Name:     name,
 		Value:    "",
+		Domain:   c.domain,
 		MaxAge:   0,
 		HttpOnly: true,
 		Expires:  time.Now().Add(-1 * time.Hour),
 	}
 }
 
-func (c CookieManager) DeleteCookies(ctx context.Context, writer http.ResponseWriter) {
-	http.SetCookie(writer, getLogoutAccessCookie())
-	http.SetCookie(writer, getLogoutRefreshCookie())
+func (c CookieManager) DeleteCookies(_ context.Context, writer http.ResponseWriter) {
+	http.SetCookie(writer, c.getLogoutCookie(accessTokenCookieName))
+	http.SetCookie(writer, c.getLogoutCookie(refreshTokenCookieName))
+	http.SetCookie(writer, c.getLogoutCookie(idTokenCookieName))
 }
 
 func (c CookieManager) getHTTPSameSitePolicy() http.SameSite {

auth/cookie_manager_test.go

@@ -1,20 +1,25 @@
 package auth
 
 import (
+	"bytes"
 	"context"
+	"encoding/base64"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 	"time"
 
-	"github.com/flyteorg/flyteadmin/auth/config"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
+
+	"github.com/flyteorg/flyteadmin/auth/config"
 )
 
-func TestCookieManager_SetTokenCookies(t *testing.T) {
+func TestCookieManager(t *testing.T) {
 	ctx := context.Background()
 	// These were generated for unit testing only.
 	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
@@ -25,124 +30,153 @@ func TestCookieManager_SetTokenCookies(t *testing.T) {
 	}
 	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
 	assert.NoError(t, err)
-
 	token := &oauth2.Token{
 		AccessToken:  "access",
 		RefreshToken: "refresh",
 	}
-
 	token = token.WithExtra(map[string]interface{}{
 		"id_token": "id token",
 	})
 
-	w := httptest.NewRecorder()
-	_, err = http.NewRequest("GET", "/api/v1/projects", nil)
-	assert.NoError(t, err)
-	err = manager.SetTokenCookies(ctx, w, token)
-	assert.NoError(t, err)
-	fmt.Println(w.Header().Get("Set-Cookie"))
-	c := w.Result().Cookies()
-	assert.Equal(t, "flyte_at", c[0].Name)
-	assert.Equal(t, "flyte_idt", c[1].Name)
-	assert.Equal(t, "flyte_rt", c[2].Name)
-}
+	t.Run("invalid_hash_key", func(t *testing.T) {
+		_, err := NewCookieManager(ctx, "wrong", blockKeyEncoded, cookieSetting)
 
-func TestCookieManager_RetrieveTokenValues(t *testing.T) {
-	ctx := context.Background()
-	// These were generated for unit testing only.
-	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
-	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
+		assert.EqualError(t, err, "[BINARY_DECODING_FAILED] Error decoding hash key bytes, caused by: illegal base64 data at input byte 4")
+	})
 
-	cookieSetting := config.CookieSettings{
-		SameSitePolicy: config.SameSiteDefaultMode,
-		Domain:         "default",
-	}
+	t.Run("invalid_block_key", func(t *testing.T) {
+		_, err := NewCookieManager(ctx, hashKeyEncoded, "wrong", cookieSetting)
 
-	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
-	assert.NoError(t, err)
+		assert.EqualError(t, err, "[BINARY_DECODING_FAILED] Error decoding block key bytes, caused by: illegal base64 data at input byte 4")
+	})
 
-	token := &oauth2.Token{
-		AccessToken:  "access",
-		RefreshToken: "refresh",
-	}
+	t.Run("set_token_cookies", func(t *testing.T) {
+		w := httptest.NewRecorder()
 
-	token = token.WithExtra(map[string]interface{}{
-		"id_token": "id token",
+		err = manager.SetTokenCookies(ctx, w, token)
+
+		assert.NoError(t, err)
+		fmt.Println(w.Header().Get("Set-Cookie"))
+		c := w.Result().Cookies()
+		assert.Equal(t, "flyte_at", c[0].Name)
+		assert.Equal(t, "flyte_idt", c[1].Name)
+		assert.Equal(t, "flyte_rt", c[2].Name)
 	})
 
-	w := httptest.NewRecorder()
-	_, err = http.NewRequest("GET", "/api/v1/projects", nil)
-	assert.NoError(t, err)
-	err = manager.SetTokenCookies(ctx, w, token)
-	assert.NoError(t, err)
+	t.Run("set_token_nil", func(t *testing.T) {
+		w := httptest.NewRecorder()
 
-	cookies := w.Result().Cookies()
-	req, err := http.NewRequest("GET", "/api/v1/projects", nil)
-	assert.NoError(t, err)
-	for _, c := range cookies {
-		req.AddCookie(c)
-	}
+		err = manager.SetTokenCookies(ctx, w, nil)
 
-	idToken, access, refresh, err := manager.RetrieveTokenValues(ctx, req)
-	assert.NoError(t, err)
-	assert.Equal(t, "id token", idToken)
-	assert.Equal(t, "access", access)
-	assert.Equal(t, "refresh", refresh)
-}
+		assert.EqualError(t, err, "[EMPTY_OAUTH_TOKEN] Attempting to set cookies with nil token")
+	})
 
-func TestGetLogoutAccessCookie(t *testing.T) {
-	cookie := getLogoutAccessCookie()
-	assert.True(t, time.Now().After(cookie.Expires))
-}
+	t.Run("set_token_cookies_wrong_key", func(t *testing.T) {
+		wrongKey := base64.RawStdEncoding.EncodeToString(bytes.Repeat([]byte("X"), 75))
+		wrongManager, err := NewCookieManager(ctx, wrongKey, wrongKey, cookieSetting)
+		require.NoError(t, err)
+		w := httptest.NewRecorder()
 
-func TestGetLogoutRefreshCookie(t *testing.T) {
-	cookie := getLogoutRefreshCookie()
-	assert.True(t, time.Now().After(cookie.Expires))
-}
+		err = wrongManager.SetTokenCookies(ctx, w, token)
 
-func TestCookieManager_DeleteCookies(t *testing.T) {
-	ctx := context.Background()
+		assert.EqualError(t, err, "[SECURE_COOKIE_ERROR] Error creating secure cookie, caused by: securecookie: error - caused by: crypto/aes: invalid key size 75")
+	})
 
-	// These were generated for unit testing only.
-	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
-	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
-	cookieSetting := config.CookieSettings{
-		SameSitePolicy: config.SameSiteDefaultMode,
-		Domain:         "default",
-	}
+	t.Run("retrieve_token_values", func(t *testing.T) {
+		w := httptest.NewRecorder()
 
-	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
-	assert.NoError(t, err)
+		err = manager.SetTokenCookies(ctx, w, token)
+		assert.NoError(t, err)
 
-	w := httptest.NewRecorder()
-	manager.DeleteCookies(ctx, w)
-	cookies := w.Result().Cookies()
-	assert.Equal(t, 2, len(cookies))
-	assert.True(t, time.Now().After(cookies[0].Expires))
-	assert.True(t, time.Now().After(cookies[1].Expires))
-}
+		cookies := w.Result().Cookies()
+		req, err := http.NewRequest("GET", "/api/v1/projects", nil)
+		assert.NoError(t, err)
+		for _, c := range cookies {
+			req.AddCookie(c)
+		}
 
-func TestGetHTTPSameSitePolicy(t *testing.T) {
-	ctx := context.Background()
+		idToken, access, refresh, err := manager.RetrieveTokenValues(ctx, req)
 
-	// These were generated for unit testing only.
-	hashKeyEncoded := "wG4pE1ccdw/pHZ2ml8wrD5VJkOtLPmBpWbKHmezWXktGaFbRoAhXidWs8OpbA3y7N8vyZhz1B1E37+tShWC7gA" //nolint:goconst
-	blockKeyEncoded := "afyABVgGOvWJFxVyOvCWCupoTn6BkNl4SOHmahho16Q"                                           //nolint:goconst
-	cookieSetting := config.CookieSettings{
-		SameSitePolicy: config.SameSiteDefaultMode,
-		Domain:         "default",
-	}
+		assert.NoError(t, err)
+		assert.Equal(t, "id token", idToken)
+		assert.Equal(t, "access", access)
+		assert.Equal(t, "refresh", refresh)
+	})
 
-	manager, err := NewCookieManager(ctx, hashKeyEncoded, blockKeyEncoded, cookieSetting)
-	assert.NoError(t, err)
-	assert.Equal(t, http.SameSiteDefaultMode, manager.getHTTPSameSitePolicy())
+	t.Run("retrieve_token_values_wrong_key", func(t *testing.T) {
+		wrongKey := base64.RawStdEncoding.EncodeToString(bytes.Repeat([]byte("X"), 75))
+		wrongManager, err := NewCookieManager(ctx, wrongKey, wrongKey, cookieSetting)
+		require.NoError(t, err)
+
+		w := httptest.NewRecorder()
+
+		err = manager.SetTokenCookies(ctx, w, token)
+		assert.NoError(t, err)
+
+		cookies := w.Result().Cookies()
+		req, err := http.NewRequest("GET", "/api/v1/projects", nil)
+		assert.NoError(t, err)
+		for _, c := range cookies {
+			req.AddCookie(c)
+		}
+
+		_, _, _, err = wrongManager.RetrieveTokenValues(ctx, req)
+
+		assert.EqualError(t, err, "[EMPTY_OAUTH_TOKEN] Error reading existing secure cookie [flyte_idt]. Error: [SECURE_COOKIE_ERROR] Error reading secure cookie flyte_idt, caused by: securecookie: error - caused by: crypto/aes: invalid key size 75")
+	})
+
+	t.Run("delete_cookies", func(t *testing.T) {
+		w := httptest.NewRecorder()
+
+		manager.DeleteCookies(ctx, w)
+
+		cookies := w.Result().Cookies()
+		require.Equal(t, 3, len(cookies))
+		assert.True(t, time.Now().After(cookies[0].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[0].Domain)
+		assert.Equal(t, accessTokenCookieName, cookies[0].Name)
+		assert.True(t, time.Now().After(cookies[1].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[1].Domain)
+		assert.Equal(t, refreshTokenCookieName, cookies[1].Name)
+		assert.True(t, time.Now().After(cookies[1].Expires))
+		assert.Equal(t, cookieSetting.Domain, cookies[1].Domain)
+		assert.Equal(t, idTokenCookieName, cookies[2].Name)
+	})
 
-	manager.sameSitePolicy = config.SameSiteLaxMode
-	assert.Equal(t, http.SameSiteLaxMode, manager.getHTTPSameSitePolicy())
+	t.Run("get_http_same_site_policy", func(t *testing.T) {
+		manager.sameSitePolicy = config.SameSiteLaxMode
+		assert.Equal(t, http.SameSiteLaxMode, manager.getHTTPSameSitePolicy())
 
-	manager.sameSitePolicy = config.SameSiteStrictMode
-	assert.Equal(t, http.SameSiteStrictMode, manager.getHTTPSameSitePolicy())
+		manager.sameSitePolicy = config.SameSiteStrictMode
+		assert.Equal(t, http.SameSiteStrictMode, manager.getHTTPSameSitePolicy())
 
-	manager.sameSitePolicy = config.SameSiteNoneMode
-	assert.Equal(t, http.SameSiteNoneMode, manager.getHTTPSameSitePolicy())
+		manager.sameSitePolicy = config.SameSiteNoneMode
+		assert.Equal(t, http.SameSiteNoneMode, manager.getHTTPSameSitePolicy())
+	})
+
+	t.Run("set_user_info", func(t *testing.T) {
+		w := httptest.NewRecorder()
+		info := &service.UserInfoResponse{
+			Subject: "sub",
+			Name:    "foo",
+		}
+
+		err := manager.SetUserInfoCookie(ctx, w, info)
+
+		assert.NoError(t, err)
+		cookies := w.Result().Cookies()
+		require.Len(t, cookies, 1)
+		assert.Equal(t, "flyte_user_info", cookies[0].Name)
+	})
+
+	t.Run("set_auth_code", func(t *testing.T) {
+		w := httptest.NewRecorder()
+
+		err := manager.SetAuthCodeCookie(ctx, w, "foo.com")
+
+		assert.NoError(t, err)
+		cookies := w.Result().Cookies()
+		require.Len(t, cookies, 1)
+		assert.Equal(t, "flyte_auth_code", cookies[0].Name)
+	})
 }