-- |
-- Copyright: © 2018-2020 IOHK
-- License: Apache-2.0
--
-- Provides functions for checking if TCP ports can be connected to, or are
-- available to listen on.
--
-- These can be used for:
--  - Waiting until a server in another process has started.
--  - Start servers for testing when there may be multiple
--    test suites running in parallel.
--

module Cardano.Wallet.Network.Ports
    (
    -- * Allocation
      PortNumber
    , getRandomPort

    -- * Status
    , isPortOpen
    , simpleSockAddr

    -- * Helpers
    , portFromURL
    , randomUnusedTCPPorts
    ) where

import Prelude

import Control.Monad
    ( filterM )
import Control.Monad.IO.Class
    ( liftIO )
import Data.List
    ( isInfixOf, sort )
import Data.Maybe
    ( fromMaybe )
import Data.Streaming.Network
    ( bindRandomPortTCP )
import Data.Word
    ( Word8 )
import Foreign.C.Error
    ( Errno (..), eCONNREFUSED )
import GHC.IO.Exception
    ( IOException (..) )
import Network.Socket
    ( Family (AF_INET)
    , PortNumber
    , SockAddr (..)
    , SocketType (Stream)
    , close'
    , connect
    , socket
    , tupleToHostAddress
    )
import Network.URI
    ( URI (..), URIAuth (..) )
import Safe
    ( readMay )
import System.Random.Shuffle
    ( shuffleM )
import UnliftIO.Exception
    ( bracket, throwIO, try )

-- | Find a TCPv4 port which is likely to be free for listening on
-- @localhost@. This binds a socket, receives an OS-assigned port, then closes
-- the socket.
--
-- Note that this is vulnerable to race conditions if another process binds the
-- port returned by 'getRandomPort' before this process does.
--
-- Do not use this unless you have no other option.
getRandomPort :: IO PortNumber
getRandomPort :: IO PortNumber
getRandomPort = do
    let hostPreference :: HostPreference
hostPreference = HostPreference
"127.0.0.1"
    (Int
port, Socket
sock) <- HostPreference -> IO (Int, Socket)
bindRandomPortTCP HostPreference
hostPreference
    IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close' Socket
sock
    PortNumber -> IO PortNumber
forall (m :: * -> *) a. Monad m => a -> m a
return (PortNumber -> IO PortNumber) -> PortNumber -> IO PortNumber
forall a b. (a -> b) -> a -> b
$ Int -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
port

-- | Checks whether @connect()@ to a given TCPv4 `SockAddr` succeeds or
-- returns `eCONNREFUSED`.
--
-- Rethrows connection exceptions in all other cases (e.g. when the host
-- is unroutable).
--
-- Code courtesy of nh2: https://stackoverflow.com/a/57022572
isPortOpen :: SockAddr -> IO Bool
isPortOpen :: SockAddr -> IO Bool
isPortOpen SockAddr
sockAddr = do
  IO Socket -> (Socket -> IO ()) -> (Socket -> IO Bool) -> IO Bool
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_INET SocketType
Stream ProtocolNumber
6 {- TCP -}) Socket -> IO ()
close' ((Socket -> IO Bool) -> IO Bool) -> (Socket -> IO Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> do
    Either IOException ()
res <- IO () -> IO (Either IOException ())
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> m (Either e a)
try (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
connect Socket
sock SockAddr
sockAddr
    case Either IOException ()
res of
      Right () -> Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
      Left IOException
e
        | (ProtocolNumber -> Errno
Errno (ProtocolNumber -> Errno) -> Maybe ProtocolNumber -> Maybe Errno
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IOException -> Maybe ProtocolNumber
ioe_errno IOException
e) Maybe Errno -> Maybe Errno -> Bool
forall a. Eq a => a -> a -> Bool
== Errno -> Maybe Errno
forall a. a -> Maybe a
Just Errno
eCONNREFUSED -> Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
        | [Char]
"WSAECONNREFUSED" [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isInfixOf` IOException -> [Char]
forall a. Show a => a -> [Char]
show IOException
e -> Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
        | Bool
otherwise -> IOException -> IO Bool
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO IOException
e

-- | Creates a `SockAttr` from host IP and port number.
--
-- Example:
-- > simpleSockAddr (127,0,0,1) 8000
simpleSockAddr :: (Word8, Word8, Word8, Word8) -> PortNumber -> SockAddr
simpleSockAddr :: (Word8, Word8, Word8, Word8) -> PortNumber -> SockAddr
simpleSockAddr (Word8, Word8, Word8, Word8)
addr PortNumber
port = PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
port ((Word8, Word8, Word8, Word8) -> HostAddress
tupleToHostAddress (Word8, Word8, Word8, Word8)
addr)

-- | Get the port from a URI, which is assumed to be a HTTP or HTTPS URL.
portFromURL :: URI -> PortNumber
portFromURL :: URI -> PortNumber
portFromURL URI
uri = PortNumber -> Maybe PortNumber -> PortNumber
forall a. a -> Maybe a -> a
fromMaybe PortNumber
fallback
    (URI -> Maybe URIAuth
uriAuthority URI
uri Maybe URIAuth -> (URIAuth -> Maybe PortNumber) -> Maybe PortNumber
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Char] -> Maybe PortNumber
forall a. Read a => [Char] -> Maybe a
readMay ([Char] -> Maybe PortNumber)
-> (URIAuth -> [Char]) -> URIAuth -> Maybe PortNumber
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Char -> Bool) -> [Char] -> [Char]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
':')) ([Char] -> [Char]) -> (URIAuth -> [Char]) -> URIAuth -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. URIAuth -> [Char]
uriPort)
  where
    fallback :: PortNumber
fallback = if URI -> [Char]
uriScheme URI
uri [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== [Char]
"https:" then PortNumber
443 else PortNumber
80

-- | Get a list of random TCPv4 ports that currently do not have any servers
-- listening on them. It may return less than the requested number of ports.
--
-- Note that this method of allocating ports is subject to race
-- conditions. Production code should use better methods such as passing a
-- listening socket to the child process.
randomUnusedTCPPorts :: Int -> IO [Int]
randomUnusedTCPPorts :: Int -> IO [Int]
randomUnusedTCPPorts Int
count = do
    [Int]
usablePorts <- [Int] -> IO [Int]
forall (m :: * -> *) a. MonadRandom m => [a] -> m [a]
shuffleM [Int
1024..Int
49151]
    [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort ([Int] -> [Int]) -> IO [Int] -> IO [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> IO Bool) -> [Int] -> IO [Int]
forall (m :: * -> *) a.
Applicative m =>
(a -> m Bool) -> [a] -> m [a]
filterM Int -> IO Bool
unused (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
count [Int]
usablePorts)
  where
    unused :: Int -> IO Bool
unused = (Bool -> Bool) -> IO Bool -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not (IO Bool -> IO Bool) -> (Int -> IO Bool) -> Int -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SockAddr -> IO Bool
isPortOpen (SockAddr -> IO Bool) -> (Int -> SockAddr) -> Int -> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8, Word8, Word8, Word8) -> PortNumber -> SockAddr
simpleSockAddr (Word8
127,Word8
0,Word8
0,Word8
1) (PortNumber -> SockAddr) -> (Int -> PortNumber) -> Int -> SockAddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral