Add local websocket proxy.

Add local websocket proxy

fmt

update websocket error logging

also stop sending http.Error on upgraded connection.

add rewrite ws connection to arch diagram

format

Author: tjholm

pkg/cloud/cloud.go

@@ -321,7 +321,7 @@ func New(projectName string, opts LocalCloudOptions) (*LocalCloud, error) {
 		return nil, err
 	}
 
-	localWebsites := websites.NewLocalWebsitesService(localGateway.GetApiAddress, opts.LocalCloudMode == StartMode)
+	localWebsites := websites.NewLocalWebsitesService(localGateway.GetApiAddress, localGateway.GetWebsocketAddress, opts.LocalCloudMode == StartMode)
 
 	return &LocalCloud{
 		servers:    make(map[string]*server.NitricServer),

pkg/cloud/gateway/gateway.go

@@ -46,7 +46,6 @@ import (
 	"github.com/nitrictech/cli/pkg/netx"
 	"github.com/nitrictech/cli/pkg/project/localconfig"
 	"github.com/nitrictech/cli/pkg/system"
-	"github.com/nitrictech/cli/pkg/view/tui"
 
 	base_http "github.com/nitrictech/nitric/cloud/common/runtime/gateway"
 
@@ -157,6 +156,19 @@ func (s *LocalGatewayService) GetApiAddress(apiName string) string {
 	return ""
 }
 
+func (s *LocalGatewayService) GetWebsocketAddress(socketName string) string {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+
+	addresses := s.GetWebsocketAddresses()
+
+	if address, ok := addresses[socketName]; ok {
+		return address
+	}
+
+	return ""
+}
+
 func (s *LocalGatewayService) GetHttpWorkerAddresses() map[string]string {
 	s.lock.RLock()
 	defer s.lock.RUnlock()
@@ -349,14 +361,14 @@ func (s *LocalGatewayService) handleWebsocketRequest(socketName string) func(ctx
 					SocketName:   socketName,
 				})
 				if err != nil {
-					tui.Error.Println(err.Error())
+					system.Logf("Websocket error: %s", err.Error())
 					return
 				}
 			}()
 
 			err = s.websocketPlugin.RegisterConnection(socketName, connectionId, ws)
 			if err != nil {
-				tui.Error.Println(err.Error())
+				system.Logf("Websocket error: %s", err.Error())
 				return
 			}
 
@@ -372,7 +384,7 @@ func (s *LocalGatewayService) handleWebsocketRequest(socketName string) func(ctx
 				if err != nil && websocket.IsCloseError(err, 1001, 1005) {
 					break
 				} else if err != nil {
-					log.Println("read:", err)
+					system.Logf("websocket read error: %v", err)
 					break
 				}
 
@@ -390,7 +402,7 @@ func (s *LocalGatewayService) handleWebsocketRequest(socketName string) func(ctx
 					},
 				})
 				if err != nil {
-					tui.Error.Println(err.Error())
+					system.Logf("Websocket error: %s", err.Error())
 					return
 				}
 			}
@@ -407,13 +419,13 @@ func (s *LocalGatewayService) handleWebsocketRequest(socketName string) func(ctx
 				},
 			})
 			if err != nil {
-				tui.Error.Println(err.Error())
+				system.Logf("Websocket error: %s", err.Error())
 				return
 			}
 		})
 		if err != nil {
 			if _, ok := err.(websocket.HandshakeError); ok {
-				tui.Error.Println(err.Error())
+				system.Logf("Websocket error: %s", err.Error())
 			}
 
 			return

pkg/cloud/websites/websites.go

@@ -30,8 +30,10 @@ import (
 	"sync"
 
 	"github.com/asaskevich/EventBus"
+	"github.com/gorilla/websocket"
 
 	"github.com/nitrictech/cli/pkg/netx"
+	"github.com/nitrictech/cli/pkg/system"
 	deploymentspb "github.com/nitrictech/nitric/core/pkg/proto/deployments/v1"
 )
 
@@ -54,11 +56,12 @@ type (
 )
 
 type LocalWebsiteService struct {
-	websiteRegLock sync.RWMutex
-	state          State
-	port           int
-	getApiAddress  GetApiAddress
-	isStartCmd     bool
+	websiteRegLock      sync.RWMutex
+	state               State
+	port                int
+	getApiAddress       GetApiAddress
+	getWebsocketAddress GetApiAddress
+	isStartCmd          bool
 
 	bus EventBus.Bus
 }
@@ -172,6 +175,23 @@ func (h staticSiteHandler) ServeHTTP(res http.ResponseWriter, req *http.Request)
 }
 
 // Start - Start the local website service
+func proxyWebSocketMessages(src, dst *websocket.Conn, errChan chan error) {
+	for {
+		messageType, message, err := src.ReadMessage()
+		if err != nil {
+			errChan <- err
+			return
+		}
+
+		err = dst.WriteMessage(messageType, message)
+		if err != nil {
+			errChan <- err
+			return
+		}
+	}
+}
+
+// Serve - Serve a website from the local filesystem
 func (l *LocalWebsiteService) Start(websites []Website) error {
 	newLis, err := netx.GetNextListener(netx.MinPort(5000))
 	if err != nil {
@@ -203,6 +223,55 @@ func (l *LocalWebsiteService) Start(websites []Website) error {
 		proxy.ServeHTTP(res, req)
 	})
 
+	// Register the API handler
+	mux.HandleFunc("/ws/{name}", func(w http.ResponseWriter, r *http.Request) {
+		// Get the WebSocket API name from the request path
+		apiName := r.PathValue("name")
+
+		// Get the address of the WebSocket API
+		apiAddress := l.getWebsocketAddress(apiName)
+		if apiAddress == "" {
+			http.Error(w, fmt.Sprintf("WebSocket API %s not found", apiName), http.StatusNotFound)
+			return
+		}
+
+		// Dial the backend WebSocket server
+		targetURL := fmt.Sprintf("ws://%s%s", apiAddress, r.URL.Path)
+		if r.URL.RawQuery != "" {
+			targetURL = fmt.Sprintf("%s?%s", targetURL, r.URL.RawQuery)
+		}
+
+		targetConn, _, err := websocket.DefaultDialer.Dial(targetURL, nil)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("Failed to connect to backend WebSocket server: %v", err), http.StatusInternalServerError)
+			return
+		}
+		defer targetConn.Close()
+
+		// Upgrade the HTTP connection to a WebSocket connection
+		upgrader := websocket.Upgrader{}
+
+		clientConn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("Failed to upgrade to WebSocket: %v", err), http.StatusInternalServerError)
+			return
+		}
+
+		defer clientConn.Close()
+
+		// Proxy messages between the client and the backend WebSocket server
+		errChan := make(chan error, 2)
+		go proxyWebSocketMessages(clientConn, targetConn, errChan)
+		go proxyWebSocketMessages(targetConn, clientConn, errChan)
+
+		// Wait for an error to occur
+		err = <-errChan
+		if err != nil && !errors.Is(err, websocket.ErrCloseSent) {
+			// Because the error is already proxied through by the connection we can just log the error here
+			system.Logf("received error on websocket %s: %v", apiName, err)
+		}
+	})
+
 	// Register the SPA handler for each website
 	for i := range websites {
 		website := &websites[i]
@@ -231,11 +300,12 @@ func (l *LocalWebsiteService) Start(websites []Website) error {
 	return nil
 }
 
-func NewLocalWebsitesService(getApiAddress GetApiAddress, isStartCmd bool) *LocalWebsiteService {
+func NewLocalWebsitesService(getApiAddress GetApiAddress, getWebsocketAddress GetApiAddress, isStartCmd bool) *LocalWebsiteService {
 	return &LocalWebsiteService{
-		state:         State{},
-		bus:           EventBus.New(),
-		getApiAddress: getApiAddress,
-		isStartCmd:    isStartCmd,
+		state:               State{},
+		bus:                 EventBus.New(),
+		getApiAddress:       getApiAddress,
+		getWebsocketAddress: getWebsocketAddress,
+		isStartCmd:          isStartCmd,
 	}
 }

pkg/dashboard/frontend/src/lib/utils/generate-architecture-data.ts

@@ -609,6 +609,23 @@ export function generateArchitectureData(data: WebSocketResponse): {
         label: `Rewrites to /api/${api.name}`,
       })
     })
+
+    data.websockets.forEach((websocket) => {
+      edges.push({
+        id: `e-${websocket.name}-websites`,
+        source: websitesNode.id,
+        target: `websocket-${websocket.name}`,
+        animated: true,
+        markerEnd: {
+          type: MarkerType.ArrowClosed,
+        },
+        markerStart: {
+          type: MarkerType.ArrowClosed,
+          orient: 'auto-start-reverse',
+        },
+        label: `Rewrites to /ws/${websocket.name}`,
+      })
+    })
   }
 
   data.services.forEach((service) => {