Container.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Hercules.Effect.Container where

import Control.Lens
import Data.Aeson (Value (String), eitherDecode, encode, object, toJSON)
import Data.Aeson.Lens
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as BS.Char8
import Data.ByteString.Lazy qualified as BL
import Data.Map qualified as M
import Data.UUID.V4 qualified as UUID
import Data.Vector qualified as V
import GHC.IO.Exception (IOErrorType (HardwareFault))
import Protolude
import System.Directory (createDirectory)
import System.FilePath ((</>))
import System.IO (hClose)
import System.IO.Error (ioeGetErrorType)
import System.Posix.IO (closeFd, fdToHandle)
import System.Posix.Terminal (openPseudoTerminal)
import System.Process (CreateProcess (..), StdStream (UseHandle), proc, waitForProcess, withCreateProcess)
import System.Process.ByteString (readCreateProcessWithExitCode)

data BindMount = BindMount
  { pathInContainer :: Text,
    pathInHost :: Text,
    readOnly :: Bool
  }

defaultBindMount :: Text -> BindMount
defaultBindMount path = BindMount {pathInContainer = path, pathInHost = path, readOnly = True}

data Config = Config
  { extraBindMounts :: [BindMount],
    executable :: Text,
    arguments :: [Text],
    environment :: Map Text Text,
    workingDirectory :: Text,
    hostname :: Text,
    rootReadOnly :: Bool,
    virtualUID :: Int,
    virtualGID :: Int,
    hostUID :: Int,
    hostGID :: Int
  }

effectToOCIRuntimeSpec :: Config -> Value -> Value
effectToOCIRuntimeSpec config spec =
  let defaultMounts = [defaultBindMount "/nix/store"]
      mounts =
        foldMap
          ( \bindMount ->
              pure $
                object
                  [ ("destination", String $ pathInContainer bindMount),
                    ("source", String $ pathInHost bindMount),
                    ("type", "bind"),
                    ( "options",
                      toJSON $
                        ["bind" :: Text]
                          <> ["ro" | readOnly bindMount]
                    )
                  ]
          )
          (defaultMounts <> extraBindMounts config)
   in spec
        & key "process" . key "args" .~ toJSON ([executable config] <> arguments config)
        & key "mounts" . _Array %~ (<> mounts)
        & key "process" . key "terminal" .~ toJSON True
        & key "process" . key "env" .~ toJSON (config & environment & M.toList <&> \(k, v) -> k <> "=" <> v)
        & key "process" . key "cwd" .~ toJSON (config & workingDirectory)
        & key "process" . key "user" . key "uid" .~ toJSON (virtualUID config)
        & key "process" . key "user" . key "gid" .~ toJSON (virtualGID config)
        & key "process" . key "user" . key "umask" .~ toJSON (0o077 :: Int)
        & key "process" . key "user" . key "additionalGids" . _Array .~ V.fromList []
        & key "linux" . _Object . at "uidMappings"
          ?~ toJSON
            [ object
                [ ("containerID", toJSON (virtualUID config)),
                  ("hostID", toJSON (hostUID config)),
                  ("size", toJSON (1 :: Int))
                ]
            ]
        & key "linux" . _Object . at "gidMappings"
          ?~ toJSON
            [ object
                [ ("containerID", toJSON (virtualGID config)),
                  ("hostID", toJSON (hostGID config)),
                  ("size", toJSON (1 :: Int))
                ]
            ]
        & key "hostname" .~ toJSON (config & hostname)
        & key "root" . key "readonly" .~ toJSON (config & rootReadOnly)
        -- TODO Use slirp? e.g. https://github.com/rootless-containers/slirp4netns or might kernel offer bridging (in the future?)
        & key "linux" . key "namespaces" . _Array %~ V.filter (\x -> x ^? key "type" . _String /= Just "network")

run :: FilePath -> Config -> IO ExitCode
run dir config = do
  let containerRuntimeExe = "crun"
      createConfigJsonSpec =
        (System.Process.proc containerRuntimeExe ["spec", "--rootless"])
          { cwd = Just dir
          }
      configJsonPath = dir </> "config.json"
      runtimeRootPath = dir </> "container-root"
      rootfsPath = dir </> "rootfs"
  (exit, _out, err) <- readCreateProcessWithExitCode createConfigJsonSpec ""
  case exit of
    ExitSuccess -> pass
    ExitFailure e -> do
      putErrText (decodeUtf8With lenientDecode err)
      panic $ "Could not create container configuration template. " <> toS containerRuntimeExe <> " terminated with exit code " <> show e
  templateBytes <- BS.readFile configJsonPath
  template <- case eitherDecode (BL.fromStrict templateBytes) of
    Right a -> pure a
    Left e -> throwIO (FatalError $ "decoding container config.json template: " <> show e)
  let configJson = effectToOCIRuntimeSpec config template
  BS.writeFile configJsonPath (BL.toStrict $ encode configJson)
  createDirectory rootfsPath
  createDirectory runtimeRootPath
  name <- do
    uuid <- UUID.nextRandom
    pure $ "hercules-ci-" <> show uuid
  (exitCode, _) <- withPseudoTerminalHandles $
    \(master, terminal) -> do
      concurrently
        ( do
            let createProcSpec =
                  (System.Process.proc containerRuntimeExe ["--root", runtimeRootPath, "run", name])
                    { std_in = UseHandle terminal, -- can't pass /dev/null :(
                      std_out = UseHandle terminal,
                      std_err = UseHandle terminal,
                      cwd = Just dir
                    }
            withCreateProcess createProcSpec \_subStdin _noOut _noErr processHandle -> do
              waitForProcess processHandle
                `onException` ( do
                                  putErrText "Terminating effect process..."
                                  _ <- System.Process.withCreateProcess (System.Process.proc containerRuntimeExe ["kill", name]) \_ _ _ kh ->
                                    waitForProcess kh
                                  threadDelay 3_000_000
                                  _ <- System.Process.withCreateProcess (System.Process.proc containerRuntimeExe ["kill", name, "KILL"]) \_ _ _ kh ->
                                    waitForProcess kh
                                  putErrText "Killed effect process."
                              )
        )
        ( do
            let shovel =
                  handleEOF (BS.Char8.hGetLine master) >>= \case
                    "" -> pass
                    someBytes | "@nix" `BS.isPrefixOf` someBytes -> do
                      -- TODO use it (example @nix { "action": "setPhase", "phase": "effectPhase" })
                      shovel
                    someBytes -> do
                      BS.hPut stderr (someBytes <> "\n")
                      shovel
                handleEOF = handle \e -> if ioeGetErrorType e == HardwareFault then pure "" else throwIO e
            shovel
        )
  pure exitCode

-- | Like 'openPseudoTerminalHandles' but closes the handles after the
-- function is done.
withPseudoTerminalHandles :: ((Handle, Handle) -> IO a) -> IO a
withPseudoTerminalHandles =
  bracket
    openPseudoTerminalHandles
    ( \(master, terminal) -> do
        hClose master `catch` \(_ :: SomeException) -> pass
        hClose terminal `catch` \(_ :: SomeException) -> pass
    )

-- | Like 'openPseudoTerminal' but returning handles, in a resource-safe manner.
openPseudoTerminalHandles :: IO (Handle, Handle)
openPseudoTerminalHandles =
  mask_ do
    (masterFd, terminalFd) <- openPseudoTerminal

    ( do
        master <- fdToHandle masterFd
        terminal <- fdToHandle terminalFd
        pure (master, terminal)
      )
      `onException` do
        closeFd masterFd
        when (terminalFd /= masterFd) (closeFd terminalFd)