Argonix.hs 13.3 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
data NixStaticInOut = Both Text
                    | Src Text
Valentin Reis's avatar
Valentin Reis committed
109
                    | SimpleTarget Text
110 111 112
                    | Target Text deriving (Eq)
data NixCommand = NixBuild | NixShell

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

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

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

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

140 141 142 143
targets :: [Text]
targets = mapMaybe toTarget nixStatic
 where
  toTarget (Target t) = Just t
Valentin Reis's avatar
Valentin Reis committed
144
  toTarget (SimpleTarget t) = Just t
145 146 147 148 149 150 151 152 153 154 155 156 157 158
  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]
Valentin Reis's avatar
Valentin Reis committed
159
nixStatic = src <> both <> simpletarget <> target
160 161 162 163 164 165 166
 where
  src = [Src "experiments"]
  both =
    Both
      <$> [ "aml"
          , "libnrm"
          , "numabench"
Valentin Reis's avatar
Valentin Reis committed
167
          , "excit"
168 169 170 171 172 173 174 175 176 177
          , "nrm"
          , "containers"
          , "argotk"
          , "amg"
          , "lammps"
          , "qmcpack"
          , "stream"
          , "openmc"
          , "argonix"
          ]
Valentin Reis's avatar
Valentin Reis committed
178
  simpletarget = SimpleTarget <$> ["numabench-check" , "excit-check"]
179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204
  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
205

Valentin Reis's avatar
Valentin Reis committed
206 207
remoteParser :: Parser ArgsRemote
remoteParser = do
Valentin Reis's avatar
Valentin Reis committed
208 209 210 211 212 213 214 215 216 217 218 219
  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
220 221 222 223 224 225 226 227 228 229 230 231 232 233 234
  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
235 236 237 238
  sandboxing <- flag
    NoSandbox
    Sandbox
    (long "sandboxing" <> short 's' <> help "Enable nix sandboxing.")
Valentin Reis's avatar
Valentin Reis committed
239 240 241 242 243 244
  grafting <- flag
    NoGraft
    Libnrm
    (long "grafting" <> short 'g' <> help
      "Enable libnrm grafting to avoid rebuilding applications."
    )
245
  overrides <- catMaybes <$> ts
Valentin Reis's avatar
Valentin Reis committed
246 247 248 249
  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
250
  pure ArgsCommon {..}
Valentin Reis's avatar
Valentin Reis committed
251
 where
252 253 254
  ts :: Parser [Maybe (Text, Text)]
  ts = traverse optSrc sources
  optSrc :: Text -> Parser (Maybe (Text, Text))
Valentin Reis's avatar
Valentin Reis committed
255 256
  optSrc longform = do
    parsed <- optional $ strOption
257 258
      (long (toS longform) <> metavar "PATH" <> help
        (toS longform <> " source folder override.")
Valentin Reis's avatar
Valentin Reis committed
259 260
      )
    pure $ mapT longform parsed
261
  mapT :: Text -> Maybe Text -> Maybe (Text, Text)
Valentin Reis's avatar
Valentin Reis committed
262 263 264 265
  mapT longform thePath = case thePath of
    Nothing -> Nothing
    Just p  -> Just (longform, p)

266 267 268 269 270
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
271
 where
272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297
  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
298
 where
299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317
  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
318
readProcessVerbose :: Verbosity -> Text -> [Text] -> [(Text, Text)] -> IO Text
319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342
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
343 344 345
  varXDG_CACHE_HOME <- rpVerbose "mktemp"
                                 ["-d", "/tmp/deletable-nix-cache-XXXX"]
                                 []
Valentin Reis's avatar
Valentin Reis committed
346 347 348
  varTMPDIR <- createAndSetTMPDIR "nixtmpdir"

  passVars  <- getVars ["PATH", "NIX_PATH", "NIX_SSL_CERT_FILE"]
349 350

  let envVars =
Valentin Reis's avatar
Valentin Reis committed
351 352
        [("XDG_CACHE_HOME", varXDG_CACHE_HOME), ("TMPDIR", varTMPDIR)]
          ++ passVars
353 354 355 356 357

  when (isTarget target) $ setupNodeOs envVars

  return envVars
 where
Valentin Reis's avatar
Valentin Reis committed
358 359 360 361 362 363 364 365
  getVars :: [Text] -> IO [(Text, Text)]
  getVars vars = catMaybes <$> mapM maybeGet vars

  maybeGet :: Text -> IO (Maybe (Text, Text))
  maybeGet var = SE.lookupEnv (toS var) >>= \case
    Nothing  -> return Nothing
    Just val -> return $ Just (var, toS val)

Valentin Reis's avatar
Valentin Reis committed
366 367
  setupNodeOs :: [(Text, Text)] -> IO ()
  setupNodeOs env = do
368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384
    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
385 386 387 388 389 390 391
      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)
392 393 394 395
  sVerbose  = shellVerbose $ verbosity sa
  rpVerbose = readProcessVerbose $ verbosity sa

  createAndSetTMPDIR :: Text -> IO Text
Valentin Reis's avatar
Valentin Reis committed
396
  createAndSetTMPDIR name = do
397 398 399 400 401 402 403
    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
404
nixArguments target ArgsCommon {..} =
405
  [argopkgs, "-A", target]
Valentin Reis's avatar
Valentin Reis committed
406
    ++ concat [ ["--arg", longform <> "-src", p] | (longform, p) <- overrides ]
Valentin Reis's avatar
Valentin Reis committed
407
    ++ (if grafting == Libnrm then ["--arg", "graftLibnrm", "true"] else [])
408 409
    ++ ["-o", "/tmp/papa"]
    ++ [ "--show-trace" | v ]
410
  where v = verbosity == Verbose
411

412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433
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
434 435 436 437 438 439

-- 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
440
  {-, "/tmp/ /etc/argo/ /var/run/ /var/lock/"-}