Matroska: Place new elements correctly when AvoidInsert moves an element (#1470)

When a save with WriteStyle::AvoidInsert moved an existing element to the
end of the segment and also added a new element, the new element was
written past the end of the file. New elements are given the old end of
the file as their offset before rendering, and that offset is then shifted
by the size change of each moved element, while the pass which lays out
the moved elements after rendering did not include the new ones.

This happens, for example, when the tags before the first cluster grow
beyond their slot and an attachment is added in the same save. The file
is left with zero bytes between the last element and the new one, the
segment size does not cover the new element, and the seek head points to
the wrong position, so the attachment cannot be read back.

The new elements are now laid out together with the moved elements,
contiguously up to the end of the segment, and their seek head entries are
updated as well.

testSaveTypesMoveAndAdd grows the tags of tags-before-cues.mkv and adds
an attachment in one AvoidInsert save, then checks that the cluster did
not move, that the children of the segment run contiguously to its
declared end, which is the end of the file, and that the tag and the
attachment read back.

taglib/matroska/matroskafile.cpp

@@ -715,9 +715,19 @@ bool Matroska::File::save(WriteStyle writeStyle)
       }
     }
     if(!movedElements.isEmpty()) {
-      // The segment end includes in-place growths AND all moved element sizes.
-      // The moved elements start right after all in-place content.
-      offset_t appendAt = d->segment->endOffset() - totalMovedSize;
+      // The segment end includes in-place growths AND all moved and new
+      // element sizes. New elements were placed at the old end before
+      // rendering and may have been shifted past it by the size changes of
+      // moved elements, so they are laid out again too: new elements first,
+      // then the moved elements, right after all in-place content.
+      offset_t totalNewSize = 0;
+      for(const auto element : newElements)
+        totalNewSize += static_cast<offset_t>(element->data().size());
+      offset_t appendAt = d->segment->endOffset() - totalMovedSize - totalNewSize;
+      for(const auto element : newElements) {
+        element->setOffset(appendAt);
+        appendAt += static_cast<offset_t>(element->data().size());
+      }
       for(const auto element : movedElements) {
         element->setOffset(appendAt);
         appendAt += static_cast<offset_t>(element->data().size());
@@ -725,12 +735,13 @@ bool Matroska::File::save(WriteStyle writeStyle)
     }
   }
 
-  // For elements that were moved to the end by AvoidInsert, update their
-  // seek head entry to reflect the new file position.
+  // For elements that were moved to the end by AvoidInsert, and new elements
+  // laid out again with them, update their seek head entry to reflect the
+  // new file position.
   if(writeStyle == WriteStyle::AvoidInsert && d->seekHead) {
     const offset_t segDataOffset = d->segment->dataOffset();
     for(const auto element : renderList) {
-      if(element->wasMoved()) {
+      if(element->wasMoved() || newElements.contains(element)) {
         d->seekHead->updateEntry(element->id(), element->offset() - segDataOffset);
       }
     }

tests/test_matroska.cpp

@@ -161,6 +161,56 @@ namespace {
     size_t bytesRead = 0;
   };
 
+  //! Reads the EBML variable size integer at \a offset into \a value, without
+  //! its length marker. Returns its length, 0 if it cannot be read.
+  int readVINT(const ByteVector &data, unsigned long long offset,
+               unsigned long long &value)
+  {
+    if(offset >= data.size())
+      return 0;
+    const auto first = static_cast<unsigned char>(data[offset]);
+    int length = 1;
+    while(length <= 8 && !(first & (0x80 >> (length - 1))))
+      ++length;
+    if(length > 8 || offset + length > data.size())
+      return 0;
+    value = first & ((0x80 >> (length - 1)) - 1);
+    for(int i = 1; i < length; ++i)
+      value = (value << 8) | static_cast<unsigned char>(data[offset + i]);
+    return length;
+  }
+
+  //! Walks the children of the Segment by their sizes and returns the offset
+  //! where the last one ends, 0 if a child cannot be parsed. \a segmentEnd is
+  //! set to the end declared by the Segment, \a clusterPos to the offset of
+  //! its first Cluster.
+  unsigned long long segmentChildrenEnd(const ByteVector &data,
+                                        unsigned long long &segmentEnd,
+                                        unsigned long long &clusterPos)
+  {
+    const ByteVector clusterId = ByteVector::fromUInt(0x1F43B675U, true);
+    const int segmentPos = data.find(ByteVector::fromUInt(0x18538067U, true));
+    unsigned long long size = 0;
+    const int sizeLength = segmentPos < 0 ? 0 : readVINT(data, segmentPos + 4, size);
+    if(!sizeLength)
+      return 0;
+    unsigned long long offset = segmentPos + 4 + sizeLength;
+    segmentEnd = offset + size;
+    clusterPos = 0;
+    while(offset < segmentEnd) {
+      unsigned long long id = 0;
+      const int idLength = readVINT(data, offset, id);
+      const int childSizeLength = idLength < 1 || idLength > 4
+        ? 0 : readVINT(data, offset + idLength, size);
+      if(!childSizeLength)
+        return 0;
+      if(clusterPos == 0 && data.containsAt(clusterId, static_cast<unsigned int>(offset)))
+        clusterPos = offset;
+      offset += idLength + childSizeLength + size;
+    }
+    return offset;
+  }
+
 }
 
 class TestMatroska : public CppUnit::TestFixture
@@ -182,6 +232,7 @@ class TestMatroska : public CppUnit::TestFixture
   CPPUNIT_TEST(testSaveTypesBeforeCues);
   CPPUNIT_TEST(testSaveTypesNoTrailingVoid);
   CPPUNIT_TEST(testSaveTypesReclaimVoid);
+  CPPUNIT_TEST(testSaveTypesMoveAndAdd);
   CPPUNIT_TEST(testUnknownSizeSegment);
   CPPUNIT_TEST(testFastReadStyleLargeSegment);
   CPPUNIT_TEST(testAttachedFileDataReadOnDemand);
@@ -1553,6 +1604,58 @@ public:
     }
   }
 
+  void testSaveTypesMoveAndAdd()
+  {
+    // tags-before-cues.mkv layout:
+    //   SeekHead | Void | SegInfo | Tracks | Tags | Cluster | Cues
+    //
+    // With AvoidInsert, growing the Tags moves them to the end of the
+    // segment. A new Attachments element added in the same save must be
+    // placed with them, so that the children of the segment still run
+    // contiguously to its end.
+    ScopedFileCopy copy("tags-before-cues", ".mkv");
+    const string newname = copy.fileName();
+
+    unsigned long long segmentEnd = 0;
+    unsigned long long origClusterPos = 0;
+    const ByteVector origData = PlainFile(TEST_FILE_PATH_C("tags-before-cues.mkv")).readAll();
+    unsigned long long childrenEnd = segmentChildrenEnd(origData, segmentEnd, origClusterPos);
+    CPPUNIT_ASSERT_EQUAL(segmentEnd, childrenEnd);
+    CPPUNIT_ASSERT(origClusterPos > 0);
+
+    const String longTitle = String(ByteVector(1000, 'x'));
+    const ByteVector attachedData(4096, 'a');
+    {
+      Matroska::File f(newname.c_str());
+      CPPUNIT_ASSERT(f.isValid());
+      CPPUNIT_ASSERT(f.tag(false));
+      CPPUNIT_ASSERT(!f.attachments(false));
+      f.tag(false)->setTitle(longTitle);
+      f.attachments(true)->addAttachedFile(Matroska::AttachedFile(
+        attachedData, "cover.jpg", "image/jpeg"));
+      CPPUNIT_ASSERT(f.save(Matroska::WriteStyle::AvoidInsert));
+    }
+    {
+      Matroska::File f(newname.c_str(), true, AudioProperties::Accurate);
+      CPPUNIT_ASSERT(f.isValid());
+      CPPUNIT_ASSERT(f.tag(false));
+      CPPUNIT_ASSERT_EQUAL(longTitle, f.tag(false)->title());
+      const auto attachments = f.attachments(false);
+      CPPUNIT_ASSERT(attachments);
+      const auto &attachedFiles = attachments->attachedFileList();
+      CPPUNIT_ASSERT_EQUAL(1U, attachedFiles.size());
+      CPPUNIT_ASSERT_EQUAL(String("cover.jpg"), attachedFiles.front().fileName());
+      CPPUNIT_ASSERT(attachedData == attachedFiles.front().data());
+    }
+
+    const ByteVector newData = PlainFile(newname.c_str()).readAll();
+    unsigned long long clusterPos = 0;
+    childrenEnd = segmentChildrenEnd(newData, segmentEnd, clusterPos);
+    CPPUNIT_ASSERT_EQUAL(origClusterPos, clusterPos);
+    CPPUNIT_ASSERT_EQUAL(static_cast<unsigned long long>(newData.size()), segmentEnd);
+    CPPUNIT_ASSERT_EQUAL(segmentEnd, childrenEnd);
+  }
+
   void testSaveTypes()
   {
     // Helper lambdas for adding data of different sizes