From bbde4b216154042f1311df11ac8a7950252a0919 Mon Sep 17 00:00:00 2001
From: Anuj Chaudhari <anuj.chaudhari@broadcom.com>
Date: Fri, 9 Aug 2024 16:37:43 -0700
Subject: [PATCH] Update TMC endpoint mapping rules for Tanzu context (#798)

---
 pkg/command/context.go         | 41 +++++++++++++++++++++++++++-------
 pkg/command/context_test.go    | 12 ++++++++++
 pkg/constants/env_variables.go |  8 +++++++
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/pkg/command/context.go b/pkg/command/context.go
index 4d0ae7c81..28073f025 100644
--- a/pkg/command/context.go
+++ b/pkg/command/context.go
@@ -12,7 +12,6 @@ import (
 	"net/url"
 	"os"
 	"reflect"
-	"regexp"
 	"sort"
 	"strconv"
 	"strings"
@@ -668,9 +667,14 @@ func globalTanzuLogin(c *configtypes.Context, generateContextNameFunc func(orgNa
 
 	// Fetch the Tanzu Hub endpoint for the Tanzu context as a best case effort
 	if tanzuHubEndpoint == "" {
-		tanzuHubEndpoint, err = csp.GetTanzuHubEndpoint(claims.OrgID, c.GlobalOpts.Auth.AccessToken, staging)
-		if err != nil {
-			log.V(7).Infof("unable to get Tanzu Hub endpoint. Error: %v", err.Error())
+		// If the TANZU_CLI_HUB_ENDPOINT is set just use the configured endpoint
+		if os.Getenv(constants.TPHubEndpoint) != "" {
+			tanzuHubEndpoint = os.Getenv(constants.TPHubEndpoint)
+		} else {
+			tanzuHubEndpoint, err = csp.GetTanzuHubEndpoint(claims.OrgID, c.GlobalOpts.Auth.AccessToken, staging)
+			if err != nil {
+				log.V(7).Infof("unable to get Tanzu Hub endpoint. Error: %v", err.Error())
+			}
 		}
 	} else {
 		log.Warningf("This tanzu context is being created with the custom Tanzu Hub endpoint: %q", tanzuHubEndpoint)
@@ -2143,14 +2147,35 @@ func renderDynamicTable(slices interface{}, tableWriter component.OutputWriter,
 }
 
 func mapTanzuEndpointToTMCEndpoint(tanzuEndpoint string) string {
-	// Define the regular expression pattern
-	apiPattern := regexp.MustCompile(`https://api\.tanzu(-\w*)?\.cloud\.vmware\.com`)
-	// Replace "api" with "tmc" in the input URL
-	tmcEndpoint := apiPattern.ReplaceAllString(tanzuEndpoint, "https://tmc.tanzu$1.cloud.vmware.com")
+	// If the TANZU_CLI_K8S_OPS_ENDPOINT is set just return the configured endpoint
+	if os.Getenv(constants.TPKubernetesOpsEndpoint) != "" {
+		return os.Getenv(constants.TPKubernetesOpsEndpoint)
+	}
+
+	tmcEndpoint := ""
+	// Define the mapping rules
+	mappingRules := []struct {
+		tanzuEndpointPrefix string
+		tmcEndpointPrefix   string
+	}{
+		{"https://api.tanzu", "https://tmc.tanzu"},
+		{"https://ucp.platform", "https://ops.platform"},
+	}
+
+	// Iterate through the mapping rules
+	for _, rule := range mappingRules {
+		if strings.HasPrefix(tanzuEndpoint, rule.tanzuEndpointPrefix) {
+			// Replace the tanzuEndpointPrefix with the tmcEndpointPrefix
+			tmcEndpoint = strings.Replace(tanzuEndpoint, rule.tanzuEndpointPrefix, rule.tmcEndpointPrefix, 1)
+			break
+		}
+	}
+
 	// Check if the transformation was successful
 	if tanzuEndpoint == tmcEndpoint {
 		return ""
 	}
+
 	return tmcEndpoint
 }
 
diff --git a/pkg/command/context_test.go b/pkg/command/context_test.go
index 29bc506db..2d9c254aa 100644
--- a/pkg/command/context_test.go
+++ b/pkg/command/context_test.go
@@ -1882,6 +1882,18 @@ func TestMapTanzuEndpointToTMCEndpoint(t *testing.T) {
 			input:    "https://api.tanzu.cloud.vmware.com",
 			expected: "https://tmc.tanzu.cloud.vmware.com",
 		},
+		{
+			input:    "https://ucp.platform.tanzu.broadcom.com",
+			expected: "https://ops.platform.tanzu.broadcom.com",
+		},
+		{
+			input:    "https://ucp.platform-dev.tanzu.broadcom.com",
+			expected: "https://ops.platform-dev.tanzu.broadcom.com",
+		},
+		{
+			input:    "https://ucp.platform-verify.tanzu.broadcom.com",
+			expected: "https://ops.platform-verify.tanzu.broadcom.com",
+		},
 		{
 			input:    "https://symphony.api.tanzu.cloud.vmware.com",
 			expected: "",
diff --git a/pkg/constants/env_variables.go b/pkg/constants/env_variables.go
index 4ae6a93d5..d35c2bf21 100644
--- a/pkg/constants/env_variables.go
+++ b/pkg/constants/env_variables.go
@@ -94,4 +94,12 @@ const (
 
 	// UseTanzuCSP uses the Tanzu CSP while login/context creation
 	UseTanzuCSP = "TANZU_CLI_USE_TANZU_CLOUD_SERVICE_PROVIDER"
+
+	// TPKubernetesOpsEndpoint specifies kubernetes ops endpoint for the Tanzu Platform
+	// This will be used as part of `tanzu login`
+	TPKubernetesOpsEndpoint = "TANZU_CLI_K8S_OPS_ENDPOINT"
+
+	// TPHubEndpoint specifies hub endpoint for the Tanzu Platform
+	// This will be used as part of `tanzu login`
+	TPHubEndpoint = "TANZU_CLI_HUB_ENDPOINT"
 )