diff --git a/asif.cabal b/asif.cabal --- a/asif.cabal +++ b/asif.cabal @@ -3,7 +3,7 @@ -- see: https://github.com/sol/hpack name: asif -version: 3.2.0 +version: 3.3.4 description: Please see the README on Github at category: Services homepage: https://github.com/packetloop/asif#readme @@ -41,6 +41,7 @@ , cpu , either , exceptions + , foldl , generic-lens , hw-bits , hw-ip @@ -48,10 +49,12 @@ , lens , network , old-locale + , profunctors , resourcet , temporary-resourcet , text , thyme + , transformers , vector exposed-modules: Arbor.File.Format.Asif @@ -70,6 +73,7 @@ Arbor.File.Format.Asif.Segment Arbor.File.Format.Asif.Type Arbor.File.Format.Asif.Whatever + Arbor.File.Format.Asif.Write other-modules: Arbor.File.Format.Asif.Format.SegmentValue Arbor.File.Format.Asif.Format.Type @@ -97,6 +101,7 @@ , cpu , either , exceptions + , foldl , generic-lens , hw-bits , hw-ip @@ -104,10 +109,12 @@ , lens , network , old-locale + , profunctors , resourcet , temporary-resourcet , text , thyme + , transformers , vector , asif , directory @@ -145,6 +152,7 @@ , cpu , either , exceptions + , foldl , generic-lens , hw-bits , hw-ip @@ -152,10 +160,12 @@ , lens , network , old-locale + , profunctors , resourcet , temporary-resourcet , text , thyme + , transformers , vector , asif , hspec @@ -166,6 +176,7 @@ Arbor.File.Format.Asif.Data.IpSpec Arbor.File.Format.Asif.Format.SegmentValueSpec Arbor.File.Format.Asif.Format.TextSpec + Arbor.File.Format.Asif.WriteSpec Arbor.File.Format.AsifSpec Arbor.TestUtils Gen.Feed diff --git a/src/Arbor/File/Format/Asif/Write.hs b/src/Arbor/File/Format/Asif/Write.hs new file mode 100644 --- /dev/null +++ b/src/Arbor/File/Format/Asif/Write.hs @@ -0,0 +1,230 @@ +{-# LANGUAGE MultiParamTypeClasses #-} + + +-- | +-- Module: Arbor.File.Format.Asif.Write +-- +-- Functions to make it easier to write ASIF files without needing to deal with +-- raw Handles and ByteStrings too much. + +module Arbor.File.Format.Asif.Write + ( + -- * Encode an entire ASIF bytestring + -- $usage + writeAsif + , buildAsifBytestring + + -- * Folds for Segments + -- $segments + , lazyByteStringSegment + , nullTerminatedStringSegment + , textSegment + , asciiSegment + , word8Segment + , word16Segment + , word32Segment + , word64Segment + , int8Segment + , int16Segment + , int32Segment + , int64Segment + , ipv4Segment + , ipv6Segment + , utcTimeMicrosSegment + + -- * Utility functions + -- $helper + , genericInitial + , genericStep + , genericExtract + , genericFold + ) +where + +import Arbor.File.Format.Asif.ByteString.Builder +import Arbor.File.Format.Asif.Data.Ip (ipv4ToWord32, ipv6ToWord32x4) +import Arbor.File.Format.Asif.Type +import Arbor.File.Format.Asif.Whatever (Whatever (..)) +import Conduit +import Control.Foldl +import Control.Lens +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Trans.Resource (MonadResource) +import Data.Int +import Data.Profunctor (lmap) +import Data.Semigroup ((<>)) +import Data.Word +import System.IO (Handle, SeekMode (AbsoluteSeek), hFlush, hSeek) +import System.IO.Temp (openTempFile) + +import qualified Arbor.File.Format.Asif.Format as F +import qualified Data.ByteString.Builder as BB +import qualified Data.ByteString.Lazy as LBS +import qualified Data.IP as IP +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.Lazy as TL +import qualified Data.Text.Lazy.Encoding as TE + +import qualified Data.Thyme.Clock.POSIX as TY +import qualified Data.Thyme.Time.Core as TY + +-- $usage +-- +-- Facilities for writing an entire ASIF file, either to an actual file or to a +-- ByteString. +-- +-- ** Usage: +-- Use the various *segment functions to produce a FoldM. These are designed to +-- allow you to take some large type (e.g. a tuple or a product type) and pull +-- out the constituent pieces to encode into 'Segment's. +-- +-- 'FoldM's are composable using '<>'. Once you have a single, provide it to +-- 'writeAsif' or 'buildAsifBytestring' along with an appropriate foldable. +-- Both these functions will stream from the input, assuming the foldable is +-- something that can be streamed from. + +-- | Write an ASIF file to the supplied handle. +-- Streams the input foldable if possible. +writeAsif :: (Foldable f, MonadResource m) + => Handle + -> String + -> Maybe TY.POSIXTime + -> FoldM m a [Segment Handle] + -> f a + -> m () +writeAsif hOutput asifType mTimestamp fld foldable = do + segments <- foldM fld foldable + contents <- segmentsC asifType mTimestamp segments + runConduit $ contents .| sinkHandle hOutput + liftIO $ hFlush hOutput + +-- | Builds a lazy ASIF bytestring. +-- Streams the input foldable if possible. +buildAsifBytestring :: (Foldable f, MonadResource m) + => String + -> Maybe TY.POSIXTime + -> FoldM m a [Segment Handle] + -> f a + -> m LBS.ByteString +buildAsifBytestring asifType mTimestamp fld foldable = do + (_, _, h) <- openTempFile Nothing "asif" + writeAsif h asifType mTimestamp fld foldable + liftIO $ hSeek h AbsoluteSeek 0 + liftIO $ LBS.hGetContents h + +----- + +-- $segments +-- +-- Use these to build 'FoldM's for the types you want to encode. +-- Compose them together using '<>'. + +-- | Builds a segment from lazy bytestrings. +-- This can in priciple cover any bytestring-y format, +-- including StringZ, Text, Binary, Bitmap, and Bitstring, as well as unknown encodings. +-- Correctly encoding the value is the responsibility of the caller. +lazyByteStringSegment :: MonadResource m => Whatever F.Format -> (a -> LBS.ByteString) -> T.Text -> FoldM m a [Segment Handle] +lazyByteStringSegment = genericFold BB.lazyByteString + +-- | Builds a segment of null-termianted strings. +-- Note that the input itself does *not* need to be null-terminated. +-- The null-termination is added by this function. +nullTerminatedStringSegment :: MonadResource m => (a -> T.Text) -> T.Text -> FoldM m a [Segment Handle] +nullTerminatedStringSegment f t = FoldM step initial extract + where + initial = genericInitial t + step h b = do + liftIO $ BB.hPutBuilder h $ BB.byteString (T.encodeUtf8 . f $ b) <> BB.word8 0 + pure h + extract = genericExtract t (Known F.StringZ) + +-- | Builds a segment of 'Text's. +textSegment :: MonadResource m => (a -> T.Text) -> T.Text -> FoldM m a [Segment Handle] +textSegment f = genericFold TE.encodeUtf8Builder (Known F.Text) (TL.fromStrict . f) + +-- | Builds a segment of 'Char's. +asciiSegment :: MonadResource m => (a -> Char) -> T.Text -> FoldM m a [Segment Handle] +asciiSegment = genericFold BB.char8 (Known F.Char) + +----- + +-- | Builds a segment of 'Word8's. +word8Segment :: MonadResource m => (a -> Word8) -> T.Text -> FoldM m a [Segment Handle] +word8Segment = genericFold BB.word8 (Known F.Word8) + +-- | Builds a segment of 'Word16's. +word16Segment :: MonadResource m => (a -> Word16) -> T.Text -> FoldM m a [Segment Handle] +word16Segment = genericFold BB.word16LE (Known F.Word16LE) + +-- | Builds a segment of 'Word32's. +word32Segment :: MonadResource m => (a -> Word32) -> T.Text -> FoldM m a [Segment Handle] +word32Segment = genericFold BB.word32LE (Known F.Word32LE) + +-- | Builds a segment of 'Word64's. +word64Segment :: MonadResource m => (a -> Word64) -> T.Text -> FoldM m a [Segment Handle] +word64Segment = genericFold BB.word64LE (Known F.Word64LE) + +----- + +-- | Builds a segment of 'Int8's. +int8Segment :: MonadResource m => (a -> Int8) -> T.Text -> FoldM m a [Segment Handle] +int8Segment = genericFold BB.int8 (Known F.Int8) + +-- | Builds a segment of 'Int16's. +int16Segment :: MonadResource m => (a -> Int16) -> T.Text -> FoldM m a [Segment Handle] +int16Segment = genericFold BB.int16LE (Known F.Int16LE) + +-- | Builds a segment of 'Int32's. +int32Segment :: MonadResource m => (a -> Int32) -> T.Text -> FoldM m a [Segment Handle] +int32Segment = genericFold BB.int32LE (Known F.Int32LE) + +-- | Builds a segment of 'Int64's. +int64Segment :: MonadResource m => (a -> Int64) -> T.Text -> FoldM m a [Segment Handle] +int64Segment = genericFold BB.int64LE (Known F.Int64LE) + +----- + +-- | Builds a segment of 'IPv4's. +ipv4Segment :: MonadResource m => (a -> IP.IPv4) -> T.Text -> FoldM m a [Segment Handle] +ipv4Segment f = genericFold BB.word32LE (Known F.Ipv4) (ipv4ToWord32 . f) + +-- | Builds a segment of 'IPv6's. +ipv6Segment :: MonadResource m => (a -> IP.IPv6) -> T.Text -> FoldM m a [Segment Handle] +ipv6Segment f = genericFold encoding (Known F.Ipv6) extract + where + -- I do not know why this is Big-Endian, when everything else is Little-Endian. + encoding = Prelude.foldMap BB.word32BE + extract = tupleToList . ipv6ToWord32x4 . f + tupleToList (w1,w2,w3,w4) = [w1,w2,w3,w4] + +----- + +-- | Builds a segment of 'UTCTime's, accurate to microseconds. +utcTimeMicrosSegment :: MonadResource m => (a -> TY.UTCTime) -> T.Text -> FoldM m a [Segment Handle] +utcTimeMicrosSegment f = genericFold BB.int64LE (Known F.TimeMicros64LE) (fromTime . f) + where + fromTime :: TY.UTCTime -> Int64 + fromTime = view (TY.posixTime . TY.microseconds) + +------------------------------------------------------------------------------- + +-- $helper +-- +-- Helper functions for creating 'FoldM's from scratch. + +genericInitial :: MonadResource m => T.Text -> m Handle +genericInitial name = do + (_, _, h) <- openTempFile Nothing (T.unpack name) + pure h + +genericStep :: MonadResource m => (a -> BB.Builder) -> Handle -> a -> m Handle +genericStep enc h b = do + liftIO $ BB.hPutBuilder h $ enc b + pure h + +genericExtract :: MonadResource m => T.Text -> Whatever F.Format -> Handle -> m [Segment Handle] +genericExtract filen typ h = pure [segment h $ metaFilename filen <> metaFormat typ] + +genericFold :: MonadResource m => (a -> BB.Builder) -> Whatever F.Format -> (b -> a) -> T.Text -> FoldM m b [Segment Handle] +genericFold enc fmt f t = lmap f $ FoldM (genericStep enc) (genericInitial t) (genericExtract t fmt) diff --git a/test/Arbor/File/Format/Asif/WriteSpec.hs b/test/Arbor/File/Format/Asif/WriteSpec.hs new file mode 100644 --- /dev/null +++ b/test/Arbor/File/Format/Asif/WriteSpec.hs @@ -0,0 +1,225 @@ +module Arbor.File.Format.Asif.WriteSpec where + +import Arbor.File.Format.Asif.Data.Ip +import Arbor.File.Format.Asif.Format +import Arbor.File.Format.Asif.Segment +import Arbor.File.Format.Asif.Write +import Control.Monad.Trans.Class +import Control.Monad.Trans.Resource +import Data.Int +import Data.Semigroup ((<>)) +import Data.Word + +import HaskellWorks.Hspec.Hedgehog +import Hedgehog +import Test.Hspec + +import qualified Data.Attoparsec.ByteString as AP +import qualified Data.IP as IP +import qualified Data.Text.Lazy as T +import qualified Data.Text.Lazy.Encoding as T +import qualified Data.Thyme.Clock.POSIX as TY +import qualified Data.Thyme.Time.Core as TY +import qualified Hedgehog.Gen as G +import qualified Hedgehog.Range as R + + +{-# ANN module ("HLint: ignore Redundant do" :: String) #-} + +spec :: Spec +spec = describe "Arbor.File.Format.Asif.Write" $ do + it "should write several segments, then read them back in" $ require $ property $ do + tplList <- forAll $ G.list (R.linear 0 50) genTriple + + let f1 = int64Segment (\(a,_,_) -> a) "first" + let f2 = word16Segment (\(_,a,_) -> a) "second" + let f3 = nullTerminatedStringSegment (\(_,_,a) -> T.toStrict a) "third" + let fld = f1 <> f2 <> f3 + + lbs <- buildAsifBytestring "wxyz" Nothing fld tplList + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + + [names, times, types, l1,l2,l3] <- forAll $ pure (segmentValues <$> segments) + + ((\(a,_,_) -> SInt64 a) <$> tplList) === l1 + ((\(_,a,_) -> SWord16 a) <$> tplList) === l2 + ((\(_,_,a) -> SString . T.encodeUtf8 $ a) <$> tplList) === l3 + +----- + + it "should write out and read back in a word8 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.word8 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (word8Segment id "word8") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SWord8 <$> lst) === seg + + it "should write out and read back in a word16 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.word16 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (word16Segment id "word16") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SWord16 <$> lst) === seg + + it "should write out and read back in a word32 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.word32 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (word32Segment id "word32") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SWord32 <$> lst) === seg + + it "should write out and read back in a word64 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.word64 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (word64Segment id "word64") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SWord64 <$> lst) === seg + +----- + + it "should write out and read back in a int8 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.int8 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (int8Segment id "int8") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SInt8 <$> lst) === seg + + it "should write out and read back in a int16 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.int16 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (int16Segment id "int16") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SInt16 <$> lst) === seg + + it "should write out and read back in a int32 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.int32 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (int32Segment id "int32") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SInt32 <$> lst) === seg + + it "should write out and read back in a int64 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (G.int64 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (int64Segment id "int64") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SInt64 <$> lst) === seg + +----- + + it "should write out and read back in a ipv4 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (word32ToIpv4 <$> G.word32 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (ipv4Segment id "ipv4") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SIpv4 <$> lst) === seg + + it "should write out and read back in a ipv6 segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) genIpv6 + + lbs <- buildAsifBytestring "wxyz" Nothing (ipv6Segment id "ipv6") lst + + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SIpv6 <$> lst) === seg + +----- + + it "should write out and read back in a time segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) (TY.posixSecondsToUTCTime . TY.fromMicroseconds <$> G.int64 R.linearBounded) + + lbs <- buildAsifBytestring "wxyz" Nothing (utcTimeMicrosSegment id "time") lst + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (STime <$> lst) === seg + +----- + + it "should write out and read back in a char segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) G.ascii + + lbs <- buildAsifBytestring "wxyz" Nothing (asciiSegment id "char") lst + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SChar <$> lst) === seg + + it "should write out and read back in a null-terminated string segment" $ require $ property $ do + lst <- forAll $ G.list (R.linear 0 50) genNonNullText + + lbs <- buildAsifBytestring "wxyz" Nothing (nullTerminatedStringSegment id "nullterminatedstring") (T.toStrict <$> lst) + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + (SString . T.encodeUtf8 <$> lst) === seg + + it "should write out and read back in a text segment" $ require $ property $ do + -- Explaination of this test: + -- We can't deliniate 'breaks' in Text when we read it back in, + -- so regardless of how many we write out, when we read it back we only + -- get a single Text value. + lst <- forAll $ G.list (R.linear 1 50) (G.text (R.linear 0 50) G.unicode) + + lbs <- buildAsifBytestring "wxyz" Nothing (textSegment id "text") lst + let Right segments = extractSegments (AP.string "seg:wxyz") lbs + [names, times, types, seg] <- forAll $ pure (segmentValues <$> segments) + + [SText . T.encodeUtf8 . T.concat $ T.fromStrict <$> lst] === seg + + +genTriple :: MonadGen m => m (Int64, Word16, T.Text) +genTriple + = (,,) + <$> G.int64 R.linearBounded + <*> G.word16 R.linearBounded + <*> genNonNullText + +genIpv6 :: MonadGen m => m IP.IPv6 +genIpv6 = word32x4ToIpv6 <$> gen4 + where + gen4 :: MonadGen m => m (Word32, Word32, Word32, Word32) + gen4 = + let g = G.word32 R.linearBounded + in (,,,) <$> g <*> g <*> g <*> g + +genNonNullText :: MonadGen m => m T.Text +genNonNullText = + T.fromStrict <$> G.text (R.linear 0 32) (G.filter (/= toEnum 0) G.unicode) + + +instance MonadResource m => MonadResource (PropertyT m) where + liftResourceT = lift . liftResourceT + +-- I have no idea how this doesn't exist, but whatever. ¯\_(ツ)_/¯ +instance MonadResource IO where + liftResourceT = runResourceT diff --git a/test/Gen/Feed.hs b/test/Gen/Feed.hs --- a/test/Gen/Feed.hs +++ b/test/Gen/Feed.hs @@ -6,15 +6,15 @@ ) where import Data.Word -import HaskellWorks.Data.Network.Ip import HaskellWorks.Data.Network.Ip.Internal +import HaskellWorks.Data.Network.Ip.Ipv4 import Hedgehog import qualified Hedgehog.Gen as G import qualified Hedgehog.Range as R -ipv4 :: MonadGen m => Word8 -> Word8 -> m Ipv4Address -ipv4 start stop = Ipv4Address <$> (fourOctetsToWord32 <$> w8 <*> w8 <*> w8 <*> w8) +ipv4 :: MonadGen m => Word8 -> Word8 -> m IpAddress +ipv4 start stop = IpAddress <$> (fourOctetsToWord32 <$> w8 <*> w8 <*> w8 <*> w8) where w8 = G.word8 (R.linear start stop) feedElemIp :: MonadGen m => Range Word32 -> m Word32