Implemented Origin checks for websockets

Author: cmaglie

cmd/arduino-app-cli/daemon/daemon.go

@@ -69,6 +69,32 @@ func NewDaemonCmd(cfg config.Configuration, version string) *cobra.Command {
 func httpHandler(ctx context.Context, cfg config.Configuration, daemonPort, version string) {
 	slog.Info("Starting HTTP server", slog.String("address", ":"+daemonPort))
 
+	corsConfig := cors.Config{
+		Origins: []string{
+			"wails://wails",
+			"wails://wails.localhost:*",
+			"http://wails.localhost",
+			"http://wails.localhost:*",
+			"http://localhost:*",
+			"https://localhost:*",
+		},
+		Methods: []string{
+			http.MethodGet,
+			http.MethodPost,
+			http.MethodPut,
+			http.MethodOptions,
+			http.MethodDelete,
+			http.MethodPatch,
+		},
+		RequestHeaders: []string{
+			"Accept",
+			"Authorization",
+			"Content-Type",
+		},
+		MaxAgeInSeconds: 86400,
+		ResponseHeaders: []string{},
+	}
+
 	apiSrv := api.NewHTTPRouter(
 		servicelocator.GetDockerClient(),
 		version,
@@ -83,40 +109,21 @@ func httpHandler(ctx context.Context, cfg config.Configuration, daemonPort, vers
 		servicelocator.GetBrickService(),
 		servicelocator.GetAppIDProvider(),
 		cfg,
+		corsConfig.Origins,
 	)
 
-	corsMiddlware, err := cors.NewMiddleware(
-		cors.Config{
-			Origins: []string{
-				"wails://wails", "wails://wails.localhost:*",
-				"http://wails.localhost", "http://wails.localhost:*",
-				"http://localhost:*", "https://localhost:*",
-			},
-			Methods: []string{
-				http.MethodGet,
-				http.MethodPost,
-				http.MethodPut,
-				http.MethodOptions,
-				http.MethodDelete,
-				http.MethodPatch,
-			},
-			RequestHeaders: []string{
-				"Accept",
-				"Authorization",
-				"Content-Type",
-			},
-			MaxAgeInSeconds: 86400,
-			ResponseHeaders: []string{},
-		},
-	)
+	// Wrap the API server with CORS middleware
+	corsMiddlware, err := cors.NewMiddleware(corsConfig)
 	if err != nil {
 		panic(err)
 	}
+	apiSrv = corsMiddlware.Wrap(apiSrv)
 
+	// Start the HTTP server
 	address := "127.0.0.1:" + daemonPort
 	httpSrv := http.Server{
 		Addr:              address,
-		Handler:           httprecover.RecoverPanic(corsMiddlware.Wrap(apiSrv)),
+		Handler:           httprecover.RecoverPanic(apiSrv),
 		ReadHeaderTimeout: 60 * time.Second,
 	}
 	go func() {

internal/api/api.go

@@ -33,6 +33,7 @@ func NewHTTPRouter(
 	brickService *bricks.Service,
 	idProvider *app.IDProvider,
 	cfg config.Configuration,
+	allowedOrigins []string,
 ) http.Handler {
 	mux := http.NewServeMux()
 	mux.Handle("GET /debug/", http.DefaultServeMux) // pprof endpoints
@@ -70,7 +71,7 @@ func NewHTTPRouter(
 
 	mux.Handle("GET /v1/docs/", http.StripPrefix("/v1/docs/", handlers.DocsServer(docsFS)))
 
-	mux.Handle("GET /v1/monitor/ws", handlers.HandleMonitorWS())
+	mux.Handle("GET /v1/monitor/ws", handlers.HandleMonitorWS(allowedOrigins))
 
 	return mux
 }

internal/api/handlers/monitor.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"net"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -15,11 +16,6 @@ import (
 	"github.com/arduino/arduino-app-cli/pkg/render"
 )
 
-var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-}
-
 func monitorStream(mon net.Conn, ws *websocket.Conn) {
 	logWebsocketError := func(msg string, err error) {
 		// Do not log simple close or interruption errors
@@ -72,7 +68,32 @@ func monitorStream(mon net.Conn, ws *websocket.Conn) {
 	}()
 }
 
-func HandleMonitorWS() http.HandlerFunc {
+func checkOrigin(origin string, allowedOrigins []string) bool {
+	for _, allowed := range allowedOrigins {
+		if strings.HasSuffix(allowed, "*") {
+			// String ends with *, match the prefix
+			if strings.HasPrefix(origin, strings.TrimSuffix(allowed, "*")) {
+				return true
+			}
+		} else {
+			// Exact match
+			if allowed == origin {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func HandleMonitorWS(allowedOrigins []string) http.HandlerFunc {
+	upgrader := websocket.Upgrader{
+		ReadBufferSize:  1024,
+		WriteBufferSize: 1024,
+		CheckOrigin: func(r *http.Request) bool {
+			return checkOrigin(r.Header.Get("Origin"), allowedOrigins)
+		},
+	}
+
 	return func(w http.ResponseWriter, r *http.Request) {
 		// Connect to monitor
 		mon, err := net.DialTimeout("tcp", "127.0.0.1:7500", time.Second)
@@ -88,7 +109,8 @@ func HandleMonitorWS() http.HandlerFunc {
 			// Remember to close monitor connection if websocket upgrade fails.
 			mon.Close()
 
-			render.EncodeResponse(w, http.StatusInternalServerError, map[string]string{"error": "Failed to upgrade connection"})
+			slog.Error("Failed to upgrade connection", slog.String("error", err.Error()))
+			render.EncodeResponse(w, http.StatusInternalServerError, map[string]string{"error": "Failed to upgrade connection: " + err.Error()})
 			return
 		}