From d8d48e42690207ad4280a95dc296a38c8e46938c Mon Sep 17 00:00:00 2001
From: realrasengan <andrew@joseon.com>
Date: Mon, 30 Mar 2026 14:35:40 -0700
Subject: [PATCH] fix WASM: self-chaining JS reader, no goroutine leak,
 per-write lock

---
 cmd/mosh-wasm/webtransport.go | 129 +++++++++++++++++-----------------
 1 file changed, 63 insertions(+), 66 deletions(-)

diff --git a/cmd/mosh-wasm/webtransport.go b/cmd/mosh-wasm/webtransport.go
index f36ea1f..51fc7c2 100644
--- a/cmd/mosh-wasm/webtransport.go
+++ b/cmd/mosh-wasm/webtransport.go
@@ -10,16 +10,11 @@ import (
 )
 
 // wtConn wraps a browser WebTransport connection as a mosh.Conn.
-// Uses WebTransport datagrams for unreliable, unordered delivery —
-// matching UDP semantics for the mosh protocol.
 type wtConn struct {
-	transport js.Value // WebTransport instance
-	reader    js.Value // datagrams.readable reader
-	writer    js.Value // datagrams.writable writer
-
-	incoming chan []byte
-	done     chan struct{}
-	once     sync.Once
+	transport js.Value
+	incoming  chan []byte
+	done      chan struct{}
+	once      sync.Once
 
 	mu       sync.Mutex
 	deadline time.Time
@@ -28,85 +23,87 @@ type wtConn struct {
 func dialWebTransport(url string) (*wtConn, error) {
 	wt := js.Global().Get("WebTransport").New(url)
 
-	// Wait for .ready promise
 	readyCh := make(chan error, 1)
-	wt.Get("ready").Call("then",
-		js.FuncOf(func(this js.Value, args []js.Value) interface{} {
-			readyCh <- nil
-			return nil
-		}),
-		js.FuncOf(func(this js.Value, args []js.Value) interface{} {
-			msg := "WebTransport connection failed"
-			if len(args) > 0 {
-				msg = args[0].Get("message").String()
+	onReady := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		readyCh <- nil
+		return nil
+	})
+	onFail := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		msg := "WebTransport connection failed"
+		if len(args) > 0 && !args[0].IsUndefined() && !args[0].IsNull() {
+			if m := args[0].Get("message"); !m.IsUndefined() {
+				msg = m.String()
 			}
-			readyCh <- errors.New(msg)
-			return nil
-		}),
-	)
+		}
+		readyCh <- errors.New(msg)
+		return nil
+	})
+	wt.Get("ready").Call("then", onReady, onFail)
 
 	select {
 	case err := <-readyCh:
 		if err != nil {
+			onReady.Release()
+			onFail.Release()
 			return nil, err
 		}
 	case <-time.After(10 * time.Second):
+		onReady.Release()
+		onFail.Release()
 		wt.Call("close")
 		return nil, errors.New("WebTransport connect timeout")
 	}
-
-	datagrams := wt.Get("datagrams")
-	reader := datagrams.Get("readable").Call("getReader")
-	writer := datagrams.Get("writable").Call("getWriter")
+	onReady.Release()
+	onFail.Release()
 
 	c := &wtConn{
 		transport: wt,
-		reader:    reader,
-		writer:    writer,
 		incoming:  make(chan []byte, 256),
 		done:      make(chan struct{}),
 	}
 
-	// Start reading datagrams in background
-	go c.readLoop()
+	// Start reading datagrams via a persistent JS callback.
+	c.startReader()
 
 	return c, nil
 }
 
-func (c *wtConn) readLoop() {
-	for {
-		ch := make(chan js.Value, 1)
-		c.reader.Call("read").Call("then",
-			js.FuncOf(func(this js.Value, args []js.Value) interface{} {
-				ch <- args[0]
-				return nil
-			}),
-			js.FuncOf(func(this js.Value, args []js.Value) interface{} {
-				close(ch)
-				return nil
-			}),
-		)
+// startReader sets up a self-chaining read loop using persistent callbacks.
+func (c *wtConn) startReader() {
+	reader := c.transport.Get("datagrams").Get("readable").Call("getReader")
 
-		select {
-		case <-c.done:
-			return
-		case result, ok := <-ch:
-			if !ok {
-				return
-			}
-			if result.Get("done").Bool() {
-				return
-			}
-			value := result.Get("value")
-			buf := make([]byte, value.Get("byteLength").Int())
-			js.CopyBytesToGo(buf, js.Global().Get("Uint8Array").New(value))
-			select {
-			case c.incoming <- buf:
-			default:
-				// drop if channel full
-			}
+	var onData, onErr js.Func
+	var readNext js.Func
+
+	readNext = js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		reader.Call("read").Call("then", onData, onErr)
+		return nil
+	})
+
+	onData = js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		result := args[0]
+		if result.Get("done").Bool() {
+			return nil
 		}
-	}
+		value := result.Get("value")
+		buf := make([]byte, value.Get("byteLength").Int())
+		js.CopyBytesToGo(buf, js.Global().Get("Uint8Array").New(value))
+		select {
+		case c.incoming <- buf:
+		default:
+		}
+		// Chain next read.
+		readNext.Invoke()
+		return nil
+	})
+
+	onErr = js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		// Stream error or closed — stop reading.
+		return nil
+	})
+
+	// Start first read.
+	readNext.Invoke()
 }
 
 func (c *wtConn) Read(b []byte) (int, error) {
@@ -135,9 +132,11 @@ func (c *wtConn) Read(b []byte) (int, error) {
 }
 
 func (c *wtConn) Write(b []byte) (int, error) {
+	writer := c.transport.Get("datagrams").Get("writable").Call("getWriter")
 	arr := js.Global().Get("Uint8Array").New(len(b))
 	js.CopyBytesToJS(arr, b)
-	c.writer.Call("write", arr)
+	writer.Call("write", arr)
+	writer.Call("releaseLock")
 	return len(b), nil
 }
 
@@ -151,8 +150,6 @@ func (c *wtConn) SetReadDeadline(t time.Time) error {
 func (c *wtConn) Close() error {
 	c.once.Do(func() {
 		close(c.done)
-		c.reader.Call("cancel")
-		c.writer.Call("close")
 		c.transport.Call("close")
 	})
 	return nil