diff --git a/README.md b/README.md index 2cd6fda..270d231 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,17 @@ GLOBAL OPTIONS: --auth-key value, -k value the authentication key to use for incoming requests. [$AUTH_KEY] + progress + + --progress-callback-timeout value the timeout for a single progress callback delivery. (default: 1s) [$PROGRESS_CALLBACK_TIMEOUT] + --progress-allowed-hosts value [ --progress-allowed-hosts value ] restrict progress callback URLs to these hosts. Entries may be an exact hostname or a "*.example.com" wildcard. Unset allows any host, subject to the private-network guard below. [$PROGRESS_ALLOWED_HOSTS] + --progress-allow-private-networks allow progress callback delivery to loopback, link-local, and private IP addresses. Leave disabled unless the callback target is known to live on a trusted private network. (default: false) [$PROGRESS_ALLOW_PRIVATE_NETWORKS] + --progress-sidecar-max-body-bytes value the maximum size, in bytes, of a single worker-authored progress event POST. (default: 16384) [$PROGRESS_SIDECAR_MAX_BODY_BYTES] + --progress-sidecar-max-events value the maximum number of worker-authored progress events relayed per evaluation. (default: 50) [$PROGRESS_SIDECAR_MAX_EVENTS] + --progress-sidecar-burst-size value how many worker-authored progress events at the start of an evaluation are exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited. (default: 5) [$PROGRESS_SIDECAR_BURST_SIZE] + --progress-sidecar-min-event-interval value the minimum spacing between worker-authored progress events relayed per evaluation, once the burst allowance above is used up. (default: 10ms) [$PROGRESS_SIDECAR_MIN_EVENT_INTERVAL] + --progress-sidecar-unbind-grace-period value how long to keep relaying worker-authored progress events after a request returns, so a fire-and-forget POST dispatched just before the result can still land. (default: 250ms) [$PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD] + function --arg value, -a value [ --arg value, -a value ] additional arguments for to the worker process. [$FUNCTION_ARGS] @@ -185,6 +196,110 @@ Example request using cases: } ``` +### Progress Events + +The shim also exposes a µEd-compatible endpoint at `POST /evaluate` (see the [µEd spec](https://mued.org/spec)), separate from the legacy `POST /` endpoint documented above. When a client calls `/evaluate` with a `callbackUrl` in the request body, the shim POSTs a small JSON event to that URL at each stage of processing — in addition to, not instead of, the normal synchronous HTTP response. + +This lets a caller show progress to the end user (e.g. "Evaluating your submission…") without polling, and without the shim needing to hold a connection open. It works identically whether the shim is deployed standalone or on AWS Lambda. + +To opt in, include `callbackUrl` in the request body and, optionally, an `X-Request-Id` header — both are part of the µEd spec's own request contract, not shim-specific additions. Every event echoes back the `X-Request-Id` value verbatim so the caller can correlate it with the original request. + +```json +{ + "submission": { "type": "TEXT", "content": { "text": "..." } }, + "task": { "referenceSolution": { "text": "..." } }, + "callbackUrl": "https://your-service.example.com/hooks/shimmy-progress" +} +``` + +Four stages are emitted, in order, for a successful evaluation: + +| Stage | Meaning | +|-------|---------| +| `preparing` | The evaluation environment is being set up (a worker is ready — freshly booted or reused). | +| `evaluating` | The evaluation function is being invoked. | +| `completed` | Feedback has been computed. `data.feedback` carries the same array returned in the synchronous response body. | +| `failed` | A terminal failure occurred at some stage. `message` is safe to show to an end user; `error` carries raw technical detail for logs only. | + +`completed` and `failed` are terminal — at most one of them is delivered per request, whichever occurs first. + +Example event body: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "evaluating", + "command": "eval", + "message": "Evaluating your submission…", + "timestamp": "2026-08-04T14:23:01.512Z" +} +``` + +Example terminal event, with the feedback payload attached: + +```json +{ + "correlationId": "req-7c193f38", + "stage": "completed", + "command": "eval", + "message": "Feedback is ready.", + "data": { + "feedback": [ + { "awardedPoints": 1, "message": "Well done" } + ] + }, + "timestamp": "2026-08-04T14:23:02.310Z" +} +``` + +Delivery is best-effort and never blocks or fails the evaluation itself: each callback POST is bounded by `--progress-callback-timeout` (default `1s`, see [Usage](#usage)); a slow, unreachable, or erroring receiver is logged and skipped, never surfaced to the caller as an evaluation failure. + +#### Callback URL safety (SSRF protection) + +Since `callbackUrl` is caller-supplied, the shim guards against it being used to reach services it shouldn't be able to reach: + +- **By default**, callback delivery refuses to dial loopback, link-local (this includes cloud metadata endpoints like `169.254.169.254`), and private (RFC1918/RFC4193) IP addresses — checked against the address actually resolved and dialed, not just the URL's literal hostname, so a public-looking domain that resolves to a private address is still blocked. Set `--progress-allow-private-networks` only if the callback target is known to live on a private network you trust (e.g. a same-VPC service). +- **`--progress-allowed-hosts`** optionally restricts callback URLs to an explicit list of hostnames (exact match, or `*.example.com` wildcards). Unset means any (non-private) host is accepted. + +A rejected callback URL behaves like any other delivery failure: it's logged and skipped, never surfaced to the caller as an evaluation failure. + +> **Note:** the µEd spec describes `callbackUrl` for asynchronous *final-result* delivery — the service may return `202 Accepted` immediately and POST the result later. The shim doesn't implement that 202 flow; it always responds synchronously with `200 OK` and the feedback body as normal. It reuses the same `callbackUrl` field to additionally deliver progress events — including the final feedback, via the `completed` event's `data` field — rather than requiring a shim-specific header for the same concept. + +#### Custom progress events from the evaluation function + +The four stages above are emitted by shimmy itself, around the evaluation function call as a whole — `evaluating` covers the entire invocation as one span. An evaluation function that does multiple steps internally (e.g. several model calls) can emit its own progress events *during* that span, which are relayed through the same `callbackUrl` alongside shimmy's own events. + +When a request opts in to progress reporting (via `callbackUrl`), shimmy starts a loopback-only HTTP listener and passes its address to the evaluation function process as the `EVAL_PROGRESS_URL` environment variable, the same way it passes `EVAL_RPC_TRANSPORT`, `EVAL_FILE_NAME_REQUEST`, etc. (see [Communication Channels](#communication-channels) below). This works identically regardless of interface (`rpc` or `file`) or RPC transport, and regardless of the evaluation function's language — it only needs to be able to make an HTTP POST. + +To emit a custom event, `POST` a small JSON body to `EVAL_PROGRESS_URL`: + +```json +{ + "message": "Checking correctness…", + "data": { "step": 2, "of": 3 } +} +``` + +- `message` (string, required): student/teacher-facing text. +- `data` (object, optional): free-form, passed through as-is. +- There is no `stage` field, by design: an evaluation function can never claim `preparing`, `evaluating`, `completed`, or `failed` — those remain exclusively shim-authored. Custom events are always delivered with `"stage": "progress"`. + +The response status is informational only — the evaluation function should treat every response as fire-and-forget and never fail on a non-2xx status. Delivery is best-effort, same as outbound callback delivery: `202` accepted (delivery to `callbackUrl` is then attempted in the background), `400` malformed body or empty `message`, `413` body too large, `429` rate limited, `503` no request currently associated with the listener (e.g. a stray POST arriving after both the request has finished and the grace period below has elapsed). + +To bound how much an evaluation function (which may be running untrusted, sandboxed code) can push through this channel, events are capped before relay: + +| Flag | Env var | Default | Description | +|------|---------|---------|-------------| +| `--progress-sidecar-max-body-bytes` | `PROGRESS_SIDECAR_MAX_BODY_BYTES` | `16384` | Maximum size, in bytes, of a single event POST. | +| `--progress-sidecar-max-events` | `PROGRESS_SIDECAR_MAX_EVENTS` | `50` | Maximum number of events relayed per evaluation. | +| `--progress-sidecar-burst-size` | `PROGRESS_SIDECAR_BURST_SIZE` | `5` | Events at the start of a span exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited. | +| `--progress-sidecar-min-event-interval` | `PROGRESS_SIDECAR_MIN_EVENT_INTERVAL` | `10ms` | Minimum spacing between relayed events, once the burst allowance is used up. | +| `--progress-sidecar-unbind-grace-period` | `PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD` | `250ms` | How long the listener keeps relaying after a request returns, so a fire-and-forget event POST dispatched by the worker just before returning its result still has a window to land. | + +> **Sandboxing note:** under `--sandbox` alone, the worker keeps the host network namespace and can reach the loopback listener normally. Only the separate, explicit `--sandbox-disable-network` flag isolates networking (and loopback specifically) — under that flag, custom progress events are silently dropped, the same as any other best-effort delivery failure. + +This is a shim-side contract only; no client library ships in this repo. Evaluation function libraries (e.g. per-language toolkits) can build a thin wrapper around reading `EVAL_PROGRESS_URL` and POSTing to it. + ### Communication Channels The shim supports two interface modes, selected with `--interface`: @@ -211,6 +326,7 @@ The shim injects the following environment variables into the evaluation functio | `EVAL_RPC_HTTP_URL` | HTTP URL (HTTP transport only) | | `EVAL_RPC_WS_URL` | WebSocket URL (WS transport only) | | `EVAL_RPC_TCP_ADDRESS` | TCP address (TCP transport only) | +| `EVAL_PROGRESS_URL` | Local URL to POST [custom progress events](#custom-progress-events-from-the-evaluation-function) to (only set when the request opted in via `callbackUrl`) | #### File System (`--interface file`) @@ -236,6 +352,7 @@ The shim also sets the following environment variables: | `EVAL_IO` | `FILE` | | `EVAL_FILE_NAME_REQUEST` | Path to the input file | | `EVAL_FILE_NAME_RESPONSE` | Path to the output file | +| `EVAL_PROGRESS_URL` | Local URL to POST [custom progress events](#custom-progress-events-from-the-evaluation-function) to (only set when the request opted in via `callbackUrl`) | > Using the file interface is recommended for large payloads such as base64-encoded images. diff --git a/cmd/root.go b/cmd/root.go index 275c31e..427f48b 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -43,6 +43,62 @@ functions on arbitrary, serverless platforms.` Category: "auth", EnvVars: []string{"AUTH_KEY"}, }, + // progress flags + &cli.DurationFlag{ + Name: "progress-callback-timeout", + Usage: "the timeout for a single progress callback delivery.", + Value: time.Second, + Category: "progress", + EnvVars: []string{"PROGRESS_CALLBACK_TIMEOUT"}, + }, + &cli.StringSliceFlag{ + Name: "progress-allowed-hosts", + Usage: "restrict progress callback URLs to these hosts. Entries may be an exact hostname or a \"*.example.com\" wildcard. Unset allows any host, subject to the private-network guard below.", + Category: "progress", + EnvVars: []string{"PROGRESS_ALLOWED_HOSTS"}, + }, + &cli.BoolFlag{ + Name: "progress-allow-private-networks", + Usage: "allow progress callback delivery to loopback, link-local, and private IP addresses. Leave disabled unless the callback target is known to live on a trusted private network.", + Value: false, + Category: "progress", + EnvVars: []string{"PROGRESS_ALLOW_PRIVATE_NETWORKS"}, + }, + &cli.Int64Flag{ + Name: "progress-sidecar-max-body-bytes", + Usage: "the maximum size, in bytes, of a single worker-authored progress event POST.", + Value: 16 * 1024, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MAX_BODY_BYTES"}, + }, + &cli.IntFlag{ + Name: "progress-sidecar-max-events", + Usage: "the maximum number of worker-authored progress events relayed per evaluation.", + Value: 50, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MAX_EVENTS"}, + }, + &cli.IntFlag{ + Name: "progress-sidecar-burst-size", + Usage: "how many worker-authored progress events at the start of an evaluation are exempt from the minimum spacing below, so a handful of legitimate back-to-back checkpoints aren't rate limited.", + Value: 5, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_BURST_SIZE"}, + }, + &cli.DurationFlag{ + Name: "progress-sidecar-min-event-interval", + Usage: "the minimum spacing between worker-authored progress events relayed per evaluation, once the burst allowance above is used up.", + Value: 10 * time.Millisecond, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_MIN_EVENT_INTERVAL"}, + }, + &cli.DurationFlag{ + Name: "progress-sidecar-unbind-grace-period", + Usage: "how long to keep relaying worker-authored progress events after a request returns, so a fire-and-forget POST dispatched just before the result can still land.", + Value: 250 * time.Millisecond, + Category: "progress", + EnvVars: []string{"PROGRESS_SIDECAR_UNBIND_GRACE_PERIOD"}, + }, // shim flags &cli.StringFlag{ Name: "interface", @@ -317,20 +373,28 @@ func parseRootConfig(ctx *cli.Context) (config.Config, error) { // map cli flags to config fields cliMap := map[string]string{ - "auth-key": "auth.key", - "max-workers": "runtime.max_workers", - "command": "runtime.cmd", - "cwd": "runtime.cwd", - "arg": "runtime.arg", - "env": "runtime.env", - "interface": "runtime.io.interface", - "rpc-transport": "runtime.io.rpc.transport", - "rpc-transport-ipc-endpoint": "runtime.io.rpc.ipc.endpoint", - "rpc-transport-http-url": "runtime.io.rpc.http.url", - "rpc-transport-ws-url": "runtime.io.rpc.ws.url", - "rpc-transport-tcp-address": "runtime.io.rpc.tcp.address", - "worker-send-timeout": "runtime.send.timeout", - "worker-stop-timeout": "runtime.stop.timeout", + "auth-key": "auth.key", + "progress-callback-timeout": "progress.callback_timeout", + "progress-allowed-hosts": "progress.allowed_hosts", + "progress-allow-private-networks": "progress.allow_private_networks", + "progress-sidecar-max-body-bytes": "progress.sidecar.max_body_bytes", + "progress-sidecar-max-events": "progress.sidecar.max_events_per_span", + "progress-sidecar-burst-size": "progress.sidecar.burst_size", + "progress-sidecar-min-event-interval": "progress.sidecar.min_event_interval", + "progress-sidecar-unbind-grace-period": "progress.sidecar.unbind_grace_period", + "max-workers": "runtime.max_workers", + "command": "runtime.cmd", + "cwd": "runtime.cwd", + "arg": "runtime.arg", + "env": "runtime.env", + "interface": "runtime.io.interface", + "rpc-transport": "runtime.io.rpc.transport", + "rpc-transport-ipc-endpoint": "runtime.io.rpc.ipc.endpoint", + "rpc-transport-http-url": "runtime.io.rpc.http.url", + "rpc-transport-ws-url": "runtime.io.rpc.ws.url", + "rpc-transport-tcp-address": "runtime.io.rpc.tcp.address", + "worker-send-timeout": "runtime.send.timeout", + "worker-stop-timeout": "runtime.stop.timeout", // sandbox "sandbox": "runtime.sandbox.enabled", "sandbox-nsjail-path": "runtime.sandbox.nsjail_path", diff --git a/config/config.go b/config/config.go index 020d38e..fc90d9b 100644 --- a/config/config.go +++ b/config/config.go @@ -1,6 +1,9 @@ package config -import "github.com/lambda-feedback/shimmy/runtime" +import ( + "github.com/lambda-feedback/shimmy/internal/progress" + "github.com/lambda-feedback/shimmy/runtime" +) type MessageEncoding string @@ -25,4 +28,7 @@ type Config struct { // Auth is the authentication configuration Auth AuthConfig `conf:"auth"` + + // Progress is the configuration for outbound progress-callback delivery + Progress progress.Config `conf:"progress"` } diff --git a/handler/module.go b/handler/module.go index a58f29f..434dbda 100644 --- a/handler/module.go +++ b/handler/module.go @@ -1,6 +1,11 @@ package handler -import "go.uber.org/fx" +import ( + "go.uber.org/fx" + + "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" +) func Module() fx.Option { return fx.Module("common", @@ -10,5 +15,7 @@ func Module() fx.Option { fx.Provide(NewHealthRoute), fx.Provide(NewMuEdEvaluateRoute), fx.Provide(NewMuEdEvaluateHealthRoute), + fx.Provide(func(cfg config.Config) progress.Config { return cfg.Progress }), + fx.Provide(progress.NewHTTPFactory), ) } diff --git a/handler/mued.go b/handler/mued.go index 53c4c73..fbcdb03 100644 --- a/handler/mued.go +++ b/handler/mued.go @@ -1,42 +1,74 @@ package handler import ( + "crypto/rand" "encoding/json" "fmt" "io" "net/http" + "time" "go.uber.org/fx" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/runtime" ) const muEdVersionHeader = "X-Api-Version" +// muEdRequestIDHeader is the µEd spec's request-tracing header (see +// https://mued.org/spec, X-Request-Id parameter). It's echoed back on every +// response, generating one if the caller didn't supply it, and progress +// events reuse the resolved value as their correlation key. +const muEdRequestIDHeader = "X-Request-Id" + +// resolveRequestID returns the caller-supplied X-Request-Id, or generates +// one if absent, so every request is traceable and correlatable even when +// the caller doesn't participate in tracing itself. +func resolveRequestID(r *http.Request) string { + if id := r.Header.Get(muEdRequestIDHeader); id != "" { + return id + } + return generateRequestID() +} + +func generateRequestID() string { + b := make([]byte, 4) + if _, err := rand.Read(b); err != nil { + // crypto/rand.Read on a real OS essentially never fails; fall back + // to a timestamp-based id rather than leaving the request untraceable. + return fmt.Sprintf("req-%08x", time.Now().UnixNano()) + } + return fmt.Sprintf("req-%x", b) +} + type MuEdHandlerParams struct { fx.In - Handler runtime.Handler - Runtime runtime.Runtime - Config config.Config - Log *zap.Logger + Handler runtime.Handler + Runtime runtime.Runtime + Config config.Config + Log *zap.Logger + ProgressFactory progress.Factory } type MuEdHandler struct { - handler runtime.Handler - runtime runtime.Runtime - config config.Config - log *zap.Logger + handler runtime.Handler + runtime runtime.Runtime + config config.Config + log *zap.Logger + progressFactory progress.Factory } func NewMuEdHandler(params MuEdHandlerParams) *MuEdHandler { return &MuEdHandler{ - handler: params.Handler, - runtime: params.Runtime, - config: params.Config, - log: params.Log, + handler: params.Handler, + runtime: params.Runtime, + config: params.Config, + log: params.Log, + progressFactory: params.ProgressFactory, } } @@ -97,6 +129,9 @@ func (h *MuEdHandler) checkAuth(w http.ResponseWriter, r *http.Request) bool { // ServeEvaluate handles POST /evaluate. func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { + requestID := resolveRequestID(r) + w.Header().Set(muEdRequestIDHeader, requestID) + if !h.checkAuth(w, r) { return } @@ -157,9 +192,28 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { Header: header, } - resp := h.handler.Handle(r.Context(), req) + var callbackURL string + if muEdReq.CallbackUrl != nil { + callbackURL = *muEdReq.CallbackUrl + } + + ctx := r.Context() + reporter, err := h.progressFactory.NewReporter(callbackURL, requestID) + if err != nil { + h.log.Warn("invalid callbackUrl, disabling progress reporting", zap.Error(err)) + } else if reporter != nil { + ctx = progress.ContextWithReporter(ctx, reporter) + } + + resp := h.handler.Handle(ctx, req) if resp.StatusCode != http.StatusOK { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: string(command), + Message: muEdErrorMessageFromBody(resp.Body), + }) + for k, v := range resp.Header { for _, vv := range v { w.Header().Add(k, vv) @@ -190,14 +244,51 @@ func (h *MuEdHandler) ServeEvaluate(w http.ResponseWriter, r *http.Request) { feedback = runtime.MuEdToEvaluateFeedback(result) } + // Carry the feedback itself on the completed event so that, when a + // caller supplies callbackUrl, that callback genuinely fulfils the + // µEd spec's "deliver feedback results to this URL" wording — even + // though shimmy always takes the synchronous 200 path rather than + // the spec's 202-Accepted deferred-delivery flow. + progress.Emit(ctx, progress.Event{ + Stage: progress.StageCompleted, + Command: string(command), + Message: "Feedback is ready.", + Data: map[string]any{"feedback": feedback}, + }) + w.Header().Set("Content-Type", "application/json") w.Header().Set(muEdVersionHeader, version) w.WriteHeader(http.StatusOK) json.NewEncoder(w).Encode(feedback) //nolint:errcheck } +// muEdErrorMessageFromBody best-effort extracts a human-readable message +// from a JSON error body of the shape {"error": {"message": "..."}}. +func muEdErrorMessageFromBody(body []byte) string { + const fallback = "We couldn't evaluate your answer. Please try again." + + var errBody map[string]any + if err := json.Unmarshal(body, &errBody); err != nil { + return fallback + } + + errObj, ok := errBody["error"].(map[string]any) + if !ok { + return fallback + } + + msg, ok := errObj["message"].(string) + if !ok || msg == "" { + return fallback + } + + return msg +} + // ServeHealth handles GET /evaluate/health. func (h *MuEdHandler) ServeHealth(w http.ResponseWriter, r *http.Request) { + w.Header().Set(muEdRequestIDHeader, resolveRequestID(r)) + if !h.checkAuth(w, r) { return } diff --git a/handler/mued_test.go b/handler/mued_test.go index afb65af..65388ac 100644 --- a/handler/mued_test.go +++ b/handler/mued_test.go @@ -8,9 +8,12 @@ import ( "io" "net/http" "net/http/httptest" + "sync" "testing" + "time" "github.com/lambda-feedback/shimmy/config" + "github.com/lambda-feedback/shimmy/internal/progress" "github.com/lambda-feedback/shimmy/runtime" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -39,18 +42,34 @@ func (m *MockRuntime) Shutdown(ctx context.Context) error { // --- Helpers --- +// newMuEdHandler builds a handler with a default, inert progress factory: +// since none of the existing tests set callbackUrl in the request body, +// NewReporter always returns (nil, nil) and behavior is unchanged. Tests +// that exercise progress reporting itself use newMuEdHandlerWithProgress. func newMuEdHandler(h runtime.Handler, r runtime.Runtime, key string) *MuEdHandler { + return newMuEdHandlerWithProgress(h, r, key, progress.NewHTTPFactory(progress.HTTPFactoryParams{ + Log: zap.NewNop(), + })) +} + +func newMuEdHandlerWithProgress(h runtime.Handler, r runtime.Runtime, key string, pf progress.Factory) *MuEdHandler { return &MuEdHandler{ - handler: h, - runtime: r, - config: config.Config{Auth: config.AuthConfig{Key: key}}, - log: zap.NewNop(), + handler: h, + runtime: r, + config: config.Config{Auth: config.AuthConfig{Key: key}}, + log: zap.NewNop(), + progressFactory: pf, } } func mathEvalBody(t *testing.T) []byte { t.Helper() - b, err := json.Marshal(map[string]any{ + return mathEvalBodyWithCallback(t, "") +} + +func mathEvalBodyWithCallback(t *testing.T, callbackURL string) []byte { + t.Helper() + body := map[string]any{ "submission": map[string]any{ "type": "MATH", "content": map[string]any{"expression": "x^2"}, @@ -60,7 +79,11 @@ func mathEvalBody(t *testing.T) []byte { "expression": "x^2", }, }, - }) + } + if callbackURL != "" { + body["callbackUrl"] = callbackURL + } + b, err := json.Marshal(body) require.NoError(t, err) return b } @@ -274,6 +297,166 @@ func TestMuEdServeEvaluate_WorkerErrorForwarded(t *testing.T) { assert.Equal(t, errorBody, bytes.TrimRight(raw, "\n")) } +// --- Progress callback tests (ServeEvaluate) --- + +// newProgressCallbackServer spins up a fake progress-callback receiver +// that records every decoded request body it receives. +func newProgressCallbackServer(t *testing.T, handlerFn http.HandlerFunc) (*httptest.Server, *[]map[string]any) { + t.Helper() + + var mu sync.Mutex + var received []map[string]any + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + mu.Lock() + received = append(received, body) + mu.Unlock() + + if handlerFn != nil { + handlerFn(w, r) + return + } + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + + return srv, &received +} + +// newProgressFactory builds a factory with SSRF protection relaxed: these +// tests use httptest.NewServer (a loopback address) to stand in for the +// caller's real, non-loopback callback receiver, so the default +// private-network guard would otherwise reject every delivery here. The +// guard itself is covered directly in internal/progress. +func newProgressFactory(t *testing.T, timeout time.Duration) progress.Factory { + t.Helper() + return progress.NewHTTPFactory(progress.HTTPFactoryParams{ + Config: progress.Config{CallbackTimeout: timeout, AllowPrivateNetworks: true}, + Log: zap.NewNop(), + }) +} + +func TestMuEdServeEvaluate_ProgressCallback_Success(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-1") + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-1", evt["correlationId"]) + assert.Equal(t, "completed", evt["stage"]) + assert.Equal(t, "eval", evt["command"]) + + data, ok := evt["data"].(map[string]any) + require.True(t, ok, "expected data field on the completed event") + feedback, ok := data["feedback"].([]any) + require.True(t, ok, "expected data.feedback array") + require.Len(t, feedback, 1) + item, ok := feedback[0].(map[string]any) + require.True(t, ok) + assert.Equal(t, "Well done", item["message"]) + assert.Equal(t, 1.0, item["awardedPoints"]) +} + +func TestMuEdServeEvaluate_ProgressCallback_Failure(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + errorBody, _ := json.Marshal(map[string]any{ + "error": map[string]any{"message": "boom"}, + }) + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything).Return(runtime.Response{ + StatusCode: http.StatusInternalServerError, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: errorBody, + }) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-2") + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusInternalServerError, w.Result().StatusCode) + + require.Len(t, *received, 1) + evt := (*received)[0] + assert.Equal(t, "corr-2", evt["correlationId"]) + assert.Equal(t, "failed", evt["stage"]) + assert.Equal(t, "boom", evt["message"]) +} + +func TestMuEdServeEvaluate_ProgressCallback_NoCallbackUrl_Unchanged(t *testing.T) { + _, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + assert.Empty(t, *received, "no callbackUrl in the request body should mean no callback requests") +} + +func TestMuEdServeEvaluate_ProgressCallback_InvalidCallbackUrl_EvaluationStillSucceeds(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, "not-a-url"))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + res := w.Result() + defer res.Body.Close() + body, _ := io.ReadAll(res.Body) + + assert.Equal(t, http.StatusOK, res.StatusCode) + + var feedback []map[string]any + require.NoError(t, json.Unmarshal(body, &feedback)) + require.Len(t, feedback, 1) +} + +func TestMuEdServeEvaluate_ProgressCallback_SlowReceiver_DoesNotBlockResponse(t *testing.T) { + srv, _ := newProgressCallbackServer(t, func(w http.ResponseWriter, r *http.Request) { + time.Sleep(200 * time.Millisecond) + w.WriteHeader(http.StatusOK) + }) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + req.Header.Set(muEdRequestIDHeader, "corr-3") + w := httptest.NewRecorder() + + start := time.Now() + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, 20*time.Millisecond)).ServeEvaluate(w, req) + elapsed := time.Since(start) + + assert.Equal(t, http.StatusOK, w.Result().StatusCode) + assert.Less(t, elapsed, 150*time.Millisecond, "ServeEvaluate should return promptly, bounded by CallbackTimeout") +} + // --- ServeHealth tests --- func TestMuEdServeHealth_Success(t *testing.T) { @@ -496,3 +679,62 @@ func TestMuEdServeHealth_UnsupportedVersionHeader(t *testing.T) { mockRuntime.AssertNotCalled(t, "Handle", mock.Anything, mock.Anything) } + +// --- Request ID tests (ServeEvaluate) --- + +func TestMuEdServeEvaluate_RequestID_EchoedWhenSupplied(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "ok")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + req.Header.Set(muEdRequestIDHeader, "caller-supplied-id") + w := httptest.NewRecorder() + + newMuEdHandler(mockHandler, nil, "").ServeEvaluate(w, req) + + assert.Equal(t, "caller-supplied-id", w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_RequestID_GeneratedWhenAbsent(t *testing.T) { + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "ok")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBody(t))) + w := httptest.NewRecorder() + + newMuEdHandler(mockHandler, nil, "").ServeEvaluate(w, req) + + assert.NotEmpty(t, w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_RequestID_EchoedOnErrorResponses(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader([]byte("not json"))) + req.Header.Set(muEdRequestIDHeader, "caller-supplied-id") + w := httptest.NewRecorder() + + newMuEdHandler(new(MockHandler), nil, "").ServeEvaluate(w, req) + + assert.Equal(t, http.StatusBadRequest, w.Result().StatusCode) + assert.Equal(t, "caller-supplied-id", w.Result().Header.Get(muEdRequestIDHeader)) +} + +func TestMuEdServeEvaluate_ProgressCallback_GeneratedRequestIDUsedAsCorrelation(t *testing.T) { + srv, received := newProgressCallbackServer(t, nil) + + mockHandler := new(MockHandler) + mockHandler.On("Handle", mock.Anything, mock.Anything). + Return(evalHandlerResponse(true, "Well done")) + + req := httptest.NewRequest(http.MethodPost, "/evaluate", bytes.NewReader(mathEvalBodyWithCallback(t, srv.URL))) + w := httptest.NewRecorder() + + newMuEdHandlerWithProgress(mockHandler, nil, "", newProgressFactory(t, time.Second)).ServeEvaluate(w, req) + + respRequestID := w.Result().Header.Get(muEdRequestIDHeader) + require.NotEmpty(t, respRequestID) + + require.Len(t, *received, 1) + assert.Equal(t, respRequestID, (*received)[0]["correlationId"]) +} diff --git a/internal/execution/dispatcher.go b/internal/execution/dispatcher.go index 300ca3f..f1a562f 100644 --- a/internal/execution/dispatcher.go +++ b/internal/execution/dispatcher.go @@ -7,6 +7,7 @@ import ( "github.com/lambda-feedback/shimmy/internal/execution/dispatcher" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type Dispatcher dispatcher.Dispatcher @@ -27,6 +28,10 @@ type Params struct { // Config is the config for the dispatcher and the underlying supervisors Config Config + // Progress configures worker-authored progress event delivery, + // passed through to the underlying supervisor(s). + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -38,8 +43,9 @@ func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { Config: dispatcher.DedicatedDispatcherConfig{ Supervisor: params.Config.Supervisor, }, - Context: params.Context, - Log: params.Log, + Context: params.Context, + Progress: params.Progress, + Log: params.Log, }, ) } else { @@ -49,8 +55,9 @@ func NewDispatcher(params Params) (dispatcher.Dispatcher, error) { Supervisor: params.Config.Supervisor, MaxWorkers: params.Config.MaxWorkers, }, - Context: params.Context, - Log: params.Log, + Context: params.Context, + Progress: params.Progress, + Log: params.Log, }, ) } diff --git a/internal/execution/dispatcher/dispatcher_dedicated.go b/internal/execution/dispatcher/dispatcher_dedicated.go index 2cb5223..842e647 100644 --- a/internal/execution/dispatcher/dispatcher_dedicated.go +++ b/internal/execution/dispatcher/dispatcher_dedicated.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type DedicatedDispatcher struct { @@ -31,6 +32,10 @@ type DedicatedDispatcherParams struct { // SupervisorFactory is the factory function to create a new supervisor SupervisorFactory SupervisorFactory + // Progress configures worker-authored progress event delivery, + // passed through to the underlying supervisor. + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -110,8 +115,9 @@ func createSupervisor( params DedicatedDispatcherParams, ) (supervisor.Supervisor, error) { return params.SupervisorFactory(supervisor.Params{ - Context: params.Context, - Config: params.Config.Supervisor, - Log: params.Log, + Context: params.Context, + Config: params.Config.Supervisor, + Progress: params.Progress, + Log: params.Log, }) } diff --git a/internal/execution/dispatcher/dispatcher_pooled.go b/internal/execution/dispatcher/dispatcher_pooled.go index 7a49429..967e26c 100644 --- a/internal/execution/dispatcher/dispatcher_pooled.go +++ b/internal/execution/dispatcher/dispatcher_pooled.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) type PooledDispatcher struct { @@ -37,6 +38,10 @@ type PooledDispatcherParams struct { // SupervisorFactory is the factory function to create a new supervisor SupervisorFactory SupervisorFactory + // Progress configures worker-authored progress event delivery, + // passed through to each pooled supervisor. + Progress progress.Config + // Log is the logger to use for the dispatcher Log *zap.Logger } @@ -157,9 +162,10 @@ func createPool( constructor := func(ctx context.Context) (supervisor.Supervisor, error) { sv, err := params.SupervisorFactory(supervisor.Params{ - Context: ctx, - Config: params.Config.Supervisor, - Log: params.Log, + Context: ctx, + Config: params.Config.Supervisor, + Progress: params.Progress, + Log: params.Log, }) if err != nil { return nil, err diff --git a/internal/execution/dispatcher/dispatcher_pooled_test.go b/internal/execution/dispatcher/dispatcher_pooled_test.go index 0fee760..723ec9a 100644 --- a/internal/execution/dispatcher/dispatcher_pooled_test.go +++ b/internal/execution/dispatcher/dispatcher_pooled_test.go @@ -3,7 +3,6 @@ package dispatcher_test import ( "context" "testing" - "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" @@ -210,8 +209,8 @@ func TestPooledDispatcher_Send_ReleaseSupervisorWaitErrorShutdown(t *testing.T) _, err := m.Send(context.Background(), "test", data) assert.NoError(t, err) - // wait for the release to happen in a goroutine - <-time.After(1 * time.Millisecond) + // wait for the background goroutine to finish by draining the pool + m.Shutdown(context.Background()) assert.Equal(t, 1, waited) } diff --git a/internal/execution/supervisor/adapter.go b/internal/execution/supervisor/adapter.go index e31eb13..7ec6803 100644 --- a/internal/execution/supervisor/adapter.go +++ b/internal/execution/supervisor/adapter.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // AdapterWorkerFactoryFn is a type alias for a function that creates a worker @@ -42,20 +43,26 @@ type Adapter interface { // MARK: - factory -// defaultAdapterFactory is the default adapter factory -// that creates an adapter based on the given IO mode. -func defaultAdapterFactory( - workerFactory AdapterWorkerFactoryFn, - config IOConfig, - log *zap.Logger, -) (Adapter, error) { - switch config.Interface { - case FileIO: - return newFileAdapter(workerFactory, log), nil - case RpcIO: - return newRpcAdapter(workerFactory, config.Rpc, log), nil - default: - return nil, ErrUnsupportedIOInterface +// newDefaultAdapterFactory returns the default AdapterFactoryFn, wiring +// each created adapter's worker-authored progress side-channel (see +// internal/progress.Sidecar) with the given limits. It's a closure rather +// than a plain function so that AdapterFactoryFn's signature - and every +// test double built against it - doesn't need to carry progress.Config +// through every caller. +func newDefaultAdapterFactory(progressCfg progress.Config) AdapterFactoryFn { + return func( + workerFactory AdapterWorkerFactoryFn, + config IOConfig, + log *zap.Logger, + ) (Adapter, error) { + switch config.Interface { + case FileIO: + return newFileAdapter(workerFactory, progressCfg.Sidecar, log), nil + case RpcIO: + return newRpcAdapter(workerFactory, config.Rpc, progressCfg.Sidecar, log), nil + default: + return nil, ErrUnsupportedIOInterface + } } } diff --git a/internal/execution/supervisor/adapter_file.go b/internal/execution/supervisor/adapter_file.go index 7917f47..1bd2838 100644 --- a/internal/execution/supervisor/adapter_file.go +++ b/internal/execution/supervisor/adapter_file.go @@ -15,6 +15,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // fileAdapter is an adapter that allows supervisors to use files to @@ -32,6 +33,9 @@ type fileAdapter struct { // worker is the worker that is managed by the adapter. worker worker.Worker + // sidecarCfg configures the worker-authored progress side-channel. + sidecarCfg progress.SidecarConfig + log *zap.Logger } @@ -39,10 +43,12 @@ var _ Adapter = (*fileAdapter)(nil) func newFileAdapter( workerFactory AdapterWorkerFactoryFn, + sidecarCfg progress.SidecarConfig, log *zap.Logger, ) *fileAdapter { return &fileAdapter{ workerFactory: workerFactory, + sidecarCfg: sidecarCfg, log: log.Named("adapter_file"), } } @@ -153,14 +159,26 @@ func (a *fileAdapter) Send( // ensure env is not nil if startParams.Env == nil { - startParams.Env = make([]string, 0, 3) + startParams.Env = make([]string, 0, 4) } + // the file interface is one process per request, so the sidecar is + // scoped entirely to this call - no Bind/Unbind swap needed, unlike + // the persistent rpcAdapter. + sidecar, err := progress.NewSidecar(a.sidecarCfg, a.log) + if err != nil { + return nil, fmt.Errorf("error starting progress sidecar: %w", err) + } + defer sidecar.Close() + + sidecar.Bind(method, progress.FromContext(ctx)) + // append req and res file names to worker env startParams.Env = append(startParams.Env, "EVAL_IO=FILE", "EVAL_FILE_NAME_REQUEST="+reqFile.Name(), "EVAL_FILE_NAME_RESPONSE="+resFile.Name(), + "EVAL_PROGRESS_URL="+sidecar.URL(), ) // create the worker with modified args and env diff --git a/internal/execution/supervisor/adapter_file_test.go b/internal/execution/supervisor/adapter_file_test.go index 380238b..db6bc09 100644 --- a/internal/execution/supervisor/adapter_file_test.go +++ b/internal/execution/supervisor/adapter_file_test.go @@ -3,17 +3,54 @@ package supervisor import ( "context" "io" + "net/http" "os" "strings" + "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) +// recordingReporter is a minimal progress.Reporter test double, local to +// this package since progress.Reporter's own test double is unexported +// in a different package. +type recordingReporter struct { + mu sync.Mutex + events []progress.Event +} + +func (r *recordingReporter) Report(_ context.Context, evt progress.Event) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, evt) +} + +func (r *recordingReporter) recorded() []progress.Event { + r.mu.Lock() + defer r.mu.Unlock() + return append([]progress.Event(nil), r.events...) +} + +// envValue returns the value of the first "key=value" entry in env, or "" +// if key isn't present. +func envValue(env []string, key string) string { + prefix := key + "=" + for _, e := range env { + if strings.HasPrefix(e, prefix) { + return strings.TrimPrefix(e, prefix) + } + } + return "" +} + func TestFileAdapter_Start_DoesNotStartWorker(t *testing.T) { a, w := createFileAdapter(t) @@ -132,6 +169,56 @@ func TestFileAdapter_Send_ReturnsInvalidDataError(t *testing.T) { w.AssertNotCalled(t, "Start") } +func TestFileAdapter_Send_InjectsProgressURLAndRelaysWorkerEvents(t *testing.T) { + w := worker.NewMockWorker(t) + + var sp *worker.StartConfig + workerFactory := func(params worker.StartConfig) (worker.Worker, error) { + sp = ¶ms + return w, nil + } + + a := &fileAdapter{ + workerFactory: workerFactory, + log: zap.NewNop(), + } + + r := &recordingReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + data := map[string]any{"foo": "bar"} + + w.EXPECT().Start(mock.Anything).RunAndReturn(func(ctx context.Context) error { + progressURL := envValue(sp.Env, "EVAL_PROGRESS_URL") + require.NotEmpty(t, progressURL, "expected EVAL_PROGRESS_URL in worker env") + + resp, err := http.Post(progressURL, "application/json", strings.NewReader(`{"message":"checking correctness"}`)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusAccepted, resp.StatusCode) + + requestFileName := sp.Args[len(sp.Args)-2] + responseFileName := sp.Args[len(sp.Args)-1] + reqData, _ := os.ReadFile(requestFileName) + _ = os.WriteFile(responseFileName, reqData, os.ModeAppend) + return nil + }) + w.EXPECT().ReadPipe().Return(io.NopCloser(strings.NewReader("")), nil) + var cell int + w.EXPECT().Wait(mock.Anything).Return(worker.ExitEvent{Code: &cell}, nil) + + _, err := a.Send(ctx, "eval", data, 10) + require.NoError(t, err) + + assert.Eventually(t, func() bool { + return len(r.recorded()) == 1 + }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") + + events := r.recorded() + assert.Equal(t, progress.StageProgress, events[0].Stage) + assert.Equal(t, "eval", events[0].Command) + assert.Equal(t, "checking correctness", events[0].Message) +} + func createFileAdapter(t *testing.T) (*fileAdapter, *worker.MockWorker) { w := worker.NewMockWorker(t) diff --git a/internal/execution/supervisor/adapter_rpc.go b/internal/execution/supervisor/adapter_rpc.go index ec40d5b..193179b 100644 --- a/internal/execution/supervisor/adapter_rpc.go +++ b/internal/execution/supervisor/adapter_rpc.go @@ -14,6 +14,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) // RpcConfig describes the configuration for the rpc interface. @@ -91,16 +92,27 @@ type rpcAdapter struct { config RpcConfig log *zap.Logger + + // sidecarCfg configures the worker-authored progress side-channel. + sidecarCfg progress.SidecarConfig + + // sidecar is the loopback HTTP listener for worker-authored progress + // events, injected into the worker's env as EVAL_PROGRESS_URL. It + // lives for this adapter's whole lifetime (one persistent worker can + // serve many requests), and is Bind/Unbind-ed around each Send call. + sidecar *progress.Sidecar } func newRpcAdapter( workerFactory AdapterWorkerFactoryFn, config RpcConfig, + sidecarCfg progress.SidecarConfig, log *zap.Logger, ) *rpcAdapter { return &rpcAdapter{ workerFactory: workerFactory, config: config, + sidecarCfg: sidecarCfg, log: log.Named("adapter_rpc"), } } @@ -113,7 +125,13 @@ func (a *rpcAdapter) Start( return errors.New("no worker factory provided") } - params.Env = buildEnv(params.Env, a.config) + sidecar, err := progress.NewSidecar(a.sidecarCfg, a.log) + if err != nil { + return fmt.Errorf("error starting progress sidecar: %w", err) + } + a.sidecar = sidecar + + params.Env = buildEnv(params.Env, a.config, sidecar.URL()) // create the worker worker, err := a.workerFactory(params) @@ -164,6 +182,18 @@ func (a *rpcAdapter) Send( return nil, errors.New("rpc client not available") } + if a.sidecar != nil { + // sendLock in the calling supervisor guarantees only one request + // is ever in flight per worker; UnbindAfterGrace closes the window + // between this call returning and the next one starting (after a + // short grace period, to give a fire-and-forget progress POST the + // worker dispatched just before returning its result a chance to + // still land), so a straggling POST can't be misattributed to an + // unrelated future request. + a.sidecar.Bind(method, progress.FromContext(ctx)) + defer a.sidecar.UnbindAfterGrace() + } + var result map[string]any ctx, cancel := context.WithTimeout(ctx, timeout) @@ -181,6 +211,12 @@ func (a *rpcAdapter) Stop() (ReleaseFunc, error) { return nil, errors.New("no worker provided") } + if a.sidecar != nil { + if err := a.sidecar.Close(); err != nil { + a.log.Warn("error closing progress sidecar", zap.Error(err)) + } + } + return stopWorker(a.worker) } @@ -283,7 +319,7 @@ func getIPCEndpoint(config IpcTransportConfig) string { } } -func buildEnv(env []string, config RpcConfig) []string { +func buildEnv(env []string, config RpcConfig, progressURL string) []string { if env == nil { env = make([]string, 0) } @@ -304,6 +340,10 @@ func buildEnv(env []string, config RpcConfig) []string { env = append(env, "EVAL_RPC_TCP_ADDRESS="+config.Tcp.Address) } + if progressURL != "" { + env = append(env, "EVAL_PROGRESS_URL="+progressURL) + } + return env } diff --git a/internal/execution/supervisor/adapter_rpc_test.go b/internal/execution/supervisor/adapter_rpc_test.go index 2ac8860..1357fe6 100644 --- a/internal/execution/supervisor/adapter_rpc_test.go +++ b/internal/execution/supervisor/adapter_rpc_test.go @@ -4,13 +4,17 @@ import ( "bytes" "context" "io" + "net/http" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) type rwc struct { @@ -38,6 +42,14 @@ func createRpcAdapter(t *testing.T) (*rpcAdapter, *worker.MockWorker) { config: RpcConfig{Transport: StdioTransport}, } + // Start (called by most tests using this helper) always spins up a + // real progress sidecar listener; close it so tests don't leak ports. + t.Cleanup(func() { + if adapter.sidecar != nil { + adapter.sidecar.Close() + } + }) + return adapter, w } @@ -142,6 +154,60 @@ func TestStdioAdapter_Stop_WaitForError(t *testing.T) { assert.ErrorIs(t, err, assert.AnError) } +func TestStdioAdapter_Start_InjectsProgressURL(t *testing.T) { + a, w := createRpcAdapter(t) + + var sp *worker.StartConfig + baseFactory := a.workerFactory + a.workerFactory = func(params worker.StartConfig) (worker.Worker, error) { + sp = ¶ms + return baseFactory(params) + } + + w.EXPECT().DuplexPipe().Return(newRwc(), nil) + w.EXPECT().Start(mock.Anything).Return(nil) + + err := a.Start(context.Background(), worker.StartConfig{}) + assert.NoError(t, err) + + assert.Contains(t, sp.Env, "EVAL_PROGRESS_URL="+a.sidecar.URL()) +} + +// TestStdioAdapter_Send_RelaysWorkerProgressEvents exercises the same +// Bind/UnbindAfterGrace path Send uses around the (separately, more fully) +// tested Sidecar, without needing a live RPC round trip - Send itself isn't +// otherwise exercised in this file (see the disabled tests below). +func TestStdioAdapter_Send_RelaysWorkerProgressEvents(t *testing.T) { + a, w := createRpcAdapter(t) + + w.EXPECT().DuplexPipe().Return(newRwc(), nil) + w.EXPECT().Start(mock.Anything).Return(nil) + + err := a.Start(context.Background(), worker.StartConfig{}) + assert.NoError(t, err) + + r := &recordingReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + // mirrors exactly what rpcAdapter.Send does with a.sidecar + a.sidecar.Bind("eval", progress.FromContext(ctx)) + defer a.sidecar.UnbindAfterGrace() + + resp, err := http.Post(a.sidecar.URL(), "application/json", strings.NewReader(`{"message":"checking correctness"}`)) + assert.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusAccepted, resp.StatusCode) + + assert.Eventually(t, func() bool { + return len(r.recorded()) == 1 + }, time.Second, 5*time.Millisecond, "expected the worker's progress event to be relayed") + + events := r.recorded() + assert.Equal(t, progress.StageProgress, events[0].Stage) + assert.Equal(t, "eval", events[0].Command) + assert.Equal(t, "checking correctness", events[0].Message) +} + // func TestStdioAdapter_Send(t *testing.T) { // a, w := createStdioAdapter(t) diff --git a/internal/execution/supervisor/adapter_test.go b/internal/execution/supervisor/adapter_test.go index c426f7a..1374d1b 100644 --- a/internal/execution/supervisor/adapter_test.go +++ b/internal/execution/supervisor/adapter_test.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) func TestDefaultAdapterFactory(t *testing.T) { @@ -16,9 +17,11 @@ func TestDefaultAdapterFactory(t *testing.T) { return w, nil } + factory := newDefaultAdapterFactory(progress.Config{}) + cases := []IOConfig{{Interface: FileIO}, {Interface: RpcIO}} for _, mode := range cases { - _, err := defaultAdapterFactory(workerFactory, mode, zap.NewNop()) + _, err := factory(workerFactory, mode, zap.NewNop()) assert.NoError(t, err) } @@ -31,7 +34,7 @@ func TestDefaultAdapterFactory_Fails(t *testing.T) { return w, nil } - _, err := defaultAdapterFactory(workerFactory, IOConfig{Interface: ""}, zap.NewNop()) + _, err := newDefaultAdapterFactory(progress.Config{})(workerFactory, IOConfig{Interface: ""}, zap.NewNop()) assert.ErrorIs(t, err, ErrUnsupportedIOInterface) } diff --git a/internal/execution/supervisor/supervisor.go b/internal/execution/supervisor/supervisor.go index f3e6587..f3d6b65 100644 --- a/internal/execution/supervisor/supervisor.go +++ b/internal/execution/supervisor/supervisor.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/worker" + "github.com/lambda-feedback/shimmy/internal/progress" ) type Supervisor interface { @@ -74,6 +75,11 @@ type Params struct { // is called when the supervisor needs to create a new worker. WorkerFactory WorkerFactoryFn + // Progress configures worker-authored progress event delivery (the + // EVAL_PROGRESS_URL side-channel). Only used when AdapterFactory is + // nil, since the default adapter factory is what wires it up. + Progress progress.Config + // Log is the logger to use for the supervisor Log *zap.Logger } @@ -99,7 +105,7 @@ func New(params Params) (Supervisor, error) { } if params.AdapterFactory == nil { - params.AdapterFactory = defaultAdapterFactory + params.AdapterFactory = newDefaultAdapterFactory(params.Progress) } createAdapter := func() (*workerRef, error) { @@ -165,12 +171,36 @@ func (s *WorkerSupervisor) Send( worker, err := s.acquireWorker(ctx) if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "We couldn't start the evaluation. Please try again.", + Error: err.Error(), + }) return nil, fmt.Errorf("failed to acquire worker: %w", err) } + progress.Emit(ctx, progress.Event{ + Stage: progress.StagePreparing, + Command: method, + Message: "Preparing your evaluation…", + }) // NOTICE: unconventional error handling ahead, as we need // to release the worker before returning the error. + progress.Emit(ctx, progress.Event{ + Stage: progress.StageEvaluating, + Command: method, + Message: "Evaluating your submission…", + }) resData, err := worker.Send(ctx, method, data, s.sendParams.Timeout) + if err != nil { + progress.Emit(ctx, progress.Event{ + Stage: progress.StageFailed, + Command: method, + Message: "Something went wrong while evaluating your answer. Please try again.", + Error: err.Error(), + }) + } release, releaseErr := s.releaseWorker() if releaseErr != nil { diff --git a/internal/execution/supervisor/supervisor_test.go b/internal/execution/supervisor/supervisor_test.go index 82bdb95..41630c1 100644 --- a/internal/execution/supervisor/supervisor_test.go +++ b/internal/execution/supervisor/supervisor_test.go @@ -9,6 +9,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution/supervisor" + "github.com/lambda-feedback/shimmy/internal/progress" ) func TestSupervisor_New_DefaultWorkerFactory(t *testing.T) { @@ -280,6 +281,101 @@ func TestSupervisor_Send_Fails(t *testing.T) { assert.NotNil(t, res) } +// MARK: - progress + +func TestSupervisor_Send_EmitsWorkerAcquiredAndRunning(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + resData := map[string]any{"result": "result"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(resData, nil) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + _, err = s.Send(ctx, "test", data) + assert.NoError(t, err) + + assert.Equal(t, []progress.Stage{ + progress.StagePreparing, + progress.StageEvaluating, + }, r.stages()) +} + +func TestSupervisor_Send_EmitsFailed_WhenAcquireFails(t *testing.T) { + mockFactory := func(supervisor.AdapterWorkerFactoryFn, supervisor.IOConfig, *zap.Logger) (supervisor.Adapter, error) { + return nil, assert.AnError + } + + s, err := createSupervisorWithFactory(supervisor.RpcIO, mockFactory) + assert.NoError(t, err) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + data := map[string]any{"data": "data"} + _, err = s.Send(ctx, "test", data) + assert.ErrorIs(t, err, assert.AnError) + + assert.Equal(t, []progress.Stage{progress.StageFailed}, r.stages()) +} + +func TestSupervisor_Send_EmitsFailed_WhenWorkerSendFails(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(nil, assert.AnError) + + r := &fakeReporter{} + ctx := progress.ContextWithReporter(context.Background(), r) + + _, err = s.Send(ctx, "test", data) + assert.ErrorIs(t, err, assert.AnError) + + assert.Equal(t, []progress.Stage{ + progress.StagePreparing, + progress.StageEvaluating, + progress.StageFailed, + }, r.stages()) +} + +func TestSupervisor_Send_NoReporterInContext_BehavesUnchanged(t *testing.T) { + s, a, err := createSupervisor(t, supervisor.RpcIO) + assert.NoError(t, err) + + data := map[string]any{"data": "data"} + resData := map[string]any{"result": "result"} + + a.EXPECT().Start(mock.Anything, mock.Anything).Return(nil) + a.EXPECT().Send(mock.Anything, "test", data, mock.Anything).Return(resData, nil) + + res, err := s.Send(context.Background(), "test", data) + assert.NoError(t, err) + assert.Equal(t, resData, res.Data) +} + +type fakeReporter struct { + events []progress.Event +} + +func (r *fakeReporter) Report(_ context.Context, evt progress.Event) { + r.events = append(r.events, evt) +} + +func (r *fakeReporter) stages() []progress.Stage { + stages := make([]progress.Stage, len(r.events)) + for i, evt := range r.events { + stages[i] = evt.Stage + } + return stages +} + // MARK: - mocks func createSupervisor(t *testing.T, mode supervisor.IOInterface) ( diff --git a/internal/progress/event.go b/internal/progress/event.go new file mode 100644 index 0000000..c362e25 --- /dev/null +++ b/internal/progress/event.go @@ -0,0 +1,68 @@ +package progress + +import "time" + +// Stage identifies a point in the lifecycle of an evaluation request that +// progress events can be emitted for. +type Stage string + +const ( + // StagePreparing indicates the evaluation environment is being set up + // (a worker is ready to receive work, whether freshly booted or reused + // from a warm pool). Deliberately named around what a student or + // teacher would recognise, not shimmy's internal "worker" concept. + StagePreparing Stage = "preparing" + + // StageEvaluating indicates the submission is being evaluated. + StageEvaluating Stage = "evaluating" + + // StageCompleted indicates feedback has been computed and is about + // to be returned to the caller. + StageCompleted Stage = "completed" + + // StageFailed indicates a terminal failure at any layer of the pipeline. + StageFailed Stage = "failed" + + // StageProgress indicates a custom, evaluation-function-authored + // progress update. Unlike the other stages, these are never emitted + // by shimmy itself — only relayed from a worker's local progress + // side-channel (see Sidecar). A worker cannot claim any other stage; + // the wire contract for that side-channel has no way to set Stage. + StageProgress Stage = "progress" +) + +// terminal reports whether the stage marks the end of an evaluation's +// progress event stream. At most one terminal event is delivered per +// Reporter instance. +func (s Stage) terminal() bool { + return s == StageCompleted || s == StageFailed +} + +// Event describes a single progress update for an evaluation request. +type Event struct { + // Stage is the lifecycle point this event describes. + Stage Stage + + // Command is the µEd command being processed (e.g. "eval", "preview"). + Command string + + // Message is a short, student/teacher-facing description of this + // event, safe to display as-is (e.g. "Evaluating your submission…"). + // It must never contain raw technical error detail — see Error. + Message string + + // Error carries raw technical error detail for StageFailed events, + // intended for logs/support diagnostics. Never display this to + // students or teachers directly; show Message instead. + Error string + + // Data is a free-form extension point. On StageCompleted it carries + // the evaluation's feedback payload (so a callbackUrl-supplying + // caller gets the final result, not just a status ping). On + // StageProgress it carries whatever the evaluation function attached + // to its custom event (see Sidecar). + Data map[string]any + + // Timestamp is set by Emit, not by callers. + Timestamp time.Time +} diff --git a/internal/progress/factory.go b/internal/progress/factory.go new file mode 100644 index 0000000..de56695 --- /dev/null +++ b/internal/progress/factory.go @@ -0,0 +1,113 @@ +package progress + +import ( + "fmt" + "net/http" + "net/url" + "time" + + "go.uber.org/fx" + "go.uber.org/zap" +) + +// defaultCallbackTimeout is used when Config.CallbackTimeout is unset. +const defaultCallbackTimeout = time.Second + +// Config is the configuration for outbound progress-callback delivery. +type Config struct { + // CallbackTimeout bounds a single progress callback POST. If unset + // (or <= 0), defaultCallbackTimeout is used. + CallbackTimeout time.Duration `conf:"callback_timeout"` + + // AllowedHosts, if non-empty, restricts callback URLs to these hosts. + // Entries may be an exact hostname (e.g. "api.example.com") or a + // "*.example.com" wildcard matching any subdomain. Empty means no + // host restriction — callback delivery is still subject to the + // private-network protection below. + AllowedHosts []string `conf:"allowed_hosts"` + + // AllowPrivateNetworks disables the default SSRF protection that + // refuses to dial loopback, link-local (including cloud metadata + // endpoints such as 169.254.169.254), and private (RFC1918/RFC4193) + // IP addresses, however the URL's hostname resolves. Only enable + // this if shimmy's callback targets are known to live on a private + // network you trust (e.g. a same-VPC service). + AllowPrivateNetworks bool `conf:"allow_private_networks"` + + // Sidecar bounds worker-authored progress events delivered via the + // EVAL_PROGRESS_URL side-channel (see sidecar.go), before they're + // relayed through the same outbound delivery path as shim-authored + // events. + Sidecar SidecarConfig `conf:"sidecar"` +} + +// Factory builds a per-request Reporter from caller-supplied callback +// coordinates. +type Factory interface { + // NewReporter returns a Reporter that delivers events to callbackURL, + // tagging each with correlationID. If callbackURL is empty, it returns + // (nil, nil) — the signal that progress reporting is disabled for this + // request. An error is returned only when callbackURL is non-empty but + // invalid. + NewReporter(callbackURL, correlationID string) (Reporter, error) +} + +type HTTPFactoryParams struct { + fx.In + + Config Config + Log *zap.Logger +} + +type HTTPFactory struct { + client *http.Client + timeout time.Duration + log *zap.Logger + allowedHosts []string +} + +var _ Factory = (*HTTPFactory)(nil) + +// NewHTTPFactory builds a Factory that delivers progress events as +// outbound HTTP POST requests. +// +// Since the callback URL is caller-supplied, delivery is guarded against +// SSRF by default: the underlying transport refuses to dial loopback, +// link-local, or private IP addresses (see Config.AllowPrivateNetworks), +// and Config.AllowedHosts can further restrict which hostnames are +// accepted at all. +func NewHTTPFactory(params HTTPFactoryParams) Factory { + timeout := params.Config.CallbackTimeout + if timeout <= 0 { + timeout = defaultCallbackTimeout + } + + client := &http.Client{} + if !params.Config.AllowPrivateNetworks { + client.Transport = newSSRFGuardedTransport() + } + + return &HTTPFactory{ + client: client, + timeout: timeout, + log: params.Log, + allowedHosts: params.Config.AllowedHosts, + } +} + +func (f *HTTPFactory) NewReporter(callbackURL, correlationID string) (Reporter, error) { + if callbackURL == "" { + return nil, nil + } + + u, err := url.ParseRequestURI(callbackURL) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") { + return nil, fmt.Errorf("invalid progress callback url: %q", callbackURL) + } + + if len(f.allowedHosts) > 0 && !hostAllowed(u.Hostname(), f.allowedHosts) { + return nil, fmt.Errorf("progress callback host %q is not in the allowed hosts list", u.Hostname()) + } + + return newHTTPReporter(f.client, callbackURL, correlationID, f.timeout, f.log.Named("progress")), nil +} diff --git a/internal/progress/factory_test.go b/internal/progress/factory_test.go new file mode 100644 index 0000000..3e6e095 --- /dev/null +++ b/internal/progress/factory_test.go @@ -0,0 +1,67 @@ +package progress + +import ( + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestFactory() *HTTPFactory { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: time.Second}, + Log: zap.NewNop(), + }) + return f.(*HTTPFactory) +} + +func TestHTTPFactory_NewReporter_EmptyURL_ReturnsNilReporterNoError(t *testing.T) { + f := newTestFactory() + + r, err := f.NewReporter("", "corr-1") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if r != nil { + t.Fatalf("expected nil reporter for empty callback url, got %v", r) + } +} + +func TestHTTPFactory_NewReporter_ValidURL_ReturnsReporter(t *testing.T) { + f := newTestFactory() + + r, err := f.NewReporter("https://example.com/callback", "corr-1") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if r == nil { + t.Fatalf("expected non-nil reporter for valid url") + } +} + +func TestHTTPFactory_NewReporter_InvalidURL_ReturnsError(t *testing.T) { + f := newTestFactory() + + cases := []string{ + "not-a-url", + "ftp://example.com/callback", + "://broken", + } + + for _, c := range cases { + if _, err := f.NewReporter(c, "corr-1"); err == nil { + t.Errorf("expected error for callback url %q, got nil", c) + } + } +} + +func TestNewHTTPFactory_DefaultsTimeoutWhenUnset(t *testing.T) { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{}, + Log: zap.NewNop(), + }).(*HTTPFactory) + + if f.timeout != defaultCallbackTimeout { + t.Errorf("expected default timeout %v, got %v", defaultCallbackTimeout, f.timeout) + } +} diff --git a/internal/progress/http_reporter.go b/internal/progress/http_reporter.go new file mode 100644 index 0000000..4d6fbb2 --- /dev/null +++ b/internal/progress/http_reporter.go @@ -0,0 +1,115 @@ +package progress + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "sync" + "time" + + "go.uber.org/zap" +) + +// payload is the JSON body POSTed to the callback URL for each event. +type payload struct { + CorrelationID string `json:"correlationId"` + Stage Stage `json:"stage"` + Command string `json:"command,omitempty"` + Message string `json:"message,omitempty"` + Error string `json:"error,omitempty"` + Data map[string]any `json:"data,omitempty"` + Timestamp time.Time `json:"timestamp"` +} + +// httpCallbackReporter delivers progress events as outbound HTTP POST +// requests to a caller-supplied URL. +type httpCallbackReporter struct { + client *http.Client + url string + correlationID string + timeout time.Duration + log *zap.Logger + + terminalOnce sync.Once +} + +var _ Reporter = (*httpCallbackReporter)(nil) + +func newHTTPReporter( + client *http.Client, + url string, + correlationID string, + timeout time.Duration, + log *zap.Logger, +) Reporter { + return &httpCallbackReporter{ + client: client, + url: url, + correlationID: correlationID, + timeout: timeout, + log: log, + } +} + +// Report POSTs evt to the configured callback URL. Delivery is best-effort: +// any error (invalid payload, dial failure, timeout, non-2xx response) is +// logged and swallowed — it must never fail or slow down the evaluation +// beyond the configured timeout. At most one terminal event (StageFailed +// or StageCompleted) is delivered per reporter instance, since both +// the supervisor and handler layers can independently detect failure. +func (r *httpCallbackReporter) Report(ctx context.Context, evt Event) { + if evt.Stage.terminal() { + sent := false + r.terminalOnce.Do(func() { + r.send(ctx, evt) + sent = true + }) + if !sent { + r.log.Debug("dropping duplicate terminal progress event", zap.String("stage", string(evt.Stage))) + } + return + } + + r.send(ctx, evt) +} + +func (r *httpCallbackReporter) send(ctx context.Context, evt Event) { + body, err := json.Marshal(payload{ + CorrelationID: r.correlationID, + Stage: evt.Stage, + Command: evt.Command, + Message: evt.Message, + Error: evt.Error, + Data: evt.Data, + Timestamp: evt.Timestamp, + }) + if err != nil { + r.log.Warn("failed to marshal progress event", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + + ctx, cancel := context.WithTimeout(ctx, r.timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.url, bytes.NewReader(body)) + if err != nil { + r.log.Warn("failed to build progress callback request", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + req.Header.Set("Content-Type", "application/json") + + resp, err := r.client.Do(req) + if err != nil { + r.log.Warn("progress callback delivery failed", zap.String("stage", string(evt.Stage)), zap.Error(err)) + return + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + r.log.Warn("progress callback returned non-2xx status", + zap.String("stage", string(evt.Stage)), + zap.Int("status", resp.StatusCode), + ) + } +} diff --git a/internal/progress/http_reporter_test.go b/internal/progress/http_reporter_test.go new file mode 100644 index 0000000..2d7ee79 --- /dev/null +++ b/internal/progress/http_reporter_test.go @@ -0,0 +1,123 @@ +package progress + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestReporter(t *testing.T, url string, timeout time.Duration) *httpCallbackReporter { + t.Helper() + return newHTTPReporter(&http.Client{}, url, "corr-1", timeout, zap.NewNop()).(*httpCallbackReporter) +} + +func TestHTTPCallbackReporter_Report_DeliversPayload(t *testing.T) { + var mu sync.Mutex + var received []payload + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var p payload + if err := json.NewDecoder(r.Body).Decode(&p); err != nil { + t.Errorf("failed to decode payload: %v", err) + } + mu.Lock() + received = append(received, p) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{Stage: StageEvaluating, Command: "eval"}) + + mu.Lock() + defer mu.Unlock() + if len(received) != 1 { + t.Fatalf("expected 1 request, got %d", len(received)) + } + if received[0].CorrelationID != "corr-1" { + t.Errorf("expected correlationId %q, got %q", "corr-1", received[0].CorrelationID) + } + if received[0].Stage != StageEvaluating { + t.Errorf("expected stage %q, got %q", StageEvaluating, received[0].Stage) + } +} + +func TestHTTPCallbackReporter_Report_TerminalEventDeliveredOnlyOnce(t *testing.T) { + var mu sync.Mutex + var count int + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + count++ + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + + // simulate both the supervisor and handler layers independently + // detecting failure and trying to emit a terminal event + r.Report(context.Background(), Event{Stage: StageFailed, Message: "boot failed"}) + r.Report(context.Background(), Event{Stage: StageFailed, Message: "handler backstop"}) + r.Report(context.Background(), Event{Stage: StageCompleted}) + + mu.Lock() + defer mu.Unlock() + if count != 1 { + t.Fatalf("expected exactly 1 terminal event delivered, got %d", count) + } +} + +func TestHTTPCallbackReporter_Report_NonTerminalEventsAllDelivered(t *testing.T) { + var mu sync.Mutex + var count int + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + count++ + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, time.Second) + r.Report(context.Background(), Event{Stage: StagePreparing}) + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + mu.Lock() + defer mu.Unlock() + if count != 2 { + t.Fatalf("expected 2 non-terminal events delivered, got %d", count) + } +} + +func TestHTTPCallbackReporter_Report_SlowReceiver_BoundedByTimeout(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(200 * time.Millisecond) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + r := newTestReporter(t, srv.URL, 20*time.Millisecond) + + start := time.Now() + r.Report(context.Background(), Event{Stage: StageEvaluating}) + elapsed := time.Since(start) + + if elapsed > 150*time.Millisecond { + t.Errorf("expected Report to return promptly bounded by timeout, took %v", elapsed) + } +} + +func TestHTTPCallbackReporter_Report_UnreachableURL_DoesNotPanic(t *testing.T) { + r := newTestReporter(t, "http://127.0.0.1:0", 50*time.Millisecond) + r.Report(context.Background(), Event{Stage: StageEvaluating}) +} diff --git a/internal/progress/reporter.go b/internal/progress/reporter.go new file mode 100644 index 0000000..e21f051 --- /dev/null +++ b/internal/progress/reporter.go @@ -0,0 +1,44 @@ +package progress + +import ( + "context" + "time" +) + +// Reporter delivers progress events for a single evaluation request. +type Reporter interface { + // Report emits a single event. Implementations MUST NOT return an + // error to the caller and MUST apply their own bounded timeout — + // progress delivery must never fail or slow down the evaluation. + Report(ctx context.Context, evt Event) +} + +type contextKey int + +var reporterKey = contextKey(0) + +// ContextWithReporter returns a copy of ctx carrying the given Reporter. +func ContextWithReporter(ctx context.Context, r Reporter) context.Context { + return context.WithValue(ctx, reporterKey, r) +} + +// FromContext returns the Reporter attached to ctx, or nil if none is +// attached. A nil Reporter is the expected, common case: most requests +// don't opt in to progress reporting. +func FromContext(ctx context.Context) Reporter { + r, _ := ctx.Value(reporterKey).(Reporter) + return r +} + +// Emit is the call-site convenience for reporting a progress event. It is +// a silent no-op when no Reporter is attached to ctx, which is what makes +// progress reporting purely opt-in/additive. +func Emit(ctx context.Context, evt Event) { + r := FromContext(ctx) + if r == nil { + return + } + + evt.Timestamp = time.Now().UTC() + r.Report(ctx, evt) +} diff --git a/internal/progress/reporter_test.go b/internal/progress/reporter_test.go new file mode 100644 index 0000000..3aee6b0 --- /dev/null +++ b/internal/progress/reporter_test.go @@ -0,0 +1,62 @@ +package progress + +import ( + "context" + "sync" + "testing" +) + +// recordingReporter is a test double shared across this package's test +// files. It's safe for concurrent use since sidecar_test.go exercises it +// from the sidecar's detached relay goroutine as well as the test +// goroutine polling for results. +type recordingReporter struct { + mu sync.Mutex + events []Event +} + +func (r *recordingReporter) Report(_ context.Context, evt Event) { + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, evt) +} + +// recorded returns a snapshot of the events received so far. +func (r *recordingReporter) recorded() []Event { + r.mu.Lock() + defer r.mu.Unlock() + return append([]Event(nil), r.events...) +} + +func TestEmit_NoReporterInContext_NoOp(t *testing.T) { + // must not panic, must not do anything observable + Emit(context.Background(), Event{Stage: StageEvaluating}) +} + +func TestEmit_WithReporter_DeliversEventAndSetsTimestamp(t *testing.T) { + r := &recordingReporter{} + ctx := ContextWithReporter(context.Background(), r) + + Emit(ctx, Event{Stage: StagePreparing, Command: "eval"}) + + events := r.recorded() + if len(events) != 1 { + t.Fatalf("expected 1 event, got %d", len(events)) + } + evt := events[0] + if evt.Stage != StagePreparing { + t.Errorf("expected stage %q, got %q", StagePreparing, evt.Stage) + } + if evt.Command != "eval" { + t.Errorf("expected command %q, got %q", "eval", evt.Command) + } + if evt.Timestamp.IsZero() { + t.Errorf("expected Emit to set a non-zero timestamp") + } +} + +func TestFromContext_NoReporter_ReturnsNil(t *testing.T) { + if r := FromContext(context.Background()); r != nil { + t.Errorf("expected nil reporter, got %v", r) + } +} diff --git a/internal/progress/sidecar.go b/internal/progress/sidecar.go new file mode 100644 index 0000000..7d73f82 --- /dev/null +++ b/internal/progress/sidecar.go @@ -0,0 +1,293 @@ +package progress + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "net/http" + "strings" + "sync" + "time" + + "go.uber.org/zap" +) + +const ( + defaultSidecarMaxBodyBytes int64 = 16 * 1024 + defaultSidecarMaxEventsPerSpan = 50 + defaultSidecarBurstSize = 5 + defaultSidecarMinEventInterval = 10 * time.Millisecond + defaultSidecarUnbindGracePeriod = 250 * time.Millisecond +) + +// SidecarConfig bounds abuse of the worker-authored progress side-channel. +// Since EVAL_PROGRESS_URL is reachable by arbitrary (and, under sandboxing, +// untrusted) worker code, delivery to the real callbackUrl must stay bounded +// regardless of how the worker behaves. +type SidecarConfig struct { + // MaxBodyBytes caps the size of a single progress event POST body. + // If unset (<= 0), defaultSidecarMaxBodyBytes is used. + MaxBodyBytes int64 `conf:"max_body_bytes"` + + // MaxEventsPerSpan caps how many progress events a single evaluation + // span (the window between Bind and the next Bind/Unbind) may relay. + // If unset (<= 0), defaultSidecarMaxEventsPerSpan is used. + MaxEventsPerSpan int `conf:"max_events_per_span"` + + // BurstSize is how many events at the start of a span are exempt from + // MinEventInterval spacing, so a handful of legitimate back-to-back + // checkpoints (e.g. a fast evaluation reporting progress at several + // points microseconds to a few ms apart) aren't rate-limited just + // because they arrive faster than any fixed spacing could accommodate. + // MinEventInterval spacing only applies once the burst is used up. + // Still bounded by MaxEventsPerSpan. If unset (== 0), + // defaultSidecarBurstSize is used; a negative value explicitly + // disables the burst allowance (spacing applies from the first event). + BurstSize int `conf:"burst_size"` + + // MinEventInterval enforces a minimum spacing between accepted events + // once a span's BurstSize allowance is used up. If unset (<= 0), + // defaultSidecarMinEventInterval is used. + MinEventInterval time.Duration `conf:"min_event_interval"` + + // UnbindGracePeriod delays detaching the bound reporter after a span + // ends, so a worker-authored progress POST that was already in flight + // (e.g. dispatched fire-and-forget just before the worker returned its + // result) still has a window to arrive and be relayed, instead of + // racing the RPC response back to shim. If unset (<= 0), + // defaultSidecarUnbindGracePeriod is used. + UnbindGracePeriod time.Duration `conf:"unbind_grace_period"` +} + +func (c SidecarConfig) withDefaults() SidecarConfig { + if c.MaxBodyBytes <= 0 { + c.MaxBodyBytes = defaultSidecarMaxBodyBytes + } + if c.MaxEventsPerSpan <= 0 { + c.MaxEventsPerSpan = defaultSidecarMaxEventsPerSpan + } + if c.BurstSize < 0 { + c.BurstSize = 0 + } else if c.BurstSize == 0 { + c.BurstSize = defaultSidecarBurstSize + } + if c.MinEventInterval <= 0 { + c.MinEventInterval = defaultSidecarMinEventInterval + } + if c.UnbindGracePeriod <= 0 { + c.UnbindGracePeriod = defaultSidecarUnbindGracePeriod + } + return c +} + +// sidecarPayload is the JSON body a worker POSTs to report a custom +// progress event. There is deliberately no "stage" field: a worker can +// never claim any stage other than StageProgress, which the sidecar +// hardcodes itself. Unknown fields (including a "stage" a worker might +// send anyway) are silently ignored by json.Decode, never merged in. +type sidecarPayload struct { + Message string `json:"message"` + Data map[string]any `json:"data,omitempty"` +} + +// Sidecar is a loopback-only HTTP listener that accepts worker-authored +// progress events and relays them, best-effort, through whichever Reporter +// is currently Bind-ed to it. It is the counterpart, on the inbound side, +// to the outbound delivery in http_reporter.go: since it only ever binds +// to 127.0.0.1, it needs no SSRF guarding, but it does need its own abuse +// limits, since the worker producing events may be untrusted. +// +// Its lifetime differs by adapter: for a persistent RPC worker, one Sidecar +// lives for the worker's whole lifetime and is Bind/Unbind-ed around each +// request; for the transient file interface, one Sidecar is created and +// Closed per request. +type Sidecar struct { + cfg SidecarConfig + log *zap.Logger + + listener net.Listener + server *http.Server + + mu sync.Mutex + command string + reporter Reporter + count int + lastSent time.Time + generation uint64 +} + +// NewSidecar starts a loopback HTTP listener on an OS-assigned port. +func NewSidecar(cfg SidecarConfig, log *zap.Logger) (*Sidecar, error) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("failed to start progress sidecar listener: %w", err) + } + + s := &Sidecar{ + cfg: cfg.withDefaults(), + log: log.Named("progress_sidecar"), + listener: ln, + } + + s.server = &http.Server{ + Handler: http.HandlerFunc(s.handle), + ReadHeaderTimeout: 2 * time.Second, + ReadTimeout: 2 * time.Second, + WriteTimeout: 2 * time.Second, + } + + go func() { + if err := s.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + s.log.Warn("progress sidecar listener stopped unexpectedly", zap.Error(err)) + } + }() + + return s, nil +} + +// URL returns the sidecar's loopback address, suitable for EVAL_PROGRESS_URL. +func (s *Sidecar) URL() string { + return "http://" + s.listener.Addr().String() +} + +// Bind associates command/reporter with the sidecar for the duration of one +// evaluation span, resetting rate-limit state so a fresh span isn't +// poisoned by the previous request's usage. Call at the start of an +// adapter's Send. A nil reporter behaves like Unbind. +func (s *Sidecar) Bind(command string, reporter Reporter) { + s.mu.Lock() + defer s.mu.Unlock() + + s.generation++ + s.command = command + s.reporter = reporter + s.count = 0 + s.lastSent = time.Time{} +} + +// Unbind detaches the current reporter immediately, so any subsequent POST +// (e.g. a straggler arriving after the bound request has already returned) +// is rejected with 503 rather than misattributed to a future, unrelated +// request. +func (s *Sidecar) Unbind() { + s.mu.Lock() + defer s.mu.Unlock() + + s.generation++ + s.command = "" + s.reporter = nil +} + +// UnbindAfterGrace schedules the detach for after cfg.UnbindGracePeriod +// instead of doing it immediately, without blocking the caller. This gives +// a worker-authored progress POST dispatched fire-and-forget just before +// the RPC response reached shim a window to still arrive and be relayed, +// rather than losing the race against Unbind and being rejected with 503. +// +// If a new span is Bind-ed (or explicitly Unbind-ed) before the grace +// period elapses, this is a no-op: the generation captured at schedule time +// will no longer match, so the stale detach never fires and never clobbers +// the newer span. +func (s *Sidecar) UnbindAfterGrace() { + s.mu.Lock() + gen := s.generation + grace := s.cfg.UnbindGracePeriod + s.mu.Unlock() + + if grace <= 0 { + s.Unbind() + return + } + + time.AfterFunc(grace, func() { + s.mu.Lock() + defer s.mu.Unlock() + + if s.generation != gen { + return + } + + s.generation++ + s.command = "" + s.reporter = nil + }) +} + +// Close shuts down the sidecar's listener. It does not wait for any +// in-flight relayed events (those run detached from the listener, see +// handle) — consistent with progress delivery never blocking shutdown. +func (s *Sidecar) Close() error { + return s.server.Close() +} + +func (s *Sidecar) handle(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + + r.Body = http.MaxBytesReader(w, r.Body, s.cfg.MaxBodyBytes) + + var body sidecarPayload + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + w.WriteHeader(http.StatusRequestEntityTooLarge) + return + } + w.WriteHeader(http.StatusBadRequest) + return + } + + if strings.TrimSpace(body.Message) == "" { + w.WriteHeader(http.StatusBadRequest) + return + } + + command, reporter, status := s.accept() + if status != 0 { + w.WriteHeader(status) + return + } + + w.WriteHeader(http.StatusAccepted) + + evt := Event{ + Stage: StageProgress, + Command: command, + Message: body.Message, + Data: body.Data, + } + + // Relay detached from the inbound request: the worker's POST must + // never be held open for the outbound callbackUrl delivery, which has + // its own bounded timeout inside Report. + go reporter.Report(context.Background(), evt) +} + +// accept reports whether a new event may be relayed right now, applying +// the bound reporter check and the abuse limits. status is 0 on success, +// or the HTTP status to reject the request with otherwise. +func (s *Sidecar) accept() (command string, reporter Reporter, status int) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.reporter == nil { + return "", nil, http.StatusServiceUnavailable + } + + now := time.Now() + if s.count >= s.cfg.MaxEventsPerSpan { + return "", nil, http.StatusTooManyRequests + } + if s.count >= s.cfg.BurstSize && !s.lastSent.IsZero() && now.Sub(s.lastSent) < s.cfg.MinEventInterval { + return "", nil, http.StatusTooManyRequests + } + + s.count++ + s.lastSent = now + + return s.command, s.reporter, 0 +} diff --git a/internal/progress/sidecar_test.go b/internal/progress/sidecar_test.go new file mode 100644 index 0000000..35548ac --- /dev/null +++ b/internal/progress/sidecar_test.go @@ -0,0 +1,232 @@ +package progress + +import ( + "bytes" + "net/http" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +func newTestSidecar(t *testing.T, cfg SidecarConfig) *Sidecar { + t.Helper() + s, err := NewSidecar(cfg, zap.NewNop()) + if err != nil { + t.Fatalf("failed to start sidecar: %v", err) + } + t.Cleanup(func() { s.Close() }) + return s +} + +func postSidecar(t *testing.T, s *Sidecar, body string) *http.Response { + t.Helper() + resp, err := http.Post(s.URL(), "application/json", bytes.NewBufferString(body)) + if err != nil { + t.Fatalf("failed to POST to sidecar: %v", err) + } + defer resp.Body.Close() + return resp +} + +// waitForEvents polls until r has at least n events or the timeout expires, +// since the sidecar relays events in a detached goroutine. +func waitForEvents(t *testing.T, r *recordingReporter, n int) []Event { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if events := r.recorded(); len(events) >= n { + return events + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %d events, got %d", n, len(r.recorded())) + return nil +} + +func TestSidecar_Accept_RelaysEventThroughBoundReporter(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + r := &recordingReporter{} + s.Bind("eval", r) + + resp := postSidecar(t, s, `{"message":"checking correctness…","data":{"step":2}}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected 202, got %d", resp.StatusCode) + } + + events := waitForEvents(t, r, 1) + evt := events[0] + if evt.Stage != StageProgress { + t.Errorf("expected stage %q, got %q", StageProgress, evt.Stage) + } + if evt.Command != "eval" { + t.Errorf("expected command %q, got %q", "eval", evt.Command) + } + if evt.Message != "checking correctness…" { + t.Errorf("unexpected message %q", evt.Message) + } + if evt.Data["step"] != float64(2) { + t.Errorf("expected data.step=2, got %v", evt.Data["step"]) + } +} + +func TestSidecar_IgnoresWorkerSuppliedStage(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + r := &recordingReporter{} + s.Bind("eval", r) + + resp := postSidecar(t, s, `{"message":"trying to spoof","stage":"completed"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected 202, got %d", resp.StatusCode) + } + + events := waitForEvents(t, r, 1) + if events[0].Stage != StageProgress { + t.Errorf("worker-supplied stage must be ignored, got %q", events[0].Stage) + } +} + +func TestSidecar_RejectsEmptyMessage(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + + resp := postSidecar(t, s, `{"message":""}`) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", resp.StatusCode) + } +} + +func TestSidecar_RejectsMalformedJSON(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + + resp := postSidecar(t, s, `not json`) + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("expected 400, got %d", resp.StatusCode) + } +} + +func TestSidecar_RejectsOversizedBody(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxBodyBytes: 16}) + s.Bind("eval", &recordingReporter{}) + + body := `{"message":"` + strings.Repeat("x", 100) + `"}` + resp := postSidecar(t, s, body) + if resp.StatusCode != http.StatusRequestEntityTooLarge { + t.Fatalf("expected 413, got %d", resp.StatusCode) + } +} + +func TestSidecar_RateLimit_MaxEventsPerSpan(t *testing.T) { + // MinEventInterval is small (not disabled - 0 means "use the default") + // and slept past between POSTs, so only MaxEventsPerSpan is under test. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 1, MinEventInterval: time.Millisecond}) + r := &recordingReporter{} + s.Bind("eval", r) + + first := postSidecar(t, s, `{"message":"one"}`) + if first.StatusCode != http.StatusAccepted { + t.Fatalf("expected first event accepted (202), got %d", first.StatusCode) + } + + time.Sleep(5 * time.Millisecond) + + second := postSidecar(t, s, `{"message":"two"}`) + if second.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected second event rate limited (429), got %d", second.StatusCode) + } +} + +func TestSidecar_RateLimit_MinEventInterval(t *testing.T) { + // BurstSize disabled so the very first event is already subject to + // interval spacing, isolating what this test exercises. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: -1, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + first := postSidecar(t, s, `{"message":"one"}`) + if first.StatusCode != http.StatusAccepted { + t.Fatalf("expected first event accepted (202), got %d", first.StatusCode) + } + + second := postSidecar(t, s, `{"message":"two"}`) + if second.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected second event rate limited (429) by min interval, got %d", second.StatusCode) + } +} + +func TestSidecar_Burst_AllowsCloselySpacedEventsWithinBurst(t *testing.T) { + // A large MinEventInterval would reject any second event immediately - + // unless it falls within the burst allowance, which is what this + // exercises: events 2 and 3 land inside BurstSize and must be accepted + // even though far less than MinEventInterval separates them. + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: 3, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + for i, msg := range []string{"one", "two", "three"} { + resp := postSidecar(t, s, `{"message":"`+msg+`"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected burst event %d accepted (202), got %d", i+1, resp.StatusCode) + } + } +} + +func TestSidecar_Burst_ThenEnforcesMinEventInterval(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 100, BurstSize: 2, MinEventInterval: time.Hour}) + r := &recordingReporter{} + s.Bind("eval", r) + + for i, msg := range []string{"one", "two"} { + resp := postSidecar(t, s, `{"message":"`+msg+`"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected burst event %d accepted (202), got %d", i+1, resp.StatusCode) + } + } + + third := postSidecar(t, s, `{"message":"three"}`) + if third.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected event past burst allowance rate limited (429), got %d", third.StatusCode) + } +} + +func TestSidecar_Bind_ResetsRateLimitState(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{MaxEventsPerSpan: 1, MinEventInterval: time.Millisecond}) + r1 := &recordingReporter{} + s.Bind("eval", r1) + + postSidecar(t, s, `{"message":"one"}`) + exhausted := postSidecar(t, s, `{"message":"two"}`) + if exhausted.StatusCode != http.StatusTooManyRequests { + t.Fatalf("expected span to be exhausted (429), got %d", exhausted.StatusCode) + } + + r2 := &recordingReporter{} + s.Bind("eval", r2) + + resp := postSidecar(t, s, `{"message":"fresh span"}`) + if resp.StatusCode != http.StatusAccepted { + t.Fatalf("expected fresh span to accept after re-Bind (202), got %d", resp.StatusCode) + } +} + +func TestSidecar_Unbound_Returns503(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + + resp := postSidecar(t, s, `{"message":"nobody home"}`) + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("expected 503 with no bound reporter, got %d", resp.StatusCode) + } +} + +func TestSidecar_Unbind_Returns503(t *testing.T) { + s := newTestSidecar(t, SidecarConfig{}) + s.Bind("eval", &recordingReporter{}) + s.Unbind() + + resp := postSidecar(t, s, `{"message":"straggler"}`) + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("expected 503 after Unbind, got %d", resp.StatusCode) + } +} diff --git a/internal/progress/ssrf.go b/internal/progress/ssrf.go new file mode 100644 index 0000000..3104dc3 --- /dev/null +++ b/internal/progress/ssrf.go @@ -0,0 +1,74 @@ +package progress + +import ( + "context" + "fmt" + "net" + "net/http" + "strings" +) + +// isDisallowedIP reports whether ip must never be a target for an outbound +// progress callback: loopback, link-local (this also covers cloud metadata +// endpoints such as AWS's 169.254.169.254), private (RFC1918/RFC4193), +// unspecified, and multicast addresses. +func isDisallowedIP(ip net.IP) bool { + return ip.IsLoopback() || + ip.IsLinkLocalUnicast() || + ip.IsLinkLocalMulticast() || + ip.IsInterfaceLocalMulticast() || + ip.IsMulticast() || + ip.IsUnspecified() || + ip.IsPrivate() +} + +// hostAllowed reports whether host matches one of the allowed patterns. +// A pattern is either an exact hostname (e.g. "api.example.com") or a +// "*.example.com" wildcard matching any subdomain of example.com (but not +// example.com itself, which must be listed separately if intended). +func hostAllowed(host string, allowed []string) bool { + host = strings.ToLower(strings.TrimSuffix(host, ".")) + for _, pattern := range allowed { + pattern = strings.ToLower(pattern) + if pattern == host { + return true + } + if suffix, ok := strings.CutPrefix(pattern, "*."); ok && strings.HasSuffix(host, "."+suffix) { + return true + } + } + return false +} + +// newSSRFGuardedTransport returns an http.Transport that resolves DNS +// itself and refuses to dial any IP address isDisallowedIP flags, rather +// than trusting the request's literal hostname string. Checking the +// hostname alone would miss the common bypass of pointing an +// innocent-looking domain at a private or link-local address. +func newSSRFGuardedTransport() *http.Transport { + transport := http.DefaultTransport.(*http.Transport).Clone() + + dialer := &net.Dialer{} + transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + host, port, err := net.SplitHostPort(addr) + if err != nil { + return nil, err + } + + ips, err := net.DefaultResolver.LookupIP(ctx, "ip", host) + if err != nil { + return nil, err + } + + for _, ip := range ips { + if isDisallowedIP(ip) { + continue + } + return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) + } + + return nil, fmt.Errorf("host %q resolves only to disallowed private/loopback/link-local addresses", host) + } + + return transport +} diff --git a/internal/progress/ssrf_test.go b/internal/progress/ssrf_test.go new file mode 100644 index 0000000..2ef5ce0 --- /dev/null +++ b/internal/progress/ssrf_test.go @@ -0,0 +1,146 @@ +package progress + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "go.uber.org/zap" +) + +func TestIsDisallowedIP(t *testing.T) { + disallowed := []string{ + "127.0.0.1", // loopback + "::1", // loopback (v6) + "169.254.169.254", // link-local: cloud metadata endpoint + "fe80::1", // link-local (v6) + "10.0.0.1", // private RFC1918 + "172.16.0.1", // private RFC1918 + "192.168.1.1", // private RFC1918 + "fc00::1", // private RFC4193 + "0.0.0.0", // unspecified + "224.0.0.1", // multicast + } + for _, s := range disallowed { + ip := net.ParseIP(s) + if ip == nil { + t.Fatalf("failed to parse test IP %q", s) + } + if !isDisallowedIP(ip) { + t.Errorf("expected %q to be disallowed", s) + } + } + + allowed := []string{ + "8.8.8.8", + "1.1.1.1", + "93.184.216.34", + } + for _, s := range allowed { + ip := net.ParseIP(s) + if ip == nil { + t.Fatalf("failed to parse test IP %q", s) + } + if isDisallowedIP(ip) { + t.Errorf("expected %q to be allowed", s) + } + } +} + +func TestHostAllowed(t *testing.T) { + allowed := []string{"api.example.com", "*.example.org"} + + cases := []struct { + host string + want bool + }{ + {"api.example.com", true}, + {"API.EXAMPLE.COM", true}, + {"other.example.com", false}, + {"foo.example.org", true}, + {"a.b.example.org", true}, + {"example.org", false}, // bare domain not covered by wildcard + {"evil.com", false}, + } + + for _, c := range cases { + if got := hostAllowed(c.host, allowed); got != c.want { + t.Errorf("hostAllowed(%q, %v) = %v, want %v", c.host, allowed, got, c.want) + } + } +} + +func TestHTTPFactory_DefaultBlocksLoopbackDelivery(t *testing.T) { + var received bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: 500 * time.Millisecond}, + Log: zap.NewNop(), + }) + + r, err := f.NewReporter(srv.URL, "corr-1") + if err != nil { + t.Fatalf("expected NewReporter to succeed (block happens at delivery time), got %v", err) + } + + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + if received { + t.Errorf("expected delivery to a loopback address to be blocked by default") + } +} + +func TestHTTPFactory_AllowPrivateNetworks_PermitsLoopbackDelivery(t *testing.T) { + var received bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{CallbackTimeout: time.Second, AllowPrivateNetworks: true}, + Log: zap.NewNop(), + }) + + r, err := f.NewReporter(srv.URL, "corr-1") + if err != nil { + t.Fatalf("expected NewReporter to succeed, got %v", err) + } + + r.Report(context.Background(), Event{Stage: StageEvaluating}) + + if !received { + t.Errorf("expected delivery to succeed with AllowPrivateNetworks: true") + } +} + +func TestHTTPFactory_AllowedHosts_RejectsUnlistedHost(t *testing.T) { + f := NewHTTPFactory(HTTPFactoryParams{ + Config: Config{ + CallbackTimeout: time.Second, + AllowedHosts: []string{"good.example.com"}, + }, + Log: zap.NewNop(), + }) + + if _, err := f.NewReporter("https://evil.example.com/hook", "corr-1"); err == nil { + t.Errorf("expected an error for a host not in AllowedHosts") + } + + r, err := f.NewReporter("https://good.example.com/hook", "corr-1") + if err != nil { + t.Fatalf("expected no error for an allowed host, got %v", err) + } + if r == nil { + t.Fatalf("expected a non-nil reporter for an allowed host") + } +} diff --git a/runtime/mued.go b/runtime/mued.go index 24e8b2f..ca33c3a 100644 --- a/runtime/mued.go +++ b/runtime/mued.go @@ -34,6 +34,16 @@ type MuEdEvaluateRequest struct { Task *MuEdTask `json:"task"` Configuration *MuEdConfiguration `json:"configuration"` PreSubmissionFeedback *MuEdPreSubmissionFeedback `json:"preSubmissionFeedback"` + + // CallbackUrl is the µEd spec's optional HTTPS callback URL (see + // https://mued.org/spec, EvaluateRequest.callbackUrl). The spec + // describes it for asynchronous final-result delivery (the service + // may return 202 Accepted and POST the result here later); shimmy + // doesn't implement that 202 flow, but reuses this same field as the + // target for progress events, since both describe "send updates + // about this request to this URL" and a caller shouldn't need a + // shimmy-specific header for something the spec already defines. + CallbackUrl *string `json:"callbackUrl"` } var SupportedMuEdVersions = []string{"0.1.0"} diff --git a/runtime/runtime.go b/runtime/runtime.go index a29ce92..8f8ee8f 100644 --- a/runtime/runtime.go +++ b/runtime/runtime.go @@ -7,6 +7,7 @@ import ( "go.uber.org/zap" "github.com/lambda-feedback/shimmy/internal/execution" + "github.com/lambda-feedback/shimmy/internal/progress" ) // Runtime is the interface for a runtime. @@ -46,6 +47,11 @@ type RuntimeParams struct { // Config is the config for the underlying runtime manager Config Config + // Progress configures worker-authored progress event delivery (the + // EVAL_PROGRESS_URL side-channel). Provided by handler.Module, shared + // with the outbound callbackUrl delivery configuration. + Progress progress.Config + // Log is the logger to use for the runtime Log *zap.Logger } @@ -53,9 +59,10 @@ type RuntimeParams struct { // NewRuntime creates a new runtime. func NewRuntime(params RuntimeParams) (Runtime, error) { dispatcher, err := execution.NewDispatcher(Params{ - Context: params.Context, - Config: params.Config, - Log: params.Log, + Context: params.Context, + Config: params.Config, + Progress: params.Progress, + Log: params.Log, }) if err != nil { return nil, err