From 0118ee20621db65402d8e270f3ec83d9fad8cadb Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Mon, 1 Jan 2024 21:31:54 +0800
Subject: [PATCH] feat: allow to set uid/gid

---
 CMakeLists.txt |  1 +
 launcher.c     |  6 +++++-
 launcher.h     |  2 ++
 library.c      |  5 +++++
 user.c         | 13 +++++++++++++
 user.h         | 12 ++++++++++++
 6 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 user.c
 create mode 100644 user.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8d6cb5b..4b6a40d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -14,6 +14,7 @@ set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fmacro-prefix-map=${CMAKE_SOURCE_DIR}=.")
 file(GLOB SRC_FILES
         ${PROJECT_SOURCE_DIR}/resource.c
         ${PROJECT_SOURCE_DIR}/sandbox.c
+        ${PROJECT_SOURCE_DIR}/user.c
         ${PROJECT_SOURCE_DIR}/rules/*.c
         ${PROJECT_SOURCE_DIR}/utils/*.c)
 set(VERSION_SCRIPT ${PROJECT_SOURCE_DIR}/version_script.txt)
diff --git a/launcher.c b/launcher.c
index 1e2734f..f08f380 100644
--- a/launcher.c
+++ b/launcher.c
@@ -24,6 +24,8 @@ void print_help(char *self) {
     LOG_WARN("  --time_limit       time limit in ms");
     LOG_WARN("  --sandbox_template sandbox template");
     LOG_WARN("  --sandbox_action   sandbox action");
+    LOG_WARN("  --uid              user id");
+    LOG_WARN("  --gid              group id");
     LOG_WARN("  --file_input       path to input file");
     LOG_WARN("  --file_output      path to output file");
     LOG_WARN("  --file_info        path to info file");
@@ -38,6 +40,8 @@ void parse(int argc, char *argv[]) {
         [CFG_TIME_LIMIT]       = {"time_limit",       required_argument, NULL, 0},
         [CFG_SANDBOX_TEMPLATE] = {"sandbox_template", required_argument, NULL, 0},
         [CFG_SANDBOX_ACTION]   = {"sandbox_action",   required_argument, NULL, 0},
+        [CFG_UID]              = {"uid",              optional_argument, NULL, 0},
+        [CFG_GID]              = {"gid",              optional_argument, NULL, 0},
         [CFG_FILE_INPUT]       = {"file_input",       required_argument, NULL, 0},
         [CFG_FILE_OUTPUT]      = {"file_output",      required_argument, NULL, 0},
         [CFG_FILE_INFO]        = {"file_info",        required_argument, NULL, 0},
@@ -59,7 +63,7 @@ void parse(int argc, char *argv[]) {
     }
 
     for (int i = 0; i < CFG_IS_VALID; i++) {
-        if (!config[i]) {
+        if (!config[i] && options[i].has_arg == required_argument) {
             print_help(argv[0]);
             LOG_ERR("Missing arguments");
             exit(ERR_ARGUMENTS);
diff --git a/launcher.h b/launcher.h
index b16a258..3fa36a3 100644
--- a/launcher.h
+++ b/launcher.h
@@ -7,6 +7,8 @@ enum ConfigIndex {
     CFG_TIME_LIMIT,
     CFG_SANDBOX_TEMPLATE,
     CFG_SANDBOX_ACTION,
+    CFG_UID,
+    CFG_GID,
     CFG_FILE_INPUT,
     CFG_FILE_OUTPUT,
     CFG_FILE_INFO,
diff --git a/library.c b/library.c
index 3483341..991ac00 100644
--- a/library.c
+++ b/library.c
@@ -1,5 +1,6 @@
 #include "resource.h"
 #include "sandbox.h"
+#include "user.h"
 #include "utils/log.h"
 
 #include <fcntl.h>
@@ -30,8 +31,12 @@ void setup_all(void) {
         config[CFG_SANDBOX_TEMPLATE] = getenv(SANDBOX_TEMPLATE);
         config[CFG_SANDBOX_ACTION]   = getenv(SANDBOX_ACTION);
         config[CFG_PROGRAM]          = getenv(SANDBOX_EXE_PATH);
+
+        config[CFG_UID] = getenv(USER_UID);
+        config[CFG_GID] = getenv(USER_GID);
     }
 
+    setup_user(config);
     setup_rlimit(config);
     setup_seccomp(config);
 }
diff --git a/user.c b/user.c
new file mode 100644
index 0000000..c140b04
--- /dev/null
+++ b/user.c
@@ -0,0 +1,13 @@
+#include "user.h"
+#include <stdlib.h>
+#include <unistd.h>
+
+void setup_user(char *config[CFG_IS_VALID + 1]) {
+    long uid = -1, gid = -1;
+
+    if (config[CFG_UID]) uid = strtol(config[CFG_UID], NULL, 10);
+    if (config[CFG_GID]) gid = strtol(config[CFG_GID], NULL, 10);
+
+    if (uid != -1) setuid(uid);
+    if (gid != -1) setgid(gid);
+}
diff --git a/user.h b/user.h
new file mode 100644
index 0000000..693203c
--- /dev/null
+++ b/user.h
@@ -0,0 +1,12 @@
+#ifndef WOJ_SANDBOX_USER_H
+#define WOJ_SANDBOX_USER_H
+
+#include "launcher.h"
+
+// Configuration Environment Variables
+#define USER_UID "USER_UID"
+#define USER_GID "USER_GID"
+
+void setup_user(char *config[CFG_IS_VALID + 1]);
+
+#endif // WOJ_SANDBOX_USER_H