module Test.Network.Ports where
import Control.Exception (IOException, SomeException, catch, onException, try)
import Control.Monad (replicateM, when)
import Data.Bits ((.|.))
import Data.Foldable (for_)
import Data.Functor (($>), (<&>))
import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef)
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import Data.Set qualified as Set
import Foreign.C.Error (throwErrnoIfMinus1)
import Foreign.C.Types (CInt (..))
import Network.Socket (
Family (AF_INET),
PortNumber,
SockAddr (SockAddrInet),
Socket,
SocketType (Stream),
bind,
close,
defaultProtocol,
mkSocket,
setCloseOnExecIfNeeded,
socket,
tupleToHostAddress,
withFdSocket,
)
import System.Directory (createDirectoryIfMissing)
import System.Environment (lookupEnv)
import System.FileLock (FileLock, SharedExclusive (Exclusive), tryLockFile)
import System.FilePath ((</>))
import System.IO (hPutStrLn, readFile', stderr)
import System.IO.Unsafe (unsafePerformIO)
import System.Info (os)
import System.Process (getCurrentPid)
import Text.Read (readMaybe)
portBandStart :: PortNumber
portBandStart :: PortNumber
portBandStart = PortNumber
24000
portBandEnd :: PortNumber
portBandEnd :: PortNumber
portBandEnd = PortNumber
32768
{-# NOINLINE reservedPortsRef #-}
reservedPortsRef :: IORef (Set PortNumber)
reservedPortsRef :: IORef (Set PortNumber)
reservedPortsRef = IO (IORef (Set PortNumber)) -> IORef (Set PortNumber)
forall a. IO a -> a
unsafePerformIO (IO (IORef (Set PortNumber)) -> IORef (Set PortNumber))
-> IO (IORef (Set PortNumber)) -> IORef (Set PortNumber)
forall a b. (a -> b) -> a -> b
$ Set PortNumber -> IO (IORef (Set PortNumber))
forall a. a -> IO (IORef a)
newIORef Set PortNumber
forall a. Monoid a => a
mempty
{-# NOINLINE nextCandidateRef #-}
nextCandidateRef :: IORef PortNumber
nextCandidateRef :: IORef PortNumber
nextCandidateRef = IO (IORef PortNumber) -> IORef PortNumber
forall a. IO a -> a
unsafePerformIO (IO (IORef PortNumber) -> IORef PortNumber)
-> IO (IORef PortNumber) -> IORef PortNumber
forall a b. (a -> b) -> a -> b
$ do
Pid
pid <- IO Pid
getCurrentPid
let bandSize :: Integer
bandSize = PortNumber -> Integer
forall a. Integral a => a -> Integer
toInteger PortNumber
portBandEnd Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- PortNumber -> Integer
forall a. Integral a => a -> Integer
toInteger PortNumber
portBandStart
PortNumber -> IO (IORef PortNumber)
forall a. a -> IO (IORef a)
newIORef (PortNumber -> IO (IORef PortNumber))
-> PortNumber -> IO (IORef PortNumber)
forall a b. (a -> b) -> a -> b
$ PortNumber
portBandStart PortNumber -> PortNumber -> PortNumber
forall a. Num a => a -> a -> a
+ Integer -> PortNumber
forall a. Num a => Integer -> a
fromInteger ((Pid -> Integer
forall a. Integral a => a -> Integer
toInteger Pid
pid Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
2657) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
bandSize)
{-# NOINLINE ephemeralRangeCheck #-}
ephemeralRangeCheck :: ()
ephemeralRangeCheck :: ()
ephemeralRangeCheck = IO () -> ()
forall a. IO a -> a
unsafePerformIO (IO () -> ()) -> IO () -> ()
forall a b. (a -> b) -> a -> b
$ do
Maybe Int
lower <-
forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (String -> IO String
readFile' String
"/proc/sys/net/ipv4/ip_local_port_range") IO (Either SomeException String)
-> (Either SomeException String -> Maybe Int) -> IO (Maybe Int)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
Right String
contents | (String
lo : [String]
_) <- String -> [String]
words String
contents -> String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe String
lo
Either SomeException String
_ -> Maybe Int
forall a. Maybe a
Nothing
Maybe Int -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Maybe Int
lower ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
lo :: Int) ->
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lo Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral PortNumber
portBandEnd) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
Handle -> String -> IO ()
hPutStrLn Handle
stderr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
String
"Test.Network.Ports: ip_local_port_range starts at "
String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
lo
String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", overlapping the test port band; OS-assigned ports may collide with allocations"
allocatePort :: Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort :: Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort Maybe (PortNumber -> PortNumber)
mDerive =
()
ephemeralRangeCheck () -> IO PortNumber -> IO PortNumber
forall a b. a -> b -> b
`seq` Int -> IO PortNumber
go (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PortNumber
portBandEnd PortNumber -> PortNumber -> PortNumber
forall a. Num a => a -> a -> a
- PortNumber
portBandStart))
where
go :: Int -> IO PortNumber
go :: Int -> IO PortNumber
go Int
0 = String -> IO PortNumber
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Test.Network.Ports: exhausted the private port band"
go Int
n = do
PortNumber
p <- IORef PortNumber
-> (PortNumber -> (PortNumber, PortNumber)) -> IO PortNumber
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef PortNumber
nextCandidateRef ((PortNumber -> (PortNumber, PortNumber)) -> IO PortNumber)
-> (PortNumber -> (PortNumber, PortNumber)) -> IO PortNumber
forall a b. (a -> b) -> a -> b
$ \PortNumber
c ->
(if PortNumber -> PortNumber
forall a. Enum a => a -> a
succ PortNumber
c PortNumber -> PortNumber -> Bool
forall a. Ord a => a -> a -> Bool
>= PortNumber
portBandEnd then PortNumber
portBandStart else PortNumber -> PortNumber
forall a. Enum a => a -> a
succ PortNumber
c, PortNumber
c)
let ps :: [PortNumber]
ps = PortNumber
p PortNumber -> [PortNumber] -> [PortNumber]
forall a. a -> [a] -> [a]
: [PortNumber]
-> ((PortNumber -> PortNumber) -> [PortNumber])
-> Maybe (PortNumber -> PortNumber)
-> [PortNumber]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (\PortNumber -> PortNumber
f -> [PortNumber -> PortNumber
f PortNumber
p]) Maybe (PortNumber -> PortNumber)
mDerive
Bool
fresh <- IORef (Set PortNumber)
-> (Set PortNumber -> (Set PortNumber, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Set PortNumber)
reservedPortsRef ((Set PortNumber -> (Set PortNumber, Bool)) -> IO Bool)
-> (Set PortNumber -> (Set PortNumber, Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Set PortNumber
rs ->
if (PortNumber -> Bool) -> [PortNumber] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (PortNumber -> Set PortNumber -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set PortNumber
rs) [PortNumber]
ps
then (Set PortNumber
rs, Bool
False)
else ((PortNumber -> Set PortNumber -> Set PortNumber)
-> Set PortNumber -> [PortNumber] -> Set PortNumber
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr PortNumber -> Set PortNumber -> Set PortNumber
forall a. Ord a => a -> Set a -> Set a
Set.insert Set PortNumber
rs [PortNumber]
ps, Bool
True)
if Bool -> Bool
not Bool
fresh
then Int -> IO PortNumber
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
else do
Bool
locked <- [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> IO [Bool] -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (PortNumber -> IO Bool) -> [PortNumber] -> IO [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM PortNumber -> IO Bool
tryLockPort [PortNumber]
ps
if Bool -> Bool
not Bool
locked
then Int -> IO PortNumber
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
else
forall e a. Exception e => IO a -> IO (Either e a)
try @IOException ([PortNumber] -> IO ()
verifyBindable [PortNumber]
ps) IO (Either IOException ())
-> (Either IOException () -> IO PortNumber) -> IO PortNumber
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Right () -> PortNumber -> IO PortNumber
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PortNumber
p
Left IOException
_ -> Int -> IO PortNumber
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
verifyBindable :: [PortNumber] -> IO ()
verifyBindable = [Socket] -> [PortNumber] -> IO ()
bindAll []
bindAll :: [Socket] -> [PortNumber] -> IO ()
bindAll [Socket]
held [] = (Socket -> IO ()) -> [Socket] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Socket -> IO ()
close [Socket]
held
bindAll [Socket]
held (PortNumber
p : [PortNumber]
rest) = do
Socket
s <- PortNumber -> IO Socket
bindSpecificLoopback PortNumber
p IO Socket -> IO () -> IO Socket
forall a b. IO a -> IO b -> IO a
`onException` (Socket -> IO ()) -> [Socket] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Socket -> IO ()
close [Socket]
held
[Socket] -> [PortNumber] -> IO ()
bindAll (Socket
s Socket -> [Socket] -> [Socket]
forall a. a -> [a] -> [a]
: [Socket]
held) [PortNumber]
rest
getRandomPort :: IO PortNumber
getRandomPort :: IO PortNumber
getRandomPort = Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort Maybe (PortNumber -> PortNumber)
forall a. Maybe a
Nothing
withFreePort :: (PortNumber -> IO a) -> IO a
withFreePort :: forall a. (PortNumber -> IO a) -> IO a
withFreePort PortNumber -> IO a
action = IO PortNumber
getRandomPort IO PortNumber -> (PortNumber -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PortNumber -> IO a
action
withFreePortAndDerived :: (PortNumber -> PortNumber) -> (PortNumber -> IO a) -> IO a
withFreePortAndDerived :: forall a.
(PortNumber -> PortNumber) -> (PortNumber -> IO a) -> IO a
withFreePortAndDerived PortNumber -> PortNumber
derive PortNumber -> IO a
action = Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort ((PortNumber -> PortNumber) -> Maybe (PortNumber -> PortNumber)
forall a. a -> Maybe a
Just PortNumber -> PortNumber
derive) IO PortNumber -> (PortNumber -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PortNumber -> IO a
action
randomUnusedTCPPorts :: Int -> IO [Int]
randomUnusedTCPPorts :: Int -> IO [Int]
randomUnusedTCPPorts Int
count =
Int -> IO Int -> IO [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
count (PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PortNumber -> Int) -> IO PortNumber -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort Maybe (PortNumber -> PortNumber)
forall a. Maybe a
Nothing)
randomUnusedTCPPortsWithDerived ::
(PortNumber -> PortNumber) ->
Int ->
IO [Int]
randomUnusedTCPPortsWithDerived :: (PortNumber -> PortNumber) -> Int -> IO [Int]
randomUnusedTCPPortsWithDerived PortNumber -> PortNumber
derive Int
count =
Int -> IO Int -> IO [Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
count (PortNumber -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PortNumber -> Int) -> IO PortNumber -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (PortNumber -> PortNumber) -> IO PortNumber
allocatePort ((PortNumber -> PortNumber) -> Maybe (PortNumber -> PortNumber)
forall a. a -> Maybe a
Just PortNumber -> PortNumber
derive))
{-# NOINLINE portLocksRef #-}
portLocksRef :: IORef [FileLock]
portLocksRef :: IORef [FileLock]
portLocksRef = IO (IORef [FileLock]) -> IORef [FileLock]
forall a. IO a -> a
unsafePerformIO (IO (IORef [FileLock]) -> IORef [FileLock])
-> IO (IORef [FileLock]) -> IORef [FileLock]
forall a b. (a -> b) -> a -> b
$ [FileLock] -> IO (IORef [FileLock])
forall a. a -> IO (IORef a)
newIORef []
tryLockPort :: PortNumber -> IO Bool
tryLockPort :: PortNumber -> IO Bool
tryLockPort PortNumber
p =
IO Bool
go IO Bool -> (IOException -> IO Bool) -> IO Bool
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
_ :: IOException) -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
where
go :: IO Bool
go = do
String
dir <- IO String
portLockDir
Bool -> String -> IO ()
createDirectoryIfMissing Bool
True String
dir
String -> SharedExclusive -> IO (Maybe FileLock)
tryLockFile (String
dir String -> String -> String
</> PortNumber -> String
forall a. Show a => a -> String
show PortNumber
p) SharedExclusive
Exclusive IO (Maybe FileLock) -> (Maybe FileLock -> IO Bool) -> IO Bool
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Just FileLock
lock -> IORef [FileLock] -> ([FileLock] -> [FileLock]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [FileLock]
portLocksRef (FileLock
lock :) IO () -> Bool -> IO Bool
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Bool
True
Maybe FileLock
Nothing -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
portLockDir :: IO FilePath
portLockDir :: IO String
portLockDir = do
String
user <- String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe String
"unknown" (Maybe String -> String) -> IO (Maybe String) -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
"USER"
String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> IO String) -> String -> IO String
forall a b. (a -> b) -> a -> b
$ String
"/tmp" String -> String -> String
</> String
"hydra-test-port-locks-" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
user
bindSpecificLoopback :: PortNumber -> IO Socket
bindSpecificLoopback :: PortNumber -> IO Socket
bindSpecificLoopback PortNumber
portNumber = do
Socket
s <- IO Socket
openStreamSocketCloexec
Socket -> SockAddr -> IO ()
bind Socket
s (PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
portNumber ((Word8, Word8, Word8, Word8) -> HostAddress
tupleToHostAddress (Word8
127, Word8
0, Word8
0, Word8
1)))
Socket -> IO Socket
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Socket
s
foreign import ccall unsafe "socket"
c_socket :: CInt -> CInt -> CInt -> IO CInt
afInet, sockStream, sockNonBlock, sockCloexec :: CInt
afInet :: CInt
afInet = CInt
2
sockStream :: CInt
sockStream = CInt
1
sockNonBlock :: CInt
sockNonBlock = CInt
0x800
sockCloexec :: CInt
sockCloexec = CInt
0x80000
openStreamSocketCloexec :: IO Socket
openStreamSocketCloexec :: IO Socket
openStreamSocketCloexec
| String
os String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"linux" = do
CInt
fd <- String -> IO CInt -> IO CInt
forall a. (Eq a, Num a) => String -> IO a -> IO a
throwErrnoIfMinus1 String
"socket" (CInt -> CInt -> CInt -> IO CInt
c_socket CInt
afInet (CInt
sockStream CInt -> CInt -> CInt
forall a. Bits a => a -> a -> a
.|. CInt
sockNonBlock CInt -> CInt -> CInt
forall a. Bits a => a -> a -> a
.|. CInt
sockCloexec) CInt
0)
CInt -> IO Socket
mkSocket CInt
fd
| Bool
otherwise = do
Socket
s <- Family -> SocketType -> CInt -> IO Socket
socket Family
AF_INET SocketType
Stream CInt
defaultProtocol
Socket -> (CInt -> IO ()) -> IO ()
forall r. Socket -> (CInt -> IO r) -> IO r
withFdSocket Socket
s CInt -> IO ()
setCloseOnExecIfNeeded
Socket -> IO Socket
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Socket
s