CLI: Add support for rlimits (#1129)

Closes #1097.
2026-09-21 07:05:57 +00:00 · 2026-02-02 13:30:19 -08:00
parent b3b5c3e609
commit 7dfe27d825
4 changed files with 368 additions and 1 deletions
@@ -61,6 +61,15 @@ public struct Flags {
            )
        )
        public var cwd: String?
+
+        @Option(
+            name: .customLong("ulimit"),
+            help: .init(
+                "Set resource limits (format: <type>=<soft>[:<hard>])",
+                valueName: "limit"
+            )
+        )
+        public var ulimits: [String] = []
    }

    public struct Resource: ParsableArguments {
@@ -281,6 +281,8 @@ public struct Parser {
            user: processFlags.user, uid: processFlags.uid,
            gid: processFlags.gid, defaultUser: defaultUser)

+        let rlimits = try Parser.rlimits(processFlags.ulimits)
+
        return .init(
            executable: commandToRun.first!,
            arguments: [String](commandToRun.dropFirst()),
@@ -288,7 +290,8 @@ public struct Parser {
            workingDirectory: workingDir,
            terminal: processFlags.tty,
            user: user,
-            supplementalGroups: additionalGroups
+            supplementalGroups: additionalGroups,
+            rlimits: rlimits
        )
    }

@@ -867,6 +870,114 @@ public struct Parser {
        return !label.ranges(of: pattern).isEmpty
    }

+    // TODO: When containerization supports all 16 (minus AS as it's not great) add
+    // them here.
+    private static let ulimitNameToRlimit: [String: String] = [
+        "core": "RLIMIT_CORE",
+        "cpu": "RLIMIT_CPU",
+        "data": "RLIMIT_DATA",
+        "fsize": "RLIMIT_FSIZE",
+        "memlock": "RLIMIT_MEMLOCK",
+        "nofile": "RLIMIT_NOFILE",
+        "nproc": "RLIMIT_NPROC",
+        "rss": "RLIMIT_RSS",
+        "stack": "RLIMIT_STACK",
+    ]
+
+    /// Parse ulimit specifications into Rlimit objects
+    /// Format: <type>=<soft>[:<hard>]
+    /// Examples:
+    ///   - nofile=1024:2048  (soft=1024, hard=2048)
+    ///   - nofile=1024       (soft=hard=1024)
+    ///   - nofile=unlimited  (soft=hard=UINT64_MAX)
+    ///   - nofile=1024:unlimited (soft=1024, hard=UINT64_MAX)
+    public static func rlimits(_ rawUlimits: [String]) throws -> [ProcessConfiguration.Rlimit] {
+        var rlimits: [ProcessConfiguration.Rlimit] = []
+        var seenTypes: Set<String> = []
+
+        for ulimit in rawUlimits {
+            let rlimit = try Parser.rlimit(ulimit)
+            if seenTypes.contains(rlimit.limit) {
+                throw ContainerizationError(
+                    .invalidArgument,
+                    message: "duplicate ulimit type: \(ulimit.split(separator: "=").first ?? "")"
+                )
+            }
+            seenTypes.insert(rlimit.limit)
+            rlimits.append(rlimit)
+        }
+
+        return rlimits
+    }
+
+    /// Parse a single ulimit specification
+    public static func rlimit(_ ulimit: String) throws -> ProcessConfiguration.Rlimit {
+        let parts = ulimit.split(separator: "=", maxSplits: 1)
+        guard parts.count == 2 else {
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "invalid ulimit format '\(ulimit)': expected <type>=<soft>[:<hard>]"
+            )
+        }
+
+        let typeName = String(parts[0]).lowercased()
+        let valuesPart = String(parts[1])
+
+        guard let rlimitType = ulimitNameToRlimit[typeName] else {
+            let validTypes = ulimitNameToRlimit.keys.sorted().joined(separator: ", ")
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "unsupported ulimit type '\(typeName)': valid types are \(validTypes)"
+            )
+        }
+
+        let valueParts = valuesPart.split(separator: ":", maxSplits: 1)
+        let soft: UInt64
+        let hard: UInt64
+
+        switch valueParts.count {
+        case 1:
+            // Single value: use for both soft and hard
+            soft = try parseRlimitValue(String(valueParts[0]), typeName: typeName)
+            hard = soft
+        case 2:
+            // Two values: soft:hard
+            soft = try parseRlimitValue(String(valueParts[0]), typeName: typeName)
+            hard = try parseRlimitValue(String(valueParts[1]), typeName: typeName)
+        default:
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "invalid ulimit format '\(ulimit)': expected <type>=<soft>[:<hard>]"
+            )
+        }
+
+        if soft > hard {
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "ulimit '\(typeName)' soft limit (\(soft)) cannot exceed hard limit (\(hard))"
+            )
+        }
+
+        return ProcessConfiguration.Rlimit(limit: rlimitType, soft: soft, hard: hard)
+    }
+
+    private static func parseRlimitValue(_ value: String, typeName: String) throws -> UInt64 {
+        let trimmed = value.trimmingCharacters(in: .whitespaces).lowercased()
+
+        if trimmed == "unlimited" || trimmed == "-1" {
+            return UInt64.max
+        }
+
+        guard let parsed = UInt64(trimmed) else {
+            throw ContainerizationError(
+                .invalidArgument,
+                message: "invalid ulimit value '\(value)' for '\(typeName)': must be a non-negative integer or 'unlimited'"
+            )
+        }
+
+        return parsed
+    }
+
    // MARK: Miscellaneous

    public static func parseBool(string: String) -> Bool? {
@@ -201,6 +201,92 @@ class TestCLIRunCommand1: CLITest {
            return
        }
    }
+
+    @Test func testRunCommandUlimitNofile() throws {
+        do {
+            let name = getTestName()
+            let softLimit = "1024"
+            let hardLimit = "2048"
+            try doLongRun(name: name, args: ["--ulimit", "nofile=\(softLimit):\(hardLimit)"])
+            defer {
+                try? doStop(name: name)
+            }
+
+            let inspectResp = try inspectContainer(name)
+            let rlimits = inspectResp.configuration.initProcess.rlimits
+            let nofileRlimit = rlimits.first { $0.limit == "RLIMIT_NOFILE" }
+            #expect(nofileRlimit != nil, "expected RLIMIT_NOFILE to be set")
+            #expect(nofileRlimit?.soft == UInt64(softLimit), "expected soft limit \(softLimit), got \(nofileRlimit?.soft ?? 0)")
+            #expect(nofileRlimit?.hard == UInt64(hardLimit), "expected hard limit \(hardLimit), got \(nofileRlimit?.hard ?? 0)")
+
+            var output = try doExec(name: name, cmd: ["sh", "-c", "ulimit -n"])
+            output = output.trimmingCharacters(in: .whitespacesAndNewlines)
+            #expect(output == softLimit, "expected ulimit -n to return \(softLimit), got \(output)")
+
+            try doStop(name: name)
+        } catch {
+            Issue.record("failed to run container \(error)")
+            return
+        }
+    }
+
+    @Test func testRunCommandUlimitNproc() throws {
+        do {
+            let name = getTestName()
+            let limit = "256"
+            try doLongRun(name: name, args: ["--ulimit", "nproc=\(limit)"])
+            defer {
+                try? doStop(name: name)
+            }
+            let inspectResp = try inspectContainer(name)
+            let rlimits = inspectResp.configuration.initProcess.rlimits
+            let nprocRlimit = rlimits.first { $0.limit == "RLIMIT_NPROC" }
+            #expect(nprocRlimit != nil, "expected RLIMIT_NPROC to be set")
+            #expect(nprocRlimit?.soft == UInt64(limit), "expected soft limit \(limit), got \(nprocRlimit?.soft ?? 0)")
+            #expect(nprocRlimit?.hard == UInt64(limit), "expected hard limit \(limit), got \(nprocRlimit?.hard ?? 0)")
+
+            var output = try doExec(name: name, cmd: ["sh", "-c", "ulimit -u"])
+            output = output.trimmingCharacters(in: .whitespacesAndNewlines)
+            #expect(output == limit, "expected ulimit -u to return \(limit), got \(output)")
+
+            try doStop(name: name)
+        } catch {
+            Issue.record("failed to run container \(error)")
+            return
+        }
+    }
+
+    @Test func testRunCommandMultipleUlimits() throws {
+        do {
+            let name = getTestName()
+            try doLongRun(
+                name: name,
+                args: [
+                    "--ulimit", "nofile=1024:2048",
+                    "--ulimit", "nproc=512",
+                    "--ulimit", "stack=8388608",
+                ])
+            defer {
+                try? doStop(name: name)
+            }
+            let inspectResp = try inspectContainer(name)
+            let rlimits = inspectResp.configuration.initProcess.rlimits
+            #expect(rlimits.count == 3, "expected 3 rlimits, got \(rlimits.count)")
+
+            let nofile = rlimits.first { $0.limit == "RLIMIT_NOFILE" }
+            let nproc = rlimits.first { $0.limit == "RLIMIT_NPROC" }
+            let stack = rlimits.first { $0.limit == "RLIMIT_STACK" }
+
+            #expect(nofile != nil && nofile?.soft == 1024 && nofile?.hard == 2048)
+            #expect(nproc != nil && nproc?.soft == 512 && nproc?.hard == 512)
+            #expect(stack != nil && stack?.soft == 8_388_608 && stack?.hard == 8_388_608)
+
+            try doStop(name: name)
+        } catch {
+            Issue.record("failed to run container \(error)")
+            return
+        }
+    }
 }

 class TestCLIRunCommand2: CLITest {
@@ -863,4 +863,165 @@ struct ParserTest {
        #expect(result.executable == "./uname")
        #expect(result.workingDirectory == "/bin")
    }
+
+    @Test
+    func testUlimitParserSoftAndHard() throws {
+        let result = try Parser.rlimits(["nofile=1024:2048"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_NOFILE")
+        #expect(result[0].soft == 1024)
+        #expect(result[0].hard == 2048)
+    }
+
+    @Test
+    func testUlimitParserSingleValue() throws {
+        let result = try Parser.rlimits(["nproc=512"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_NPROC")
+        #expect(result[0].soft == 512)
+        #expect(result[0].hard == 512)
+    }
+
+    @Test
+    func testUlimitParserUnlimited() throws {
+        let result = try Parser.rlimits(["memlock=unlimited"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_MEMLOCK")
+        #expect(result[0].soft == UInt64.max)
+        #expect(result[0].hard == UInt64.max)
+    }
+
+    @Test
+    func testUlimitParserUnlimitedHardOnly() throws {
+        let result = try Parser.rlimits(["stack=8192:unlimited"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_STACK")
+        #expect(result[0].soft == 8192)
+        #expect(result[0].hard == UInt64.max)
+    }
+
+    @Test
+    func testUlimitParserMinusOneAsUnlimited() throws {
+        let result = try Parser.rlimits(["core=-1"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_CORE")
+        #expect(result[0].soft == UInt64.max)
+        #expect(result[0].hard == UInt64.max)
+    }
+
+    @Test
+    func testUlimitParserMultipleUlimits() throws {
+        let result = try Parser.rlimits(["nofile=1024:2048", "nproc=256", "cpu=60:120"])
+        #expect(result.count == 3)
+        #expect(result[0].limit == "RLIMIT_NOFILE")
+        #expect(result[1].limit == "RLIMIT_NPROC")
+        #expect(result[2].limit == "RLIMIT_CPU")
+    }
+
+    @Test
+    func testUlimitParserAllSupportedTypes() throws {
+        let types = ["core", "cpu", "data", "fsize", "memlock", "nofile", "nproc", "rss", "stack"]
+        let expectedRlimits = [
+            "RLIMIT_CORE", "RLIMIT_CPU", "RLIMIT_DATA", "RLIMIT_FSIZE",
+            "RLIMIT_MEMLOCK", "RLIMIT_NOFILE", "RLIMIT_NPROC", "RLIMIT_RSS", "RLIMIT_STACK",
+        ]
+
+        for (i, type) in types.enumerated() {
+            let result = try Parser.rlimits(["\(type)=100"])
+            #expect(result.count == 1)
+            #expect(result[0].limit == expectedRlimits[i])
+        }
+    }
+
+    @Test
+    func testUlimitParserCaseInsensitive() throws {
+        let result = try Parser.rlimits(["NOFILE=1024", "Nproc=512"])
+        #expect(result.count == 2)
+        #expect(result[0].limit == "RLIMIT_NOFILE")
+        #expect(result[1].limit == "RLIMIT_NPROC")
+    }
+
+    @Test
+    func testUlimitParserInvalidFormat() throws {
+        #expect {
+            _ = try Parser.rlimits(["nofile"])
+        } throws: { error in
+            guard let error = error as? ContainerizationError else {
+                return false
+            }
+            return error.description.contains("invalid ulimit format")
+        }
+    }
+
+    @Test
+    func testUlimitParserUnsupportedType() throws {
+        #expect {
+            _ = try Parser.rlimits(["foo=100"])
+        } throws: { error in
+            guard let error = error as? ContainerizationError else {
+                return false
+            }
+            return error.description.contains("unsupported ulimit type")
+        }
+    }
+
+    @Test
+    func testUlimitParserSoftExceedsHard() throws {
+        #expect {
+            _ = try Parser.rlimits(["nofile=2048:1024"])
+        } throws: { error in
+            guard let error = error as? ContainerizationError else {
+                return false
+            }
+            return error.description.contains("soft limit") && error.description.contains("cannot exceed hard limit")
+        }
+    }
+
+    @Test
+    func testUlimitParserDuplicateType() throws {
+        #expect {
+            _ = try Parser.rlimits(["nofile=1024", "nofile=2048"])
+        } throws: { error in
+            guard let error = error as? ContainerizationError else {
+                return false
+            }
+            return error.description.contains("duplicate ulimit type")
+        }
+    }
+
+    @Test
+    func testUlimitParserInvalidValue() throws {
+        #expect {
+            _ = try Parser.rlimits(["nofile=abc"])
+        } throws: { error in
+            guard let error = error as? ContainerizationError else {
+                return false
+            }
+            return error.description.contains("invalid ulimit value")
+        }
+    }
+
+    @Test
+    func testUlimitParserEmptyArray() throws {
+        let result = try Parser.rlimits([])
+        #expect(result.isEmpty)
+    }
+
+    @Test
+    func testUlimitParserZeroValue() throws {
+        let result = try Parser.rlimits(["core=0"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_CORE")
+        #expect(result[0].soft == 0)
+        #expect(result[0].hard == 0)
+    }
+
+    @Test
+    func testUlimitParserLargeValues() throws {
+        let result = try Parser.rlimits(["nproc=\(UInt64.max - 1):\(UInt64.max)"])
+        #expect(result.count == 1)
+        #expect(result[0].limit == "RLIMIT_NPROC")
+        #expect(result[0].soft == UInt64.max - 1)
+        #expect(result[0].hard == UInt64.max)
+    }
 }