xk-hy-forward-go/internal/forward/proxy_test.go

package forward

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
)

func TestSuperviseProxyUpstreamHeadersAndBody(t *testing.T) {
	const cipher = "BASE64_CIPHER_EXAMPLE"
	var upstreamBody []byte
	var upstreamHdr http.Header

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamBody, _ = io.ReadAll(r.Body)
		upstreamHdr = r.Header.Clone()
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"code":200}`))
	}))
	defer upstream.Close()

	proxy, err := newReverseProxy(upstream.URL, "supervise")
	if err != nil {
		t.Fatal(err)
	}

	body := bytes.NewBufferString(cipher)
	req := httptest.NewRequest(http.MethodPost, "http://forward.local/province/supervise/data", body)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Ca-Signature", "sig")
	req.Header.Set("requestBody", cipher)
	req.Header.Set("X-Ca-Appkey", "key")
	req.Header.Set("X-Ca-Encryption", "AES")
	req.Header.Set("X-Ca-Nonce", "n")
	req.Header.Set("X-Ca-Timestamp", "1716200000")
	req.Header.Set("X-Service-Id", "his.provinceDataUploadService")
	req.Header.Set("X-Service-Method", "uploadConsultIndicators")
	req.Header.Set("Cookie", "junk")
	req.Header.Set("User-Agent", "test-agent")

	rec := httptest.NewRecorder()
	proxy.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("upstream status=%d body=%s", rec.Code, rec.Body.String())
	}
	if string(upstreamBody) != cipher {
		t.Fatalf("body mismatch: got %q want %q", upstreamBody, cipher)
	}
	if upstreamHdr.Get("Requestbody") != cipher {
		t.Fatalf("requestBody header missing or wrong: %v", upstreamHdr)
	}
	if upstreamHdr.Get("X-Ca-Signature") != "sig" {
		t.Fatal("X-Ca-Signature stripped")
	}
	if upstreamHdr.Get("Cookie") != "" || upstreamHdr.Get("User-Agent") != "" {
		t.Fatalf("junk headers leaked: %v", upstreamHdr)
	}
}

func TestFileProxyStripsJunkKeepsMultipart(t *testing.T) {
	var upstreamHdr http.Header
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamHdr = r.Header.Clone()
		w.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	proxy, err := newReverseProxy(upstream.URL, "file")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodPost, "http://forward.local/mng/file/auth/upload", bytes.NewReader([]byte("x")))
	req.Header.Set("Content-Type", "multipart/form-data; boundary=abc")
	req.Header.Set("X-Authorization", "tok")
	req.Header.Set("Cookie", "c")

	rec := httptest.NewRecorder()
	proxy.ServeHTTP(rec, req)

	if upstreamHdr.Get("Content-Type") == "" || upstreamHdr.Get("X-Authorization") != "tok" {
		t.Fatalf("required headers missing: %v", upstreamHdr)
	}
	if upstreamHdr.Get("Cookie") != "" {
		t.Fatal("cookie must not reach upstream")
	}
}

func TestHandleForwardRejectsBadForwardToken(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	proxy, err := newReverseProxy(upstream.URL, "supervise")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodPost, "/province/supervise/data", bytes.NewReader([]byte("x")))
	req.Header.Set("X-Forward-Token", "wrong")

	rec := httptest.NewRecorder()
	handleForward(rec, req, "supervise", proxy, upstream.URL, "", "expected-secret", true)

	if rec.Code != http.StatusForbidden {
		t.Fatalf("want 403 got %d", rec.Code)
	}
}

func TestHandleForwardDryRunReturnsOutboundPayload(t *testing.T) {
	const cipher = "BASE64_CIPHER_EXAMPLE"
	var upstreamHits atomic.Int32

	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamHits.Add(1)
		_, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusOK)
	}))
	defer upstream.Close()

	proxy, err := newReverseProxy(upstream.URL, "supervise")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodPost, "/province/supervise/data", bytes.NewReader([]byte(cipher)))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Ca-Signature", "sig")
	req.Header.Set("Requestbody", cipher)
	req.Header.Set("Cookie", "junk")

	rec := httptest.NewRecorder()
	handleForward(rec, req, "supervise", proxy, upstream.URL, "", "", false)

	if rec.Code != http.StatusOK {
		t.Fatalf("want 200 got %d body=%s", rec.Code, rec.Body.String())
	}
	if upstreamHits.Load() != 0 {
		t.Fatal("dry run must not call upstream")
	}

	var resp dryRunResponse
	if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
		t.Fatal(err)
	}
	if !resp.DryRun || resp.Kind != "supervise" || resp.TargetURL != upstream.URL {
		t.Fatalf("unexpected dry run meta: %+v", resp)
	}
	if resp.Headers["X-Ca-Signature"][0] != "sig" {
		t.Fatalf("signature header missing: %v", resp.Headers)
	}
	if _, ok := resp.Headers["Cookie"]; ok {
		t.Fatal("cookie must be stripped in dry run headers")
	}
	if resp.Body != cipher {
		t.Fatalf("plain body mismatch: %q", resp.Body)
	}
	decoded, err := base64.StdEncoding.DecodeString(resp.BodyBase64)
	if err != nil || string(decoded) != cipher {
		t.Fatalf("base64 body mismatch: %v", err)
	}
}

func TestHandleForwardDryRunFileOmitsPlainBody(t *testing.T) {
	proxy, err := newReverseProxy("http://example.com/upload", "file")
	if err != nil {
		t.Fatal(err)
	}

	req := httptest.NewRequest(http.MethodPost, "/mng/file/auth/upload", bytes.NewReader([]byte("binary")))
	req.Header.Set("Content-Type", "multipart/form-data; boundary=abc")
	req.Header.Set("X-Authorization", "tok")

	rec := httptest.NewRecorder()
	handleForward(rec, req, "file", proxy, "http://example.com/upload", "", "", false)

	var resp dryRunResponse
	if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
		t.Fatal(err)
	}
	if resp.Body != "" {
		t.Fatalf("file channel should not include plain body, got %q", resp.Body)
	}
	if resp.BodyBase64 == "" {
		t.Fatal("body_base64 required")
	}
}