From 24c3e87d09c1ca461024fb5f6c633e96dae35eab Mon Sep 17 00:00:00 2001
From: tom <tom@ceebs.dev>
Date: Sat, 13 Jul 2024 23:33:20 +1000
Subject: [PATCH] Add flag injection

---
 vendor/nsjail/context/inject-flag.sh        | 12 ++++++++++++
 vendor/nsjail/context/nsjail-user.sh        |  9 +++++++--
 vendor/nsjail/dockerfiles/Dockerfile.debian |  3 ++-
 vendor/nsjail/dockerfiles/Dockerfile.ubuntu |  3 ++-
 4 files changed, 23 insertions(+), 4 deletions(-)
 create mode 100644 vendor/nsjail/context/inject-flag.sh

diff --git a/vendor/nsjail/context/inject-flag.sh b/vendor/nsjail/context/inject-flag.sh
new file mode 100644
index 0000000..d82e000
--- /dev/null
+++ b/vendor/nsjail/context/inject-flag.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+export MODE=ONCE
+read FLAG
+export TMP_PARAMS="$TMP_PARAMS"$(cat <<-END
+, {
+    dst: "/flag",
+    src_content: $(echo "${FLAG@Q}")
+  }
+END
+)
+nsjail --config <(echo "$CONFIG" | envsubst)
diff --git a/vendor/nsjail/context/nsjail-user.sh b/vendor/nsjail/context/nsjail-user.sh
index 6b10a75..931fca1 100644
--- a/vendor/nsjail/context/nsjail-user.sh
+++ b/vendor/nsjail/context/nsjail-user.sh
@@ -2,7 +2,7 @@
 
 # idempotency ftw
 CONFIG_FILE="/home/ctf/nsjail.cfg"
-CONFIG=`cat "$CONFIG_FILE"`
+export CONFIG=`cat "$CONFIG_FILE"`
 
 # check and set default env vars
 export MODE=${MODE:-LISTEN}
@@ -29,6 +29,7 @@ if [ $TMP_ENABLED -eq 1 ]; then
     is_bind: false,
     rw: true
   }
+
 END
 )
 fi
@@ -41,5 +42,9 @@ if [ -f "/sys/fs/cgroup/cgroup.controllers" ]; then
 	fi
 fi
 
-nsjail --config <(echo "$CONFIG" | envsubst) --env FLAG
+if [ "$MODE" == "LISTEN_INJECT_FLAG" ]; then
+	socat tcp-listen:$PORT,reuseaddr,fork "exec:/docker-init/inject-flag.sh"
+else
+	nsjail --config <(echo "$CONFIG" | envsubst) --env FLAG
+fi
 
diff --git a/vendor/nsjail/dockerfiles/Dockerfile.debian b/vendor/nsjail/dockerfiles/Dockerfile.debian
index 315fd82..7c23671 100644
--- a/vendor/nsjail/dockerfiles/Dockerfile.debian
+++ b/vendor/nsjail/dockerfiles/Dockerfile.debian
@@ -32,6 +32,7 @@ RUN dpkg --add-architecture i386 \
      python3-venv \
      python3-gmpy2 \
      nano \
+     socat \
      $LIBPROTOBUF_VERSION \
      libnl-route-3-200 \
      libc6:i386 \
@@ -41,7 +42,7 @@ RUN dpkg --add-architecture i386 \
 COPY --from=build /nsjail/nsjail /usr/bin/nsjail
 RUN useradd -r -m ctf
 
-COPY docker-entrypoint.sh nsjail-launcher.sh nsjail-user.sh nsjail-config-cgroups.sh /docker-init/
+COPY docker-entrypoint.sh nsjail-launcher.sh nsjail-user.sh nsjail-config-cgroups.sh inject-flag.sh /docker-init/
 RUN chmod +x /docker-init/*
 
 ENTRYPOINT ["/docker-init/docker-entrypoint.sh"]
diff --git a/vendor/nsjail/dockerfiles/Dockerfile.ubuntu b/vendor/nsjail/dockerfiles/Dockerfile.ubuntu
index 390f870..24e80f8 100644
--- a/vendor/nsjail/dockerfiles/Dockerfile.ubuntu
+++ b/vendor/nsjail/dockerfiles/Dockerfile.ubuntu
@@ -32,6 +32,7 @@ RUN dpkg --add-architecture i386 \
      python3-gmpy2 \
      python3-pip \
      nano \
+     socat \
      $LIBPROTOBUF_VERSION \
      libnl-route-3-200 \
      libc6:i386 \
@@ -41,7 +42,7 @@ RUN dpkg --add-architecture i386 \
 COPY --from=build /nsjail/nsjail /usr/bin/nsjail
 RUN useradd -r -m ctf
 
-COPY docker-entrypoint.sh nsjail-launcher.sh nsjail-user.sh nsjail-config-cgroups.sh /docker-init/
+COPY docker-entrypoint.sh nsjail-launcher.sh nsjail-user.sh nsjail-config-cgroups.sh inject-flag.sh /docker-init/
 RUN chmod +x /docker-init/*
 
 ENTRYPOINT ["/docker-init/docker-entrypoint.sh"]