1package tinygit23import (4 "context"5 "crypto/tls"6 "io"7 "net/http"8 "net/http/httptest"9 "net/url"10 "strings"11 "sync/atomic"12 "testing"13 "time"14)1516func repairProxyClient(t *testing.T, ctx context.Context, target string, budget *byteBudget, limit int64) (*repairProxy, *http.Client) {17 t.Helper()18 u, err := url.Parse(target)19 if err != nil {20 t.Fatal(err)21 }22 proxy, err := newRepairProxy(ctx, u.Hostname(), u.Hostname(), u.Port(), budget, limit)23 if err != nil {24 t.Fatal(err)25 }26 t.Cleanup(func() { _ = proxy.Close() })27 proxyURL, _ := url.Parse(proxy.URL())28 transport := &http.Transport{Proxy: http.ProxyURL(proxyURL), TLSClientConfig: &tls.Config{InsecureSkipVerify: true}} // Local test server's self-signed certificate.29 t.Cleanup(transport.CloseIdleConnections)30 return proxy, &http.Client{Transport: transport, Timeout: 5 * time.Second}31}3233func TestRepairProxyEnforcesSharedResponseBudget(t *testing.T) {34 var requests atomic.Int3235 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {36 requests.Add(1)37 if r.URL.Path == "/small" {38 _, _ = io.WriteString(w, "first")39 return40 }41 _, _ = io.WriteString(w, strings.Repeat("x", 8192))42 }))43 defer server.Close()44 budget := &byteBudget{}45 _, client := repairProxyClient(t, context.Background(), server.URL, budget, 2048)46 response, err := client.Get(server.URL + "/small")47 if err != nil {48 t.Fatal(err)49 }50 body, err := io.ReadAll(response.Body)51 response.Body.Close()52 if err != nil || string(body) != "first" {53 t.Fatalf("first response = %q, %v", body, err)54 }55 before := budget.used.Load()56 if before == 0 || before >= 2048 {57 t.Fatalf("first request consumed %d bytes", before)58 }59 response, err = client.Get(server.URL + "/large")60 if err == nil {61 body, err = io.ReadAll(response.Body)62 response.Body.Close()63 if len(body) >= 2048 {64 t.Fatalf("oversized response leaked %d bytes", len(body))65 }66 }67 if err == nil {68 t.Fatal("oversized response completed")69 }70 if budget.used.Load() != 2048 {71 t.Fatalf("budget = %d, want 2048", budget.used.Load())72 }73 if requests.Load() != 2 {74 t.Fatalf("upstream requests = %d, want two", requests.Load())75 }76 if response, err = client.Get(server.URL + "/small"); err == nil {77 response.Body.Close()78 t.Fatal("request passed after shared budget exhausted")79 }80 if requests.Load() != 2 {81 t.Fatal("exhausted proxy contacted upstream")82 }83}8485func TestRepairProxyLimitsHTTPSTunnel(t *testing.T) {86 server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {87 _, _ = io.WriteString(w, strings.Repeat("x", 256<<10))88 }))89 defer server.Close()90 budget := &byteBudget{}91 _, client := repairProxyClient(t, context.Background(), server.URL, budget, 64<<10)92 response, err := client.Get(server.URL)93 if err == nil {94 body, bodyErr := io.ReadAll(response.Body)95 response.Body.Close()96 if len(body) >= 64<<10 {97 t.Fatalf("HTTPS body exceeded wire budget: %d", len(body))98 }99 err = bodyErr100 }101 if err == nil {102 t.Fatal("oversized HTTPS body completed")103 }104 if budget.used.Load() != 64<<10 {105 t.Fatalf("HTTPS budget = %d", budget.used.Load())106 }107}108109func TestRepairProxyCancellationInterruptsStalledUpstream(t *testing.T) {110 started := make(chan struct{})111 server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {112 close(started)113 <-r.Context().Done()114 }))115 defer server.Close()116 ctx, cancel := context.WithCancel(context.Background())117 defer cancel()118 _, client := repairProxyClient(t, ctx, server.URL, &byteBudget{}, 4096)119 done := make(chan error, 1)120 go func() {121 response, err := client.Get(server.URL)122 if response != nil {123 response.Body.Close()124 }125 done <- err126 }()127 select {128 case <-started:129 case <-time.After(5 * time.Second):130 t.Fatal("upstream did not receive request")131 }132 cancel()133 select {134 case err := <-done:135 if err == nil {136 t.Fatal("canceled request succeeded")137 }138 case <-time.After(time.Second):139 t.Fatal("cancel did not interrupt stalled request")140 }141}142143func TestRepairProxyRejectsAnotherPort(t *testing.T) {144 var contacted atomic.Bool145 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { contacted.Store(true) }))146 defer server.Close()147 _, client := repairProxyClient(t, context.Background(), "http://127.0.0.1:1", &byteBudget{}, 4096)148 if response, err := client.Get(server.URL); err == nil {149 response.Body.Close()150 t.Fatal("proxy accepted a different port")151 }152 if contacted.Load() {153 t.Fatal("proxy contacted a different target")154 }155}