Implemented UDP chunking.

2 files changed, 35 insertions, 24 deletions

diff --git a/src/Graylog/Types.hs b/src/Graylog/Types.hs
index 8fad8cc..5c32a8b 100644
--- a/src/Graylog/Types.hs
+++ b/src/Graylog/Types.hs
@@ -8,6 +8,7 @@ module Graylog.Types
    , _graylogAddress
    , _graylogSocket
    , _graylogHostName
+   , _graylogChunkSize
    , ChunkSize
 
    , defaultChunkSize
@@ -38,7 +39,9 @@ defaultChunkSize = 8192
 
 openGraylog
    :: HostName -> ServiceName -> ChunkSize -> IO (Either String Graylog)
-openGraylog h p cksize = getAddrInfo Nothing (Just h) (Just p) >>= \case
+openGraylog h p cksize
+   | cksize < 1024 = return $ Left "ChunkSize must be at least 1024."
+   | otherwise     = getAddrInfo Nothing (Just h) (Just p) >>= \case
    []     -> return $ Left "No address info found."
    infos ->
       case find (\i -> addrSocketType i == Datagram) infos of
diff --git a/src/Graylog/UDP.hs b/src/Graylog/UDP.hs
index c925d7e..cd7fe5a 100644
--- a/src/Graylog/UDP.hs
+++ b/src/Graylog/UDP.hs
@@ -5,37 +5,45 @@ module Graylog.UDP
    ) where
 
 import           Data.Aeson
-{-import           Data.ByteString.Builder-}
-{-import qualified Data.ByteString.Lazy           as LBS-}
-{-import           Data.Word-}
+import           Data.ByteString.Builder
+import qualified Data.ByteString.Lazy           as LBS
+import           Data.Monoid
+import           Data.Word
 import           Network.Socket.ByteString.Lazy
-{-import           System.Random-}
+import           System.Random
 
 import           Graylog.Gelf                   as Export
 import           Graylog.Types                  as Export
 
 sendLog :: Graylog -> GELF -> IO ()
 sendLog glog msg = do
-   _ <- send (_graylogSocket glog) raw
-   print raw
-   return ()
+   cks <- chunky glog raw
+   mapM_ (send $ _graylogSocket glog) cks
    where
       raw = encode msg
 
-{-sendLog :: Graylog -> GELF -> IO ()-}
-{-sendLog glog msg = do-}
-   {-cks <- chunky glog raw-}
-   {-mapM_ (send $ _graylogSocket glog) cks-}
-   {-where-}
-      {-raw = encode msg-}
+chunky :: Graylog -> LBS.ByteString -> IO [LBS.ByteString]
+chunky glog raw = do
+   groupId <- randomIO
+   let groups = divide totalNum raw
+   return $ append groupId groups seqNums
+   where
+      magic           = word8 0x1e <> word8 0x0f
+      seqNums         = [0..] :: [Word8]
+      totalNum        = if excess > 0 then count + 1 else count
+      (count, excess) = quotRem (LBS.length raw) gsize
+      hlen            = 12
+      gsize           = (fromIntegral (_graylogChunkSize glog)) - hlen
+
+      divide   0 dat = [dat]
+      divide num dat = let (pre, post) = LBS.splitAt gsize dat
+                        in pre : divide (num - 1) post
 
-{-chunky :: Graylog -> LBS.ByteString -> IO [LBS.ByteString]-}
-{-chunky glog raw = do-}
-   {-groupId <- randomIO-}
-   {-splitAt gsize-}
-   {-where-}
-      {-magic           = word8 0x1e <> word8 0x0f-}
-      {-seqNum          = undefined-}
-      {-(count, excess) = quotRem (LBS.length raw) gzie-}
-      {-hlen            = 12-}
-      {-gsize           = (fromIntegral (_graylogChunkSize glog)) - hlen-}
+      append _   []     _      = []
+      append _   _      []     = error "the impossible has happened."
+      append gid (g:gs) (s:ss) = (toLazyByteString
+         $ magic
+        <> word64BE gid
+        <> word8 s
+        <> word8 (fromIntegral totalNum)
+        <> lazyByteString g) : append gid gs ss