Fix chunk offsets (#110)

Author: varsill

lib/membrane_mp4/container/header.ex

@@ -4,6 +4,7 @@ defmodule Membrane.MP4.Container.Header do
 
   The `content_size` field is equal to the box size minus the size of the header (8 bytes).
   """
+  use Bunch.Access
 
   @enforce_keys [:name, :content_size, :header_size]
 

lib/membrane_mp4/container/parse_helper.ex

@@ -18,23 +18,24 @@ defmodule Membrane.MP4.Container.ParseHelper do
 
   def parse_boxes(data, schema, context, acc) do
     withl header_content:
-            {:ok, %{name: name, content_size: content_size}, rest} <- Header.parse(data),
+            {:ok, %{name: name, content_size: content_size, header_size: header_size}, rest} <-
+              Header.parse(data),
           header_content: <<content::binary-size(content_size), data::binary>> <- rest,
           do: box_schema = schema[name],
           known?: true <- box_schema && not box_schema.black_box?,
           try:
             {:ok, {fields, rest}, context} <- parse_fields(content, box_schema.fields, context),
           try:
             {:ok, children, <<>>, context} <- parse_boxes(rest, box_schema.children, context, []) do
-      box = %{fields: fields, children: children}
+      box = %{fields: fields, children: children, size: content_size, header_size: header_size}
       parse_boxes(data, schema, context, [{name, box} | acc])
     else
       header_content: _error ->
         # more data needed
         {:ok, Enum.reverse(acc), data, context}
 
       known?: _ ->
-        box = %{content: content}
+        box = %{content: content, size: content_size, header_size: header_size}
         parse_boxes(data, schema, context, [{name, box} | acc])
 
       try: {:error, context} ->

lib/membrane_mp4/container/schema.ex

@@ -289,7 +289,7 @@ defmodule Membrane.MP4.Container.Schema do
                                   {:list,
                                    [
                                      sample_count: :uint32,
-                                     sample_offset: :uint32
+                                     sample_composition_offset: :uint32
                                    ]}
                               ]
                         ],
@@ -448,7 +448,8 @@ defmodule Membrane.MP4.Container.Schema do
                                sample_duration: :uint32,
                                sample_size: :uint32,
                                sample_flags: :bin32,
-                               sample_offset: {:uint32, when: {0x800, :fo_flags, 0x800}}
+                               sample_composition_offset:
+                                 {:uint32, when: {0x800, :fo_flags, 0x800}}
                              ]}
                         ]
                   ]

lib/membrane_mp4/demuxer/isom.ex

@@ -142,7 +142,10 @@ defmodule Membrane.MP4.Demuxer.ISOM do
         } = state
       ) do
     {samples, rest, samples_info} =
-      SamplesInfo.get_samples(state.samples_info, state.partial <> buffer.payload)
+      SamplesInfo.get_samples(
+        state.samples_info,
+        state.partial <> buffer.payload
+      )
 
     buffers = get_buffer_actions(samples)
 
@@ -157,7 +160,10 @@ defmodule Membrane.MP4.Demuxer.ISOM do
       ) do
     # Until all pads are connected we are storing all the samples
     {samples, rest, samples_info} =
-      SamplesInfo.get_samples(state.samples_info, state.partial <> buffer.payload)
+      SamplesInfo.get_samples(
+        state.samples_info,
+        state.partial <> buffer.payload
+      )
 
     state = store_samples(state, samples)
 
@@ -176,40 +182,43 @@ defmodule Membrane.MP4.Demuxer.ISOM do
 
     maybe_header = parse_header(rest)
 
-    state =
-      if maybe_header,
-        do: %{
-          state
-          | mdat_size: maybe_header.content_size,
-            mdat_header_size: maybe_header.header_size
-        },
-        else: state
-
     update_fsm_state_ctx =
       if :mdat in Keyword.keys(state.boxes) or
            (maybe_header != nil and maybe_header.name == :mdat) do
         :started_parsing_mdat
       end
 
-    state = update_fsm_state(state, update_fsm_state_ctx) |> set_partial(rest)
+    state =
+      set_mdat_metadata(state, update_fsm_state_ctx, maybe_header)
+      |> update_fsm_state(update_fsm_state_ctx)
+      |> set_partial(rest)
 
     cond do
       state.fsm_state == :mdat_reading ->
         handle_can_read_mdat_box(ctx, state)
 
       state.optimize_for_non_fast_start? ->
-        state =
-          if state.fsm_state == :skip_mdat,
-            do: %{state | mdat_beginning: state.boxes_size},
-            else: state
-
         handle_non_fast_start_optimization(state)
 
       true ->
         {[], state}
     end
   end
 
+  defp set_mdat_metadata(state, context, maybe_header) do
+    if context == :started_parsing_mdat do
+      %{
+        state
+        | mdat_beginning: state.mdat_beginning || get_mdat_header_beginning(state.boxes),
+          mdat_header_size:
+            state.mdat_header_size || maybe_header[:header_size] || state.boxes[:mdat].header_size,
+          mdat_size: state.mdat_size || maybe_header[:content_size] || state.boxes[:mdat].size
+      }
+    else
+      state
+    end
+  end
+
   defp set_partial(state, rest) do
     partial = if state.fsm_state in [:skip_mdat, :go_back_to_mdat], do: <<>>, else: rest
     %{state | partial: partial}
@@ -285,7 +294,12 @@ defmodule Membrane.MP4.Demuxer.ISOM do
   end
 
   defp handle_non_fast_start_optimization(%{fsm_state: :go_back_to_mdat} = state) do
-    seek(state, state.mdat_beginning, state.mdat_size + state.mdat_header_size, false)
+    seek(
+      state,
+      state.mdat_beginning,
+      state.mdat_size + state.mdat_header_size,
+      false
+    )
   end
 
   defp handle_non_fast_start_optimization(state) do
@@ -315,7 +329,14 @@ defmodule Membrane.MP4.Demuxer.ISOM do
       end
 
     state =
-      %{state | samples_info: SamplesInfo.get_samples_info(state.boxes[:moov])}
+      %{
+        state
+        | samples_info:
+            SamplesInfo.get_samples_info(
+              state.boxes[:moov],
+              state.mdat_beginning + state.mdat_header_size
+            )
+      }
       |> update_fsm_state()
 
     # Parse the data we received so far (partial or the whole mdat box in a single buffer) and
@@ -330,7 +351,9 @@ defmodule Membrane.MP4.Demuxer.ISOM do
         content
       end
 
-    {samples, rest, samples_info} = SamplesInfo.get_samples(state.samples_info, data)
+    {samples, rest, samples_info} =
+      SamplesInfo.get_samples(state.samples_info, data)
+
     state = %{state | samples_info: samples_info, partial: rest}
 
     all_pads_connected? = all_pads_connected?(ctx, state)
@@ -472,4 +495,16 @@ defmodule Membrane.MP4.Demuxer.ISOM do
       {:end_of_stream, pad_ref}
     end)
   end
+
+  defp get_mdat_header_beginning([]) do
+    0
+  end
+
+  defp get_mdat_header_beginning([{:mdat, _box} | _rest]) do
+    0
+  end
+
+  defp get_mdat_header_beginning([{_other_name, box} | rest]) do
+    box.header_size + box.size + get_mdat_header_beginning(rest)
+  end
 end

lib/membrane_mp4/demuxer/isom/samples_info.ex

@@ -15,7 +15,8 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
     :tracks_number,
     :timescales,
     :last_dts,
-    :sample_tables
+    :sample_tables,
+    :mdat_iterator
   ]
 
   defstruct @enforce_keys
@@ -42,7 +43,8 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
             (track_id :: pos_integer()) => last_dts :: Ratio.t() | nil
           },
           tracks_number: pos_integer(),
-          sample_tables: %{(track_id :: pos_integer()) => SampleTable.t()}
+          sample_tables: %{(track_id :: pos_integer()) => SampleTable.t()},
+          mdat_iterator: non_neg_integer()
         }
 
   @doc """
@@ -51,9 +53,10 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
   a whole sample, and has yet to be parsed.
   """
   @spec get_samples(t, data :: binary()) ::
-          {[{Buffer.t(), track_id :: pos_integer()}], rest :: binary, t}
+          {[{Buffer.t(), track_id :: pos_integer()}], rest :: binary(), t()}
   def get_samples(samples_info, data) do
-    {samples_info, rest, buffers} = do_get_samples(samples_info, data, [])
+    {samples_info, rest, buffers} =
+      do_get_samples(samples_info, data, [])
 
     {buffers, rest, samples_info}
   end
@@ -63,24 +66,32 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
   end
 
   defp do_get_samples(samples_info, data, buffers) do
-    [%{size: size, track_id: track_id} = sample | samples] = samples_info.samples
+    [%{size: size, track_id: track_id, sample_offset: sample_offset} = sample | samples] =
+      samples_info.samples
 
-    if size <= byte_size(data) do
-      <<payload::binary-size(size), rest::binary>> = data
+    to_skip = sample_offset - samples_info.mdat_iterator
 
-      {dts, pts, samples_info} = get_dts_and_pts(samples_info, sample)
+    case data do
+      <<_to_skip::binary-size(to_skip), payload::binary-size(size), rest::binary>> ->
+        {dts, pts, samples_info} = get_dts_and_pts(samples_info, sample)
 
-      buffer =
-        {%Buffer{
-           payload: payload,
-           dts: dts,
-           pts: pts
-         }, track_id}
+        buffer =
+          {%Buffer{
+             payload: payload,
+             dts: dts,
+             pts: pts
+           }, track_id}
 
-      samples_info = %{samples_info | samples: samples}
-      do_get_samples(samples_info, rest, [buffer | buffers])
-    else
-      {samples_info, data, Enum.reverse(buffers)}
+        samples_info = %{samples_info | samples: samples}
+
+        do_get_samples(
+          %{samples_info | mdat_iterator: samples_info.mdat_iterator + to_skip + size},
+          rest,
+          [buffer | buffers]
+        )
+
+      _other ->
+        {samples_info, data, Enum.reverse(buffers)}
     end
   end
 
@@ -118,8 +129,8 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
   present in the `mdat` box.
   The list of samples in the returned struct is used to extract data from the `mdat` box and get output buffers.
   """
-  @spec get_samples_info(%{children: boxes :: Container.t()}) :: t
-  def get_samples_info(%{children: boxes}) do
+  @spec get_samples_info(%{children: boxes :: Container.t()}, non_neg_integer()) :: t
+  def get_samples_info(%{children: boxes}, mdat_beginning) do
     tracks =
       boxes
       |> Enum.filter(fn {type, _content} -> type == :trak end)
@@ -159,15 +170,18 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
            :decoding_deltas,
            :sample_sizes,
            :samples_per_chunk,
-           :composition_offsets
+           :composition_offsets,
+           :chunk_offset
          ])}
       end)
 
     # Create a samples' description list for each chunk and flatten it
     {samples, _acc} =
       chunk_offsets
       |> Enum.flat_map_reduce(tracks_data, fn %{track_id: track_id} = chunk, tracks_data ->
-        {new_samples, track_data} = get_chunk_samples(chunk, tracks_data[track_id])
+        {new_samples, {track_data, _sample_offset}} =
+          get_chunk_samples(chunk, tracks_data[track_id])
+
         {new_samples, %{tracks_data | track_id => track_data}}
       end)
 
@@ -183,19 +197,28 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
       tracks_number: map_size(tracks),
       timescales: timescales,
       last_dts: last_dts,
-      sample_tables: sample_tables
+      sample_tables: sample_tables,
+      mdat_iterator: mdat_beginning
     }
   end
 
   defp get_chunk_samples(chunk, track_data) do
-    %{chunk_no: chunk_no, track_id: track_id} = chunk
+    %{chunk_no: chunk_no, track_id: track_id, chunk_offset: chunk_offset} = chunk
 
     {track_data, samples_no} = get_samples_no(chunk_no, track_data)
 
-    Enum.map_reduce(1..samples_no, track_data, fn _no, track_data ->
+    Enum.map_reduce(1..samples_no, {track_data, chunk_offset}, fn _no,
+                                                                  {track_data, sample_offset} ->
       {sample, track_data} = get_sample_description(track_data)
-      sample = Map.put(sample, :track_id, track_id)
-      {sample, track_data}
+
+      sample =
+        Map.merge(sample, %{
+          track_id: track_id,
+          chunk_offset: chunk_offset,
+          sample_offset: sample_offset
+        })
+
+      {sample, {track_data, sample_offset + sample.size}}
     end)
   end
 
@@ -250,14 +273,19 @@ defmodule Membrane.MP4.Demuxer.ISOM.SamplesInfo do
 
     {sample_composition_offset, composition_offsets} =
       case composition_offsets do
-        [%{sample_count: 1, sample_offset: offset} | composition_offsets] ->
+        [%{sample_count: 1, sample_composition_offset: offset} | composition_offsets] ->
           {offset, composition_offsets}
 
-        [%{sample_count: count, sample_offset: offset} | composition_offsets] ->
-          {offset, [%{sample_count: count - 1, sample_offset: offset} | composition_offsets]}
+        [%{sample_count: count, sample_composition_offset: offset} | composition_offsets] ->
+          {offset,
+           [%{sample_count: count - 1, sample_composition_offset: offset} | composition_offsets]}
       end
 
-    {%{size: size, sample_delta: delta, sample_composition_offset: sample_composition_offset},
+    {%{
+       size: size,
+       sample_delta: delta,
+       sample_composition_offset: sample_composition_offset
+     },
      %{
        track_data
        | decoding_deltas: deltas,

lib/membrane_mp4/movie_box/sample_table_box.ex

@@ -255,7 +255,7 @@ defmodule Membrane.MP4.MovieBox.SampleTableBox do
       # if no :ctts box is available, assume that the offset between
       # composition time and the decoding time is equal to 0
       Enum.map(boxes[:stts].fields.entry_list, fn entry ->
-        %{sample_count: entry.sample_count, sample_offset: 0}
+        %{sample_count: entry.sample_count, sample_composition_offset: 0}
       end)
     end
   end

lib/membrane_mp4/muxer/cmaf.ex

@@ -568,7 +568,7 @@ defmodule Membrane.MP4.Muxer.CMAF do
           sample.metadata.duration
           |> Helper.timescalify(timescale)
           |> Ratio.trunc(),
-        sample_offset: Helper.timescalify(sample.pts - sample.dts, timescale)
+        sample_composition_offset: Helper.timescalify(sample.pts - sample.dts, timescale)
       }
     end)
   end

lib/membrane_mp4/track/sample_table.ex

@@ -28,7 +28,7 @@ defmodule Membrane.MP4.Track.SampleTable do
           ],
           composition_offsets: [
             %{
-              sample_offset: Ratio.t(),
+              sample_composition_offset: Ratio.t(),
               sample_count: pos_integer
             }
           ],