module Hans.Socket.Udp where
import Hans.Addr (toAddr,fromAddr,isBroadcastAddr)
import qualified Hans.Buffer.Datagram as DGram
import Hans.Device.Types (Device)
import Hans.Lens (view,to)
import Hans.Network (Network,RouteInfo(..))
import Hans.Socket.Types
import Hans.Types
import Hans.Udp.Output (primSendUdp)
import Control.Exception (throwIO)
import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.IORef (IORef,newIORef,readIORef)
data SockState addr = KnownRoute !(RouteInfo addr) !addr !SockPort !SockPort
| KnownSource !(Maybe Device) !addr !SockPort
data UdpSocket addr = UdpSocket { udpNS :: !NetworkStack
, udpBuffer :: !UdpBuffer
, udpSockState :: !(IORef (SockState addr))
, udpClose :: !(IO ())
}
instance HasNetworkStack (UdpSocket addr) where
networkStack = to udpNS
instance Socket UdpSocket where
sClose UdpSocket { .. } = udpClose
newUdpSocket :: (HasNetworkStack ns, Network addr)
=> ns
-> SocketConfig
-> Maybe Device
-> addr
-> Maybe SockPort
-> IO (UdpSocket addr)
newUdpSocket ns SocketConfig { .. } mbDev src mbSrcPort =
do let udpNS = view networkStack ns
srcPort <- case mbSrcPort of
Nothing -> do mb <- nextUdpPort udpNS (toAddr src)
case mb of
Just port -> return port
Nothing -> throwIO NoPortAvailable
Just p -> return p
udpSockState <- newIORef (KnownSource mbDev src srcPort)
udpBuffer <- DGram.newBuffer scRecvBufferSize
mbClose <- registerRecv udpNS (toAddr src) srcPort udpBuffer
udpClose <- case mbClose of
Just unreg -> return unreg
Nothing -> throwIO AlreadyListening
return $! UdpSocket { .. }
instance DataSocket UdpSocket where
sConnect ns SocketConfig { .. } mbDev src mbPort dst dstPort =
do let udpNS = view networkStack ns
ri <- route udpNS mbDev src dst
srcPort <- case mbPort of
Just p -> return p
Nothing -> do mb <- nextUdpPort udpNS (toAddr (riSource ri))
case mb of
Just port -> return port
Nothing -> throwIO NoPortAvailable
udpSockState <- newIORef (KnownRoute ri dst srcPort dstPort)
udpBuffer <- DGram.newBuffer scRecvBufferSize
mbClose <- registerRecv udpNS (toAddr src) srcPort udpBuffer
udpClose <- case mbClose of
Just unreg -> return unreg
Nothing -> throwIO AlreadyConnected
return UdpSocket { .. }
sCanWrite UdpSocket { .. } =
do path <- readIORef udpSockState
case path of
KnownRoute _ _ _ _ -> return True
KnownSource{} -> return False
sWrite UdpSocket { .. } bytes =
do path <- readIORef udpSockState
case path of
KnownRoute ri dst srcPort dstPort ->
do sent <- primSendUdp udpNS ri dst srcPort dstPort bytes
if sent
then return (fromIntegral (L.length bytes))
else return (1)
KnownSource{} ->
throwIO NoConnection
sCanRead UdpSocket { .. } =
not `fmap` DGram.isEmptyBuffer udpBuffer
sRead UdpSocket { .. } len =
do (_,buf) <- DGram.readChunk udpBuffer
return (L.fromStrict (S.take len buf))
sTryRead UdpSocket { .. } len =
do mb <- DGram.tryReadChunk udpBuffer
case mb of
Just (_,buf) -> return $! Just $! L.fromStrict $! S.take len buf
Nothing -> return Nothing
recvfrom :: Network addr
=> UdpSocket addr -> IO (Device,addr,SockPort,L.ByteString)
recvfrom sock = do
(dev,srcAddr,srcPort,_,msg) <- recvfrom' sock
return (dev,srcAddr,srcPort,msg)
recvfrom' :: Network addr
=> UdpSocket addr -> IO (Device,addr,SockPort,addr,L.ByteString)
recvfrom' UdpSocket { .. } = loop
where
loop =
do ((dev,srcAddr,srcPort,dstAddr,_), chunk) <- DGram.readChunk udpBuffer
case (fromAddr srcAddr, fromAddr dstAddr) of
(Just srcAddr', Just dstAddr') ->
return (dev,srcAddr',srcPort,dstAddr',L.fromStrict chunk)
_ -> loop
sendto :: Network addr
=> UdpSocket addr -> addr -> SockPort -> L.ByteString -> IO ()
sendto UdpSocket { .. } = \ dst dstPort bytes ->
do state <- readIORef udpSockState
case state of
KnownSource mbDev src srcPort ->
do mbRoute <- route' udpNS mbDev src dst
case mbRoute of
Just ri ->
do _ <- primSendUdp udpNS ri dst srcPort dstPort bytes
return ()
Nothing
| Just dev <- mbDev, isBroadcastAddr dst ->
do let ri = RouteInfo { riSource = src
, riNext = dst
, riDev = dev }
_ <- primSendUdp udpNS ri dst srcPort dstPort bytes
return ()
_ ->
throwIO NoRouteToHost
KnownRoute{} ->
throwIO AlreadyConnected