feat: correct process routines

http.go

@@ -10,6 +10,8 @@ import (
 	"net"
 	"net/http"
 	"strings"
+
+	"github.com/sourcegraph/conc"
 )
 
 const proxyAuthHeaderKey = "Proxy-Authorization"
@@ -91,17 +93,19 @@ func (s *HTTPServer) handle(req *http.Request) (peer net.Conn, err error) {
 	return
 }
 
-func (s *HTTPServer) serve(conn net.Conn) error {
+func (s *HTTPServer) serve(conn net.Conn) {
 	var rd = bufio.NewReader(conn)
 	req, err := http.ReadRequest(rd)
 	if err != nil {
-		return fmt.Errorf("read request failed: %w", err)
+		log.Printf("read request failed: %s\n", err)
+		return
 	}
 
 	code, err := s.authenticate(req)
 	if err != nil {
 		_ = responseWith(req, code).Write(conn)
-		return err
+		log.Println(err)
+		return
 	}
 
 	var peer net.Conn
@@ -112,23 +116,29 @@ func (s *HTTPServer) serve(conn net.Conn) error {
 		peer, err = s.handle(req)
 	default:
 		_ = responseWith(req, http.StatusMethodNotAllowed).Write(conn)
-		return fmt.Errorf("unsupported protocol: %s", req.Method)
+		log.Printf("unsupported protocol: %s\n", req.Method)
+		return
 	}
 	if err != nil {
-		return fmt.Errorf("dial proxy failed: %w", err)
+		log.Printf("dial proxy failed: %s\n", err)
+		return
 	}
 	if peer == nil {
-		return fmt.Errorf("dial proxy failed: peer nil")
+		log.Println("dial proxy failed: peer nil")
+		return
 	}
-	defer peer.Close()
-
 	go func() {
-		_, _ = io.Copy(conn, peer)
+		wg := conc.NewWaitGroup()
+		wg.Go(func() {
+			_, err = io.Copy(conn, peer)
+		})
+		wg.Go(func() {
+			_, err = io.Copy(peer, conn)
+		})
+		wg.Wait()
+		_ = peer.Close()
+		_ = conn.Close()
 	}()
-
-	_, err = io.Copy(peer, conn)
-
-	return err
 }
 
 // ListenAndServe is used to create a listener and serve on it
@@ -137,20 +147,16 @@ func (s *HTTPServer) ListenAndServe(network, addr string) error {
 	if err != nil {
 		return fmt.Errorf("listen tcp failed: %w", err)
 	}
-	defer server.Close()
-	var conn net.Conn
+	defer func(server net.Listener) {
+		_ = server.Close()
+	}(server)
 	for {
-		conn, err = server.Accept()
+		conn, err := server.Accept()
 		if err != nil {
 			return fmt.Errorf("accept request failed: %w", err)
 		}
 		go func(conn net.Conn) {
-			err = s.serve(conn)
-			if err != nil {
-				log.Println(err)
-			}
-			_ = conn.Close()
-			conn = nil
+			s.serve(conn)
 		}(conn)
 	}
 }