{-# LANGUAGE TypeApplications #-}

module Hedgehog.Extras.Stock.IO.Network.Port
  ( randomPort
  , reserveRandomPort
  , portInUse
  ) where

import           Control.Exception
import           Control.Monad (Monad (..), MonadFail (..))
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import           Data.Bool
import           Data.Either
import           Data.Function
import           Network.Socket

-- | Return a random available port on a specified host address
randomPort :: ()
  => MonadIO m
  => MonadFail m
  => HostAddress
  -> m PortNumber
randomPort :: forall (m :: * -> *).
(MonadIO m, MonadFail m) =>
HostAddress -> m PortNumber
randomPort HostAddress
hostAddress = do
  Socket
sock <- IO Socket -> m Socket
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Socket -> m Socket) -> IO Socket -> m Socket
forall a b. (a -> b) -> a -> b
$ Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_INET SocketType
Stream ProtocolNumber
defaultProtocol
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
bind Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
defaultPort HostAddress
hostAddress
  SockAddrInet PortNumber
port HostAddress
_ <- IO SockAddr -> m SockAddr
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO SockAddr -> m SockAddr) -> IO SockAddr -> m SockAddr
forall a b. (a -> b) -> a -> b
$ Socket -> IO SockAddr
getSocketName Socket
sock
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close Socket
sock
  PortNumber -> m PortNumber
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return PortNumber
port

reserveRandomPort :: ()
  => MonadFail m
  => MonadResource m
  => HostAddress
  -> m (ReleaseKey, PortNumber)
reserveRandomPort :: forall (m :: * -> *).
(MonadFail m, MonadResource m) =>
HostAddress -> m (ReleaseKey, PortNumber)
reserveRandomPort HostAddress
hostAddress = do
  Socket
sock <- IO Socket -> m Socket
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Socket -> m Socket) -> IO Socket -> m Socket
forall a b. (a -> b) -> a -> b
$ Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_INET SocketType
Stream ProtocolNumber
defaultProtocol
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
bind Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
defaultPort HostAddress
hostAddress
  SockAddrInet PortNumber
port HostAddress
_ <- IO SockAddr -> m SockAddr
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO SockAddr -> m SockAddr) -> IO SockAddr -> m SockAddr
forall a b. (a -> b) -> a -> b
$ Socket -> IO SockAddr
getSocketName Socket
sock
  ReleaseKey
releaseKey <- IO () -> m ReleaseKey
forall (m :: * -> *). MonadResource m => IO () -> m ReleaseKey
register (IO () -> m ReleaseKey) -> IO () -> m ReleaseKey
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close Socket
sock
  (ReleaseKey, PortNumber) -> m (ReleaseKey, PortNumber)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (ReleaseKey
releaseKey, PortNumber
port)

-- | Check if a port is in use on a specified host address
portInUse :: ()
  => MonadIO m
  => HostAddress
  -> PortNumber
  -> m Bool
portInUse :: forall (m :: * -> *).
MonadIO m =>
HostAddress -> PortNumber -> m Bool
portInUse HostAddress
hostAddress PortNumber
pn = do
  Socket
sock <- IO Socket -> m Socket
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Socket -> m Socket) -> IO Socket -> m Socket
forall a b. (a -> b) -> a -> b
$ Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
AF_INET SocketType
Stream ProtocolNumber
defaultProtocol
  Either SomeException ()
result <- IO (Either SomeException ()) -> m (Either SomeException ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException ()) -> m (Either SomeException ()))
-> IO (Either SomeException ()) -> m (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ Socket -> SockAddr -> IO ()
bind Socket
sock (PortNumber -> HostAddress -> SockAddr
SockAddrInet PortNumber
pn HostAddress
hostAddress)
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close Socket
sock
  Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ (SomeException -> Bool)
-> (() -> Bool) -> Either SomeException () -> Bool
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Bool -> SomeException -> Bool
forall a b. a -> b -> a
const Bool
False) (Bool -> () -> Bool
forall a b. a -> b -> a
const Bool
True) Either SomeException ()
result