Merge pull request #123 from SenseUnit/optimism

handler: postpone conn acknowledgement if client uses fast start

Author: Snawoot

handler/handler.go

@@ -103,9 +103,6 @@ func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request, u
 		}
 		defer localconn.Close()
 
-		// Inform client connection is built
-		fmt.Fprintf(localconn, "HTTP/%d.%d 200 OK\r\n\r\n", req.ProtoMajor, req.ProtoMinor)
-
 		if buffered := rw.Reader.Buffered(); buffered > 0 {
 			s.logger.Debug("saving %d bytes buffered in bufio.ReadWriter", buffered)
 			s.forward(
@@ -114,11 +111,20 @@ func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request, u
 				wrapH1ReqBody(io.NopCloser(io.LimitReader(rw.Reader, int64(buffered)))),
 				wrapH1RespWriter(conn),
 			)
+			s.forward(
+				req.Context(),
+				username,
+				wrapPendingWrite(
+					[]byte(fmt.Sprintf("HTTP/%d.%d 200 OK\r\n\r\n", req.ProtoMajor, req.ProtoMinor)),
+					localconn,
+				),
+				conn,
+			)
 		} else {
 			s.logger.Debug("not rescuing remaining data in bufio.ReadWriter")
+			fmt.Fprintf(localconn, "HTTP/%d.%d 200 OK\r\n\r\n", req.ProtoMajor, req.ProtoMinor)
+			s.forward(req.Context(), username, localconn, conn)
 		}
-
-		s.forward(req.Context(), username, localconn, conn)
 	} else if req.ProtoMajor == 2 {
 		wr.Header()["Date"] = nil
 		wr.WriteHeader(http.StatusOK)

handler/proxy.go

@@ -62,3 +62,28 @@ func flush(flusher interface{}) bool {
 	f.Flush()
 	return true
 }
+
+func wrapPendingWrite(data []byte, c net.Conn) *pendingWriteConn {
+	return &pendingWriteConn{
+		data: data,
+		Conn: c,
+	}
+}
+
+type pendingWriteConn struct {
+	net.Conn
+	data []byte
+	done bool
+}
+
+func (p *pendingWriteConn) Write(b []byte) (n int, err error) {
+	if !p.done {
+		buf := append(append(make([]byte, 0, len(p.data)+len(b)), p.data...), b...)
+		n, err := p.Conn.Write(buf)
+		n = max(0, n - len(p.data))
+		p.done = true
+		p.data = nil
+		return n, err
+	}
+	return p.Conn.Write(b)
+}