From a6c2f123621c0ac40b16f21e5e0e994974ab7167 Mon Sep 17 00:00:00 2001
From: samzong <samzong.lu@gmail.com>
Date: Fri, 26 Jun 2026 12:41:50 -0400
Subject: [PATCH] fix(runtime): support absolute polling locations

Signed-off-by: samzong <samzong.lu@gmail.com>
---
 pkg/runtime/poll.go      | 33 +++++++++++++++++++++++++++
 pkg/runtime/poll_test.go | 48 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+)

diff --git a/pkg/runtime/poll.go b/pkg/runtime/poll.go
index e963f33..f93f5db 100644
--- a/pkg/runtime/poll.go
+++ b/pkg/runtime/poll.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
+	"strings"
 	"time"
 )
 
@@ -17,6 +19,27 @@ func PollUntilDone(ctx context.Context, hostname, location string, opts ClientOp
 	if timeout <= 0 {
 		timeout = DefaultPollTimeout
 	}
+	base, err := BaseURL(hostname)
+	if err != nil {
+		return nil, err
+	}
+	baseURL, err := url.Parse(base)
+	if err != nil {
+		return nil, err
+	}
+	port := func(u *url.URL) string {
+		if p := u.Port(); p != "" {
+			return p
+		}
+		switch strings.ToLower(u.Scheme) {
+		case "http":
+			return "80"
+		case "https":
+			return "443"
+		default:
+			return ""
+		}
+	}
 	deadline := time.Now().Add(timeout)
 	backoff := pollInitBackoff
 
@@ -25,6 +48,16 @@ func PollUntilDone(ctx context.Context, hostname, location string, opts ClientOp
 			return nil, fmt.Errorf("polling timed out after %s", timeout)
 		}
 
+		loc, err := url.Parse(location)
+		if err != nil {
+			return nil, fmt.Errorf("parse polling location: %w", err)
+		}
+		if loc.IsAbs() || loc.Host != "" {
+			if !strings.EqualFold(loc.Scheme, baseURL.Scheme) || !strings.EqualFold(loc.Hostname(), baseURL.Hostname()) || port(loc) != port(baseURL) {
+				return nil, fmt.Errorf("cross-host polling location %q", location)
+			}
+			location = loc.RequestURI()
+		}
 		r, err := DoRawFull(ctx, hostname, "GET", location, nil, opts)
 		if err != nil {
 			return nil, err
diff --git a/pkg/runtime/poll_test.go b/pkg/runtime/poll_test.go
index be09dfe..2875dea 100644
--- a/pkg/runtime/poll_test.go
+++ b/pkg/runtime/poll_test.go
@@ -61,6 +61,54 @@ func TestPollUntilDone_EventualSuccess(t *testing.T) {
 	}
 }
 
+func TestPollUntilDone_AcceptsSameHostAbsoluteLocation(t *testing.T) {
+	var gotPath string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.RequestURI()
+		_, _ = w.Write([]byte(`{}`))
+	}))
+	defer srv.Close()
+
+	if _, err := PollUntilDone(context.Background(), srv.URL, srv.URL+"/status?job=1", ClientOptions{Timeout: 5 * time.Second}, 30*time.Second); err != nil {
+		t.Fatalf("PollUntilDone: %v", err)
+	}
+	if gotPath != "/status?job=1" {
+		t.Errorf("poll path = %q, want /status?job=1", gotPath)
+	}
+}
+
+func TestPollUntilDone_AcceptsDefaultPortAbsoluteLocation(t *testing.T) {
+	var gotHost string
+	opts := ClientOptions{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+		gotHost = r.URL.Host
+		return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header)}, nil
+	})}
+
+	if _, err := PollUntilDone(context.Background(), "api.example.com", "https://api.example.com:443/status", opts, 30*time.Second); err != nil {
+		t.Fatalf("PollUntilDone: %v", err)
+	}
+	if gotHost != "api.example.com" {
+		t.Errorf("host = %q, want api.example.com", gotHost)
+	}
+}
+
+func TestPollUntilDone_RejectsCrossHostAbsoluteLocation(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Location", "https://example.com/status")
+		w.WriteHeader(http.StatusAccepted)
+	}))
+	defer srv.Close()
+
+	_, err := PollUntilDone(context.Background(), srv.URL, "/status", ClientOptions{Timeout: 5 * time.Second}, 30*time.Second)
+	if err == nil || err.Error() != `cross-host polling location "https://example.com/status"` {
+		t.Fatalf("PollUntilDone error = %v, want cross-host polling location", err)
+	}
+}
+
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
+
 func TestPollUntilDone_Timeout(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Location", "/status")