Merge pull request #256 from buzzfeed/jusshersmith-test-websockets

sso_proxy: add test for websockets and update docs

.circleci/config.yml

@@ -1,6 +1,6 @@
 defaults: &defaults
   docker:
-    - image: circleci/golang:1.11
+    - image: circleci/golang:1.12
   working_directory: /go/src/github.com/buzzfeed/sso
 
 attach_workspace: &attach_workspace

docs/sso_config.md

@@ -178,6 +178,9 @@ honored as valid. The grace period ends either after the TTL expires or when
 * The grace period defined by `grace_period_ttl` is granted on a per-user basis,
   starting from the first failure to authenticate.
 
+### Websockets
+SSO supports upstreams that use websockets, providing the upstream has a positive flush interval (`flush_interval`) set.
+
 
 ### `sso_proxy` Endpoints
 * `/` - Begins the proxy process, attempting to authenticate the session cookie, redirecting to `sso-authenticator` if there is no cookie or an invalid one.

go.mod

@@ -5,6 +5,7 @@ require (
 	github.com/benbjohnson/clock v0.0.0-20161215174838-7dc76406b6d3
 	github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
 	github.com/datadog/datadog-go v0.0.0-20180822151419-281ae9f2d895
+	github.com/gorilla/websocket v1.4.0
 	github.com/imdario/mergo v0.3.7
 	github.com/kelseyhightower/envconfig v1.3.0
 	github.com/mccutchen/go-httpbin v1.1.1

go.sum

@@ -100,6 +100,7 @@ github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/gorilla/handlers v1.4.0/go.mod h1:Qkdc/uu4tH4g6mTK6auzZ766c4CA0Ng8+o/OAirnOIQ=
 github.com/gorilla/mux v1.7.2/go.mod h1:1lud6UwP+6orDFRuTfBEV8e9/aOM/c4fVVCaMa2zaAs=
+github.com/gorilla/websocket v1.4.0 h1:WDFjx/TMzVgy9VdMMQi2K2Emtwi2QcUQsztZ/zLaH/Q=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=

internal/proxy/reverse_proxy_test.go

@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"github.com/18F/hmacauth"
+	"github.com/gorilla/websocket"
 )
 
 func TestDeleteCookie(t *testing.T) {
@@ -835,3 +836,112 @@ func TestHMACSignatures(t *testing.T) {
 		})
 	}
 }
+
+func TestWebsocketSupport(t *testing.T) {
+	testCases := []struct {
+		Name          string
+		FlushInterval time.Duration
+		Timeout       time.Duration
+		ExpectedMsg   string
+		ExpectedErr   error
+	}{
+		{
+			Name:          "valid websocket support",
+			FlushInterval: 100 * time.Millisecond,
+			Timeout:       10 * time.Second,
+			ExpectedMsg:   "websocket-support-enabled",
+		},
+		{
+			// this test currently exists to mimic and record the behavior
+			// that occurs when a non-positive flush interval is set.
+			// (we set a default timeout, but no default flush interval)
+			// this test should be updated if this is combatted against.
+			Name:          "invalid websocket support, non-positive flush interval",
+			FlushInterval: 0,
+			Timeout:       10 * time.Second,
+			ExpectedErr:   websocket.ErrBadHandshake,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			var upgrader = websocket.Upgrader{
+				ReadBufferSize:  1024,
+				WriteBufferSize: 1024,
+			}
+			// set up http server with websocket support,
+			// and simple functionality of reading and respoding
+			// with a message.
+			backend := httptest.NewServer(http.HandlerFunc(
+				func(w http.ResponseWriter, r *http.Request) {
+					conn, err := upgrader.Upgrade(w, r, nil)
+					if err != nil {
+						return
+					}
+					for {
+						messageType, p, err := conn.ReadMessage()
+						if err != nil {
+							return
+						}
+						if err := conn.WriteMessage(messageType, p); err != nil {
+							return
+						}
+					}
+				},
+			))
+			defer backend.Close()
+
+			// set up upstream with FlushInterval (required for websocket)
+			u, _ := url.Parse(backend.URL)
+			config := &UpstreamConfig{
+				Route: &SimpleRoute{
+					ToURL: u,
+				},
+				FlushInterval: tc.FlushInterval,
+				Timeout:       tc.Timeout,
+			}
+
+			reverseProxy, err := NewUpstreamReverseProxy(config, nil)
+			if err != nil {
+				t.Fatalf("unexpected err: %v", err)
+			}
+
+			frontend := httptest.NewServer(reverseProxy)
+			defer frontend.Close()
+
+			// test websocket connectivity
+			// send message, and expect it to be returned afterwards.
+			websocketURL, _ := url.Parse(frontend.URL)
+			websocketURL.Scheme = "ws"
+			ws, _, err := websocket.DefaultDialer.Dial(websocketURL.String(), nil)
+			if err != nil {
+				if err == tc.ExpectedErr {
+					// if the expected error is returned we want to set the test as passed
+					// and stop further execution. this is purely related to the
+					// 'invalid websocket support, non-positive flush interval' test.
+					return
+				} else {
+					t.Logf("expected: %q", tc.ExpectedErr)
+					t.Logf("     got: %q", err)
+					t.Fatalf("unexpected error returned")
+				}
+			}
+			defer ws.Close()
+
+			if err := ws.WriteMessage(websocket.TextMessage, []byte(tc.ExpectedMsg)); err != nil {
+				t.Fatalf("%v", err)
+			}
+
+			_, byteResp, err := ws.ReadMessage()
+			if err != nil {
+				t.Fatalf("%v", err)
+			}
+
+			resp := string(byteResp)
+			if resp != tc.ExpectedMsg {
+				t.Logf("expected: %q", tc.ExpectedMsg)
+				t.Logf("     got: %q", resp)
+				t.Fatalf("unexpected message returned")
+			}
+		})
+	}
+}