rebase websockets proxy

Author: tjholm

pkg/cloud/websites/websites.go

@@ -31,8 +31,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"
 )
 
@@ -55,10 +57,11 @@ type (
 )
 
 type LocalWebsiteService struct {
-	websiteRegLock sync.RWMutex
-	state          State
-	getApiAddress  GetApiAddress
-	isStartCmd     bool
+	websiteRegLock      sync.RWMutex
+	state               State
+	getApiAddress       GetApiAddress
+	getWebsocketAddress GetApiAddress
+	isStartCmd          bool
 
 	bus EventBus.Bus
 }
@@ -74,6 +77,22 @@ func (l *LocalWebsiteService) SubscribeToState(fn func(State)) {
 	_ = l.bus.Subscribe(localWebsitesTopic, fn)
 }
 
+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
+		}
+	}
+}
+
 // register - Register a new website
 func (l *LocalWebsiteService) register(website Website, port int) {
 	l.websiteRegLock.Lock()
@@ -182,6 +201,55 @@ func (h staticSiteHandler) ServeHTTP(res http.ResponseWriter, req *http.Request)
 	h.serveStatic(res, req)
 }
 
+// createWebsocketPathHandler creates a handler for WebSocket proxy requests
+func (l *LocalWebsiteService) createWebsocketPathHandler(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)
+	}
+}
+
 // createAPIPathHandler creates a handler for API proxy requests
 func (l *LocalWebsiteService) createAPIPathHandler() http.HandlerFunc {
 	return func(res http.ResponseWriter, req *http.Request) {
@@ -252,6 +320,9 @@ func (l *LocalWebsiteService) Start(websites []Website) error {
 			// Register the API proxy handler for this website
 			mux.HandleFunc("/api/{name}/", l.createAPIPathHandler())
 
+			// Register the WebSocket proxy handler for this website
+			mux.HandleFunc("/ws/{name}", l.createWebsocketPathHandler)
+
 			// Create the SPA handler for this website
 			spa := staticSiteHandler{
 				website:    website,
@@ -289,6 +360,9 @@ func (l *LocalWebsiteService) Start(websites []Website) error {
 		// Register the API proxy handler
 		mux.HandleFunc("/api/{name}/", l.createAPIPathHandler())
 
+		// Register the WebSocket proxy handler for this website
+		mux.HandleFunc("/ws/{name}", l.createWebsocketPathHandler)
+
 		// Register the SPA handler for each website
 		for i := range websites {
 			website := &websites[i]
@@ -325,11 +399,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,
 	}
 }