Argonix.hs 13.2 KB
Newer Older
Valentin Reis's avatar
Valentin Reis committed
1 2
{-# language RecordWildCards #-}
{-# language ApplicativeDo #-}
3
{-# language NoImplicitPrelude #-}
Valentin Reis's avatar
Valentin Reis committed
4
{-# language LambdaCase #-}
Valentin Reis's avatar
Valentin Reis committed
5 6
{-# language OverloadedStrings #-}

Valentin Reis's avatar
Valentin Reis committed
7 8 9 10 11
{-|
Module      : Main
Description : argonix
Copyright   : (c) Valentin Reis, 2018
License     : MIT
Valentin Reis's avatar
Valentin Reis committed
12 13
Maintainer  : fre@freux.fr
-}
Valentin Reis's avatar
Valentin Reis committed
14

15 16 17 18
module Main
  ( main
  )
where
Valentin Reis's avatar
Valentin Reis committed
19 20

import           Data.Default
21
import           Protolude
Valentin Reis's avatar
Valentin Reis committed
22
import           Options.Applicative
23 24 25 26
import qualified Data.Text                     as T
                                                ( lines
                                                , isInfixOf
                                                , unwords
Valentin Reis's avatar
Valentin Reis committed
27
                                                )
Valentin Reis's avatar
Valentin Reis committed
28
import           System.Console.ANSI
29
import qualified System.Directory              as SD
Valentin Reis's avatar
Valentin Reis committed
30
import           System.Console.ANSI.Types      ( Color )
Valentin Reis's avatar
Valentin Reis committed
31
import           System.Posix.Process
32 33 34 35 36
import qualified System.Process.Typed          as P
                                                ( runProcess
                                                , runProcess_
                                                , setEnv
                                                , readProcessStdout_
37
                                                , proc
38
                                                , shell
39
                                                )
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
import           System.Exit                    ( ExitCode )
import qualified System.IO                     as SIO
                                                ( hSetBuffering
                                                , stdout
                                                , BufferMode(..)
                                                )
import qualified System.Environment            as SE

-- | color printing
colorIn :: Color -> IO () -> IO ()
colorIn color she = setC color *> she *> setC White
  where setC c = setSGR [SetColor Foreground Dull c]

checkFsAttributes :: Verbosity -> FilePath -> IO ()
checkFsAttributes verbosity workingDirectory = do
  putVerbose verbosity "Checking filesystem attributes on /tmp"
  findmnt <-
    P.readProcessStdout_
    $  P.shell
    $  "findmnt -T "
    <> workingDirectory
    <> " -o OPTIONS"
  when (T.isInfixOf "nosuid" $ toS findmnt) $ die
    (  "The output directory, "
    <> toS workingDirectory
    <> " must not mounted with \"nosuid\""
    )
67 68

main :: IO ()
69 70
main = SIO.hSetBuffering SIO.stdout SIO.NoBuffering
  <> void (join (execParser (info (opts <**> helper) idm)))
71 72 73 74 75
 where
  opts :: Parser (IO ())
  opts = hsubparser
    (  command
        "build"
Valentin Reis's avatar
Valentin Reis committed
76
        (info (wrap NixBuild <$> targetParser <*> commonParser)
77 78 79
              (progDesc "Run an argo-compatible nix-build.")
        )
    <> command
Valentin Reis's avatar
Valentin Reis committed
80
         "shell"
Valentin Reis's avatar
Valentin Reis committed
81
         (info (wrap NixShell <$> targetParser <*> commonParser)
Valentin Reis's avatar
Valentin Reis committed
82
               (progDesc "Enter an argo-compatible nix-shell")
83
         )
Valentin Reis's avatar
Valentin Reis committed
84 85 86 87 88 89 90 91
    <> command
         "remote-shell"
         (info
           (remotely <$> targetParser <*> commonParser <*> remoteParser)
           (progDesc
             "Enter an argo-compatible nix-shell on a remote machine with nix enabled"
           )
         )
92 93 94
    <> help "Type of operation to run."
    )

Valentin Reis's avatar
Valentin Reis committed
95
data ArgsCommon = ArgsCommon
Valentin Reis's avatar
Valentin Reis committed
96 97 98
  { argopkgs   :: Text
   , verbosity :: Verbosity
   , run       :: Maybe Text
99
   , overrides :: [(Text, Text)]
Valentin Reis's avatar
Valentin Reis committed
100
   , grafting  :: Grafting
Valentin Reis's avatar
Valentin Reis committed
101
   , sandboxing :: Sandboxing
Valentin Reis's avatar
Valentin Reis committed
102
  } deriving (Show)
Valentin Reis's avatar
Valentin Reis committed
103
data Verbosity = Verbose | Normal deriving (Show, Eq)
Valentin Reis's avatar
Valentin Reis committed
104
data Sandboxing = Sandbox | NoSandbox deriving (Show, Eq)
Valentin Reis's avatar
Valentin Reis committed
105
data Grafting = Libnrm | NoGraft deriving (Show, Eq)
Valentin Reis's avatar
Valentin Reis committed
106

107 108 109 110 111
data NixStaticInOut = Both Text
                    | Src Text
                    | Target Text deriving (Eq)
data NixCommand = NixBuild | NixShell

Valentin Reis's avatar
Valentin Reis committed
112
data ArgsRemote = ArgsRemote
Valentin Reis's avatar
style  
Valentin Reis committed
113 114
  {  targetMachine
   , retreive
Valentin Reis's avatar
Valentin Reis committed
115 116 117
   , retreiveAs    :: Maybe Text
  } deriving (Show)

118 119 120 121
toCommand :: NixCommand -> Text
toCommand NixBuild = "nix-build"
toCommand NixShell = "nix-shell"

Valentin Reis's avatar
Valentin Reis committed
122 123
instance Default ArgsCommon where
  def = ArgsCommon
Valentin Reis's avatar
Valentin Reis committed
124 125 126 127 128 129
    { verbosity  = Normal,
      argopkgs   = "<argopkgs>",
      run        = Nothing,
      overrides  = [],
      grafting   = NoGraft,
      sandboxing = NoSandbox
Valentin Reis's avatar
Valentin Reis committed
130 131
    }

Valentin Reis's avatar
Valentin Reis committed
132 133 134
instance Default ArgsRemote where
  def = ArgsRemote
    { targetMachine = Nothing,
Valentin Reis's avatar
Valentin Reis committed
135 136
      retreive      = Nothing,
      retreiveAs    = Nothing
Valentin Reis's avatar
Valentin Reis committed
137 138
    }

139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
targets :: [Text]
targets = mapMaybe toTarget nixStatic
 where
  toTarget (Target t) = Just t
  toTarget (Both   t) = Just t
  toTarget _          = Nothing

sources :: [Text]
sources = mapMaybe toSrc nixStatic
 where
  toSrc (Src  t) = Just t
  toSrc (Both t) = Just t
  toSrc _        = Nothing

isTarget :: Text -> Bool
isTarget x = Target x `elem` nixStatic

nixStatic :: [NixStaticInOut]
nixStatic = src <> both <> target
 where
  src = [Src "experiments"]
  both =
    Both
      <$> [ "aml"
          , "libnrm"
          , "numabench"
          , "nrm"
          , "containers"
          , "argotk"
          , "amg"
          , "lammps"
          , "qmcpack"
          , "stream"
          , "openmc"
          , "argonix"
          ]
  target =
    Target
      <$> [ "powerexpe"
          , "test"
          , "report"
          , "testHello"
          , "testListen"
          , "testListen"
          , "testHello"
          , "testListen"
          , "testPerfwrapper"
          , "testPower"
          , "testSTREAM"
          , "testAMG"
          , "testOpenMC"
          , "testLAMMPS"
          , "testQMCPack"
          , "testAll"
          ]

targetParser :: Parser Text
targetParser = strArgument
  (metavar "TARGET" <> showDefault <> help
    (toS ("The build target, in " <> mconcat ts))
  )
  where ts = intersperse " " targets
201

Valentin Reis's avatar
Valentin Reis committed
202 203
remoteParser :: Parser ArgsRemote
remoteParser = do
Valentin Reis's avatar
Valentin Reis committed
204 205 206 207 208 209 210 211 212 213 214 215
  retreive <- optional $ strOption
    (long "retreive" <> metavar "RELATIVE_PATH" <> help
      "File/Folder to retreive from the remote machine."
    )
  retreiveAs <- optional $ strOption
    (long "retreiveAs" <> metavar "LOCAL_PATH" <> help
      "File/Folder to save retreived data as (locally)."
    )
  targetMachine <- optional $ strOption
    (long "target_machine" <> metavar "USER@HOST" <> help
      "Target machine. defaults to localhost via forking."
    )
Valentin Reis's avatar
Valentin Reis committed
216 217 218 219 220 221 222 223 224 225 226 227 228 229 230
  pure ArgsRemote {..}

commonParser :: Parser ArgsCommon
commonParser = do
  argopkgs <- strOption
    (  long "argopkgs"
    <> metavar "ARGOPKGS"
    <> showDefault
    <> value (argopkgs def)
    <> help "Nix expression that produces the argopkgs source path."
    )
  verbosity <- flag
    Normal
    Verbose
    (long "verbose" <> short 'v' <> help "Enable verbose mode")
Valentin Reis's avatar
Valentin Reis committed
231 232 233 234
  sandboxing <- flag
    NoSandbox
    Sandbox
    (long "sandboxing" <> short 's' <> help "Enable nix sandboxing.")
Valentin Reis's avatar
Valentin Reis committed
235 236 237 238 239 240
  grafting <- flag
    NoGraft
    Libnrm
    (long "grafting" <> short 'g' <> help
      "Enable libnrm grafting to avoid rebuilding applications."
    )
241
  overrides <- catMaybes <$> ts
Valentin Reis's avatar
Valentin Reis committed
242 243 244 245
  run       <- optional $ strOption
    (long "run" <> metavar "COMMAND" <> help
      "Command to run the environment instead of an interactive shell"
    )
Valentin Reis's avatar
Valentin Reis committed
246
  pure ArgsCommon {..}
Valentin Reis's avatar
Valentin Reis committed
247
 where
248 249 250
  ts :: Parser [Maybe (Text, Text)]
  ts = traverse optSrc sources
  optSrc :: Text -> Parser (Maybe (Text, Text))
Valentin Reis's avatar
Valentin Reis committed
251 252
  optSrc longform = do
    parsed <- optional $ strOption
253 254
      (long (toS longform) <> metavar "PATH" <> help
        (toS longform <> " source folder override.")
Valentin Reis's avatar
Valentin Reis committed
255 256
      )
    pure $ mapT longform parsed
257
  mapT :: Text -> Maybe Text -> Maybe (Text, Text)
Valentin Reis's avatar
Valentin Reis committed
258 259 260 261
  mapT longform thePath = case thePath of
    Nothing -> Nothing
    Just p  -> Just (longform, p)

262 263 264 265 266
sudoRemoveFile :: Text -> FilePath -> IO ()
sudoRemoveFile desc filePath = do
  P.runProcess_ (P.shell $ "test " <> filePath)
  putText $ "found file " <> desc <> " at " <> toS filePath <> ".. "
  go False
Valentin Reis's avatar
Valentin Reis committed
267
 where
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293
  go useSudo = P.runProcess (sudorm useSudo) >>= \case
    ExitSuccess   -> colorIn Green $ putText "Successfully removed."
    ExitFailure _ -> if useSudo then reportFailure else trySudo
  sudorm useSudo =
    P.shell $ (if useSudo then "sudo " else "") <> "rm -rf " <> filePath
  trySudo = do
    putText
      $  "Failed to remove "
      <> desc
      <> " at "
      <> toS filePath
      <> ". Trying sudo..\n"
    go True
  reportFailure =
    colorIn Red
      $  putText
      $  "Failed to remove "
      <> desc
      <> " at "
      <> toS filePath
      <> ", even with sudo."


cleanSockets :: Verbosity -> IO ()
cleanSockets verbosity =
  putVerbose verbosity "cleaning sockets" >> for_ socketList cleanSocket
Valentin Reis's avatar
Valentin Reis committed
294
 where
295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
  cleanSocket :: FilePath -> IO ()
  cleanSocket = sudoRemoveFile "socket"

  socketList :: [FilePath]
  socketList =
    [ "/tmp/nrm-downstream-in"
    , "/tmp/nrm-downstream-event"
    , "/tmp/nrm-upstream-in"
    , "/tmp/nrm-upstream-event"
    ]

putVerbose :: MonadIO m => Verbosity -> Text -> m ()
putVerbose verbosity = when (verbosity == Verbose) . putText

shellVerbose :: MonadIO m => Verbosity -> Text -> m ExitCode
shellVerbose verbosity s = do
  putVerbose verbosity s
  P.runProcess (P.shell $ toS s)

Valentin Reis's avatar
Valentin Reis committed
314
readProcessVerbose :: Verbosity -> Text -> [Text] -> [(Text, Text)] -> IO Text
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338
readProcessVerbose verbosity name arguments envVars = do
  putVerbose verbosity $ name <> mconcat (intersperse " " arguments)
  head
    .   T.lines
    .   toS
    <$> P.readProcessStdout_
          ( P.setEnv (fmap (\(x, y) -> (toS x, toS y)) envVars)
          $ P.proc (toS name) (fmap toS arguments)
          )
    >>= \case
          Just r  -> return r
          Nothing -> die "readProcess returned more than one line."

setupSystem :: Text -> ArgsCommon -> IO [(Text, Text)]
setupSystem target sa = do

  sVerbose "sudo rm -rf result" >>= \case
    ExitSuccess   -> putText "removed ./result"
    ExitFailure n -> die ("Failed to remove ./result " <> show n)

  sVerbose "sudo rm -rf nixtmpdir" >>= \case
    ExitSuccess   -> putText "removed ./nixtmpdir"
    ExitFailure n -> die ("Failed to remove ./nixtmpdir " <> show n)

Valentin Reis's avatar
Valentin Reis committed
339 340 341 342 343 344 345
  varXDG_CACHE_HOME <- rpVerbose "mktemp"
                                 ["-d", "/tmp/deletable-nix-cache-XXXX"]
                                 []
  varTMPDIR            <- createAndSetTMPDIR "nixtmpdir"
  varNIX_PATH          <- toS <$> SE.getEnv "NIX_PATH"
  varPATH              <- toS <$> SE.getEnv "PATH"
  varNIX_SSL_CERT_FILE <- toS <$> SE.getEnv "NIX_SSL_CERT_FILE"
346 347

  let envVars =
Valentin Reis's avatar
Valentin Reis committed
348 349 350 351 352
        [ ("XDG_CACHE_HOME"   , varXDG_CACHE_HOME)
        , ("TMPDIR"           , varTMPDIR)
        , ("NIX_PATH"         , varNIX_PATH)
        , ("PATH"             , varPATH)
        , ("NIX_SSL_CERT_FILE", varNIX_SSL_CERT_FILE)
353 354 355 356 357 358
        ]

  when (isTarget target) $ setupNodeOs envVars

  return envVars
 where
Valentin Reis's avatar
Valentin Reis committed
359 360
  setupNodeOs :: [(Text, Text)] -> IO ()
  setupNodeOs env = do
361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
    cleanSockets $ verbosity sa
    checkFsAttributes (verbosity sa) "/tmp"
    nodeos_config <- rpVerbose
      "nix-build"
      (fmap toS (nixArguments "containers" sa) ++ ["--no-out-link"])
      env
    sVerbose "sudo rm -rf /tmp/argo_nodeos_config"
    sVerbose
      $  "cp "
      <> toS nodeos_config
      <> "/bin/argo_nodeos_config /tmp/argo_nodeos_config"
    sVerbose "sudo chown root:root /tmp/argo_nodeos_config" >>= \case
      ExitSuccess -> return ()
      ExitFailure n ->
        die ("Failed to set argo_nodeos_config permissions " <> show n)
    sVerbose "sudo chmod u+sw /tmp/argo_nodeos_config" >>= \case
      ExitSuccess -> return ()
Valentin Reis's avatar
Valentin Reis committed
378 379 380 381 382 383 384
      ExitFailure n ->
        die ("Setting suid bit failed with exit code " <> show n)
    sVerbose "sudo /tmp/argo_nodeos_config --clean_config=kill_content:true"
      >>= \case
            ExitSuccess -> return ()
            ExitFailure n ->
              die ("NodeOS config cleaning failed with exit code" <> show n)
385 386 387 388
  sVerbose  = shellVerbose $ verbosity sa
  rpVerbose = readProcessVerbose $ verbosity sa

  createAndSetTMPDIR :: Text -> IO Text
Valentin Reis's avatar
Valentin Reis committed
389
  createAndSetTMPDIR name = do
390 391 392 393 394 395 396
    let p = "/tmp/" <> name
    SD.createDirectoryIfMissing True (toS p)
    sVerbose $ "chmod 777 " <> p
    sVerbose $ "ln -s " <> p <> " " <> name
    return $ toS p

nixArguments :: Text -> ArgsCommon -> [Text]
Valentin Reis's avatar
Valentin Reis committed
397
nixArguments target ArgsCommon {..} =
398
  [argopkgs, "-A", target]
Valentin Reis's avatar
Valentin Reis committed
399
    ++ concat [ ["--arg", longform <> "-src", p] | (longform, p) <- overrides ]
Valentin Reis's avatar
Valentin Reis committed
400
    ++ (if grafting == Libnrm then ["--arg", "graftLibnrm", "true"] else [])
401 402
    ++ ["-o", "/tmp/papa"]
    ++ [ "--show-trace" | v ]
403
  where v = verbosity == Verbose
404

405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426
wrap :: NixCommand -> Text -> ArgsCommon -> IO ()
wrap nixCommand target sa@ArgsCommon {..} = do
  when (verbosity == Verbose) $ print sa
  envVars <- setupSystem target sa
  putText $ nixc <> " " <> T.unwords arglist
  case nixCommand of
    NixShell -> executeFile (toS nixc) True (fmap toS arglist) Nothing
    NixBuild ->
      P.runProcess_
          (P.setEnv (fmap (\(x, y) -> (toS x, toS y)) envVars)
                    (P.proc (toS nixc) $ fmap toS arglist)
          )
        <> P.runProcess_ (P.proc "cp" ["-r", "/tmp/papa", "./result"])
 where
  nixc = toCommand nixCommand
  arglist =
    nixArguments target sa
      ++ ["--allow-new-privileges", "-K", "--option", "build-use-sandbox"]
      ++ [if sandboxing == Sandbox then "true" else "false"]

remotely :: Text -> ArgsCommon -> ArgsRemote -> IO ()
remotely _ _ _ = putText "unsupported in this version" >> undefined
Valentin Reis's avatar
Valentin Reis committed
427 428 429 430 431 432

-- no sandbox use is possible. If sandbox paths were read only we could add:
   {-, "--option"-}
   {-, "build-use-sandbox"-}
   {-, "--option"-}
   {-, "extra-sandbox-paths"-}
Valentin Reis's avatar
Valentin Reis committed
433
  {-, "/tmp/ /etc/argo/ /var/run/ /var/lock/"-}