content: {-# LANGUAGE RecordWildCards #-}
  
  module Hans.IP4.ArpTable (
      -- * Arp Table
      ArpTable(), newArpTable,
  
      -- ** Update
      addEntry,
      markUnreachable,
  
      -- ** Query
      lookupEntry,
      resolveAddr, QueryResult(..),
      WaitStrategy(), blockingStrategy, writeChanStrategy
  
    ) where
  
  import           Hans.Config (Config(..))
  import           Hans.Device.Types (DeviceStats,updateError,statTX)
  import           Hans.Ethernet (Mac)
  import qualified Hans.HashTable as HT
  import           Hans.IP4.Packet (IP4)
  import           Hans.Threads (forkNamed)
  import           Hans.Time (toUSeconds)
  
  import           Control.Concurrent
                       (threadDelay,MVar,newEmptyMVar,ThreadId,tryPutMVar)
  import qualified Control.Concurrent.BoundedChan as BC
  import           Control.Monad (forever)
  import           Data.Time.Clock
                       (UTCTime,NominalDiffTime,addUTCTime,getCurrentTime)
  
  
  -- | The Arp table consists of a map of IP4 to Mac, as well as a heap that
  -- orders the IP4 addresses according to when their corresponding entries should
  -- be expired.
  --
  -- NOTE: There's currently no way to limit the memory use of the arp table, but
  -- that might not be a huge problem as the entries are added based on requests
  -- from higher layers.
  --
  -- INVARIANT: there should never be entries in the map that aren't also in the
  -- heap.
  data ArpTable = ArpTable { atMacs        :: !(HT.HashTable IP4 Entry)
                           , atLifetime    :: !NominalDiffTime
                           , atPurgeThread :: !ThreadId
                           }
  
  data Entry = Waiting [Maybe Mac -> IO ()]
             | Present !UTCTime !Mac
  
  
  newArpTable :: Config -> IO ArpTable
  newArpTable Config { .. } =
    do atMacs        <- HT.newHashTable cfgArpTableSize
       atPurgeThread <- forkNamed "Arp Purge Thread"
                            (purgeArpTable cfgArpTableLifetime atMacs)
       return ArpTable { atLifetime = cfgArpTableLifetime, .. }
  
  
  -- | Loops forever, delaying until the next arp table entry needs to be purged.
  -- If no entries exist, it waits for the maximum entry lifetime before checking
  -- again.
  purgeArpTable :: NominalDiffTime -> HT.HashTable IP4 Entry -> IO ()
  purgeArpTable lifetime table = forever $
    do now <- getCurrentTime
       HT.filterHashTable (update now) table
       threadDelay delay
  
    where
  
    update _   _ Waiting{}          = False
    update now _ (Present expire _) = expire < now
  
    delay = toUSeconds lifetime
  
  
  -- | Lookup an entry in the Arp table.
  lookupEntry :: ArpTable -> IP4 -> IO (Maybe Mac)
  lookupEntry ArpTable { .. } spa =
    do mb <- HT.lookup spa atMacs
       case mb of
         Just (Present _ mac) -> return (Just mac)
         _                    -> return Nothing
  
  
  -- | Insert an entry into the Arp table, and unblock any waiting actions.
  addEntry :: ArpTable -> IP4 -> Mac -> IO ()
  addEntry ArpTable { .. } spa sha =
    do now <- getCurrentTime
       let end = addUTCTime atLifetime now
  
       waiters <- HT.alter (update end) spa atMacs
  
       -- NOTE: as we don't allow user-supplied IO actions to be registered as
       -- callbacks, we don't catch any exceptions here; the only way that an
       -- action can end up in this list is through the constructors for the
       -- WaitStrategy type defined below.
       --
       -- If it turns out that it's significantly impacting the performance of the
       -- fast path, we should consider forking a thread to run the waiters.
       mapM_ ($ Just sha) waiters
  
    where
  
    update expire (Just (Waiting ks)) = (Just (Present expire sha), ks)
    update expire (Just Present{})    = (Just (Present expire sha), [])
    update expire Nothing             = (Just (Present expire sha), [])
  
  
  -- | If nobody has responded to queries for this address, notify any waiters
  -- that there is no mac associated.
  --
  -- NOTE: in the future, it might be nice to keep an entry in the table that
  -- indicates that this host is unreachable.
  markUnreachable :: ArpTable -> IP4 -> IO ()
  markUnreachable ArpTable { .. } addr =
    do waiters <- HT.alter update addr atMacs
  
       -- See the note in addEntry about why we don't need to sandbox the
       -- callbacks.
       mapM_ ($ Nothing) waiters
  
    where
  
    update (Just (Waiting ks))  = (Nothing, ks)
    update ent@(Just Present{}) = (ent, [])
    update Nothing              = (Nothing,[])
  
  
  
  newtype WaitStrategy res = WaitStrategy { getWaiter :: IO (Maybe Mac -> IO (), res) }
  
  
  -- | Have the entry block on a 
  blockingStrategy :: WaitStrategy (MVar (Maybe Mac))
  blockingStrategy  = WaitStrategy $
    do mvar <- newEmptyMVar
  
       let write mb = do _ <- tryPutMVar mvar mb
                         return ()
  
       return (write, mvar)
  
  
  -- | Write the discovered Mac to a bounded channel, passing it through a
  -- filtering function first.
  writeChanStrategy :: Maybe DeviceStats -> (Maybe Mac -> Maybe msg) -> BC.BoundedChan msg
                    -> WaitStrategy ()
  writeChanStrategy mbStats f chan = WaitStrategy (return (handler,()))
    where
    handler mb =
      case f mb of
        Just msg -> do written <- BC.tryWriteChan chan msg
                       case mbStats of
                         Just stats | not written -> updateError statTX stats
                         _                        -> return ()
  
        -- XXX should this update a stat?
        Nothing  -> return ()
  
  
  data QueryResult res = Known !Mac
                       | Unknown !Bool res
  
  
  -- | Returns either the address, or an empty 'MVar' that will eventually contain
  -- the 'Mac', or 'Nothing' if the Arp request fails.
  resolveAddr :: ArpTable -> IP4 -> WaitStrategy res -> IO (QueryResult res)
  resolveAddr arp addr strategy =
    do mb <- lookupEntry arp addr
       case mb of
         Just mac -> return (Known mac)
         Nothing  -> registerWaiter arp addr strategy
  
  
  -- | Register to wait on an entry in the table.
  registerWaiter :: ArpTable -> IP4 -> WaitStrategy res -> IO (QueryResult res)
  registerWaiter ArpTable { .. } addr strategy =
    do waiter <- getWaiter strategy
       HT.alter (update waiter) addr atMacs
    where
    update (w,r) (Just (Waiting ws))        = (Just (Waiting (w:ws)), Unknown False r)
    update _     ent@(Just (Present _ mac)) = (ent, Known mac)
    update (w,r) Nothing                    = (Just (Waiting [w]), Unknown True r)
  