Add websocket.Upgrader option.

sockjs/options.go

@@ -7,6 +7,8 @@ import (
 	"net/http"
 	"sync"
 	"time"
+
+	"github.com/gorilla/websocket"
 )
 
 var (
@@ -38,6 +40,9 @@ type Options struct {
 	Websocket bool
 	// This option can be used to enable raw websockets support by the server. By default raw websockets are disabled.
 	RawWebsocket bool
+	// Provide a custom Upgrader for Websocket connections to enable features like compression.
+	// See https://godoc.org/github.com/gorilla/websocket#Upgrader for more details.
+	WebsocketUpgrader *websocket.Upgrader
 	// In order to keep proxies and load balancers from closing long running http requests we need to pretend that the connection is active
 	// and send a heartbeat packet once in a while. This setting controls how often this is done.
 	// By default a heartbeat packet is sent every 25 seconds.
@@ -54,13 +59,14 @@ type Options struct {
 
 // DefaultOptions is a convenient set of options to be used for sockjs
 var DefaultOptions = Options{
-	Websocket:       true,
-	RawWebsocket:    false,
-	JSessionID:      nil,
-	SockJSURL:       "//cdnjs.cloudflare.com/ajax/libs/sockjs-client/0.3.4/sockjs.min.js",
-	HeartbeatDelay:  25 * time.Second,
-	DisconnectDelay: 5 * time.Second,
-	ResponseLimit:   128 * 1024,
+	Websocket:         true,
+	RawWebsocket:      false,
+	JSessionID:        nil,
+	SockJSURL:         "//cdnjs.cloudflare.com/ajax/libs/sockjs-client/0.3.4/sockjs.min.js",
+	HeartbeatDelay:    25 * time.Second,
+	DisconnectDelay:   5 * time.Second,
+	ResponseLimit:     128 * 1024,
+	WebsocketUpgrader: nil,
 }
 
 type info struct {

sockjs/rawwebsocket.go

@@ -8,7 +8,15 @@ import (
 )
 
 func (h *handler) rawWebsocket(rw http.ResponseWriter, req *http.Request) {
-	conn, err := websocket.Upgrade(rw, req, nil, WebSocketReadBufSize, WebSocketWriteBufSize)
+	var conn *websocket.Conn
+	var err error
+	if h.options.WebsocketUpgrader != nil {
+		conn, err = h.options.WebsocketUpgrader.Upgrade(rw, req, nil)
+	} else {
+		// use default as before, so that those 2 buffer size variables are used as before
+		conn, err = websocket.Upgrade(rw, req, nil, WebSocketReadBufSize, WebSocketWriteBufSize)
+	}
+
 	if _, ok := err.(websocket.HandshakeError); ok {
 		http.Error(rw, `Can "Upgrade" only to "WebSocket".`, http.StatusBadRequest)
 		return

sockjs/rawwebsocket_test.go

@@ -119,3 +119,37 @@ func TestHandler_RawWebSocketCommunication(t *testing.T) {
 	}
 	<-done
 }
+
+func TestHandler_RawCustomWebSocketCommunication(t *testing.T) {
+	h := newTestHandler()
+	h.options.WebsocketUpgrader = &websocket.Upgrader{
+		ReadBufferSize:  0,
+		WriteBufferSize: 0,
+		CheckOrigin:     func(_ *http.Request) bool { return true },
+		Error:           func(w http.ResponseWriter, r *http.Request, status int, reason error) {},
+	}
+	server := httptest.NewServer(http.HandlerFunc(h.rawWebsocket))
+	url := "ws" + server.URL[4:]
+	var done = make(chan struct{})
+	h.handlerFunc = func(conn Session) {
+		conn.Send("message 1")
+		conn.Send("message 2")
+		expected := "[\"message 3\"]\n"
+		msg, err := conn.Recv()
+		if msg != expected || err != nil {
+			t.Errorf("Got '%s', expected '%s'", msg, expected)
+		}
+		conn.Close(123, "close")
+		close(done)
+	}
+	conn, _, _ := websocket.DefaultDialer.Dial(url, map[string][]string{"Origin": []string{server.URL}})
+	conn.WriteJSON([]string{"message 3"})
+	var expected = []string{"message 1", "message 2"}
+	for _, exp := range expected {
+		_, msg, err := conn.ReadMessage()
+		if string(msg) != exp || err != nil {
+			t.Errorf("Wrong frame, got '%s' and error '%v', expected '%s' without error", msg, err, exp)
+		}
+	}
+	<-done
+}

sockjs/websocket.go

@@ -17,7 +17,14 @@ var WebSocketReadBufSize = 4096
 var WebSocketWriteBufSize = 4096
 
 func (h *handler) sockjsWebsocket(rw http.ResponseWriter, req *http.Request) {
-	conn, err := websocket.Upgrade(rw, req, nil, WebSocketReadBufSize, WebSocketWriteBufSize)
+	var conn *websocket.Conn
+	var err error
+	if h.options.WebsocketUpgrader != nil {
+		conn, err = h.options.WebsocketUpgrader.Upgrade(rw, req, nil)
+	} else {
+		// use default as before, so that those 2 buffer size variables are used as before
+		conn, err = websocket.Upgrade(rw, req, nil, WebSocketReadBufSize, WebSocketWriteBufSize)
+	}
 	if _, ok := err.(websocket.HandshakeError); ok {
 		http.Error(rw, `Can "Upgrade" only to "WebSocket".`, http.StatusBadRequest)
 		return

sockjs/websocket_test.go

@@ -120,3 +120,36 @@ func TestHandler_WebSocketCommunication(t *testing.T) {
 	}
 	<-done
 }
+
+func TestHandler_CustomWebSocketCommunication(t *testing.T) {
+	h := newTestHandler()
+	h.options.WebsocketUpgrader = &websocket.Upgrader{
+		ReadBufferSize:  0,
+		WriteBufferSize: 0,
+		CheckOrigin:     func(_ *http.Request) bool { return true },
+		Error:           func(w http.ResponseWriter, r *http.Request, status int, reason error) {},
+	}
+	server := httptest.NewServer(http.HandlerFunc(h.sockjsWebsocket))
+	url := "ws" + server.URL[4:]
+	var done = make(chan struct{})
+	h.handlerFunc = func(conn Session) {
+		conn.Send("message 1")
+		conn.Send("message 2")
+		msg, err := conn.Recv()
+		if msg != "message 3" || err != nil {
+			t.Errorf("Got '%s', expected '%s'", msg, "message 3")
+		}
+		conn.Close(123, "close")
+		close(done)
+	}
+	conn, _, _ := websocket.DefaultDialer.Dial(url, map[string][]string{"Origin": []string{server.URL}})
+	conn.WriteJSON([]string{"message 3"})
+	var expected = []string{"o", `a["message 1"]`, `a["message 2"]`, `c[123,"close"]`}
+	for _, exp := range expected {
+		_, msg, err := conn.ReadMessage()
+		if string(msg) != exp || err != nil {
+			t.Errorf("Wrong frame, got '%s' and error '%v', expected '%s' without error", msg, err, exp)
+		}
+	}
+	<-done
+}