tsconnect: harden JS bridge callbacks

Contain synchronous and asynchronous JS failures, make Drive cancellation fence response callbacks, and keep netmap snapshots current across peer changes and watcher restarts.

Co-Authored-By: gpt-5.6-sol <noreply@openai.com>

cmd/tsconnect/wasm/drive.go

@@ -30,6 +30,73 @@ type jsFileSystemForRemote struct {
 	fn js.Value
 }
 
+type driveResponse struct {
+	mu                   sync.Mutex
+	w                    http.ResponseWriter
+	doneCh               chan error
+	live                 bool
+	responseStarted      bool
+	testBeforeWriteCheck func()
+}
+
+func (r *driveResponse) isLive() bool {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.live
+}
+
+func (r *driveResponse) writeHead(args []js.Value) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if !r.live || len(args) < 1 {
+		return
+	}
+	status := args[0].Int()
+	if len(args) > 1 && !args[1].IsUndefined() && !args[1].IsNull() {
+		for k, vs := range jsHeadersToGo(args[1]) {
+			for _, v := range vs {
+				r.w.Header().Add(k, v)
+			}
+		}
+	}
+	r.w.WriteHeader(status)
+	r.responseStarted = true
+}
+
+func (r *driveResponse) write(args []js.Value) {
+	if r.testBeforeWriteCheck != nil {
+		r.testBeforeWriteCheck()
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if !r.live || len(args) < 1 {
+		return
+	}
+	data := args[0]
+	buf := make([]byte, data.Get("length").Int())
+	js.CopyBytesToGo(buf, data)
+	r.responseStarted = true
+	if _, err := r.w.Write(buf); err != nil {
+		select {
+		case r.doneCh <- err:
+		default:
+		}
+		return
+	}
+	if f, ok := r.w.(http.Flusher); ok {
+		f.Flush()
+	}
+}
+
+func (r *driveResponse) finish(resultErr, contextErr error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if resultErr != nil && contextErr == nil && !r.responseStarted {
+		http.Error(r.w, "drive handler failed", http.StatusInternalServerError)
+	}
+	r.live = false
+}
+
 func (fs *jsFileSystemForRemote) setHandler(fn js.Value) {
 	fs.mu.Lock()
 	fs.fn = fn
@@ -50,7 +117,8 @@ func (fs *jsFileSystemForRemote) Close() error { return nil }
 // response body back via write()/end() callbacks, so no full-body buffering
 // occurs regardless of file size.
 //
-// The call blocks until JS calls end() (or a write error occurs).
+// The call blocks until JS calls end(), a write or handler error occurs, or
+// the request context is cancelled.
 func (fs *jsFileSystemForRemote) ServeHTTPWithPerms(
 	perms drive.Permissions, w http.ResponseWriter, r *http.Request,
 ) {
@@ -58,14 +126,21 @@ func (fs *jsFileSystemForRemote) ServeHTTPWithPerms(
 	fn := fs.fn
 	fs.mu.RUnlock()
 
-	if fn.IsUndefined() || fn.IsNull() {
+	if fn.Type() != js.TypeFunction {
 		http.NotFound(w, r)
 		return
 	}
 
+	response := &driveResponse{w: w, doneCh: make(chan error, 1), live: true}
+
 	// readBodyChunk is exposed to JS as req.readBodyChunk().
 	// Each call returns a Promise<Uint8Array|null>: null signals EOF.
 	readBodyChunk := js.FuncOf(func(_ js.Value, _ []js.Value) any {
+		if !response.isLive() {
+			return makePromise(func() (any, error) {
+				return nil, errors.New("drive request is closed")
+			})
+		}
 		return makePromise(func() (any, error) {
 			buf := make([]byte, 65536)
 			n, err := r.Body.Read(buf)
@@ -81,51 +156,33 @@ func (fs *jsFileSystemForRemote) ServeHTTPWithPerms(
 		})
 	})
 
-	// doneCh receives nil when JS calls end(), or a write error if Write fails.
-	doneCh := make(chan error, 1)
-
 	// writeHead sets response headers and status code. Must be called before write().
 	writeHead := js.FuncOf(func(_ js.Value, args []js.Value) any {
-		if len(args) < 1 {
-			return nil
-		}
-		status := args[0].Int()
-		if len(args) > 1 && !args[1].IsUndefined() && !args[1].IsNull() {
-			for k, vs := range jsHeadersToGo(args[1]) {
-				for _, v := range vs {
-					w.Header().Add(k, v)
-				}
-			}
-		}
-		w.WriteHeader(status)
+		response.writeHead(args)
 		return nil
 	})
 
 	// write streams a single response body chunk to the client.
 	write := js.FuncOf(func(_ js.Value, args []js.Value) any {
-		if len(args) < 1 {
-			return nil
-		}
-		data := args[0]
-		buf := make([]byte, data.Get("length").Int())
-		js.CopyBytesToGo(buf, data)
-		if _, werr := w.Write(buf); werr != nil {
-			select {
-			case doneCh <- werr:
-			default:
-			}
-			return nil
-		}
-		if f, ok := w.(http.Flusher); ok {
-			f.Flush()
-		}
+		response.write(args)
 		return nil
 	})
 
 	// end signals that the response is complete.
 	end := js.FuncOf(func(_ js.Value, _ []js.Value) any {
 		select {
-		case doneCh <- nil:
+		case response.doneCh <- nil:
+		default:
+		}
+		return nil
+	})
+	rejected := js.FuncOf(func(_ js.Value, args []js.Value) any {
+		err := errors.New("JavaScript drive handler rejected")
+		if len(args) > 0 && args[0].Type() == js.TypeString {
+			err = fmt.Errorf("JavaScript drive handler rejected: %s", args[0].String())
+		}
+		select {
+		case response.doneCh <- err:
 		default:
 		}
 		return nil
@@ -136,6 +193,7 @@ func (fs *jsFileSystemForRemote) ServeHTTPWithPerms(
 		writeHead.Release()
 		write.Release()
 		end.Release()
+		rejected.Release()
 	}()
 
 	jsReq := map[string]any{
@@ -151,11 +209,29 @@ func (fs *jsFileSystemForRemote) ServeHTTPWithPerms(
 		"end":       end,
 	}
 
-	fn.Invoke(jsReq, jsRes, drivePermsToJS(perms))
+	result, handlerErr := callJSFunction(fn, jsReq, jsRes, drivePermsToJS(perms))
+	if handlerErr == nil {
+		if hasThen, err := hasJSFunctionProperty(result, "then"); err != nil {
+			handlerErr = err
+		} else if hasThen {
+			_, handlerErr = callJSMethod(result, "then", js.Undefined(), rejected)
+		}
+	}
+	if handlerErr != nil {
+		select {
+		case response.doneCh <- handlerErr:
+		default:
+		}
+	}
 
-	// Block this goroutine until JS calls end() or a write error occurs.
-	// The Go WASM scheduler yields back to JS while we wait.
-	<-doneCh
+	var resultErr error
+	select {
+	case resultErr = <-response.doneCh:
+	case <-r.Context().Done():
+		resultErr = r.Context().Err()
+	}
+
+	response.finish(resultErr, r.Context().Err())
 }
 
 // drivePermsToJS converts drive.Permissions to a plain JS-friendly object.