correctly handle keep-alive in http server

proxy/http/server.go

@@ -69,15 +69,17 @@ func parseHost(rawHost string, defaultPort v2net.Port) (v2net.Destination, error
 }
 
 func (s *Server) Process(ctx context.Context, network v2net.Network, conn internet.Connection, dispatcher dispatcher.Interface) error {
+Start:
 	conn.SetReadDeadline(time.Now().Add(time.Second * 8))
 	reader := bufio.NewReaderSize(conn, 2048)
 
 	request, err := http.ReadRequest(reader)
 	if err != nil {
+		trace := newError("failed to read http request").Base(err)
 		if errors.Cause(err) != io.EOF {
-			log.Trace(newError("failed to read http request").Base(err).AtWarning())
+			trace.AtWarning()
 		}
-		return err
+		return trace
 	}
 	log.Trace(newError("request to Method [", request.Method, "] Host [", request.Host, "] with URL [", request.URL, "]"))
 	conn.SetReadDeadline(time.Time{})
@@ -98,9 +100,19 @@ func (s *Server) Process(ctx context.Context, network v2net.Network, conn intern
 
 	if strings.ToUpper(request.Method) == "CONNECT" {
 		return s.handleConnect(ctx, request, reader, conn, dest, dispatcher)
-	} else {
-		return s.handlePlainHTTP(ctx, request, reader, conn, dest, dispatcher)
 	}
+
+	keepAlive := (strings.TrimSpace(strings.ToLower(request.Header.Get("Proxy-Connection"))) == "keep-alive")
+
+	err = s.handlePlainHTTP(ctx, request, reader, conn, dest, dispatcher)
+	if err == errWaitAnother {
+		if keepAlive {
+			goto Start
+		}
+		err = nil
+	}
+
+	return err
 }
 
 func (s *Server) handleConnect(ctx context.Context, request *http.Request, reader io.Reader, writer io.Writer, dest v2net.Destination, dispatcher dispatcher.Interface) error {
@@ -115,6 +127,7 @@ func (s *Server) handleConnect(ctx context.Context, request *http.Request, reade
 		ContentLength: -1, // Don't send Content-Length in CONNECT.
 		Close:         false,
 	}
+	response.Header.Set("Proxy-Connection", "close")
 	if err := response.Write(writer); err != nil {
 		return newError("failed to write back OK response").Base(err)
 	}
@@ -159,54 +172,51 @@ func (s *Server) handleConnect(ctx context.Context, request *http.Request, reade
 }
 
 // @VisibleForTesting
-func StripHopByHopHeaders(request *http.Request) {
+func StripHopByHopHeaders(header http.Header) {
 	// Strip hop-by-hop header basaed on RFC:
 	// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.5.1
 	// https://www.mnot.net/blog/2011/07/11/what_proxies_must_do
 
-	request.Header.Del("Proxy-Connection")
-	request.Header.Del("Proxy-Authenticate")
-	request.Header.Del("Proxy-Authorization")
-	request.Header.Del("TE")
-	request.Header.Del("Trailers")
-	request.Header.Del("Transfer-Encoding")
-	request.Header.Del("Upgrade")
-
-	// TODO: support keep-alive
-	connections := request.Header.Get("Connection")
-	request.Header.Set("Connection", "close")
+	header.Del("Proxy-Connection")
+	header.Del("Proxy-Authenticate")
+	header.Del("Proxy-Authorization")
+	header.Del("TE")
+	header.Del("Trailers")
+	header.Del("Transfer-Encoding")
+	header.Del("Upgrade")
+
+	connections := header.Get("Connection")
 	if len(connections) == 0 {
 		return
 	}
 	for _, h := range strings.Split(connections, ",") {
-		request.Header.Del(strings.TrimSpace(h))
+		header.Del(strings.TrimSpace(h))
 	}
+	header.Del("Connection")
 }
 
-func generateResponse(statusCode int, status string) *http.Response {
-	hdr := http.Header(make(map[string][]string))
-	hdr.Set("Connection", "close")
-	return &http.Response{
-		Status:        status,
-		StatusCode:    statusCode,
-		Proto:         "HTTP/1.1",
-		ProtoMajor:    1,
-		ProtoMinor:    1,
-		Header:        hdr,
-		Body:          nil,
-		ContentLength: 0,
-		Close:         true,
-	}
-}
+var errWaitAnother = newError("keep alive")
 
 func (s *Server) handlePlainHTTP(ctx context.Context, request *http.Request, reader io.Reader, writer io.Writer, dest v2net.Destination, dispatcher dispatcher.Interface) error {
 	if len(request.URL.Host) <= 0 {
-		response := generateResponse(400, "Bad Request")
+		response := &http.Response{
+			Status:        "Bad Request",
+			StatusCode:    400,
+			Proto:         "HTTP/1.1",
+			ProtoMajor:    1,
+			ProtoMinor:    1,
+			Header:        http.Header(make(map[string][]string)),
+			Body:          nil,
+			ContentLength: 0,
+			Close:         true,
+		}
+		response.Header.Set("Proxy-Connection", "close")
+		response.Header.Set("Connection", "close")
 		return response.Write(writer)
 	}
 
 	request.Host = request.URL.Host
-	StripHopByHopHeaders(request)
+	StripHopByHopHeaders(request.Header)
 
 	ray, err := dispatcher.Dispatch(ctx, dest)
 	if err != nil {
@@ -217,6 +227,8 @@ func (s *Server) handlePlainHTTP(ctx context.Context, request *http.Request, rea
 	defer input.Close()
 
 	requestDone := signal.ExecuteAsync(func() error {
+		request.Header.Set("Connection", "close")
+
 		requestWriter := buf.NewBufferedWriter(buf.ToBytesWriter(ray.InboundInput()))
 		err := request.Write(requestWriter)
 		if err != nil {
@@ -231,9 +243,26 @@ func (s *Server) handlePlainHTTP(ctx context.Context, request *http.Request, rea
 	responseDone := signal.ExecuteAsync(func() error {
 		responseReader := bufio.NewReader(buf.ToBytesReader(ray.InboundOutput()))
 		response, err := http.ReadResponse(responseReader, request)
-		if err != nil {
+		if err == nil {
+			StripHopByHopHeaders(response.Header)
+			response.Header.Set("Proxy-Connection", "keep-alive")
+			response.Header.Set("Connection", "Keep-Alive")
+			response.Header.Set("Keep-Alive", "timeout=4")
+		} else {
 			log.Trace(newError("failed to read response").Base(err).AtWarning())
-			response = generateResponse(503, "Service Unavailable")
+			response = &http.Response{
+				Status:        "Service Unavailable",
+				StatusCode:    503,
+				Proto:         "HTTP/1.1",
+				ProtoMajor:    1,
+				ProtoMinor:    1,
+				Header:        http.Header(make(map[string][]string)),
+				Body:          nil,
+				ContentLength: 0,
+				Close:         true,
+			}
+			response.Header.Set("Connection", "close")
+			response.Header.Set("Proxy-Connection", "close")
 		}
 		responseWriter := buf.NewBufferedWriter(writer)
 		if err := response.Write(responseWriter); err != nil {
@@ -252,7 +281,7 @@ func (s *Server) handlePlainHTTP(ctx context.Context, request *http.Request, rea
 		return newError("connection ends").Base(err)
 	}
 
-	return nil
+	return errWaitAnother
 }
 
 func init() {

proxy/http/server_test.go

@@ -38,7 +38,7 @@ Accept-Language: de,en;q=0.7,en-us;q=0.3
 	assert.String(req.Header.Get("Proxy-Connection")).Equals("keep-alive")
 	assert.String(req.Header.Get("Proxy-Authenticate")).Equals("abc")
 
-	StripHopByHopHeaders(req)
+	StripHopByHopHeaders(req.Header)
 	assert.String(req.Header.Get("Connection")).Equals("close")
 	assert.String(req.Header.Get("Foo")).Equals("")
 	assert.String(req.Header.Get("Bar")).Equals("")