#!/bin/bash
# =====================================================================
# Job runner for the Octave sandbox.
#
# WHY FILES AND NOT HTTP
# This container executes code supplied by whoever is talking to the
# model, so it runs with network_mode "none" - there is no network stack
# to serve HTTP on. Work arrives as files on a small tmpfs shared with
# the api container:
#
#   /jobs/in/<id>/        directory the api fills with main.m, any extra
#                         .m files, and a `meta` file carrying a nonce
#   /jobs/in/<id>.ready   created LAST, so its appearance means the
#                         directory is complete. Without it the runner
#                         races the api and executes half-written code.
#   /jobs/out/<id>.out    result, moved into place atomically
#
# THE NONCE
# The executed script shares this uid and can therefore see /jobs. An
# adversarial review pointed out that this lets a script do more than the
# denial-of-service originally documented: it could WRITE another job's
# result file and hand that user a forged answer. So every result is
# stamped with a per-job nonce that the api generates and checks, and
# `meta` is deleted before Octave starts - the nonce lives only in this
# shell's memory, and the interpreter never has anything to read it from.
# The runner also writes the section markers itself, so a script cannot
# emit one.
#
# WHAT REMAINS, stated plainly: a hostile script can still delete or read
# another queued job, so it can cause a timeout or see someone else's
# source. It cannot escape the container, reach the network, forge a
# result, or read anything of the operator's. Closing the last gap needs a
# per-job container and therefore the docker socket, which would be a far
# worse trade than the problem.
# =====================================================================
set -u

# One line per sandboxed run, with every other component's log in ./logs.
# `exec` rather than a tee per echo, so anything the loop reports - including
# a job that failed in a way nobody anticipated - lands in the file too.
LOG_DIR="${OWUITOOLS_LOG_DIR:-/logs}"
if [ -d "$LOG_DIR" ] && [ -w "$LOG_DIR" ]; then
  exec > >(tee -a "$LOG_DIR/octave.log") 2>&1
fi

JOBS=${OWUITOOLS_JOBS_DIR:-/jobs}
DEFAULT_TIMEOUT=${OWUITOOLS_OCTAVE_TIMEOUT:-25}
MAX_TIMEOUT=${OWUITOOLS_OCTAVE_MAX_TIMEOUT:-60}
MAX_OUT=${OWUITOOLS_OCTAVE_MAX_OUTPUT:-200000}
REAP_MINUTES=${OWUITOOLS_OCTAVE_REAP_MINUTES:-10}

mkdir -p "$JOBS/in" "$JOBS/run" "$JOBS/out" "$HOME" 2>/dev/null

# Wrapping the user script in try/catch is what lets us hand the model a
# structured error with a line number instead of raw stderr, which varies
# by Octave version and buries the location.
#
# `message` is emitted LAST and everything after it is the message. An
# earlier version put it first, so an error whose text happened to contain
# a line starting "at:" - entirely possible, Octave quotes source - was
# truncated there and the rest parsed as a garbage stack frame.
DRIVER='
__ws = warning();
warning("off", "all");
try
  pkg load signal;  pkg load control;
  pkg load statistics; pkg load optim;
catch
end
try
  graphics_toolkit("gnuplot");
catch
end
try
  set(0, "defaultfigurevisible", "off");
catch
end
% Two MATLAB-compatibility preferences, both measured against MATLAB`s
% documented behaviour rather than assumed:
%   print_empty_dimensions - Octave prints `x = [](0x0)` where MATLAB
%     prints `x = []`, and the model reads that output back to the user.
%   save_default_options - `save("data.mat", "x")` wrote an Octave TEXT
%     file with a `# Created by Octave` header. A user opening it in
%     MATLAB gets nothing. -mat7-binary writes a real MAT-file (verified:
%     the header reads "MATLAB 5.0 MAT-file"), which both read.
% The rest of Octave`s --traditional set is deliberately not taken: it
% also turns struct printing off and the terminal bell on.
try
  print_empty_dimensions(false);
catch
end
try
  save_default_options("-mat7-binary");
catch
end
warning(__ws);
clear __ws;
% Ahead of Octave`s own print. Renders PNGs through the cairo terminal so
% thin lines are anti-aliased instead of dropping pixels; see shim/print.m.
% Suppressed BEFORE the addpath, not after: the warning is emitted at the
% moment the shim joins the path, so silencing it afterwards was one line
% too late. The shadow is deliberate and permanent, and this is the same
% stderr the model reads for its real errors.
warning("off", "Octave:shadowed-function");
try
  addpath("/opt/owuitools-shim", "-begin");
catch
end
more off;
% Octave`s own parser, used as a second compatibility checker: it names
% the Octave-only constructs in the file it is about to run - `x += 1`,
% `!=`, `.*=`, a bare newline inside parentheses - with a line number, and
% it is never wrong about its own dialect the way a regex can be.
%
% This was OFF for a long time and the comment here said it was unusable,
% because Octave`s own library is written in the same dialect and the real
% error disappeared under hundreds of warnings about files nobody can
% edit. The missing piece was that every warning names its file: the
% library ones are dropped by path in run.sh, and what is left is the
% user`s code. Turned on LAST, after the pkg loads and the addpath, so
% none of that machinery is parsed with it enabled.
warning("on", "Octave:language-extension");
try
  source("main.m");
catch err
  fflush(stdout);
  fprintf(stderr, "\n__OWUITOOLS_ERR__\n");
  fprintf(stderr, "identifier:%s\n", err.identifier);
  if (isfield(err, "stack") && numel(err.stack) > 0)
    for k = 1:numel(err.stack)
      fprintf(stderr, "at:%s:%d\n", err.stack(k).name, err.stack(k).line);
    end
  end
  fprintf(stderr, "message:%s\n", err.message);
  fflush(stderr);
  exit(3);
end
fflush(stdout);
exit(0);
'

run_job() {
  local id="$1" dir="$JOBS/run/$id"
  local work timeout rc started ended nonce figure
  work="$dir/work"

  timeout="$DEFAULT_TIMEOUT"
  nonce=""
  figure=""
  if [ -f "$dir/meta" ]; then
    local t
    t=$(sed -n 's/^TIMEOUT=\([0-9]\{1,3\}\)$/\1/p' "$dir/meta" | head -1)
    if [ -n "$t" ] && [ "$t" -ge 1 ] 2>/dev/null && [ "$t" -le "$MAX_TIMEOUT" ]; then
      timeout="$t"
    fi
    nonce=$(sed -n 's/^NONCE=\([A-Za-z0-9]\{8,64\}\)$/\1/p' "$dir/meta" | head -1)
    figure=$(sed -n 's/^FIGURE=\([A-Za-z0-9_.-]\{1,64\}\)$/\1/p' "$dir/meta" | head -1)
  fi

  mkdir -p "$work"
  # Only .m files move into the working directory. Anything else the api
  # left there is metadata, and letting it onto the Octave path would be a
  # way to smuggle in a file the linter never saw.
  local inputs="$dir/__inputs"
  : > "$inputs"
  for f in "$dir"/*.m; do
    [ -e "$f" ] || continue
    cp "$f" "$work/"
    basename "$f" >> "$inputs"
  done
  cp "$inputs" "$work/../__inputs.keep" 2>/dev/null

  printf '%s' "$DRIVER" > "$work/__driver.m"

  # The nonce must not survive into anything the interpreter can read.
  rm -f "$dir/meta" "$inputs" 2>/dev/null

  started=$(date +%s%3N)
  ( cd "$work" && timeout -s KILL "$timeout" \
      octave --no-window-system --norc --no-history --quiet \
             --eval "run('__driver.m')" ) \
      > "$work/__stdout" 2> "$work/__stderr"
  rc=$?
  ended=$(date +%s%3N)

  # Files the script itself created: what is in the working directory now,
  # minus what we put there. Newline separated, because a filename may
  # legitimately contain a comma and a comma-joined list then reads back
  # as two files that do not exist.
  local artefacts
  artefacts=$(cd "$work" && ls -1 2>/dev/null \
    | grep -v -x -F -f "$work/../__inputs.keep" \
    | grep -v -x -E '__driver\.m|__stdout|__stderr' \
    | head -20)

  {
    echo "NONCE=$nonce"
    echo "RC=$rc"
    echo "TIMEOUT=$timeout"
    echo "DURATION_MS=$((ended - started))"
    echo "__OWUITOOLS_ARTEFACTS__"
    printf '%s\n' "$artefacts"
    echo "__OWUITOOLS_STDOUT__"
    head -c "$MAX_OUT" "$work/__stdout" 2>/dev/null
    echo
    echo "__OWUITOOLS_STDERR__"
    # Octave:language-extension is on (see the driver) so the parser tells
    # us which Octave-only constructs the USER wrote. Octave`s own library
    # is written in the same dialect and would bury that under hundreds of
    # lines about files nobody can edit, so those are dropped here, by
    # path, before anything is measured against the output cap.
    #
    # A warning is three parts - the line itself, a "called from", and its
    # indented frames - and dropping only the first left the frames behind
    # as orphans, which read like a stack trace for an error that never
    # happened. The file is matched rather than the spacing because 7.3
    # prints "offile" for "of file".
    awk '
      /^warning: Octave language extension used:.*(\/usr\/share\/octave|\/opt\/owuitools-shim)/ {
        skip = 1; next
      }
      skip && /^warning: called from[ \t]*$/ { next }
      skip && /^[ \t]+/ { next }
      skip && /^[ \t]*$/ { next }
      { skip = 0; print }
    ' "$work/__stderr" 2>/dev/null | head -c "$MAX_OUT"
    # A figure the caller asked for, base64 on one line. Open WebUI turns a
    # tool result that is a data: URI into a file shown to the user and
    # never sends the bytes to the model, so this costs no context.
    if [ -n "$figure" ] && [ -f "$work/$figure" ]; then
      sz=$(stat -c %s "$work/$figure" 2>/dev/null || echo 0)
      if [ "$sz" -gt 0 ] && [ "$sz" -le 3000000 ]; then
        echo
        echo "__OWUITOOLS_FIGURE__"
        base64 -w 0 "$work/$figure" 2>/dev/null
        echo
      fi
    fi
  } > "$JOBS/out/$id.tmp" 2>/dev/null

  mv -f "$JOBS/out/$id.tmp" "$JOBS/out/$id.out" 2>/dev/null
  rm -rf "$dir" 2>/dev/null
}

# Anything left over from a previous container start would otherwise wedge
# its caller until timeout.
rm -rf "${JOBS:?}/run"/* 2>/dev/null

reap_counter=0

# `claimed`, not `found`: a job we can SEE but cannot CLAIM - a permissions
# problem on the volume, say - would otherwise spin this loop at full CPU
# forever, because the marker never disappears.
while true; do
  claimed=0
  for ready in "$JOBS/in"/*.ready; do
    [ -e "$ready" ] || continue
    id=$(basename "$ready" .ready)
    case "$id" in
      *[!A-Za-z0-9_-]*) rm -f "$ready" 2>/dev/null; continue ;;
    esac
    if mv "$JOBS/in/$id" "$JOBS/run/$id" 2>/dev/null; then
      rm -f "$ready" 2>/dev/null
      claimed=1
      # Touched here as well as below: the runner is serial, and a long
      # job used to starve the heartbeat and report the container
      # unhealthy in exactly the case the healthcheck exists to survive.
      touch "$JOBS/.heartbeat" 2>/dev/null
      run_job "$id"
    else
      rm -f "$ready" 2>/dev/null
    fi
  done

  # Nothing else reclaims these. A result the api gave up waiting for, or
  # a job directory whose .ready never arrived, would sit on a 64 MB tmpfs
  # until the container restarted.
  reap_counter=$((reap_counter + 1))
  if [ "$reap_counter" -ge 400 ]; then
    reap_counter=0
    find "$JOBS/out" -maxdepth 1 -type f -mmin "+$REAP_MINUTES" -delete 2>/dev/null
    find "$JOBS/in" -maxdepth 1 -mmin "+$REAP_MINUTES" \
         \( -type f -o -type d \) ! -path "$JOBS/in" -exec rm -rf {} + 2>/dev/null
  fi

  touch "$JOBS/.heartbeat" 2>/dev/null
  [ "$claimed" -eq 1 ] || sleep 0.15
done