From 7babb885d49d61a24a3829d960006ec0c33339a0 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katroganGH@gmail.com>
Date: Thu, 9 Jun 2022 14:43:54 -0600
Subject: [PATCH] Fix userinfo metadata base64 encoding (#443)

* Use grpc base64 encoding

Signed-off-by: Katrina Rogan <katroganGH@gmail.com>

* remove check for empty user info

Signed-off-by: Katrina Rogan <katroganGH@gmail.com>
---
 flyteadmin/auth/constants.go           |  3 ++-
 flyteadmin/auth/encoding_utils.go      | 22 ------------------
 flyteadmin/auth/encoding_utils_test.go | 31 --------------------------
 flyteadmin/auth/handlers.go            |  2 +-
 flyteadmin/auth/handlers_test.go       |  1 -
 flyteadmin/auth/token.go               |  4 ++--
 6 files changed, 5 insertions(+), 58 deletions(-)
 delete mode 100644 flyteadmin/auth/encoding_utils.go
 delete mode 100644 flyteadmin/auth/encoding_utils_test.go

diff --git a/flyteadmin/auth/constants.go b/flyteadmin/auth/constants.go
index 1b7366f54d..7b64227561 100644
--- a/flyteadmin/auth/constants.go
+++ b/flyteadmin/auth/constants.go
@@ -9,7 +9,8 @@ const (
 	DefaultAuthorizationHeader    = "authorization"
 	BearerScheme                  = "Bearer"
 	IDTokenScheme                 = "IDToken"
-	UserInfoMDKey                 = "UserInfo"
+	// Add the -bin suffix so that the header value is automatically base64 encoded
+	UserInfoMDKey = "UserInfo-bin"
 
 	// https://tools.ietf.org/html/rfc8414
 	// This should be defined without a leading slash. If there is one, the url library's ResolveReference will make it a root path
diff --git a/flyteadmin/auth/encoding_utils.go b/flyteadmin/auth/encoding_utils.go
deleted file mode 100644
index e0f296e2f1..0000000000
--- a/flyteadmin/auth/encoding_utils.go
+++ /dev/null
@@ -1,22 +0,0 @@
-package auth
-
-import (
-	"context"
-	"encoding/base64"
-
-	"github.com/flyteorg/flytestdlib/logger"
-)
-
-// EncodeBase64 returns the base64 encoded version of the data
-func EncodeBase64(raw []byte) string {
-	return base64.RawStdEncoding.EncodeToString(raw)
-}
-
-// DecodeFromBase64 returns the original encoded bytes and logs warning in case of error
-func DecodeFromBase64(encodedData string) ([]byte, error) {
-	decodedData, err := base64.StdEncoding.DecodeString(encodedData)
-	if err != nil {
-		logger.Warnf(context.TODO(), "Unable to decode %v due to %v", encodedData, err)
-	}
-	return decodedData, err
-}
diff --git a/flyteadmin/auth/encoding_utils_test.go b/flyteadmin/auth/encoding_utils_test.go
deleted file mode 100644
index 2918d76ccb..0000000000
--- a/flyteadmin/auth/encoding_utils_test.go
+++ /dev/null
@@ -1,31 +0,0 @@
-package auth
-
-import (
-	"encoding/base64"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestEncodeAscii(t *testing.T) {
-	assert.Equal(t, "bmls", EncodeBase64([]byte("nil")))
-	assert.Equal(t, "w4RwZmVs", EncodeBase64([]byte("Äpfel")))
-}
-
-func TestDecodeFromAscii(t *testing.T) {
-	type data struct {
-		decoded     string
-		encoded     string
-		expectedErr error
-	}
-	tt := []data{
-		{decoded: "nil", encoded: "bmls", expectedErr: nil},
-		{decoded: "Äpfel", encoded: "w4RwZmVs", expectedErr: nil},
-		{decoded: "", encoded: "Äpfel", expectedErr: base64.CorruptInputError(0)},
-	}
-	for _, testdata := range tt {
-		actualDecoded, actualErr := DecodeFromBase64(testdata.encoded)
-		assert.Equal(t, []byte(testdata.decoded), actualDecoded)
-		assert.Equal(t, testdata.expectedErr, actualErr)
-	}
-}
diff --git a/flyteadmin/auth/handlers.go b/flyteadmin/auth/handlers.go
index 9ddca70d8e..4f1f35f34d 100644
--- a/flyteadmin/auth/handlers.go
+++ b/flyteadmin/auth/handlers.go
@@ -323,7 +323,7 @@ func GetHTTPRequestCookieToMetadataHandler(authCtx interfaces.AuthenticationCont
 		}
 
 		if len(raw) > 0 {
-			meta.Set(UserInfoMDKey, EncodeBase64(raw))
+			meta.Set(UserInfoMDKey, string(raw))
 		}
 
 		return meta
diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go
index 40d5609c34..5a210b0a59 100644
--- a/flyteadmin/auth/handlers_test.go
+++ b/flyteadmin/auth/handlers_test.go
@@ -234,7 +234,6 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) {
 	req.AddCookie(&idCookie)
 
 	assert.Equal(t, "IDToken a.b.c", handler(ctx, req)["authorization"][0])
-	assert.Equal(t, "bnVsbA", handler(ctx, req).Get(UserInfoMDKey)[0])
 }
 
 func TestGetHTTPMetadataTaggingHandler(t *testing.T) {
diff --git a/flyteadmin/auth/token.go b/flyteadmin/auth/token.go
index 7b97f637da..9433ea6ffe 100644
--- a/flyteadmin/auth/token.go
+++ b/flyteadmin/auth/token.go
@@ -111,10 +111,10 @@ func GRPCGetIdentityFromIDToken(ctx context.Context, clientID string, provider *
 	}
 
 	meta := metautils.ExtractIncoming(ctx)
-	userInfoDecoded, _ := DecodeFromBase64(meta.Get(UserInfoMDKey))
+	userInfoDecoded := meta.Get(UserInfoMDKey)
 	userInfo := &service.UserInfoResponse{}
 	if len(userInfoDecoded) > 0 {
-		err = json.Unmarshal(userInfoDecoded, userInfo)
+		err = json.Unmarshal([]byte(userInfoDecoded), userInfo)
 		if err != nil {
 			logger.Infof(ctx, "Could not unmarshal user info from metadata %v", err)
 		}