Merge websocket proxy feature from openshift/oauth-proxy

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Changes since v3.1.0
 
+- [#92](https://github.com/pusher/oauth2_proxy/pull/92) Merge websocket proxy feature from openshift/oauth-proxy (@butzist)
 - [#57](https://github.com/pusher/oauth2_proxy/pull/57) Fall back to using OIDC Subject instead of Email (@aigarius)
 - [#85](https://github.com/pusher/oauth2_proxy/pull/85) Use non-root user in docker images (@kskewes)
 - [#68](https://github.com/pusher/oauth2_proxy/pull/68) forward X-Auth-Access-Token header (@davidholsgrove)

README.md

@@ -242,6 +242,7 @@ Usage of oauth2_proxy:
   -profile-url string: Profile access endpoint
   -provider string: OAuth provider (default "google")
   -proxy-prefix string: the url root path that this proxy should be nested under (e.g. /<oauth2>/sign_in) (default "/oauth2")
+  -proxy-websockets: enables WebSocket proxying (default true)
   -redeem-url string: Token redemption endpoint
   -redirect-url string: the OAuth Redirect URL. ie: "https://internalapp.yourcompany.com/oauth2/callback"
   -request-logging: Log requests to stdout (default true)

main.go

@@ -10,7 +10,7 @@ import (
 	"time"
 
 	"github.com/BurntSushi/toml"
-	"github.com/mreiferson/go-options"
+	options "github.com/mreiferson/go-options"
 )
 
 func main() {
@@ -62,6 +62,7 @@ func main() {
 	flagSet.String("custom-templates-dir", "", "path to custom html templates")
 	flagSet.String("footer", "", "custom footer string. Use \"-\" to disable default footer.")
 	flagSet.String("proxy-prefix", "/oauth2", "the url root path that this proxy should be nested under (e.g. /<oauth2>/sign_in)")
+	flagSet.Bool("proxy-websockets", true, "enables WebSocket proxying")
 
 	flagSet.String("cookie-name", "_oauth2_proxy", "the name of the cookie that the oauth_proxy creates")
 	flagSet.String("cookie-secret", "", "the seed string for secure cookies (optionally base64 encoded)")

oauthproxy.go

@@ -17,6 +17,7 @@ import (
 	"github.com/mbland/hmacauth"
 	"github.com/pusher/oauth2_proxy/cookie"
 	"github.com/pusher/oauth2_proxy/providers"
+	"github.com/yhat/wsutil"
 )
 
 const (
@@ -95,9 +96,10 @@ type OAuthProxy struct {
 
 // UpstreamProxy represents an upstream server to proxy to
 type UpstreamProxy struct {
-	upstream string
-	handler  http.Handler
-	auth     hmacauth.HmacAuth
+	upstream  string
+	handler   http.Handler
+	wsHandler http.Handler
+	auth      hmacauth.HmacAuth
 }
 
 // ServeHTTP proxies requests to the upstream provider while signing the
@@ -108,7 +110,12 @@ func (u *UpstreamProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		r.Header.Set("GAP-Auth", w.Header().Get("GAP-Auth"))
 		u.auth.SignRequest(r)
 	}
-	u.handler.ServeHTTP(w, r)
+	if u.wsHandler != nil && r.Header.Get("Connection") == "Upgrade" && r.Header.Get("Upgrade") == "websocket" {
+		u.wsHandler.ServeHTTP(w, r)
+	} else {
+		u.handler.ServeHTTP(w, r)
+	}
+
 }
 
 // NewReverseProxy creates a new reverse proxy for proxying requests to upstream
@@ -145,6 +152,26 @@ func NewFileServer(path string, filesystemPath string) (proxy http.Handler) {
 	return http.StripPrefix(path, http.FileServer(http.Dir(filesystemPath)))
 }
 
+// NewWebSocketOrRestReverseProxy creates a reverse proxy for REST or websocket based on url
+func NewWebSocketOrRestReverseProxy(u *url.URL, opts *Options, auth hmacauth.HmacAuth) (restProxy http.Handler) {
+	u.Path = ""
+	proxy := NewReverseProxy(u, opts.FlushInterval)
+	if !opts.PassHostHeader {
+		setProxyUpstreamHostHeader(proxy, u)
+	} else {
+		setProxyDirector(proxy)
+	}
+
+	// this should give us a wss:// scheme if the url is https:// based.
+	var wsProxy *wsutil.ReverseProxy
+	if opts.ProxyWebSockets {
+		wsScheme := "ws" + strings.TrimPrefix(u.Scheme, "http")
+		wsURL := &url.URL{Scheme: wsScheme, Host: u.Host}
+		wsProxy = wsutil.NewSingleHostReverseProxy(wsURL)
+	}
+	return &UpstreamProxy{u.Host, proxy, wsProxy, auth}
+}
+
 // NewOAuthProxy creates a new instance of OOuthProxy from the options provided
 func NewOAuthProxy(opts *Options, validator func(string) bool) *OAuthProxy {
 	serveMux := http.NewServeMux()
@@ -157,23 +184,17 @@ func NewOAuthProxy(opts *Options, validator func(string) bool) *OAuthProxy {
 		path := u.Path
 		switch u.Scheme {
 		case httpScheme, httpsScheme:
-			u.Path = ""
 			log.Printf("mapping path %q => upstream %q", path, u)
-			proxy := NewReverseProxy(u, opts.FlushInterval)
-			if !opts.PassHostHeader {
-				setProxyUpstreamHostHeader(proxy, u)
-			} else {
-				setProxyDirector(proxy)
-			}
-			serveMux.Handle(path,
-				&UpstreamProxy{u.Host, proxy, auth})
+			proxy := NewWebSocketOrRestReverseProxy(u, opts, auth)
+			serveMux.Handle(path, proxy)
+
 		case "file":
 			if u.Fragment != "" {
 				path = u.Fragment
 			}
 			log.Printf("mapping path %q => file system %q", path, u.Path)
 			proxy := NewFileServer(path, u.Path)
-			serveMux.Handle(path, &UpstreamProxy{path, proxy, nil})
+			serveMux.Handle(path, &UpstreamProxy{path, proxy, nil, nil})
 		default:
 			panic(fmt.Sprintf("unknown upstream protocol %s", u.Scheme))
 		}

oauthproxy_test.go

@@ -19,13 +19,91 @@ import (
 	"github.com/pusher/oauth2_proxy/providers"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/net/websocket"
 )
 
 func init() {
 	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
 
 }
 
+type WebSocketOrRestHandler struct {
+	restHandler http.Handler
+	wsHandler   http.Handler
+}
+
+func (h *WebSocketOrRestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		h.wsHandler.ServeHTTP(w, r)
+	} else {
+		h.restHandler.ServeHTTP(w, r)
+	}
+}
+
+func TestWebSocketProxy(t *testing.T) {
+	handler := WebSocketOrRestHandler{
+		restHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(200)
+			hostname, _, _ := net.SplitHostPort(r.Host)
+			w.Write([]byte(hostname))
+		}),
+		wsHandler: websocket.Handler(func(ws *websocket.Conn) {
+			defer ws.Close()
+			var data []byte
+			err := websocket.Message.Receive(ws, &data)
+			if err != nil {
+				t.Fatalf("err %s", err)
+				return
+			}
+			err = websocket.Message.Send(ws, data)
+			if err != nil {
+				t.Fatalf("err %s", err)
+			}
+			return
+		}),
+	}
+	backend := httptest.NewServer(&handler)
+	defer backend.Close()
+
+	backendURL, _ := url.Parse(backend.URL)
+
+	options := NewOptions()
+	var auth hmacauth.HmacAuth
+	options.PassHostHeader = true
+	proxyHandler := NewWebSocketOrRestReverseProxy(backendURL, options, auth)
+	frontend := httptest.NewServer(proxyHandler)
+	defer frontend.Close()
+
+	frontendURL, _ := url.Parse(frontend.URL)
+	frontendWSURL := "ws://" + frontendURL.Host + "/"
+
+	ws, err := websocket.Dial(frontendWSURL, "", "http://localhost/")
+	if err != nil {
+		t.Fatalf("err %s", err)
+	}
+	request := []byte("hello, world!")
+	err = websocket.Message.Send(ws, request)
+	if err != nil {
+		t.Fatalf("err %s", err)
+	}
+	var response = make([]byte, 1024)
+	websocket.Message.Receive(ws, &response)
+	if err != nil {
+		t.Fatalf("err %s", err)
+	}
+	if g, e := string(request), string(response); g != e {
+		t.Errorf("got body %q; expected %q", g, e)
+	}
+
+	getReq, _ := http.NewRequest("GET", frontend.URL, nil)
+	res, _ := http.DefaultClient.Do(getReq)
+	bodyBytes, _ := ioutil.ReadAll(res.Body)
+	backendHostname, _, _ := net.SplitHostPort(backendURL.Host)
+	if g, e := string(bodyBytes), backendHostname; g != e {
+		t.Errorf("got body %q; expected %q", g, e)
+	}
+}
+
 func TestNewReverseProxy(t *testing.T) {
 	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(200)

options.go

@@ -21,14 +21,15 @@ import (
 // Options holds Configuration Options that can be set by Command Line Flag,
 // or Config File
 type Options struct {
-	ProxyPrefix  string `flag:"proxy-prefix" cfg:"proxy-prefix"`
-	HTTPAddress  string `flag:"http-address" cfg:"http_address"`
-	HTTPSAddress string `flag:"https-address" cfg:"https_address"`
-	RedirectURL  string `flag:"redirect-url" cfg:"redirect_url"`
-	ClientID     string `flag:"client-id" cfg:"client_id" env:"OAUTH2_PROXY_CLIENT_ID"`
-	ClientSecret string `flag:"client-secret" cfg:"client_secret" env:"OAUTH2_PROXY_CLIENT_SECRET"`
-	TLSCertFile  string `flag:"tls-cert" cfg:"tls_cert_file"`
-	TLSKeyFile   string `flag:"tls-key" cfg:"tls_key_file"`
+	ProxyPrefix     string `flag:"proxy-prefix" cfg:"proxy-prefix"`
+	ProxyWebSockets bool   `flag:"proxy-websockets" cfg:"proxy_websockets"`
+	HTTPAddress     string `flag:"http-address" cfg:"http_address"`
+	HTTPSAddress    string `flag:"https-address" cfg:"https_address"`
+	RedirectURL     string `flag:"redirect-url" cfg:"redirect_url"`
+	ClientID        string `flag:"client-id" cfg:"client_id" env:"OAUTH2_PROXY_CLIENT_ID"`
+	ClientSecret    string `flag:"client-secret" cfg:"client_secret" env:"OAUTH2_PROXY_CLIENT_SECRET"`
+	TLSCertFile     string `flag:"tls-cert" cfg:"tls_cert_file"`
+	TLSKeyFile      string `flag:"tls-key" cfg:"tls_key_file"`
 
 	AuthenticatedEmailsFile  string   `flag:"authenticated-emails-file" cfg:"authenticated_emails_file"`
 	AzureTenant              string   `flag:"azure-tenant" cfg:"azure_tenant"`
@@ -105,6 +106,7 @@ type SignatureData struct {
 func NewOptions() *Options {
 	return &Options{
 		ProxyPrefix:          "/oauth2",
+		ProxyWebSockets:      true,
 		HTTPAddress:          "127.0.0.1:4180",
 		HTTPSAddress:         ":443",
 		DisplayHtpasswdForm:  true,