Skip to content

Commit 187d563

Browse files
superfashisuperfashi
committed
partial finish
1 parent ae8d5d5 commit 187d563

File tree

6 files changed

+289
-183
lines changed

6 files changed

+289
-183
lines changed

‎.gitignore

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
.stack-work/
2-
*~
2+
*~
3+
*.EXE.prof

‎app/Main.hs

+179-23
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,32 @@
11
{-# LANGUAGE LambdaCase #-}
2+
{-# LANGUAGE TupleSections #-}
23

34
module Main (main) where
45

56
import Common (FragmentationIndicator (..), consumeAll)
67
import Control.Monad (unless, when)
7-
import Control.Monad.IO.Class (MonadIO (liftIO))
8-
import Data.Binary (Binary, Get, Word16, Word32, Word64, get)
8+
import Data.Binary (Word16, Word32, Word64, get)
99
import Data.Binary.Get (runGetOrFail)
10-
import qualified Data.ByteString.Lazy as L (readFile)
10+
import qualified Data.ByteString.Lazy as L (appendFile, drop, empty, hPut, length, pack, putStr, readFile, writeFile)
1111
import Data.ByteString.Lazy.Internal (ByteString (Empty))
1212
import Data.Int (Int64)
1313
import qualified Data.Map as Map
14-
import Data.Maybe (isJust)
15-
import Lib
16-
import Message (ControlMessage, ControlMessages (..))
14+
import Data.Maybe (isJust, isNothing)
15+
import Lib hiding (MFUNonTimed, MFUTimed)
16+
import qualified Lib as MFUNonTimed (MFUNonTimed (mfuData))
17+
import qualified Lib as MFUTimed (MFUTimed (mfuData))
18+
import Message
1719
import Net (NetworkTimeProtocolData)
1820
import Streamly.Data.Stream (Stream, mapMaybe, unfoldrM)
1921
import qualified Streamly.Data.Stream as S (foldr, foldrM, mapM, take, toList)
2022
import qualified Streamly.Internal.Data.Pipe.Type as Pipe
2123
import Streamly.Internal.Data.SVar.Type (adaptState)
2224
import Streamly.Internal.Data.Stream (Step (..), Stream (Stream))
25+
import qualified Streamly.Internal.Data.Stream as IS (mapM_)
26+
import qualified Streamly.Internal.Data.Stream.StreamD.Container as D (nub)
2327
import System.Environment (getArgs)
28+
import System.IO (IOMode (WriteMode), hPutStrLn, withBinaryFile)
29+
import Table
2430
import Text.Printf (printf)
2531
import Text.Show.Pretty (pPrint)
2632

@@ -126,39 +132,168 @@ reassembleControlMessages (Stream step state) = do
126132
miss :: Word32 -> IO ()
127133
miss s = putStrLn $ printf "Sequence number missing for Context ID 0x%X, Packet ID 0x%X, %d -> %d" cid pid s pseq
128134

135+
type MFUState = Map.Map (Word16, Word16) (Word32, Maybe ByteString)
136+
137+
reassembleMFU :: Stream IO ((Word16, Word16), Word32, MPU) -> Stream IO ((Word16, Word16), ByteString)
138+
reassembleMFU (Stream step state) = do
139+
Stream step' ([], Map.empty, state)
140+
where
141+
step' _ ((cpid, x) : xs, m, st) = do
142+
return $ Yield (cpid, x) (xs, m, st)
143+
step' gst ([], m, st) = do
144+
r <- step (adaptState gst) st
145+
case r of
146+
Yield a s -> do
147+
res <- consume m a
148+
case res of
149+
Pipe.Yield a' s' -> return $ Skip (a', s', s)
150+
Pipe.Continue s' -> return $ Skip ([], s', s)
151+
Skip s -> return $ Skip ([], m, s)
152+
Stop -> do
153+
let m' = Map.filter (isJust . snd) m
154+
unless (Map.null m') $ print $ "Unfinished fragmented messages " ++ show (Map.keys m')
155+
return Stop
156+
157+
consume :: MFUState -> ((Word16, Word16), Word32, MPU) -> IO (Pipe.Step MFUState [((Word16, Word16), ByteString)])
158+
consume m ((cid, pid), pseq, MPU ind _ _ frag) = case ind of
159+
FragmentationIndicatorUndivided -> do
160+
handleHead
161+
let t = extractMulti frag
162+
return $ Pipe.Yield (((cid, pid),) <$> t) $ Map.insert (cid, pid) (pseq, Nothing) m
163+
FragmentationIndicatorDividedHead -> do
164+
handleHead
165+
let t = extractOne frag
166+
return $ Pipe.Continue $ Map.insert (cid, pid) (pseq, Just t) m
167+
FragmentationIndicatorDividedBody -> do
168+
prev <- handleEnd
169+
case prev of
170+
Nothing -> return $ Pipe.Continue m
171+
Just bs -> do
172+
let t = extractOne frag
173+
return $ Pipe.Continue $ Map.insert (cid, pid) (pseq, Just $ bs <> t) m
174+
FragmentationIndicatorDividedEnd -> do
175+
prev <- handleEnd
176+
case prev of
177+
Nothing -> return $ Pipe.Continue m
178+
Just bs -> do
179+
let t = extractOne frag
180+
return $ Pipe.Yield [((cid, pid), bs <> t)] $ Map.insert (cid, pid) (pseq, Nothing) m
181+
where
182+
extractOne :: MPUFragment -> ByteString
183+
extractOne (MPUFragmentMFU (MFUNonTimedType d)) = MFUNonTimed.mfuData d
184+
extractOne (MPUFragmentMFU (MFUTimedType d)) = MFUTimed.mfuData d
185+
extractOne c = error $ show (cid, pid, pseq, c)
186+
187+
extractMulti :: MPUFragment -> [ByteString]
188+
extractMulti (MPUFragmentMFU (MFUNonTimedType d)) = [MFUNonTimed.mfuData d]
189+
extractMulti (MPUFragmentMFU (MFUTimedType d)) = [MFUTimed.mfuData d]
190+
extractMulti (MPUFragmentMFU (MFUNonTimedAggregatedType d)) = MFUNonTimed.mfuData <$> d
191+
extractMulti (MPUFragmentMFU (MFUTimedAggregatedType d)) = MFUTimed.mfuData <$> d
192+
extractMulti c = error $ show (cid, pid, pseq, c)
193+
194+
handleHead :: IO ()
195+
handleHead =
196+
case Map.lookup (cid, pid) m of
197+
Nothing -> return ()
198+
Just (s, _) | s + 1 /= pseq -> miss s
199+
Just (_, Just _) -> fail "Incorrect fragmentation indicator"
200+
_ -> return ()
201+
202+
handleEnd :: IO (Maybe ByteString)
203+
handleEnd =
204+
case Map.lookup (cid, pid) m of
205+
Nothing -> do
206+
discard
207+
return Nothing
208+
Just (s, _) | s + 1 /= pseq -> miss s
209+
Just (_, Just b) -> return $ Just b
210+
Just (_, Nothing) -> fail "Incorrect fragmentation indicator"
211+
212+
discard :: IO ()
213+
discard = putStrLn $ "\tDiscarding sequence " ++ show pseq ++ ", indicator " ++ show ind
214+
215+
miss :: Word32 -> IO a
216+
miss s = do
217+
when (s >= pseq) $ fail $ printf "Sequence number regression for Context ID 0x%X, Packet ID 0x%X, %d -> %d" cid pid s pseq
218+
fail $ printf "Sequence number missing for Context ID 0x%X, Packet ID 0x%X, %d -> %d" cid pid s pseq
219+
129220
parseTLVPackets :: ParseState -> IO (Maybe (TLVPacket, ParseState))
130221
parseTLVPackets (ParseState Empty _) = return Nothing
131222
parseTLVPackets (ParseState bs o) = do
132223
case runGetOrFail get bs of
133224
Left (_, off, err) -> fail $ "error parse at " ++ show (o + off) ++ ": " ++ err
134225
Right (r, off, p) -> return $ Just (p, ParseState r (o + off))
135226

136-
extractControlMessages :: TLVPacket -> Maybe ((Word16, Word16), Word32, ControlMessages)
137-
extractControlMessages
227+
extractComponent :: (MMTPPayload -> Maybe a) -> TLVPacket -> Maybe ((Word16, Word16), Word32, a)
228+
extractComponent
229+
e
138230
( TLVPacketHeaderCompressedIP
139231
( CompressedIPPacket
140232
{ contextId = cid,
141233
contextIdentificationHeader = h
142234
}
143235
)
144-
) = case h of
145-
ContextIdentificationNoCompressedHeader h' -> case h' of
146-
( MMTPPacket
147-
{ packetId = pid,
148-
packetSequenceNumber = pseq,
149-
mmtpPayload = MMTPPayloadControlMessages m
150-
}
151-
) -> Just ((cid, pid), pseq, m)
236+
) =
237+
case h of
238+
ContextIdentificationHeaderPartialIPv6UDP _ _ m -> extractMMTP m
239+
ContextIdentificationNoCompressedHeader m -> extractMMTP m
152240
_ -> Nothing
153-
_ -> Nothing
154-
extractControlMessages _ = Nothing
241+
where
242+
extractMMTP
243+
( MMTPPacket
244+
{ Lib.packetId = pid,
245+
packetSequenceNumber = pseq,
246+
mmtpPayload = p
247+
}
248+
) = case e p of Just a -> Just ((cid, pid), pseq, a); _ -> Nothing
249+
extractComponent _ _ = Nothing
250+
251+
extractControlMessages :: TLVPacket -> Maybe ((Word16, Word16), Word32, ControlMessages)
252+
extractControlMessages = extractComponent (\case MMTPPayloadControlMessages m -> Just m; _ -> Nothing)
253+
254+
extractMPU :: TLVPacket -> Maybe ((Word16, Word16), Word32, MPU)
255+
extractMPU = extractComponent (\case MMTPPayloadMPU m -> Just m; _ -> Nothing)
155256

156257
extractNTP :: TLVPacket -> Maybe NetworkTimeProtocolData
157258
extractNTP (TLVPacketIPv6 (IPv6Packet {payload = p})) = case p of
158259
Right d -> Just d
159260
_ -> Nothing
160261
extractNTP _ = Nothing
161262

263+
extractMPT :: ControlMessage -> Maybe MMTPackageTable
264+
extractMPT (ControlMessagePA (PAMessage {tables = tbls})) = extract tbls
265+
where
266+
extract :: [Table] -> Maybe MMTPackageTable
267+
extract [] = Nothing
268+
extract (MPT x : _) = Just x
269+
extract (_ : xs) = extract xs
270+
extractMPT _ = Nothing
271+
272+
extractPLT :: ControlMessage -> Maybe PackageListTable
273+
extractPLT (ControlMessagePA (PAMessage {tables = tbls})) = extract tbls
274+
where
275+
extract :: [Table] -> Maybe PackageListTable
276+
extract [] = Nothing
277+
extract (PLT x : _) = Just x
278+
extract (_ : xs) = extract xs
279+
extractPLT _ = Nothing
280+
281+
extractMFU :: MPU -> Maybe MFU
282+
extractMFU (MPU {mpuFragment = MPUFragmentMFU m}) = do
283+
case m of
284+
MFUNonTimedType m' -> Just $ MFUNonTimedType $ m' {MFUNonTimed.mfuData = L.empty}
285+
MFUTimedType m' -> Just $ MFUTimedType $ m' {MFUTimed.mfuData = L.empty}
286+
_ -> Nothing
287+
extractMFU _ = Nothing
288+
289+
type CollectMap = Map.Map (Word16, Word16) ByteString
290+
291+
collectMFU :: ((Word16, Word16), ByteString) -> CollectMap -> CollectMap
292+
collectMFU (cpid, bs) = Map.insertWith ((<>) . fix) cpid bs
293+
where
294+
fix :: ByteString -> ByteString
295+
fix bs = if L.length bs < 4 then error "Invalid MFU" else L.pack [0, 0, 0, 1] <> L.drop 4 bs
296+
162297
main :: IO ()
163298
main = do
164299
args <- getArgs
@@ -167,14 +302,35 @@ main = do
167302
x : _ -> return x
168303
file <- L.readFile filename
169304
let packets = unfoldrM parseTLVPackets (ParseState file 0)
170-
messages = reassembleControlMessages $ mapMaybe extractControlMessages packets
305+
let messages = reassembleControlMessages $ mapMaybe extractControlMessages packets
171306
-- (r, o) <- S.foldrM reassembleControlMessages (pure (mempty, [])) fragments
172307

173308
-- _ <- S.toList $ S.mapM pPrint $ (\(pid, pseq, m) -> (pid, pseq, Message.fragmentationIndicator m)) <$> fragments
174-
len <- S.foldr (\_ c -> c + 1) (0 :: Word64) $ mapMaybe extractNTP packets
175-
print len
176-
len <- S.foldr (\_ c -> c + 1) (0 :: Word64) messages
177-
print len
309+
-- len <- S.foldr (\_ c -> c + 1) (0 :: Word64) $ mapMaybe extractNTP packets
310+
-- print len
311+
-- len <- S.foldr (\_ c -> c + 1) (0 :: Word64) messages
312+
-- print len
313+
314+
let mpuPackets = mapMaybe extractMPU packets
315+
reasembledMFU = reassembleMFU (mapMaybe (\case w@((1, 0xF100), _, _) -> Just w; _ -> Nothing) mpuPackets)
316+
-- len <- S.foldr (\_ c -> c + 1) (0 :: Word64) mpuPackets
317+
-- print len
318+
-- S.toList $ S.mapM pPrint $ S.take 3 mpuPackets
319+
-- return ()
320+
collected <- S.foldr collectMFU Map.empty reasembledMFU
321+
L.writeFile "dump" $ collected Map.! (1, 0xF100)
322+
323+
-- IS.mapM_ (\(c, b) -> pPrint (c, b)) reasembledMFU
324+
325+
-- withBinaryFile "dump" WriteMode $ \h ->
326+
-- IS.mapM_ (hPutStrLn h) $
327+
-- mapMaybe
328+
-- ( \(s, p) ->
329+
-- case extractMFU p of
330+
-- Just c -> Just $ show (Lib.fragmentationIndicator p, s, mpuSequenceNumber p, c)
331+
-- _ -> Nothing
332+
-- )
333+
-- (mapMaybe (\case ((1, 0xF100), s, h) -> Just (s, h); _ -> Nothing) mpuPackets)
178334

179335
-- list <- toList $ parseTLVPackets (runGetIncremental get) file
180336
-- print list

0 commit comments

Comments
 (0)