#!/usr/bin/env python3 # -*- coding: utf-8 -*- import sys import os import time import subprocess import json import psutil ResultJsonPath = "runner.json" StdinPath = "stdin.txt" StdoutPath = "stdout.txt" StderrPath = "stderr.txt" RunAs = "runner" # set to None if use root or not on linux ReserveMemory = 50*1024*1024 # 50MB MaxMemory = 1*1024*1024*1024 # 1GB def updateMaxMemory(): global MaxMemory memAvailable = MaxMemory cgroupLimit = MaxMemory with open("/proc/meminfo", "rb") as result: for line in result: if line.startswith(b"MemAvailable:"): memAvailable = int(line.split()[1]) * 1024 with open("/sys/fs/cgroup/memory/memory.stat", "rb") as result: for line in result: if line.startswith(b"hierarchical_memory_limit"): cgroupLimit = int(line.split()[1]) MaxMemory = min(MaxMemory, memAvailable, cgroupLimit) - ReserveMemory def clearFiles(): if os.path.isfile(ResultJsonPath): os.unlink(ResultJsonPath) if os.path.isfile(StdoutPath): os.unlink(StdoutPath) if os.path.isfile(StderrPath): os.unlink(StderrPath) def limitVirtualMemory(virtualMemoryLimitMb): """this would cause some program like java or nodejs fail, because they require large VMS""" if not RunAs: return virtualMemoryLimit = None if virtualMemoryLimitMb < 0: # default limit virtualMemoryLimit = MaxMemory elif virtualMemoryLimitMb == 0: # no limit pass else: # manual limit virtualMemoryLimit = virtualMemoryLimitMb * 1024 * 1024 # commit to /etc/security/limits.conf, unit is kb # example: runner hard as 512000 limitsText = b"" limitLine = ("%s hard as"%RunAs).encode("utf-8") with open("/etc/security/limits.conf", "rb") as result: for line in result: if line.startswith(limitLine): continue limitsText += line.strip() + b"\n" if virtualMemoryLimit is not None: limitsText += limitLine limitsText += (" %d\n"%int(virtualMemoryLimit / 1024)).encode("utf-8") with open("/etc/security/limits.conf", "wb") as result: result.write(limitsText) def writeResult(obj): with open(ResultJsonPath, "wb") as result: jsonBytes = json.dumps(obj).encode("utf-8") result.write(jsonBytes) result.write(b"\n") def nowInMillisecond(): return int(time.time() * 1000) def killRecursive(psutilProc): """kill process and it's childs""" try: for childProc in psutilProc.children(recursive=True): try: childProc.kill() except psutil.NoSuchProcess: pass psutilProc.kill() except psutil.NoSuchProcess: pass def killUser(username): """kill processes under specific user""" if username is None or username == "root": return # stop first, then kill, it's very effective for fork bomb with nproc limit killSet = set() killSetSize = None while (killSetSize is None) or (killSetSize < len(killSet)): killSetSize = len(killSet) for proc in psutil.process_iter(): try: if proc.username() == username: proc.suspend() killSet.add(proc.pid) # print("stopped %s"%proc.pid) except psutil.NoSuchProcess: pass killCount = None while (killCount is None) or (killCount > 0): killCount = 0 for proc in psutil.process_iter(): try: if proc.username() == username: proc.kill() killCount += 1 # print("killed %s"%proc.pid) except psutil.NoSuchProcess: pass def getUsedMemoryInBytes(psutilProc): """get maximum used memory in bytes from process and it's childs, return RSS not VMS""" try: rss = psutilProc.memory_info().rss for childProc in psutilProc.children(recursive=True): try: rss = max(rss, childProc.memory_info().rss) except psutil.NoSuchProcess: pass return rss except psutil.NoSuchProcess: return 0 def getUserTimeInMilliseconds(psutilProc): """get maximum user time in milliseconds from process and it's childs""" try: userTime = psutilProc.cpu_times().user for childProc in psutilProc.children(recursive=True): try: userTime = max(userTime, childProc.cpu_times().user) except psutil.NoSuchProcess: pass return int(userTime * 1000) except psutil.NoSuchProcess: return 0 def run(userLimitMs, totalLimitMs, cmdline): userTime = 0 totalTime = 0 peakMemory = 0 # RSS, not VMS exitCode = None isTimeout = False isOOM = False error = "" cmdOption = [ "su", RunAs, "-c", cmdline ] if RunAs else cmdline shellOption = False if RunAs else cmdline if not os.path.isfile(StdinPath): open(StdinPath, "wb").close() with open(StdinPath, "rb") as stdin, open(StdoutPath, "wb") as stdout, open(StderrPath, "wb") as stderr: startTime = nowInMillisecond() proc = subprocess.Popen( cmdOption, shell=shellOption, stdin=stdin, stdout=stdout, stderr=stderr) try: psutilProc = psutil.Process(proc.pid) except psutil.NoSuchProcess: pass while True: exitCode = proc.poll() if exitCode is not None: break if psutilProc is None: error = "process exist but get it's information failed" break peakMemory = max(peakMemory, getUsedMemoryInBytes(psutilProc)) userTime = max(userTime, getUserTimeInMilliseconds(psutilProc)) if userTime > userLimitMs: isTimeout = True error = "killed by timeout(user)" killRecursive(psutilProc) break if nowInMillisecond() - startTime > totalLimitMs: isTimeout = True error = "killed by timeout(total)" killRecursive(psutilProc) break if peakMemory >= MaxMemory: isOOM = True error = "killed by OOM" killRecursive(psutilProc) break time.sleep(0.001) exitCode = proc.poll() totalTime = nowInMillisecond() - startTime if exitCode is None: exitCode = -1 # kill signal sent but still exist killUser(RunAs) # cleanup if exitCode != 0 and not error: with open(StderrPath, "rb") as stderr: errorBytes = stderr.read() try: error = errorBytes.decode("utf-8") except UnicodeDecodeError: pass writeResult({ "Command": cmdline, "UserTime": userTime, "TotalTime": totalTime, "PeakMemory": peakMemory, # maybe 0 if process exited so quick "ExitCode": exitCode, "IsTimeout": isTimeout, "IsOOM": isOOM, "Error": error }) def main(): """ Input: First Line: user_time_limit_in_ms [total_time_limit_in_ms] [virtual_memory_limit_in_mb] Second Line: command_line stdin.txt (optional) Output: stdout.txt stderr.txt runner.json """ limitTimeMs = 0 cmdline = "" try: updateMaxMemory() clearFiles() limits = input().split() cmdline = input() userLimitMs = int(limits[0]) totalLimitMs = int(limits[1]) if (len(limits) > 1) else (userLimitMs * 2) virtualMemoryLimitMb = int(limits[2]) if (len(limits) > 2) else -1 limitVirtualMemory(virtualMemoryLimitMb) run(userLimitMs, totalLimitMs, cmdline) except Exception as e: writeResult({ "Command": cmdline, "UserTime": 0, "TotalTime": 0, "PeakMemory": 0, "ExitCode": -1, "IsTimeout": False, "Error": repr(e) }) if __name__ == "__main__": main()