From e92c9b278e7090ad6f54595de0633ef0a013fa9d Mon Sep 17 00:00:00 2001
From: Rohith Jayawardene <gambol99@gmail.com>
Date: Thu, 8 Feb 2018 11:54:55 +0000
Subject: [PATCH] Websocket Headers (#311)

- fixing up the headers for websockets
---
 forwarding.go | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/forwarding.go b/forwarding.go
index 6851a10..2ea61f6 100644
--- a/forwarding.go
+++ b/forwarding.go
@@ -30,7 +30,7 @@ func (r *oauthProxy) proxyMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		next.ServeHTTP(w, req)
 
-		// step: retrieve the request scope
+		// @step: retrieve the request scope
 		scope := req.Context().Value(contextScopeName)
 		if scope != nil {
 			sc := scope.(*RequestScope)
@@ -39,23 +39,12 @@ func (r *oauthProxy) proxyMiddleware(next http.Handler) http.Handler {
 			}
 		}
 
-		if isUpgradedConnection(req) {
-			r.log.Debug("upgrading the connnection", zap.String("client_ip", req.RemoteAddr))
-			if err := tryUpdateConnection(req, w, r.endpoint); err != nil {
-				r.log.Error("failed to upgrade connection", zap.Error(err))
-				w.WriteHeader(http.StatusInternalServerError)
-				return
-			}
-			return
-		}
-
-		// add any custom headers to the request
+		// @step: add any custom headers to the request
 		for k, v := range r.config.Headers {
 			req.Header.Set(k, v)
 		}
 
-		// By default goproxy only provides a forwarding proxy, thus all requests have to be absolute
-		// and we must update the host headers
+		// @note: by default goproxy only provides a forwarding proxy, thus all requests have to be absolute and we must update the host headers
 		req.URL.Host = r.endpoint.Host
 		req.URL.Scheme = r.endpoint.Scheme
 		if v := req.Header.Get("Host"); v != "" {
@@ -65,10 +54,21 @@ func (r *oauthProxy) proxyMiddleware(next http.Handler) http.Handler {
 			req.Host = r.endpoint.Host
 		}
 
+		// @step: add the proxy forwarding headers
 		req.Header.Add("X-Forwarded-For", realIP(req))
 		req.Header.Set("X-Forwarded-Host", req.URL.Host)
 		req.Header.Set("X-Forwarded-Proto", req.Header.Get("X-Forwarded-Proto"))
 
+		if isUpgradedConnection(req) {
+			r.log.Debug("upgrading the connnection", zap.String("client_ip", req.RemoteAddr))
+			if err := tryUpdateConnection(req, w, r.endpoint); err != nil {
+				r.log.Error("failed to upgrade connection", zap.Error(err))
+				w.WriteHeader(http.StatusInternalServerError)
+				return
+			}
+			return
+		}
+
 		r.upstream.ServeHTTP(w, req)
 	})
 }
-- 
GitLab