Argonix.hs 13.5 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 29
import           System.Console.ANSI
import           System.Console.ANSI.Types      ( Color )
Valentin Reis's avatar
Valentin Reis committed
30
import           System.Posix.Process
31 32 33 34 35
import qualified System.Process.Typed          as P
                                                ( runProcess
                                                , runProcess_
                                                , setEnv
                                                , readProcessStdout_
36
                                                , proc
37
                                                , shell
38
                                                )
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
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\""
    )
66 67

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

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

106 107
data NixStaticInOut = Both Text
                    | Src Text
Valentin Reis's avatar
Valentin Reis committed
108
                    | SimpleTarget Text
109 110 111
                    | 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
targets :: [Text]
targets = mapMaybe toTarget nixStatic
 where
  toTarget (Target t) = Just t
Valentin Reis's avatar
Valentin Reis committed
143
  toTarget (SimpleTarget t) = Just t
144 145 146 147 148 149 150 151 152 153 154 155 156 157
  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
158
nixStatic = src <> both <> simpletarget <> target
159 160 161 162 163 164 165
 where
  src = [Src "experiments"]
  both =
    Both
      <$> [ "aml"
          , "libnrm"
          , "numabench"
166
          , "repoquality"
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
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."

Valentin Reis's avatar
Valentin Reis committed
332 333 334
runProcessVerbose :: Verbosity -> Text -> [Text] -> [(Text, Text)] -> IO ()
runProcessVerbose a b c d = void $ readProcessVerbose a b c d

Valentin Reis's avatar
Valentin Reis committed
335
setupSystem :: Text -> ArgsCommon -> IO ([(Text, Text)],Text)
336 337 338 339 340 341 342 343 344 345
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
346 347 348
  varXDG_CACHE_HOME <- rpVerbose "mktemp"
                                 ["-d", "/tmp/deletable-nix-cache-XXXX"]
                                 []
Valentin Reis's avatar
Valentin Reis committed
349 350 351 352 353

  varTMPDIR <- rpVerbose "mktemp"
                                 ["-d", "/tmp/tmpdir-XXXX"]
                                 []
  sVerbose $ "chmod 777 " <> varTMPDIR
Valentin Reis's avatar
Valentin Reis committed
354 355

  passVars  <- getVars ["PATH", "NIX_PATH", "NIX_SSL_CERT_FILE"]
356 357

  let envVars =
Valentin Reis's avatar
Valentin Reis committed
358 359
        [("XDG_CACHE_HOME", varXDG_CACHE_HOME), ("TMPDIR", varTMPDIR)]
          ++ passVars
360 361 362

  when (isTarget target) $ setupNodeOs envVars

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

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

411 412 413
wrap :: NixCommand -> Text -> ArgsCommon -> IO ()
wrap nixCommand target sa@ArgsCommon {..} = do
  when (verbosity == Verbose) $ print sa
Valentin Reis's avatar
Valentin Reis committed
414
  (envVars,tmpXDGHOME) <- setupSystem target sa
415 416 417 418 419 420 421 422
  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)
          )
Valentin Reis's avatar
Valentin Reis committed
423 424
        <> runProcessVerbose verbosity "cp" ["-r", "/tmp/papa", "./result"] []
        <> runProcessVerbose verbosity "cp" ["-r", toS tmpXDGHOME, "./nixtmpdir"] []
425 426 427 428 429 430 431 432 433
 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/"-}