diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog + +## 0.0.1.0 + +- Initial version diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2023, Tristan de Cacqueray +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,34 @@ +module Main (main) where + +import Codec.EBML qualified as EBML +import Data.ByteString qualified as BS +import Data.Text qualified as Text +import Data.Text.IO qualified +import System.Environment (getArgs) +import System.IO (Handle, IOMode (ReadMode), stdin, withBinaryFile) + +main :: IO () +main = + getArgs >>= \case + [fp] -> do + let schemas = EBML.webmSchemas + Right ebml <- EBML.decodeEBMLFile schemas fp + Data.Text.IO.putStrLn (EBML.prettyEBMLDocument schemas ebml) + ["split", fp] -> do + let ir = EBML.newStreamReader + withBinaryFile fp ReadMode (printSplit ir) + [] -> printSplit EBML.newStreamReader stdin + _ -> error "usage: haskell-ebml FILE" + +printSplit :: EBML.StreamReader -> Handle -> IO () +printSplit ir handl = do + putStr "Reading 2048 bytes... " + buf <- BS.hGet handl 2048 + case EBML.feedReader buf ir of + Left e -> error (Text.unpack e) + Right (Nothing, _) | buf == "" -> putStrLn "Done!" + Right (mFrame, newIR) -> do + case mFrame of + Nothing -> putStrLn "Need more data" + Just frame -> putStrLn $ "Got a new frame: " <> show (BS.length frame.media) <> " " <> show (BS.take 8 frame.media) + printSplit newIR handl diff --git a/data/firefox-mrec-opus.webm b/data/firefox-mrec-opus.webm new file mode 100644 Binary files /dev/null and b/data/firefox-mrec-opus.webm differ diff --git a/data/var-int.golden b/data/var-int.golden new file mode 100644 --- /dev/null +++ b/data/var-int.golden @@ -0,0 +1,4 @@ + 1000 0010 => 2 + 0100 0000 0000 0010 => 2 + 0010 0000 0000 0000 0000 0010 => 2 + 0001 0000 0000 0000 0000 0000 0000 0010 => 2 diff --git a/ebml.cabal b/ebml.cabal new file mode 100644 --- /dev/null +++ b/ebml.cabal @@ -0,0 +1,115 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.35.1. +-- +-- see: https://github.com/sol/hpack + +name: ebml +version: 0.1.0.0 +synopsis: A pure EBML parser +description: Pure decoder for the Extensible Binary Meta Language (EBML) format. + . + Use this library to parse mkv/webm file or split a webm stream segments. +category: Codec +homepage: https://github.com/TristanCacqueray/haskell-ebml#readme +bug-reports: https://github.com/TristanCacqueray/haskell-ebml/issues +author: Tristan Cacqueray +maintainer: tdecacqu@redhat.com +license: BSD3 +license-file: LICENSE +build-type: Simple +extra-source-files: + CHANGELOG.md + data/firefox-mrec-opus.webm + data/var-int.golden + +source-repository head + type: git + location: https://github.com/TristanCacqueray/haskell-ebml + +library + exposed-modules: + Codec.EBML + other-modules: + Codec.EBML.Element + Codec.EBML.Get + Codec.EBML.Header + Codec.EBML.Pretty + Codec.EBML.Schema + Codec.EBML.Stream + Codec.EBML.WebM + Paths_ebml + hs-source-dirs: + src + default-extensions: + OverloadedStrings + ImportQualifiedPost + LambdaCase + DerivingStrategies + DuplicateRecordFields + OverloadedRecordDot + BlockArguments + StrictData + MultiWayIf + ghc-options: -Wall -fwarn-incomplete-uni-patterns -Wno-partial-type-signatures -fwrite-ide-info -hiedir=.hie -Wno-missing-methods + build-depends: + base <5 + , binary + , bytestring + , containers + , text + default-language: GHC2021 + +executable haskell-ebml + main-is: Main.hs + other-modules: + Paths_ebml + hs-source-dirs: + app + default-extensions: + OverloadedStrings + ImportQualifiedPost + LambdaCase + DerivingStrategies + DuplicateRecordFields + OverloadedRecordDot + BlockArguments + StrictData + MultiWayIf + ghc-options: -Wall -fwarn-incomplete-uni-patterns -Wno-partial-type-signatures -fwrite-ide-info -hiedir=.hie -Wno-missing-methods -threaded -rtsopts -with-rtsopts=-T + build-depends: + base <5 + , bytestring + , ebml + , text + default-language: GHC2021 + +test-suite spec + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Paths_ebml + hs-source-dirs: + test + default-extensions: + OverloadedStrings + ImportQualifiedPost + LambdaCase + DerivingStrategies + DuplicateRecordFields + OverloadedRecordDot + BlockArguments + StrictData + MultiWayIf + ghc-options: -Wall -fwarn-incomplete-uni-patterns -Wno-partial-type-signatures -fwrite-ide-info -hiedir=.hie -Wno-missing-methods -threaded -rtsopts -with-rtsopts=-T + build-depends: + base <5 + , binary + , bytestring + , ebml + , split + , tasty + , tasty-golden + , tasty-hunit + , text + default-language: GHC2021 diff --git a/src/Codec/EBML.hs b/src/Codec/EBML.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML.hs @@ -0,0 +1,99 @@ +{- | This module is intended to be imported qualified: + +> import qualified Codec.EBML as EBML + +Decode a webm file with: + +> EBML.decodeWebMFile "path/file.webm" + +Split a webm stream segments with: + +> let streamReader = EBML.newStreamReader +> buf <- acquire data +> EBML.feedReader buf streamReader + +References: + + - EBML specification introduction: https://matroska-org.github.io/libebml/specs.html + - Document layout: https://www.matroska.org/technical/diagram.html + - The matroska schema: https://github.com/ietf-wg-cellar/matroska-specification/blob/master/ebml_matroska.xml + - The matroska schema doc: https://www.matroska.org/technical/elements.html + - The webm guidelines: https://www.webmproject.org/docs/container/ + - The MSE byte stream format spec: https://w3c.github.io/media-source/index.html#byte-stream-format-specs + - The MSE byte stream format for webm: https://w3c.github.io/mse-byte-stream-format-webm/ +-} +module Codec.EBML ( + -- * WebM decoder + decodeWebM, + WebMDocument (..), + WebMCluster (..), + + -- * Raw EBML decoder + decodeEBMLDocument, + webmSchemas, + + -- * EBML stream reader + module Codec.EBML.Stream, + + -- * EBML data types + EBMLDocument (..), + EBMLElement (..), + EBMLValue (..), + EBMLElementHeader (..), + EBMLID (..), + + -- * EBML schema data types + EBMLSchema (..), + + -- * Helpers + decodeEBMLFile, + decodeWebMFile, + prettyEBMLDocument, + + -- * Low-level API, mostly for testing + EBMLSchemas, + compileSchemas, + getDocument, + getElementHeader, + getElementID, + getDataSize, + getElement, +) where + +import Data.Binary.Get (runGetOrFail) +import Data.ByteString.Lazy qualified as LBS +import Data.Text (Text) +import Data.Text qualified as Text + +import Codec.EBML.Element +import Codec.EBML.Get +import Codec.EBML.Header +import Codec.EBML.Pretty +import Codec.EBML.Schema +import Codec.EBML.Stream +import Codec.EBML.WebM + +-- | Lazy decode a 'WebMDocument'. +decodeWebM :: LBS.ByteString -> Either Text WebMDocument +decodeWebM lbs = decodeWebMDocument =<< decodeEBMLDocument webmSchemas lbs + +-- | The webm document schemas. +webmSchemas :: [EBMLSchema] +webmSchemas = schemaHeader + +-- | Lazy decode a 'EBMLDocument'. +decodeEBMLDocument :: [EBMLSchema] -> LBS.ByteString -> Either Text EBMLDocument +decodeEBMLDocument schemas lbs = case runGetOrFail (getDocument (compileSchemas schemas)) lbs of + Left (_, _, err) -> Left (Text.pack err) + Right ("", _, x) -> Right x + Right (_rest, l, _) -> Left ("Left over data at " <> Text.pack (show l)) + +-- | Decode a raw EBML file. +decodeEBMLFile :: [EBMLSchema] -> FilePath -> IO (Either Text EBMLDocument) +decodeEBMLFile schemas fp = decodeEBMLDocument schemas <$> LBS.readFile fp + +-- | Decode a webm file. +decodeWebMFile :: FilePath -> IO (Either Text WebMDocument) +decodeWebMFile fp = do + ebml <- decodeEBMLFile webmSchemas fp + pure $ decodeWebMDocument =<< ebml diff --git a/src/Codec/EBML/Element.hs b/src/Codec/EBML/Element.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Element.hs @@ -0,0 +1,88 @@ +-- | EBML core data decoder, see: https://matroska-org.github.io/libebml/specs.html +module Codec.EBML.Element where + +import Data.Binary.Get (Get, getWord8) +import Data.Bits (Bits (shift, testBit, (.|.)), (.&.)) +import Data.ByteString (ByteString) +import Data.Int (Int64) +import Data.Text (Text) +import Data.Word (Word32, Word64) + +-- | EBML document structure, including the Header and Body Root. +newtype EBMLDocument = EBMLDocument [EBMLElement] + +-- | EBML element id. +newtype EBMLID = EBMLID Word32 + deriving (Show) + deriving newtype (Num, Eq, Ord) + +-- | EBML element header. +data EBMLElementHeader = EBMLElementHeader + { eid :: EBMLID + , size :: Maybe Word64 + -- ^ size is Nothing for unknown-sized element. + } + deriving (Eq, Show) + +-- | EBML element. +data EBMLElement = EBMLElement + { header :: EBMLElementHeader + , value :: EBMLValue + } + deriving (Eq, Show) + +-- | EBML element value. +data EBMLValue + = EBMLRoot [EBMLElement] + | EBMLSignedInteger Int64 + | EBMLUnsignedInteger Word64 + | EBMLFloat Double + | EBMLText Text + | EBMLDate Text + | EBMLBinary ByteString + deriving (Eq, Show) + +getElementHeader :: Get EBMLElementHeader +getElementHeader = EBMLElementHeader <$> getElementID <*> getMaybeDataSize + +getElementID :: Get EBMLID +getElementID = + EBMLID <$> do + b1 <- getWord8 + let w1 = fromIntegral b1 + if + | b1 `testBit` 7 -> getVar 0 w1 + | b1 `testBit` 6 -> getVar 1 w1 + | b1 `testBit` 5 -> getVar 2 w1 + | b1 `testBit` 4 -> getVar 3 w1 + | otherwise -> fail ("Invalid width: " <> show b1) + +getMaybeDataSize :: Get (Maybe Word64) +getMaybeDataSize = do + sz <- getDataSize + pure $ + -- TODO: better check for unknown-sized for different VINT_DATA size. + -- though, it seems like this is the common value. + if sz == 0xFFFFFFFFFFFFFF + then Nothing + else Just sz + +getDataSize :: Get Word64 +getDataSize = do + b1 <- getWord8 + if + | b1 `testBit` 7 -> getVar 0 (fromIntegral (b1 .&. 127)) + | b1 `testBit` 6 -> getVar 1 (fromIntegral (b1 .&. 63)) + | b1 `testBit` 5 -> getVar 2 (fromIntegral (b1 .&. 31)) + | b1 `testBit` 4 -> getVar 3 (fromIntegral (b1 .&. 15)) + | b1 `testBit` 3 -> getVar 4 (fromIntegral (b1 .&. 7)) + | b1 `testBit` 2 -> getVar 5 (fromIntegral (b1 .&. 3)) + | b1 `testBit` 1 -> getVar 6 (fromIntegral (b1 .&. 1)) + | b1 `testBit` 0 -> getVar 7 0 + | otherwise -> pure 0 + +getVar :: (Num a, Bits a) => Int -> a -> Get a +getVar 0 acc = pure acc +getVar n acc = do + b <- getWord8 + getVar (n - 1) ((acc `shift` 8) .|. fromIntegral b) diff --git a/src/Codec/EBML/Get.hs b/src/Codec/EBML/Get.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Get.hs @@ -0,0 +1,105 @@ +module Codec.EBML.Get where + +import Data.Binary.Get (Get, bytesRead, getByteString, isEmpty, lookAheadM) +import Data.Text.Encoding (decodeUtf8) +import Data.Word (Word64) + +import Codec.EBML.Element +import Codec.EBML.Schema +import Data.Bits (Bits) + +getElement :: EBMLSchemas -> Get EBMLElement +getElement schemas = do + elth <- getElementHeader + getElementValue schemas elth + +getElementValue :: EBMLSchemas -> EBMLElementHeader -> Get EBMLElement +getElementValue schemas elth = do + -- here is a good place to add traceM debug + val <- case lookupSchema elth.eid schemas of + Nothing -> getBinary elth + Just schema -> schema.decode schemas elth + pure $ EBMLElement elth val + +getDocument :: EBMLSchemas -> Get EBMLDocument +getDocument schemas = EBMLDocument <$> go + where + go = do + empty <- isEmpty + if empty + then pure [] + else do + elt <- getElement schemas + elts <- go + pure (elt : elts) + +getBinary :: EBMLElementHeader -> Get EBMLValue +getBinary elth = case elth.size of + Nothing -> fail "Invalid binary header size" + Just sz -> EBMLBinary <$> getByteString (fromIntegral sz) + +getText :: EBMLElementHeader -> Get EBMLValue +getText elth = case elth.size of + Nothing -> fail "Invalid text header size" + Just sz -> EBMLText . decodeUtf8 <$> getByteString (fromIntegral sz) + +getUnsignedInteger :: EBMLElementHeader -> Get EBMLValue +getUnsignedInteger elth = EBMLUnsignedInteger <$> getInt elth.size + +getInteger :: EBMLElementHeader -> Get EBMLValue +getInteger elth = EBMLUnsignedInteger <$> getInt elth.size + +getInt :: (Bits a, Integral a) => Maybe Word64 -> Get a +getInt size = getVar sz 0 + where + -- TODO: check the value is in the [0..8] range + sz = maybe 0 fromIntegral size + +getRoot :: EBMLSchemas -> EBMLElementHeader -> Get EBMLValue +getRoot schemas elth = case elth.size of + Nothing -> EBMLRoot <$> getUntil schemas elth.eid + Just sz -> getRootFixed schemas sz + +getUntil :: EBMLSchemas -> EBMLID -> Get [EBMLElement] +getUntil schemas eid = go + where + getChild :: Get (Maybe EBMLElement) + getChild = do + -- This is not exactly correct. The rfc-8794 spec (chapter 6.2) says we should decode until + -- any valid parent or global element. Because the EBMLSchema doesn't yet contain this information, + -- and because in practice such unknown-sized element are segment/cluster, we simply decode until + -- we find a matching element. + elth <- getElementHeader + if elth.eid == eid + then pure Nothing + else Just <$> getElementValue schemas elth + + go = do + empty <- isEmpty + if empty + then pure [] + else goGet + + goGet = + lookAheadM getChild >>= \case + Just elt -> do + elts <- go + pure (elt : elts) + Nothing -> pure [] + +getRootFixed :: EBMLSchemas -> Word64 -> Get EBMLValue +getRootFixed schemas sz = do + startPosition <- bytesRead + let maxPosition = startPosition + fromIntegral sz + getChilds = do + currentPosition <- bytesRead + if + | currentPosition > maxPosition -> + fail $ "Element decode position " <> show currentPosition <> " exceed parent size " <> show sz + | currentPosition == maxPosition -> + pure [] + | otherwise -> do + elt <- getElement schemas + elts <- getChilds + pure (elt : elts) + EBMLRoot <$> getChilds diff --git a/src/Codec/EBML/Header.hs b/src/Codec/EBML/Header.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Header.hs @@ -0,0 +1,33 @@ +-- | Header schema definition, see: https://github.com/ietf-wg-cellar/ebml-specification/blob/master/specification.markdown#ebml-header-elements +module Codec.EBML.Header where + +import Data.Text (Text) + +import Codec.EBML.Element +import Codec.EBML.Get +import Codec.EBML.Schema + +schemaHeader :: [EBMLSchema] +schemaHeader = + [ EBMLSchema "EBML" 0x1A45DFA3 getRoot + , EBMLSchema "DocType" 0x4282 (const getText) + , EBMLSchema "Segment" 0x18538067 getRoot + , EBMLSchema "Info" 0x1549A966 getRoot + , EBMLSchema "Cluster" 0x1F43B675 getRoot + ] + <> map fromUints uints + +fromUints :: (Text, EBMLID) -> EBMLSchema +fromUints (n, i) = EBMLSchema n i (const getUnsignedInteger) + +uints :: [(Text, EBMLID)] +uints = + [ ("EBMLVersion", 0x4286) + , ("EBMLReadVersion", 0x42F7) + , ("EBMLMaxIDLength", 0x42F2) + , ("EBMLMaxSizeLength", 0x42F3) + , ("DocTypeVersion", 0x4287) + , ("DocTypeReadVersion", 0x4285) + , ("TimestampScale", 0x2AD7B1) + , ("Timestamp", 0xE7) + ] diff --git a/src/Codec/EBML/Pretty.hs b/src/Codec/EBML/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Pretty.hs @@ -0,0 +1,28 @@ +module Codec.EBML.Pretty where + +import Data.ByteString qualified as BS +import Data.Text (Text) +import Data.Text qualified as Text +import Numeric.Natural + +import Codec.EBML.Element +import Codec.EBML.Schema + +-- | Pretty-print a 'EBMLDocument'. +prettyEBMLDocument :: [EBMLSchema] -> EBMLDocument -> Text +prettyEBMLDocument schemas (EBMLDocument xs) = mconcat $ map (prettyElement (compileSchemas schemas) 0) xs + +prettyElement :: EBMLSchemas -> Natural -> EBMLElement -> Text +prettyElement schemas indent elt = indentTxt <> eltIDTxt <> ": " <> eltValueTxt + where + indentTxt = Text.replicate (fromIntegral indent) " " + eltIDTxt = case lookupSchema elt.header.eid schemas of + Just schema -> schema.name + Nothing -> Text.pack (show elt.header.eid) + eltValueTxt = case elt.value of + EBMLRoot xs -> "\n" <> mconcat (map (prettyElement schemas (indent + 2)) xs) + EBMLText txt -> txt <> "\n" + EBMLBinary bs -> "[raw:" <> Text.pack (show $ BS.length bs) <> " " <> bsTxt bs <> "]\n" + EBMLUnsignedInteger x -> Text.pack (show x) <> "\n" + _ -> "value\n" + bsTxt = Text.pack . show . BS.take 32 diff --git a/src/Codec/EBML/Schema.hs b/src/Codec/EBML/Schema.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Schema.hs @@ -0,0 +1,36 @@ +module Codec.EBML.Schema where + +import Data.Binary.Get (Get) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) + +import Codec.EBML.Element + +{- | EBML schema definition. + +Note that this is missing: + +- min/max occurrence constraint. +- default value. +- element path. +-} +data EBMLSchema = EBMLSchema + { name :: Text + -- ^ The element name. + , eid :: EBMLID + -- ^ The element id. + , decode :: EBMLSchemas -> EBMLElementHeader -> Get EBMLValue + -- ^ How to decode the element value. + } + +newtype EBMLSchemas = EBMLSchemas (Map EBMLID EBMLSchema) + +-- | Combine a list of schema for decoder. +compileSchemas :: [EBMLSchema] -> EBMLSchemas +compileSchemas = EBMLSchemas . Map.fromList . map toKV + where + toKV schema = (schema.eid, schema) + +lookupSchema :: EBMLID -> EBMLSchemas -> Maybe EBMLSchema +lookupSchema eid (EBMLSchemas schemas) = Map.lookup eid schemas diff --git a/src/Codec/EBML/Stream.hs b/src/Codec/EBML/Stream.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/Stream.hs @@ -0,0 +1,123 @@ +module Codec.EBML.Stream (StreamReader, newStreamReader, StreamFrame (..), feedReader) where + +import Control.Monad (void, when) +import Data.Binary.Get qualified as Get +import Data.ByteString qualified as BS +import Data.Text (Text) +import Data.Text qualified as Text + +import Codec.EBML.Element +import Codec.EBML.Get +import Codec.EBML.Header +import Codec.EBML.Schema +import Codec.EBML.WebM qualified as WebM + +-- | A valid frame that can be served. +data StreamFrame = StreamFrame + { initialization :: BS.ByteString + -- ^ The initialization segments, to be provided before the first media segment. + , media :: BS.ByteString + -- ^ The begining of the last media segment found in the input buffer. + } + +-- | Create a stream reader with 'newStreamReader', and decode media segments with 'feedReader'. +data StreamReader = StreamReader + { acc :: [BS.ByteString] + -- ^ Accumulate data in case the header is not completed in the first buffer. + , consumed :: Int + -- ^ Keep track of the decoder position accross multiple buffers. + , header :: Maybe BS.ByteString + -- ^ The stream initialization segments. + , decoder :: Get.Decoder () + -- ^ The current decoder. + } + +streamSchema :: EBMLSchemas +streamSchema = compileSchemas schemaHeader + +-- | Read the initialization frame. +getInitialization :: Get.Get () +getInitialization = do + -- Read the EBML header element + elt <- getElement streamSchema + when (elt.header.eid /= 0x1A45DFA3) do + fail $ "Invalid magic: " <> show elt.header + + -- Read the begining of the first segment, until the first cluster + segmentHead <- getElementHeader + when (segmentHead.eid /= 0x18538067) do + fail $ "Invalid segment: " <> show segmentHead + elts <- getUntil streamSchema 0x1F43B675 + case WebM.decodeSegment elts of + Right _webmDocument -> pure () + Left err -> fail (Text.unpack err) + +-- | Read a cluster frame. +getCluster :: Get.Get () +getCluster = do + clusterHead <- getElementHeader + when (clusterHead.eid /= 0x1F43B675) do + fail $ "Invalid cluster: " <> show clusterHead + getClusterBody + +getClusterBody :: Get.Get () +getClusterBody = do + elts <- getUntil streamSchema 0x1F43B675 + case elts of + (elt : _) | elt.header.eid == 0xE7 -> pure () + _ -> fail "Cluster first element is not a timestamp" + +getClusterRemaining :: Get.Get () +getClusterRemaining = do + elth <- getElementHeader + if elth.eid == 0x1F43B675 + then -- This is in fact a new cluster, get its body + getClusterBody + else -- This is a cluster left-over, let's keep on reading until a new start + void (getUntil streamSchema 0x1F43B675) + +-- | Initialize a stream reader. +newStreamReader :: StreamReader +newStreamReader = StreamReader [] 0 Nothing (Get.runGetIncremental getInitialization) + +-- | Feed data into a stream reader. Returns either an error, or maybe a new 'StreamFrame' and an updated StreamReader. +feedReader :: BS.ByteString -> StreamReader -> Either Text (Maybe StreamFrame, StreamReader) +feedReader = go Nothing + where + -- This is the end + go Nothing "" sr = case Get.pushEndOfInput sr.decoder of + Get.Fail _ _ s -> Left (Text.pack s) + Get.Partial _ -> Left "Missing data" + Get.Done "" _ _ -> Right (Nothing, sr) + Get.Done{} -> Left "Left-over data" + -- Feed the decoder + go mFrame bs sr = + case Get.pushChunk sr.decoder bs of + Get.Fail _ _ s -> Left (Text.pack s) + newDecoder@(Get.Partial _) -> + let newAcc = case sr.header of + Nothing -> bs : sr.acc + -- We don't need to accumulate data once the header is known. + Just _ -> [] + newSR = StreamReader newAcc (sr.consumed + BS.length bs) sr.header newDecoder + in Right (mFrame, newSR) + Get.Done leftover consumed _ + | BS.null leftover -> + -- We might have ended on a in-cluster element, use the remainingDecoder next time + Right (mFrame, newIR remainingDecoder) + | otherwise -> + -- There might be a new frame after, keep on decoding + go newFrame leftover (newIR segmentDecoder) + where + -- The header is either the one already parsed, or the current complete decoded buffer. + newHeader = case sr.header of + Just header -> header + Nothing -> + let currentPos = fromIntegral consumed - sr.consumed + in mconcat $ reverse (BS.take currentPos bs : sr.acc) + -- The new frame starts after what was decoded. + newFrame = Just (StreamFrame newHeader leftover) + newIR = StreamReader [] 0 (Just newHeader) + + remainingDecoder = Get.runGetIncremental getClusterRemaining + segmentDecoder = Get.runGetIncremental getCluster diff --git a/src/Codec/EBML/WebM.hs b/src/Codec/EBML/WebM.hs new file mode 100644 --- /dev/null +++ b/src/Codec/EBML/WebM.hs @@ -0,0 +1,96 @@ +{- | This module contains the logic to convert a raw EBMLDocument into a WebMDocument + +See: https://www.matroska.org/technical/diagram.html +-} +module Codec.EBML.WebM where + +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Word (Word64) + +import Codec.EBML.Element +import Data.Foldable (find) +import Data.Maybe (catMaybes) + +-- | A WebM document. +data WebMDocument = WebMDocument + { timestampScale :: Word64 + -- ^ Base unit for Segment Ticks and Track Ticks, in nanoseconds. A TimestampScale of 1_000_000 means segments' timestamps are expressed in milliseconds; + , clusters :: [WebMCluster] + -- ^ The list of clusters. + } + +-- | A WebM cluster, e.g. a media segment. +data WebMCluster = WebMCluster + { timestamp :: Word64 + -- ^ Absolute timestamp of the cluster. + , content :: [EBMLElement] + -- ^ The cluster elements. + } + +decodeWebMDocument :: EBMLDocument -> Either Text WebMDocument +decodeWebMDocument = \case + (EBMLDocument [header, segment]) -> do + headerElements <- getChilds header + segmentElements <- getChilds segment + docType <- getText =<< getElt headerElements 0x4282 + docVersion <- getUInt =<< getElt headerElements 0x4287 + if docType /= "webm" || docVersion /= 2 + then Left ("Invalid doctype: " <> Text.pack (show (docType, docVersion))) + else decodeSegment segmentElements + _ -> Left "Invalid EBML file structure" + +decodeSegment :: [EBMLElement] -> Either Text WebMDocument +decodeSegment = go 0 + where + go scale xs@(x : rest) + | x.header.eid == 0x1F43B675 = WebMDocument scale . catMaybes <$> traverse decodeWebMCluster xs + | x.header.eid == 0x1549A966 = do + info <- getChilds x + scaleValue <- getUInt =<< getElt info 0x2AD7B1 + go scaleValue rest + | otherwise = go scale rest + go scale [] = Right (WebMDocument scale []) + +decodeWebMCluster :: EBMLElement -> Either Text (Maybe WebMCluster) +decodeWebMCluster elt + | elt.header.eid == 0x1F43B675 = + Just <$> do + childs <- getChilds elt + case childs of + (tsElt : xs) + | tsElt.header.eid == 0xE7 -> do + timestamp <- getUInt tsElt + Right $ WebMCluster timestamp xs + _ -> Left "Cluster first element is not a timestamp" + | otherwise = Right Nothing + +-- | Extract the document type, version and the segment elements. +documentSegment :: EBMLDocument -> Either Text (Text, Word64, [EBMLElement]) +documentSegment (EBMLDocument [header, segment]) = do + headerElements <- getChilds header + segmentElements <- getChilds segment + docType <- getText =<< getElt headerElements 0x4282 + docVersion <- getUInt =<< getElt headerElements 0x4287 + pure (docType, docVersion, segmentElements) +documentSegment _ = Left "Invalid EBML file structure" + +getElt :: [EBMLElement] -> EBMLID -> Either Text EBMLElement +getElt xs eid = case find (\elt -> elt.header.eid == eid) xs of + Just elt -> Right elt + Nothing -> Left ("Element " <> Text.pack (show eid) <> " not found") + +getText :: EBMLElement -> Either Text Text +getText elt = case elt.value of + EBMLText txt -> Right txt + _ -> Left "Invalid text value" + +getUInt :: EBMLElement -> Either Text Word64 +getUInt elt = case elt.value of + EBMLUnsignedInteger x -> Right x + _ -> Left ("Invalid uint value " <> Text.pack (show elt.value)) + +getChilds :: EBMLElement -> Either Text [EBMLElement] +getChilds elt = case elt.value of + EBMLRoot xs -> Right xs + _ -> Left "Element is not a root" diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,97 @@ +module Main (main) where + +import Control.Exception (SomeException, try) +import Data.Binary.Get (runGet) +import Data.Binary.Put (putWord8, runPut) +import Data.ByteString.Char8 qualified as BS +import Data.Char (digitToInt) +import Data.Either (fromRight) +import Data.Foldable (forM_, traverse_) +import Data.List (foldl') +import Data.List.Split (chunksOf) +import Data.Text qualified as Text +import Data.Word (Word8) +import Test.Tasty +import Test.Tasty.Golden +import Test.Tasty.HUnit + +import Codec.EBML qualified as EBML + +main :: IO () +main = do + sampleFile <- readFileMaybe "./data/Volcano_Lava_Sample.webm" + streamFile <- readFileMaybe "./data/firefox-mrec-opus.webm" + defaultMain $ testGroup "Codec.EBML" [unitTests, integrationTests sampleFile streamFile] + where + readFileMaybe fp = fromRight "" <$> try @SomeException (BS.readFile fp) + +unitTests :: TestTree +unitTests = + testGroup + "Unit tests" + [ goldenVsString "VarInt" "data/var-int.golden" + . pure + . BS.fromStrict + . BS.unlines + . map testVarInts + $ [ "1000 0010" + , "0100 0000 0000 0010" + , "0010 0000 0000 0000 0000 0010" + , "0001 0000 0000 0000 0000 0000 0000 0010" + ] + ] + +testVarInts :: String -> BS.ByteString +testVarInts str = BS.pack $ padStr <> " => " <> padVal + where + padStr = replicate (42 - length str) ' ' <> str + padVal = replicate (8 - length val) ' ' <> val + val = show $ runGet EBML.getDataSize bs + bs = runPut $ traverse_ (putWord8 . readOctet) $ chunksOf 8 $ filter (/= ' ') str + +readOctet :: String -> Word8 +readOctet s + | length s /= 8 = error $ "Invalid length: " <> s + | otherwise = fromIntegral $ foldl' (\acc x -> acc * 2 + digitToInt x) 0 s + +integrationTests :: BS.ByteString -> BS.ByteString -> TestTree +integrationTests sampleFile streamFile = + testGroup + "Integration tests" + [ testGroup "sample" (decodeFile sampleFile 53054906 4458 49 0 2794) + , testGroup "stream" (decodeFile streamFile 3499 260 6 6 1006) + ] + where + decodeFile bs size headerSize clusterCount clusterTs1 clusterTs2 + | BS.length bs /= size = + -- the file was not checkout + [testCase "skip" (pure ())] + | otherwise = + [ testCase "decode" do + -- lazy decode + case EBML.decodeWebM (BS.fromStrict bs) of + Left e -> error (Text.unpack e) + Right webM -> do + webM.timestampScale @?= 1_000_000 + length webM.clusters @?= clusterCount + (head webM.clusters).timestamp @?= clusterTs1 + (webM.clusters !! 1).timestamp @?= clusterTs2 + , testCase "stream" do + -- incremental decode + let go buf sr acc = + let (cur, next) = BS.splitAt 256 buf + in case EBML.feedReader cur sr of + Left e -> error (Text.unpack e) + Right (mFrame, nextSR) + | cur == "" -> reverse newAcc + | otherwise -> go next nextSR newAcc + where + newAcc = maybe acc (: acc) mFrame + frames = go bs EBML.newStreamReader [] + -- this works because the chunk size is small enough to get every segment. + length frames @?= clusterCount + BS.length (head frames).initialization @?= headerSize + BS.take 4 (head frames).initialization @?= "\x1A\x45\xdf\xa3" + forM_ frames $ \frame -> do + BS.take 4 frame.media @?= "\x1f\x43\xb6\x75" + ]