proxy/internal/probe/probe_test.go

package probe_test

import (
	"bufio"
	"context"
	"io"
	"net"
	"testing"
	"time"

	"proxyctl/internal/probe"
)

// fakeEndpoint starts a listener on 127.0.0.1:0 and serves each accepted
// connection with handler until the test ends.
func fakeEndpoint(t *testing.T, handler func(net.Conn)) string {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { listener.Close() })
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}
			go handler(conn)
		}
	}()
	return listener.Addr().String()
}

// serveSOCKS5 answers the SOCKS5 greeting and closes anything else so the
// HTTP check fails fast instead of waiting out its deadline.
func serveSOCKS5(conn net.Conn) {
	defer conn.Close()
	first := make([]byte, 1)
	if _, err := io.ReadFull(conn, first); err != nil {
		return
	}
	if first[0] != 0x05 {
		return
	}
	if _, err := io.ReadFull(conn, make([]byte, 2)); err != nil {
		return
	}
	conn.Write([]byte{0x05, 0x00})
}

func TestDetectRecognizesSOCKS5Endpoint(t *testing.T) {
	address := fakeEndpoint(t, serveSOCKS5)
	caps, err := probe.Detect(context.Background(), address)
	if err != nil {
		t.Fatalf("Detect returned error: %v", err)
	}
	if !caps.SOCKS5 || caps.HTTP {
		t.Fatalf("Detect = %#v, want SOCKS5 only", caps)
	}
}

func TestDetectSucceedsWhenSOCKS5RejectsNoAuth(t *testing.T) {
	// A SOCKS5 server may answer 0xFF ("no acceptable methods") to the
	// no-auth greeting; the version byte still identifies the protocol.
	address := fakeEndpoint(t, func(conn net.Conn) {
		defer conn.Close()
		first := make([]byte, 1)
		if _, err := io.ReadFull(conn, first); err != nil {
			return
		}
		if first[0] != 0x05 {
			return
		}
		if _, err := io.ReadFull(conn, make([]byte, 2)); err != nil {
			return
		}
		conn.Write([]byte{0x05, 0xFF})
	})
	caps, err := probe.Detect(context.Background(), address)
	if err != nil {
		t.Fatalf("Detect returned error: %v", err)
	}
	if !caps.SOCKS5 {
		t.Fatalf("Detect = %#v, want SOCKS5", caps)
	}
}

func TestDetectHonorsContextCancellation(t *testing.T) {
	// A silent endpoint accepts connections but never answers; cancellation
	// must cut the attempt short instead of waiting out the full timeout.
	address := fakeEndpoint(t, func(conn net.Conn) {
		defer conn.Close()
		time.Sleep(10 * time.Second)
	})
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	start := time.Now()
	if _, err := probe.Detect(ctx, address); err == nil {
		t.Fatal("Detect on silent endpoint with cancelled context returned nil error")
	}
	if elapsed := time.Since(start); elapsed > 3*time.Second {
		t.Fatalf("Detect ignored context cancellation, took %s", elapsed)
	}
}

// serveHTTPProxy closes the SOCKS5 attempt fast, reads the CONNECT request,
// and answers with an HTTP error status line like a real proxy refusing its
// own loopback discard port would.
func serveHTTPProxy(conn net.Conn) {
	defer conn.Close()
	first := make([]byte, 1)
	if _, err := io.ReadFull(conn, first); err != nil {
		return
	}
	if first[0] == 0x05 {
		return
	}
	reader := bufio.NewReader(conn)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			return
		}
		if line == "\r\n" {
			break
		}
	}
	conn.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
}

func TestDetectRecognizesHTTPProxyEndpoint(t *testing.T) {
	address := fakeEndpoint(t, serveHTTPProxy)
	caps, err := probe.Detect(context.Background(), address)
	if err != nil {
		t.Fatalf("Detect returned error: %v", err)
	}
	if !caps.HTTP || caps.SOCKS5 {
		t.Fatalf("Detect = %#v, want HTTP only", caps)
	}
}

func TestDetectReportsBothProtocolsOnMixedPort(t *testing.T) {
	// Both checks always run so a port answering both handshakes is
	// reported truthfully instead of stopping at the first hit.
	address := fakeEndpoint(t, func(conn net.Conn) {
		defer conn.Close()
		first := make([]byte, 1)
		if _, err := io.ReadFull(conn, first); err != nil {
			return
		}
		if first[0] == 0x05 {
			if _, err := io.ReadFull(conn, make([]byte, 2)); err != nil {
				return
			}
			conn.Write([]byte{0x05, 0x00})
			return
		}
		conn.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))
	})
	caps, err := probe.Detect(context.Background(), address)
	if err != nil {
		t.Fatalf("Detect returned error: %v", err)
	}
	if !caps.HTTP || !caps.SOCKS5 {
		t.Fatalf("Detect = %#v, want both HTTP and SOCKS5", caps)
	}
}

func TestDetectReturnsDialErrorForClosedPort(t *testing.T) {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	address := listener.Addr().String()
	listener.Close()
	caps, err := probe.Detect(context.Background(), address)
	if err == nil {
		t.Fatalf("Detect on closed port = %#v, nil error; want dial error", caps)
	}
}

func TestDetectReturnsZeroCapabilitiesForUnrecognizedEndpoint(t *testing.T) {
	address := fakeEndpoint(t, func(conn net.Conn) {
		defer conn.Close()
		conn.Write([]byte("garbage\r\n"))
	})
	caps, err := probe.Detect(context.Background(), address)
	if err != nil {
		t.Fatalf("Detect on unrecognized endpoint returned error: %v", err)
	}
	if caps.HTTP || caps.SOCKS5 {
		t.Fatalf("Detect = %#v, want zero Capabilities", caps)
	}
}

func TestURLMapsCapabilitiesToScheme(t *testing.T) {
	tests := []struct {
		name string
		raw  string
		caps probe.Capabilities
		want string
	}{
		{"http detected", "localhost:8888", probe.Capabilities{HTTP: true}, "http://localhost:8888"},
		{"http preferred over socks5", "localhost:8888", probe.Capabilities{HTTP: true, SOCKS5: true}, "http://localhost:8888"},
		{"socks5 only uses socks5h", "localhost:8888", probe.Capabilities{SOCKS5: true}, "socks5h://localhost:8888"},
		{"userinfo preserved", "user:secret@localhost:8888", probe.Capabilities{SOCKS5: true}, "socks5h://user:secret@localhost:8888"},
		{"typed scheme replaced", "socks5://localhost:1080", probe.Capabilities{HTTP: true}, "http://localhost:1080"},
		{"typed scheme and userinfo", "http://user:secret@localhost:8080", probe.Capabilities{SOCKS5: true}, "socks5h://user:secret@localhost:8080"},
		{"IPv6 endpoint", "[::1]:8888", probe.Capabilities{HTTP: true}, "http://[::1]:8888"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := probe.URL(tt.raw, tt.caps)
			if err != nil {
				t.Fatalf("URL(%q) returned error: %v", tt.raw, err)
			}
			if got != tt.want {
				t.Fatalf("URL(%q) = %q, want %q", tt.raw, got, tt.want)
			}
		})
	}
}

func TestURLPreservesRawUserinfoBytesWithoutEscaping(t *testing.T) {
	got, err := probe.URL("user:p@ss word@proxy.example:8080", probe.Capabilities{HTTP: true})
	if err != nil {
		t.Fatalf("URL returned error: %v", err)
	}
	want := "http://user:p@ss word@proxy.example:8080"
	if got != want {
		t.Fatalf("URL = %q, want %q", got, want)
	}
}

func TestURLRejectsInputWithoutHostPort(t *testing.T) {
	for _, raw := range []string{"", "localhost", "http://localhost", "user:secret@"} {
		if got, err := probe.URL(raw, probe.Capabilities{HTTP: true}); err == nil {
			t.Fatalf("URL(%q) = %q, want error", raw, got)
		}
	}
	if got, err := probe.URL("localhost:8888", probe.Capabilities{}); err == nil {
		t.Fatalf("URL with no detected protocol = %q, want error", got)
	}
}