{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeApplications #-}

module Test.Sandwich.Contexts.Kubernetes.KubectlPortForward (
  withKubectlPortForward
  , withKubectlPortForward'
  , KubectlPortForwardContext (..)
  ) where

import Control.Monad
import Control.Monad.IO.Unlift
import Control.Retry
import Data.String.Interpolate
import qualified Data.Text as T
import Network.Socket (PortNumber)
import Relude hiding (withFile)
import System.FilePath
import System.Process (getPid)
import Test.Sandwich
import Test.Sandwich.Contexts.Files
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Contexts.Kubernetes.Util.Ports
import Test.Sandwich.Contexts.Kubernetes.Util.SocketUtil
import Test.Sandwich.Util.Process (gracefullyStopProcess)
import UnliftIO.Async
import UnliftIO.Directory
import UnliftIO.Exception
import UnliftIO.IO
import UnliftIO.Process


-- * Types

newtype KubectlPortForwardContext = KubectlPortForwardContext {
  KubectlPortForwardContext -> PortNumber
kubectlPortForwardPort :: PortNumber
  }

-- * Implementation

-- | Run a @kubectl port-forward@ process, making the port available in the 'KubectlPortForwardContext'.
--
-- Note that this will stop working if the pod you're talking to goes away (even if you do it against a service).
-- If this happens, a rerun of the command is needed to resume port forwarding.
withKubectlPortForward :: (
  HasCallStack, KubectlBasic context m
  )
  -- | Path to kubeconfig file
  => FilePath
  -- | Namespace
  -> Text
  -- | Target name (pod, service, etc.)
  -> Text
  -- | Target port number
  -> PortNumber
  -> (KubectlPortForwardContext -> m a)
  -> m a
withKubectlPortForward :: forall context (m :: * -> *) a.
(HasCallStack, KubectlBasic context m) =>
FilePath
-> Text
-> Text
-> PortNumber
-> (KubectlPortForwardContext -> m a)
-> m a
withKubectlPortForward FilePath
kubeConfigFile Text
namespace Text
targetName PortNumber
targetPort KubectlPortForwardContext -> m a
action = do
  kubectlBinary <- forall (a :: Symbol) context (m :: * -> *).
(MonadReader context m, HasFile context a) =>
m FilePath
askFile @"kubectl"
  withKubectlPortForward' kubectlBinary kubeConfigFile namespace (const True) Nothing targetName targetPort action

-- | Same as 'withKubectlPortForward', but allows you to pass in the @kubectl@ binary path.
withKubectlPortForward' :: (
  HasCallStack, KubernetesBasic context m
  )
  => FilePath
  -- | Path to kubeconfig file
  -> FilePath
  -- | Namespace
  -> Text
  -- | Callback to check if the proposed local port is acceptable
  -> (PortNumber -> Bool)
  -> Maybe PortNumber
  -- | Target name (pod, service, etc.)
  -> Text
  -- | Target port number
  -> PortNumber
  -> (KubectlPortForwardContext -> m a)
  -> m a
withKubectlPortForward' :: forall context (m :: * -> *) a.
(HasCallStack, KubernetesBasic context m) =>
FilePath
-> FilePath
-> Text
-> (PortNumber -> Bool)
-> Maybe PortNumber
-> Text
-> PortNumber
-> (KubectlPortForwardContext -> m a)
-> m a
withKubectlPortForward' FilePath
kubectlBinary FilePath
kubeConfigFile Text
namespace PortNumber -> Bool
isAcceptablePort Maybe PortNumber
maybeHostPort Text
targetName PortNumber
targetPort KubectlPortForwardContext -> m a
action = do
  port <- m PortNumber
-> (PortNumber -> m PortNumber) -> Maybe PortNumber -> m PortNumber
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ((PortNumber -> Bool) -> m PortNumber
forall (m :: * -> *).
MonadUnliftIO m =>
(PortNumber -> Bool) -> m PortNumber
findFreePortOrException' PortNumber -> Bool
isAcceptablePort) PortNumber -> m PortNumber
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe PortNumber
maybeHostPort

  let args = [FilePath
"port-forward", Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
targetName, [i|#{port}:#{targetPort}|]
             , FilePath
"--namespace", Text -> FilePath
forall a. ToString a => a -> FilePath
toString Text
namespace
             , FilePath
"--kubeconfig", FilePath
kubeConfigFile]

  debug [i|Running kubectl #{unwords $ fmap toText args}|]

  dir <- getCurrentFolder >>= \case
    Just FilePath
x -> FilePath -> m FilePath
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FilePath
x FilePath -> FilePath -> FilePath
</> FilePath
"port-forwarding-logs-kubectl")
    Maybe FilePath
Nothing -> FilePath -> m FilePath
forall (m :: * -> *) a.
(HasCallStack, MonadIO m) =>
FilePath -> m a
expectationFailure [i|Expected a current folder in withKubectlPortForward'|]

  createDirectoryIfMissing True dir
  let logPath = FilePath
dir FilePath -> FilePath -> FilePath
</> Text -> FilePath
forall a. ToString a => a -> FilePath
toString (HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"/" Text
"_" Text
targetName) FilePath -> FilePath -> FilePath
<.> FilePath
"port-forwarding.log"

  withFile logPath WriteMode $ \Handle
h -> do
    let restarterThread :: m b
restarterThread = m () -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (m () -> m b) -> m () -> m b
forall a b. (a -> b) -> a -> b
$ do
          m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
-> ((Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
    -> m ())
-> ((Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
    -> m ())
-> m ()
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
forall (m :: * -> *).
MonadIO m =>
CreateProcess
-> m (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess ((FilePath -> [FilePath] -> CreateProcess
proc FilePath
kubectlBinary [FilePath]
args) {
                                     std_out = UseHandle h
                                     , std_err = UseHandle h
                                     , create_group = True
                                     }))
                  (\(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
ps) -> ProcessHandle -> Int -> m ()
forall (m :: * -> *).
(MonadIO m, MonadLogger m) =>
ProcessHandle -> Int -> m ()
gracefullyStopProcess ProcessHandle
ps Int
30000000)
                  (\(Maybe Handle
_, Maybe Handle
_, Maybe Handle
_, ProcessHandle
ps) -> do
                      pid <- IO (Maybe Pid) -> m (Maybe Pid)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Pid) -> m (Maybe Pid))
-> IO (Maybe Pid) -> m (Maybe Pid)
forall a b. (a -> b) -> a -> b
$ ProcessHandle -> IO (Maybe Pid)
getPid ProcessHandle
ps
                      info [i|Got pid for kubectl port forward: #{pid}|]

                      code <- waitForProcess ps
                      warn [i|kubectl port-forward #{targetName} #{port}:#{targetPort} exited with code: #{code}. Restarting...|]
                  )

    m (ZonkAny 0) -> (Async (ZonkAny 0) -> m a) -> m a
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync m (ZonkAny 0)
forall {b}. m b
restarterThread ((Async (ZonkAny 0) -> m a) -> m a)
-> (Async (ZonkAny 0) -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \Async (ZonkAny 0)
_ -> do
      let policy :: RetryPolicyM IO
policy = Int -> RetryPolicyM IO
forall (m :: * -> *). Monad m => Int -> RetryPolicyM m
constantDelay Int
100000 RetryPolicyM IO -> RetryPolicyM IO -> RetryPolicyM IO
forall a. Semigroup a => a -> a -> a
<> Int -> RetryPolicy
limitRetries Int
100
      m Bool -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Bool -> m ()) -> m Bool -> m ()
forall a b. (a -> b) -> a -> b
$ IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ RetryPolicyM IO
-> (RetryStatus -> Bool -> IO Bool)
-> (RetryStatus -> IO Bool)
-> IO Bool
forall (m :: * -> *) b.
MonadIO m =>
RetryPolicyM m
-> (RetryStatus -> b -> m Bool) -> (RetryStatus -> m b) -> m b
retrying RetryPolicyM IO
policy (\RetryStatus
_ Bool
ret -> Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
ret) ((RetryStatus -> IO Bool) -> IO Bool)
-> (RetryStatus -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \RetryStatus
_ -> do
        Bool -> Bool
not (Bool -> Bool) -> IO Bool -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SockAddr -> IO Bool
isPortOpen ((Word8, Word8, Word8, Word8) -> PortNumber -> SockAddr
simpleSockAddr (Word8
127, Word8
0, Word8
0, Word8
1) PortNumber
port)

      KubectlPortForwardContext -> m a
action (KubectlPortForwardContext -> m a)
-> KubectlPortForwardContext -> m a
forall a b. (a -> b) -> a -> b
$ KubectlPortForwardContext { kubectlPortForwardPort :: PortNumber
kubectlPortForwardPort = PortNumber
port }