snapshot of https://github.com/ory/kratos.git (master, 2026-09-15; blobs > 131072B omitted)

oryx/jsonnetsecure/jsonnet_pool.go

@@ -0,0 +1,313 @@
+// Copyright © 2024 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package jsonnetsecure
+
+// Known limitations/edge cases:
+// - The child process exiting early (e.g. crashing) or getting killed (e.g. reaching some OS limit)
+//   is not detected and no error will be returned in this case from `eval()`.
+// - Misbehaving jsonnet scripts in the middle of a batch being passed to the child process for evaluation may result in
+//   no error (as mentioned above), and other valid scripts in this batch may result
+//   in an error (because the output from the child process is truncated).
+//
+// Possible remediations:
+// - Do not pass a batch of scripts to a worker, only pass one script at a time (to isolate misbehaving scripts)
+// - Validate that the output is valid JSON (to detect truncated output)
+// - Detect the child process exiting (to return an error)
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"io"
+	"math"
+	"os/exec"
+	"strings"
+	"time"
+
+	"github.com/jackc/puddle/v2"
+	"github.com/pkg/errors"
+	"go.opentelemetry.io/otel/attribute"
+	semconv "go.opentelemetry.io/otel/semconv/v1.27.0"
+	"go.opentelemetry.io/otel/trace"
+
+	"github.com/ory/x/otelx"
+)
+
+const (
+	KiB                = 1024
+	jsonnetOutputLimit = 256 * KiB
+	jsonnetErrLimit    = 1 * KiB
+)
+
+type (
+	processPoolVM struct {
+		path   string
+		args   []string
+		ctx    context.Context
+		params processParameters
+		pool   *pool
+	}
+	Pool interface {
+		Close()
+		// Stat returns a snapshot of the worker pool's statistics.
+		Stat() *puddle.Stat
+		private()
+	}
+	pool struct {
+		puddle *puddle.Pool[worker]
+	}
+	worker struct {
+		cmd    *exec.Cmd
+		stdin  chan<- []byte
+		stdout <-chan string
+		stderr <-chan string
+	}
+	contextKeyType string
+)
+
+var (
+	ErrProcessPoolClosed = errors.New("jsonnetsecure: process pool closed")
+
+	_ VM   = (*processPoolVM)(nil)
+	_ Pool = (*pool)(nil)
+
+	contextValuePath contextKeyType = "argc"
+	contextValueArgs contextKeyType = "argv"
+)
+
+func NewProcessPool(size int) Pool {
+	size = max(5, min(size, math.MaxInt32))
+	pud, err := puddle.NewPool(&puddle.Config[worker]{
+		MaxSize:     int32(size), //nolint:gosec // disable G115 // because of the previous min/max, 5 <= size <= math.MaxInt32
+		Constructor: newWorker,
+		Destructor:  worker.destroy,
+	})
+	if err != nil {
+		panic(err) // this should never happen, see implementation of puddle.NewPool
+	}
+	for range size {
+		// warm pool
+		go pud.CreateResource(context.Background())
+	}
+	go func() {
+		for {
+			time.Sleep(10 * time.Second)
+			for _, proc := range pud.AcquireAllIdle() {
+				if proc.Value().cmd.ProcessState != nil {
+					proc.Destroy()
+				} else {
+					proc.Release()
+				}
+			}
+		}
+	}()
+	return &pool{pud}
+}
+
+func (*pool) private() {}
+
+func (p *pool) Close() {
+	p.puddle.Close()
+}
+
+func (p *pool) Stat() *puddle.Stat {
+	return p.puddle.Stat()
+}
+
+func newWorker(ctx context.Context) (_ worker, err error) {
+	tracer := trace.SpanFromContext(ctx).TracerProvider().Tracer("")
+	ctx, span := tracer.Start(ctx, "jsonnetsecure.newWorker")
+	defer otelx.End(span, &err)
+
+	path, _ := ctx.Value(contextValuePath).(string)
+	if path == "" {
+		return worker{}, errors.New("newWorker: missing binary path in context")
+	}
+	args, _ := ctx.Value(contextValueArgs).([]string)
+	cmd := exec.Command(path, append(args, "-0")...)
+	cmd.Env = []string{"GOMAXPROCS=1"}
+	cmd.WaitDelay = 100 * time.Millisecond
+
+	span.SetAttributes(semconv.ProcessCommand(cmd.Path), semconv.ProcessCommandArgs(cmd.Args...))
+
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return worker{}, errors.Wrap(err, "newWorker: failed to create stdin pipe")
+	}
+
+	in := make(chan []byte, 1)
+	go func(c <-chan []byte) {
+		for input := range c {
+			if _, err := stdin.Write(append(input, 0)); err != nil {
+				stdin.Close()
+				return
+			}
+		}
+	}(in)
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return worker{}, errors.Wrap(err, "newWorker: failed to create stdout pipe")
+	}
+	stderr, err := cmd.StderrPipe()
+	if err != nil {
+		return worker{}, errors.Wrap(err, "newWorker: failed to create stderr pipe")
+	}
+
+	if err := cmd.Start(); err != nil {
+		return worker{}, errors.Wrap(err, "newWorker: failed to start process")
+	}
+
+	span.SetAttributes(semconv.ProcessPID(cmd.Process.Pid))
+
+	scan := func(c chan<- string, r io.Reader, maxTokenSize int) {
+		defer close(c)
+		scanner := bufio.NewScanner(r)
+		scanner.Buffer(make([]byte, 0, 64*KiB), maxTokenSize)
+
+		scanner.Split(splitNull)
+		for scanner.Scan() {
+			c <- scanner.Text()
+		}
+		if err := scanner.Err(); err != nil {
+			c <- "ERROR: scan: " + err.Error()
+		}
+	}
+	out := make(chan string, 1)
+	go scan(out, stdout, jsonnetOutputLimit)
+	errs := make(chan string, 1)
+	go scan(errs, stderr, jsonnetErrLimit)
+
+	w := worker{
+		cmd:    cmd,
+		stdin:  in,
+		stdout: out,
+		stderr: errs,
+	}
+	_, err = w.eval(ctx, []byte("{}")) // warm up
+	if err != nil {
+		w.destroy()
+		return worker{}, errors.Wrap(err, "newWorker: warm up failed")
+	}
+
+	return w, nil
+}
+
+func (w worker) destroy() {
+	close(w.stdin)
+	w.cmd.Process.Kill()
+	w.cmd.Wait()
+}
+
+func (w worker) eval(ctx context.Context, processParams []byte) (output string, err error) {
+	tracer := trace.SpanFromContext(ctx).TracerProvider().Tracer("")
+	ctx, span := tracer.Start(ctx, "jsonnetsecure.worker.eval", trace.WithAttributes(
+		semconv.ProcessPID(w.cmd.Process.Pid)))
+	defer otelx.End(span, &err)
+
+	// The worker is exclusively acquired, so the schedstat delta over this
+	// evaluation is attributable to it. cpu_time_us is time actually spent on
+	// a CPU; runqueue_wait_us is time runnable but waiting for one. Span wall
+	// time far exceeding their sum means the process was not even runnable,
+	// e.g. dequeued by cgroup CPU throttling.
+	//
+	// A kernel without CONFIG_SCHED_INFO reports "0 0 0" instead of failing
+	// the read. The worker has already spent CPU time on its warm-up
+	// evaluation, so a zero cpuTime means the fields are dead — skip them
+	// rather than record misleading zeros.
+	if before, ok := readSchedstat(w.cmd.Process.Pid); ok && before.cpuTime > 0 {
+		defer func() {
+			after, ok := readSchedstat(w.cmd.Process.Pid)
+			if !ok {
+				return
+			}
+			span.SetAttributes(
+				attribute.Int64("jsonnet.worker.cpu_time_us", (after.cpuTime-before.cpuTime).Microseconds()),
+				attribute.Int64("jsonnet.worker.runqueue_wait_us", (after.runqueueWait-before.runqueueWait).Microseconds()),
+			)
+		}()
+	}
+
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+	case w.stdin <- processParams:
+		break
+	}
+
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+	case output := <-w.stdout:
+		return output, nil
+	case err := <-w.stderr:
+		return "", errors.New(err)
+	}
+}
+
+func (vm *processPoolVM) EvaluateAnonymousSnippet(filename string, snippet string) (_ string, err error) {
+	tracer := trace.SpanFromContext(vm.ctx).TracerProvider().Tracer("")
+	ctx, span := tracer.Start(vm.ctx, "jsonnetsecure.processPoolVM.EvaluateAnonymousSnippet", trace.WithAttributes(attribute.String("filename", filename)))
+	defer otelx.End(span, &err)
+
+	params := vm.params
+	params.Filename = filename
+	params.Snippet = snippet
+	pp, err := json.Marshal(params)
+	if err != nil {
+		return "", errors.Wrap(err, "jsonnetsecure: marshal")
+	}
+
+	ctx = context.WithValue(ctx, contextValuePath, vm.path)
+	ctx = context.WithValue(ctx, contextValueArgs, vm.args)
+	worker, err := vm.pool.puddle.Acquire(ctx)
+	if err != nil {
+		return "", errors.Wrap(err, "jsonnetsecure: acquire")
+	}
+
+	ctx, cancel := context.WithTimeoutCause(ctx, 1*time.Second, errors.Errorf("failed to run jsonnet within 1s: filename=%s", filename))
+	defer cancel()
+	result, err := worker.Value().eval(ctx, pp)
+	if err != nil {
+		worker.Destroy()
+		return "", errors.Wrap(err, "jsonnetsecure: eval")
+	} else {
+		worker.Release()
+	}
+
+	if strings.HasPrefix(result, "ERROR: ") {
+		return "", errors.New("jsonnetsecure: " + result)
+	}
+
+	return result, nil
+}
+
+func NewProcessPoolVM(opts *vmOptions) VM {
+	ctx := opts.ctx
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	return &processPoolVM{
+		path: opts.jsonnetBinaryPath,
+		args: opts.args,
+		ctx:  ctx,
+		pool: opts.pool,
+	}
+}
+
+func (vm *processPoolVM) ExtCode(key string, val string) {
+	vm.params.ExtCodes = append(vm.params.ExtCodes, kv{key, val})
+}
+
+func (vm *processPoolVM) ExtVar(key string, val string) {
+	vm.params.ExtVars = append(vm.params.ExtVars, kv{key, val})
+}
+
+func (vm *processPoolVM) TLACode(key string, val string) {
+	vm.params.TLACodes = append(vm.params.TLACodes, kv{key, val})
+}
+
+func (vm *processPoolVM) TLAVar(key string, val string) {
+	vm.params.TLAVars = append(vm.params.TLAVars, kv{key, val})
+}