vulcand websocket diff

notes.md

• The WebsocketUpgrader has to be put so far up the chain within frontend.go, otherwise you have a non-hijackable ResponseWriter.
• The RoundRobin class is used to determine which backend to net.Dial against.
• I could not yet figure out how to shove the RoundRobin rebalencer into this. The API for it does not have a NextServer().
• The whole Bidir stuff on the bottom of the upgrader is by someone else. (vulcand/vulcand#78 (comment))
• The bidir code causes some errors to be logged when the connection is closed.
• The code isn't using an http.Handler errorHandler like it should.
• NO TLS???

proxy-frontend.diff

diff --git a/proxy/frontend.go b/proxy/frontend.go
index 69565a7..48a35b3 100644
--- a/proxy/frontend.go
+++ b/proxy/frontend.go
@@ -177,8 +177,10 @@ func (f *frontend) rebuild() error {
 		return err
 	}
 
+	upg := newWebsocketUpgrader(rr, str)
+
 	// Add the frontend to the router
-	if err := f.mux.router.Handle(f.frontend.Route, str); err != nil {
+	if err := f.mux.router.Handle(f.frontend.Route, upg); err != nil {
 		return err
 	}
 

proxy-upgrader.diff

diff --git a/proxy/upgrader.go b/proxy/upgrader.go
new file mode 100644
index 0000000..f3d79a9
--- /dev/null
+++ b/proxy/upgrader.go
@@ -0,0 +1,101 @@
+package proxy
+
+import (
+	"bufio"
+	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/log"
+	"github.com/mailgun/vulcand/Godeps/_workspace/src/github.com/mailgun/oxy/roundrobin"
+	"io"
+	"net"
+	"net/http"
+	"strings"
+)
+
+// WebsocketUpgrader is an HTTP middleware that detects for websocket upgrade requests
+// and establishes an HTTP connection via a chosen backend server
+type WebsocketUpgrader struct {
+	next http.Handler
+	rr   *roundrobin.RoundRobin
+}
+
+// creat the upgrader via a roundrobin and the expected next handler (if not websocket)
+func newWebsocketUpgrader(rr *roundrobin.RoundRobin, next http.Handler) *WebsocketUpgrader {
+	return &WebsocketUpgrader{
+		next: next,
+		rr:   rr,
+	}
+}
+
+// ServeHTTP waits for a websocket upgrade request and creates a TCP connection between
+// the backend server and the frontend
+func (u *WebsocketUpgrader) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	if strings.Join(req.Header["Upgrade"], "") == "websocket" {
+
+		url, err := u.rr.NextServer()
+		if err != nil {
+			log.Errorf("Can't round robin")
+			return
+		}
+
+		hj, ok := w.(http.Hijacker)
+
+		if !ok {
+			log.Errorf("Webserver doesn't support hijacking")
+		}
+
+		conn, bufrw, err := hj.Hijack()
+		defer conn.Close()
+
+		conn2, err := net.Dial("tcp", url.Host)
+		if err != nil {
+			log.Errorf("Couldn't connect to backend server: %v", err)
+			return
+		}
+		defer conn2.Close()
+
+		err = req.Write(conn2)
+		if err != nil {
+			log.Errorf("writing request to backend server failed: %v", err)
+			return
+		}
+
+		copyBidir(conn, bufrw, conn2, bufio.NewReadWriter(bufio.NewReader(conn2), bufio.NewWriter(conn2)))
+
+		return
+	}
+
+	u.next.ServeHTTP(w, req)
+}
+
+func copyBetween(dest *bufio.ReadWriter, src *bufio.ReadWriter) {
+	buf := make([]byte, 40*1024)
+	for {
+		n, err := src.Read(buf)
+		if err != nil && err != io.EOF {
+			log.Errorf("Read failed: %v", err)
+			return
+		}
+		if n == 0 {
+			return
+		}
+		dest.Write(buf[0:n])
+		dest.Flush()
+	}
+}
+
+func copyBidir(conn1 io.ReadWriteCloser, rw1 *bufio.ReadWriter, conn2 io.ReadWriteCloser, rw2 *bufio.ReadWriter) {
+	finished := make(chan bool)
+
+	go func() {
+		copyBetween(rw2, rw1)
+		conn2.Close()
+		finished <- true
+	}()
+	go func() {
+		copyBetween(rw1, rw2)
+		conn1.Close()
+		finished <- true
+	}()
+
+	<-finished
+	<-finished
+}