让插件支持监听ws

internal/plugin/route_websocket.go

@@ -0,0 +1,507 @@
+package plugin
+
+import (
+	"errors"
+	"fmt"
+	"math"
+	"net"
+	"sync"
+	"time"
+	"unicode/utf8"
+
+	"github.com/Calcium-Ion/moejs"
+	"github.com/Calcium-Ion/moejs/engine"
+	"github.com/gorilla/websocket"
+	"github.com/komari-monitor/komari/pkg/jsruntime"
+	"github.com/komari-monitor/komari/pkg/jsruntime/native"
+	"github.com/komari-monitor/komari/web/connection"
+)
+
+const (
+	routeWSOpen            = 1
+	routeWSClosing         = 2
+	routeWSClosed          = 3
+	routeWSMaxQueuedFrames = 64
+	routeWSMaxQueuedBytes  = 16 << 20
+	routeWSWriteTimeout    = 5 * time.Second
+	routeWSCloseTimeout    = 5 * time.Second
+)
+
+type routeWSFrame struct {
+	kind int
+	data []byte
+}
+
+type routeWSCloseEvent struct {
+	code    int
+	reason  string
+	clean   bool
+	message string
+}
+
+// routeWebSocket is registered before hijacking: runtime shutdown must be able
+// to close the network connection even while UpgradeSafeConn waits for a
+// synchronous wsConnect hook. No SafeConn operation or hook runs on the VM loop.
+// mu protects transport state only; jsObject belongs exclusively to the loop.
+type routeWebSocket struct {
+	host *jsruntime.Host
+	mu   sync.Mutex
+
+	resourceID uint64
+	netConn    net.Conn
+	conn       *connection.SafeConn
+	raw        *websocket.Conn
+	state      int
+	started    bool
+	finished   bool
+	notified   bool
+	result     routeWSCloseEvent
+	closeTimer *time.Timer
+
+	queue       []routeWSFrame
+	queuedCount int // includes the writer's in-flight frame
+	queuedBytes int
+	wake        chan struct{}
+	done        chan struct{}
+
+	closeRequested bool
+	closeCode      int
+	closeReason    string
+	closeWritten   bool
+	peerClosed     bool
+	peerCode       int
+	peerReason     string
+	pong           string
+	pongPending    bool
+
+	jsObject *engine.Object
+}
+
+func newRouteWebSocket(host *jsruntime.Host) *routeWebSocket {
+	ws := &routeWebSocket{
+		host: host, state: routeWSClosed,
+		wake: make(chan struct{}, 1), done: make(chan struct{}),
+	}
+	// AddResource may call abort synchronously, or shutdown may race with its
+	// return. Never publish the ID without checking whether teardown already ran.
+	id := host.AddResource(ws.abort)
+	ws.mu.Lock()
+	ws.resourceID = id
+	finished := ws.finished
+	ws.mu.Unlock()
+	if finished {
+		host.RemoveResource(id)
+	}
+	return ws
+}
+
+// captureConn is called by the HTTP writer's Hijack method, before the
+// upgrader can enter any hooks. It also handles shutdown winning that race.
+func (ws *routeWebSocket) captureConn(conn net.Conn) {
+	ws.mu.Lock()
+	finished := ws.finished
+	if !finished {
+		ws.netConn = conn
+	}
+	ws.mu.Unlock()
+	if finished && conn != nil {
+		_ = conn.Close()
+	}
+}
+
+// attach runs on the HTTP pump, never the VM loop: GetConn takes SafeConn's
+// write mutex. The underlying Gorilla connection is private to these pumps.
+func (ws *routeWebSocket) attach(conn *connection.SafeConn) bool {
+	if conn == nil {
+		return false
+	}
+	raw := conn.GetConn()
+	raw.SetReadLimit(maxWSFrameBytes)
+	raw.SetCloseHandler(func(code int, reason string) error {
+		ws.mu.Lock()
+		if !ws.finished {
+			ws.peerClosed, ws.peerCode, ws.peerReason = true, code, reason
+			// The peer cannot consume further data. Drop queued frames and
+			// prioritize the close reply (the in-flight writer stays counted).
+			for _, frame := range ws.queue {
+				ws.queuedCount--
+				ws.queuedBytes -= len(frame.data)
+			}
+			ws.queue = nil
+			ws.pong, ws.pongPending = "", false
+			if !ws.closeRequested {
+				ws.requestCloseLocked(code, reason)
+			}
+		}
+		ws.mu.Unlock()
+		ws.signal()
+		return nil // the writer, not the reader, sends the close reply
+	})
+	raw.SetPingHandler(func(data string) error {
+		ws.mu.Lock()
+		if !ws.finished {
+			// A ping flood must not allocate an unbounded control queue. RFC 6455
+			// allows replying only to the most recently processed ping.
+			ws.pong, ws.pongPending = data, true
+		}
+		ws.mu.Unlock()
+		ws.signal()
+		return nil
+	})
+
+	ws.mu.Lock()
+	if ws.finished || ws.conn != nil {
+		ws.mu.Unlock()
+		_ = raw.Close()
+		go func() { _ = conn.Close() }()
+		return false
+	}
+	ws.conn, ws.raw, ws.state = conn, raw, routeWSOpen
+	if ws.netConn == nil {
+		ws.netConn = raw.UnderlyingConn()
+	}
+	ws.mu.Unlock()
+	return true
+}
+
+// abort is a non-waiting resource callback. Close the captured network socket
+// first to unblock I/O; SafeConn.Close may wait for same-plugin hooks, so it
+// must happen separately, after the socket is already closed.
+func (ws *routeWebSocket) abort() {
+	ws.finish(routeWSCloseEvent{code: websocket.CloseAbnormalClosure})
+}
+
+func (ws *routeWebSocket) finish(result routeWSCloseEvent) {
+	ws.mu.Lock()
+	if ws.finished {
+		ws.mu.Unlock()
+		return
+	}
+	ws.finished, ws.state, ws.result = true, routeWSClosed, result
+	ws.queue = nil
+	ws.queuedCount, ws.queuedBytes = 0, 0
+	if ws.closeTimer != nil {
+		ws.closeTimer.Stop()
+	}
+	netConn, conn, id := ws.netConn, ws.conn, ws.resourceID
+	notify := ws.started && !ws.notified
+	if notify {
+		ws.notified = true
+	}
+	ws.mu.Unlock()
+
+	if netConn != nil {
+		_ = netConn.Close()
+	}
+	ws.host.RemoveResource(id)
+	close(ws.done)
+	if conn != nil {
+		go func() { _ = conn.Close() }()
+	}
+	if notify {
+		ws.notifyClose(result)
+	}
+}
+
+// start is called on a later loop turn than promise resolution. This barrier
+// lets .then/await reactions install handlers before any data or close event.
+func (ws *routeWebSocket) start() {
+	ws.mu.Lock()
+	if ws.started {
+		ws.mu.Unlock()
+		return
+	}
+	ws.started = true
+	if ws.finished {
+		notify, result := !ws.notified, ws.result
+		ws.notified = true
+		ws.mu.Unlock()
+		if notify {
+			ws.notifyClose(result)
+		}
+		return
+	}
+	attached := ws.conn != nil
+	ws.mu.Unlock()
+	if !attached || !ws.host.ResourcesOpen() {
+		ws.abort()
+		return
+	}
+	go ws.writePump()
+	go ws.readPump()
+}
+
+func (ws *routeWebSocket) signal() {
+	select {
+	case ws.wake <- struct{}{}:
+	default:
+	}
+}
+
+func (ws *routeWebSocket) requestCloseLocked(code int, reason string) {
+	ws.state, ws.closeRequested = routeWSClosing, true
+	ws.closeCode, ws.closeReason = code, reason
+	ws.closeTimer = time.AfterFunc(routeWSCloseTimeout, func() {
+		ws.fail(errors.New("WebSocket close handshake timed out"))
+	})
+}
+
+func (ws *routeWebSocket) fail(err error) {
+	ws.finish(routeWSCloseEvent{code: websocket.CloseAbnormalClosure, message: err.Error()})
+}
+
+func (ws *routeWebSocket) writePump() {
+	for {
+		ws.mu.Lock()
+		if ws.finished {
+			ws.mu.Unlock()
+			return
+		}
+		if ws.closeWritten && ws.peerClosed {
+			result := routeWSCloseEvent{code: ws.peerCode, reason: ws.peerReason, clean: true}
+			ws.mu.Unlock()
+			ws.finish(result)
+			return
+		}
+		var frame routeWSFrame
+		control, haveFrame := false, false
+		switch {
+		case ws.pongPending && !ws.closeWritten:
+			frame = routeWSFrame{kind: websocket.PongMessage, data: []byte(ws.pong)}
+			ws.pong, ws.pongPending, control, haveFrame = "", false, true, true
+		case len(ws.queue) != 0:
+			frame, haveFrame = ws.queue[0], true
+			ws.queue[0] = routeWSFrame{}
+			ws.queue = ws.queue[1:]
+		case ws.closeRequested && !ws.closeWritten:
+			frame = routeWSFrame{kind: websocket.CloseMessage, data: websocket.FormatCloseMessage(ws.closeCode, ws.closeReason)}
+			control, haveFrame = true, true
+		}
+		conn, raw := ws.conn, ws.raw
+		ws.mu.Unlock()
+		if !haveFrame {
+			select {
+			case <-ws.wake:
+			case <-ws.done:
+				return
+			}
+			continue
+		}
+
+		deadline := time.Now().Add(routeWSWriteTimeout)
+		var err error
+		if control {
+			err = conn.WriteControl(frame.kind, frame.data, deadline)
+		} else {
+			err = raw.SetWriteDeadline(deadline)
+			if err == nil {
+				err = conn.WriteMessage(frame.kind, frame.data)
+			}
+		}
+		if err != nil {
+			ws.fail(err)
+			return
+		}
+		ws.mu.Lock()
+		if !ws.finished {
+			if frame.kind == websocket.CloseMessage {
+				ws.closeWritten = true
+			} else if !control {
+				ws.queuedCount--
+				ws.queuedBytes -= len(frame.data)
+			}
+		}
+		ws.mu.Unlock()
+	}
+}
+
+func (ws *routeWebSocket) readPump() {
+	for {
+		kind, data, err := ws.conn.ReadMessage()
+		if err != nil {
+			ws.mu.Lock()
+			peerClosed := ws.peerClosed
+			ws.mu.Unlock()
+			if peerClosed {
+				ws.signal() // the writer completes the close handshake
+			} else {
+				ws.fail(err)
+			}
+			return
+		}
+		if len(data) > maxWSFrameBytes {
+			ws.fail(fmt.Errorf("WebSocket message exceeds %d bytes", maxWSFrameBytes))
+			return
+		}
+		if kind != websocket.TextMessage && kind != websocket.BinaryMessage {
+			continue
+		}
+		if kind == websocket.TextMessage && !utf8.Valid(data) {
+			ws.fail(errors.New("WebSocket text message is not valid UTF-8"))
+			return
+		}
+		ack := make(chan struct{})
+		if !ws.host.RunOnLoop(func(vm *moejs.Runtime) {
+			defer close(ack)
+			jobErr := ws.host.RunJob(vm, "plugin WebSocket.onmessage", func() error {
+				ws.mu.Lock()
+				open := ws.state == routeWSOpen
+				ws.mu.Unlock()
+				if !open || ws.jsObject == nil {
+					return nil
+				}
+				var value any = string(data)
+				if kind == websocket.BinaryMessage {
+					buffer, err := vm.Realm().NewArrayBuffer(data)
+					if err != nil {
+						return err
+					}
+					value = engine.ObjectValue(buffer)
+				}
+				event := vm.Realm().NewObject()
+				if err := native.Set(vm, event, "data", value); err != nil {
+					return err
+				}
+				return ws.callHandler(vm, "onmessage", event)
+			})
+			if jobErr != nil {
+				ws.fail(jobErr)
+			}
+		}) {
+			ws.abort()
+			return
+		}
+		// Only one inbound payload may wait for JavaScript. Cancellation must
+		// also release this reader when queued jobs are discarded on shutdown.
+		select {
+		case <-ack:
+		case <-ws.done:
+			return
+		}
+	}
+}
+
+// All methods below that touch jsObject or JS values run only on the host loop.
+func (ws *routeWebSocket) object(vm *moejs.Runtime) *engine.Object {
+	if ws.jsObject != nil {
+		return ws.jsObject
+	}
+	obj := vm.Realm().NewObject()
+	ws.jsObject = obj
+	for _, name := range []string{"onmessage", "onerror", "onclose"} {
+		_ = native.Set(vm, obj, name, moejs.Null())
+	}
+	getter := native.Bind(vm, "get readyState", 0, func(_ *moejs.Realm, _ moejs.Value, _ []moejs.Value) (moejs.Value, error) {
+		ws.mu.Lock()
+		state := ws.state
+		ws.mu.Unlock()
+		return native.From(vm, state), nil
+	})
+	desc := engine.AccessorDescriptor(getter.AsObject(), nil, 0)
+	desc.SetEnumerable(true)
+	if err := obj.DefinePropertyOrThrow(vm.Realm(), vm.Realm().KeyFromGoString("readyState"), desc); err != nil {
+		panic(err)
+	}
+	_ = native.Function(vm, obj, "send", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
+		value := moejs.Arg(args, 0)
+		kind := websocket.BinaryMessage
+		var data []byte
+		if value.IsString() {
+			kind = websocket.TextMessage
+			text := native.String(vm, value)
+			if len(text) > maxWSFrameBytes {
+				return moejs.Undefined(), native.TypeError(vm, "WebSocket message exceeds %d bytes", maxWSFrameBytes)
+			}
+			data = []byte(text)
+		} else if value.IsObject() {
+			var ok bool
+			data, ok = value.AsObject().BufferData()
+			if !ok {
+				return moejs.Undefined(), native.TypeError(vm, "WebSocket.send requires a string, ArrayBuffer or ArrayBuffer view")
+			}
+		} else {
+			return moejs.Undefined(), native.TypeError(vm, "WebSocket.send requires a string, ArrayBuffer or ArrayBuffer view")
+		}
+		if len(data) > maxWSFrameBytes {
+			return moejs.Undefined(), native.TypeError(vm, "WebSocket message exceeds %d bytes", maxWSFrameBytes)
+		}
+		ws.mu.Lock()
+		if ws.state != routeWSOpen {
+			ws.mu.Unlock()
+			return moejs.Undefined(), native.TypeError(vm, "WebSocket is not open")
+		}
+		if ws.queuedCount >= routeWSMaxQueuedFrames || len(data) > routeWSMaxQueuedBytes-ws.queuedBytes {
+			ws.mu.Unlock()
+			return moejs.Undefined(), native.TypeError(vm, "WebSocket send queue is full")
+		}
+		ws.queue = append(ws.queue, routeWSFrame{kind: kind, data: append([]byte(nil), data...)})
+		ws.queuedCount++
+		ws.queuedBytes += len(data)
+		ws.mu.Unlock()
+		ws.signal()
+		return moejs.Undefined(), nil
+	})
+	_ = native.Function(vm, obj, "close", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
+		code := websocket.CloseNormalClosure
+		if v := moejs.Arg(args, 0); !v.IsUndefined() {
+			number := native.Float(vm, v)
+			if math.IsNaN(number) || math.IsInf(number, 0) || math.Trunc(number) != number || (number != 1000 && (number < 3000 || number > 4999)) {
+				return moejs.Undefined(), native.TypeError(vm, "WebSocket close code must be 1000 or between 3000 and 4999")
+			}
+			code = int(number)
+		}
+		reason := ""
+		if v := moejs.Arg(args, 1); !v.IsUndefined() {
+			reason = native.String(vm, v)
+		}
+		if len(reason) > 123 || !utf8.ValidString(reason) {
+			return moejs.Undefined(), native.TypeError(vm, "WebSocket close reason must be valid UTF-8 and at most 123 bytes")
+		}
+		ws.mu.Lock()
+		if ws.state == routeWSOpen {
+			ws.requestCloseLocked(code, reason)
+		}
+		ws.mu.Unlock()
+		ws.signal()
+		return moejs.Undefined(), nil
+	})
+	return obj
+}
+
+func (ws *routeWebSocket) callHandler(vm *moejs.Runtime, name string, event *engine.Object) error {
+	if ws.jsObject == nil {
+		return nil
+	}
+	callback, err := vm.Get(engine.ObjectValue(ws.jsObject), name)
+	if err != nil || !engine.IsCallable(callback) {
+		return err
+	}
+	_, err = vm.CallFunction(callback, engine.ObjectValue(ws.jsObject), engine.ObjectValue(event))
+	return err
+}
+
+func (ws *routeWebSocket) notifyClose(result routeWSCloseEvent) {
+	ws.host.RunOnLoop(func(vm *moejs.Runtime) {
+		if ws.jsObject == nil {
+			return
+		}
+		defer func() { ws.jsObject = nil }()
+		if result.message != "" {
+			_ = ws.host.RunJob(vm, "plugin WebSocket.onerror", func() error {
+				event := vm.Realm().NewObject()
+				if err := native.Set(vm, event, "message", result.message); err != nil {
+					return err
+				}
+				return ws.callHandler(vm, "onerror", event)
+			})
+		}
+		// A throwing onerror must not suppress the final onclose callback.
+		_ = ws.host.RunJob(vm, "plugin WebSocket.onclose", func() error {
+			event := vm.Realm().NewObject()
+			_ = native.Set(vm, event, "code", result.code)
+			_ = native.Set(vm, event, "reason", result.reason)
+			_ = native.Set(vm, event, "wasClean", result.clean)
+			return ws.callHandler(vm, "onclose", event)
+		})
+	})
+}

internal/plugin/server.go

@@ -1,11 +1,13 @@
 package plugin
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -17,12 +19,14 @@ import (
 	"github.com/Calcium-Ion/moejs"
 	"github.com/Calcium-Ion/moejs/engine"
 	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
 	"github.com/komari-monitor/komari/database/models"
 	"github.com/komari-monitor/komari/pkg/jsruntime"
 	"github.com/komari-monitor/komari/pkg/jsruntime/httpbody"
 	"github.com/komari-monitor/komari/pkg/jsruntime/native"
 	"github.com/komari-monitor/komari/pkg/jsruntime/require"
 	"github.com/komari-monitor/komari/pkg/rpc"
+	"github.com/komari-monitor/komari/web/api"
 )
 
 // registerServerModule registers the "server" native module for one plugin
@@ -30,7 +34,8 @@ import (
 //
 //	server.route(method, path, handler)   register an HTTP route on the host
 //	                                      engine; handler receives (req, res)
-//	                                      and must call res.end() to finish
+//	                                      and finishes via res.end() or an
+//	                                      awaited req.upgradeWebSocket()
 //	server.static(path, dir, opts)        serve a folder from the plugin
 //	                                      directory at a mount path; opts.spa
 //	                                      makes unmatched paths fall back to
@@ -401,9 +406,21 @@ func (m *Manager) routeHandler(short string) gin.HandlerFunc {
 			chunks:     make(chan []byte, 4),
 			abortCh:    make(chan struct{}),
 		}
+		state.upgrades = make(chan *routeUpgrade, 1)
+		defer state.seal()
+		// Gin recycles c when the request returns, including on a timeout while
+		// this turn is still queued. Only the HTTP pump uses the original.
+		snapshot := c.Copy()
+		snapshot.Request = c.Request.Clone(c.Request.Context())
+		if principal := rpcPrincipalFromGin(c); principal != nil {
+			copy := *principal
+			copy.Roles = append([]string(nil), principal.Roles...)
+			snapshot.Set(principalContextKey, &copy)
+		}
 		queued := host.RunOnLoop(func(vm *moejs.Runtime) {
-			req := routeRequest(vm, c, body)
+			req := routeRequest(vm, snapshot, body)
 			res := routeResponseObject(vm, state)
+			bindRouteUpgrade(vm, req, res, state, host, snapshot.Request)
 			runErr := host.RunJob(vm, "plugin route "+short, func() error {
 				_, callErr := vm.CallFunction(handler, moejs.Undefined(), engine.ObjectValue(req), engine.ObjectValue(res))
 				return callErr
@@ -446,6 +463,36 @@ func (m *Manager) routeHandler(short string) gin.HandlerFunc {
 		committed := false
 		for {
 			select {
+			case upgrade := <-state.upgrades:
+				state.mu.Lock()
+				state.sealed = true
+				state.upgradeClaimed = true
+				state.mu.Unlock()
+				idle.Stop()
+				if c.Request.Context().Err() != nil || !host.ResourcesOpen() {
+					upgrade.ws.abort()
+					upgrade.complete(errors.New("WebSocket upgrade canceled"))
+					return
+				}
+				stopCancel := context.AfterFunc(c.Request.Context(), upgrade.ws.abort)
+				original := c.Writer
+				c.Writer = &routeUpgradeWriter{ResponseWriter: original, ws: upgrade.ws}
+				socket, upgradeErr := api.UpgradeSafeConn(c, func(upgrader *websocket.Upgrader) {
+					upgrader.HandshakeTimeout = host.Timeout()
+				})
+				c.Writer = original
+				stopCancel()
+				if upgradeErr != nil {
+					upgrade.ws.abort()
+					upgrade.complete(upgradeErr)
+					return
+				}
+				if !upgrade.ws.attach(socket) {
+					upgrade.complete(errors.New("WebSocket upgrade canceled"))
+					return
+				}
+				upgrade.complete(nil)
+				return
 			case chunk := <-state.chunks:
 				if !committed {
 					state.mu.Lock()
@@ -471,19 +518,18 @@ func (m *Manager) routeHandler(short string) gin.HandlerFunc {
 				}
 				idle.Reset(host.Timeout())
 			case <-state.done:
+				state.seal()
 				if !committed {
 					writeFinal(0, nil)
 				}
 				return
 			case <-c.Request.Context().Done():
-				state.aborted.Store(true)
-				close(state.abortCh)
+				state.seal()
 				return
 			case <-idle.C:
+				state.seal()
 				if state.isStreaming() {
 					// Streaming idle timeout: end the stream and notify JS
-					state.aborted.Store(true)
-					close(state.abortCh)
 					return
 				}
 				writeFinal(http.StatusGatewayTimeout, []byte("plugin route handler timed out"))
@@ -626,12 +672,100 @@ type routeResponse struct {
 	chunks    chan []byte
 	aborted   atomic.Bool
 	abortCh   chan struct{}
+
+	wrote          bool
+	sealed         bool
+	upgradeClaimed bool
+	upgrade        *routeUpgrade
+	upgrades       chan *routeUpgrade
 }
 
-func (state *routeResponse) markStreaming() {
+// routeUpgrade keeps the Promise on its original loop. No network goroutine
+// touches the resolver or socket object directly.
+type routeUpgrade struct {
+	ws       *routeWebSocket
+	complete func(error)
+}
+
+type routeUpgradeWriter struct {
+	gin.ResponseWriter
+	ws *routeWebSocket
+}
+
+func (w *routeUpgradeWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	conn, rw, err := w.ResponseWriter.Hijack()
+	if err == nil {
+		w.ws.captureConn(conn)
+	}
+	return conn, rw, err
+}
+
+func bindRouteUpgrade(vm *moejs.Runtime, req, res *engine.Object, state *routeResponse, host *jsruntime.Host, request *http.Request) {
+	_ = native.Function(vm, req, "upgradeWebSocket", func(_ *moejs.Realm, _ moejs.Value, _ []moejs.Value) (moejs.Value, error) {
+		promise, resolve, reject := vm.NewPromise()
+		fail := func(message string) (moejs.Value, error) {
+			_ = reject(native.GoError(vm, errors.New(message)))
+			return native.From(vm, promise), nil
+		}
+		streaming := native.Bool(native.Get(vm, res, "streaming"))
+		state.mu.Lock()
+		if state.sealed || state.ended || state.upgrade != nil || state.wrote || state.streaming || streaming {
+			state.mu.Unlock()
+			return fail("HTTP response is no longer available for WebSocket upgrade")
+		}
+		if request.Method != http.MethodGet || !websocket.IsWebSocketUpgrade(request) {
+			state.mu.Unlock()
+			return fail("require a GET WebSocket upgrade request")
+		}
+		ws := newRouteWebSocket(host)
+		u := &routeUpgrade{ws: ws}
+		u.complete = func(upgradeErr error) {
+			if !host.RunOnLoop(func(vm *moejs.Runtime) {
+				if !host.ResourcesOpen() {
+					ws.abort()
+					return
+				}
+				err := host.RunJob(vm, "WebSocket upgrade", func() error {
+					if upgradeErr != nil {
+						return reject(native.GoError(vm, upgradeErr))
+					}
+					return resolve(engine.ObjectValue(ws.object(vm)))
+				})
+				if err != nil || upgradeErr != nil {
+					ws.abort()
+					return
+				}
+				// Promise continuations install handlers before a later loop turn
+				// starts delivering messages (including an immediate peer send).
+				if !host.RunOnLoop(func(_ *moejs.Runtime) { ws.start() }) {
+					ws.abort()
+				}
+			}) {
+				ws.abort()
+			}
+		}
+		state.upgrade = u
+		state.mu.Unlock()
+		state.upgrades <- u
+		return native.From(vm, promise), nil
+	})
+}
+
+// seal prevents queued JavaScript from writing or upgrading a returned request.
+func (state *routeResponse) seal() {
 	state.mu.Lock()
-	state.streaming = true
+	wasSealed := state.sealed
+	state.sealed = true
+	pending := state.upgrade
+	claimed := state.upgradeClaimed
 	state.mu.Unlock()
+	if state.aborted.CompareAndSwap(false, true) {
+		close(state.abortCh)
+	}
+	if pending != nil && !claimed && !wasSealed {
+		pending.ws.abort()
+		pending.complete(errors.New("WebSocket upgrade canceled"))
+	}
 }
 
 func (state *routeResponse) isStreaming() bool {
@@ -654,7 +788,7 @@ func (state *routeResponse) pushChunk(p []byte) bool {
 func (state *routeResponse) finish(status int, body []byte) {
 	state.mu.Lock()
 	defer state.mu.Unlock()
-	if state.ended {
+	if state.ended || state.sealed || state.upgrade != nil {
 		return
 	}
 	if status != 0 {
@@ -683,13 +817,17 @@ func routeResponseObject(vm *moejs.Runtime, state *routeResponse) *engine.Object
 	})
 	_ = native.Function(vm, res, "setHeader", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
 		name, value := native.String(vm, moejs.Arg(args, 0)), moejs.Arg(args, 1)
+		var values []string
+		if err := native.ExportTo(vm, value, &values); err != nil {
+			values = []string{native.String(vm, value)}
+		}
 		state.mu.Lock()
 		defer state.mu.Unlock()
-		var values []string
-		if err := native.ExportTo(vm, value, &values); err == nil {
-			state.headers[name] = values
-		} else {
-			state.headers.Set(name, native.String(vm, value))
+		if state.upgrade != nil {
+			return moejs.Undefined(), native.TypeError(vm, "HTTP response is unavailable after WebSocket upgrade")
+		}
+		if !state.sealed {
+			state.headers[http.CanonicalHeaderKey(name)] = values
 		}
 		return engine.ObjectValue(res), nil
 	})
@@ -706,9 +844,15 @@ func routeResponseObject(vm *moejs.Runtime, state *routeResponse) *engine.Object
 		return native.From(vm, values), nil
 	})
 	_ = native.Function(vm, res, "removeHeader", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
+		name := native.String(vm, moejs.Arg(args, 0))
 		state.mu.Lock()
-		state.headers.Del(native.String(vm, moejs.Arg(args, 0)))
-		state.mu.Unlock()
+		defer state.mu.Unlock()
+		if state.upgrade != nil {
+			return moejs.Undefined(), native.TypeError(vm, "HTTP response is unavailable after WebSocket upgrade")
+		}
+		if !state.sealed {
+			state.headers.Del(name)
+		}
 		return moejs.Undefined(), nil
 	})
 	_ = native.Function(vm, res, "write", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
@@ -717,23 +861,42 @@ func routeResponseObject(vm *moejs.Runtime, state *routeResponse) *engine.Object
 		if err := native.ExportTo(vm, value, &data); err != nil {
 			data = []byte(native.String(vm, value))
 		}
-		if native.Bool(native.Get(vm, res, "streaming")) {
-			state.markStreaming()
-			_ = state.pushChunk(data)
-			return native.From(vm, true), nil
-		}
+		streaming := native.Bool(native.Get(vm, res, "streaming"))
 		state.mu.Lock()
+		if state.upgrade != nil {
+			state.mu.Unlock()
+			return moejs.Undefined(), native.TypeError(vm, "HTTP response is unavailable after WebSocket upgrade")
+		}
+		if state.sealed || state.ended {
+			state.mu.Unlock()
+			return native.From(vm, false), nil
+		}
+		state.wrote = true
+		if streaming {
+			state.streaming = true
+			state.mu.Unlock()
+			return native.From(vm, state.pushChunk(data)), nil
+		}
 		_, _ = state.body.Write(data)
 		state.mu.Unlock()
 		return native.From(vm, true), nil
 	})
 	_ = native.Function(vm, res, "end", func(_ *moejs.Realm, _ moejs.Value, args []moejs.Value) (moejs.Value, error) {
+		var tail string
+		if len(args) > 0 && !native.IsNil(args[0]) {
+			tail = native.String(vm, args[0])
+		}
+		status := int(native.Int(vm, native.Get(vm, res, "statusCode")))
 		state.mu.Lock()
 		defer state.mu.Unlock()
-		if len(args) > 0 && !native.IsNil(args[0]) {
-			_, _ = state.body.WriteString(native.String(vm, args[0]))
+		if state.upgrade != nil {
+			return moejs.Undefined(), native.TypeError(vm, "HTTP response is unavailable after WebSocket upgrade")
 		}
-		state.statusCode = int(native.Int(vm, native.Get(vm, res, "statusCode")))
+		if state.sealed || state.ended {
+			return engine.ObjectValue(res), nil
+		}
+		_, _ = state.body.WriteString(tail)
+		state.statusCode = status
 		if !state.ended {
 			state.ended = true
 			close(state.done)

pkg/jsruntime/host.go

@@ -37,3 +37,24 @@ func (h *Host) RunJob(vm *moe.Runtime, name string, job func() error) error {
 func (h *Host) Timeout() time.Duration {
 	return h.runtime.Timeout()
 }
+
+// AddResource registers a closer to run when the runtime shuts down. If shutdown
+// has already started, it runs the closer immediately and returns 0. Closers
+// must be nonblocking and must not access the JavaScript VM or wait for its event
+// loop; they may run outside the loop during shutdown.
+func (h *Host) AddResource(close func()) uint64 {
+	return h.runtime.AddResource(close)
+}
+
+// RemoveResource unregisters a resource without calling its closer. Unknown IDs
+// and 0 are ignored.
+func (h *Host) RemoveResource(id uint64) {
+	h.runtime.RemoveResource(id)
+}
+
+// ResourcesOpen reports whether the runtime still accepts resources, not whether
+// any resources are registered. The result is a snapshot; AddResource handles
+// registration racing with shutdown by immediately closing the new resource.
+func (h *Host) ResourcesOpen() bool {
+	return h.runtime.ResourcesOpen()
+}