From d3e52bb457e39ceace9657824e1451b0206b5b2d Mon Sep 17 00:00:00 2001
From: "rui.zheng" <rui.zheng@tcl.com>
Date: Sun, 25 Sep 2016 18:31:54 +0800
Subject: [PATCH] http2 support with http1.x compatible

---
 conn.go  | 10 +++++--
 http2.go | 90 +++++++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 76 insertions(+), 24 deletions(-)

diff --git a/conn.go b/conn.go
index 281792f..c49f5be 100644
--- a/conn.go
+++ b/conn.go
@@ -89,9 +89,13 @@ func listenAndServe(arg Args) error {
 
 func listenAndServeHttp2(arg Args) error {
 	srv := http.Server{
-		Addr:      arg.Addr,
-		Handler:   http.HandlerFunc(handlerHttp2Request),
-		TLSConfig: &tls.Config{Certificates: []tls.Certificate{arg.Cert}},
+		Addr:    arg.Addr,
+		Handler: http.HandlerFunc(handlerHttp2Request),
+		TLSConfig: &tls.Config{
+			Certificates: []tls.Certificate{arg.Cert},
+			//MinVersion:               tls.VersionTLS12,
+			//PreferServerCipherSuites: true,
+		},
 	}
 	http2.ConfigureServer(&srv, nil)
 	return srv.ListenAndServeTLS("", "")
diff --git a/http2.go b/http2.go
index 63cefff..994063c 100644
--- a/http2.go
+++ b/http2.go
@@ -1,8 +1,10 @@
 package main
 
 import (
+	"bufio"
 	"github.com/golang/glog"
 	"golang.org/x/net/http2"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httputil"
@@ -27,6 +29,8 @@ func handlerHttp2Request(w http.ResponseWriter, req *http.Request) {
 	var c net.Conn
 	var err error
 
+	fw := flushWriter{w}
+
 	c, err = Connect(req.Host)
 	if err != nil {
 		glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, req.Host, err)
@@ -34,40 +38,84 @@ func handlerHttp2Request(w http.ResponseWriter, req *http.Request) {
 			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
 		glog.V(LDEBUG).Infof("[http2] %s <- %s\n%s", req.RemoteAddr, req.Host, string(b))
 		//w.WriteHeader(http.StatusServiceUnavailable)
-		w.Write(b)
+		fw.Write(b)
 		return
 	}
 	defer c.Close()
 
+	rChan := make(chan error, 1)
+	wChan := make(chan error, 1)
+
 	if req.Method == http.MethodConnect {
-		b := []byte("HTTP/2.0 200 Connection established\r\n" +
-			"Proxy-Agent: gost/" + Version + "\r\n\r\n")
-		glog.V(LDEBUG).Infof("[http2] %s <- %s\n%s", req.RemoteAddr, req.Host, string(b))
-		w.Write(b)
-		if f, ok := w.(http.Flusher); ok {
-			f.Flush()
+		w.Header().Set("Proxy-Agent", "gost/"+Version)
+		w.WriteHeader(http.StatusOK)
+
+		if fw, ok := w.(http.Flusher); ok {
+			fw.Flush()
 		}
+
+		// compatible with HTTP 1.x
+		if hj, ok := w.(http.Hijacker); ok && req.ProtoMajor == 1 {
+			conn, _, err := hj.Hijack()
+			if err != nil {
+				glog.V(LWARNING).Infoln(err)
+				return
+			}
+			defer conn.Close()
+
+			go Pipe(conn, c, rChan)
+			go Pipe(c, conn, wChan)
+		} else {
+			go Pipe(req.Body, c, rChan)
+			go Pipe(c, fw, wChan)
+		}
+
+		select {
+		case err := <-rChan:
+			glog.V(LWARNING).Infoln("r exit", err)
+		case err := <-wChan:
+			glog.V(LWARNING).Infoln("w exit", err)
+		}
+
 	} else {
 		req.Header.Set("Connection", "Keep-Alive")
 		if err = req.Write(c); err != nil {
 			glog.V(LWARNING).Infof("[http2] %s -> %s : %s", req.RemoteAddr, req.Host, err)
 			return
 		}
+
+		resp, err := http.ReadResponse(bufio.NewReader(c), req)
+		if err != nil {
+			glog.V(LWARNING).Infoln(err)
+			return
+		}
+		defer resp.Body.Close()
+
+		for k, v := range resp.Header {
+			for _, vv := range v {
+				w.Header().Add(k, vv)
+			}
+		}
+
+		w.WriteHeader(resp.StatusCode)
+		if _, err := io.Copy(fw, resp.Body); err != nil {
+			glog.V(LWARNING).Infoln(err)
+		}
 	}
 
-	glog.V(LINFO).Infof("[http2] %s <-> %s", req.RemoteAddr, req.Host)
+	//glog.V(LINFO).Infof("[http2] %s <-> %s", req.RemoteAddr, req.Host)
 
-	rChan := make(chan error, 1)
-	wChan := make(chan error, 1)
-	go Pipe(c, w, wChan)
-	go Pipe(req.Body, c, rChan)
-
-	select {
-	case err = <-wChan:
-		glog.V(LWARNING).Infoln("w exit", err)
-	case err = <-rChan:
-		glog.V(LWARNING).Infoln("r exit", err)
-	}
-
-	glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, req.Host)
+	//glog.V(LINFO).Infof("[http2] %s >-< %s", req.RemoteAddr, req.Host)
+}
+
+type flushWriter struct {
+	w io.Writer
+}
+
+func (fw flushWriter) Write(p []byte) (n int, err error) {
+	n, err = fw.w.Write(p)
+	if f, ok := fw.w.(http.Flusher); ok {
+		f.Flush()
+	}
+	return
 }