Move proxy target authorization handling to middleware instead of internals

- Add support and test for WebSockect authorization handling
- Cleanup

Author: v1k45

middleware/proxy.go

@@ -375,6 +375,13 @@ func ProxyWithConfig(config ProxyConfig) echo.MiddlewareFunc {
 				// that Balancer may have replaced with c.SetRequest.
 				req = c.Request()
 
+				// Handle authorization from target URL (for both HTTP and WebSocket)
+				if tgt.URL.User != nil {
+					username := tgt.URL.User.Username()
+					password, _ := tgt.URL.User.Password()
+					req.SetBasicAuth(username, password)
+				}
+
 				// Proxy
 				switch {
 				case c.IsWebSocket():
@@ -431,15 +438,5 @@ func proxyHTTP(tgt *ProxyTarget, c echo.Context, config ProxyConfig) http.Handle
 	proxy.Transport = config.Transport
 	proxy.ModifyResponse = config.ModifyResponse
 
-	// Handle authorization from target URL
-	if tgt.URL.User != nil {
-		originalDirector := proxy.Director
-		proxy.Director = func(req *http.Request) {
-			originalDirector(req)
-			password, _ := tgt.URL.User.Password()
-			req.SetBasicAuth(tgt.URL.User.Username(), password)
-		}
-	}
-
 	return proxy
 }

middleware/proxy_test.go

@@ -1093,3 +1093,68 @@ func TestProxyWithAuthorizationHeader(t *testing.T) {
 	assert.Equal(t, "target 2", rec.Body.String())
 	assert.Equal(t, "", receivedAuthHeader)
 }
+
+func TestProxyWithConfigWebSocketAuthorizationHeader(t *testing.T) {
+	// Capture the authorization header received by the WebSocket server
+	var receivedAuthHeader string
+	var authHeaderMutex sync.Mutex
+
+	// Create a WebSocket server that captures the Authorization header
+	wsServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		authHeaderMutex.Lock()
+		receivedAuthHeader = r.Header.Get("Authorization")
+		authHeaderMutex.Unlock()
+
+		wsHandler := func(conn *websocket.Conn) {
+			defer conn.Close()
+			for {
+				var msg string
+				err := websocket.Message.Receive(conn, &msg)
+				if err != nil {
+					return
+				}
+				// Echo message back to the client
+				websocket.Message.Send(conn, msg)
+			}
+		}
+		websocket.Server{Handler: wsHandler}.ServeHTTP(w, r)
+	}))
+	defer wsServer.Close()
+
+	// Create proxy server with target URL containing user:pass credentials
+	targetURL, _ := url.Parse(wsServer.URL)
+	targetURL.User = url.UserPassword("wsuser", "wspass")
+
+	e := echo.New()
+	balancer := NewRandomBalancer([]*ProxyTarget{{URL: targetURL}})
+	e.Use(ProxyWithConfig(ProxyConfig{Balancer: balancer}))
+
+	proxyServer := httptest.NewServer(e)
+	defer proxyServer.Close()
+
+	// Connect to the proxy WebSocket
+	proxyURL, _ := url.Parse(proxyServer.URL)
+	proxyURL.Scheme = "ws"
+	proxyURL.Path = "/"
+
+	wsConn, err := websocket.Dial(proxyURL.String(), "", "http://localhost/")
+	assert.NoError(t, err)
+	defer wsConn.Close()
+
+	// Send message to verify WebSocket connection works
+	sendMsg := "Hello, WebSocket with Auth!"
+	err = websocket.Message.Send(wsConn, sendMsg)
+	assert.NoError(t, err)
+
+	// Read response
+	var recvMsg string
+	err = websocket.Message.Receive(wsConn, &recvMsg)
+	assert.NoError(t, err)
+	assert.Equal(t, sendMsg, recvMsg)
+
+	// Verify authorization header was forwarded
+	authHeaderMutex.Lock()
+	expectedAuth := "Basic d3N1c2VyOndzcGFzcw==" // base64 of "wsuser:wspass"
+	assert.Equal(t, expectedAuth, receivedAuthHeader)
+	authHeaderMutex.Unlock()
+}