{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}

{-|
Helper module for working with containers.
-}

module Test.Sandwich.Contexts.Container (
  ContainerOptions (..)
  , defaultContainerOptions

  , ContainerSystem (..)
  , waitForHealth

  -- * Container/host conversions
  , containerPortToHostPort
  , containerNameToContainerId

  -- * Misc
  , isInContainer
  ) where

import Control.Monad.Catch
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Control.Retry
import Data.Aeson as A
import Data.Aeson.TH as A
import qualified Data.List as L
import qualified Data.Map as M
import Data.String.Interpolate
import qualified Data.Text as T
import Network.Socket (PortNumber)
import Relude
import Safe
import System.Exit
import Test.Sandwich
import qualified Text.Show
import UnliftIO.Process


-- | Type to represent generic options for launching containers.
data ContainerOptions = ContainerOptions {
  ContainerOptions -> ContainerSystem
containerOptionsSystem :: ContainerSystem
  , ContainerOptions -> Maybe Text
containerOptionsName :: Maybe Text
  }
  deriving (Int -> ContainerOptions -> ShowS
[ContainerOptions] -> ShowS
ContainerOptions -> String
(Int -> ContainerOptions -> ShowS)
-> (ContainerOptions -> String)
-> ([ContainerOptions] -> ShowS)
-> Show ContainerOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ContainerOptions -> ShowS
showsPrec :: Int -> ContainerOptions -> ShowS
$cshow :: ContainerOptions -> String
show :: ContainerOptions -> String
$cshowList :: [ContainerOptions] -> ShowS
showList :: [ContainerOptions] -> ShowS
Show, ContainerOptions -> ContainerOptions -> Bool
(ContainerOptions -> ContainerOptions -> Bool)
-> (ContainerOptions -> ContainerOptions -> Bool)
-> Eq ContainerOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ContainerOptions -> ContainerOptions -> Bool
== :: ContainerOptions -> ContainerOptions -> Bool
$c/= :: ContainerOptions -> ContainerOptions -> Bool
/= :: ContainerOptions -> ContainerOptions -> Bool
Eq)

defaultContainerOptions :: ContainerOptions
defaultContainerOptions :: ContainerOptions
defaultContainerOptions = ContainerOptions {
  containerOptionsSystem :: ContainerSystem
containerOptionsSystem = ContainerSystem
ContainerSystemPodman
  , containerOptionsName :: Maybe Text
containerOptionsName = Maybe Text
forall a. Maybe a
Nothing
  }

-- | Type to represent which container system we're using.
data ContainerSystem =
  ContainerSystemDocker
  | ContainerSystemPodman
  deriving (ContainerSystem -> ContainerSystem -> Bool
(ContainerSystem -> ContainerSystem -> Bool)
-> (ContainerSystem -> ContainerSystem -> Bool)
-> Eq ContainerSystem
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ContainerSystem -> ContainerSystem -> Bool
== :: ContainerSystem -> ContainerSystem -> Bool
$c/= :: ContainerSystem -> ContainerSystem -> Bool
/= :: ContainerSystem -> ContainerSystem -> Bool
Eq)

instance Show ContainerSystem where
  show :: ContainerSystem -> String
show ContainerSystem
ContainerSystemDocker = String
"docker"
  show ContainerSystem
ContainerSystemPodman = String
"podman"

-- | Test if the test process is currently running in a container.
isInContainer :: MonadIO m => m Bool
isInContainer :: forall (m :: * -> *). MonadIO m => m Bool
isInContainer = do
  Text
output <- String -> Text
forall a. ToText a => a -> Text
toText (String -> Text) -> m String -> m Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CreateProcess -> String -> m String
forall (m :: * -> *).
MonadIO m =>
CreateProcess -> String -> m String
readCreateProcess (String -> CreateProcess
shell String
"cat /proc/1/sched | head -n 1") String
""
  Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
    (Text
"init" Text -> Text -> Bool
`T.isInfixOf` Text
output)
    Bool -> Bool -> Bool
|| (Text
"systemd" Text -> Text -> Bool
`T.isInfixOf` Text
output)
    Bool -> Bool -> Bool
|| (Text
"bwrap" Text -> Text -> Bool
`T.isInfixOf` Text
output)

-- | Wait for a container to be in a healthy state.
waitForHealth :: forall m. (HasCallStack, MonadLoggerIO m, MonadMask m) => ContainerSystem -> Text -> m ()
waitForHealth :: forall (m :: * -> *).
(HasCallStack, MonadLoggerIO m, MonadMask m) =>
ContainerSystem -> Text -> m ()
waitForHealth ContainerSystem
containerSystem Text
containerID = do
  let policy :: RetryPolicyM m
policy = Int -> RetryPolicyM m -> RetryPolicyM m
forall (m :: * -> *).
Monad m =>
Int -> RetryPolicyM m -> RetryPolicyM m
limitRetriesByCumulativeDelay (Int
60 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1_000_000) (RetryPolicyM m -> RetryPolicyM m)
-> RetryPolicyM m -> RetryPolicyM m
forall a b. (a -> b) -> a -> b
$ Int -> RetryPolicyM m -> RetryPolicyM m
forall (m :: * -> *).
Monad m =>
Int -> RetryPolicyM m -> RetryPolicyM m
capDelay Int
1_000_000 (RetryPolicyM m -> RetryPolicyM m)
-> RetryPolicyM m -> RetryPolicyM m
forall a b. (a -> b) -> a -> b
$ Int -> RetryPolicyM m
forall (m :: * -> *). Monad m => Int -> RetryPolicyM m
exponentialBackoff Int
1000
  RetryPolicyM m -> (RetryStatus -> m ()) -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
RetryPolicyM m -> (RetryStatus -> m a) -> m a
recoverAll RetryPolicyM m
policy ((RetryStatus -> m ()) -> m ()) -> (RetryStatus -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \RetryStatus
_ -> do
    Text
health <- (Text -> Text
T.strip (Text -> Text) -> (String -> Text) -> String -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
forall a. ToText a => a -> Text
toText) (String -> Text) -> m String -> m Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (CreateProcess -> String -> m String
forall (m :: * -> *).
MonadIO m =>
CreateProcess -> String -> m String
readCreateProcess (
      String -> CreateProcess
shell [i|#{containerSystem} inspect --format "{{json .State.Health.Status }}" #{containerID}|]) String
""
      )

    case Text
health of
      Text
"\"healthy\"" -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      Text
_ -> do
        -- Try running the health check manually, when possible.
        -- This is a workaround for rootless podman failing to have working healthchecks.
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ContainerSystem
containerSystem ContainerSystem -> ContainerSystem -> Bool
forall a. Eq a => a -> a -> Bool
== ContainerSystem
ContainerSystemPodman) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
          -- TODO: use createProcessWithLogging here?
          (ExitCode
exitCode, String
sout, String
serr) <- CreateProcess -> String -> m (ExitCode, String, String)
forall (m :: * -> *).
MonadIO m =>
CreateProcess -> String -> m (ExitCode, String, String)
readCreateProcessWithExitCode (String -> [String] -> CreateProcess
proc String
"podman" [String
"healthcheck", String
"run", Text -> String
forall a. ToString a => a -> String
toString Text
containerID]) String
""
          Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ExitCode
exitCode ExitCode -> ExitCode -> Bool
forall a. Eq a => a -> a -> Bool
/= ExitCode
ExitSuccess) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
            Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|Failed to manually run healthcheck. Code: #{exitCode}. Stdout: '#{sout}'. Stderr: '#{serr}'.|]

        String -> m ()
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Health was: #{health}|]


data HostPortInfo = HostPortInfo {
  HostPortInfo -> Text
hostPortInfoHostIp :: Text
  , HostPortInfo -> Text
hostPortInfoHostPort :: Text
  }
deriveJSON (A.defaultOptions { A.fieldLabelModifier = L.drop (L.length ("hostPortInfo" :: String)) }) ''HostPortInfo

-- | Map a port number inside a container to a port number on the host.
containerPortToHostPort :: (HasCallStack, MonadIO m) => ContainerSystem -> Text -> PortNumber -> m PortNumber
containerPortToHostPort :: forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
ContainerSystem -> Text -> PortNumber -> m PortNumber
containerPortToHostPort ContainerSystem
containerSystem Text
containerName PortNumber
containerPort = do
  let inspectPortCmd :: String
inspectPortCmd = [i|#{containerSystem} inspect --format='{{json .NetworkSettings.Ports}}' #{containerName}|]

  Text
rawNetworkSettings <- IO (ExitCode, String, String) -> m (ExitCode, String, String)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (CreateProcess -> String -> IO (ExitCode, String, String)
forall (m :: * -> *).
MonadIO m =>
CreateProcess -> String -> m (ExitCode, String, String)
readCreateProcessWithExitCode (String -> CreateProcess
shell String
inspectPortCmd) String
"") m (ExitCode, String, String)
-> ((ExitCode, String, String) -> m Text) -> m Text
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    (ExitCode
ExitSuccess, String
sout, String
_serr) -> Text -> m Text
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text -> m Text) -> Text -> m Text
forall a b. (a -> b) -> a -> b
$ Text -> Text
T.strip (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ String -> Text
forall a. ToText a => a -> Text
toText String
sout
    (ExitFailure Int
n, String
sout, String
serr) -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Failed to read container ports (error code #{n}). Stdout: '#{sout}'. Stderr: '#{serr}'.|]

  Map Text [HostPortInfo]
networkSettings :: Map Text [HostPortInfo] <- case ByteString -> Either String (Map Text [HostPortInfo])
forall a. FromJSON a => ByteString -> Either String a
A.eitherDecode (Text -> ByteString
forall a b. ConvertUtf8 a b => a -> b
encodeUtf8 Text
rawNetworkSettings) of
    Left String
err -> String -> m (Map Text [HostPortInfo])
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Failed to decode network settings: #{err}. Settings were #{rawNetworkSettings}.|]
    Right Map Text [HostPortInfo]
x -> Map Text [HostPortInfo] -> m (Map Text [HostPortInfo])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Text [HostPortInfo]
x

  Text
rawPort <- case Text -> Map Text [HostPortInfo] -> Maybe [HostPortInfo]
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup [i|#{containerPort}/tcp|] Map Text [HostPortInfo]
networkSettings of
    Just (HostPortInfo
x:[HostPortInfo]
_) -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> m Text) -> Text -> m Text
forall a b. (a -> b) -> a -> b
$ HostPortInfo -> Text
hostPortInfoHostPort HostPortInfo
x
    Maybe [HostPortInfo]
_ -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Couldn't find any host ports corresponding to container port #{containerPort}. Network settings: #{A.encode networkSettings}|]

  case String -> Maybe PortNumber
forall a. Read a => String -> Maybe a
readMay (Text -> String
forall a. ToString a => a -> String
toString Text
rawPort) of
    Just PortNumber
x -> PortNumber -> m PortNumber
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PortNumber
x
    Maybe PortNumber
Nothing -> String -> m PortNumber
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Couldn't read container port number: '#{rawPort}'|]

-- | Convert a container name to a container ID.
containerNameToContainerId :: (HasCallStack, MonadIO m) => ContainerSystem -> Text -> m Text
containerNameToContainerId :: forall (m :: * -> *).
(HasCallStack, MonadIO m) =>
ContainerSystem -> Text -> m Text
containerNameToContainerId ContainerSystem
containerSystem Text
containerName = do
  let cmd :: String
cmd = [i|#{containerSystem} inspect --format='{{.Id}}' #{containerName}|]
  IO (ExitCode, String, String) -> m (ExitCode, String, String)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (CreateProcess -> String -> IO (ExitCode, String, String)
forall (m :: * -> *).
MonadIO m =>
CreateProcess -> String -> m (ExitCode, String, String)
readCreateProcessWithExitCode (String -> CreateProcess
shell String
cmd) String
"") m (ExitCode, String, String)
-> ((ExitCode, String, String) -> m Text) -> m Text
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    (ExitCode
ExitSuccess, String
sout, String
_serr) -> Text -> m Text
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text -> m Text) -> Text -> m Text
forall a b. (a -> b) -> a -> b
$ Text -> Text
T.strip (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ String -> Text
forall a. ToText a => a -> Text
toText String
sout
    (ExitFailure Int
n, String
sout, String
serr) -> String -> m Text
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Failed to obtain container ID for container named '#{containerName}'. Code: #{n}. Stdout: '#{sout}'. Stderr: '#{serr}'.|]