From a4a691b7b8930633dee879044d7817039356f605 Mon Sep 17 00:00:00 2001
From: Barna Csorogi <barnacs@justletit.be>
Date: Sat, 26 Mar 2016 18:52:06 +0100
Subject: [PATCH] add http2 support

Add http2 using the go-1.6 net/http built-in support.

net/http's http2 doesn't support hijacking, so instead of hijacking the
inital CONNECT tcp connection, tunnel proxied requests/responses over
the reader and writer streams of the CONNECT request.
---
 proxy/mitmconn.go | 74 +++++++++++++++++++++++++++++++++++++++++++++++
 proxy/proxy.go    | 19 +++++-------
 2 files changed, 81 insertions(+), 12 deletions(-)
 create mode 100644 proxy/mitmconn.go

diff --git a/proxy/mitmconn.go b/proxy/mitmconn.go
new file mode 100644
index 0000000..443b606
--- /dev/null
+++ b/proxy/mitmconn.go
@@ -0,0 +1,74 @@
+package proxy
+
+import (
+	"io"
+	"net"
+	"net/http"
+	"time"
+)
+
+type mitmConn struct {
+	w          FlushWriter
+	r          io.Reader
+	remoteAddr addr
+	closed     chan struct{}
+}
+
+type FlushWriter interface {
+	io.Writer
+	http.Flusher
+}
+
+func newMitmConn(w FlushWriter, r io.Reader, remoteAddr string) *mitmConn {
+	return &mitmConn{
+		w:          w,
+		r:          r,
+		remoteAddr: addr(remoteAddr),
+		closed:     make(chan struct{}),
+	}
+}
+
+func (c *mitmConn) Read(b []byte) (int, error) {
+	return c.r.Read(b)
+}
+
+func (c *mitmConn) Write(b []byte) (int, error) {
+	n, err := c.w.Write(b)
+	c.w.Flush()
+	return n, err
+}
+
+func (c *mitmConn) Close() error {
+	close(c.closed)
+	return nil
+}
+
+func (c *mitmConn) RemoteAddr() net.Addr {
+	return c.remoteAddr
+}
+
+func (c *mitmConn) LocalAddr() net.Addr {
+	panic("not implemented")
+}
+
+func (c *mitmConn) SetDeadline(t time.Time) error {
+	panic("not implemented")
+}
+
+func (c *mitmConn) SetReadDeadline(t time.Time) error {
+	panic("not implemented")
+}
+
+func (c *mitmConn) SetWriteDeadline(t time.Time) error {
+	panic("not implemented")
+}
+
+type addr string
+
+func (a addr) String() string {
+	return string(a)
+}
+
+func (a addr) Network() string {
+	return "tcp"
+}
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 1609d25..dc3232c 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -53,7 +53,7 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func (p *Proxy) handle(w http.ResponseWriter, r *http.Request) error {
 	if r.Method == "CONNECT" {
-		return p.handleConnect(w, r.Host)
+		return p.handleConnect(w, r)
 	}
 	resp, err := forward(r)
 	if err != nil {
@@ -88,24 +88,19 @@ func (p *Proxy) proxyResponse(w *ResponseWriter, r *ResponseReader) error {
 	return nil
 }
 
-func (p *Proxy) handleConnect(w http.ResponseWriter, host string) error {
+func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) error {
 	if p.ml == nil {
 		return fmt.Errorf("CONNECT received but mitm is not enabled")
 	}
-	h, ok := w.(http.Hijacker)
-	if !ok {
-		return fmt.Errorf("connection cannot be hijacked")
-	}
 	w.WriteHeader(http.StatusOK)
-	conn, _, err := h.Hijack()
+	fw := w.(FlushWriter)
+	fw.Flush()
+	conn := newMitmConn(fw, r.Body, r.RemoteAddr)
+	sconn, err := p.ml.Serve(conn, r.Host)
 	if err != nil {
 		return err
 	}
-	sconn, err := p.ml.Serve(conn, host)
-	if err != nil {
-		conn.Close()
-		return err
-	}
 	sconn.Close() // TODO: reuse this connection for https requests
+	<-conn.closed
 	return nil
 }