1424 lines · 40.5 KB · Go
RawDownload
package pluginhost import (	"bufio"	"context"	"crypto/tls"	"encoding/base64"	"io"	"net"	"net/http"	"net/http/httptest"	"net/url"	"strings"	"sync/atomic"	"testing"	"time" 	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi") func TestHostHTTPClientMarksUpstreamAttempt(t *testing.T) {	t.Parallel() 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {		_, _ = io.WriteString(w, "ok")	}))	t.Cleanup(server.Close) 	client := New().newHTTPClient(nil)	for _, test := range []struct {		name string		do   func(context.Context) error	}{		{			name: "buffered",			do: func(ctx context.Context) error {				_, errDo := client.Do(ctx, pluginapi.HTTPRequest{URL: server.URL})				return errDo			},		},		{			name: "stream",			do: func(ctx context.Context) error {				response, errDo := client.DoStream(ctx, pluginapi.HTTPRequest{URL: server.URL})				if errDo != nil {					return errDo				}				for range response.Chunks {				}				return nil			},		},	} {		t.Run(test.name, func(t *testing.T) {			ctx := cliproxyexecutor.WithUpstreamAttemptTracker(context.Background())			if errDo := test.do(ctx); errDo != nil {				t.Fatalf("host HTTP request error = %v", errDo)			}			if !cliproxyexecutor.UpstreamAttempted(ctx) {				t.Fatal("host HTTP request did not mark an upstream attempt")			}		})	}} func TestHostHTTPClientAppliesWireProfile(t *testing.T) {	t.Parallel() 	listener, errListen := net.Listen("tcp", "127.0.0.1:0")	if errListen != nil {		t.Fatalf("listen tcp: %v", errListen)	}	defer func() {		if errClose := listener.Close(); errClose != nil {			t.Logf("close listener error: %v", errClose)		}	}() 	type capturedRequest struct {		rawHeader string	}	captured := make(chan capturedRequest, 2) 	go func() {		for {			conn, errAccept := listener.Accept()			if errAccept != nil {				return			}			go func(c net.Conn) {				defer func() {					if errClose := c.Close(); errClose != nil {						t.Logf("close connection error: %v", errClose)					}				}()				buf := make([]byte, 4096)				n, errRead := c.Read(buf)				if errRead != nil && errRead != io.EOF {					return				}				data := string(buf[:n])				headerEnd := strings.Index(data, "\r\n\r\n")				if headerEnd >= 0 {					captured <- capturedRequest{rawHeader: data[:headerEnd]}				}				resp := "HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"				_, _ = c.Write([]byte(resp))			}(conn)		}	}() 	client := New().newHTTPClient(nil)	profile := &pluginapi.HTTPWireProfile{		HTTP1Only:              true,		DisableAutoCompression: true,		HeaderProfile:          []string{"x-custom-b", "X-Custom-A", "User-Agent", "Host"},	} 	for _, mode := range []string{"buffered", "stream"} {		req := pluginapi.HTTPRequest{			Method: http.MethodGet,			URL:    "http://" + listener.Addr().String() + "/test",			Headers: http.Header{				"X-Custom-A": []string{"value-a"},				"x-custom-b": []string{"value-b"},				"User-Agent": []string{"test-agent"},			},			WireProfile: profile,		} 		if mode == "buffered" {			resp, errDo := client.Do(context.Background(), req)			if errDo != nil {				t.Fatalf("Do error: %v", errDo)			}			if string(resp.Body) != "ok" {				t.Fatalf("unexpected body: %q", string(resp.Body))			}		} else {			resp, errDo := client.DoStream(context.Background(), req)			if errDo != nil {				t.Fatalf("DoStream error: %v", errDo)			}			var body []byte			for chunk := range resp.Chunks {				if chunk.Err != nil {					t.Fatalf("stream chunk error: %v", chunk.Err)				}				body = append(body, chunk.Payload...)			}			if string(body) != "ok" {				t.Fatalf("unexpected stream body: %q", string(body))			}		} 		select {		case capReq := <-captured:			lines := strings.Split(capReq.rawHeader, "\r\n")			if len(lines) == 0 || !strings.HasPrefix(lines[0], "GET /test HTTP/1.1") {				t.Fatalf("unexpected request line: %q", capReq.rawHeader)			}			expectedHeaders := []string{"x-custom-b:", "X-Custom-A:", "User-Agent:", "Host:"}			headerIndex := 0			for _, line := range lines[1:] {				if headerIndex < len(expectedHeaders) && strings.HasPrefix(line, expectedHeaders[headerIndex]) {					headerIndex++				}				if strings.HasPrefix(strings.ToLower(line), "accept-encoding:") {					t.Fatalf("expected no automatic Accept-Encoding when DisableAutoCompression is set, got: %s", line)				}			}			if headerIndex != len(expectedHeaders) {				t.Fatalf("headers did not match expected order/casing: matched %d/%d; raw headers:\n%s",					headerIndex, len(expectedHeaders), capReq.rawHeader)			}		case <-time.After(3 * time.Second):			t.Fatalf("timeout waiting for %s request to reach server", mode)		}	}} func TestHostHTTPClientWireProfile_DisableAutoCompression_PreservesHTTP2(t *testing.T) {	t.Parallel() 	client := New().newHTTPClient(nil).(*hostHTTPClient)	profile := &pluginapi.HTTPWireProfile{		DisableAutoCompression: true,	} 	httpReq, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "https://example.com", nil)	httpClient, cleanup, errClient := client.newHTTPClientForRequest(context.Background(), nil, pluginapi.HTTPRequest{		URL:         "https://example.com",		WireProfile: profile,	}, httpReq)	if errClient != nil {		t.Fatalf("newHTTPClientForRequest error = %v", errClient)	}	defer cleanup() 	transport, ok := httpClient.Transport.(*http.Transport)	if !ok {		t.Fatalf("transport type = %T, want *http.Transport", httpClient.Transport)	}	if !transport.DisableCompression {		t.Fatal("expected DisableCompression = true")	}	if !transport.ForceAttemptHTTP2 {		t.Fatal("expected ForceAttemptHTTP2 = true when HTTP1Only is false")	}	if transport.TLSNextProto != nil {		t.Fatal("expected TLSNextProto = nil when HTTP1Only is false")	}} type dummyRoundTripper struct{} func (d *dummyRoundTripper) RoundTrip(*http.Request) (*http.Response, error) {	return nil, nil} func TestHostHTTPClientWireProfile_CustomRoundTripperValidation(t *testing.T) {	t.Parallel() 	client := New().newHTTPClient(nil)	ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", &dummyRoundTripper{}) 	_, errDo := client.Do(ctx, pluginapi.HTTPRequest{		URL: "http://example.com",		WireProfile: &pluginapi.HTTPWireProfile{			HeaderProfile: []string{"Host"},		},	})	if errDo == nil {		t.Fatal("expected error with custom non-*http.Transport roundtripper, got nil")	}	if !strings.Contains(errDo.Error(), "wire profile is not supported with custom context RoundTripper") {		t.Fatalf("unexpected error message: %v", errDo)	}} func TestHostHTTPClientWireProfile_PlainHTTPProxyUsesStandardProxy(t *testing.T) {	t.Parallel() 	receivedMethod := make(chan string, 1)	proxyServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		receivedMethod <- r.Method		w.WriteHeader(http.StatusOK)		_, _ = w.Write([]byte("proxied"))	}))	t.Cleanup(proxyServer.Close) 	client := New().newHTTPClient(nil).(*hostHTTPClient)	profile := &pluginapi.HTTPWireProfile{		HTTP1Only:     true,		HeaderProfile: []string{"Host", "User-Agent"},	} 	cfg := &config.Config{		SDKConfig: config.SDKConfig{			ProxyURL: proxyServer.URL,		},	} 	httpReq, errReq := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://example.org/resource", nil)	if errReq != nil {		t.Fatalf("new request error = %v", errReq)	} 	httpClient, cleanup, errClient := client.newHTTPClientForRequest(context.Background(), cfg, pluginapi.HTTPRequest{		URL:         "http://example.org/resource",		WireProfile: profile,	}, httpReq)	if errClient != nil {		t.Fatalf("newHTTPClientForRequest error = %v", errClient)	}	defer cleanup() 	resp, errDo := httpClient.Do(httpReq)	if errDo != nil {		t.Fatalf("execute request error = %v", errDo)	}	defer func() {		_ = resp.Body.Close()	}() 	select {	case method := <-receivedMethod:		if method != http.MethodGet {			t.Fatalf("proxy received method %q, want %q (should not use CONNECT for plain http)", method, http.MethodGet)		}	case <-time.After(3 * time.Second):		t.Fatal("timeout waiting for proxy request")	}} func TestHostHTTPClientWireProfile_ClosesIdleConnectionsOnCompletion(t *testing.T) {	t.Parallel() 	closedConns := make(chan struct{}, 1)	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {		_, _ = io.WriteString(w, "ok")	}))	server.Config.ConnState = func(c net.Conn, state http.ConnState) {		if state == http.StateClosed {			select {			case closedConns <- struct{}{}:			default:			}		}	}	server.Start()	t.Cleanup(server.Close) 	client := New().newHTTPClient(nil)	profile := &pluginapi.HTTPWireProfile{		HTTP1Only: true,	} 	resp, errDo := client.Do(context.Background(), pluginapi.HTTPRequest{		URL:         server.URL,		WireProfile: profile,	})	if errDo != nil {		t.Fatalf("Do error = %v", errDo)	}	if string(resp.Body) != "ok" {		t.Fatalf("body = %q, want ok", string(resp.Body))	} 	select {	case <-closedConns:		// Confirmed idle connection was actively closed by cleanup	case <-time.After(3 * time.Second):		t.Fatal("expected idle connection to be closed by transport cleanup, but timed out")	}} func TestHostHTTPClientWireProfile_ProxyPriorityOverContextRoundTripper(t *testing.T) {	t.Parallel() 	proxyReceived := make(chan struct{}, 1)	proxyServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		select {		case proxyReceived <- struct{}{}:		default:		}		w.WriteHeader(http.StatusOK)		_, _ = w.Write([]byte("proxied"))	}))	t.Cleanup(proxyServer.Close) 	// Auth has proxy configured	auth := &coreauth.Auth{		ProxyURL: proxyServer.URL,	}	client := New().newHTTPClient(auth) 	// Context contains a dummy roundtripper	ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", &dummyRoundTripper{}) 	resp, errDo := client.Do(ctx, pluginapi.HTTPRequest{		URL: "http://example.org/priority-test",		WireProfile: &pluginapi.HTTPWireProfile{			HeaderProfile: []string{"Host"},		},	})	if errDo != nil {		t.Fatalf("Do error = %v", errDo)	}	if string(resp.Body) != "proxied" {		t.Fatalf("body = %q, want proxied", string(resp.Body))	} 	select {	case <-proxyReceived:		// Confirmed auth proxy was used instead of context roundtripper	case <-time.After(3 * time.Second):		t.Fatal("expected auth proxy to receive request, timed out")	}} func TestHostHTTPClientWireProfile_HTTPToHTTPSRedirect(t *testing.T) {	t.Parallel() 	var httpsServer *httptest.Server	httpsServer = httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		w.WriteHeader(http.StatusOK)		_, _ = w.Write([]byte("https-ok"))	}))	t.Cleanup(httpsServer.Close) 	httpServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		http.Redirect(w, r, httpsServer.URL+"/redirected", http.StatusFound)	}))	t.Cleanup(httpServer.Close) 	// Pass the TLS test server's client transport certificate to the host	host := New()	client := host.newHTTPClient(nil).(*hostHTTPClient) 	// Provide a base transport that trusts the test TLS certificate	customTransport := httpsServer.Client().Transport.(*http.Transport).Clone()	ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", customTransport) 	profile := &pluginapi.HTTPWireProfile{		HTTP1Only:     true,		HeaderProfile: []string{"Host", "User-Agent"},	} 	resp, errDo := client.Do(ctx, pluginapi.HTTPRequest{		URL:         httpServer.URL + "/start",		WireProfile: profile,	})	if errDo != nil {		t.Fatalf("Do error across HTTP->HTTPS redirect: %v", errDo)	}	if string(resp.Body) != "https-ok" {		t.Fatalf("body = %q, want https-ok", string(resp.Body))	}} func TestHostHTTPClientWireProfile_SOCKS5ProxyWithHeaderProfile(t *testing.T) {	t.Parallel() 	backendReceived := make(chan string, 1)	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {		backendReceived <- r.Header.Get("X-Custom-A")		w.WriteHeader(http.StatusOK)		_, _ = w.Write([]byte("socks5-ok"))	}))	t.Cleanup(backend.Close) 	socksListener, errListen := net.Listen("tcp", "127.0.0.1:0")	if errListen != nil {		t.Fatalf("listen socks5 error: %v", errListen)	}	t.Cleanup(func() { _ = socksListener.Close() }) 	go func() {		conn, errAccept := socksListener.Accept()		if errAccept != nil {			return		}		defer conn.Close() 		buf := make([]byte, 256)		n, errRead := conn.Read(buf)		if errRead != nil || n < 3 || buf[0] != 0x05 {			return		}		_, _ = conn.Write([]byte{0x05, 0x00}) 		n, errRead = conn.Read(buf)		if errRead != nil || n < 4 || buf[1] != 0x01 {			return		} 		backendConn, errDial := net.Dial("tcp", backend.Listener.Addr().String())		if errDial != nil {			_, _ = conn.Write([]byte{0x05, 0x01, 0x00, 0x01, 0, 0, 0, 0, 0, 0})			return		}		defer backendConn.Close() 		_, _ = conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}) 		done := make(chan struct{})		go func() {			_, _ = io.Copy(backendConn, conn)			close(done)		}()		_, _ = io.Copy(conn, backendConn)		<-done	}() 	auth := &coreauth.Auth{		ProxyURL: "socks5://" + socksListener.Addr().String(),	}	client := New().newHTTPClient(auth)	profile := &pluginapi.HTTPWireProfile{		HTTP1Only:     true,		HeaderProfile: []string{"X-Custom-A", "Host"},	} 	resp, errDo := client.Do(context.Background(), pluginapi.HTTPRequest{		URL:         backend.URL + "/test",		Headers:     http.Header{"X-Custom-A": []string{"custom-val"}},		WireProfile: profile,	})	if errDo != nil {		t.Fatalf("Do error through SOCKS5: %v", errDo)	}	if string(resp.Body) != "socks5-ok" {		t.Fatalf("body = %q, want socks5-ok", string(resp.Body))	} 	select {

Lines 1–500 of 1424.