Merge pull request #900 from Pangjiping/fix/sdk-go-forward-endpoint-headers

fix(sdk-go): forward all GetEndpoint headers on subsequent requests

Author: ninan-nn

sdks/sandbox/go/config.go

@@ -177,22 +177,24 @@ func (c *ConnectionConfig) lifecycleClient() *LifecycleClient {
 	return NewLifecycleClient(c.GetBaseURL()+"/"+APIVersion, c.GetAPIKey(), c.clientOpts(true)...)
 }
 
-// execdClient creates an ExecdClient for a resolved endpoint.
-// endpointHeaders are additional headers from the endpoint resolution (e.g. routing headers).
-func (c *ConnectionConfig) execdClient(endpointURL, token string, endpointHeaders map[string]string) *ExecdClient {
+// execdClient creates an ExecdClient for a resolved endpoint. All headers
+// returned by the lifecycle GetEndpoint call (auth tokens, routing hints,
+// sticky-session keys, etc.) are forwarded as-is on every subsequent request.
+func (c *ConnectionConfig) execdClient(endpointURL string, endpointHeaders map[string]string) *ExecdClient {
 	opts := c.clientOpts(true)
 	if len(endpointHeaders) > 0 {
 		opts = append(opts, WithHeaders(endpointHeaders))
 	}
-	return NewExecdClient(endpointURL, token, opts...)
+	return NewExecdClient(endpointURL, "", opts...)
 }
 
-// egressClient creates an EgressClient for a resolved endpoint.
-// endpointHeaders are additional headers from the endpoint resolution (e.g. routing headers).
-func (c *ConnectionConfig) egressClient(endpointURL, token string, endpointHeaders map[string]string) *EgressClient {
+// egressClient creates an EgressClient for a resolved endpoint. All headers
+// returned by the lifecycle GetEndpoint call are forwarded as-is on every
+// subsequent request.
+func (c *ConnectionConfig) egressClient(endpointURL string, endpointHeaders map[string]string) *EgressClient {
 	opts := c.clientOpts(false)
 	if len(endpointHeaders) > 0 {
 		opts = append(opts, WithHeaders(endpointHeaders))
 	}
-	return NewEgressClient(endpointURL, token, opts...)
+	return NewEgressClient(endpointURL, "", opts...)
 }

sdks/sandbox/go/egress.go

@@ -22,14 +22,17 @@ type EgressClient struct {
 	*Client
 }
 
+// egressAuthHeader is the authentication header used by the Egress sidecar API.
+const egressAuthHeader = "OPENSANDBOX-EGRESS-AUTH"
+
 // NewEgressClient creates a new EgressClient.
 // baseURL is the sandbox-specific egress sidecar endpoint
 // (e.g. "http://localhost:18080").
 // authToken is the value for the OPENSANDBOX-EGRESS-AUTH header; pass ""
 // if the sidecar does not require authentication.
 func NewEgressClient(baseURL, authToken string, opts ...Option) *EgressClient {
 	return &EgressClient{
-		Client: NewClient(baseURL, authToken, "OPENSANDBOX-EGRESS-AUTH", opts...),
+		Client: NewClient(baseURL, authToken, egressAuthHeader, opts...),
 	}
 }
 

sdks/sandbox/go/opensandbox_test.go

@@ -1022,6 +1022,92 @@ func TestExecdAuthHeader(t *testing.T) {
 	require.NoErrorf(t, err, "Ping")
 }
 
+// TestResolveExecdForwardsAllEndpointHeaders verifies that every header
+// returned by GetEndpoint (auth tokens, routing hints, sticky-session keys,
+// etc.) is forwarded as-is on subsequent execd requests, mirroring the
+// Python SDK behavior.
+func TestResolveExecdForwardsAllEndpointHeaders(t *testing.T) {
+	endpointHeaders := map[string]string{
+		"X-EXECD-ACCESS-TOKEN": "execd-tok",
+		"X-Route-Hint":         "vip-pool",
+		"X-Sticky-Session":     "sess-abc",
+	}
+
+	execdSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for k, want := range endpointHeaders {
+			if got := r.Header.Get(k); got != want {
+				assert.Fail(t, fmt.Sprintf("header %s = %q, want %q", k, got, want))
+			}
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer execdSrv.Close()
+
+	lifecycleSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/endpoints/") {
+			jsonResponse(w, http.StatusOK, Endpoint{
+				Endpoint: execdSrv.URL,
+				Headers:  endpointHeaders,
+			})
+			return
+		}
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer lifecycleSrv.Close()
+
+	config := ConnectionConfig{Domain: lifecycleSrv.URL}
+	sb := &Sandbox{
+		id:        "sbx-headers",
+		config:    &config,
+		lifecycle: config.lifecycleClient(),
+	}
+
+	require.NoErrorf(t, sb.resolveExecd(context.Background()), "resolveExecd")
+	require.NoErrorf(t, sb.execd.Ping(context.Background()), "Ping")
+}
+
+// TestResolveEgressForwardsAllEndpointHeaders verifies the same forwarding
+// behavior for the egress sidecar client.
+func TestResolveEgressForwardsAllEndpointHeaders(t *testing.T) {
+	endpointHeaders := map[string]string{
+		"OPENSANDBOX-EGRESS-AUTH": "egress-tok",
+		"X-Route-Hint":            "egress-vip",
+	}
+
+	egressSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for k, want := range endpointHeaders {
+			if got := r.Header.Get(k); got != want {
+				assert.Fail(t, fmt.Sprintf("header %s = %q, want %q", k, got, want))
+			}
+		}
+		jsonResponse(w, http.StatusOK, PolicyStatusResponse{Status: "ok"})
+	}))
+	defer egressSrv.Close()
+
+	lifecycleSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == http.MethodGet && strings.Contains(r.URL.Path, "/endpoints/") {
+			jsonResponse(w, http.StatusOK, Endpoint{
+				Endpoint: egressSrv.URL,
+				Headers:  endpointHeaders,
+			})
+			return
+		}
+		w.WriteHeader(http.StatusNotFound)
+	}))
+	defer lifecycleSrv.Close()
+
+	config := ConnectionConfig{Domain: lifecycleSrv.URL}
+	sb := &Sandbox{
+		id:        "sbx-egress-headers",
+		config:    &config,
+		lifecycle: config.lifecycleClient(),
+	}
+
+	require.NoErrorf(t, sb.resolveEgress(context.Background()), "resolveEgress")
+	_, err := sb.egress.GetPolicy(context.Background())
+	require.NoErrorf(t, err, "GetPolicy")
+}
+
 func TestSandboxManager_ListFilter(t *testing.T) {
 	now := time.Now().UTC().Truncate(time.Second)
 	want := ListSandboxesResponse{

sdks/sandbox/go/sandbox.go

@@ -411,23 +411,19 @@ func (s *Sandbox) resolveExecd(ctx context.Context) error {
 		execdURL = s.config.GetProtocol() + "://" + execdURL
 	}
 
-	token := ""
-	var extraHeaders map[string]string
-	if endpoint.Headers != nil {
-		token = endpoint.Headers["X-EXECD-ACCESS-TOKEN"]
-		// Preserve all endpoint headers (e.g. routing headers) except the auth token
-		extraHeaders = make(map[string]string, len(endpoint.Headers))
-		for k, v := range endpoint.Headers {
-			if k != "X-EXECD-ACCESS-TOKEN" {
-				extraHeaders[k] = v
+	headers := make(map[string]string, len(endpoint.Headers)+1)
+	for k, v := range endpoint.Headers {
+		headers[k] = v
+	}
+	if s.config.UseServerProxy {
+		if _, ok := headers[execdAuthHeader]; !ok {
+			if apiKey := s.config.GetAPIKey(); apiKey != "" {
+				headers[execdAuthHeader] = apiKey
 			}
 		}
 	}
-	if s.config.UseServerProxy && token == "" {
-		token = s.config.GetAPIKey()
-	}
 
-	s.execd = s.config.execdClient(execdURL, token, extraHeaders)
+	s.execd = s.config.execdClient(execdURL, headers)
 	return nil
 }
 
@@ -451,18 +447,18 @@ func (s *Sandbox) resolveEgress(ctx context.Context) error {
 		egressURL = s.config.GetProtocol() + "://" + egressURL
 	}
 
-	token := ""
-	var extraHeaders map[string]string
-	if endpoint.Headers != nil {
-		token = endpoint.Headers["OPENSANDBOX-EGRESS-AUTH"]
-		extraHeaders = make(map[string]string, len(endpoint.Headers))
-		for k, v := range endpoint.Headers {
-			if k != "OPENSANDBOX-EGRESS-AUTH" {
-				extraHeaders[k] = v
+	headers := make(map[string]string, len(endpoint.Headers)+1)
+	for k, v := range endpoint.Headers {
+		headers[k] = v
+	}
+	if s.config.UseServerProxy {
+		if _, ok := headers[egressAuthHeader]; !ok {
+			if apiKey := s.config.GetAPIKey(); apiKey != "" {
+				headers[egressAuthHeader] = apiKey
 			}
 		}
 	}
 
-	s.egress = s.config.egressClient(egressURL, token, extraHeaders)
+	s.egress = s.config.egressClient(egressURL, headers)
 	return nil
 }