Fix compatibility

Author: wzshiming

compatibility_read_deadline.go

@@ -26,7 +26,20 @@ func (w listenerCompatibilityReadDeadline) Accept() (net.Conn, error) {
 	if err != nil {
 		return nil, err
 	}
-	return connCompatibilityReadDeadline{c}, nil
+	return NewConnCompatibilityReadDeadline(c), nil
+}
+
+// NewConnCompatibilityReadDeadline this is a wrapper used to be compatible with
+// the net.Conn after wrapping it so that it can be hijacked properly.
+// there is no effect if the content is not manipulated.
+func NewConnCompatibilityReadDeadline(conn net.Conn) net.Conn {
+	if conn == nil {
+		return nil
+	}
+	if conn, ok := conn.(connCompatibilityReadDeadline); ok {
+		return conn
+	}
+	return connCompatibilityReadDeadline{conn}
 }
 
 type connCompatibilityReadDeadline struct {
@@ -35,7 +48,7 @@ type connCompatibilityReadDeadline struct {
 
 func (d connCompatibilityReadDeadline) SetReadDeadline(t time.Time) error {
 	if aLongTimeAgo == t {
-		t = time.Now().Add(time.Second)
+		t = time.Now().Add(1 * time.Second)
 	}
 	return d.Conn.SetReadDeadline(t)
 }

go.mod

@@ -2,6 +2,6 @@ module github.com/wzshiming/httpproxy
 
 go 1.18
 
-require golang.org/x/net v0.2.0
+require golang.org/x/net v0.5.0
 
-require golang.org/x/text v0.4.0 // indirect
+require golang.org/x/text v0.6.0 // indirect

go.sum

@@ -1,4 +1,4 @@
-golang.org/x/net v0.2.0 h1:sZfSu1wtKLGlWI4ZZayP0ck9Y73K1ynO6gqzTdBVdPU=
-golang.org/x/net v0.2.0/go.mod h1:KqCZLdyyvdV855qA2rE3GC2aiw5xGR5TEjj8smXukLY=
-golang.org/x/text v0.4.0 h1:BrVqGRd7+k1DiOgtnFvAkoQEWQvBc25ouMJM6429SFg=
-golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
+golang.org/x/net v0.5.0 h1:GyT4nK/YDHSqa1c4753ouYCDajOYKTja9Xb/OHtgvSw=
+golang.org/x/net v0.5.0/go.mod h1:DivGGAXEgPSlEBzxGzZI+ZLohi+xUj054jfeKui00ws=
+golang.org/x/text v0.6.0 h1:3XmdazWV+ubf7QgHSTWeykHOci5oeekaGJBLkrkaw4k=
+golang.org/x/text v0.6.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=

proxy.go

@@ -1,6 +1,7 @@
 package httpproxy
 
 import (
+	"bufio"
 	"context"
 	"fmt"
 	"io"
@@ -96,23 +97,22 @@ func (p *ProxyHandler) proxyConnect(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, e, http.StatusInternalServerError)
 		return
 	}
+	defer targetConn.Close()
 
-	if flusher, ok := w.(http.Flusher); ok {
-		flusher.Flush()
-	} else {
-		w.WriteHeader(http.StatusOK)
-	}
+	w.WriteHeader(http.StatusOK)
 
-	clientConn, _, err := hijacker.Hijack()
+	conn, rw, err := hijacker.Hijack()
 	if err != nil {
-		e := err.Error()
+		e := fmt.Sprintf("hijack failed: %v", err)
 		if p.Logger != nil {
 			p.Logger.Println(e)
 		}
 		http.Error(w, e, http.StatusInternalServerError)
 		return
 	}
 
+	clientConn := newBufConn(conn, rw)
+
 	var buf1, buf2 []byte
 	if p.BytesPool != nil {
 		buf1 = p.BytesPool.Get()
@@ -151,3 +151,34 @@ func (p *ProxyHandler) proxyDial(ctx context.Context, network, address string) (
 	}
 	return proxyDial(ctx, network, address)
 }
+
+func newBufConn(conn net.Conn, rw *bufio.ReadWriter) net.Conn {
+	rw.Flush()
+	if rw.Reader.Buffered() == 0 {
+		// If there's no buffered data to be read,
+		// we can just discard the bufio.ReadWriter.
+		return conn
+	}
+	return &bufConn{conn, rw.Reader}
+}
+
+// bufConn wraps a net.Conn, but reads drain the bufio.Reader first.
+type bufConn struct {
+	net.Conn
+	*bufio.Reader
+}
+
+func (c *bufConn) Read(p []byte) (int, error) {
+	if c.Reader == nil {
+		return c.Conn.Read(p)
+	}
+	n := c.Reader.Buffered()
+	if n == 0 {
+		c.Reader = nil
+		return c.Conn.Read(p)
+	}
+	if n < len(p) {
+		p = p[:n]
+	}
+	return c.Reader.Read(p)
+}

tunnel.go

@@ -5,7 +5,6 @@ import (
 	"io"
 )
 
-
 // tunnel create tunnels for two io.ReadWriteCloser
 func tunnel(ctx context.Context, c1, c2 io.ReadWriteCloser, buf1, buf2 []byte) error {
 	ctx, cancel := context.WithCancel(ctx)