Improvements and tests.

sockjs/handler.go

@@ -7,8 +7,6 @@ import (
 	"regexp"
 	"strings"
 	"sync"
-
-	"github.com/gorilla/websocket"
 )
 
 var (
@@ -21,7 +19,6 @@ type handler struct {
 	options     Options
 	handlerFunc func(Session)
 	mappings    []*mapping
-	upgrader    websocket.Upgrader
 
 	sessionsMux sync.Mutex
 	sessions    map[string]*session
@@ -39,7 +36,6 @@ func newHandler(prefix string, opts Options, handlerFunc func(Session)) *handler
 		options:     opts,
 		handlerFunc: handlerFunc,
 		sessions:    make(map[string]*session),
-		upgrader:    opts.WebsocketUpgrader,
 	}
 
 	sessionPrefix := prefix + "/[^/.]+/[^/.]+"

sockjs/options.go

@@ -42,7 +42,7 @@ type Options struct {
 	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
+	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.
@@ -59,24 +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,
-	WebsocketUpgrader: websocket.Upgrader{
-		ReadBufferSize:  WebSocketReadBufSize,
-		WriteBufferSize: WebSocketWriteBufSize,
-		CheckOrigin: func(_ *http.Request) bool {
-			// Allow all connections by default.
-			return true
-		},
-		Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) {
-			// Don't return errors to maintain backwards compatibility.
-		},
-	},
+	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 := h.upgrader.Upgrade(rw, req, nil)
+	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 := h.upgrader.Upgrade(rw, req, nil)
+	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
+}