From f18c79ca473b50d5abc700b5071cd6cd89a5cc98 Mon Sep 17 00:00:00 2001
From: Nian Tang <niantang@tesla.com>
Date: Mon, 23 Sep 2024 10:43:49 +0800
Subject: [PATCH] Add custom header to specify fleetapi base url

---
 examples/unlock/unlock.go   |  2 +-
 pkg/account/account.go      | 13 +++++++++----
 pkg/account/account_test.go | 27 ++++++++++++++++++++++++---
 pkg/cli/config.go           |  2 +-
 pkg/proxy/proxy.go          |  4 +++-
 5 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/examples/unlock/unlock.go b/examples/unlock/unlock.go
index 13656da7..e47bc441 100644
--- a/examples/unlock/unlock.go
+++ b/examples/unlock/unlock.go
@@ -61,7 +61,7 @@ func main() {
 
 	// This example program sends commands over the Internet, which requires a Tesla account login
 	// token. The protocol can also work over BLE; see other programs in the example directory.
-	acct, err := account.New(string(oauthToken), userAgent)
+	acct, err := account.New(string(oauthToken), userAgent, "")
 	if err != nil {
 		logger.Printf("Authentication error: %s", err)
 		return
diff --git a/pkg/account/account.go b/pkg/account/account.go
index e038a49f..e5c518d7 100644
--- a/pkg/account/account.go
+++ b/pkg/account/account.go
@@ -114,7 +114,7 @@ func (p *oauthPayload) domain() string {
 
 // New returns an [Account] that can be used to fetch a [vehicle.Vehicle].
 // Optional userAgent can be passed in - otherwise it will be generated from code
-func New(oauthToken, userAgent string) (*Account, error) {
+func New(oauthToken, userAgent, fleetapiHost string) (*Account, error) {
 	parts := strings.Split(oauthToken, ".")
 	if len(parts) != 3 {
 		return nil, fmt.Errorf("client provided malformed OAuth token")
@@ -128,9 +128,14 @@ func New(oauthToken, userAgent string) (*Account, error) {
 		return nil, fmt.Errorf("client provided malformed OAuth token: %s", err)
 	}
 
-	domain := payload.domain()
-	if domain == "" {
-		return nil, fmt.Errorf("client provided OAuth token with invalid audiences")
+	domain := ""
+	if fleetapiHost != "" {
+		domain, _ = strings.CutPrefix(fleetapiHost, "https://")
+	} else {
+		domain = payload.domain()
+		if domain == "" {
+			return nil, fmt.Errorf("client provided OAuth token with invalid audiences")
+		}
 	}
 	return &Account{
 		UserAgent:  buildUserAgent(userAgent),
diff --git a/pkg/account/account_test.go b/pkg/account/account_test.go
index 9a99d4ad..f441651e 100644
--- a/pkg/account/account_test.go
+++ b/pkg/account/account_test.go
@@ -33,7 +33,7 @@ func TestNewAccount(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.description, func(t *testing.T) {
-			acct, err := New(test.jwt, "")
+			acct, err := New(test.jwt, "", "")
 			if (err != nil) != test.shouldError {
 				t.Errorf("Unexpected result: err = %v, shouldError = %v", err, test.shouldError)
 			}
@@ -50,7 +50,7 @@ func TestDomainDefault(t *testing.T) {
 		Audiences: []string{"https://auth.tesla.com/nts"},
 	}
 
-	acct, err := New(makeTestJWT(payload), "")
+	acct, err := New(makeTestJWT(payload), "", "")
 	if err != nil {
 		t.Fatalf("Returned error on valid JWT: %s", err)
 	}
@@ -70,7 +70,7 @@ func TestDomainExtraction(t *testing.T) {
 		OUCode: "EU",
 	}
 
-	acct, err := New(makeTestJWT(payload), "")
+	acct, err := New(makeTestJWT(payload), "", "")
 	if err != nil {
 		t.Fatalf("Returned error on valid JWT: %s", err)
 	}
@@ -80,6 +80,27 @@ func TestDomainExtraction(t *testing.T) {
 	}
 }
 
+// TestDomainExtraction tests the extraction of the correct domain based on OUCode.
+func TestTargetHost(t *testing.T) {
+	payload := &oauthPayload{
+		Audiences: []string{
+			"https://auth.tesla.com/nts",
+			"https://fleet-api.prd.na.vn.cloud.tesla.com",
+			"https://fleet-api.prd.eu.vn.cloud.tesla.com",
+		},
+		OUCode: "EU",
+	}
+
+	acct, err := New(makeTestJWT(payload), "", "https://fleet-api.prd.na.vn.cloud.tesla.com")
+	if err != nil {
+		t.Fatalf("Returned error on valid JWT: %s", err)
+	}
+	expectedHost := "fleet-api.prd.na.vn.cloud.tesla.com"
+	if acct == nil || acct.Host != expectedHost {
+		t.Errorf("acct = %+v, expected Host = %s", acct, expectedHost)
+	}
+}
+
 // makeTestJWT creates a JWT string with the given payload.
 func makeTestJWT(payload *oauthPayload) string {
 	jwtBody, _ := json.Marshal(payload)
diff --git a/pkg/cli/config.go b/pkg/cli/config.go
index 0a64ea0a..b4741338 100644
--- a/pkg/cli/config.go
+++ b/pkg/cli/config.go
@@ -421,7 +421,7 @@ func (c *Config) Account() (*account.Account, error) {
 	if err != nil {
 		return nil, err
 	}
-	return account.New(token, "")
+	return account.New(token, "", "")
 }
 
 // SavePrivateKey writes skey to the system keyring or file, depending on what options are
diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go
index bae76c7b..bf73e7a6 100644
--- a/pkg/proxy/proxy.go
+++ b/pkg/proxy/proxy.go
@@ -37,7 +37,9 @@ func getAccount(req *http.Request) (*account.Account, error) {
 	if !ok {
 		return nil, fmt.Errorf("client did not provide an OAuth token")
 	}
-	return account.New(token, proxyProtocolVersion)
+	fleetapiHost := req.Header.Get("Fleetapi-Host")
+
+	return account.New(token, proxyProtocolVersion, fleetapiHost)
 }
 
 // Proxy exposes an HTTP API for sending vehicle commands.