support tls to non-tls and non-tls to tls websocket proxy

Author: t-ibayashi-safie

middleware/proxy.go

@@ -132,40 +132,31 @@ var DefaultProxyConfig = ProxyConfig{
 }
 
 func proxyRaw(t *ProxyTarget, c echo.Context, config ProxyConfig) http.Handler {
+	var dialFunc func(ctx context.Context, network, addr string) (net.Conn, error)
+	if transport, ok := config.Transport.(*http.Transport); ok {
+		if transport.TLSClientConfig != nil {
+			d := tls.Dialer{
+				Config: transport.TLSClientConfig,
+			}
+			dialFunc = d.DialContext
+		}
+	}
+	if dialFunc == nil {
+		var d net.Dialer
+		dialFunc = d.DialContext
+	}
+
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		in, _, err := c.Response().Hijack()
 		if err != nil {
 			c.Set("_error", fmt.Errorf("proxy raw, hijack error=%w, url=%s", err, t.URL))
 			return
 		}
 		defer in.Close()
-
-		var out net.Conn
-		if c.IsTLS() {
-			transport, ok := config.Transport.(*http.Transport)
-			if !ok {
-				c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, "proxy raw, invalid transport type"))
-				return
-			}
-
-			if transport.TLSClientConfig == nil {
-				c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, "proxy raw, TLSClientConfig is not set"))
-				return
-			}
-
-			out, err = tls.Dial("tcp", t.URL.Host, transport.TLSClientConfig)
-			if err != nil {
-				c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, dial error=%v, url=%s", err, t.URL)))
-				return
-			}
-			defer out.Close()
-		} else {
-			out, err = net.Dial("tcp", t.URL.Host)
-			if err != nil {
-				c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, dial error=%v, url=%s", err, t.URL)))
-				return
-			}
-			defer out.Close()
+		out, err := dialFunc(c.Request().Context(), "tcp", t.URL.Host)
+		if err != nil {
+			c.Set("_error", echo.NewHTTPError(http.StatusBadGateway, fmt.Sprintf("proxy raw, dial error=%v, url=%s", err, t.URL)))
+			return
 		}
 
 		// Write header

middleware/proxy_test.go

@@ -813,14 +813,8 @@ func TestModifyResponseUseContext(t *testing.T) {
 	assert.Equal(t, "CUSTOM_BALANCER", rec.Header().Get("FROM_BALANCER"))
 }
 
-func TestProxyWithConfigWebSocketTCP(t *testing.T) {
-	/*
-		Arrange
-	*/
-	e := echo.New()
-
-	// Create a WebSocket test server
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+func createSimpleWebSocketServer(serveTLS bool) *httptest.Server {
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		wsHandler := func(conn *websocket.Conn) {
 			defer conn.Close()
 			for {
@@ -834,15 +828,59 @@ func TestProxyWithConfigWebSocketTCP(t *testing.T) {
 			}
 		}
 		websocket.Server{Handler: wsHandler}.ServeHTTP(w, r)
-	}))
-	defer srv.Close()
+	})
+	if serveTLS {
+		return httptest.NewTLSServer(handler)
+	}
+	return httptest.NewServer(handler)
+}
 
-	tgtURL, _ := url.Parse(srv.URL)
-	balancer := NewRandomBalancer([]*ProxyTarget{{URL: tgtURL}})
+func createSimpleProxyServer(t *testing.T, srv *httptest.Server, serveTLS bool, toTLS bool) *httptest.Server {
+	e := echo.New()
 
-	e.Use(ProxyWithConfig(ProxyConfig{Balancer: balancer}))
+	if toTLS {
+		// proxy to tls target
+		tgtURL, _ := url.Parse(srv.URL)
+		tgtURL.Scheme = "wss"
+		balancer := NewRandomBalancer([]*ProxyTarget{{URL: tgtURL}})
 
+		defaultTransport, ok := http.DefaultTransport.(*http.Transport)
+		if !ok {
+			t.Fatal("Default transport is not of type *http.Transport")
+		}
+		transport := defaultTransport.Clone()
+		transport.TLSClientConfig = &tls.Config{
+			InsecureSkipVerify: true,
+		}
+		e.Use(ProxyWithConfig(ProxyConfig{Balancer: balancer, Transport: transport}))
+	} else {
+		// proxy to non-TLS target
+		tgtURL, _ := url.Parse(srv.URL)
+		balancer := NewRandomBalancer([]*ProxyTarget{{URL: tgtURL}})
+		e.Use(ProxyWithConfig(ProxyConfig{Balancer: balancer}))
+	}
+
+	if serveTLS {
+		// serve proxy server with TLS
+		ts := httptest.NewTLSServer(e)
+		return ts
+	}
+	// serve proxy server without TLS
 	ts := httptest.NewServer(e)
+	return ts
+}
+
+// TestProxyWithConfigWebSocketNonTLS2NonTLS tests the proxy with non-TLS to non-TLS WebSocket connection.
+func TestProxyWithConfigWebSocketNonTLS2NonTLS(t *testing.T) {
+	/*
+		Arrange
+	*/
+	// Create a WebSocket test server (non-TLS)
+	srv := createSimpleWebSocketServer(false)
+	defer srv.Close()
+
+	// create proxy server (non-TLS to non-TLS)
+	ts := createSimpleProxyServer(t, srv, false, false)
 	defer ts.Close()
 
 	tsURL, _ := url.Parse(ts.URL)
@@ -859,7 +897,7 @@ func TestProxyWithConfigWebSocketTCP(t *testing.T) {
 	defer wsConn.Close()
 
 	// Send message
-	sendMsg := "Hello, WebSocket!"
+	sendMsg := "Hello, Non TLS WebSocket!"
 	err = websocket.Message.Send(wsConn, sendMsg)
 	assert.NoError(t, err)
 
@@ -873,48 +911,103 @@ func TestProxyWithConfigWebSocketTCP(t *testing.T) {
 	assert.Equal(t, sendMsg, recvMsg)
 }
 
-func TestProxyWithConfigWebSocketTLS(t *testing.T) {
+// TestProxyWithConfigWebSocketTLS2TLS tests the proxy with TLS to TLS WebSocket connection.
+func TestProxyWithConfigWebSocketTLS2TLS(t *testing.T) {
 	/*
 		Arrange
 	*/
-	e := echo.New()
-
-	// Create a WebSocket test server
-	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		wsHandler := func(conn *websocket.Conn) {
-			defer conn.Close()
-			for {
-				var msg string
-				err := websocket.Message.Receive(conn, &msg)
-				if err != nil {
-					return
-				}
-				// message back to the client
-				websocket.Message.Send(conn, msg)
-			}
-		}
-		websocket.Server{Handler: wsHandler}.ServeHTTP(w, r)
-	}))
+	// Create a WebSocket test server (TLS)
+	srv := createSimpleWebSocketServer(true)
 	defer srv.Close()
 
-	// create proxy server
-	tgtURL, _ := url.Parse(srv.URL)
-	tgtURL.Scheme = "wss"
+	// create proxy server (TLS to TLS)
+	ts := createSimpleProxyServer(t, srv, true, true)
+	defer ts.Close()
 
-	balancer := NewRandomBalancer([]*ProxyTarget{{URL: tgtURL}})
+	tsURL, _ := url.Parse(ts.URL)
+	tsURL.Scheme = "wss"
+	tsURL.Path = "/"
 
-	defaultTransport, ok := http.DefaultTransport.(*http.Transport)
-	if !ok {
-		t.Fatal("Default transport is not of type *http.Transport")
-	}
-	transport := defaultTransport.Clone()
-	transport.TLSClientConfig = &tls.Config{
-		InsecureSkipVerify: true,
+	/*
+		Act
+	*/
+	origin, err := url.Parse(ts.URL)
+	assert.NoError(t, err)
+	config := &websocket.Config{
+		Location:  tsURL,
+		Origin:    origin,
+		TlsConfig: &tls.Config{InsecureSkipVerify: true}, // skip verify for testing
+		Version:   websocket.ProtocolVersionHybi13,
 	}
-	e.Use(ProxyWithConfig(ProxyConfig{Balancer: balancer, Transport: transport}))
+	wsConn, err := websocket.DialConfig(config)
+	assert.NoError(t, err)
+	defer wsConn.Close()
+
+	// Send message
+	sendMsg := "Hello, TLS to TLS WebSocket!"
+	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)
+}
+
+// TestProxyWithConfigWebSocketNonTLS2TLS tests the proxy with non-TLS to TLS WebSocket connection.
+func TestProxyWithConfigWebSocketNonTLS2TLS(t *testing.T) {
+	/*
+		Arrange
+	*/
+
+	// Create a WebSocket test server (TLS)
+	srv := createSimpleWebSocketServer(true)
+	defer srv.Close()
+
+	// create proxy server (Non-TLS to TLS)
+	ts := createSimpleProxyServer(t, srv, false, true)
+	defer ts.Close()
+
+	tsURL, _ := url.Parse(ts.URL)
+	tsURL.Scheme = "ws"
+	tsURL.Path = "/"
+
+	/*
+		Act
+	*/
+	// Connect to the proxy WebSocket
+	wsConn, err := websocket.Dial(tsURL.String(), "", "http://localhost/")
+	assert.NoError(t, err)
+	defer wsConn.Close()
+
+	// Send message
+	sendMsg := "Hello, Non TLS to TLS WebSocket!"
+	err = websocket.Message.Send(wsConn, sendMsg)
+	assert.NoError(t, err)
+
+	/*
+		Assert
+	*/
+	// Read response
+	var recvMsg string
+	err = websocket.Message.Receive(wsConn, &recvMsg)
+	assert.NoError(t, err)
+	assert.Equal(t, sendMsg, recvMsg)
+}
+
+// TestProxyWithConfigWebSocketTLSToNoneTLS tests the proxy with TLS to non-TLS WebSocket connection. (TLS termination)
+func TestProxyWithConfigWebSocketTLS2NonTLS(t *testing.T) {
+	/*
+		Arrange
+	*/
+
+	// Create a WebSocket test server (non-TLS)
+	srv := createSimpleWebSocketServer(false)
+	defer srv.Close()
 
-	// Start test server
-	ts := httptest.NewTLSServer(e)
+	// create proxy server (TLS to non-TLS)
+	ts := createSimpleProxyServer(t, srv, true, false)
 	defer ts.Close()
 
 	tsURL, _ := url.Parse(ts.URL)
@@ -937,7 +1030,7 @@ func TestProxyWithConfigWebSocketTLS(t *testing.T) {
 	defer wsConn.Close()
 
 	// Send message
-	sendMsg := "Hello, TLS WebSocket!"
+	sendMsg := "Hello, TLS to NoneTLS WebSocket!"
 	err = websocket.Message.Send(wsConn, sendMsg)
 	assert.NoError(t, err)