//===----------------------------------------------------------------------===//
// Copyright © 2025 Apple Inc. and the Containerization project authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//

/// NOTE: This binary implements a very small subset of the OCI runtime spec, mostly just
/// the process configurations. Mounts are somewhat functional, but masked and read only paths
/// aren't checked today. Today the namespaces are also ignored, and we always spawn a new pid
/// and mount namespace.

import ArgumentParser
import ContainerizationError
import ContainerizationOCI
import Foundation
import LCShim
import Logging
import Musl

@main
struct App: ParsableCommand {
    static let configuration = CommandConfiguration(
        commandName: "vmexec",
        version: "0.1.0",
        subcommands: [
            ExecCommand.self,
            RunCommand.self,
        ]
    )

    static let standardErrorLock = NSLock()

    @Sendable
    static func standardError(label: String) -> StreamLogHandler {
        standardErrorLock.withLock {
            StreamLogHandler.standardError(label: label)
        }
    }
}

extension App {
    /// Applies O_CLOEXEC to all file descriptors currently open for
    /// the process except the stdio fd values
    static func applyCloseExecOnFDs() throws {
        let minFD = 2  // stdin, stdout, stderr should be preserved

        let fdList = try FileManager.default.contentsOfDirectory(atPath: "/proc/self/fd")

        for fdStr in fdList {
            guard let fd = Int(fdStr) else {
                continue
            }
            if fd <= minFD {
                continue
            }

            _ = fcntl(Int32(fd), F_SETFD, FD_CLOEXEC)
        }
    }

    static func exec(process: ContainerizationOCI.Process) throws {
        let executable = strdup(process.args[0])
        var argv = process.args.map { strdup($0) }
        argv += [nil]

        let env = process.env.map { strdup($0) } + [nil]
        let cwd = process.cwd

        // switch cwd
        guard chdir(cwd) == 0 else {
            throw App.Errno(stage: "chdir(cwd)", info: "Failed to change directory to '\(cwd)'")
        }

        guard execvpe(executable, argv, env) != -1 else {
            throw App.Errno(stage: "execvpe(\(String(describing: executable)))", info: "Failed to exec [\(process.args.joined(separator: " "))]")
        }
        fatalError("execvpe failed")
    }

    static func setPermissions(user: ContainerizationOCI.User) throws {
        if user.additionalGids.count > 0 {
            guard setgroups(user.additionalGids.count, user.additionalGids) == 0 else {
                throw App.Errno(stage: "setgroups()")
            }
        }
        guard setgid(user.gid) == 0 else {
            throw App.Errno(stage: "setgid()")
        }
        // NOTE: setuid has to be done last because once the uid has been
        // changed, then the process will lose privilege to set the group
        // and supplementary groups
        guard setuid(user.uid) == 0 else {
            throw App.Errno(stage: "setuid()")
        }
    }

    static func setRLimits(rlimits: [ContainerizationOCI.POSIXRlimit]) throws {
        for rl in rlimits {
            var limit = rlimit(rlim_cur: rl.soft, rlim_max: rl.hard)
            let resource: Int32
            switch rl.type {
            case "RLIMIT_AS":
                resource = RLIMIT_AS
            case "RLIMIT_CORE":
                resource = RLIMIT_CORE
            case "RLIMIT_CPU":
                resource = RLIMIT_CPU
            case "RLIMIT_DATA":
                resource = RLIMIT_DATA
            case "RLIMIT_FSIZE":
                resource = RLIMIT_FSIZE
            case "RLIMIT_NOFILE":
                resource = RLIMIT_NOFILE
            case "RLIMIT_STACK":
                resource = RLIMIT_STACK
            case "RLIMIT_NPROC":
                resource = RLIMIT_NPROC
            case "RLIMIT_RSS":
                resource = RLIMIT_RSS
            case "RLIMIT_MEMLOCK":
                resource = RLIMIT_MEMLOCK
            default:
                errno = EINVAL
                throw App.Errno(stage: "rlimit key unknown")
            }
            guard setrlimit(resource, &limit) == 0 else {
                throw App.Errno(stage: "setrlimit()")
            }
        }
    }

    static func Errno(stage: String, info: String = "") -> ContainerizationError {
        let posix = POSIXError(.init(rawValue: errno)!, userInfo: ["stage": stage])
        return ContainerizationError(.internalError, message: "\(info) \(String(describing: posix))")
    }
}