spark-mcp/internal/httpclient/redirect_test.go

package httpclient

import (
	"context"
	"encoding/base64"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestDoWithRedirect_SameHost(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/start":
			w.Header().Set("Location", "/final")
			w.WriteHeader(http.StatusFound)
		case "/final":
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("ok"))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer srv.Close()

	c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
	req, err := http.NewRequest(http.MethodGet, srv.URL+"/start", nil)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	resp, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1")
	if err != nil {
		t.Fatalf("do with redirect: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Fatalf("status=%d, want 200", resp.StatusCode)
	}
	body, _ := io.ReadAll(resp.Body)
	if string(body) != "ok" {
		t.Fatalf("body=%q, want ok", string(body))
	}
}

func TestDoWithRedirect_CrossHost_PreservesAuth(t *testing.T) {
	srv2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(r.Header.Get("Authorization")))
	}))
	defer srv2.Close()

	srv1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/start" {
			w.Header().Set("Location", srv2.URL+"/final")
			w.WriteHeader(http.StatusTemporaryRedirect)
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer srv1.Close()

	c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
	req, err := http.NewRequest(http.MethodGet, srv1.URL+"/start", nil)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	req.SetBasicAuth("u", "p")

	resp, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1")
	if err != nil {
		t.Fatalf("do with redirect: %v", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read body: %v", err)
	}
	want := "Basic " + base64.StdEncoding.EncodeToString([]byte("u:p"))
	if string(body) != want {
		t.Fatalf("body=%q, want %q", string(body), want)
	}
}

func TestDoWithRedirect_MaxLimit(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Location", "/loop")
		w.WriteHeader(http.StatusFound)
	}))
	defer srv.Close()

	c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
	req, _ := http.NewRequest(http.MethodGet, srv.URL+"/loop", nil)
	_, err := c.DoWithRedirect(context.Background(), req, 2, "127.0.0.1")
	if err == nil {
		t.Fatal("expected too many redirects error")
	}
	if !strings.Contains(err.Error(), "too many redirects") {
		t.Fatalf("expected too many redirects, got %v", err)
	}
}

func TestDoWithRedirect_NoLocation(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusFound)
	}))
	defer srv.Close()

	c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
	req, _ := http.NewRequest(http.MethodGet, srv.URL, nil)
	_, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1")
	if err == nil {
		t.Fatal("expected error for missing Location")
	}
	if !strings.Contains(err.Error(), "without Location header") {
		t.Fatalf("expected missing Location error, got %v", err)
	}
}

func TestDoWithRedirect_ZeroDisallows(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Location", "/elsewhere")
		w.WriteHeader(http.StatusFound)
	}))
	defer srv.Close()

	c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
	req, _ := http.NewRequest(http.MethodGet, srv.URL, nil)
	resp, err := c.DoWithRedirect(context.Background(), req, 0, "127.0.0.1")
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusFound {
		t.Fatalf("status=%d, want 302", resp.StatusCode)
	}
}

// TestDoWithRedirect_DeferredRFC9110Gaps documents two known
// non-conformances with RFC 9110 §15.4 that we have NOT fixed because
// the current upstream YARN / SHS API surface only triggers 307 on
// GET amContainerLogs, where neither gap is fired.
//
// The assertions below verify the CURRENT (gap) behavior.  Any future
// fix must update these assertions deliberately — a silent flip would
// mean the gap fix passed without anyone thinking about 303 method
// switch or 307-POST body resend.  See the TODO comment in
// redirect.go for the full discussion.
//
//	(1) 303 See Other: we keep the original method (RFC requires switch
//	    to GET).  The body IS correctly dropped — that's RFC-compliant
//	    even with the wrong method.
//	(2) 307 Temporary Redirect / 308 Permanent Redirect on a POST: we
//	    keep the method (correct) but drop the body (RFC requires
//	    resend with the same method).
func TestDoWithRedirect_DeferredRFC9110Gaps(t *testing.T) {
	t.Run("303 keeps method, should switch to GET", func(t *testing.T) {
		var hitMethod, hitBody string
		target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			hitMethod = r.Method
			b, _ := io.ReadAll(r.Body)
			hitBody = string(b)
			w.WriteHeader(http.StatusOK)
		}))
		defer target.Close()

		origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Location", target.URL+"/see-other")
			w.WriteHeader(http.StatusSeeOther)
		}))
		defer origin.Close()

		c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
		req, _ := http.NewRequest(http.MethodPost, origin.URL+"/start", strings.NewReader("payload-303"))
		if _, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1"); err != nil {
			t.Fatalf("DoWithRedirect: %v", err)
		}

		// CURRENT (gap) behavior.  RFC-correct behavior: hitMethod == "GET".
		if hitMethod != http.MethodPost {
			t.Errorf("303 arrived as %s, want POST (current gap; RFC says GET)", hitMethod)
		}
		// Body must be empty — that part is RFC-compliant (303 always
		// drops body, even when the method is wrong).
		if hitBody != "" {
			t.Errorf("303 body=%q, want empty (RFC requires drop regardless of method)", hitBody)
		}
	})

	t.Run("307 on POST drops body, should resend", func(t *testing.T) {
		var hitMethod, hitBody string
		target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			hitMethod = r.Method
			b, _ := io.ReadAll(r.Body)
			hitBody = string(b)
			w.WriteHeader(http.StatusOK)
		}))
		defer target.Close()

		origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Location", target.URL+"/temp")
			w.WriteHeader(http.StatusTemporaryRedirect)
		}))
		defer origin.Close()

		c := New(Config{Timeout: 5 * time.Second, MaxResponseBytes: 1024})
		req, _ := http.NewRequest(http.MethodPost, origin.URL+"/start", strings.NewReader("payload-307"))
		if _, err := c.DoWithRedirect(context.Background(), req, 5, "127.0.0.1"); err != nil {
			t.Fatalf("DoWithRedirect: %v", err)
		}

		// Method preserved (correct).
		if hitMethod != http.MethodPost {
			t.Errorf("307 arrived as %s, want POST", hitMethod)
		}
		// CURRENT (gap) behavior.  RFC-correct behavior: hitBody == "payload-307".
		if hitBody != "" {
			t.Errorf("307 body=%q, want empty (current gap; RFC says resend payload-307)", hitBody)
		}
	})
}