Merge c30263dfa1 into 487da737b9

proxy/http/client.go

@@ -25,13 +25,20 @@ func (s *HTTP) Addr() string {
 }
 
 // Dial connects to the address addr on the network net via the proxy.
-func (s *HTTP) Dial(network, addr string) (net.Conn, error) {
+func (s *HTTP) Dial(network, addr string) (_ net.Conn, err error) {
 	rc, err := s.dialer.Dial(network, s.addr)
 	if err != nil {
 		log.F("[http] dial to %s error: %s", s.addr, err)
 		return nil, err
 	}
 
+	c := proxy.NewConn(rc)
+	defer func() {
+		if err != nil {
+			c.Close()
+		}
+	}()
+
 	buf := pool.GetBytesBuffer()
 	defer pool.PutBytesBuffer(buf)
 
@@ -51,17 +58,18 @@ func (s *HTTP) Dial(network, addr string) (net.Conn, error) {
 		return nil, err
 	}
 
-	c := proxy.NewConn(rc)
 	tpr := textproto.NewReader(c.Reader())
 	line, err := tpr.ReadLine()
 	if err != nil {
-		return c, err
+		return nil, err
 	}
 
 	_, code, _, ok := parseStartLine(line)
 	if ok && code == "200" {
-		tpr.ReadMIMEHeader()
+		if _, err = tpr.ReadMIMEHeader(); err != nil {
-		return c, err
+			return nil, err
+		}
+		return c, nil
 	}
 
 	switch code {

proxy/http/client_test.go

@@ -0,0 +1,126 @@
+package http
+
+import (
+	"bufio"
+	"errors"
+	"io"
+	"net"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/nadoo/glider/proxy"
+)
+
+type connectTestConn struct {
+	net.Conn
+	closed     atomic.Bool
+	writeError error
+}
+
+func (c *connectTestConn) Close() error {
+	c.closed.Store(true)
+	return c.Conn.Close()
+}
+
+func (c *connectTestConn) Write(p []byte) (int, error) {
+	if c.writeError != nil {
+		return 0, c.writeError
+	}
+	return c.Conn.Write(p)
+}
+
+type connectTestDialer struct{ conn net.Conn }
+
+func (d connectTestDialer) Addr() string                          { return "test-proxy" }
+func (d connectTestDialer) Dial(string, string) (net.Conn, error) { return d.conn, nil }
+func (d connectTestDialer) DialUDP(string, string) (net.PacketConn, error) {
+	return nil, proxy.ErrNotSupported
+}
+
+func TestDialClosesFailedCONNECT(t *testing.T) {
+	for _, tt := range []struct {
+		name       string
+		response   string
+		writeError error
+	}{
+		{name: "request write", writeError: errors.New("write failed")},
+		{name: "response EOF"},
+		{name: "invalid status", response: "invalid\r\n\r\n"},
+		{name: "rejected status", response: "HTTP/1.1 407 Proxy Authentication Required\r\n\r\n"},
+		{name: "invalid headers", response: "HTTP/1.1 200 OK\r\ninvalid\r\n\r\n"},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			client, server := net.Pipe()
+			t.Cleanup(func() { client.Close(); server.Close() })
+			client.SetDeadline(time.Now().Add(3 * time.Second))
+			tracked := &connectTestConn{Conn: client, writeError: tt.writeError}
+			if tt.writeError == nil {
+				go func() {
+					r := bufio.NewReader(server)
+					for {
+						line, err := r.ReadString('\n')
+						if err != nil {
+							return
+						}
+						if line == "\r\n" {
+							break
+						}
+					}
+					if tt.response == "" {
+						server.Close()
+						return
+					}
+					io.WriteString(server, tt.response)
+				}()
+			}
+			h := &HTTP{addr: "proxy.example:8080", dialer: connectTestDialer{tracked}}
+			conn, err := h.Dial("tcp", "target.example:443")
+			if err == nil {
+				t.Error("Dial succeeded for a failed CONNECT response")
+			}
+			if conn != nil {
+				t.Error("Dial returned a connection after a failed CONNECT")
+			}
+			if !tracked.closed.Load() {
+				t.Error("failed CONNECT left its connection open")
+			}
+		})
+	}
+}
+
+func TestDialPreservesBufferedTunnelData(t *testing.T) {
+	client, server := net.Pipe()
+	t.Cleanup(func() { client.Close(); server.Close() })
+	client.SetDeadline(time.Now().Add(3 * time.Second))
+	tracked := &connectTestConn{Conn: client}
+	go func() {
+		r := bufio.NewReader(server)
+		for {
+			line, err := r.ReadString('\n')
+			if err != nil {
+				return
+			}
+			if line == "\r\n" {
+				break
+			}
+		}
+		io.WriteString(server, "HTTP/1.1 200 Connection established\r\nProxy-Agent: test\r\n\r\nhello")
+	}()
+	h := &HTTP{addr: "proxy.example:8080", dialer: connectTestDialer{tracked}}
+	conn, err := h.Dial("tcp", "target.example:443")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+	if tracked.closed.Load() {
+		t.Fatal("successful tunnel was closed")
+	}
+	data := make([]byte, 5)
+	if _, err := io.ReadFull(conn, data); err != nil {
+		t.Fatal(err)
+	}
+	if string(data) != "hello" {
+		t.Fatalf("tunnel data = %q", data)
+	}
+}