WASM: support WebTransport reconnection with session token

cmd/mosh-wasm/main.go

@@ -136,6 +136,23 @@ func (s *session) jsObject() js.Value {
 		return string(out)
 	}))
 
+	// reconnect() — swap WebTransport connection using session token.
+	obj.Set("reconnect", js.FuncOf(func(this js.Value, args []js.Value) interface{} {
+		handler := js.FuncOf(func(this js.Value, pargs []js.Value) interface{} {
+			resolve := pargs[0]
+			rejectFn := pargs[1]
+			go func() {
+				if err := s.conn.Reconnect(); err != nil {
+					rejectFn.Invoke(err.Error())
+					return
+				}
+				resolve.Invoke(true)
+			}()
+			return nil
+		})
+		return js.Global().Get("Promise").New(handler)
+	}))
+
 	obj.Set("close", js.FuncOf(func(this js.Value, args []js.Value) interface{} {
 		s.mu.Lock()
 		s.closed = true

cmd/mosh-wasm/webtransport.go

@@ -4,21 +4,27 @@ package main
 
 import (
 	"errors"
+	"strings"
 	"sync"
 	"syscall/js"
 	"time"
 )
 
 // wtConn wraps a browser WebTransport connection as a mosh.Conn.
+// Supports transparent reconnection: when the underlying WebTransport
+// dies, reads return timeouts while the worker re-dials. Once a new
+// connection is swapped in via Reconnect(), reads resume.
 type wtConn struct {
+	mu        sync.Mutex
 	transport js.Value
 	writer    js.Value // persistent writable stream writer
 	incoming  chan []byte
 	done      chan struct{}
 	once      sync.Once
+	deadline  time.Time
 
-	mu       sync.Mutex
-	deadline time.Time
+	url   string
+	token string // session token from relay, used for reconnect
 }
 
 func dialWebTransport(url string) (*wtConn, error) {
@@ -64,6 +70,7 @@ func dialWebTransport(url string) (*wtConn, error) {
 		writer:    writer,
 		incoming:  make(chan []byte, 256),
 		done:      make(chan struct{}),
+		url:       url,
 	}
 
 	// Start reading datagrams via a persistent JS callback.
@@ -72,9 +79,94 @@ func dialWebTransport(url string) (*wtConn, error) {
 	return c, nil
 }
 
+// Reconnect dials a new WebTransport connection using the session token
+// and swaps it into this conn. The mosh client keeps its SSP state and
+// resumes transparently.
+func (c *wtConn) Reconnect() error {
+	c.mu.Lock()
+	token := c.token
+	url := c.url
+	c.mu.Unlock()
+
+	if token == "" {
+		return errors.New("no session token")
+	}
+
+	// Add token to URL for reconnect.
+	sep := "&"
+	if !strings.Contains(url, "?") {
+		sep = "?"
+	}
+	reconnURL := url + sep + "token=" + token
+
+	wt := js.Global().Get("WebTransport").New(reconnURL)
+
+	readyCh := make(chan error, 1)
+	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 := "reconnect 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
+	})
+	wt.Get("ready").Call("then", onReady, onFail)
+
+	select {
+	case err := <-readyCh:
+		if err != nil {
+			onReady.Release()
+			onFail.Release()
+			return err
+		}
+	case <-time.After(10 * time.Second):
+		onReady.Release()
+		onFail.Release()
+		wt.Call("close")
+		return errors.New("reconnect timeout")
+	}
+	onReady.Release()
+	onFail.Release()
+
+	writer := wt.Get("datagrams").Get("writable").Call("getWriter")
+
+	// Swap in the new connection.
+	c.mu.Lock()
+	oldTransport := c.transport
+	c.transport = wt
+	c.writer = writer
+	c.mu.Unlock()
+
+	// Close old transport (best-effort).
+	if !oldTransport.IsUndefined() && !oldTransport.IsNull() {
+		defer func() { recover() }()
+		oldTransport.Call("close")
+	}
+
+	// Start reading from the new connection.
+	c.startReader()
+
+	return nil
+}
+
+// Token returns the session token received from the relay.
+func (c *wtConn) Token() string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.token
+}
+
 // startReader sets up a self-chaining read loop using persistent callbacks.
 func (c *wtConn) startReader() {
+	c.mu.Lock()
 	reader := c.transport.Get("datagrams").Get("readable").Call("getReader")
+	c.mu.Unlock()
 
 	var onData, onErr js.Func
 	var readNext js.Func
@@ -92,10 +184,19 @@ func (c *wtConn) startReader() {
 		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:
+
+		// Check for TOKEN: prefix (session token from relay).
+		if len(buf) > 6 && string(buf[:6]) == "TOKEN:" {
+			c.mu.Lock()
+			c.token = string(buf[6:])
+			c.mu.Unlock()
+		} else {
+			select {
+			case c.incoming <- buf:
+			default:
+			}
 		}
+
 		// Defer next read to next event loop tick so Go goroutines can process.
 		js.Global().Call("setTimeout", readNext, 0)
 		return nil
@@ -103,6 +204,7 @@ func (c *wtConn) startReader() {
 
 	onErr = js.FuncOf(func(this js.Value, args []js.Value) interface{} {
 		// Stream error or closed — stop reading.
+		// The reconnect logic in the worker will handle this.
 		return nil
 	})
 
@@ -136,9 +238,17 @@ func (c *wtConn) Read(b []byte) (int, error) {
 }
 
 func (c *wtConn) Write(b []byte) (int, error) {
+	c.mu.Lock()
+	writer := c.writer
+	c.mu.Unlock()
+
+	if writer.IsUndefined() || writer.IsNull() {
+		return 0, errors.New("not connected")
+	}
+
 	arr := js.Global().Get("Uint8Array").New(len(b))
 	js.CopyBytesToJS(arr, b)
-	c.writer.Call("write", arr)
+	writer.Call("write", arr)
 	return len(b), nil
 }
 
@@ -152,7 +262,12 @@ func (c *wtConn) SetReadDeadline(t time.Time) error {
 func (c *wtConn) Close() error {
 	c.once.Do(func() {
 		close(c.done)
-		c.transport.Call("close")
+		c.mu.Lock()
+		t := c.transport
+		c.mu.Unlock()
+		if !t.IsUndefined() && !t.IsNull() {
+			t.Call("close")
+		}
 	})
 	return nil
 }