fix: harden local proxy and import boundaries

internal/vncclient/vncclient.go

@@ -9,6 +9,8 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/url"
+	"strings"
 	"sync"
 	"time"
 
@@ -53,11 +55,13 @@ type Config struct {
 }
 
 var upgrader = websocket.Upgrader{
-	CheckOrigin:     func(r *http.Request) bool { return true },
+	CheckOrigin:     allowProxyOrigin,
 	ReadBufferSize:  65536,
 	WriteBufferSize: 65536,
 }
 
+const maxWebSocketMessageBytes = 8 << 20
+
 // NewClient creates a new VNC proxy client.
 func NewClient(cfg Config) *Client {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -186,6 +190,7 @@ func (c *Client) startProxy(vncConn net.Conn) (string, error) {
 			return
 		}
 		defer ws.Close()
+		ws.SetReadLimit(maxWebSocketMessageBytes)
 
 		// Bridge WebSocket ↔ VNC TCP in both directions
 		done := make(chan struct{}, 2)
@@ -252,6 +257,34 @@ func (c *Client) startProxy(vncConn net.Conn) (string, error) {
 	return wsURL, nil
 }
 
+func allowProxyOrigin(r *http.Request) bool {
+	origin := r.Header.Get("Origin")
+	if origin == "" {
+		return true
+	}
+	u, err := url.Parse(origin)
+	if err != nil {
+		return false
+	}
+	switch strings.ToLower(u.Scheme) {
+	case "wails":
+		return true
+	case "http", "https":
+		return isLoopbackHost(u.Hostname())
+	default:
+		return false
+	}
+}
+
+func isLoopbackHost(host string) bool {
+	h := strings.ToLower(strings.TrimSpace(host))
+	if h == "localhost" {
+		return true
+	}
+	ip := net.ParseIP(h)
+	return ip != nil && ip.IsLoopback()
+}
+
 func newProxyToken() (string, error) {
 	token := make([]byte, 16)
 	if _, err := rand.Read(token); err != nil {

internal/vncclient/vncclient_test.go

@@ -1,6 +1,9 @@
 package vncclient
 
-import "testing"
+import (
+	"net/http"
+	"testing"
+)
 
 func TestNewProxyTokenIsRandomHex(t *testing.T) {
 	first, err := newProxyToken()
@@ -24,3 +27,35 @@ func TestNewProxyTokenIsRandomHex(t *testing.T) {
 		}
 	}
 }
+
+func TestAllowProxyOrigin(t *testing.T) {
+	tests := []struct {
+		name   string
+		origin string
+		want   bool
+	}{
+		{name: "missing origin", origin: "", want: true},
+		{name: "wails origin", origin: "wails://wails", want: true},
+		{name: "localhost http", origin: "http://localhost:34115", want: true},
+		{name: "loopback ipv4", origin: "http://127.0.0.1:34115", want: true},
+		{name: "loopback ipv6", origin: "http://[::1]:34115", want: true},
+		{name: "remote http", origin: "http://example.com", want: false},
+		{name: "lookalike host", origin: "http://127.0.0.1.example.com", want: false},
+		{name: "bad origin", origin: "://bad", want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req, err := http.NewRequest(http.MethodGet, "http://127.0.0.1/ws", nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if tt.origin != "" {
+				req.Header.Set("Origin", tt.origin)
+			}
+			if got := allowProxyOrigin(req); got != tt.want {
+				t.Fatalf("allowProxyOrigin(%q) = %v, want %v", tt.origin, got, tt.want)
+			}
+		})
+	}
+}