diff --git a/proxy/http/client.go b/proxy/http/client.go index 6562ccc..cb3d185 100644 --- a/proxy/http/client.go +++ b/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() - return c, err + if _, err = tpr.ReadMIMEHeader(); err != nil { + return nil, err + } + return c, nil } switch code { diff --git a/proxy/http/client_test.go b/proxy/http/client_test.go new file mode 100644 index 0000000..2100a3c --- /dev/null +++ b/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) + } +}