TrafficGenerator.hs

{-# LANGUAGE Arrows, TypeOperators #-}

-- | Module that can be used to generate UDP traffic using a signal 
-- function that outputs packets and displays information on the local
-- console.
module Nettle.FRPControl.TrafficGenerator 
    (
     TrafficGenerator
    , PacketSendCommand
    , ConsoleMessage
    , UDPPort
    , driveTrafficGenerator

    ) where 

import Network.Socket
import Data.Word
import Data.Bits
import Control.Exception
import Nettle.IPv4.IPAddress
import Nettle.FRPControl.AFRP
import Control.Concurrent.MVar
import Control.Concurrent
import Data.Time.Clock (getCurrentTime, diffUTCTime)
import Control.Monad

-- |The type of signal function that can be used to generate UDP traffic.
type TrafficGenerator    = SF () (Event PacketSendCommand) 

-- |A console message is simply a string.
type ConsoleMessage = String

-- |A packet send command is a list of triples including the destination address, 
-- the UDP port, and the payload (as a String).
type PacketSendCommand   = [(IPAddress, UDPPort, String)]

-- |A UDP port is a 16 bit number.
type UDPPort = Word16


-- |Drive (i.e. run) a traffic generator.
driveTrafficGenerator :: SF () (Event PacketSendCommand, Event ConsoleMessage) -> IO ()
driveTrafficGenerator sf = do 

  cntVar           <- newMVar 0
  lastSenseTimeVar <- newEmptyMVar
  inCh             <- newChan 

  let initiator = do forkIO $ genClockSignal inCh
                     t <- getCurrentTime
                     putMVar lastSenseTimeVar t
                     return ()

  let sensor _ = do ma <- readChan inCh
                    t' <- getCurrentTime
                    t  <- swapMVar lastSenseTimeVar t'
                    let delta = fromRational (toRational (diffUTCTime t' t))
                    return (delta, ma)

  let actuator sock _ (e1,e2) = 
          do event (return ()) (mapM_ send) e1
             event (return ()) putStr e2
             return False
          where send (addr,port,msg) = let ipAddr     = toHostAddress addr 
                                           udpPortNum = PortNum (flipBytes port) 
                                           dest       = SockAddrInet udpPortNum ipAddr
                                       in do sendTo sock msg dest
                                             modifyMVar_ cntVar (return . (+1)) 
                                             return ()
  sock   <- socket AF_INET Datagram defaultProtocol
  finally (reactimate initiator sensor (actuator sock) sf)
          (do sClose sock
              cnt <- readMVar cntVar
              putStrLn ("Number of packets sent: " ++ show cnt)
          )

    where 

      genClockSignal ch  = 
          forever (threadDelay clockCycle >> writeChan ch Nothing)
          where clockCycle = 1000

flipBytes :: Word16 -> Word16
flipBytes word = shiftL lbyte 8 + hbyte
    where hbyte = shiftR word 8 .&. (2^8 - 1)
          lbyte = word .&. (2^8 - 1)

toHostAddress :: IPAddress -> HostAddress
toHostAddress addr = shiftL (fromIntegral b4) 24 + 
                     shiftL (fromIntegral b3) 16 + 
                     shiftL (fromIntegral b2) 8 + 
                     fromIntegral b1
    where (b1,b2,b3,b4) = addressToOctets addr


  
{-
-- |This function can be used to run a traffic generator.
driveTrafficGenerator :: TrafficGenerator -> IO ()
driveTrafficGenerator sf = do 
  cntVar <- newMVar 0
  let senseChans = ()
  sock <- socket AF_INET Datagram defaultProtocol
  let send (addr,port,msg) = let ipAddr     = toHostAddress addr 
                                 udpPortNum = PortNum (flipBytes port) 
                                 dest       = SockAddrInet udpPortNum ipAddr
                             in do sendTo sock msg dest
                                   modifyMVar_ cntVar (return . (+1)) 
                                   return ()
  let packetActuator e = case e of 
                           Event cs -> mapM_ send cs
                           NoEvent  -> return ()
  let actuators = (packetActuator, ())
  let sf' = arr (const ()) >>> sf >>> arr (\e -> OutCons (e, (OutNil ())))
  finally (sfDriver inputRep senseChans outputRep actuators sf') (sClose sock >> readMVar cntVar >>= \cnt -> putStrLn ("Number sent: " ++ show cnt))



type InputVector     = TNil
type ControllerInput = SFInput InputVector -- ESumNil () 
inputRep :: Rep InputVector
inputRep = RepNil

type OutputVector      = Event PacketSendCommand ::: TNil
type ControllerOutput  = SFOutput OutputVector -- == OutCons (Event PacketSendCommand) (OutNil ())

outputRep :: Rep OutputVector
outputRep = RepCons RepNil
-}