OGGSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.OGGSpec (tests) where

import Test.Tasty
import Test.Tasty.HUnit
import Control.Monad.Except (runExceptT)
import Data.Bits (shiftL, shiftR, testBit, xor)
import qualified Data.ByteString as BS
import Data.Word (Word32)
import qualified Data.ByteString.Base64 as B64
import qualified Data.HashMap.Strict as HM
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import System.Directory (getTemporaryDirectory, removeFile)
import System.FilePath ((</>))
import System.OsPath (encodeFS)

import Monatone.Metadata
import Monatone.OGG (parseOGG, loadAlbumArtOGG)
import Monatone.Writer (updateMetadata, setTitle, emptyUpdate)

tests :: TestTree
tests = testGroup "OGG Parser"
  [ testCase "parses comment packet spanning multiple pages" testSpanningComment
  , testCase "parses Opus tags" testOpusTags
  , testCase "loads album art with base64 padding from spanning packet" testSpanningAlbumArt
  , testCase "computes Vorbis duration from the last granule" testVorbisDuration
  , testCase "computes Opus duration accounting for pre-skip" testOpusDuration
  , testCase "updates Vorbis tags in place" testVorbisWrite
  , testCase "updates Opus tags in place" testOpusWrite
  ]

-- | A comment packet larger than one page (> 64KB) must be reassembled
-- from its continuation pages, not parsed per-page
testSpanningComment :: IO ()
testSpanningComment = do
  let bigValue = T.replicate 100000 "x"  -- forces the packet across pages
      commentPacket = vorbisCommentPacket
        [ ("TITLE", "Spanning Title")
        , ("BIGTAG", bigValue)
        , ("ARTIST", "Spanning Artist")
        ]
  metadata <- parseSynthetic "monatone-test-span.ogg" [vorbisIdentPacket, commentPacket]
  assertEqual "title" (Just "Spanning Title") (title metadata)
  assertEqual "artist" (Just "Spanning Artist") (artist metadata)
  assertEqual "big tag survives reassembly" (Just [bigValue]) (HM.lookup "BIGTAG" (rawTags metadata))
  assertEqual "sample rate" (Just 44100) (sampleRate (audioProperties metadata))

-- | OpusTags packets were previously never matched at all
testOpusTags :: IO ()
testOpusTags = do
  let commentPacket = opusTagsPacket [("TITLE", "Opus Title"), ("ARTIST", "Opus Artist")]
  metadata <- parseSynthetic "monatone-test-opus.ogg" [opusIdentPacket, commentPacket]
  assertEqual "codec" (Just CodecOpus) (codec (audioProperties metadata))
  assertEqual "title" (Just "Opus Title") (title metadata)
  assertEqual "artist" (Just "Opus Artist") (artist metadata)
  -- Opus always decodes at 48kHz regardless of the header's input rate
  assertEqual "sample rate" (Just 48000) (sampleRate (audioProperties metadata))

-- | Album art forces both a spanning packet and base64 '=' padding
testSpanningAlbumArt :: IO ()
testSpanningAlbumArt = do
  let imageData = BS.replicate 80000 0xAB  -- forces the packet across pages
      pictureBlock = flacPictureBlock "image/png" imageData
      encoded = TE.decodeUtf8 (B64.encode pictureBlock)
      commentPacket = vorbisCommentPacket
        [ ("TITLE", "Art Test")
        , ("METADATA_BLOCK_PICTURE", encoded)
        ]
  path <- writeSynthetic "monatone-test-art.ogg" [vorbisIdentPacket, commentPacket]
  osPath <- encodeFS path
  -- The encoded block ends in '=' padding; the value must not be truncated
  assertBool "test data has base64 padding" ("=" `T.isSuffixOf` encoded)
  result <- runExceptT $ loadAlbumArtOGG osPath
  removeFile path
  case result of
    Left err -> assertFailure $ "loadAlbumArtOGG failed: " ++ show err
    Right Nothing -> assertFailure "Expected album art, got Nothing"
    Right (Just art) -> do
      assertEqual "mime type" "image/png" (albumArtMimeType art)
      assertEqual "image data" imageData (albumArtData art)

-- | Duration comes from the granule of the file's last page
testVorbisDuration :: IO ()
testVorbisDuration = do
  let bytes = mkOggFile [vorbisIdentPacket, vorbisCommentPacket [("TITLE", "T")]]
        <> lastPage 441000  -- 10s at 44100Hz
  metadata <- parseSyntheticBytes "monatone-dur.ogg" bytes
  assertEqual "duration" (Just 10000) (duration (audioProperties metadata))

-- | Opus granules run at 48kHz and include the pre-skip
testOpusDuration :: IO ()
testOpusDuration = do
  let bytes = mkOggFile [opusIdentPacket, opusTagsPacket [("TITLE", "T")]]
        <> lastPage (480000 + 312)  -- 10s at 48kHz + the header's 312 pre-skip
  metadata <- parseSyntheticBytes "monatone-dur.opus.ogg" bytes
  assertEqual "duration" (Just 10000) (duration (audioProperties metadata))

-- | Updating tags rewrites the comment packet, keeps the setup packet and
-- audio pages, and leaves every page with a valid CRC
testVorbisWrite :: IO ()
testVorbisWrite = do
  let commentPacket = vorbisCommentPacket [("ARTIST", "Keep Me"), ("TITLE", "Old")]
      setupPacket = "\x05vorbis" <> BS.replicate 24 0xAB
      audioPacket = "AUDIOPACKETDATA" <> BS.replicate 300 0x77
      bytes = mkOggPageGroups [[vorbisIdentPacket], [commentPacket, setupPacket], [audioPacket]]
  path <- writeSyntheticBytes "monatone-write.ogg" bytes
  osPath <- encodeFS path

  runExceptT (updateMetadata osPath (setTitle "New Title" emptyUpdate))
    >>= either (assertFailure . show) return

  metadata <- runExceptT (parseOGG osPath) >>= either (assertFailure . show) return
  final <- BS.readFile path
  removeFile path
  assertEqual "title updated" (Just "New Title") (title metadata)
  assertEqual "artist kept" (Just "Keep Me") (artist metadata)
  assertBool "audio kept" ("AUDIOPACKETDATA" `BS.isInfixOf` final)
  assertBool "setup packet kept" (BS.replicate 24 0xAB `BS.isInfixOf` final)
  assertBool "all page CRCs valid" (validOggCrcs final)

testOpusWrite :: IO ()
testOpusWrite = do
  let audioPacket = "OPUSAUDIODATA" <> BS.replicate 100 0x55
      bytes = mkOggPageGroups
        [ [opusIdentPacket]
        , [opusTagsPacket [("ARTIST", "Keep Me"), ("TITLE", "Old")]]
        , [audioPacket]
        ]
  path <- writeSyntheticBytes "monatone-write-opus.ogg" bytes
  osPath <- encodeFS path

  runExceptT (updateMetadata osPath (setTitle "Opus New" emptyUpdate))
    >>= either (assertFailure . show) return

  metadata <- runExceptT (parseOGG osPath) >>= either (assertFailure . show) return
  final <- BS.readFile path
  removeFile path
  assertEqual "title updated" (Just "Opus New") (title metadata)
  assertEqual "artist kept" (Just "Keep Me") (artist metadata)
  assertEqual "codec" (Just CodecOpus) (codec (audioProperties metadata))
  assertBool "audio kept" ("OPUSAUDIODATA" `BS.isInfixOf` final)
  assertBool "all page CRCs valid" (validOggCrcs final)

-- Helpers

parseSynthetic :: FilePath -> [BS.ByteString] -> IO Metadata
parseSynthetic name packets = parseSyntheticBytes name (mkOggFile packets)

parseSyntheticBytes :: FilePath -> BS.ByteString -> IO Metadata
parseSyntheticBytes name bytes = do
  path <- writeSyntheticBytes name bytes
  osPath <- encodeFS path
  result <- runExceptT $ parseOGG osPath
  removeFile path
  either (assertFailure . show) return result

writeSynthetic :: FilePath -> [BS.ByteString] -> IO FilePath
writeSynthetic name packets = writeSyntheticBytes name (mkOggFile packets)

writeSyntheticBytes :: FilePath -> BS.ByteString -> IO FilePath
writeSyntheticBytes name bytes = do
  tmpDir <- getTemporaryDirectory
  let path = tmpDir </> name
  BS.writeFile path bytes
  return path

-- | A final page with the given granule position (EOS, one small segment)
lastPage :: Int -> BS.ByteString
lastPage granule = BS.concat
  [ "OggS"
  , BS.singleton 0            -- version
  , BS.singleton 0x04         -- EOS
  , w64le granule
  , BS.replicate 4 0          -- serial
  , BS.replicate 4 0          -- page sequence
  , BS.replicate 4 0          -- CRC (not verified by the parser)
  , BS.singleton 1
  , BS.singleton 10
  , BS.replicate 10 0x33
  ]

-- | Verify the OGG CRC of every page in the stream
validOggCrcs :: BS.ByteString -> Bool
validOggCrcs = go
  where
    go bs
      | BS.null bs = True
      | BS.length bs < 27 || BS.take 4 bs /= "OggS" = False
      | otherwise =
          let numSegments = fromIntegral $ BS.index bs 26
              lacings = map fromIntegral $ BS.unpack $ BS.take numSegments $ BS.drop 27 bs
              pageLen = 27 + numSegments + sum lacings
              page = BS.take pageLen bs
              stored = BS.take 4 $ BS.drop 22 page
              zeroed = BS.take 22 page <> BS.replicate 4 0 <> BS.drop 26 page
          in BS.length page == pageLen
             && stored == w32le' (fromIntegral (oggCrc zeroed))
             && go (BS.drop pageLen bs)

    w32le' :: Int -> BS.ByteString
    w32le' n = BS.pack [fromIntegral (n `shiftR` s) | s <- [0, 8, 16, 24]]

    oggCrc :: BS.ByteString -> Word32
    oggCrc = BS.foldl' step 0
      where
        step crc byte = spin (8 :: Int) (crc `xor` (fromIntegral byte `shiftL` 24))
        spin 0 crc = crc
        spin n crc = spin (n - 1) $
          if testBit crc 31 then (crc `shiftL` 1) `xor` 0x04c11db7 else crc `shiftL` 1

-- | Serialize logical packets into an OGG page stream. Packets are split
-- into 255-byte lacing segments; pages hold at most 255 segments, and a
-- page starting mid-packet gets the continuation flag (0x01).
mkOggFile :: [BS.ByteString] -> BS.ByteString
mkOggFile packets = mkOggPageGroups [packets]

-- | Like mkOggFile, but each group of packets starts on a fresh page -
-- the spec-conformant layout the writer expects:
-- [[ident], [comment, setup], [audio]]
mkOggPageGroups :: [[BS.ByteString]] -> BS.ByteString
mkOggPageGroups groups = BS.concat (concatMap renderGroup groups)
  where
    renderGroup packets =
      renderPages False (chunksOf 255 (concatMap packetSegments packets))

    -- Segments of one packet: all 255 bytes except the final short one.
    -- A packet whose length is a multiple of 255 ends with an empty segment.
    packetSegments packet =
      let (chunk, rest) = BS.splitAt 255 packet
      in if BS.length chunk == 255
           then chunk : packetSegments rest
           else [chunk]

    chunksOf _ [] = []
    chunksOf n xs = let (h, t) = splitAt n xs in h : chunksOf n t

    renderPages _ [] = []
    renderPages continued (pageSegments : rest) =
      let nextContinued = case reverse pageSegments of
            (lastSegment:_) -> BS.length lastSegment == 255
            [] -> False
      in renderPage continued pageSegments : renderPages nextContinued rest

    renderPage continued pageSegments = BS.concat
      [ "OggS"
      , BS.singleton 0                                   -- version
      , BS.singleton (if continued then 0x01 else 0x00)  -- header type flags
      , BS.replicate 8 0                                 -- granule position
      , BS.replicate 4 0                                 -- serial number
      , BS.replicate 4 0                                 -- page sequence
      , BS.replicate 4 0                                 -- CRC (not verified)
      , BS.singleton (fromIntegral (length pageSegments))
      , BS.pack (map (fromIntegral . BS.length) pageSegments)
      , BS.concat pageSegments
      ]

vorbisIdentPacket :: BS.ByteString
vorbisIdentPacket = BS.concat
  [ "\x01vorbis"
  , w32le 0                 -- vorbis version
  , BS.singleton 2          -- channels
  , w32le 44100             -- sample rate
  , w32le 0                 -- bitrate maximum
  , w32le 128000            -- bitrate nominal
  , w32le 0                 -- bitrate minimum
  , BS.singleton 0xB8       -- blocksizes
  , BS.singleton 0x01       -- framing bit
  ]

opusIdentPacket :: BS.ByteString
opusIdentPacket = BS.concat
  [ "OpusHead"
  , BS.singleton 1          -- version
  , BS.singleton 2          -- channels
  , BS.pack [0x38, 0x01]    -- pre-skip
  , w32le 44100             -- input sample rate (not the decode rate)
  , BS.pack [0, 0]          -- output gain
  , BS.singleton 0          -- mapping family
  ]

vorbisCommentPacket :: [(Text, Text)] -> BS.ByteString
vorbisCommentPacket comments =
  "\x03vorbis" <> commentBody comments <> BS.singleton 0x01  -- framing bit

opusTagsPacket :: [(Text, Text)] -> BS.ByteString
opusTagsPacket comments = "OpusTags" <> commentBody comments

commentBody :: [(Text, Text)] -> BS.ByteString
commentBody comments = BS.concat $
  [ w32le 0                                    -- vendor string length
  , w32le (fromIntegral (length comments))
  ] ++ map renderComment comments
  where
    renderComment (key, value) =
      let bytes = TE.encodeUtf8 (key <> "=" <> value)
      in w32le (fromIntegral (BS.length bytes)) <> bytes

flacPictureBlock :: BS.ByteString -> BS.ByteString -> BS.ByteString
flacPictureBlock mimeType imageData = BS.concat
  [ w32be 3                                    -- picture type: front cover
  , w32be (fromIntegral (BS.length mimeType))
  , mimeType
  , w32be 0                                    -- description length
  , w32be 0, w32be 0, w32be 0, w32be 0         -- width/height/depth/colors
  , w32be (fromIntegral (BS.length imageData))
  , imageData
  ]

w32le :: Int -> BS.ByteString
w32le n = BS.pack [fromIntegral (n `div` (256 ^ i)) | i <- [0 :: Int, 1, 2, 3]]

w64le :: Int -> BS.ByteString
w64le n = BS.pack [fromIntegral (n `div` (256 ^ i)) | i <- [0 :: Int .. 7]]

w32be :: Int -> BS.ByteString
w32be n = BS.pack [fromIntegral (n `div` (256 ^ i)) | i <- [3, 2, 1, 0 :: Int]]