Merge pull request #96 from acmacalister/feature/ws-proxy

add supported for ws in reverse proxy

config/setup/rewrite_test.go

@@ -4,8 +4,9 @@ import (
 	"testing"
 
 	"fmt"
-	"github.com/mholt/caddy/middleware/rewrite"
 	"regexp"
+
+	"github.com/mholt/caddy/middleware/rewrite"
 )
 
 func TestRewrite(t *testing.T) {

middleware/proxy/reverseproxy.go

@@ -147,14 +147,43 @@ func (p *ReverseProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request, extr
 	}
 	defer res.Body.Close()
 
-	for _, h := range hopHeaders {
-		res.Header.Del(h)
-	}
+	if res.StatusCode == http.StatusSwitchingProtocols && res.Header.Get("Upgrade") == "websocket" {
+		hj, ok := rw.(http.Hijacker)
+		if !ok {
+			return nil
+		}
+
+		conn, _, err := hj.Hijack()
+		if err != nil {
+			return err
+		}
+
+		backendConn, err := net.Dial("tcp", outreq.Host)
+		if err != nil {
+			conn.Close()
+			return err
+		}
 
-	copyHeader(rw.Header(), res.Header)
+		outreq.Write(backendConn)
+
+		go func() {
+			io.Copy(backendConn, conn) // write tcp stream to backend.
+			backendConn.Close()
+		}()
+
+		io.Copy(conn, backendConn) // read tcp stream from backend.
+		conn.Close()
+	} else {
+		for _, h := range hopHeaders {
+			res.Header.Del(h)
+		}
+
+		copyHeader(rw.Header(), res.Header)
+
+		rw.WriteHeader(res.StatusCode)
+		p.copyResponse(rw, res.Body)
+	}
 
-	rw.WriteHeader(res.StatusCode)
-	p.copyResponse(rw, res.Body)
 	return nil
 }
 

middleware/proxy/upstream.go

@@ -12,7 +12,10 @@ import (
 	"github.com/mholt/caddy/config/parse"
 )
 
-var supportedPolicies map[string]func() Policy = make(map[string]func() Policy)
+var (
+	supportedPolicies map[string]func() Policy = make(map[string]func() Policy)
+	proxyHeaders      http.Header              = make(http.Header)
+)
 
 type staticUpstream struct {
 	from   string
@@ -40,7 +43,7 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {
 			FailTimeout: 10 * time.Second,
 			MaxFails:    1,
 		}
-		var proxyHeaders http.Header
+
 		if !c.Args(&upstream.from) {
 			return upstreams, c.ArgErr()
 		}
@@ -97,10 +100,10 @@ func NewStaticUpstreams(c parse.Dispenser) ([]Upstream, error) {
 				if !c.Args(&header, &value) {
 					return upstreams, c.ArgErr()
 				}
-				if proxyHeaders == nil {
-					proxyHeaders = make(map[string][]string)
-				}
 				proxyHeaders.Add(header, value)
+			case "websocket":
+				proxyHeaders.Add("Connection", "{>Connection}")
+				proxyHeaders.Add("Upgrade", "{>Upgrade}")
 			}
 		}