mirror of
https://github.com/DeBrosOfficial/orama.git
synced 2026-06-16 22:54:12 +00:00
- Attach InvocationContext to the execution context in Engine.Execute to ensure host functions resolve identity from the request context. - Fixes a race condition where concurrent stateless invocations would overwrite the global singleton, causing cross-tenant leaks or nil namespace errors. - Added a regression test to verify per-invocation isolation under load.
66 lines
3.0 KiB
Go
66 lines
3.0 KiB
Go
package serverless
|
|
|
|
import "context"
|
|
|
|
// invCtxKey is the unexported context-value key used to attach an
|
|
// InvocationContext to a Go context. The empty struct is the standard
|
|
// Go pattern for context keys (avoids string-collision risk).
|
|
type invCtxKey struct{}
|
|
|
|
// WithInvocationContext returns a derived ctx that carries invCtx. Host
|
|
// function accessors check the ctx FIRST and only fall back to the
|
|
// HostFunctions singleton field when nothing is carried on ctx.
|
|
//
|
|
// Why this exists: HostFunctions is a process-wide singleton (one per
|
|
// gateway engine). Its `invCtx` field is shared across all WASM instances.
|
|
// For STATELESS functions the gateway sets/clears that field per-call
|
|
// (executor contextSetter/contextClearer), but the lock is released
|
|
// before WASM runs — two concurrent invocations CAN race on the field,
|
|
// and one's host call CAN read the other's identity.
|
|
//
|
|
// For PERSISTENT WS functions the race is far worse: the field used to be
|
|
// bound ONCE at instantiation and reused for the connection's lifetime.
|
|
// Two simultaneous persistent WS connections from different users
|
|
// overwrote each other's invCtx, and every subsequent function_invoke /
|
|
// GetCallerJWTSubject / GetSecret call from inside the WASM read whatever
|
|
// was bound LAST — silently leaking identity across tenants.
|
|
//
|
|
// The fix is per-call invCtx propagation through Go's context.Context.
|
|
// wazero passes the ctx given to api.Function.Call all the way through
|
|
// to host function callbacks (engine.go's host-function wrappers receive
|
|
// it), so every WASM-host hop carries its own invCtx and never reads the
|
|
// shared field.
|
|
//
|
|
// Persistent WS uses this exclusively (see persistent.Instance, which
|
|
// wraps every export call's ctx with the per-instance invCtx).
|
|
//
|
|
// Stateless Engine.Execute also attaches invCtx via this helper since
|
|
// bugboard #348 — AnChat's pubsub-triggered message-push-handler
|
|
// confirmed the "microseconds" race window was actually observable
|
|
// under production fan-out load: concurrent invocations either
|
|
// cross-tenant-leaked the namespace (silent) or saw a nil singleton
|
|
// during the brief window between contextSetter on one goroutine and
|
|
// contextClearer on another, producing "no namespace in invocation
|
|
// context" errors at host-fn entry. The singleton SetInvocationContext
|
|
// path remains in place as defense-in-depth — every host fn resolves
|
|
// via currentInvocationContext, which prefers ctx-attached over the
|
|
// singleton field, so the race is closed for the live path.
|
|
func WithInvocationContext(ctx context.Context, invCtx *InvocationContext) context.Context {
|
|
if invCtx == nil {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, invCtxKey{}, invCtx)
|
|
}
|
|
|
|
// InvocationContextFromCtx extracts the invCtx attached via
|
|
// WithInvocationContext, or nil if none is present. Exported so the
|
|
// hostfunctions package and any other consumer can read it without
|
|
// duplicating the key type.
|
|
func InvocationContextFromCtx(ctx context.Context) *InvocationContext {
|
|
if ctx == nil {
|
|
return nil
|
|
}
|
|
v, _ := ctx.Value(invCtxKey{}).(*InvocationContext)
|
|
return v
|
|
}
|