spark-mcp/internal/httpclient/redirect.go

package httpclient

import (
	"context"
	"errors"
	"fmt"
	"net/http"
)

const defaultMaxRedirects = 5

// DoWithRedirect 手动跟随 3xx,跨主机跳转保留 Authorization header。
// maxRedirects == 0 禁止 redirect(等同 Do)
// maxRedirects < 0 用 defaultMaxRedirects
func (c *Client) DoWithRedirect(ctx context.Context, req *http.Request, maxRedirects int, allowedHosts ...string) (*http.Response, error) {
	if maxRedirects == 0 {
		return c.Do(ctx, req, allowedHosts...)
	}
	if maxRedirects < 0 {
		maxRedirects = defaultMaxRedirects
	}

	current := req
	remaining := maxRedirects
	for {
		resp, err := c.Do(ctx, current, allowedHosts...)
		if err != nil {
			return nil, err
		}
		if resp.StatusCode < 300 || resp.StatusCode >= 400 {
			return resp, nil
		}

		// 3xx
		if remaining == 0 {
			_ = resp.Body.Close()
			return nil, errors.New("httpclient: too many redirects")
		}
		remaining--

		loc := resp.Header.Get("Location")
		if loc == "" {
			_ = resp.Body.Close()
			return nil, fmt.Errorf("httpclient: %d %s without Location header", resp.StatusCode, resp.Status)
		}
		next, err := current.URL.Parse(loc)
		if err != nil {
			_ = resp.Body.Close()
			return nil, fmt.Errorf("httpclient: parse redirect Location %q: %w", loc, err)
		}
		_ = resp.Body.Close()

		newReq, err := http.NewRequestWithContext(ctx, current.Method, next.String(), nil)
		if err != nil {
			return nil, fmt.Errorf("httpclient: build redirect request: %w", err)
		}
		// **保留** header (Authorization 等)
		newReq.Header = current.Header.Clone()
		newReq.Body = nil
		current = newReq
	}
}