{-# LANGUAGE FlexibleContexts #-}

module Net.UDP_Client(
    initialize,UDP_API,Interface(..),Packet(..),template,Port(..)
  ) where

import Net.Concurrent
import Control.Monad.State
import Control.Monad.Trans(lift)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.List((\\))
--import System.Random(randomRIO)

import Net.UDP
import qualified Net.IPv4 as IPv4
import qualified Net.Interface as Net
import Net.Utils(doReq)
import Net.Packet(InPacket,OutPacket)
import Net.Wire
import Monad.Util

type UDP_API m =
  Net.TimedInterface m (IPv4.Addr,Packet InPacket) (IPv4.Addr,Packet OutPacket)

data Interface m
   = Interface {
       forall (m :: * -> *). Interface m -> Port -> m (UDP_API m)
listen :: Port -> m (UDP_API m),
       forall (m :: * -> *). Interface m -> m (Port, UDP_API m)
listenAny :: m (Port,UDP_API m),
       forall (m :: * -> *). Interface m -> Port -> m ()
unlisten :: Port -> m ()
   }

data Req m
  = Listen Port (UDP_API m->m ())
  | ListenAny ((Port,UDP_API m)->m ())
  | Unlisten Port
  | FromNetwork (IPv4.Packet (Packet InPacket))

type Clients m = Map Port ((IPv4.Addr,Packet InPacket)->m ())

initialize :: ([Char] -> m ())
-> Addr
-> Interface
     m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
-> m (Interface m)
initialize [Char] -> m ()
putStrLn Addr
myIP Interface m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
iface =
  do c (Req m)
reqChan <- m (c (Req m))
forall a. m (c a)
forall (c :: * -> *) (io :: * -> *) a. ChannelIO c io => io (c a)
newChan
     m () -> m ThreadId
forall (io :: * -> *). ForkIO io => io () -> io ThreadId
fork (m () -> m ThreadId) -> m () -> m ThreadId
forall a b. (a -> b) -> a -> b
$ m () -> m ()
forall {m :: * -> *} {a} {b}. Monad m => m a -> m b
loop (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ c (Req m) -> Req m -> m ()
forall a. c a -> a -> m ()
forall (c :: * -> *) (io :: * -> *) a.
ChannelIO c io =>
c a -> a -> io ()
writeChan c (Req m)
reqChan (Req m -> m ())
-> (Packet (Packet InPacket) -> Req m)
-> Packet (Packet InPacket)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Packet (Packet InPacket) -> Req m
forall (m :: * -> *). Packet (Packet InPacket) -> Req m
FromNetwork (Packet (Packet InPacket) -> m ())
-> m (Packet (Packet InPacket)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Interface m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
-> m (Packet (Packet InPacket))
forall (m :: * -> *) i o. Interface m i o -> m i
Net.rx Interface m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
iface
     m () -> m ThreadId
forall (io :: * -> *). ForkIO io => io () -> io ThreadId
fork (m () -> m ThreadId) -> m () -> m ThreadId
forall a b. (a -> b) -> a -> b
$ ([Char] -> m ())
-> Addr
-> Interface
     m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
-> c (Req m)
-> m ()
forall (r :: * -> *) (m :: * -> *) (c :: * -> *) i a.
(Eq (r ()), RefIO r m, ChannelIO c m, DelayIO m, ForkIO m) =>
([Char] -> m ())
-> Addr
-> Interface m i (Packet (Packet OutPacket))
-> c (Req m)
-> m a
server [Char] -> m ()
debug Addr
myIP Interface m (Packet (Packet InPacket)) (Packet (Packet OutPacket))
iface c (Req m)
reqChan
     let listen :: Port -> m (UDP_API m)
listen = c (Req m) -> ((UDP_API m -> m ()) -> Req m) -> m (UDP_API m)
forall {c :: * -> *} {m :: * -> *} {v :: * -> *} {io :: * -> *} {a}
       {b}.
(ChannelIO c m, MVarIO v m, MVarIO v io) =>
c a -> ((b -> io ()) -> a) -> m b
doReq c (Req m)
reqChan (((UDP_API m -> m ()) -> Req m) -> m (UDP_API m))
-> (Port -> (UDP_API m -> m ()) -> Req m) -> Port -> m (UDP_API m)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Port -> (UDP_API m -> m ()) -> Req m
forall (m :: * -> *). Port -> (UDP_API m -> m ()) -> Req m
Listen
	 listenAny :: m (Port, UDP_API m)
listenAny = c (Req m)
-> (((Port, UDP_API m) -> m ()) -> Req m) -> m (Port, UDP_API m)
forall {c :: * -> *} {m :: * -> *} {v :: * -> *} {io :: * -> *} {a}
       {b}.
(ChannelIO c m, MVarIO v m, MVarIO v io) =>
c a -> ((b -> io ()) -> a) -> m b
doReq c (Req m)
reqChan ((Port, UDP_API m) -> m ()) -> Req m
forall (m :: * -> *). ((Port, UDP_API m) -> m ()) -> Req m
ListenAny
	 unlisten :: Port -> m ()
unlisten = c (Req m) -> Req m -> m ()
forall a. c a -> a -> m ()
forall (c :: * -> *) (io :: * -> *) a.
ChannelIO c io =>
c a -> a -> io ()
writeChan c (Req m)
reqChan (Req m -> m ()) -> (Port -> Req m) -> Port -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Port -> Req m
forall (m :: * -> *). Port -> Req m
Unlisten
     Interface m -> m (Interface m)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Interface m -> m (Interface m)) -> Interface m -> m (Interface m)
forall a b. (a -> b) -> a -> b
$ (Port -> m (UDP_API m))
-> m (Port, UDP_API m) -> (Port -> m ()) -> Interface m
forall (m :: * -> *).
(Port -> m (UDP_API m))
-> m (Port, UDP_API m) -> (Port -> m ()) -> Interface m
Interface Port -> m (UDP_API m)
listen m (Port, UDP_API m)
listenAny Port -> m ()
unlisten
  where
    debug :: [Char] -> m ()
debug = [Char] -> m ()
putStrLn ([Char] -> m ()) -> ([Char] -> [Char]) -> [Char] -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Char]
"UDP: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++)

-- I copied this type signature ghci.
server :: (Eq (r ()), RefIO r m, ChannelIO c m, DelayIO m, ForkIO m)
       => ([Char] -> m ())
       -> IPv4.Addr
       -> Net.Interface m i (IPv4.Packet (Packet OutPacket))
       -> c (Req m)
       -> m a
server :: forall (r :: * -> *) (m :: * -> *) (c :: * -> *) i a.
(Eq (r ()), RefIO r m, ChannelIO c m, DelayIO m, ForkIO m) =>
([Char] -> m ())
-> Addr
-> Interface m i (Packet (Packet OutPacket))
-> c (Req m)
-> m a
server [Char] -> m ()
debug Addr
myIP Interface m i (Packet (Packet OutPacket))
iface c (Req m)
reqChan =
    (StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a
 -> Map Port ((Addr, Packet InPacket) -> m ()) -> m a)
-> Map Port ((Addr, Packet InPacket) -> m ())
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a
-> m a
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a
-> Map Port ((Addr, Packet InPacket) -> m ()) -> m a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT Map Port ((Addr, Packet InPacket) -> m ())
forall {k} {a}. Map k a
init (StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a -> m a)
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a -> m a
forall a b. (a -> b) -> a -> b
$ StateT (Map Port ((Addr, Packet InPacket) -> m ())) m ()
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a
forall {m :: * -> *} {a} {b}. Monad m => m a -> m b
loop (Req m -> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m ()
forall {t :: (* -> *) -> * -> *}.
(MonadTrans t,
 MonadState (Map Port ((Addr, Packet InPacket) -> m ())) (t m)) =>
Req m -> t m ()
handle(Req m -> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m ())
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m (Req m)
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<c (Req m)
-> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m (Req m)
forall a.
c a -> StateT (Map Port ((Addr, Packet InPacket) -> m ())) m a
forall (c :: * -> *) (io :: * -> *) a.
ChannelIO c io =>
c a -> io a
readChan c (Req m)
reqChan)
  where
    init :: Map k a
init = Map k a
forall {k} {a}. Map k a
Map.empty

    handle :: Req m -> t m ()
handle Req m
req =
      case Req m
req of
	Listen Port
port UDP_API m -> m ()
reply -> Port -> (UDP_API m -> m ()) -> t m ()
forall {t :: (* -> *) -> * -> *} {m :: * -> *} {k} {i} {a}.
(MonadTrans t, ChannelIO c m, MonadState (Map k (i -> m ())) (t m),
 Ord k) =>
k -> (TimedInterface m i (Addr, Packet OutPacket) -> m a) -> t m ()
listen Port
port UDP_API m -> m ()
reply
	ListenAny (Port, UDP_API m) -> m ()
reply -> do Port
port <- t m Port
pickPort
			      let reply' :: UDP_API m -> m ()
reply' UDP_API m
iface = (Port, UDP_API m) -> m ()
reply (Port
port,UDP_API m
iface)
			      Port -> (UDP_API m -> m ()) -> t m ()
forall {t :: (* -> *) -> * -> *} {m :: * -> *} {k} {i} {a}.
(MonadTrans t, ChannelIO c m, MonadState (Map k (i -> m ())) (t m),
 Ord k) =>
k -> (TimedInterface m i (Addr, Packet OutPacket) -> m a) -> t m ()
listen Port
port UDP_API m -> m ()
reply'
	Unlisten Port
port -> (Map Port ((Addr, Packet InPacket) -> m ())
 -> Map Port ((Addr, Packet InPacket) -> m ()))
-> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Port
-> Map Port ((Addr, Packet InPacket) -> m ())
-> Map Port ((Addr, Packet InPacket) -> m ())
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Port
port)
	FromNetwork Packet (Packet InPacket)
ipPack ->
	   do let udpPacket :: Packet InPacket
udpPacket = Packet (Packet InPacket) -> Packet InPacket
forall content. Packet content -> content
IPv4.content Packet (Packet InPacket)
ipPack
		  src :: Addr
src = Packet (Packet InPacket) -> Addr
forall content. Packet content -> Addr
IPv4.source Packet (Packet InPacket)
ipPack
		  dst :: Addr
dst = Packet (Packet InPacket) -> Addr
forall content. Packet content -> Addr
IPv4.dest Packet (Packet InPacket)
ipPack
		  sp :: Port
sp = Packet InPacket -> Port
forall content. Packet content -> Port
sourcePort Packet InPacket
udpPacket
		  port :: Port
port = Packet InPacket -> Port
forall content. Packet content -> Port
destPort Packet InPacket
udpPacket
	      Map Port ((Addr, Packet InPacket) -> m ())
clients <- t m (Map Port ((Addr, Packet InPacket) -> m ()))
forall s (m :: * -> *). MonadState s m => m s
get
	      case Port
-> Map Port ((Addr, Packet InPacket) -> m ())
-> Maybe ((Addr, Packet InPacket) -> m ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Port
port Map Port ((Addr, Packet InPacket) -> m ())
clients of
		Just (Addr, Packet InPacket) -> m ()
toClient -> m () -> t m ()
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> m () -> t m ()
forall a b. (a -> b) -> a -> b
$ (Addr, Packet InPacket) -> m ()
toClient (Addr
src,Packet InPacket
udpPacket)
		Maybe ((Addr, Packet InPacket) -> m ())
_ -> m () -> t m ()
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> t m ()) -> m () -> t m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
debug ([Char] -> m ()) -> [Char] -> m ()
forall a b. (a -> b) -> a -> b
$ [Char]
"Dropped packet from "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++(Addr, Port) -> [Char]
forall a. Show a => a -> [Char]
show (Addr
src,Port
sp)
			            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" to "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++(Addr, Port) -> [Char]
forall a. Show a => a -> [Char]
show (Packet (Packet InPacket) -> Addr
forall content. Packet content -> Addr
IPv4.dest Packet (Packet InPacket)
ipPack,Port
port)
      where
        listen :: k -> (TimedInterface m i (Addr, Packet OutPacket) -> m a) -> t m ()
listen k
port TimedInterface m i (Addr, Packet OutPacket) -> m a
reply =
	  do Net.TimedInterface{rxT :: forall (m :: * -> *) i o.
TimedInterface m i o -> Maybe Int -> m (Maybe i)
Net.rxT=Maybe Int -> m (Maybe i)
rx,txT :: forall (m :: * -> *) i o. TimedInterface m i o -> o -> m ()
Net.txT=i -> m ()
toClient} <- () -> t m (TimedInterface m i i)
forall {r :: * -> *} {m1 :: * -> *} {c :: * -> *} {m2 :: * -> *}
       {o}.
(Eq (r ()), RefIO r m1, ForkIO m1, DelayIO m1, ChannelIO c m2,
 ChannelIO c m1) =>
() -> m2 (TimedInterface m1 o o)
timedWire()
	     m a -> t m a
forall (m :: * -> *) a. Monad m => m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> t m a) -> m a -> t m a
forall a b. (a -> b) -> a -> b
$ TimedInterface m i (Addr, Packet OutPacket) -> m a
reply ((Maybe Int -> m (Maybe i))
-> ((Addr, Packet OutPacket) -> m ())
-> TimedInterface m i (Addr, Packet OutPacket)
forall (m :: * -> *) i o.
(Maybe Int -> m (Maybe i)) -> (o -> m ()) -> TimedInterface m i o
Net.TimedInterface Maybe Int -> m (Maybe i)
rx (Addr, Packet OutPacket) -> m ()
tx)
	     (Map k (i -> m ()) -> Map k (i -> m ())) -> t m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\Map k (i -> m ())
clients -> k -> (i -> m ()) -> Map k (i -> m ()) -> Map k (i -> m ())
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
port i -> m ()
toClient Map k (i -> m ())
clients)
	pickPort :: t m Port
pickPort = do [Port]
inuse <- (Map Port ((Addr, Packet InPacket) -> m ()) -> [Port])
-> t m [Port]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets Map Port ((Addr, Packet InPacket) -> m ()) -> [Port]
forall k a. Map k a -> [k]
Map.keys
		      Port -> t m Port
forall a. a -> t m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Port -> t m Port) -> Port -> t m Port
forall a b. (a -> b) -> a -> b
$ [Port] -> Port
forall a. HasCallStack => [a] -> a
head ((Word16 -> Port) -> [Word16] -> [Port]
forall a b. (a -> b) -> [a] -> [b]
map Word16 -> Port
Port [Word16
32768..Word16
65535][Port] -> [Port] -> [Port]
forall a. Eq a => [a] -> [a] -> [a]
\\[Port]
inuse)

    tx :: (Addr, Packet OutPacket) -> m ()
tx (Addr
destIP,Packet OutPacket
updPacket) =
      Interface m i (Packet (Packet OutPacket))
-> Packet (Packet OutPacket) -> m ()
forall (m :: * -> *) i o. Interface m i o -> o -> m ()
Net.tx Interface m i (Packet (Packet OutPacket))
iface (Protocol
-> Addr -> Addr -> Packet OutPacket -> Packet (Packet OutPacket)
forall {content}.
Protocol -> Addr -> Addr -> content -> Packet content
IPv4.template Protocol
IPv4.UDP Addr
myIP Addr
destIP Packet OutPacket
updPacket)