Ports.hs

{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

module Test.Sandwich.Contexts.Kubernetes.Util.Ports where

import Control.Monad.IO.Unlift
import Control.Retry
import Network.Socket
import Relude
import System.Random (randomRIO)
import UnliftIO.Exception


-- | Find an unused port in a given range
findFreePortInRange' :: forall m. (
  MonadUnliftIO m
  ) => RetryPolicy -> (PortNumber, PortNumber) -> [PortNumber] -> m (Maybe PortNumber)
findFreePortInRange' retryPolicy (start, end) blacklist = retrying retryPolicy (\_retryStatus result -> return $ isNothing result) (const findFreePortInRange')
  where
    getAcceptableCandidate :: m PortNumber
    getAcceptableCandidate = do
      candidate <- liftIO (fromInteger <$> randomRIO (fromIntegral start, fromIntegral end))
      if | candidate `elem` blacklist -> getAcceptableCandidate
         | otherwise -> return candidate

    findFreePortInRange' :: m (Maybe PortNumber)
    findFreePortInRange' = do
      candidate <- getAcceptableCandidate
      isPortFree candidate >>= \case
        False -> return Nothing
        True -> return $ Just candidate

isPortFree :: (MonadUnliftIO m) => PortNumber -> m Bool
isPortFree candidate = catch (tryOpenAndClosePort candidate >> return True)
                             (\(_ :: SomeException) -> return False)

tryOpenAndClosePort :: MonadIO m => PortNumber -> m PortNumber
tryOpenAndClosePort port = liftIO $ do
  sock <- socket AF_INET Stream 0
  setSocketOption sock ReuseAddr 1
  let hints = defaultHints { addrSocketType = Stream, addrFamily = AF_INET }
  getAddrInfo (Just hints) (Just "127.0.0.1") (Just $ show port) >>= \case
    ((AddrInfo {addrAddress=addr}):_) -> do
      bind sock addr
      close sock
      return $ fromIntegral port
    [] -> error "Couldn't resolve address 127.0.0.1"

findFreePortInRange :: (
  MonadUnliftIO m
  ) => (PortNumber, PortNumber) -> [PortNumber] -> m (Maybe PortNumber)
findFreePortInRange = findFreePortInRange' (limitRetries 50)

-- | Find an unused port in the ephemeral port range.
-- See https://en.wikipedia.org/wiki/List_of_TCP_and_UDP_port_numbers
-- This works without a timeout since there should always be a port in there somewhere;
-- it might be advisable to wrap in a timeout anyway.
findFreePort :: (MonadUnliftIO m) => m (Maybe PortNumber)
findFreePort = findFreePortInRange (49152, 65535) []

findFreePortOrException :: (MonadUnliftIO m) => m PortNumber
findFreePortOrException = findFreePortOrException' (const True)

findFreePortOrException' :: (MonadUnliftIO m) => (PortNumber -> Bool) -> m PortNumber
findFreePortOrException' isAcceptable = findFreePort >>= \case
  Just port
    | isAcceptable port -> return port
    | otherwise -> findFreePortOrException' isAcceptable
  Nothing -> error "Couldn't find free port"

findFreePortNotIn :: (MonadUnliftIO m) => [PortNumber] -> m (Maybe PortNumber)
findFreePortNotIn = findFreePortInRange (49152, 65535)