Add subject to UserInfoResponse header (flyteorg#488)

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

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

flyteadmin/auth/handlers.go

@@ -8,22 +8,19 @@ import (
 	"strings"
 	"time"
 
-	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
-
-	"golang.org/x/oauth2"
-
-	"github.com/flyteorg/flyteadmin/pkg/common"
-	"google.golang.org/grpc/peer"
-
-	"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
-
 	"github.com/flyteorg/flyteadmin/auth/interfaces"
+	"github.com/flyteorg/flyteadmin/pkg/common"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
 	"github.com/flyteorg/flytestdlib/errors"
 	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/grpc-ecosystem/go-grpc-middleware/util/metautils"
+	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/runtime/protoiface"
 )
 
 const (
@@ -33,6 +30,7 @@ const (
 )
 
 type HTTPRequestToMetadataAnnotator func(ctx context.Context, request *http.Request) metadata.MD
+type UserInfoForwardResponseHandler func(ctx context.Context, w http.ResponseWriter, m protoiface.MessageV1) error
 
 type AuthenticatedClientMeta struct {
 	ClientIds     []string
@@ -443,3 +441,13 @@ func GetLogoutEndpointHandler(ctx context.Context, authCtx interfaces.Authentica
 		}
 	}
 }
+
+func GetUserInfoForwardResponseHandler() UserInfoForwardResponseHandler {
+	return func(ctx context.Context, w http.ResponseWriter, m protoiface.MessageV1) error {
+		info, ok := m.(*service.UserInfoResponse)
+		if ok {
+			w.Header().Set("X-User-Subject", info.Subject)
+		}
+		return nil
+	}
+}

flyteadmin/auth/handlers_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/flyteorg/flyteadmin/auth/config"
 	"github.com/flyteorg/flyteadmin/auth/interfaces/mocks"
 	"github.com/flyteorg/flyteadmin/pkg/common"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
 	stdConfig "github.com/flyteorg/flytestdlib/config"
 
 	"github.com/coreos/go-oidc"
@@ -290,3 +291,20 @@ func TestGetOIdCMetadataEndpointRedirectHandler(t *testing.T) {
 	assert.Equal(t, http.StatusSeeOther, w.Code)
 	assert.Equal(t, "http://www.google.com/.well-known/openid-configuration", w.Header()["Location"][0])
 }
+
+func TestUserInfoForwardResponseHander(t *testing.T) {
+	ctx := context.Background()
+	handler := GetUserInfoForwardResponseHandler()
+	w := httptest.NewRecorder()
+	resp := service.UserInfoResponse{
+		Subject: "user-id",
+	}
+	assert.NoError(t, handler(ctx, w, &resp))
+	assert.Contains(t, w.Result().Header, "X-User-Subject")
+	assert.Equal(t, w.Result().Header["X-User-Subject"], []string{"user-id"})
+
+	w = httptest.NewRecorder()
+	unrelatedResp := service.OAuth2MetadataResponse{}
+	assert.NoError(t, handler(ctx, w, &unrelatedResp))
+	assert.NotContains(t, w.Result().Header, "X-User-Subject")
+}

flyteadmin/pkg/server/service.go

@@ -178,6 +178,9 @@ func newHTTPServer(ctx context.Context, cfg *config.ServerConfig, _ *authConfig.
 	// This option means that http requests are served with protobufs, instead of json. We always want this.
 	gwmuxOptions = append(gwmuxOptions, runtime.WithMarshalerOption("application/octet-stream", &runtime.ProtoMarshaller{}))
 
+	// This option sets subject in the user info response
+	gwmuxOptions = append(gwmuxOptions, runtime.WithForwardResponseOption(auth.GetUserInfoForwardResponseHandler()))
+
 	if cfg.Security.UseAuth {
 		// Add HTTP handlers for OIDC endpoints
 		auth.RegisterHandlers(ctx, mux, authCtx)