diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index dd3ea8e..7b98291 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -50,7 +50,7 @@ jobs: - name: golangci-lint uses: golangci/golangci-lint-action@v6 with: - version: latest + version: v1.64.8 vulncheck: runs-on: ubuntu-latest diff --git a/README.md b/README.md index 0b10c9d..5bbc66f 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,13 @@ err := ops.Reindex(ctx, "input.mkv", "output.mkv") `reader.ReadStream` parses metadata and returns a `*BlockReader` from an `io.Reader` without ever calling Seek. `writer.NewStreamWriter` writes a live MKV stream to an `io.Writer` using unknown-size Segment and Clusters. See [docs/library.md](docs/library.md) for details. +**Remux a file to WebM:** +```go +// Validates the codecs (VP8/VP9/AV1, Vorbis/Opus, WebVTT), copies the media +// verbatim into a webm-DocType container, rejects non-WebM codecs. +err := matroska.RemuxToWebM(ctx, "in.mkv", "out.webm") +``` + **Edit metadata with custom FS (S3, HTTP, etc.):** ```go s3fs := &matroska.FS{ @@ -178,12 +185,12 @@ err := matroska.EditMetadata(ctx, "s3://bucket/movie.mkv", "s3://bucket/out.mkv" cmd/mkvgo/ CLI binary commands/ one file per command group -matroska/ facade -- re-exports everything, backward compat +matroska/ facade -- stable public API, re-exports everything -mkv/ core types, FS port, EBML IDs +mkv/ core types, FS port, EBML IDs (experimental, may change) reader/ parse MKV → Container writer/ Container → MKV bytes - ops/ high-level operations (mux, split, merge, edit...) + ops/ high-level operations (mux, split, merge, edit, remux-webm...) subtitle/ SRT/ASS parsing ebml/ low-level EBML encoding/decoding (no Matroska knowledge) diff --git a/docs/cli.md b/docs/cli.md index 2527836..9d89f2e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -205,7 +205,7 @@ mkvgo edit -o - The JSON is a partial `Container` struct. Only fields you include are changed. ```bash -mkvgo edit movie.mkv -o out.mkv '{"title":"New Title"}}' +mkvgo edit movie.mkv -o out.mkv '{"title":"New Title"}' cat patch.json | mkvgo edit movie.mkv -o out.mkv - ``` @@ -255,7 +255,7 @@ mkvgo edit-inplace '' ``` ```bash -mkvgo edit-inplace movie.mkv '{"title":"Quick Fix"}}' +mkvgo edit-inplace movie.mkv '{"title":"Quick Fix"}' ``` ### remove-track diff --git a/docs/library.md b/docs/library.md index 57e56ac..8c7627d 100644 --- a/docs/library.md +++ b/docs/library.md @@ -12,7 +12,9 @@ | `matroska` | `github.com/gravity-zero/mkvgo/matroska` | Facade -- re-exports everything | | `ebml` | `github.com/gravity-zero/mkvgo/ebml` | Low-level EBML codec | -For most use cases, import `matroska` (the facade). Import sub-packages directly when you need fine-grained control. +`matroska` is the stable public API -- import it for most use cases. The `mkv`, `mkv/reader`, `mkv/writer`, `mkv/ops` and `mkv/subtitle` packages are lower-level and experimental: their APIs may change between minor versions. Import them directly when you need capabilities the facade does not expose (streaming, `NewWebMStreamWriter`). + +Operations process the container incrementally -- they read and write block by block (or cluster by cluster) and never hold the whole file in memory, so multi-gigabyte inputs run with bounded memory. --- @@ -71,6 +73,34 @@ err := writer.Write(&buf, container) --- +## WebM Output + +WebM is a constrained Matroska profile: the `webm` DocType and a small codec set (VP8/VP9/AV1 video, Vorbis/Opus audio, WebVTT subtitles). + +Check whether a `Container` can be written as WebM: + +```go +if err := matroska.ValidateWebM(container); err != nil { + // Names each track whose codec is outside the WebM subset, or which is + // missing mandatory init data (Opus OpusHead, Vorbis headers, AV1 av1C). + return err +} +``` + +`matroska.WriteWebM` writes the `webm` DocType (version 4 when an AV1 track is present, else 2) plus Info and Tracks. Like `writer.Write`, it writes metadata only -- no clusters. + +For a complete, playable WebM with frames, remux a source file: + +```go +err := matroska.RemuxToWebM(ctx, "in.mkv", "out.webm") +``` + +`RemuxToWebM` validates the codecs, copies every block verbatim into time-bounded `webm` clusters, and rejects sources with non-WebM codecs. Elements outside the WebM subset (Chapters, Attachments, Tags) are dropped; list them beforehand with `matroska.WebMNonSubsetElements(container)`. + +To write a WebM stream live (no source file), use `writer.NewWebMStreamWriter` -- the WebM counterpart of `NewStreamWriter`. + +--- + ## Mux / Demux **Mux** -- combine tracks from multiple sources: @@ -344,16 +374,16 @@ err := matroska.RemoveTrack(ctx, "in.mkv", "out.mkv", []uint64{3}, opts) ```go import "github.com/gravity-zero/mkvgo/mkv/subtitle" -entries, err := subtitle.ParseSRT(srtReader) -// []subtitle.SRTEntry{Index, StartMs, EndMs, Text} +entries, err := subtitle.ParseSRT("subs.srt") +// []subtitle.SRTEntry{StartMs, EndMs, Text} ``` ### ASS/SSA ```go -assFile, err := subtitle.ParseASS(assReader) -// assFile.ScriptInfo, assFile.Styles, assFile.Events -// Each event: Layer, Start, End, Style, Name, Text +assFile, err := subtitle.ParseASS("subs.ass") +// assFile.Header (raw [Script Info] + [V4+ Styles] block), assFile.Events +// Each subtitle.ASSEvent{StartMs, EndMs, Fields} ``` ### Extract from MKV @@ -382,6 +412,8 @@ err := matroska.MergeASS(ctx, "movie.mkv", "subs.ass", "out.mkv", "jpn", "Japane All functions return `error`. No panics, no logging. +The reader tolerates corrupted bodies: a zeroed or padded region between clusters (seen in some real-world rips) does not abort the read -- the parser resyncs to the next valid Cluster and returns the metadata gathered so far. A damaged EBML/Segment header still returns an error. Malformed input never panics. + ```go c, err := matroska.Open(ctx, path) if err != nil { diff --git a/ebml/primitives.go b/ebml/primitives.go index 8afcb13..ed3e93a 100644 --- a/ebml/primitives.go +++ b/ebml/primitives.go @@ -53,13 +53,36 @@ func ReadFloat(r io.Reader, size int64) (float64, error) { return math.Float64frombits(binary.BigEndian.Uint64(buf)), nil } -// ReadString reads a UTF-8/ASCII string, trimming trailing nulls. -func ReadString(r io.Reader, size int64) (string, error) { +// readExact reads exactly size bytes. For sizes above a small threshold it +// grows the buffer incrementally (via io.LimitReader) instead of allocating +// size bytes upfront, so a malformed element that declares a huge size but +// supplies little data cannot force a giant allocation (memory DoS). +func readExact(r io.Reader, size int64) ([]byte, error) { if err := checkSize(size); err != nil { - return "", err + return nil, err } - buf := make([]byte, size) - if _, err := io.ReadFull(r, buf); err != nil { + const maxUpfront = 1 << 20 // 1 MiB: allocate exactly for the common small case + if size <= maxUpfront { + buf := make([]byte, size) + if _, err := io.ReadFull(r, buf); err != nil { + return nil, err + } + return buf, nil + } + buf, err := io.ReadAll(io.LimitReader(r, size)) + if err != nil { + return nil, err + } + if int64(len(buf)) != size { + return nil, io.ErrUnexpectedEOF + } + return buf, nil +} + +// ReadString reads a UTF-8/ASCII string, trimming trailing nulls. +func ReadString(r io.Reader, size int64) (string, error) { + buf, err := readExact(r, size) + if err != nil { return "", err } for len(buf) > 0 && buf[len(buf)-1] == 0 { @@ -70,12 +93,5 @@ func ReadString(r io.Reader, size int64) (string, error) { // ReadBytes reads raw bytes. func ReadBytes(r io.Reader, size int64) ([]byte, error) { - if err := checkSize(size); err != nil { - return nil, err - } - buf := make([]byte, size) - if _, err := io.ReadFull(r, buf); err != nil { - return nil, err - } - return buf, nil + return readExact(r, size) } diff --git a/ebml/primitives_test.go b/ebml/primitives_test.go index 867e2c7..7ed877d 100644 --- a/ebml/primitives_test.go +++ b/ebml/primitives_test.go @@ -104,6 +104,30 @@ func TestReadBytes_ExceedsMaxSize(t *testing.T) { } } +// TestReadBytes_HugeDeclaredSizeNoAlloc guards against a memory DoS: a tiny +// input that declares a near-MaxElementSize element must fail with the data +// available, NOT allocate the declared size upfront. (If it over-allocated, this +// test would spike ~256 MB / be killed under the race detector's memory limits.) +func TestReadBytes_HugeDeclaredSizeNoAlloc(t *testing.T) { + r := bytes.NewReader([]byte{1, 2, 3, 4, 5}) + _, err := ReadBytes(r, 256*1024*1024) // 256 MB declared, 5 bytes available + if err == nil { + t.Fatal("expected error for truncated huge element") + } +} + +func TestCheckSizeBoundary(t *testing.T) { + if err := checkSize(MaxElementSize); err != nil { + t.Errorf("checkSize(MaxElementSize) = %v, want nil (cap is inclusive)", err) + } + if err := checkSize(MaxElementSize + 1); err == nil { + t.Error("checkSize(MaxElementSize+1) = nil, want error") + } + if err := checkSize(-1); err == nil { + t.Error("checkSize(-1) = nil, want error") + } +} + func TestReadBytes_NegativeSize(t *testing.T) { r := strings.NewReader("x") _, err := ReadBytes(r, -1) diff --git a/ebml/reader.go b/ebml/reader.go index 0e8010b..8bb5617 100644 --- a/ebml/reader.go +++ b/ebml/reader.go @@ -26,11 +26,11 @@ func ReadVINT(r io.Reader) (uint64, int, error) { val := uint64(b) if width > 1 { - rest := make([]byte, width-1) - if _, err := io.ReadFull(r, rest); err != nil { + var rest [7]byte // width is 1..8, so width-1 fits; avoids a heap alloc per VINT + if _, err := io.ReadFull(r, rest[:width-1]); err != nil { return 0, 0, err } - for _, rb := range rest { + for _, rb := range rest[:width-1] { val = (val << 8) | uint64(rb) } } @@ -44,6 +44,11 @@ func ReadElementID(r io.Reader) (uint32, int, error) { if err != nil { return 0, 0, err } + if n > 4 { + // EBML element IDs are 1-4 octets; a wider VINT would silently truncate + // into uint32. Reject it rather than corrupt the parse. + return 0, n, fmt.Errorf("invalid element ID: %d-octet VINT exceeds 4-octet limit", n) + } return uint32(val), n, nil } diff --git a/ebml/reader_test.go b/ebml/reader_test.go index 2ec3bd0..4fb0272 100644 --- a/ebml/reader_test.go +++ b/ebml/reader_test.go @@ -102,3 +102,17 @@ func TestReadVINT_MultiByteValues(t *testing.T) { } } } + +func TestReadElementID_RejectsWideVINT(t *testing.T) { + // A 5-octet VINT (first set bit at position 3 -> width 5) is not a valid + // element ID and would silently truncate into uint32. + _, _, err := ReadElementID(bytes.NewReader([]byte{0x08, 0x00, 0x00, 0x00, 0x01})) + if err == nil { + t.Fatal("expected error for 5-octet element ID") + } + // A valid 4-octet ID (EBML header) must still parse. + id, n, err := ReadElementID(bytes.NewReader([]byte{0x1A, 0x45, 0xDF, 0xA3})) + if err != nil || n != 4 || id != IDEBMLHeader { + t.Fatalf("4-octet ID: id=0x%X n=%d err=%v", id, n, err) + } +} diff --git a/matroska/matroska.go b/matroska/matroska.go index 77c9742..42a7305 100644 --- a/matroska/matroska.go +++ b/matroska/matroska.go @@ -1,5 +1,12 @@ -// Package matroska provides backward-compatible access to the mkvgo toolkit. -// New code should import mkv, mkv/reader, mkv/writer, mkv/ops, mkv/subtitle directly. +// Package matroska is the stable, supported public API of mkvgo: a small, +// curated facade over the lower-level building blocks in mkv and its +// subpackages. Prefer it for application code — its exported surface is the one +// kept backward-compatible. +// +// The mkv, mkv/reader, mkv/writer, mkv/ops and mkv/subtitle packages are +// lower-level and EXPERIMENTAL: their APIs may change between minor versions. +// Import them directly only for capabilities this facade does not expose yet +// (e.g. streaming readers/writers, NewWebMStreamWriter). package matroska import ( @@ -79,18 +86,71 @@ func Write(w io.Writer, c *Container) error { return writer.Write(w, c) } +// WriteWebM writes c as a WebM file (validates WebM codec compatibility, then +// writes with the "webm" DocType). See mkv.ValidateWebM. +func WriteWebM(w io.Writer, c *Container) error { + return writer.WriteWebM(w, c) +} + +// ValidateWebM reports whether c can be written as WebM, naming any track whose +// codec falls outside the WebM subset (VP8/VP9/AV1, Vorbis/Opus, WebVTT). +func ValidateWebM(c *Container) error { + return mkv.ValidateWebM(c) +} + +// IsWebMCodec reports whether a codec (short name "vp9" or Matroska id "V_VP9") +// is permitted in WebM. +func IsWebMCodec(codec string) bool { + return mkv.IsWebMCodec(codec) +} + +// WebMDocTypeVersion returns the EBML DocTypeVersion needed for c as WebM +// (4 when an AV1 track is present, else 2). +func WebMDocTypeVersion(c *Container) uint64 { + return mkv.WebMDocTypeVersion(c) +} + +// RemuxToWebM reads srcPath and writes a complete, playable WebM file to +// dstPath, copying the media verbatim. Rejects sources with non-WebM codecs. +// Non-subset elements (Chapters/Attachments/Tags) are dropped — see +// WebMNonSubsetElements to detect that loss beforehand. +func RemuxToWebM(ctx context.Context, srcPath, dstPath string, opts ...Options) error { + return ops.RemuxToWebM(ctx, srcPath, dstPath, opts...) +} + +// WebMNonSubsetElements lists the elements in c (Chapters/Attachments/Tags) that +// a WebM remux will drop; empty means nothing is lost. +func WebMNonSubsetElements(c *Container) []string { + return mkv.WebMNonSubsetElements(c) +} + // --- Operations --- -func Mux(ctx context.Context, opts MuxOptions) error { return ops.Mux(ctx, opts) } -func Demux(ctx context.Context, opts DemuxOptions) error { return ops.Demux(ctx, opts) } -func Split(ctx context.Context, opts SplitOptions) ([]string, error) { return ops.Split(ctx, opts) } -func Join(ctx context.Context, sources []string, dstPath string) error { - return ops.Join(ctx, sources, dstPath) +func Mux(ctx context.Context, opts MuxOptions, extra ...Options) error { + return ops.Mux(ctx, opts, extra...) +} +func Demux(ctx context.Context, opts DemuxOptions, extra ...Options) error { + return ops.Demux(ctx, opts, extra...) +} +func Split(ctx context.Context, opts SplitOptions, extra ...Options) ([]string, error) { + return ops.Split(ctx, opts, extra...) +} +func Join(ctx context.Context, sources []string, dstPath string, opts ...Options) error { + return ops.Join(ctx, sources, dstPath, opts...) } -func Merge(ctx context.Context, opts MergeOptions) error { return ops.Merge(ctx, opts) } -func Validate(ctx context.Context, path string) ([]Issue, error) { return ops.Validate(ctx, path) } -func Compare(ctx context.Context, pathA, pathB string) ([]Diff, error) { - return ops.Compare(ctx, pathA, pathB) +func Merge(ctx context.Context, opts MergeOptions, extra ...Options) error { + return ops.Merge(ctx, opts, extra...) +} +func Validate(ctx context.Context, path string, opts ...Options) ([]Issue, error) { + return ops.Validate(ctx, path, opts...) +} +func Compare(ctx context.Context, pathA, pathB string, opts ...Options) ([]Diff, error) { + return ops.Compare(ctx, pathA, pathB, opts...) +} + +// Reindex rebuilds the seek index (Cues) of a file. See ops.Reindex. +func Reindex(ctx context.Context, srcPath, dstPath string, opts ...Options) error { + return ops.Reindex(ctx, srcPath, dstPath, opts...) } func RemoveTrack(ctx context.Context, srcPath, dstPath string, removeIDs []uint64, opts ...Options) error { diff --git a/mkv/doc.go b/mkv/doc.go new file mode 100644 index 0000000..70249a2 --- /dev/null +++ b/mkv/doc.go @@ -0,0 +1,10 @@ +// Package mkv holds the core Matroska model (Container, Track, Block, …) shared +// by its subpackages reader, writer, ops and subtitle, which implement the +// parsing, serialisation and high-level operations. +// +// STABILITY: mkv and its subpackages are lower-level building blocks and are +// considered EXPERIMENTAL — their exported APIs may change between minor +// versions. For a stable, backward-compatible surface, use the top-level +// matroska package; reach for these packages directly only when you need +// capabilities matroska does not expose (streaming, custom operations). +package mkv diff --git a/mkv/ops/demux.go b/mkv/ops/demux.go index dbe0cf7..d7ea8fd 100644 --- a/mkv/ops/demux.go +++ b/mkv/ops/demux.go @@ -1,6 +1,7 @@ package ops import ( + "bufio" "context" "fmt" "io" @@ -67,6 +68,13 @@ func Demux(ctx context.Context, opts mkv.DemuxOptions, extra ...mkv.Options) err return fmt.Errorf("write track %d: %w", blk.TrackNumber, err) } } + // Flush buffered writers so a final write error (e.g. disk full) is surfaced + // rather than swallowed by the deferred Close. + for id, w := range writers { + if err := w.Flush(); err != nil { + return fmt.Errorf("flush track %d: %w", id, err) + } + } return nil } @@ -90,8 +98,8 @@ func buildTrackSet(c *mkv.Container, trackIDs []uint64) map[uint64]mkv.Track { return m } -func openOutputFiles(tracks map[uint64]mkv.Track, dir string, fs *mkv.FS) (map[uint64]io.Writer, []io.Closer, error) { - writers := make(map[uint64]io.Writer, len(tracks)) +func openOutputFiles(tracks map[uint64]mkv.Track, dir string, fs *mkv.FS) (map[uint64]*bufio.Writer, []io.Closer, error) { + writers := make(map[uint64]*bufio.Writer, len(tracks)) var closers []io.Closer for id, t := range tracks { ext := sanitizeCodec(t.Codec) @@ -110,7 +118,7 @@ func openOutputFiles(tracks map[uint64]mkv.Track, dir string, fs *mkv.FS) (map[u } return nil, nil, err } - writers[id] = f + writers[id] = bufio.NewWriterSize(f, 64<<10) // batch block writes closers = append(closers, f) } return writers, closers, nil diff --git a/mkv/ops/edit.go b/mkv/ops/edit.go index 44200c6..2f52aab 100644 --- a/mkv/ops/edit.go +++ b/mkv/ops/edit.go @@ -3,7 +3,6 @@ package ops import ( "context" "fmt" - "io" "github.com/gravity-zero/mkvgo/mkv" "github.com/gravity-zero/mkvgo/mkv/reader" @@ -84,18 +83,6 @@ func AddTrack(ctx context.Context, srcPath, dstPath string, input mkv.TrackInput newID := uint64(len(c.Tracks) + 1) remap := identityRemap(c.Tracks) - blocks, err := readFilteredBlocks(ctx, srcPath, c.Info.TimecodeScale, remap, fs) - if err != nil { - return err - } - - addRemap := map[uint64]uint64{input.TrackID: newID} - addBlocks, err := readFilteredBlocks(ctx, input.SourcePath, srcAdd.Info.TimecodeScale, addRemap, fs) - if err != nil { - return err - } - blocks = mergeBlocks(blocks, addBlocks) - t := *addedTrack t.ID = newID if input.Language != "" { @@ -107,9 +94,9 @@ func AddTrack(ctx context.Context, srcPath, dstPath string, input mkv.TrackInput t.IsDefault = input.IsDefault tracks := append(c.Tracks, t) - var durationMs int64 - if len(blocks) > 0 { - durationMs = blocks[len(blocks)-1].Timecode + durationMs := c.DurationMs + if srcAdd.DurationMs > durationMs { + durationMs = srcAdd.DurationMs } out, err := fs.DoCreate(dstPath) @@ -125,7 +112,11 @@ func AddTrack(ctx context.Context, srcPath, dstPath string, input mkv.TrackInput if err := mw.WriteMetadata(c, tracks, durationMs); err != nil { return err } - if err := writeBlocksAsClusters(mw, blocks, c.Info.TimecodeScale); err != nil { + sources := []mergeSource{ + {path: srcPath, scale: c.Info.TimecodeScale, remap: remap}, + {path: input.SourcePath, scale: srcAdd.Info.TimecodeScale, remap: map[uint64]uint64{input.TrackID: newID}}, + } + if err := streamMergeToWriter(ctx, mw, c.Info.TimecodeScale, fs, sources); err != nil { return err } return mw.Finalize() @@ -196,78 +187,3 @@ func ExtractAttachment(ctx context.Context, srcPath string, attachID uint64, out } return fmt.Errorf("attachment %d not found", attachID) } - -func writeBlocksAsClusters(mw *writer.MKVWriter, blocks []mkv.Block, timecodeScale int64) error { - if len(blocks) == 0 { - return nil - } - var cluster []mkv.Block - clusterTS := blocks[0].Timecode - - for i := range blocks { - b := &blocks[i] - if b.Timecode-clusterTS >= defaultClusterDurationMs && len(cluster) > 0 { - if err := mw.WriteClusterWithCues(clusterTS, timecodeScale, cluster); err != nil { - return err - } - cluster = cluster[:0] - clusterTS = b.Timecode - } - cluster = append(cluster, *b) - } - if len(cluster) > 0 { - return mw.WriteClusterWithCues(clusterTS, timecodeScale, cluster) - } - return nil -} - -func readFilteredBlocks(ctx context.Context, path string, timecodeScale int64, remap map[uint64]uint64, fs *mkv.FS) ([]mkv.Block, error) { - f, err := fs.DoOpen(path) - if err != nil { - return nil, err - } - defer f.Close() - - br, err := reader.NewBlockReader(f, timecodeScale) - if err != nil { - return nil, err - } - - var blocks []mkv.Block - for { - if ctx.Err() != nil { - return nil, ctx.Err() - } - blk, err := br.Next() - if err == io.EOF { - break - } - if err != nil { - return nil, err - } - newID, ok := remap[blk.TrackNumber] - if !ok { - continue - } - blk.TrackNumber = newID - blocks = append(blocks, blk) - } - return blocks, nil -} - -func mergeBlocks(a, b []mkv.Block) []mkv.Block { - merged := make([]mkv.Block, 0, len(a)+len(b)) - i, j := 0, 0 - for i < len(a) && j < len(b) { - if a[i].Timecode <= b[j].Timecode { - merged = append(merged, a[i]) - i++ - } else { - merged = append(merged, b[j]) - j++ - } - } - merged = append(merged, a[i:]...) - merged = append(merged, b[j:]...) - return merged -} diff --git a/mkv/ops/inplace.go b/mkv/ops/inplace.go index 8e0f875..9eb968b 100644 --- a/mkv/ops/inplace.go +++ b/mkv/ops/inplace.go @@ -14,6 +14,11 @@ import ( "github.com/gravity-zero/mkvgo/mkv/writer" ) +// EditInPlace rewrites only the metadata region of path on disk (instant, no +// cluster copy). It writes directly over the source via Seek+Write, so it is +// NOT crash-safe: a crash mid-write can corrupt the file, and it fails if the +// new metadata does not fit the existing region. For precious or untrusted data, +// prefer EditMetadata, which writes a fresh file the caller can atomically rename. func EditInPlace(ctx context.Context, path string, edit func(*mkv.Container), opts ...mkv.Options) error { fs := mkv.FSFrom(opts) c, err := reader.OpenWithFS(ctx, path, fs) @@ -84,6 +89,14 @@ func EditInPlace(ctx context.Context, path string, edit func(*mkv.Container), op } } + // Flush to stable storage. NOTE: the write above is in-place and NOT atomic — + // a crash mid-write can corrupt the source file. Use EditMetadata (full + // rewrite to a new file) when crash safety matters. + if s, ok := f.(interface{ Sync() error }); ok { + if err := s.Sync(); err != nil { + return fmt.Errorf("sync: %w", err) + } + } return nil } @@ -142,6 +155,9 @@ func findMetadataRegion(path string, fs *mkv.FS) (metadataRegion, error) { switch eh.ID { case mkv.IDInfo, mkv.IDTracks, mkv.IDChapters, mkv.IDAttachments, mkv.IDTags, mkv.IDSeekHead, mkv.IDVoid: + if eh.Size < 0 { // unknown-size metadata: a negative Seek would corrupt the scan + return metadataRegion{}, fmt.Errorf("unknown-size metadata element 0x%X", eh.ID) + } if region.start < 0 { region.start = pos } diff --git a/mkv/ops/join.go b/mkv/ops/join.go index 673b7c5..991be33 100644 --- a/mkv/ops/join.go +++ b/mkv/ops/join.go @@ -48,7 +48,10 @@ func Join(ctx context.Context, sources []string, dstPath string, opts ...mkv.Opt var totalDurationMs int64 for _, src := range sources { - c, _ := reader.OpenWithFS(ctx, src, fs) + c, err := reader.OpenWithFS(ctx, src, fs) + if err != nil { + return err + } totalDurationMs += c.DurationMs } diff --git a/mkv/ops/legacy_blocks_test.go b/mkv/ops/legacy_blocks_test.go new file mode 100644 index 0000000..f61482e --- /dev/null +++ b/mkv/ops/legacy_blocks_test.go @@ -0,0 +1,90 @@ +package ops + +// Test-only block helpers. Production code uses streamToWriter / +// streamMergeToWriter (stream.go); these load all blocks in memory and are kept +// only because their unit tests still exercise the cluster-batching, merge and +// filter logic directly. + +import ( + "context" + "io" + + "github.com/gravity-zero/mkvgo/mkv" + "github.com/gravity-zero/mkvgo/mkv/reader" + "github.com/gravity-zero/mkvgo/mkv/writer" +) + +func writeBlocksAsClusters(mw *writer.MKVWriter, blocks []mkv.Block, timecodeScale int64) error { + if len(blocks) == 0 { + return nil + } + var cluster []mkv.Block + clusterTS := blocks[0].Timecode + + for i := range blocks { + b := &blocks[i] + if b.Timecode-clusterTS >= defaultClusterDurationMs && len(cluster) > 0 { + if err := mw.WriteClusterWithCues(clusterTS, timecodeScale, cluster); err != nil { + return err + } + cluster = cluster[:0] + clusterTS = b.Timecode + } + cluster = append(cluster, *b) + } + if len(cluster) > 0 { + return mw.WriteClusterWithCues(clusterTS, timecodeScale, cluster) + } + return nil +} + +func readFilteredBlocks(ctx context.Context, path string, timecodeScale int64, remap map[uint64]uint64, fs *mkv.FS) ([]mkv.Block, error) { + f, err := fs.DoOpen(path) + if err != nil { + return nil, err + } + defer f.Close() + + br, err := reader.NewBlockReader(f, timecodeScale) + if err != nil { + return nil, err + } + + var blocks []mkv.Block + for { + if ctx.Err() != nil { + return nil, ctx.Err() + } + blk, err := br.Next() + if err == io.EOF { + break + } + if err != nil { + return nil, err + } + newID, ok := remap[blk.TrackNumber] + if !ok { + continue + } + blk.TrackNumber = newID + blocks = append(blocks, blk) + } + return blocks, nil +} + +func mergeBlocks(a, b []mkv.Block) []mkv.Block { + merged := make([]mkv.Block, 0, len(a)+len(b)) + i, j := 0, 0 + for i < len(a) && j < len(b) { + if a[i].Timecode <= b[j].Timecode { + merged = append(merged, a[i]) + i++ + } else { + merged = append(merged, b[j]) + j++ + } + } + merged = append(merged, a[i:]...) + merged = append(merged, b[j:]...) + return merged +} diff --git a/mkv/ops/merge.go b/mkv/ops/merge.go index 7a6f397..9930072 100644 --- a/mkv/ops/merge.go +++ b/mkv/ops/merge.go @@ -46,7 +46,10 @@ func Merge(ctx context.Context, opts mkv.MergeOptions, extra ...mkv.Options) err return fmt.Errorf("no tracks selected") } - first, _ := reader.OpenWithFS(ctx, opts.Inputs[0].SourcePath, fs) + first, err := reader.OpenWithFS(ctx, opts.Inputs[0].SourcePath, fs) + if err != nil { + return err + } muxOpts := mkv.MuxOptions{ OutputPath: opts.OutputPath, Tracks: trackInputs, diff --git a/mkv/ops/merge_stream_test.go b/mkv/ops/merge_stream_test.go new file mode 100644 index 0000000..40142ff --- /dev/null +++ b/mkv/ops/merge_stream_test.go @@ -0,0 +1,145 @@ +package ops + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "testing" + + "github.com/gravity-zero/mkvgo/mkv" + "github.com/gravity-zero/mkvgo/mkv/reader" + "github.com/gravity-zero/mkvgo/mkv/writer" +) + +// writeTrackMKV writes a seekable single-track MKV with one block per timecode, +// grouped into <=1s clusters so relative timecodes stay in int16 range. +func writeTrackMKV(t *testing.T, path, codec string, timecodes []int64) { + t.Helper() + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: codec}} + + var seg bytes.Buffer + mustNil(t, writer.WriteSegmentInfo(&seg, &info, 0)) + mustNil(t, writer.WriteTracks(&seg, tracks)) + for i := 0; i < len(timecodes); { + base := timecodes[i] + var blocks []mkv.Block + for i < len(timecodes) && timecodes[i]-base < 1000 { + blocks = append(blocks, mkv.Block{ + TrackNumber: 1, Timecode: timecodes[i], Keyframe: true, + Data: []byte{byte(timecodes[i])}, + }) + i++ + } + mustNil(t, writer.WriteCluster(&seg, base, info.TimecodeScale, blocks)) + } + + var buf bytes.Buffer + mustNil(t, writer.WriteEBMLHeader(&buf)) + mustNil(t, writer.WriteMasterElement(&buf, mkv.IDSegment, seg.Bytes())) + mustNil(t, os.WriteFile(path, buf.Bytes(), 0o644)) +} + +func collectAllBlocks(t *testing.T, path string) []mkv.Block { + t.Helper() + c, err := reader.Open(context.Background(), path) + mustNil(t, err) + f, err := os.Open(path) + mustNil(t, err) + defer f.Close() + br, err := reader.NewBlockReader(f, c.Info.TimecodeScale) + mustNil(t, err) + var out []mkv.Block + for { + b, err := br.Next() + if errors.Is(err, io.EOF) { + break + } + mustNil(t, err) + out = append(out, b) + } + return out +} + +func assertMonotonic(t *testing.T, blocks []mkv.Block) { + t.Helper() + for i := 1; i < len(blocks); i++ { + if blocks[i].Timecode < blocks[i-1].Timecode { + t.Errorf("blocks not timecode-ordered at %d: %d after %d", i, blocks[i].Timecode, blocks[i-1].Timecode) + } + } +} + +// TestMuxStreamMergePreservesAndInterleaves checks the streaming k-way merge: +// every block from every source survives, tracks are remapped, and the output is +// ordered by timecode across sources -- including across cluster boundaries. +func TestMuxStreamMergePreservesAndInterleaves(t *testing.T) { + dir := t.TempDir() + srcA := filepath.Join(dir, "a.mkv") + srcB := filepath.Join(dir, "b.mkv") + // interleaved across sources and spanning multiple 1s clusters. + writeTrackMKV(t, srcA, "vp9", []int64{0, 100, 1200, 2400}) + writeTrackMKV(t, srcB, "opus", []int64{50, 1100, 1300, 2000}) + dst := filepath.Join(dir, "muxed.mkv") + + mustNil(t, Mux(context.Background(), mkv.MuxOptions{ + OutputPath: dst, + Tracks: []mkv.TrackInput{ + {SourcePath: srcA, TrackID: 1}, + {SourcePath: srcB, TrackID: 1}, + }, + })) + + blocks := collectAllBlocks(t, dst) + + if len(blocks) != 8 { + t.Fatalf("got %d blocks, want 8 (all preserved across both sources)", len(blocks)) + } + counts := map[uint64]int{} + for _, b := range blocks { + counts[b.TrackNumber]++ + } + if counts[1] != 4 || counts[2] != 4 { + t.Errorf("per-track counts = %v, want map[1:4 2:4]", counts) + } + assertMonotonic(t, blocks) + + want := []int64{0, 50, 100, 1100, 1200, 1300, 2000, 2400} + for i, b := range blocks { + if b.Timecode != want[i] { + t.Errorf("block %d timecode = %d, want %d (merge order)", i, b.Timecode, want[i]) + } + } +} + +// TestAddTrackStreamMergeInterleaves checks that AddTrack (the 2-source variant) +// interleaves the added track's blocks with the base file's by timecode and +// keeps every block. +func TestAddTrackStreamMergeInterleaves(t *testing.T) { + dir := t.TempDir() + base := filepath.Join(dir, "base.mkv") + add := filepath.Join(dir, "add.mkv") + writeTrackMKV(t, base, "vp9", []int64{0, 100, 200}) + writeTrackMKV(t, add, "opus", []int64{50, 150, 250}) + dst := filepath.Join(dir, "out.mkv") + + mustNil(t, AddTrack(context.Background(), base, dst, mkv.TrackInput{ + SourcePath: add, TrackID: 1, Language: "eng", + })) + + blocks := collectAllBlocks(t, dst) + if len(blocks) != 6 { + t.Fatalf("got %d blocks, want 6 (base + added)", len(blocks)) + } + counts := map[uint64]int{} + for _, b := range blocks { + counts[b.TrackNumber]++ + } + if counts[1] != 3 || counts[2] != 3 { + t.Errorf("per-track counts = %v, want map[1:3 2:3]", counts) + } + assertMonotonic(t, blocks) +} diff --git a/mkv/ops/mux.go b/mkv/ops/mux.go index fcdb041..a265230 100644 --- a/mkv/ops/mux.go +++ b/mkv/ops/mux.go @@ -3,37 +3,36 @@ package ops import ( "context" "fmt" - "io" - "sort" "github.com/gravity-zero/mkvgo/mkv" "github.com/gravity-zero/mkvgo/mkv/reader" "github.com/gravity-zero/mkvgo/mkv/writer" ) -func Mux(ctx context.Context, opts mkv.MuxOptions, extra ...mkv.Options) error { +func Mux(ctx context.Context, opts mkv.MuxOptions, extra ...mkv.Options) (err error) { fs := mkv.FSFrom(extra) out, err := fs.DoCreate(opts.OutputPath) if err != nil { return fmt.Errorf("create output: %w", err) } - defer out.Close() + // Surface a Close error on the success path (e.g. a custom FS that finalises + // the write on Close) instead of silently dropping it. + defer func() { + if cerr := out.Close(); cerr != nil && err == nil { + err = cerr + } + }() tracks, trackMap, err := buildMuxTracks(ctx, opts.Tracks, fs) if err != nil { return err } - blocks, timecodeScale, err := collectBlocks(ctx, opts.Tracks, trackMap, fs) + sources, timecodeScale, durationMs, err := buildMuxSources(ctx, opts.Tracks, trackMap, fs) if err != nil { return err } - var durationMs int64 - if len(blocks) > 0 { - durationMs = blocks[len(blocks)-1].Timecode - } - c := &mkv.Container{ Info: mkv.SegmentInfo{ TimecodeScale: timecodeScale, @@ -53,7 +52,7 @@ func Mux(ctx context.Context, opts mkv.MuxOptions, extra ...mkv.Options) error { if err := mw.WriteMetadata(c, tracks, durationMs); err != nil { return err } - if err := writeBlocksAsClusters(mw, blocks, timecodeScale); err != nil { + if err := streamMergeToWriter(ctx, mw, timecodeScale, fs, sources); err != nil { return err } return mw.Finalize() @@ -102,74 +101,40 @@ type trackKey struct { trackID uint64 } -func collectBlocks(ctx context.Context, inputs []mkv.TrackInput, trackMap map[trackKey]uint64, fs *mkv.FS) ([]mkv.Block, int64, error) { - type sourceReq struct { - path string - wantTracks map[uint64]uint64 - } - sourceMap := make(map[string]*sourceReq) +// buildMuxSources groups the track inputs by source file (each becomes one +// mergeSource with a track remap) and returns them along with the output +// TimecodeScale and duration -- read from source metadata, so no blocks are +// loaded. streamMergeToWriter then merges them with bounded memory. +func buildMuxSources(ctx context.Context, inputs []mkv.TrackInput, trackMap map[trackKey]uint64, fs *mkv.FS) ([]mergeSource, int64, int64, error) { + order := make([]string, 0, len(inputs)) + remaps := make(map[string]map[uint64]uint64) for _, inp := range inputs { - sr, ok := sourceMap[inp.SourcePath] + rm, ok := remaps[inp.SourcePath] if !ok { - sr = &sourceReq{path: inp.SourcePath, wantTracks: make(map[uint64]uint64)} - sourceMap[inp.SourcePath] = sr + rm = make(map[uint64]uint64) + remaps[inp.SourcePath] = rm + order = append(order, inp.SourcePath) } - newID := trackMap[trackKey{inp.SourcePath, inp.TrackID}] - sr.wantTracks[inp.TrackID] = newID + rm[inp.TrackID] = trackMap[trackKey{inp.SourcePath, inp.TrackID}] } - var allBlocks []mkv.Block - var timecodeScale int64 - - for _, sr := range sourceMap { - c, err := reader.OpenWithFS(ctx, sr.path, fs) + var timecodeScale, durationMs int64 + sources := make([]mergeSource, 0, len(order)) + for _, path := range order { + c, err := reader.OpenWithFS(ctx, path, fs) if err != nil { - return nil, 0, err + return nil, 0, 0, err } if timecodeScale == 0 { timecodeScale = c.Info.TimecodeScale } - - f, err := fs.DoOpen(sr.path) - if err != nil { - return nil, 0, err - } - - br, err := reader.NewBlockReader(f, c.Info.TimecodeScale) - if err != nil { - f.Close() - return nil, 0, err + if c.DurationMs > durationMs { + durationMs = c.DurationMs } - - for { - if ctx.Err() != nil { - f.Close() - return nil, 0, ctx.Err() - } - blk, err := br.Next() - if err == io.EOF { - break - } - if err != nil { - f.Close() - return nil, 0, err - } - newID, ok := sr.wantTracks[blk.TrackNumber] - if !ok { - continue - } - blk.TrackNumber = newID - allBlocks = append(allBlocks, blk) - } - f.Close() + sources = append(sources, mergeSource{path: path, scale: c.Info.TimecodeScale, remap: remaps[path]}) } - - sort.Slice(allBlocks, func(i, j int) bool { - return allBlocks[i].Timecode < allBlocks[j].Timecode - }) - if timecodeScale == 0 { timecodeScale = 1000000 } - return allBlocks, timecodeScale, nil + return sources, timecodeScale, durationMs, nil } diff --git a/mkv/ops/reindex.go b/mkv/ops/reindex.go index 3c194c4..80df5ca 100644 --- a/mkv/ops/reindex.go +++ b/mkv/ops/reindex.go @@ -390,6 +390,9 @@ func Reindex(ctx context.Context, srcPath, dstPath string, opts ...mkv.Options) if ebmlHdr.Size < 0 { return fmt.Errorf("reindex: EBML header has unknown size") } + if ebmlHdr.Size > maxReindexClusterSize { + return fmt.Errorf("reindex: EBML header size %d exceeds limit (%d bytes)", ebmlHdr.Size, maxReindexClusterSize) + } ebmlBody := make([]byte, ebmlHdr.Size) if _, err := io.ReadFull(r, ebmlBody); err != nil { return fmt.Errorf("reindex: read EBML body: %w", err) @@ -475,6 +478,9 @@ func Reindex(ctx context.Context, srcPath, dstPath string, opts ...mkv.Options) if h.Size < 0 { return fmt.Errorf("reindex: unknown-size metadata element 0x%X", h.ID) } + if h.Size > maxReindexClusterSize { + return fmt.Errorf("reindex: metadata element 0x%X size %d exceeds limit (%d bytes)", h.ID, h.Size, maxReindexClusterSize) + } // Buffer the body so we can (a) scan Info for timecodeScale and // (b) write it out verbatim in one step. metaBuf := make([]byte, h.Size) diff --git a/mkv/ops/reindex_dos_test.go b/mkv/ops/reindex_dos_test.go new file mode 100644 index 0000000..cea3cef --- /dev/null +++ b/mkv/ops/reindex_dos_test.go @@ -0,0 +1,65 @@ +package ops + +import ( + "bytes" + "context" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/gravity-zero/mkvgo/ebml" +) + +// TestReindexRejectsHugeEBMLHeader is the exact attack the audit flagged: a +// ~6-byte file that declares a 10 GiB EBML header. Reindex must reject it via +// the size cap instead of make([]byte, 10GiB) (allocation bomb / OOM). +func TestReindexRejectsHugeEBMLHeader(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "evil.mkv") + + var buf bytes.Buffer + ebml.WriteElementID(&buf, ebml.IDEBMLHeader) + ebml.WriteDataSize(&buf, 10<<30) // declares 10 GiB; no body follows + if err := os.WriteFile(src, buf.Bytes(), 0o644); err != nil { + t.Fatal(err) + } + + err := Reindex(context.Background(), src, filepath.Join(dir, "out.mkv")) + if err == nil { + t.Fatal("Reindex accepted a 10 GiB EBML header from a tiny file (would OOM)") + } + if !strings.Contains(err.Error(), "exceeds limit") { + t.Fatalf("expected the size-cap error, got: %v", err) + } +} + +// TestReindexRejectsHugeMetadataElement checks the metadata-buffer cap (reindex +// reads Info/Tracks/... bodies into memory): a tiny file declaring a huge Info +// element must be rejected, not allocated. +func TestReindexRejectsHugeMetadataElement(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "evil2.mkv") + + var seg bytes.Buffer + // An Info element declaring 10 GiB, no body. + ebml.WriteElementID(&seg, 0x1549A966) // IDInfo + ebml.WriteDataSize(&seg, 10<<30) + + var buf bytes.Buffer + ebml.WriteElementHeader(&buf, ebml.IDEBMLHeader, 0) + ebml.WriteElementID(&buf, 0x18538067) // IDSegment + ebml.WriteDataSize(&buf, int64(seg.Len())) + buf.Write(seg.Bytes()) + if err := os.WriteFile(src, buf.Bytes(), 0o644); err != nil { + t.Fatal(err) + } + + err := Reindex(context.Background(), src, filepath.Join(dir, "out2.mkv")) + if err == nil { + t.Fatal("Reindex accepted a 10 GiB Info element (would OOM)") + } + if !strings.Contains(err.Error(), "exceeds limit") { + t.Fatalf("expected the size-cap error, got: %v", err) + } +} diff --git a/mkv/ops/stream.go b/mkv/ops/stream.go index 49533a4..5c2472c 100644 --- a/mkv/ops/stream.go +++ b/mkv/ops/stream.go @@ -122,3 +122,119 @@ func identityRemap(tracks []mkv.Track) map[uint64]uint64 { } return remap } + +// mergeSource is one input to streamMergeToWriter: a file, its TimecodeScale, +// and the map of its source track numbers to the output track numbers to keep. +type mergeSource struct { + path string + scale int64 + remap map[uint64]uint64 +} + +// streamMergeToWriter k-way merges the wanted blocks of several sources by +// timecode and writes them as time-bounded clusters. It holds only one block +// per source plus the current cluster in memory -- bounded regardless of file +// size, so it is safe for very large inputs (unlike collecting + sorting every +// block). Block.Timecode is in milliseconds for every source, so cross-source +// comparison needs no scale normalisation. +func streamMergeToWriter(ctx context.Context, mw *writer.MKVWriter, outScale int64, fs *mkv.FS, sources []mergeSource) error { + type state struct { + br *reader.BlockReader + f mkv.ReadSeekCloser + head mkv.Block + ok bool + } + states := make([]*state, len(sources)) + defer func() { + for _, s := range states { + if s != nil && s.f != nil { + s.f.Close() + } + } + }() + + advance := func(i int) error { + st := states[i] + for { + blk, err := st.br.Next() + if err == io.EOF { + st.ok = false + return nil + } + if err != nil { + return err + } + newID, ok := sources[i].remap[blk.TrackNumber] + if !ok { + continue + } + blk.TrackNumber = newID + st.head, st.ok = blk, true + return nil + } + } + + for i, src := range sources { + f, err := fs.DoOpen(src.path) + if err != nil { + return err + } + br, err := reader.NewBlockReader(f, src.scale) + if err != nil { + f.Close() + return err + } + states[i] = &state{br: br, f: f} + if err := advance(i); err != nil { + return err + } + } + + var cluster []mkv.Block + clusterTS := int64(-1) + flush := func() error { + if len(cluster) == 0 { + return nil + } + err := mw.WriteClusterWithCues(clusterTS, outScale, cluster) + cluster = cluster[:0] + return err + } + + for { + if ctx.Err() != nil { + return ctx.Err() + } + // Pick the source whose head block has the smallest timecode. Each + // source is read in file order; for the common single-track-per-source + // case that order is already sorted, so this reproduces a global sort. + mi := -1 + for i, st := range states { + if !st.ok { + continue + } + if mi < 0 || st.head.Timecode < states[mi].head.Timecode { + mi = i + } + } + if mi < 0 { + break + } + blk := states[mi].head + if err := advance(mi); err != nil { + return err + } + + if clusterTS < 0 { + clusterTS = blk.Timecode + } + if blk.Timecode-clusterTS >= defaultClusterDurationMs && len(cluster) > 0 { + if err := flush(); err != nil { + return err + } + clusterTS = blk.Timecode + } + cluster = append(cluster, blk) + } + return flush() +} diff --git a/mkv/ops/webm.go b/mkv/ops/webm.go new file mode 100644 index 0000000..867b8df --- /dev/null +++ b/mkv/ops/webm.go @@ -0,0 +1,86 @@ +package ops + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + + "github.com/gravity-zero/mkvgo/mkv" + "github.com/gravity-zero/mkvgo/mkv/reader" + "github.com/gravity-zero/mkvgo/mkv/writer" +) + +// RemuxToWebM reads srcPath and writes a complete, playable WebM file to +// dstPath: it validates that every track uses a WebM-compatible codec, then +// copies the media (every block, verbatim) into a "webm"-DocType container with +// keyframe-aligned clusters. It does NOT transcode — a source whose codecs fall +// outside the WebM subset is rejected with an error and no output is produced. +// +// Unlike writer.WriteWebM (metadata only), this produces a file with frames. +func RemuxToWebM(ctx context.Context, srcPath, dstPath string, extra ...mkv.Options) error { + fs := mkv.FSFrom(extra) + + c, err := reader.OpenWithFS(ctx, srcPath, fs) + if err != nil { + return err + } + if err := mkv.ValidateWebM(c); err != nil { + return err + } + + src, err := fs.DoOpen(srcPath) + if err != nil { + return err + } + defer src.Close() + br, err := reader.NewBlockReader(src, c.Info.TimecodeScale) + if err != nil { + return err + } + + dst, err := fs.DoCreate(dstPath) + if err != nil { + return err + } + defer dst.Close() + + // Buffer the output: the StreamWriter emits several small writes per block, + // which would otherwise be a syscall each on this file-copy hot path. (The + // StreamWriter itself stays unbuffered so live-streaming callers keep low latency.) + buf := bufio.NewWriterSize(dst, 256<<10) + sw, err := writer.NewWebMStreamWriter(buf, c.Info, c.Tracks) + if err != nil { + return err + } + + // Group blocks into time-bounded clusters (~1s) rather than splitting on + // every keyframe: in multiplexed A/V every Opus frame is a "keyframe", so a + // keyframe-per-cluster policy would emit one tiny cluster per audio frame. + const clusterDurationMs = 1000 + clusterStart := int64(-1) + for { + if ctx.Err() != nil { + return ctx.Err() + } + b, err := br.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return fmt.Errorf("remux webm: read block: %w", err) + } + if clusterStart < 0 || b.Timecode-clusterStart >= clusterDurationMs { + sw.FlushCluster() // force the next write to open a fresh cluster + clusterStart = b.Timecode + } + if err := sw.WriteBlockInCurrentCluster(b); err != nil { + return fmt.Errorf("remux webm: write block: %w", err) + } + } + if err := buf.Flush(); err != nil { + return fmt.Errorf("remux webm: flush: %w", err) + } + return nil +} diff --git a/mkv/ops/webm_test.go b/mkv/ops/webm_test.go new file mode 100644 index 0000000..0ba9085 --- /dev/null +++ b/mkv/ops/webm_test.go @@ -0,0 +1,254 @@ +package ops + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "testing" + + "github.com/gravity-zero/mkvgo/mkv" + "github.com/gravity-zero/mkvgo/mkv/reader" + "github.com/gravity-zero/mkvgo/mkv/writer" +) + +// buildSource writes a small seekable MKV (one video track + a 2-block cluster) +// to path, using the given codec. +func buildSource(t *testing.T, path, codec string) { + t.Helper() + w, h := uint32(320), uint32(240) + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: codec, Width: &w, Height: &h}} + blocks := []mkv.Block{ + {TrackNumber: 1, Timecode: 0, Keyframe: true, Data: []byte{0x01, 0x02, 0x03}}, + {TrackNumber: 1, Timecode: 40, Keyframe: false, Data: []byte{0x04, 0x05}}, + } + + var seg bytes.Buffer + if err := writer.WriteSegmentInfo(&seg, &info, 0); err != nil { + t.Fatal(err) + } + if err := writer.WriteTracks(&seg, tracks); err != nil { + t.Fatal(err) + } + if err := writer.WriteCluster(&seg, 0, info.TimecodeScale, blocks); err != nil { + t.Fatal(err) + } + + var buf bytes.Buffer + if err := writer.WriteEBMLHeader(&buf); err != nil { + t.Fatal(err) + } + if err := writer.WriteMasterElement(&buf, mkv.IDSegment, seg.Bytes()); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, buf.Bytes(), 0o644); err != nil { + t.Fatal(err) + } +} + +func TestRemuxToWebM(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "src.mkv") + dst := filepath.Join(dir, "out.webm") + buildSource(t, src, "vp9") + + if err := RemuxToWebM(context.Background(), src, dst); err != nil { + t.Fatalf("RemuxToWebM: %v", err) + } + + raw, err := os.ReadFile(dst) + if err != nil { + t.Fatal(err) + } + if !bytes.Contains(raw[:64], []byte("webm")) { + t.Error("output DocType is not webm") + } + + // The output is an unknown-size stream — read it back with ReadStream. + c, br, err := reader.ReadStream(context.Background(), bytes.NewReader(raw)) + if err != nil { + t.Fatalf("ReadStream output: %v", err) + } + if len(c.Tracks) != 1 || c.Tracks[0].Codec != "vp9" { + t.Errorf("tracks = %+v, want one vp9 track", c.Tracks) + } + if err := mkv.ValidateWebM(c); err != nil { + t.Errorf("remuxed output is not WebM-valid: %v", err) + } + n := 0 + for { + _, err := br.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("block read: %v", err) + } + n++ + } + if n != 2 { + t.Errorf("remuxed %d blocks, want 2", n) + } +} + +func TestRemuxToWebMRejectsNonWebMCodec(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "src.mkv") + dst := filepath.Join(dir, "out.webm") + buildSource(t, src, "h264") + + if err := RemuxToWebM(context.Background(), src, dst); err == nil { + t.Fatal("RemuxToWebM accepted an h264 source") + } +} + +// TestRemuxToWebMStripsChapters verifies the WebM element-subset behaviour: a +// source carrying a Chapter (not part of the WebM streaming output) is remuxed +// to a WebM that contains no chapters, and WebMNonSubsetElements flags the loss. +func TestRemuxToWebMStripsChapters(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "src.mkv") + dst := filepath.Join(dir, "out.webm") + + w, h := uint32(320), uint32(240) + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: "vp9", Width: &w, Height: &h}} + blocks := []mkv.Block{{TrackNumber: 1, Timecode: 0, Keyframe: true, Data: []byte{1, 2, 3}}} + chapters := []mkv.Chapter{{ID: 1, Title: "intro", StartMs: 0}} + + var seg bytes.Buffer + mustNil(t, writer.WriteSegmentInfo(&seg, &info, 0)) + mustNil(t, writer.WriteTracks(&seg, tracks)) + mustNil(t, writer.WriteChapters(&seg, chapters)) + mustNil(t, writer.WriteCluster(&seg, 0, info.TimecodeScale, blocks)) + var buf bytes.Buffer + mustNil(t, writer.WriteEBMLHeader(&buf)) + mustNil(t, writer.WriteMasterElement(&buf, mkv.IDSegment, seg.Bytes())) + if err := os.WriteFile(src, buf.Bytes(), 0o644); err != nil { + t.Fatal(err) + } + + sc, err := reader.Open(context.Background(), src) + if err != nil { + t.Fatal(err) + } + if len(sc.Chapters) != 1 { + t.Fatalf("source build: expected 1 chapter, got %d", len(sc.Chapters)) + } + if got := mkv.WebMNonSubsetElements(sc); len(got) != 1 || got[0] != "Chapters" { + t.Errorf("WebMNonSubsetElements = %v, want [Chapters]", got) + } + + if err := RemuxToWebM(context.Background(), src, dst); err != nil { + t.Fatalf("RemuxToWebM: %v", err) + } + raw, err := os.ReadFile(dst) + if err != nil { + t.Fatal(err) + } + out, _, err := reader.ReadStream(context.Background(), bytes.NewReader(raw)) + if err != nil { + t.Fatalf("ReadStream: %v", err) + } + if len(out.Chapters) != 0 { + t.Errorf("WebM output should carry no chapters, got %d", len(out.Chapters)) + } +} + +// TestRemuxToWebMClustersByTime guards against the keyframe-per-cluster trap: +// every Opus audio frame is a keyframe, so RemuxToWebM must group by time, not +// open a fresh cluster per block. +func TestRemuxToWebMClustersByTime(t *testing.T) { + dir := t.TempDir() + src := filepath.Join(dir, "src.mkv") + dst := filepath.Join(dir, "out.webm") + + opusHead := append([]byte("OpusHead"), 0x01, 0x02, 0x38, 0x01, 0x80, 0xbb, 0x00, 0x00, 0x00, 0x00, 0x00) + sr := 48000.0 + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{{ID: 1, Type: mkv.AudioTrack, Codec: "opus", CodecPrivate: opusHead, SampleRate: &sr}} + var blocks []mkv.Block + for i := 0; i < 100; i++ { // 100 "keyframe" frames over 2000 ms + blocks = append(blocks, mkv.Block{TrackNumber: 1, Timecode: int64(i) * 20, Keyframe: true, Data: []byte{0xAA}}) + } + + var seg bytes.Buffer + mustNil(t, writer.WriteSegmentInfo(&seg, &info, 0)) + mustNil(t, writer.WriteTracks(&seg, tracks)) + mustNil(t, writer.WriteCluster(&seg, 0, info.TimecodeScale, blocks)) + var buf bytes.Buffer + mustNil(t, writer.WriteEBMLHeader(&buf)) + mustNil(t, writer.WriteMasterElement(&buf, mkv.IDSegment, seg.Bytes())) + mustNil(t, os.WriteFile(src, buf.Bytes(), 0o644)) + + mustNil(t, RemuxToWebM(context.Background(), src, dst)) + + raw, err := os.ReadFile(dst) + if err != nil { + t.Fatal(err) + } + clusters := bytes.Count(raw, []byte{0x1F, 0x43, 0xB6, 0x75}) + if clusters < 2 || clusters > 4 { + t.Errorf("got %d clusters for 100 keyframe frames over 2s, want time-based (~2-3, not per-frame)", clusters) + } +} + +func mustNil(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} + +// BenchmarkRemuxToWebM remuxes a ~5 MB VP9 source (5000 blocks of 1 KiB, +// keyframe every 50) so the cost is dominated by block copy, not setup. +func BenchmarkRemuxToWebM(b *testing.B) { + dir := b.TempDir() + src := filepath.Join(dir, "src.mkv") + dst := filepath.Join(dir, "out.webm") + + w, h := uint32(640), uint32(360) + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: "vp9", Width: &w, Height: &h}} + payload := make([]byte, 1024) + const n = 5000 + + var clusters bytes.Buffer + for i := 0; i < n; i += 50 { + var blocks []mkv.Block + for j := i; j < i+50 && j < n; j++ { + blocks = append(blocks, mkv.Block{TrackNumber: 1, Timecode: int64(j) * 40, Keyframe: j == i, Data: payload}) + } + if err := writer.WriteCluster(&clusters, int64(i)*40, info.TimecodeScale, blocks); err != nil { + b.Fatal(err) + } + } + var seg bytes.Buffer + mustNilB(b, writer.WriteSegmentInfo(&seg, &info, 0)) + mustNilB(b, writer.WriteTracks(&seg, tracks)) + seg.Write(clusters.Bytes()) + var buf bytes.Buffer + mustNilB(b, writer.WriteEBMLHeader(&buf)) + mustNilB(b, writer.WriteMasterElement(&buf, mkv.IDSegment, seg.Bytes())) + if err := os.WriteFile(src, buf.Bytes(), 0o644); err != nil { + b.Fatal(err) + } + + b.SetBytes(int64(buf.Len())) + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := RemuxToWebM(context.Background(), src, dst); err != nil { + b.Fatal(err) + } + } +} + +func mustNilB(b *testing.B, err error) { + b.Helper() + if err != nil { + b.Fatal(err) + } +} diff --git a/mkv/reader/blocks.go b/mkv/reader/blocks.go index 7cc633f..b541061 100644 --- a/mkv/reader/blocks.go +++ b/mkv/reader/blocks.go @@ -323,7 +323,10 @@ func (br *BlockReader) parseBlock(size int64, simple bool) (mkv.Block, error) { return mkv.Block{}, err } - frameCount := int(raw[0]) + 1 + if len(raw) == 0 { + return mkv.Block{}, fmt.Errorf("laced block missing lacing header byte") + } + frameCount := int(raw[0]) + 1 // Matroska lace count byte = number of frames minus 1 raw = raw[1:] frameSizes, err := decodeLacingSizes(lacing, raw, frameCount) @@ -332,6 +335,9 @@ func (br *BlockReader) parseBlock(size int64, simple bool) (mkv.Block, error) { } headerBytes := lacingHeaderLen(lacing, frameSizes) + if headerBytes < 0 || headerBytes > len(raw) { + return mkv.Block{}, fmt.Errorf("laced block header (%d bytes) exceeds data (%d bytes)", headerBytes, len(raw)) + } raw = raw[headerBytes:] tc, err := safeTimecodeMs(br.clusterTS+int64(relTC), br.timecodeScale) @@ -362,10 +368,7 @@ func (br *BlockReader) parseBlockGroup(size int64) (mkv.Block, error) { var block mkv.Block var found bool - for { - if br.r.tell() >= end { - break - } + for br.r.tell() < end { h, _, err := ebml.ReadElementHeader(br.r) if err != nil { return mkv.Block{}, err @@ -427,7 +430,13 @@ func decodeLacingSizes(lacing byte, raw []byte, frameCount int) ([]int, error) { pos += width total := firstSize for i := 1; i < frameCount-1; i++ { + if pos > len(raw) { + return nil, fmt.Errorf("ebml lacing: header truncated at frame %d", i) + } val, w := readVINTFromBuf(raw[pos:]) + if w == 0 { + return nil, fmt.Errorf("ebml lacing: invalid size vint at frame %d", i) + } pos += w dataBits := uint(w * 7) bias := int64(1) << (dataBits - 1) diff --git a/mkv/reader/blocks_fuzz_test.go b/mkv/reader/blocks_fuzz_test.go new file mode 100644 index 0000000..cfce54d --- /dev/null +++ b/mkv/reader/blocks_fuzz_test.go @@ -0,0 +1,134 @@ +package reader + +import ( + "bytes" + "context" + "testing" + + "github.com/gravity-zero/mkvgo/ebml" + "github.com/gravity-zero/mkvgo/mkv" +) + +// clusterWithSimpleBlock wraps a raw SimpleBlock payload (track + relTC + flags + +// lacing data) in a cluster (Timestamp + SimpleBlock) for buildBlockReaderInput. +func clusterWithSimpleBlock(blockPayload []byte) []byte { + var cluster bytes.Buffer + ebml.WriteElementHeader(&cluster, mkv.IDTimestamp, 1) + ebml.WriteUint(&cluster, 0, 1) + ebml.WriteElementHeader(&cluster, mkv.IDSimpleBlock, int64(len(blockPayload))) + cluster.Write(blockPayload) + return cluster.Bytes() +} + +// TestLacedBlockMalformedNoPanic covers the two lacing panics the audit found: +// a laced SimpleBlock with zero payload (raw[0] out of range) and a lacing header +// longer than the data (raw[headerBytes:] out of range). The parser must return +// an error, never panic. +func TestLacedBlockMalformedNoPanic(t *testing.T) { + cases := []struct { + name string + payload []byte + }{ + // track=1, relTC=0, lacing flag set, then ZERO lacing bytes -> dataSize==0. + {"xiph zero payload", []byte{0x81, 0x00, 0x00, 0x02}}, + {"fixed zero payload", []byte{0x81, 0x00, 0x00, 0x04}}, + {"ebml zero payload", []byte{0x81, 0x00, 0x00, 0x06}}, + // frame-count byte present but the lacing header overruns the data. + {"xiph header overflow", []byte{0x81, 0x00, 0x00, 0x02, 0x02, 0xFF, 0xFF}}, + {"ebml header overflow", []byte{0x81, 0x00, 0x00, 0x06, 0x05}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Fatalf("parseBlock panicked on %q: %v", tc.name, r) + } + }() + br, err := NewBlockReader(buildBlockReaderInput(clusterWithSimpleBlock(tc.payload)), 1000000) + if err != nil { + return + } + for { + if _, err := br.Next(); err != nil { + return // an error is the expected, safe outcome + } + } + }) + } +} + +// TestLacedBlockKeyframeFlag verifies that only the FIRST frame of a laced +// keyframe block carries the keyframe flag (guards parseBlock's `keyframe && i == 0`). +func TestLacedBlockKeyframeFlag(t *testing.T) { + // track=1, relTC=0, flags=keyframe(0x80)|fixed-lacing(0x04), frameCount byte=1 + // (=> 2 frames), then 2 bytes of data (1 byte per frame). + payload := []byte{0x81, 0x00, 0x00, 0x84, 0x01, 0xAA, 0xBB} + br, err := NewBlockReader(buildBlockReaderInput(clusterWithSimpleBlock(payload)), 1000000) + if err != nil { + t.Fatal(err) + } + b0, err := br.Next() + if err != nil { + t.Fatalf("frame 0: %v", err) + } + b1, err := br.Next() + if err != nil { + t.Fatalf("frame 1: %v", err) + } + if !b0.Keyframe { + t.Error("frame 0 of a keyframe laced block should be a keyframe") + } + if b1.Keyframe { + t.Error("frame 1 of a laced block must not be a keyframe") + } + if len(b0.Data) != 1 || b0.Data[0] != 0xAA || len(b1.Data) != 1 || b1.Data[0] != 0xBB { + t.Errorf("frame data wrong: b0=%v b1=%v", b0.Data, b1.Data) + } +} + +// FuzzBlockReader drives the block/lacing parser (parseBlock/decodeLacingSizes) +// with arbitrary cluster payloads -- the path FuzzRead does not reach. Contract: +// never panic; return blocks or an error. +func FuzzBlockReader(f *testing.F) { + f.Add(clusterWithSimpleBlock([]byte{0x81, 0x00, 0x00, 0x02})) + f.Add(clusterWithSimpleBlock([]byte{0x81, 0x00, 0x00, 0x06, 0x05})) + f.Add(clusterWithSimpleBlock([]byte{0x81, 0x00, 0x00, 0x00, 0xAA})) // unlaced + f.Add([]byte{}) + + f.Fuzz(func(t *testing.T, clusterPayload []byte) { + br, err := NewBlockReader(buildBlockReaderInput(clusterPayload), 1000000) + if err != nil { + return + } + for i := 0; i < 1000; i++ { // bound iterations against a degenerate loop + if _, err := br.Next(); err != nil { + return + } + } + }) +} + +// TestReaderNormalizesZeroTimecodeScale: a file declaring TimecodeScale=0 must be +// normalised to the default, otherwise downstream WriteCluster/Cues divide by zero. +func TestReaderNormalizesZeroTimecodeScale(t *testing.T) { + var info bytes.Buffer + ebml.WriteElementHeader(&info, mkv.IDTimecodeScale, 1) + ebml.WriteUint(&info, 0, 1) // explicit TimecodeScale = 0 (malformed) + + var seg bytes.Buffer + ebml.WriteElementHeader(&seg, mkv.IDInfo, int64(info.Len())) + seg.Write(info.Bytes()) + + var buf bytes.Buffer + ebml.WriteElementHeader(&buf, ebml.IDEBMLHeader, 0) + ebml.WriteElementHeader(&buf, mkv.IDSegment, int64(seg.Len())) + buf.Write(seg.Bytes()) + + c, err := Read(context.Background(), bytes.NewReader(buf.Bytes()), "zerots.mkv") + if err != nil { + t.Fatal(err) + } + if c.Info.TimecodeScale != 1000000 { + t.Errorf("TimecodeScale = %d, want normalised to 1000000", c.Info.TimecodeScale) + } +} diff --git a/mkv/reader/fuzz_test.go b/mkv/reader/fuzz_test.go new file mode 100644 index 0000000..10dc37e --- /dev/null +++ b/mkv/reader/fuzz_test.go @@ -0,0 +1,27 @@ +package reader + +import ( + "bytes" + "context" + "testing" +) + +// FuzzRead exercises the metadata parser against arbitrary bytes. Its only +// contract: Read must never panic or hang on malformed input — it may only +// return a Container or an error. Allocations are bounded by ebml.MaxElementSize, +// so the fuzzer cannot OOM the process. The seed corpus covers valid files and +// the corruption shapes the resync handles. +func FuzzRead(f *testing.F) { + f.Add(buildMinimalMKV().Bytes()) + f.Add(buildGappedMKV(4096, true)) + f.Add(buildGappedMKV(200<<10, false)) + f.Add(append(realCluster(), 0x00, 0x00, 0x00, 0x00)) // cluster then padding + f.Add([]byte{0x1A, 0x45, 0xDF, 0xA3}) // bare EBML id + f.Add([]byte{0x18, 0x53, 0x80, 0x67, 0x01, 0xFF}) // segment, truncated size + f.Add([]byte{0x1F, 0x43, 0xB6, 0x75}) // bare cluster magic + f.Add([]byte{}) + + f.Fuzz(func(t *testing.T, data []byte) { + _, _ = Read(context.Background(), bytes.NewReader(data), "fuzz.mkv") + }) +} diff --git a/mkv/reader/reader.go b/mkv/reader/reader.go index 8a0c76d..dfe7c4f 100644 --- a/mkv/reader/reader.go +++ b/mkv/reader/reader.go @@ -1,6 +1,7 @@ package reader import ( + "bytes" "context" "errors" "fmt" @@ -32,7 +33,7 @@ func OpenWithFS(ctx context.Context, path string, fs *mkv.FS) (*mkv.Container, e } func Read(ctx context.Context, r io.ReadSeeker, path string) (*mkv.Container, error) { - p := &parser{r: r} + p := &parser{r: r, metaBudget: maxMetadataBytes} c := &mkv.Container{Path: path} if err := p.parseEBMLHeader(); err != nil { @@ -52,7 +53,22 @@ func Read(ctx context.Context, r io.ReadSeeker, path string) (*mkv.Container, er } type parser struct { - r io.ReadSeeker + r io.ReadSeeker + metaBudget int64 // remaining bytes allowed for in-memory metadata +} + +// maxMetadataBytes caps the TOTAL bytes a single parse pulls into the Container +// (attachments, codec-private, binary tags). The 512MB per-element cap does not +// bound a file with many large metadata elements; this does. Untrusted-input +// callers that ingest concurrently should still bound their own parallelism. +const maxMetadataBytes = 1 << 30 // 1 GiB + +func (p *parser) chargeMeta(n int64) error { + p.metaBudget -= n + if p.metaBudget < 0 { + return fmt.Errorf("in-memory metadata exceeds %d-byte budget", maxMetadataBytes) + } + return nil } func (p *parser) readHeader() (ebml.ElementHeader, int, error) { @@ -105,7 +121,20 @@ func (p *parser) parseSegment(ctx context.Context, c *mkv.Container) error { if errors.Is(err, io.EOF) { break } - return err + // A corrupted or zero-padded region in the body (seen in some real + // rips: a multi-MB run of 0x00 between clusters) makes the next + // element header undecodable. Rather than abort the whole read like + // a strict parser, resync to the next Cluster and keep going, as + // ffmpeg/mkvtoolnix do. If nothing recognizable remains, stop with + // the metadata gathered so far. + off, rerr := p.resyncToCluster(endPos) + if rerr != nil { + return rerr + } + if off < 0 { + break + } + continue } switch eh.ID { case mkv.IDInfo: @@ -144,6 +173,120 @@ func (p *parser) parseSegment(ctx context.Context, c *mkv.Container) error { return nil } +// clusterMagic is the 4-byte Cluster element ID (0x1F43B675), the anchor used +// to resync after a corrupted/padded region in the body. +var clusterMagic = []byte{0x1F, 0x43, 0xB6, 0x75} + +// clusterChildIDs are element IDs that can legitimately be the first child of a +// Cluster. A real cluster opens with one of these; requiring it rejects a +// clusterMagic byte-sequence that occurs by chance inside corrupted data. +var clusterChildIDs = map[uint32]bool{ + mkv.IDTimestamp: true, // 0xE7 + mkv.IDSimpleBlock: true, // 0xA3 + mkv.IDBlockGroup: true, // 0xA0 + mkv.IDVoid: true, // 0xEC + 0xA7: true, // Position + 0xAB: true, // PrevSize + 0xBF: true, // CRC-32 + 0x5854: true, // SilentTracks +} + +// resyncToCluster scans forward from the current position for the next *valid* +// Cluster, bounded by limit (the segment end, or -1 to scan until EOF). A +// candidate is accepted only if isClusterAt confirms it (real Cluster ID + a +// recognizable first child), so a magic sequence occurring by chance inside +// corruption is skipped rather than trusted. On success it positions the reader +// at the Cluster ID and returns its offset; it returns -1 (with a nil error) +// when no valid Cluster remains before limit. Only genuine I/O errors are returned. +func (p *parser) resyncToCluster(limit int64) (int64, error) { + from, err := p.r.Seek(0, io.SeekCurrent) + if err != nil { + return -1, err + } + for { + off, err := p.scanForMagic(from, limit) + if err != nil || off < 0 { + return -1, err + } + valid, err := p.isClusterAt(off, limit) + if err != nil { + return -1, err + } + if valid { + if _, err := p.r.Seek(off, io.SeekStart); err != nil { + return -1, err + } + return off, nil + } + from = off + 1 // false positive: resume scanning just past it + } +} + +// scanForMagic returns the absolute offset of the next clusterMagic at or after +// `from` and before `limit` (-1 = until EOF), or -1 if none. It reads forward in +// windows, carrying the last few bytes so a magic split across a read boundary +// is still found. +func (p *parser) scanForMagic(from, limit int64) (int64, error) { + if _, err := p.r.Seek(from, io.SeekStart); err != nil { + return -1, err + } + const window = 64 << 10 + buf := make([]byte, len(clusterMagic)-1+window) + tail := 0 + next := from + for { + base := next - int64(tail) // absolute offset of buf[0] + n, rerr := p.r.Read(buf[tail : tail+window]) + end := tail + n + + search := buf[:end] + if limit >= 0 { + if max := limit - base; max < int64(end) { + if max < 0 { + max = 0 + } + search = buf[:max] + } + } + if i := bytes.Index(search, clusterMagic); i >= 0 { + return base + int64(i), nil + } + + next += int64(n) + if (limit >= 0 && next >= limit) || rerr != nil { + return -1, nil + } + keep := len(clusterMagic) - 1 + if end < keep { + keep = end + } + copy(buf[:keep], buf[end-keep:end]) + tail = keep + } +} + +// isClusterAt reports whether a real Cluster begins at off: the element ID must +// be Cluster, its declared size must decode and (when known) fit within limit, +// and its first child must be a recognizable cluster-level element. +func (p *parser) isClusterAt(off, limit int64) (bool, error) { + if _, err := p.r.Seek(off, io.SeekStart); err != nil { + return false, err + } + h, _, err := p.readHeader() + if err != nil || h.ID != mkv.IDCluster { + return false, nil + } + bodyStart, _ := p.r.Seek(0, io.SeekCurrent) + if h.Size >= 0 && limit >= 0 && bodyStart+h.Size > limit { + return false, nil // declared size overruns the segment — not a real cluster + } + child, _, err := p.readHeader() + if err != nil { + return false, nil + } + return clusterChildIDs[child.ID], nil +} + func (p *parser) parseInfo(size int64, c *mkv.Container) error { cur, _ := p.r.Seek(0, io.SeekCurrent) end := cur + size @@ -164,7 +307,9 @@ func (p *parser) parseInfo(size int64, c *mkv.Container) error { if err != nil { return err } - c.Info.TimecodeScale = int64(v) + if v > 0 { // keep the 1000000 default; a 0 scale would divide-by-zero downstream + c.Info.TimecodeScale = int64(v) + } case mkv.IDDuration: v, err := ebml.ReadFloat(p.r, eh.Size) if err != nil { @@ -302,6 +447,9 @@ func (p *parser) parseTrackEntry(size int64) (mkv.Track, error) { t.Codec = v } case mkv.IDCodecPrivate: + if err := p.chargeMeta(eh.Size); err != nil { + return t, err + } v, err := ebml.ReadBytes(p.r, eh.Size) if err != nil { return t, err @@ -477,7 +625,7 @@ func (p *parser) parseEditionEntry(size int64, c *mkv.Container) error { } ordered = v == 1 case mkv.IDChapterAtom: - ch, err := p.parseChapterAtom(eh.Size) + ch, err := p.parseChapterAtom(eh.Size, 0) if err != nil { return err } @@ -493,7 +641,10 @@ func (p *parser) parseEditionEntry(size int64, c *mkv.Container) error { return nil } -func (p *parser) parseChapterAtom(size int64) (mkv.Chapter, error) { +func (p *parser) parseChapterAtom(size int64, depth int) (mkv.Chapter, error) { + if depth > maxChapterDepth { + return mkv.Chapter{}, fmt.Errorf("chapter nesting exceeds %d levels", maxChapterDepth) + } cur, _ := p.r.Seek(0, io.SeekCurrent) end := cur + size ch := mkv.Chapter{} @@ -530,7 +681,7 @@ func (p *parser) parseChapterAtom(size int64) (mkv.Chapter, error) { return ch, err } case mkv.IDChapterAtom: - sub, err := p.parseChapterAtom(eh.Size) + sub, err := p.parseChapterAtom(eh.Size, depth+1) if err != nil { return ch, err } @@ -637,6 +788,9 @@ func (p *parser) parseAttachedFile(size int64) (mkv.Attachment, error) { } att.MIMEType = v case mkv.IDFileData: + if err := p.chargeMeta(eh.Size); err != nil { + return att, err + } data, err := ebml.ReadBytes(p.r, eh.Size) if err != nil { return att, err @@ -746,6 +900,8 @@ func (p *parser) parseTargets(size int64, tag *mkv.Tag) error { return nil } +const maxChapterDepth = 64 + const maxTagDepth = 64 func (p *parser) parseSimpleTagDepth(size int64, depth int) (mkv.SimpleTag, error) { @@ -784,6 +940,9 @@ func (p *parser) parseSimpleTagDepth(size int64, depth int) (mkv.SimpleTag, erro } st.Language = v case mkv.IDTagBinary: + if err := p.chargeMeta(eh.Size); err != nil { + return st, err + } v, err := ebml.ReadBytes(p.r, eh.Size) if err != nil { return st, err diff --git a/mkv/reader/reader_test.go b/mkv/reader/reader_test.go index dbf4667..02935cd 100644 --- a/mkv/reader/reader_test.go +++ b/mkv/reader/reader_test.go @@ -5,6 +5,7 @@ import ( "context" "fmt" "io" + "os" "strings" "testing" @@ -61,6 +62,233 @@ func TestInvalidVINT(t *testing.T) { } } +// realCluster builds a minimal but realistic Cluster element: it opens with a +// Timestamp (the first child the resync validator looks for) and carries one +// tiny SimpleBlock, so isClusterAt accepts it as a genuine cluster. +func realCluster() []byte { + var body bytes.Buffer + ebml.WriteElementHeader(&body, mkv.IDTimestamp, 1) + body.WriteByte(0x00) + var sb bytes.Buffer + ebml.WriteDataSize(&sb, 1) // track number VINT = 1 + sb.Write([]byte{0x00, 0x00, 0x80, 0xAA}) // relTC=0, flags=keyframe, 1 data byte + ebml.WriteElementHeader(&body, mkv.IDSimpleBlock, int64(sb.Len())) + body.Write(sb.Bytes()) + + var clu bytes.Buffer + ebml.WriteElementHeader(&clu, mkv.IDCluster, int64(body.Len())) + clu.Write(body.Bytes()) + return clu.Bytes() +} + +func gappedInfo() []byte { + var info bytes.Buffer + var tcs bytes.Buffer + ebml.WriteUint(&tcs, 500000, 4) // non-default TimecodeScale + ebml.WriteElementHeader(&info, mkv.IDTimecodeScale, int64(tcs.Len())) + info.Write(tcs.Bytes()) + var dur bytes.Buffer + ebml.WriteFloat(&dur, 1234) + ebml.WriteElementHeader(&info, mkv.IDDuration, int64(dur.Len())) + info.Write(dur.Bytes()) + + var out bytes.Buffer + ebml.WriteElementHeader(&out, mkv.IDInfo, int64(info.Len())) + out.Write(info.Bytes()) + return out.Bytes() +} + +func wrapSegment(inner []byte) []byte { + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(len(inner))) + buf.Write(inner) + return buf.Bytes() +} + +func oneCuePoint() []byte { + var cue bytes.Buffer + ebml.WriteElementHeader(&cue, mkv.IDCuePoint, 0) + var out bytes.Buffer + ebml.WriteElementHeader(&out, mkv.IDCues, int64(cue.Len())) + out.Write(cue.Bytes()) + return out.Bytes() +} + +// buildGappedMKV builds a valid MKV whose body contains a run of zero bytes +// between two clusters — the corruption pattern seen in some real-world rips, +// where a multi-MB region is zeroed out. If postCluster is true a recoverable +// Cluster (followed by a Cues element) is placed after the gap. +func buildGappedMKV(gap int, postCluster bool) []byte { + var inner bytes.Buffer + inner.Write(gappedInfo()) + inner.Write(realCluster()) // first (good) cluster + inner.Write(make([]byte, gap)) // zeroed/corrupted region + if postCluster { + inner.Write(realCluster()) // resync anchor + inner.Write(oneCuePoint()) + } + return wrapSegment(inner.Bytes()) +} + +func TestRecoverFromZeroPaddingGap(t *testing.T) { + tests := []struct { + name string + gap int + postCluster bool + wantCues int + }{ + // gap > 64 KiB resync window: exercises the boundary-carry path. + {"recoverable gap", 200 << 10, true, 1}, + // corruption runs to the segment end: nothing to resync to, but the + // metadata read before the gap must still succeed (no error). + {"gap to end", 4096, false, 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data := buildGappedMKV(tt.gap, tt.postCluster) + c, err := Read(context.Background(), bytes.NewReader(data), tt.name+".mkv") + if err != nil { + t.Fatalf("Read returned error on recoverable corruption: %v", err) + } + if c.Info.TimecodeScale != 500000 { + t.Errorf("pre-gap Info lost: TimecodeScale = %d, want 500000", c.Info.TimecodeScale) + } + if c.Info.Duration != 1234 { + t.Errorf("pre-gap Info lost: Duration = %g, want 1234", c.Info.Duration) + } + if len(c.Cues) != tt.wantCues { + t.Errorf("post-gap recovery: got %d cues, want %d", len(c.Cues), tt.wantCues) + } + }) + } +} + +// TestResyncSkipsFalseClusterMagic plants the 4-byte Cluster ID by chance inside +// the corrupted region. The false anchor declares a huge size: a naive resync +// that trusted it would skip clean past the genuine cluster AND the Cues (losing +// them), so this test fails if the validation in isClusterAt is removed. With +// the guard, the candidate is rejected (its declared size overruns the segment) +// and recovery lands on the real cluster, so the Cues survive. +func TestResyncSkipsFalseClusterMagic(t *testing.T) { + // Cluster magic + a 3-byte size VINT encoding 65535 — far past segment end. + falseAnchor := []byte{0x1F, 0x43, 0xB6, 0x75, 0x20, 0xFF, 0xFF} + + var gap bytes.Buffer + gap.Write(make([]byte, 4096)) + gap.Write(falseAnchor) + gap.Write(make([]byte, 4096)) + + var inner bytes.Buffer + inner.Write(gappedInfo()) + inner.Write(realCluster()) // first good cluster + inner.Write(gap.Bytes()) // corruption containing a false magic + inner.Write(realCluster()) // the genuine resync target + inner.Write(oneCuePoint()) + + data := wrapSegment(inner.Bytes()) + c, err := Read(context.Background(), bytes.NewReader(data), "false-magic.mkv") + if err != nil { + t.Fatalf("Read errored instead of skipping the false magic: %v", err) + } + if len(c.Cues) != 1 { + t.Fatalf("expected recovery at the genuine cluster (1 cue), got %d — false magic was likely trusted", len(c.Cues)) + } +} + +// TestRecoveryMultipleGaps deliberately corrupts the body with TWO separate +// zero-padding regions; the parser must resync past both and still reach the +// trailing Cues. +func TestRecoveryMultipleGaps(t *testing.T) { + var inner bytes.Buffer + inner.Write(gappedInfo()) + inner.Write(realCluster()) + inner.Write(make([]byte, 70<<10)) // gap 1 (> 64 KiB window) + inner.Write(realCluster()) + inner.Write(make([]byte, 70<<10)) // gap 2 + inner.Write(realCluster()) + inner.Write(oneCuePoint()) + + c, err := Read(context.Background(), bytes.NewReader(wrapSegment(inner.Bytes())), "multigap.mkv") + if err != nil { + t.Fatalf("Read across 2 gaps: %v", err) + } + if c.Info.TimecodeScale != 500000 { + t.Errorf("pre-corruption Info lost: %d", c.Info.TimecodeScale) + } + if len(c.Cues) != 1 { + t.Errorf("recovery across 2 gaps failed: got %d cues, want 1", len(c.Cues)) + } +} + +// TestRecoveryTruncatedInGap simulates a truncated download: the Segment claims +// more bytes than the file actually holds, and the data ends inside a corrupted +// gap. The parser must return the metadata it already gathered, without erroring +// or panicking. +func TestRecoveryTruncatedInGap(t *testing.T) { + var inner bytes.Buffer + inner.Write(gappedInfo()) + inner.Write(realCluster()) + inner.Write(make([]byte, 8192)) // gap — then the file just stops here + + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(inner.Len())+100000) // Segment claims more than exists + buf.Write(inner.Bytes()) + + c, err := Read(context.Background(), bytes.NewReader(buf.Bytes()), "truncated.mkv") + if err != nil { + t.Fatalf("truncated file should recover gracefully, got: %v", err) + } + if c.Info.TimecodeScale != 500000 { + t.Errorf("pre-truncation Info lost: %d", c.Info.TimecodeScale) + } +} + +// TestReadCorruptionNoPanic feeds many deliberately-corrupted variants of a real +// muxer fixture (truncations, zeroed regions, garbage injection, header +// byte-flips) to Read. The contract under any corruption: never panic — only +// return data or an error. (A real-corpus complement to FuzzRead.) +func TestReadCorruptionNoPanic(t *testing.T) { + orig, err := os.ReadFile("../../internal/testdata/sample.mkv") + if err != nil { + t.Skipf("sample.mkv not available: %v", err) + } + n := len(orig) + + read := func(label string, data []byte) { + defer func() { + if r := recover(); r != nil { + t.Errorf("panic on %s: %v", label, r) + } + }() + _, _ = Read(context.Background(), bytes.NewReader(data), "corrupt.mkv") + } + + for _, p := range []int{1, 5, 10, 25, 50, 75, 90, 99} { + read(fmt.Sprintf("trunc-%d%%", p), orig[:n*p/100]) + } + for _, p := range []int{5, 20, 40, 60, 80} { + d := append([]byte(nil), orig...) + for i := n * p / 100; i < n*p/100+4096 && i < n; i++ { + d[i] = 0 + } + read(fmt.Sprintf("zerogap-%d%%", p), d) + } + for _, p := range []int{5, 30, 60, 90} { + d := append([]byte(nil), orig...) + for i := n * p / 100; i < n*p/100+512 && i < n; i++ { + d[i] = 0xFF + } + read(fmt.Sprintf("garbage-%d%%", p), d) + } + for off := 0; off < 2048 && off < n; off += 8 { + d := append([]byte(nil), orig...) + d[off] ^= 0xFF + read(fmt.Sprintf("flip@%d", off), d) + } +} + func TestHugeElementSize(t *testing.T) { var buf bytes.Buffer writeEBMLHeader(&buf) @@ -1236,3 +1464,68 @@ func TestTruncatedParsers(t *testing.T) { Read(context.Background(), tr, "trunc.mkv") } } + +// TestChapterRecursionDepthLimit guards against unbounded ChapterAtom recursion +// (a crafted deeply-nested chapters file would otherwise stack-overflow). +func TestChapterRecursionDepthLimit(t *testing.T) { + atom := []byte{} // innermost empty ChapterAtom body + for i := 0; i < 70; i++ { + var wrap bytes.Buffer + ebml.WriteElementHeader(&wrap, mkv.IDChapterAtom, int64(len(atom))) + wrap.Write(atom) + atom = wrap.Bytes() + } + // Chapters > EditionEntry > (70 nested ChapterAtoms) + var edition bytes.Buffer + ebml.WriteElementHeader(&edition, mkv.IDEditionEntry, int64(len(atom))) + edition.Write(atom) + var seg bytes.Buffer + ebml.WriteElementHeader(&seg, mkv.IDChapters, int64(edition.Len())) + seg.Write(edition.Bytes()) + + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(seg.Len())) + buf.Write(seg.Bytes()) + + _, err := Read(context.Background(), bytes.NewReader(buf.Bytes()), "deepchap.mkv") + if err == nil { + t.Fatal("expected error for deep chapter nesting") + } + if !strings.Contains(err.Error(), "chapter nesting") { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestChargeMeta(t *testing.T) { + p := &parser{metaBudget: 100} + if err := p.chargeMeta(60); err != nil { + t.Fatalf("60/100: %v", err) + } + if err := p.chargeMeta(40); err != nil { + t.Fatalf("exactly-at-budget should pass: %v", err) + } + if err := p.chargeMeta(1); err == nil { + t.Fatal("expected budget-exceeded error") + } +} + +func TestMetadataBudgetRejectsHugeAttachment(t *testing.T) { + // AttachedFile whose FileData declares 2 GiB (> 1 GiB budget), no body. + var fd bytes.Buffer + ebml.WriteElementHeader(&fd, mkv.IDFileData, 2<<30) + atts := wrapEl(mkv.IDAttachments, wrapEl(mkv.IDAttachedFile, fd.Bytes())) + + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(len(atts))) + buf.Write(atts) + + _, err := Read(context.Background(), bytes.NewReader(buf.Bytes()), "huge-att.mkv") + if err == nil { + t.Fatal("expected metadata-budget error") + } + if !strings.Contains(err.Error(), "budget") { + t.Fatalf("expected budget error, got: %v", err) + } +} diff --git a/mkv/reader/stream.go b/mkv/reader/stream.go index b095fe8..d744366 100644 --- a/mkv/reader/stream.go +++ b/mkv/reader/stream.go @@ -257,7 +257,9 @@ func (p *streamParser) parseStreamInfo(size int64, c *mkv.Container) error { if err != nil { return err } - c.Info.TimecodeScale = int64(v) + if v > 0 { // keep default; 0 scale would divide-by-zero downstream + c.Info.TimecodeScale = int64(v) + } case mkv.IDDuration: v, err := p.readFloat(h.Size) if err != nil { @@ -376,6 +378,8 @@ func (p *streamParser) parseStreamTrackEntry(size int64) (mkv.Track, error) { return p.parseStreamVideo(h.Size, &t) case mkv.IDAudio: return p.parseStreamAudio(h.Size, &t) + case mkv.IDContentEncodings: + return p.parseStreamContentEncodings(h.Size, &t) default: return p.skip(h.Size) } @@ -384,6 +388,41 @@ func (p *streamParser) parseStreamTrackEntry(size int64) (mkv.Track, error) { return t, err } +// parseStreamContentEncodings mirrors the seekable parser: it extracts the +// header-stripping bytes (ContentCompSettings) so the streaming reader restores +// blocks identically to the seekable one. +func (p *streamParser) parseStreamContentEncodings(size int64, t *mkv.Track) error { + return p.boundedLoop(size, func(h ebml.ElementHeader) error { + if h.ID == mkv.IDContentEncoding { + return p.parseStreamContentEncoding(h.Size, t) + } + return p.skip(h.Size) + }) +} + +func (p *streamParser) parseStreamContentEncoding(size int64, t *mkv.Track) error { + return p.boundedLoop(size, func(h ebml.ElementHeader) error { + if h.ID == mkv.IDContentCompression { + return p.parseStreamContentCompression(h.Size, t) + } + return p.skip(h.Size) + }) +} + +func (p *streamParser) parseStreamContentCompression(size int64, t *mkv.Track) error { + return p.boundedLoop(size, func(h ebml.ElementHeader) error { + if h.ID == mkv.IDContentCompSettings { + v, err := p.readBytes(h.Size) + if err != nil { + return err + } + t.HeaderStripping = v + return nil + } + return p.skip(h.Size) + }) +} + func (p *streamParser) parseStreamVideo(size int64, t *mkv.Track) error { return p.boundedLoop(size, func(h ebml.ElementHeader) error { switch h.ID { @@ -450,7 +489,7 @@ func (p *streamParser) parseStreamChapters(size int64, c *mkv.Container) error { func (p *streamParser) parseStreamEdition(size int64, c *mkv.Container) error { return p.boundedLoop(size, func(h ebml.ElementHeader) error { if h.ID == mkv.IDChapterAtom { - ch, err := p.parseStreamChapterAtom(h.Size) + ch, err := p.parseStreamChapterAtom(h.Size, 0) if err != nil { return err } @@ -461,7 +500,10 @@ func (p *streamParser) parseStreamEdition(size int64, c *mkv.Container) error { }) } -func (p *streamParser) parseStreamChapterAtom(size int64) (mkv.Chapter, error) { +func (p *streamParser) parseStreamChapterAtom(size int64, depth int) (mkv.Chapter, error) { + if depth > maxChapterDepth { + return mkv.Chapter{}, fmt.Errorf("chapter nesting exceeds %d levels", maxChapterDepth) + } ch := mkv.Chapter{} err := p.boundedLoop(size, func(h ebml.ElementHeader) error { switch h.ID { @@ -485,6 +527,12 @@ func (p *streamParser) parseStreamChapterAtom(size int64) (mkv.Chapter, error) { ch.EndMs = int64(v / 1_000_000) case mkv.IDChapterDisplay: return p.parseStreamChapterDisplay(h.Size, &ch) + case mkv.IDChapterAtom: + sub, err := p.parseStreamChapterAtom(h.Size, depth+1) + if err != nil { + return err + } + ch.SubChapters = append(ch.SubChapters, sub) default: return p.skip(h.Size) } @@ -579,7 +627,7 @@ func (p *streamParser) parseStreamTag(size int64) (mkv.Tag, error) { case mkv.IDTargets: return p.parseStreamTargets(h.Size, &tag) case mkv.IDSimpleTag: - st, err := p.parseStreamSimpleTag(h.Size) + st, err := p.parseStreamSimpleTag(h.Size, 0) if err != nil { return err } @@ -614,7 +662,10 @@ func (p *streamParser) parseStreamTargets(size int64, tag *mkv.Tag) error { }) } -func (p *streamParser) parseStreamSimpleTag(size int64) (mkv.SimpleTag, error) { +func (p *streamParser) parseStreamSimpleTag(size int64, depth int) (mkv.SimpleTag, error) { + if depth > maxTagDepth { + return mkv.SimpleTag{}, fmt.Errorf("SimpleTag nesting exceeds %d levels", maxTagDepth) + } st := mkv.SimpleTag{} err := p.boundedLoop(size, func(h ebml.ElementHeader) error { switch h.ID { @@ -636,6 +687,18 @@ func (p *streamParser) parseStreamSimpleTag(size int64) (mkv.SimpleTag, error) { return err } st.Language = v + case mkv.IDTagBinary: + v, err := p.readBytes(h.Size) + if err != nil { + return err + } + st.Binary = v + case mkv.IDSimpleTag: + sub, err := p.parseStreamSimpleTag(h.Size, depth+1) + if err != nil { + return err + } + st.SubTags = append(st.SubTags, sub) default: return p.skip(h.Size) } diff --git a/mkv/reader/stream_parity_test.go b/mkv/reader/stream_parity_test.go new file mode 100644 index 0000000..00044e6 --- /dev/null +++ b/mkv/reader/stream_parity_test.go @@ -0,0 +1,154 @@ +package reader + +import ( + "bytes" + "context" + "testing" + + "github.com/gravity-zero/mkvgo/ebml" + "github.com/gravity-zero/mkvgo/mkv" +) + +// TestStreamSeekableParity verifies the streaming parser (ReadStream) captures +// the same track ContentEncodings (HeaderStripping) and nested SimpleTags as the +// seekable parser (Read) — the drift the audit flagged. +func TestStreamSeekableParity(t *testing.T) { + headerStrip := []byte{0xDE, 0xAD, 0xBE, 0xEF} + + // Tracks > TrackEntry > ContentEncodings > ContentEncoding > ContentCompression > ContentCompSettings + comp := wrapEl(mkv.IDContentCompSettings, headerStrip) + enc := wrapEl(mkv.IDContentCompression, comp) + enc = wrapEl(mkv.IDContentEncoding, enc) + ce := wrapEl(mkv.IDContentEncodings, enc) + + var te bytes.Buffer + ebml.WriteElementHeader(&te, mkv.IDTrackNumber, 1) + ebml.WriteUint(&te, 1, 1) + ebml.WriteElementHeader(&te, mkv.IDTrackType, 1) + ebml.WriteUint(&te, mkv.TrackTypeVideo, 1) + ebml.WriteElementHeader(&te, mkv.IDCodecID, 5) + ebml.WriteString(&te, "V_VP9") + te.Write(ce) + tracks := wrapEl(mkv.IDTracks, wrapEl(mkv.IDTrackEntry, te.Bytes())) + + // Tags > Tag > SimpleTag{ TagName, TagBinary, SimpleTag{ TagName } } + innerST := wrapEl(mkv.IDSimpleTag, mkvTagName("SUB")) + var st bytes.Buffer + st.Write(mkvTagName("OUTER")) + ebml.WriteElementHeader(&st, mkv.IDTagBinary, 2) + ebml.WriteBytes(&st, []byte{0x01, 0x02}) + st.Write(innerST) + tags := wrapEl(mkv.IDTags, wrapEl(mkv.IDTag, wrapEl(mkv.IDSimpleTag, st.Bytes()))) + + var seg bytes.Buffer + ebml.WriteElementHeader(&seg, mkv.IDInfo, 0) + seg.Write(tracks) + seg.Write(tags) + seg.Write(realCluster()) // so ReadStream reaches a cluster and returns + + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(seg.Len())) + buf.Write(seg.Bytes()) + data := buf.Bytes() + + cs, err := Read(context.Background(), bytes.NewReader(data), "x.mkv") + if err != nil { + t.Fatalf("Read: %v", err) + } + cst, _, err := ReadStream(context.Background(), bytes.NewReader(data)) + if err != nil { + t.Fatalf("ReadStream: %v", err) + } + + // HeaderStripping parity (and non-empty in the seekable baseline) + if len(cs.Tracks) != 1 || !bytes.Equal(cs.Tracks[0].HeaderStripping, headerStrip) { + t.Fatalf("seekable HeaderStripping = %v, want %v", trackStrip(cs), headerStrip) + } + if !bytes.Equal(trackStrip(cst), trackStrip(cs)) { + t.Errorf("streaming HeaderStripping = %v, seekable = %v (parity broken)", trackStrip(cst), trackStrip(cs)) + } + + // Nested SimpleTag parity + if nSub(cs) != 1 { + t.Fatalf("seekable nested subtags = %d, want 1", nSub(cs)) + } + if nSub(cst) != nSub(cs) { + t.Errorf("streaming nested subtags = %d, seekable = %d (parity broken)", nSub(cst), nSub(cs)) + } +} + +func wrapEl(id uint32, body []byte) []byte { + var b bytes.Buffer + ebml.WriteElementHeader(&b, id, int64(len(body))) + b.Write(body) + return b.Bytes() +} + +func mkvTagName(s string) []byte { + var b bytes.Buffer + ebml.WriteElementHeader(&b, mkv.IDTagName, int64(len(s))) + ebml.WriteString(&b, s) + return b.Bytes() +} + +func trackStrip(c *mkv.Container) []byte { + if len(c.Tracks) == 0 { + return nil + } + return c.Tracks[0].HeaderStripping +} + +func nSub(c *mkv.Container) int { + if len(c.Tags) == 0 || len(c.Tags[0].SimpleTags) == 0 { + return -1 + } + return len(c.Tags[0].SimpleTags[0].SubTags) +} + +// TestStreamSeekableParitySubChapters verifies the streaming parser captures +// nested ChapterAtoms (SubChapters) identically to the seekable parser. +func TestStreamSeekableParitySubChapters(t *testing.T) { + chapUID := func(id byte) []byte { + var b bytes.Buffer + ebml.WriteElementHeader(&b, mkv.IDChapterUID, 1) + ebml.WriteUint(&b, uint64(id), 1) + return b.Bytes() + } + // ChapterAtom{ UID, ChapterAtom{ UID } } + var outer bytes.Buffer + outer.Write(chapUID(1)) + outer.Write(wrapEl(mkv.IDChapterAtom, chapUID(2))) + chapters := wrapEl(mkv.IDChapters, wrapEl(mkv.IDEditionEntry, wrapEl(mkv.IDChapterAtom, outer.Bytes()))) + + var seg bytes.Buffer + ebml.WriteElementHeader(&seg, mkv.IDInfo, 0) + seg.Write(chapters) + seg.Write(realCluster()) + var buf bytes.Buffer + writeEBMLHeader(&buf) + writeSegmentStart(&buf, int64(seg.Len())) + buf.Write(seg.Bytes()) + data := buf.Bytes() + + cs, err := Read(context.Background(), bytes.NewReader(data), "x.mkv") + if err != nil { + t.Fatalf("Read: %v", err) + } + cst, _, err := ReadStream(context.Background(), bytes.NewReader(data)) + if err != nil { + t.Fatalf("ReadStream: %v", err) + } + nSub := func(c *mkv.Container) int { + if len(c.Chapters) == 0 { + return -1 + } + return len(c.Chapters[0].SubChapters) + } + if nSub(cs) != 1 { + t.Fatalf("seekable sub-chapters = %d, want 1", nSub(cs)) + } + if nSub(cst) != nSub(cs) { + t.Errorf("streaming sub-chapters = %d, seekable = %d (parity broken)", nSub(cst), nSub(cs)) + } +} diff --git a/mkv/reader/testdata/fuzz/FuzzBlockReader/1e19e37543caa16f b/mkv/reader/testdata/fuzz/FuzzBlockReader/1e19e37543caa16f new file mode 100644 index 0000000..4c15643 --- /dev/null +++ b/mkv/reader/testdata/fuzz/FuzzBlockReader/1e19e37543caa16f @@ -0,0 +1,2 @@ +go test fuzz v1 +[]byte("\xa3\x88A0007B00") diff --git a/mkv/webm.go b/mkv/webm.go new file mode 100644 index 0000000..9f8af0a --- /dev/null +++ b/mkv/webm.go @@ -0,0 +1,123 @@ +package mkv + +import ( + "fmt" + "strings" +) + +// WebM is a constrained profile of Matroska: it declares the "webm" DocType and +// permits only a small set of codecs. webmCodecs is that allowlist (the WebVTT +// subtitle family is matched separately, by prefix, in IsWebMCodec). +var webmCodecs = map[string]bool{ + "V_VP8": true, // video + "V_VP9": true, + "V_AV1": true, + "A_VORBIS": true, // audio + "A_OPUS": true, + "S_TEXT/WEBVTT": true, // subtitles (Matroska-style WebVTT id) +} + +// canonicalCodecID maps a codec string to its full Matroska codec ID. mkvgo's +// reader stores the short name (e.g. "vp9") in Track.Codec, while the WebM +// allowlist is keyed by full IDs (e.g. "V_VP9"); this accepts either form. +func canonicalCodecID(codec string) string { + if webmCodecs[codec] { + return codec // already a full ID + } + for full, short := range CodecShortName { + if short == codec { + return full + } + } + return codec +} + +// IsWebMCodec reports whether a codec is permitted in a WebM file. It accepts +// either the full Matroska codec ID ("V_VP9") or mkvgo's short name ("vp9"). +func IsWebMCodec(codec string) bool { + id := canonicalCodecID(codec) + if webmCodecs[id] { + return true + } + // WebM's own subtitle/caption ids: D_WEBVTT/SUBTITLES, /CAPTIONS, /DESCRIPTIONS, /METADATA. + return strings.HasPrefix(id, "D_WEBVTT/") +} + +// WebMDocTypeVersion returns the EBML DocTypeVersion needed to write c as WebM: +// 4 when an AV1 track is present (AV1-in-WebM is a DocTypeVersion-4 feature), +// otherwise 2 (the classic WebM baseline). +func WebMDocTypeVersion(c *Container) uint64 { + for _, t := range c.Tracks { + if canonicalCodecID(t.Codec) == "V_AV1" { + return 4 + } + } + return 2 +} + +// ValidateWebM checks that c can be written as WebM: every track must use a +// codec from the WebM allowlist AND carry the codec initialisation data that +// strict players (browsers) require. mkvgo is a container tool — it cannot +// transcode or synthesise init data — so either is a hard error. It returns an +// error naming each offending track, or nil when c is WebM-compatible. +func ValidateWebM(c *Container) error { + var bad []string + for _, t := range c.Tracks { + if !IsWebMCodec(t.Codec) { + bad = append(bad, fmt.Sprintf("track %d (%s %q): codec outside the WebM subset", t.ID, t.Type, t.Codec)) + continue + } + if msg := webmCodecInitError(t); msg != "" { + bad = append(bad, fmt.Sprintf("track %d (%s): %s", t.ID, t.Codec, msg)) + } + } + if len(bad) > 0 { + return fmt.Errorf("not WebM-compatible: %s (WebM allows only VP8/VP9/AV1 video, Vorbis/Opus audio, WebVTT subtitles)", strings.Join(bad, "; ")) + } + return nil +} + +// WebMNonSubsetElements lists the top-level metadata elements present in c that +// the WebM streaming output does NOT carry — RemuxToWebM / NewWebMStreamWriter +// emit only Info + Tracks + Clusters (+ Cues when seekable). Attachments are not +// part of WebM at all; Chapters and Tags are dropped because the streaming +// writer does not serialise them. An empty result means nothing will be lost. +// Callers can use this to detect data loss before remuxing. +func WebMNonSubsetElements(c *Container) []string { + var out []string + if len(c.Chapters) > 0 { + out = append(out, "Chapters") + } + if len(c.Attachments) > 0 { + out = append(out, "Attachments") + } + if len(c.Tags) > 0 { + out = append(out, "Tags") + } + return out +} + +// webmCodecInitError reports, for codecs whose CodecPrivate is mandatory in +// WebM, why a track's setup data is missing or malformed — or "" if acceptable. +// Browsers reject tracks lacking this initialisation data even when the codec +// itself is allowed. +func webmCodecInitError(t Track) string { + switch canonicalCodecID(t.Codec) { + case "A_OPUS": + if len(t.CodecPrivate) < 8 || string(t.CodecPrivate[:8]) != "OpusHead" { + return "Opus track is missing its OpusHead CodecPrivate" + } + case "A_VORBIS": + if len(t.CodecPrivate) == 0 { + return "Vorbis track is missing its setup-header CodecPrivate" + } + if t.CodecPrivate[0] != 0x02 { + return "Vorbis CodecPrivate is malformed (expected 3 packed setup headers)" + } + case "V_AV1": + if len(t.CodecPrivate) == 0 { + return "AV1 track is missing its av1C CodecPrivate" + } + } + return "" +} diff --git a/mkv/webm_test.go b/mkv/webm_test.go new file mode 100644 index 0000000..064e665 --- /dev/null +++ b/mkv/webm_test.go @@ -0,0 +1,150 @@ +package mkv + +import ( + "strings" + "testing" +) + +// minimal valid-enough codec setup data for the validator's prefix/shape checks. +var ( + opusHead = append([]byte("OpusHead"), 0x01, 0x02, 0x38, 0x01, 0x80, 0xbb, 0x00, 0x00, 0x00, 0x00, 0x00) + vorbisInit = []byte{0x02, 0x00, 0x00, 0x01} // 3 packed headers marker + av1Config = []byte{0x81, 0x00, 0x00, 0x00} // av1C marker byte +) + +func TestIsWebMCodec(t *testing.T) { + tests := []struct { + codec string + want bool + }{ + {"V_VP8", true}, + {"V_VP9", true}, + {"V_AV1", true}, + {"A_VORBIS", true}, + {"A_OPUS", true}, + // short names, as the reader actually stores them in Track.Codec: + {"vp8", true}, + {"vp9", true}, + {"av1", true}, + {"vorbis", true}, + {"opus", true}, + {"h264", false}, + {"aac", false}, + {"S_TEXT/WEBVTT", true}, + {"D_WEBVTT/SUBTITLES", true}, + {"D_WEBVTT/CAPTIONS", true}, + {"V_MPEG4/ISO/AVC", false}, // h264 + {"V_MPEGH/ISO/HEVC", false}, // hevc + {"A_AC3", false}, + {"S_TEXT/ASS", false}, + {"", false}, + } + for _, tt := range tests { + if got := IsWebMCodec(tt.codec); got != tt.want { + t.Errorf("IsWebMCodec(%q) = %v, want %v", tt.codec, got, tt.want) + } + } +} + +func TestValidateWebM(t *testing.T) { + tests := []struct { + name string + tracks []Track + wantErr bool + mentions string // substring expected in the error + }{ + { + name: "vp9 + opus ok", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "V_VP9"}, {ID: 2, Type: AudioTrack, Codec: "A_OPUS", CodecPrivate: opusHead}}, + wantErr: false, + }, + { + name: "av1 with config ok", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "V_AV1", CodecPrivate: av1Config}}, + wantErr: false, + }, + { + name: "vorbis with setup ok", + tracks: []Track{{ID: 1, Type: AudioTrack, Codec: "A_VORBIS", CodecPrivate: vorbisInit}}, + wantErr: false, + }, + { + // short names, as a real parsed Container carries them. + name: "short-name vp9 + opus ok", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "vp9"}, {ID: 2, Type: AudioTrack, Codec: "opus", CodecPrivate: opusHead}}, + wantErr: false, + }, + { + name: "empty container ok", + tracks: nil, + wantErr: false, + }, + { + name: "h264 rejected", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "V_MPEG4/ISO/AVC"}}, + wantErr: true, + mentions: "V_MPEG4/ISO/AVC", + }, + { + name: "ac3 audio rejected", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "V_VP9"}, {ID: 2, Type: AudioTrack, Codec: "A_AC3"}}, + wantErr: true, + mentions: "A_AC3", + }, + { + name: "opus without OpusHead rejected", + tracks: []Track{{ID: 1, Type: AudioTrack, Codec: "A_OPUS"}}, + wantErr: true, + mentions: "OpusHead", + }, + { + name: "vorbis without setup rejected", + tracks: []Track{{ID: 1, Type: AudioTrack, Codec: "A_VORBIS"}}, + wantErr: true, + mentions: "setup-header", + }, + { + name: "av1 without config rejected", + tracks: []Track{{ID: 1, Type: VideoTrack, Codec: "V_AV1"}}, + wantErr: true, + mentions: "av1C", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateWebM(&Container{Tracks: tt.tracks}) + if (err != nil) != tt.wantErr { + t.Fatalf("ValidateWebM err = %v, wantErr %v", err, tt.wantErr) + } + if tt.mentions != "" && !strings.Contains(err.Error(), tt.mentions) { + t.Errorf("error %q does not mention %q", err, tt.mentions) + } + }) + } +} + +func TestWebMNonSubsetElements(t *testing.T) { + clean := &Container{Tracks: []Track{{Codec: "vp9"}}} + if got := WebMNonSubsetElements(clean); len(got) != 0 { + t.Errorf("clean container: got %v, want none", got) + } + dirty := &Container{ + Chapters: []Chapter{{Title: "ch"}}, + Attachments: []Attachment{{}}, + Tags: []Tag{{}}, + } + if got := WebMNonSubsetElements(dirty); len(got) != 3 { + t.Errorf("dirty container: got %v, want 3 (Chapters, Attachments, Tags)", got) + } +} + +func TestWebMDocTypeVersion(t *testing.T) { + vp9opus := &Container{Tracks: []Track{{Codec: "V_VP9"}, {Codec: "A_OPUS"}}} + if got := WebMDocTypeVersion(vp9opus); got != 2 { + t.Errorf("WebMDocTypeVersion(VP9/Opus) = %d, want 2", got) + } + withAV1 := &Container{Tracks: []Track{{Codec: "V_AV1"}, {Codec: "A_OPUS"}}} + if got := WebMDocTypeVersion(withAV1); got != 4 { + t.Errorf("WebMDocTypeVersion(AV1) = %d, want 4", got) + } +} diff --git a/mkv/writer/stream_writer.go b/mkv/writer/stream_writer.go index f03be15..784066c 100644 --- a/mkv/writer/stream_writer.go +++ b/mkv/writer/stream_writer.go @@ -49,12 +49,35 @@ type StreamWriter struct { // precision). The Duration field of info is intentionally ignored: a live // stream has unknown duration. func NewStreamWriter(w io.Writer, info mkv.SegmentInfo, tracks []mkv.Track) (*StreamWriter, error) { + return newStreamWriter(w, info, tracks, WriteEBMLHeader) +} + +// NewWebMStreamWriter is like NewStreamWriter but produces a complete WebM +// stream: it first validates that every track uses a WebM-compatible codec +// (mkv.ValidateWebM), then writes the "webm" DocType header at the right version +// (4 if AV1, else 2). The streaming layout it emits — Info + Tracks + Clusters, +// with no chapters/attachments/tags/SeekHead — is already within the WebM +// element subset, so the result is a real, playable .webm (frames included), +// unlike the metadata-only writer.WriteWebM. +func NewWebMStreamWriter(w io.Writer, info mkv.SegmentInfo, tracks []mkv.Track) (*StreamWriter, error) { + probe := &mkv.Container{Tracks: tracks} + if err := mkv.ValidateWebM(probe); err != nil { + return nil, err + } + version := mkv.WebMDocTypeVersion(probe) + header := func(w io.Writer) error { + return writeEBMLHeaderDocType(w, "webm", version, 2) + } + return newStreamWriter(w, info, tracks, header) +} + +func newStreamWriter(w io.Writer, info mkv.SegmentInfo, tracks []mkv.Track, writeHeader func(io.Writer) error) (*StreamWriter, error) { if info.TimecodeScale <= 0 { info.TimecodeScale = 1_000_000 } // EBML header. - if err := WriteEBMLHeader(w); err != nil { + if err := writeHeader(w); err != nil { return nil, fmt.Errorf("stream writer: EBML header: %w", err) } diff --git a/mkv/writer/webm_stream_test.go b/mkv/writer/webm_stream_test.go new file mode 100644 index 0000000..d6b51e2 --- /dev/null +++ b/mkv/writer/webm_stream_test.go @@ -0,0 +1,74 @@ +package writer + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + + "github.com/gravity-zero/mkvgo/mkv" + "github.com/gravity-zero/mkvgo/mkv/reader" +) + +func TestNewWebMStreamWriter(t *testing.T) { + opusHead := append([]byte("OpusHead"), 0x01, 0x02, 0x38, 0x01, 0x80, 0xbb, 0x00, 0x00, 0x00, 0x00, 0x00) + info := mkv.SegmentInfo{TimecodeScale: 1_000_000} + tracks := []mkv.Track{ + {ID: 1, Type: mkv.VideoTrack, Codec: "vp9"}, + {ID: 2, Type: mkv.AudioTrack, Codec: "opus", CodecPrivate: opusHead}, + } + blocks := []mkv.Block{ + {TrackNumber: 1, Timecode: 0, Keyframe: true, Data: []byte{0x01, 0x02}}, + {TrackNumber: 2, Timecode: 0, Keyframe: false, Data: []byte{0xAA}}, + {TrackNumber: 1, Timecode: 40, Keyframe: false, Data: []byte{0x03}}, + } + + var buf bytes.Buffer + sw, err := NewWebMStreamWriter(&buf, info, tracks) + if err != nil { + t.Fatalf("NewWebMStreamWriter: %v", err) + } + for _, b := range blocks { + if err := sw.WriteBlock(b); err != nil { + t.Fatalf("WriteBlock: %v", err) + } + } + if err := sw.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + data := buf.Bytes() + + if dt, ver := ebmlHeaderInfo(t, data); dt != "webm" || ver != 2 { + t.Errorf("header = (%q, v%d), want (webm, v2)", dt, ver) + } + + // Read the complete stream back (frames included) via the streaming reader. + c, br, err := reader.ReadStream(context.Background(), bytes.NewReader(data)) + if err != nil { + t.Fatalf("ReadStream: %v", err) + } + if len(c.Tracks) != 2 { + t.Errorf("read back %d tracks, want 2", len(c.Tracks)) + } + n := 0 + for { + _, err := br.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("block read: %v", err) + } + n++ + } + if n != len(blocks) { + t.Errorf("read back %d blocks, want %d", n, len(blocks)) + } + + // A non-WebM codec is rejected up front. + bad := []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: "h264"}} + if _, err := NewWebMStreamWriter(io.Discard, info, bad); err == nil { + t.Error("NewWebMStreamWriter accepted a non-WebM codec") + } +} diff --git a/mkv/writer/writer.go b/mkv/writer/writer.go index b58bfdd..5e6dd5b 100644 --- a/mkv/writer/writer.go +++ b/mkv/writer/writer.go @@ -65,6 +65,37 @@ func Write(w io.Writer, c *mkv.Container) error { if err := WriteEBMLHeader(w); err != nil { return err } + return writeSegment(w, c) +} + +// WriteWebM writes c as a WebM file: it first verifies that every track uses a +// WebM-compatible codec (mkv.ValidateWebM), then writes the container with the +// "webm" DocType. If any track is incompatible it returns an error without +// writing anything — mkvgo cannot transcode, so this is a hard failure. +func WriteWebM(w io.Writer, c *mkv.Container) error { + if err := mkv.ValidateWebM(c); err != nil { + return err + } + if err := writeEBMLHeaderDocType(w, "webm", mkv.WebMDocTypeVersion(c), 2); err != nil { + return err + } + // WebM permits only a restricted element set, so write just Info + Tracks + // (no chapters/attachments/tags). NOTE: like writer.Write, this writes + // metadata only — NO clusters. For a complete, playable .webm use + // writer.NewWebMStreamWriter or ops.RemuxToWebM. + var seg bytes.Buffer + if err := WriteSegmentInfo(&seg, &c.Info, c.DurationMs); err != nil { + return err + } + if len(c.Tracks) > 0 { + if err := WriteTracks(&seg, c.Tracks); err != nil { + return err + } + } + return WriteMasterElement(w, mkv.IDSegment, seg.Bytes()) +} + +func writeSegment(w io.Writer, c *mkv.Container) error { var seg bytes.Buffer if err := WriteSegmentInfo(&seg, &c.Info, c.DurationMs); err != nil { return err @@ -134,14 +165,23 @@ func WriteMasterElement(w io.Writer, id uint32, children []byte) error { } func WriteEBMLHeader(w io.Writer) error { + return writeEBMLHeaderDocType(w, "matroska", 4, 2) +} + +// WriteEBMLHeaderWebM writes an EBML header declaring the "webm" DocType. +func WriteEBMLHeaderWebM(w io.Writer) error { + return writeEBMLHeaderDocType(w, "webm", 2, 2) +} + +func writeEBMLHeaderDocType(w io.Writer, docType string, version, readVersion uint64) error { var e ew e.uint(ebml.IDEBMLVersion, 1) e.uint(ebml.IDEBMLReadVersion, 1) e.uint(ebml.IDEBMLMaxIDLength, 4) e.uint(ebml.IDEBMLMaxSizeLength, 8) - e.str(ebml.IDDocType, "matroska") - e.uint(ebml.IDDocTypeVersion, 4) - e.uint(ebml.IDDocTypeReadVersion, 2) + e.str(ebml.IDDocType, docType) + e.uint(ebml.IDDocTypeVersion, version) + e.uint(ebml.IDDocTypeReadVersion, readVersion) return e.flush(w, ebml.IDEBMLHeader) } @@ -378,6 +418,9 @@ func WriteSimpleBlock(w io.Writer, trackNum uint64, relTC int16, keyframe bool, } func WriteCluster(w io.Writer, clusterTS int64, timecodeScale int64, blocks []mkv.Block) error { + if timecodeScale <= 0 { // guard against divide-by-zero from a malformed source + timecodeScale = 1000000 + } rawTS := uint64(clusterTS * 1000000 / timecodeScale) var e ew e.uint(mkv.IDTimestamp, rawTS) @@ -393,6 +436,9 @@ func WriteCluster(w io.Writer, clusterTS int64, timecodeScale int64, blocks []mk } func WriteCues(w io.Writer, cues []mkv.CuePoint, timecodeScale int64) error { + if timecodeScale <= 0 { // guard against divide-by-zero from a malformed source + timecodeScale = 1000000 + } var e ew for i := range cues { cp := &cues[i] diff --git a/mkv/writer/writer_test.go b/mkv/writer/writer_test.go index a2f7a0d..1f0341d 100644 --- a/mkv/writer/writer_test.go +++ b/mkv/writer/writer_test.go @@ -7,10 +7,124 @@ import ( "testing" "time" + "github.com/gravity-zero/mkvgo/ebml" "github.com/gravity-zero/mkvgo/mkv" "github.com/gravity-zero/mkvgo/mkv/reader" ) +// ebmlHeaderInfo parses the EBML header of data and returns its DocType and +// DocTypeVersion. +func ebmlHeaderInfo(t *testing.T, data []byte) (docType string, docTypeVersion uint64) { + t.Helper() + r := bytes.NewReader(data) + h, _, err := ebml.ReadElementHeader(r) + if err != nil || h.ID != ebml.IDEBMLHeader { + t.Fatalf("not an EBML header: id=0x%X err=%v", h.ID, err) + } + body := make([]byte, h.Size) + if _, err := io.ReadFull(r, body); err != nil { + t.Fatal(err) + } + br := bytes.NewReader(body) + for { + eh, _, err := ebml.ReadElementHeader(br) + if err != nil { + break // end of header body + } + switch eh.ID { + case ebml.IDDocType: + if docType, err = ebml.ReadString(br, eh.Size); err != nil { + t.Fatal(err) + } + case ebml.IDDocTypeVersion: + if docTypeVersion, err = ebml.ReadUint(br, eh.Size); err != nil { + t.Fatal(err) + } + default: + br.Seek(eh.Size, io.SeekCurrent) + } + } + return docType, docTypeVersion +} + +// TestWriteZeroTimecodeScaleNoPanic: a malformed source can yield TimecodeScale=0; +// WriteCluster/WriteCues must not divide-by-zero (they default to 1000000). +func TestWriteZeroTimecodeScaleNoPanic(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Fatalf("divide-by-zero panic on TimecodeScale=0: %v", r) + } + }() + var cbuf bytes.Buffer + if err := WriteCluster(&cbuf, 0, 0, []mkv.Block{{TrackNumber: 1, Timecode: 0, Keyframe: true, Data: []byte{1}}}); err != nil { + t.Fatalf("WriteCluster: %v", err) + } + var qbuf bytes.Buffer + if err := WriteCues(&qbuf, []mkv.CuePoint{{TimeMs: 0, Track: 1, ClusterPos: 0}}, 0); err != nil { + t.Fatalf("WriteCues: %v", err) + } +} + +func TestWriteWebM(t *testing.T) { + opusHead := append([]byte("OpusHead"), 0x01, 0x02, 0x38, 0x01, 0x80, 0xbb, 0x00, 0x00, 0x00, 0x00, 0x00) + webmOK := &mkv.Container{ + Info: mkv.SegmentInfo{TimecodeScale: 1000000}, + Tracks: []mkv.Track{ + {ID: 1, Type: mkv.VideoTrack, Codec: "V_VP9"}, + {ID: 2, Type: mkv.AudioTrack, Codec: "A_OPUS", CodecPrivate: opusHead}, + }, + } + + // Default matroska writer keeps the matroska DocType. + var mkvBuf bytes.Buffer + if err := Write(&mkvBuf, webmOK); err != nil { + t.Fatalf("Write: %v", err) + } + if dt, _ := ebmlHeaderInfo(t, mkvBuf.Bytes()); dt != "matroska" { + t.Errorf("Write DocType = %q, want matroska", dt) + } + + // WebM writer emits the webm DocType at version 2 for VP9/Opus. + var webmBuf bytes.Buffer + if err := WriteWebM(&webmBuf, webmOK); err != nil { + t.Fatalf("WriteWebM: %v", err) + } + if dt, ver := ebmlHeaderInfo(t, webmBuf.Bytes()); dt != "webm" || ver != 2 { + t.Errorf("WriteWebM header = (%q, v%d), want (webm, v2)", dt, ver) + } + // And it still reads back as a valid container. + c, err := reader.Read(context.Background(), bytes.NewReader(webmBuf.Bytes()), "out.webm") + if err != nil { + t.Fatalf("read back WebM: %v", err) + } + if len(c.Tracks) != 2 { + t.Errorf("read back %d tracks, want 2", len(c.Tracks)) + } + + // AV1 bumps DocTypeVersion to 4. + av1 := &mkv.Container{ + Info: mkv.SegmentInfo{TimecodeScale: 1000000}, + Tracks: []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: "V_AV1", CodecPrivate: []byte{0x81, 0x00, 0x00, 0x00}}}, + } + var av1Buf bytes.Buffer + if err := WriteWebM(&av1Buf, av1); err != nil { + t.Fatalf("WriteWebM AV1: %v", err) + } + if dt, ver := ebmlHeaderInfo(t, av1Buf.Bytes()); dt != "webm" || ver != 4 { + t.Errorf("WriteWebM AV1 header = (%q, v%d), want (webm, v4)", dt, ver) + } + + // Incompatible codec is rejected and nothing is written. + bad := &mkv.Container{Tracks: []mkv.Track{{ID: 1, Type: mkv.VideoTrack, Codec: "V_MPEG4/ISO/AVC"}}} + var badBuf bytes.Buffer + if err := WriteWebM(&badBuf, bad); err == nil { + t.Fatal("WriteWebM accepted a non-WebM codec") + } + if badBuf.Len() != 0 { + t.Errorf("WriteWebM wrote %d bytes despite rejecting the container", badBuf.Len()) + } +} + func TestWriteAndReadBack(t *testing.T) { w := uint32(1920) h := uint32(1080)