From e0bb1f1dfad0693aa4d237dfb627b361c2f020a2 Mon Sep 17 00:00:00 2001
From: Anders Schwartz <anders.schwartz@kurtosistech.com>
Date: Fri, 1 Dec 2023 16:53:21 -0500
Subject: [PATCH] fix: look up apikey by JWT token to avoid any random jwt
 token to access the API key

---
 enclave-manager/server/server.go | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/enclave-manager/server/server.go b/enclave-manager/server/server.go
index 064ab9a957..4f8ac1e58b 100644
--- a/enclave-manager/server/server.go
+++ b/enclave-manager/server/server.go
@@ -44,7 +44,7 @@ type WebServer struct {
 	engineServiceClient *kurtosis_engine_rpc_api_bindingsconnect.EngineServiceClient
 	enforceAuth         bool
 	instanceConfig      *kurtosis_backend_server_rpc_api_bindings.GetCloudInstanceConfigResponse
-	apiKey              *string
+	apiKeyMap           map[string]*string
 }
 
 func NewWebserver(enforceAuth bool) (*WebServer, error) {
@@ -57,6 +57,7 @@ func NewWebserver(enforceAuth bool) (*WebServer, error) {
 		enforceAuth:         enforceAuth,
 		instanceConfigMutex: &sync.RWMutex{},
 		apiKeyMutex:         &sync.RWMutex{},
+		apiKeyMap:           map[string]*string{},
 	}, nil
 }
 
@@ -468,21 +469,21 @@ func (c *WebServer) ConvertJwtTokenToApiKey(
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "Failed to create the Cloud backend client")
 	}
-	request := &connect.Request[kurtosis_backend_server_rpc_api_bindings.GetOrCreateApiKeyRequest]{
-		Msg: &kurtosis_backend_server_rpc_api_bindings.GetOrCreateApiKeyRequest{
-			AccessToken: jwtToken,
-		},
-	}
 
-	if c.apiKey != nil {
+	if c.apiKeyMap[jwtToken] != nil {
 		return &Authentication{
-			ApiKey:   *c.apiKey,
+			ApiKey:   *c.apiKeyMap[jwtToken],
 			JwtToken: jwtToken,
 		}, nil
 	} else {
 		c.apiKeyMutex.Lock()
 		defer c.apiKeyMutex.Unlock()
 
+		request := &connect.Request[kurtosis_backend_server_rpc_api_bindings.GetOrCreateApiKeyRequest]{
+			Msg: &kurtosis_backend_server_rpc_api_bindings.GetOrCreateApiKeyRequest{
+				AccessToken: jwtToken,
+			},
+		}
 		result, err := (*client).GetOrCreateApiKey(ctx, request)
 		if err != nil {
 			return nil, stacktrace.Propagate(err, "Failed to get the API key")
@@ -494,7 +495,7 @@ func (c *WebServer) ConvertJwtTokenToApiKey(
 		}
 
 		if len(result.Msg.ApiKey) > 0 {
-			c.apiKey = &result.Msg.ApiKey
+			c.apiKeyMap[jwtToken] = &result.Msg.ApiKey
 			return &Authentication{
 				ApiKey:   result.Msg.ApiKey,
 				JwtToken: jwtToken,