1 files changed, 139 insertions, 0 deletions
diff --git a/internal/proxy/dialer_test.go b/internal/proxy/dialer_test.go
new file mode 100644
index 0000000..14746c9
--- /dev/null
+++ b/internal/proxy/dialer_test.go
@@ -0,0 +1,139 @@
+package proxy
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+)
+
+// startStubConnectProxy runs a minimal HTTP CONNECT proxy for the
+// duration of the test: it accepts one CONNECT request, replies with the
+// given status, and if status is 200 splices the tunnel through to a
+// real dial of the requested host. Returns the proxy's address.
+func startStubConnectProxy(t *testing.T, status int) string {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	t.Cleanup(func() { ln.Close() })
+
+	go func() {
+		c, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer c.Close()
+		br := bufio.NewReader(c)
+		req, err := http.ReadRequest(br)
+		if err != nil {
+			return
+		}
+		if req.Method != http.MethodConnect {
+			c.Write([]byte("HTTP/1.1 405 Method Not Allowed\r\n\r\n"))
+			return
+		}
+		if status != http.StatusOK {
+			c.Write([]byte("HTTP/1.1 403 Forbidden\r\n\r\n"))
+			return
+		}
+		target, err := net.Dial("tcp", req.Host)
+		if err != nil {
+			c.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
+			return
+		}
+		defer target.Close()
+		c.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
+		done := make(chan struct{}, 2)
+		go func() { io.Copy(target, br); done <- struct{}{} }()
+		go func() { io.Copy(c, target); done <- struct{}{} }()
+		<-done
+	}()
+	return ln.Addr().String()
+}
+
+// startEchoServer runs a TCP server that echoes back whatever it reads,
+// standing in for "the origin" on the far side of a CONNECT tunnel.
+func startEchoServer(t *testing.T) string {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	t.Cleanup(func() { ln.Close() })
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				io.Copy(c, c)
+			}(c)
+		}
+	}()
+	return ln.Addr().String()
+}
+
+func TestDialViaProxyDirect(t *testing.T) {
+	echoAddr := startEchoServer(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	conn, err := dialViaProxy(ctx, echoAddr, "")
+	if err != nil {
+		t.Fatalf("dialViaProxy direct: %v", err)
+	}
+	defer conn.Close()
+
+	if _, err := conn.Write([]byte("hello")); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	buf := make([]byte, 5)
+	if _, err := io.ReadFull(conn, buf); err != nil {
+		t.Fatalf("read: %v", err)
+	}
+	if string(buf) != "hello" {
+		t.Errorf("got %q, want %q", buf, "hello")
+	}
+}
+
+func TestDialViaProxyTunneled(t *testing.T) {
+	echoAddr := startEchoServer(t)
+	proxyAddr := startStubConnectProxy(t, http.StatusOK)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	conn, err := dialViaProxy(ctx, echoAddr, proxyAddr)
+	if err != nil {
+		t.Fatalf("dialViaProxy via proxy: %v", err)
+	}
+	defer conn.Close()
+
+	if _, err := conn.Write([]byte("world")); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	buf := make([]byte, 5)
+	if _, err := io.ReadFull(conn, buf); err != nil {
+		t.Fatalf("read: %v", err)
+	}
+	if string(buf) != "world" {
+		t.Errorf("got %q, want %q", buf, "world")
+	}
+}
+
+func TestDialViaProxyRejected(t *testing.T) {
+	proxyAddr := startStubConnectProxy(t, http.StatusForbidden)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+
+	_, err := dialViaProxy(ctx, "example.invalid:443", proxyAddr)
+	if err == nil {
+		t.Fatal("expected an error when the upstream proxy refuses CONNECT, got nil")
+	}
+}