{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}

module Test.Sandwich.Contexts.Kubernetes.KindCluster.Images (
  getLoadedImagesKind
  , clusterContainsImageKind
  , loadImageKind
  ) where

import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Data.Aeson as A
import qualified Data.Set as Set
import Data.String.Interpolate
import qualified Data.Vector as V
import Relude
import System.Exit
import System.FilePath
import Test.Sandwich
import Test.Sandwich.Contexts.Kubernetes.KindCluster.Setup
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Contexts.Kubernetes.Util.Aeson
import Test.Sandwich.Contexts.Kubernetes.Util.Images
import UnliftIO.Directory
import UnliftIO.Process
import UnliftIO.Temporary


-- | Load an image into a Kind cluster.
loadImageKind :: (
  HasCallStack, MonadUnliftIO m, MonadLoggerIO m
  )
  -- | Path to @kind@ binary
  => FilePath
  -- | Cluster name
  -> Text
  -- | Image load spec
  -> ImageLoadSpec
  -- | Extra environment variables
  -> Maybe [(String, String)]
  -- | Returns transformed image name
  -> m Text
loadImageKind :: forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLoggerIO m) =>
FilePath
-> Text -> ImageLoadSpec -> Maybe [(FilePath, FilePath)] -> m Text
loadImageKind FilePath
kindBinary Text
clusterName ImageLoadSpec
imageLoadSpec Maybe [(FilePath, FilePath)]
env = do
  case ImageLoadSpec
imageLoadSpec of
    ImageLoadSpecTarball FilePath
image -> do
      FilePath -> m Bool
forall (m :: * -> *). MonadIO m => FilePath -> m Bool
doesDirectoryExist (FilePath -> FilePath
forall a. ToString a => a -> FilePath
toString FilePath
image) m Bool -> (Bool -> 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
        Bool
True ->
          -- Uncompressed directory: tar it up (but don't zip).
          -- TODO: don't depend on external tar binary
          FilePath -> (FilePath -> m Text) -> m Text
forall (m :: * -> *) a.
MonadUnliftIO m =>
FilePath -> (FilePath -> m a) -> m a
withSystemTempDirectory FilePath
"kind-image-zip" ((FilePath -> m Text) -> m Text) -> (FilePath -> m Text) -> m Text
forall a b. (a -> b) -> a -> b
$ \FilePath
dir -> do
            let tarFile :: FilePath
tarFile = FilePath
dir FilePath -> FilePath -> FilePath
</> FilePath
"test.tar"
            _ <- CreateProcess -> FilePath -> m FilePath
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> FilePath -> m FilePath
readCreateProcessWithLogging (FilePath -> CreateProcess
shell [i|tar -C #{image} --dereference --hard-dereference --xform s:'^./':: -c . > #{tarFile}|]) FilePath
""
            imageLoad tarFile
            readUncompressedImageName (toString image)

        Bool
False -> case FilePath -> FilePath
takeExtension (FilePath -> FilePath
forall a. ToString a => a -> FilePath
toString FilePath
image) of
          FilePath
".tar" -> do
            FilePath -> m ()
imageLoad (FilePath -> FilePath
forall a. ToString a => a -> FilePath
toString FilePath
image)
            FilePath -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
FilePath -> m Text
readImageName (FilePath -> FilePath
forall a. ToString a => a -> FilePath
toString FilePath
image)
          FilePath
".gz" -> do
            FilePath -> (FilePath -> m Text) -> m Text
forall (m :: * -> *) a.
MonadUnliftIO m =>
FilePath -> (FilePath -> m a) -> m a
withSystemTempDirectory FilePath
"image-tarball" ((FilePath -> m Text) -> m Text) -> (FilePath -> m Text) -> m Text
forall a b. (a -> b) -> a -> b
$ \FilePath
tempDir -> do
              let tarFile :: FilePath
tarFile = FilePath
tempDir FilePath -> FilePath -> FilePath
</> FilePath
"image.tar"
              -- TODO: don't depend on external gzip binary
              CreateProcess -> m ProcessHandle
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> m ProcessHandle
createProcessWithLogging (FilePath -> CreateProcess
shell [i|cat "#{image}" | gzip -d > "#{tarFile}"|])
                m ProcessHandle -> (ProcessHandle -> m ExitCode) -> m ExitCode
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProcessHandle -> m ExitCode
forall (m :: * -> *). MonadIO m => ProcessHandle -> m ExitCode
waitForProcess m ExitCode -> (ExitCode -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ExitCode -> ExitCode -> m ()
forall (m :: * -> *) a.
(HasCallStack, MonadIO m, Eq a, Show a) =>
a -> a -> m ()
`shouldBe` ExitCode
ExitSuccess)
              FilePath -> m ()
imageLoad FilePath
tarFile
              FilePath -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
FilePath -> m Text
readImageName (FilePath -> FilePath
forall a. ToString a => a -> FilePath
toString FilePath
image)
          FilePath
_ -> FilePath -> m Text
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
FilePath -> m a
expectationFailure [i|Unexpected image extension in #{image}. Wanted .tar, .tar.gz, or uncompressed directory.|]

    ImageLoadSpecDocker Text
image ImagePullPolicy
pullPolicy -> do
      _ <- Text -> ImagePullPolicy -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> ImagePullPolicy -> m Bool
dockerPullIfNecessary Text
image ImagePullPolicy
pullPolicy

      createProcessWithLogging (
        (shell [i|#{kindBinary} load docker-image #{image} --name #{clusterName}|]) {
            env = env
            }) >>= waitForProcess >>= (`shouldBe` ExitSuccess)

      return image
    ImageLoadSpecPodman Text
image ImagePullPolicy
pullPolicy -> do
      _ <- Text -> ImagePullPolicy -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLoggerIO m) =>
Text -> ImagePullPolicy -> m Bool
podmanPullIfNecessary Text
image ImagePullPolicy
pullPolicy

      _ <- expectationFailure [i|Not implemented yet.|]

      return image
  where
    imageLoad :: FilePath -> m ()
imageLoad FilePath
tarFile =
      CreateProcess -> m ProcessHandle
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
CreateProcess -> m ProcessHandle
createProcessWithLogging (
        (FilePath -> CreateProcess
shell [i|#{kindBinary} load image-archive #{tarFile} --name #{clusterName}|]) {
            env = env
            }) m ProcessHandle -> (ProcessHandle -> m ExitCode) -> m ExitCode
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ProcessHandle -> m ExitCode
forall (m :: * -> *). MonadIO m => ProcessHandle -> m ExitCode
waitForProcess m ExitCode -> (ExitCode -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ExitCode -> ExitCode -> m ()
forall (m :: * -> *) a.
(HasCallStack, MonadIO m, Eq a, Show a) =>
a -> a -> m ()
`shouldBe` ExitCode
ExitSuccess)

-- | Get the set of loaded images on the given Kind cluster.
getLoadedImagesKind :: (
  HasCallStack, MonadUnliftIO m, MonadLogger m
  )
  => KubernetesClusterContext
  -- | Driver (should be "docker" or "podman")
  -> Text
  -- | Path to @kind@ binary
  -> FilePath
  -- | Extra environment variables
  -> Maybe [(String, String)]
  -> m (Set Text)
getLoadedImagesKind :: forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext
-> Text -> FilePath -> Maybe [(FilePath, FilePath)] -> m (Set Text)
getLoadedImagesKind KubernetesClusterContext
kcc Text
driver FilePath
kindBinary Maybe [(FilePath, FilePath)]
env = do
  chosenNode <- KubernetesClusterContext
-> FilePath -> Maybe [(FilePath, FilePath)] -> m [Text]
forall (m :: * -> *).
MonadUnliftIO m =>
KubernetesClusterContext
-> FilePath -> Maybe [(FilePath, FilePath)] -> m [Text]
getNodes KubernetesClusterContext
kcc FilePath
kindBinary Maybe [(FilePath, FilePath)]
env m [Text] -> ([Text] -> 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
    (Text
x:[Text]
_) -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
x
    [] -> FilePath -> m Text
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
FilePath -> m a
expectationFailure [i|Couldn't identify a Kind node.|]

  output <- readCreateProcessWithLogging (
    (proc (toString driver) [
        "exec"
        , toString chosenNode
        , "crictl", "images", "-o", "json"
        ]) { env = env }
    ) ""

  case A.eitherDecode (encodeUtf8 output) of
    Left FilePath
err -> FilePath -> m (Set Text)
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
FilePath -> m a
expectationFailure [i|Couldn't decode JSON (#{err}): #{output}|]
    Right (A.Object (Text -> Object -> Maybe Value
forall v. Text -> KeyMap v -> Maybe v
aesonLookup Text
"images" -> Just (A.Array Array
images))) -> Set Text -> m (Set Text)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Set Text -> m (Set Text)) -> Set Text -> m (Set Text)
forall a b. (a -> b) -> a -> b
$ [Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList ([Text] -> Set Text) -> [Text] -> Set Text
forall a b. (a -> b) -> a -> b
$ (Value -> [Text]) -> Array -> [Text]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Value -> [Text]
extractRepoTags Array
images
    Either FilePath Value
_ -> FilePath -> m (Set Text)
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
FilePath -> m a
expectationFailure [i|Unexpected format in JSON: #{output}|]

  where
    extractRepoTags :: A.Value -> [Text]
    extractRepoTags :: Value -> [Text]
extractRepoTags (A.Object (Text -> Object -> Maybe Value
forall v. Text -> KeyMap v -> Maybe v
aesonLookup Text
"repoTags" -> Just (A.Array Array
xs))) = [Text
t | A.String Text
t <- Array -> [Value]
forall a. Vector a -> [a]
V.toList Array
xs]
    extractRepoTags Value
_ = []

-- | Test if the Kind cluster contains a given image.
clusterContainsImageKind :: (
  HasCallStack, MonadUnliftIO m, MonadLogger m
  )
  => KubernetesClusterContext
  -- | Driver (should be "docker" or "podman")
  -> Text
  -- | Path to @kind@ binary
  -> FilePath
  -- | Extra environment variables
  -> Maybe [(String, String)]
  -> Text
  -> m Bool
clusterContainsImageKind :: forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext
-> Text
-> FilePath
-> Maybe [(FilePath, FilePath)]
-> Text
-> m Bool
clusterContainsImageKind KubernetesClusterContext
kcc Text
driver FilePath
kindBinary Maybe [(FilePath, FilePath)]
env Text
image = do
  imageName <- case FilePath -> Bool
isAbsolute (Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
image) of
    Bool
False -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
image
    Bool
True -> FilePath -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
FilePath -> m Text
readImageName (Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
image)

  loadedImages <- getLoadedImagesKind kcc driver kindBinary env

  return (
    imageName `Set.member` loadedImages

    -- Deal with weird prefixing Minikube does; see
    -- https://github.com/kubernetes/minikube/issues/19343
    || ("docker.io/" <> imageName) `Set.member` loadedImages
    || ("docker.io/library/" <> imageName) `Set.member` loadedImages
    )