From 47c84b58bc66759b26812e309eed38beef85a06a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samuel=20Jir=C3=A9nius?= <samuel@jirenius.com>
Date: Mon, 5 Aug 2019 11:43:36 +0200
Subject: [PATCH] GH-105: Added option to enable WebSocket compression.

---
 main.go                              |  3 +++
 server/config.go                     |  2 ++
 server/service.go                    |  6 ++++--
 server/wsHandler.go                  | 18 +++++++++---------
 test/19websocket_compression_test.go | 17 +++++++++++++++++
 5 files changed, 35 insertions(+), 11 deletions(-)
 create mode 100644 test/19websocket_compression_test.go

diff --git a/main.go b/main.go
index 5fa56fe..c2f6484 100644
--- a/main.go
+++ b/main.go
@@ -36,6 +36,8 @@ Server Options:
 
 Common Options:
     -h, --help                       Show this message
+
+Configuration Documentation:         https://resgate.io/docs/get-started/configuration/
 `
 
 // Config holds server configuration
@@ -90,6 +92,7 @@ func (c *Config) Init(fs *flag.FlagSet, args []string) error {
 	fs.StringVar(&c.APIEncoding, "apiencoding", "", "Encoding for web resources.")
 	fs.IntVar(&c.RequestTimeout, "r", 0, "Timeout in milliseconds for NATS requests.")
 	fs.IntVar(&c.RequestTimeout, "reqtimeout", 0, "Timeout in milliseconds for NATS requests.")
+	fs.BoolVar(&c.Debug, "debug", false, "Enable debugging.")
 
 	if err := fs.Parse(args); err != nil {
 		printAndDie(fmt.Sprintf("error parsing arguments: %s", err.Error()), true)
diff --git a/server/config.go b/server/config.go
index 63a36a5..c6a4503 100644
--- a/server/config.go
+++ b/server/config.go
@@ -21,6 +21,8 @@ type Config struct {
 	TLSCert string `json:"certFile"`
 	TLSKey  string `json:"keyFile"`
 
+	WSCompression bool `json:"wsCompression"`
+
 	NoHTTP bool `json:"-"` // Disable start of the HTTP server. Used for testing
 
 	scheme           string
diff --git a/server/service.go b/server/service.go
index 8883ab7..0daf134 100644
--- a/server/service.go
+++ b/server/service.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/gorilla/websocket"
 	"github.com/resgateio/resgate/logger"
 	"github.com/resgateio/resgate/server/mq"
 	"github.com/resgateio/resgate/server/rescache"
@@ -26,8 +27,9 @@ type Service struct {
 	enc APIEncoder
 
 	// wsListener/wsConn
-	conns map[string]*wsConn // Connections by wsConn Id's
-	wg    sync.WaitGroup     // Wait for all connections to be disconnected
+	upgrader websocket.Upgrader
+	conns    map[string]*wsConn // Connections by wsConn Id's
+	wg       sync.WaitGroup     // Wait for all connections to be disconnected
 }
 
 // NewService creates a new Service
diff --git a/server/wsHandler.go b/server/wsHandler.go
index 934060f..b8662f7 100644
--- a/server/wsHandler.go
+++ b/server/wsHandler.go
@@ -8,17 +8,17 @@ import (
 	"github.com/gorilla/websocket"
 )
 
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-	CheckOrigin: func(r *http.Request) bool {
-		return true
-	},
-}
-
 var wsDisconnectTimeout = 3 * time.Second
 
 func (s *Service) initWSHandler() error {
+	s.upgrader = websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			return true
+		},
+		EnableCompression: s.cfg.WSCompression,
+	}
 	s.conns = make(map[string]*wsConn)
 	return nil
 }
@@ -31,7 +31,7 @@ func (s *Service) GetWSHandlerFunc() http.Handler {
 
 func (s *Service) wsHandler(w http.ResponseWriter, r *http.Request) {
 	// Upgrade to gorilla websocket
-	ws, err := upgrader.Upgrade(w, r, nil)
+	ws, err := s.upgrader.Upgrade(w, r, nil)
 	if err != nil {
 		s.Debugf("Failed to upgrade connection from %s: %s", r.RemoteAddr, err.Error())
 		return
diff --git a/test/19websocket_compression_test.go b/test/19websocket_compression_test.go
new file mode 100644
index 0000000..07d1229
--- /dev/null
+++ b/test/19websocket_compression_test.go
@@ -0,0 +1,17 @@
+package test
+
+import (
+	"testing"
+
+	"github.com/resgateio/resgate/server"
+)
+
+// Test subscribing to a resource with WebSocket compression enabled
+func TestWebSocketCompressionEnabled(t *testing.T) {
+	runTest(t, func(s *Session) {
+		c := s.Connect()
+		subscribeToTestModel(t, s, c)
+	}, func(c *server.Config) {
+		c.WSCompression = true
+	})
+}