stremio/node/media-node/internal/mediainfo/matroska.go

package mediainfo

import (
	"encoding/binary"
	"io"
	"os"
	"strings"
)

// Track is a media stream with optional Matroska language/title.
type Track struct {
	Type     string `json:"type"` // video | audio | subtitle | unknown
	Number   int    `json:"number"`
	Language string `json:"language,omitempty"`
	Title    string `json:"title,omitempty"`
	Codec    string `json:"codec,omitempty"`
}

// ProbeFile returns embedded tracks. For MKV/WebM we parse EBML; other containers return nil.
func ProbeFile(path string) ([]Track, error) {
	lower := strings.ToLower(path)
	if strings.HasSuffix(lower, ".mkv") || strings.HasSuffix(lower, ".webm") || strings.HasSuffix(lower, ".mka") {
		return probeMatroska(path)
	}
	return nil, nil
}

func probeMatroska(path string) ([]Track, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		return nil, err
	}
	// Tracks element is normally near the start; cap read for huge files.
	limit := info.Size()
	if limit > 8<<20 {
		limit = 8 << 20
	}
	r := &ebmlReader{r: io.LimitReader(f, limit)}

	id, _, err := r.readIDSize()
	if err != nil {
		return nil, err
	}
	// EBML header
	if id != 0x1A45DFA3 {
		return nil, nil
	}
	if err := r.skip(r.lastSize); err != nil {
		return nil, err
	}

	var tracks []Track
	for {
		id, size, err := r.readIDSize()
		if err == io.EOF {
			break
		}
		if err != nil {
			break
		}
		switch id {
		case 0x18538067: // Segment — descend
			continue
		case 0x1654AE6B: // Tracks
			tracks, err = parseTracksElement(r, size)
			return tracks, err
		case 0x114D9B74, 0x1549A966, 0x1C53BB6B, 0x1F43B675, 0x1043A770, 0x1254C367, 0x1941A469:
			// SeekHead, Info, Cues, Cluster, Chapters, Tags, Attachments
			if err := r.skip(size); err != nil {
				return tracks, nil
			}
		default:
			if err := r.skip(size); err != nil {
				return tracks, nil
			}
		}
	}
	return tracks, nil
}

func parseTracksElement(r *ebmlReader, size uint64) ([]Track, error) {
	end := r.pos + size
	var out []Track
	for r.pos < end {
		id, sz, err := r.readIDSize()
		if err != nil {
			return out, nil
		}
		if id == 0xAE { // TrackEntry
			t, err := parseTrackEntry(r, sz)
			if err == nil && t.Type != "" {
				out = append(out, t)
			}
		} else {
			_ = r.skip(sz)
		}
	}
	return out, nil
}

func parseTrackEntry(r *ebmlReader, size uint64) (Track, error) {
	end := r.pos + size
	t := Track{Type: "unknown"}
	for r.pos < end {
		id, sz, err := r.readIDSize()
		if err != nil {
			return t, err
		}
		switch id {
		case 0xD7: // TrackNumber
			v, err := r.readUint(sz)
			if err != nil {
				return t, err
			}
			t.Number = int(v)
		case 0x83: // TrackType
			v, err := r.readUint(sz)
			if err != nil {
				return t, err
			}
			switch v {
			case 1:
				t.Type = "video"
			case 2:
				t.Type = "audio"
			case 17:
				t.Type = "subtitle"
			default:
				t.Type = "unknown"
			}
		case 0x22B59C, 0x22B59D: // Language / LanguageIETF
			s, err := r.readString(sz)
			if err != nil {
				return t, err
			}
			s = strings.TrimSpace(s)
			if s != "" && !strings.EqualFold(s, "und") {
				t.Language = s
			}
		case 0x536E: // Name
			s, err := r.readString(sz)
			if err != nil {
				return t, err
			}
			t.Title = strings.TrimSpace(s)
		case 0x86: // CodecID
			s, err := r.readString(sz)
			if err != nil {
				return t, err
			}
			t.Codec = s
		default:
			if err := r.skip(sz); err != nil {
				return t, err
			}
		}
	}
	return t, nil
}

type ebmlReader struct {
	r        io.Reader
	pos      uint64
	lastSize uint64
	buf      [8]byte
}

func (e *ebmlReader) readByte() (byte, error) {
	n, err := e.r.Read(e.buf[:1])
	if n == 1 {
		e.pos++
		return e.buf[0], nil
	}
	if err == nil {
		err = io.EOF
	}
	return 0, err
}

func (e *ebmlReader) readIDSize() (id uint32, size uint64, err error) {
	b0, err := e.readByte()
	if err != nil {
		return 0, 0, err
	}
	var idLen int
	switch {
	case b0&0x80 != 0:
		idLen = 1
		id = uint32(b0)
	case b0&0x40 != 0:
		idLen = 2
		id = uint32(b0)
	case b0&0x20 != 0:
		idLen = 3
		id = uint32(b0)
	case b0&0x10 != 0:
		idLen = 4
		id = uint32(b0)
	default:
		return 0, 0, io.ErrUnexpectedEOF
	}
	for i := 1; i < idLen; i++ {
		b, err := e.readByte()
		if err != nil {
			return 0, 0, err
		}
		id = (id << 8) | uint32(b)
	}

	b0, err = e.readByte()
	if err != nil {
		return 0, 0, err
	}
	var sizeLen int
	var mask byte
	switch {
	case b0&0x80 != 0:
		sizeLen, mask = 1, 0x7F
	case b0&0x40 != 0:
		sizeLen, mask = 2, 0x3F
	case b0&0x20 != 0:
		sizeLen, mask = 3, 0x1F
	case b0&0x10 != 0:
		sizeLen, mask = 4, 0x0F
	case b0&0x08 != 0:
		sizeLen, mask = 5, 0x07
	case b0&0x04 != 0:
		sizeLen, mask = 6, 0x03
	case b0&0x02 != 0:
		sizeLen, mask = 7, 0x01
	case b0&0x01 != 0:
		sizeLen, mask = 8, 0x00
	default:
		return 0, 0, io.ErrUnexpectedEOF
	}
	size = uint64(b0 & mask)
	for i := 1; i < sizeLen; i++ {
		b, err := e.readByte()
		if err != nil {
			return 0, 0, err
		}
		size = (size << 8) | uint64(b)
	}
	// Unknown size marker
	if size == (uint64(1)<<(7*sizeLen))-1 {
		size = 0
	}
	e.lastSize = size
	return id, size, nil
}

func (e *ebmlReader) skip(n uint64) error {
	if n == 0 {
		return nil
	}
	buf := make([]byte, 32*1024)
	left := n
	for left > 0 {
		chunk := uint64(len(buf))
		if chunk > left {
			chunk = left
		}
		read, err := io.ReadFull(e.r, buf[:chunk])
		e.pos += uint64(read)
		if err != nil {
			return err
		}
		left -= chunk
	}
	return nil
}

func (e *ebmlReader) readUint(n uint64) (uint64, error) {
	if n == 0 || n > 8 {
		return 0, io.ErrUnexpectedEOF
	}
	var raw [8]byte
	if _, err := io.ReadFull(e.r, raw[8-n:]); err != nil {
		return 0, err
	}
	e.pos += n
	return binary.BigEndian.Uint64(raw[:]), nil
}

func (e *ebmlReader) readString(n uint64) (string, error) {
	if n > 1<<20 {
		return "", io.ErrUnexpectedEOF
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(e.r, buf); err != nil {
		return "", err
	}
	e.pos += n
	// Trim trailing NULs common in Matroska strings
	return strings.TrimRight(string(buf), "\x00"), nil
}