From 67a4ffaf54755698fe43d31033c5ef99fcff16ce Mon Sep 17 00:00:00 2001
From: Eder Ignatowicz <ignatowicz@gmail.com>
Date: Mon, 11 Nov 2024 11:28:08 -0500
Subject: [PATCH] feat(bff): load x-forwarded-access-token in cluster setup

Signed-off-by: Eder Ignatowicz <ignatowicz@gmail.com>
---
 clients/ui/bff/internal/api/middleware.go     | 19 +++++++++++--------
 clients/ui/bff/internal/integrations/k8s.go   |  7 +++++++
 .../ui/frontend/src/shared/api/apiUtils.ts    |  5 ++++-
 3 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/clients/ui/bff/internal/api/middleware.go b/clients/ui/bff/internal/api/middleware.go
index 04cfc7806..743cd2567 100644
--- a/clients/ui/bff/internal/api/middleware.go
+++ b/clients/ui/bff/internal/api/middleware.go
@@ -5,13 +5,13 @@ import (
 	"fmt"
 	"github.com/julienschmidt/httprouter"
 	"github.com/kubeflow/model-registry/ui/bff/internal/integrations"
-	"k8s.io/client-go/rest"
 	"net/http"
 )
 
 type contextKey string
 
 const httpClientKey contextKey = "httpClientKey"
+const userAccessToken = "x-forwarded-access-token"
 
 func (app *App) RecoverPanic(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -47,7 +47,7 @@ func (app *App) AttachRESTClient(handler func(http.ResponseWriter, *http.Request
 			return
 		}
 		var bearerToken string
-		bearerToken, err = resolveBearerToken(app.kubernetesClient)
+		bearerToken, err = resolveBearerToken(app.kubernetesClient, r.Header)
 		if err != nil {
 			app.serverErrorResponse(w, r, fmt.Errorf("failed to resolve BearerToken): %v", err))
 			return
@@ -63,21 +63,24 @@ func (app *App) AttachRESTClient(handler func(http.ResponseWriter, *http.Request
 	}
 }
 
-func resolveBearerToken(k8s integrations.KubernetesClientInterface) (string, error) {
+func resolveBearerToken(k8s integrations.KubernetesClientInterface, header http.Header) (string, error) {
 	var bearerToken string
-	_, err := rest.InClusterConfig()
-	if err == nil {
+	//check if I'm inside cluster
+	if k8s.IsInCluster() {
 		//in cluster
-		//TODO (eder) load bearerToken probably from x-forwarded-access-bearerToken
-		return "", fmt.Errorf("failed to create Rest client (not implemented yet - inside cluster): %v", err)
+		bearerToken = header.Get(userAccessToken)
+		if bearerToken == "" {
+			return "", fmt.Errorf("failed to create Rest client (not able to get bearerToken on cluster)")
+		}
 	} else {
 		//off cluster (development)
+		var err error
 		bearerToken, err = k8s.BearerToken()
 		if err != nil {
 			return "", fmt.Errorf("failed to fetch BearerToken in development mode: %v", err)
 		}
 	}
-	return bearerToken, err
+	return bearerToken, nil
 }
 
 func resolveModelRegistryURL(id string, client integrations.KubernetesClientInterface) (string, error) {
diff --git a/clients/ui/bff/internal/integrations/k8s.go b/clients/ui/bff/internal/integrations/k8s.go
index 60729a100..2dd2c4c2d 100644
--- a/clients/ui/bff/internal/integrations/k8s.go
+++ b/clients/ui/bff/internal/integrations/k8s.go
@@ -3,6 +3,7 @@ package integrations
 import (
 	"context"
 	"fmt"
+	"k8s.io/client-go/rest"
 	"log/slog"
 	"os"
 	"time"
@@ -22,6 +23,7 @@ type KubernetesClientInterface interface {
 	GetServiceDetails() ([]ServiceDetails, error)
 	BearerToken() (string, error)
 	Shutdown(ctx context.Context, logger *slog.Logger) error
+	IsInCluster() bool
 }
 
 type ServiceDetails struct {
@@ -126,6 +128,11 @@ func (kc *KubernetesClient) Shutdown(ctx context.Context, logger *slog.Logger) e
 	}
 }
 
+func (kc *KubernetesClient) IsInCluster() bool {
+	_, err := rest.InClusterConfig()
+	return err == nil
+}
+
 func (kc *KubernetesClient) BearerToken() (string, error) {
 	return kc.Token, nil
 }
diff --git a/clients/ui/frontend/src/shared/api/apiUtils.ts b/clients/ui/frontend/src/shared/api/apiUtils.ts
index fc0aa71cb..b59807186 100644
--- a/clients/ui/frontend/src/shared/api/apiUtils.ts
+++ b/clients/ui/frontend/src/shared/api/apiUtils.ts
@@ -76,7 +76,10 @@ const callRestJSON = <T>(
 
   return fetch(`${host}${path}${searchParams ? `?${searchParams}` : ''}`, {
     ...otherOptions,
-    ...(contentType && { headers: { 'Content-Type': contentType } }),
+    headers: {
+      ...otherOptions.headers,
+      ...(contentType && { 'Content-Type': contentType }),
+    },
     method,
     body: formData ?? requestData,
   }).then((response) =>