Files
xk-hy-forward-go/internal/forward/proxy_test.go
2026-05-22 08:42:27 +08:00

113 lines
3.4 KiB
Go

package forward
import (
"bytes"
"io"
"net/http"
"net/http/httptest"
"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")
if rec.Code != http.StatusForbidden {
t.Fatalf("want 403 got %d", rec.Code)
}
}