- 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???
Last active
September 22, 2015 16:32
-
-
Save sheenobu/5aad1b1c07e9ff52eefe to your computer and use it in GitHub Desktop.
vulcand websocket diff
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
} | |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
+} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment