From 5dfe7f8561da277f5dafa3b7e086700b7d824fec Mon Sep 17 00:00:00 2001
From: Dreamacro <8615343+Dreamacro@users.noreply.github.com>
Date: Thu, 24 Dec 2020 14:54:48 +0800
Subject: [PATCH] Fix: handle keep alive on http connect proxy

---
 proxy/http/server.go | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/proxy/http/server.go b/proxy/http/server.go
index 8eecf1e7..7cf2b1c1 100644
--- a/proxy/http/server.go
+++ b/proxy/http/server.go
@@ -72,21 +72,29 @@ func canActivate(loginStr string, authenticator auth.Authenticator, cache *cache
 
 func HandleConn(conn net.Conn, cache *cache.Cache) {
 	br := bufio.NewReader(conn)
+
+keepAlive:
 	request, err := http.ReadRequest(br)
 	if err != nil || request.URL.Host == "" {
 		conn.Close()
 		return
 	}
 
+	keepAlive := strings.TrimSpace(strings.ToLower(request.Header.Get("Proxy-Connection"))) == "keep-alive"
 	authenticator := authStore.Authenticator()
 	if authenticator != nil {
 		if authStrings := strings.Split(request.Header.Get("Proxy-Authorization"), " "); len(authStrings) != 2 {
 			conn.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic\r\n\r\n"))
-			conn.Close()
+			if keepAlive {
+				goto keepAlive
+			}
 			return
 		} else if !canActivate(authStrings[1], authenticator, cache) {
 			conn.Write([]byte("HTTP/1.1 403 Forbidden\r\n\r\n"))
 			log.Infoln("Auth failed from %s", conn.RemoteAddr().String())
+			if keepAlive {
+				goto keepAlive
+			}
 			conn.Close()
 			return
 		}
@@ -95,6 +103,7 @@ func HandleConn(conn net.Conn, cache *cache.Cache) {
 	if request.Method == http.MethodConnect {
 		_, err := conn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n"))
 		if err != nil {
+			conn.Close()
 			return
 		}
 		tunnel.Add(adapters.NewHTTPS(request, conn))