@@ -17,6 +17,7 @@ import (
1717 "github.com/mbland/hmacauth"
1818 "github.com/pusher/oauth2_proxy/cookie"
1919 "github.com/pusher/oauth2_proxy/providers"
20+ "github.com/yhat/wsutil"
2021)
2122
2223const (
@@ -95,9 +96,10 @@ type OAuthProxy struct {
9596
9697// UpstreamProxy represents an upstream server to proxy to
9798type UpstreamProxy struct {
98− upstream string
99− handler http.Handler
100− auth hmacauth.HmacAuth
99+ upstream string
100+ handler http.Handler
101+ wsHandler http.Handler
102+ auth hmacauth.HmacAuth
101103}
102104
103105// ServeHTTP proxies requests to the upstream provider while signing the
@@ -108,7 +110,12 @@ func (u *UpstreamProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
108110 r.Header.Set("GAP-Auth", w.Header().Get("GAP-Auth"))
109111 u.auth.SignRequest(r)
110112 }
111− u.handler.ServeHTTP(w, r)
113+ if u.wsHandler != nil && r.Header.Get("Connection") == "Upgrade" && r.Header.Get("Upgrade") == "websocket" {
114+ u.wsHandler.ServeHTTP(w, r)
115+ } else {
116+ u.handler.ServeHTTP(w, r)
117+ }
118+
112119}
113120
114121// NewReverseProxy creates a new reverse proxy for proxying requests to upstream
@@ -145,6 +152,26 @@ func NewFileServer(path string, filesystemPath string) (proxy http.Handler) {
145152 return http.StripPrefix(path, http.FileServer(http.Dir(filesystemPath)))
146153}
147154
155+// NewWebSocketOrRestReverseProxy creates a reverse proxy for REST or websocket based on url
156+func NewWebSocketOrRestReverseProxy(u *url.URL, opts *Options, auth hmacauth.HmacAuth) (restProxy http.Handler) {
157+ u.Path = ""
158+ proxy := NewReverseProxy(u, opts.FlushInterval)
159+ if !opts.PassHostHeader {
160+ setProxyUpstreamHostHeader(proxy, u)
161+ } else {
162+ setProxyDirector(proxy)
163+ }
164+
165+ // this should give us a wss:// scheme if the url is https:// based.
166+ var wsProxy *wsutil.ReverseProxy
167+ if opts.ProxyWebSockets {
168+ wsScheme := "ws" + strings.TrimPrefix(u.Scheme, "http")
169+ wsURL := &url.URL{Scheme: wsScheme, Host: u.Host}
170+ wsProxy = wsutil.NewSingleHostReverseProxy(wsURL)
171+ }
172+ return &UpstreamProxy{u.Host, proxy, wsProxy, auth}
173+}
174+
148175// NewOAuthProxy creates a new instance of OOuthProxy from the options provided
149176func NewOAuthProxy(opts *Options, validator func(string) bool) *OAuthProxy {
150177 serveMux := http.NewServeMux()
@@ -157,23 +184,17 @@ func NewOAuthProxy(opts *Options, validator func(string) bool) *OAuthProxy {
157184 path := u.Path
158185 switch u.Scheme {
159186 case httpScheme, httpsScheme:
160− u.Path = ""
161187 log.Printf("mapping path %q => upstream %q", path, u)
162− proxy := NewReverseProxy(u, opts.FlushInterval)
163− if !opts.PassHostHeader {
164− setProxyUpstreamHostHeader(proxy, u)
165− } else {
166− setProxyDirector(proxy)
167− }
168− serveMux.Handle(path,
169− &UpstreamProxy{u.Host, proxy, auth})
188+ proxy := NewWebSocketOrRestReverseProxy(u, opts, auth)
189+ serveMux.Handle(path, proxy)
190+
170191 case "file":
171192 if u.Fragment != "" {
172193 path = u.Fragment
173194 }
174195 log.Printf("mapping path %q => file system %q", path, u.Path)
175196 proxy := NewFileServer(path, u.Path)
176− serveMux.Handle(path, &UpstreamProxy{path, proxy, nil})
197+ serveMux.Handle(path, &UpstreamProxy{path, proxy, nil, nil})
177198 default:
178199 panic(fmt.Sprintf("unknown upstream protocol %s", u.Scheme))
179200 }