DPL Analysis: harden the native AOD reader against corrupt input (#16031)

Author: ktf

Framework/AODNative/src/Container.cxx

@@ -27,8 +27,6 @@ constexpr size_t kHeaderReadSize = 128;
 constexpr size_t kDirectoryReadSize = 64;
 // Largest gap between the TTrees of a DF still read together with them.
 constexpr uint64_t kMetadataGap = 4 * 1024;
-// Fixed part of a TBasket key, after which come the class name, name and title.
-constexpr size_t kBasketHeaderTail = 19;
 // Smallest possible key header: 32-bit seeks and three empty strings.
 constexpr size_t kMinKeySize = 29;
 // Metadata objects (trees, maps) in AO2Ds are a few MB at most. A larger
@@ -184,6 +182,9 @@ std::vector<std::string> RootFile::dataFrames()
   for (auto& k : topKeys()) {
     std::smatch m;
     if (std::regex_match(k.name, m, tfRegex)) {
+      if (m[1].length() > 19) {
+        throw CorruptData(mUrl, k.seekKey, "DF number of " + k.name + " does not fit 64 bits");
+      }
       auto n = std::stoull(m[1].str());
       if (seen.insert(n).second) {
         numbers.push_back(n);
@@ -325,13 +326,8 @@ TreeInfo RootFile::readTree(const KeyInfo& key)
         continue;
       }
       auto buf = readRange(*mSource, basket.seek, std::min<uint64_t>(basket.nbytes, kHeaderReadSize * 4));
-      auto k = parseKey(buf, basket.seek);
       Cursor c{buf, mUrl, basket.seek};
-      c.seek(k.keylen - kBasketHeaderTail);
-      c.u16(); // fVersion
-      c.i32(); // fBufferSize
-      c.i32(); // fNevBufSize
-      basket.nEntries = c.i32();
+      basket.nEntries = readBasketHeader(c).nevbuf;
     }
   }
   return tree;

Framework/AODNative/src/Cursor.h

@@ -127,6 +127,48 @@ class Cursor
   size_t mPos = 0;
 };
 
+/// The key and header of a TBasket.
+struct BasketHeader {
+  uint32_t nbytes;
+  uint32_t objlen;
+  uint16_t keylen;
+  uint16_t version;
+  int32_t nevbuf;
+  int32_t last;
+  uint8_t flag;
+};
+
+/// Reads a TBasket key and header forwards from @a c, as TBasket::Streamer
+/// writes them: a negative fNevBufSize is followed by one more byte, fIOBits.
+/// Leaves @a c at the end of the key.
+inline BasketHeader readBasketHeader(Cursor& c)
+{
+  auto start = c.pos();
+  BasketHeader h;
+  h.nbytes = c.u32();
+  bool big = c.i16() > 1000;
+  h.objlen = c.u32();
+  c.u32(); // fDatime
+  h.keylen = c.u16();
+  c.u16();              // fCycle
+  c.skip(big ? 16 : 8); // fSeekKey, fSeekPdir
+  c.string();           // fClassName
+  c.string();           // fName
+  c.string();           // fTitle
+  h.version = c.u16();
+  c.i32(); // fBufferSize
+  if (c.i32() < 0) {
+    c.u8(); // fIOBits
+  }
+  h.nevbuf = c.i32();
+  h.last = c.i32();
+  h.flag = c.u8();
+  if (c.pos() != start + h.keylen) {
+    c.fail("TBasket key of " + std::to_string(c.pos() - start) + " bytes, but fKeylen is " + std::to_string(h.keylen));
+  }
+  return h;
+}
+
 } // namespace o2::aod::native
 
 #endif // O2_FRAMEWORK_AODNATIVE_CURSOR_H_

Framework/AODNative/src/Decoder.cxx

@@ -28,8 +28,6 @@ namespace o2::aod::native
 
 namespace
 {
-// Fixed part of a TBasket key, after which come the class name, name and title.
-constexpr size_t kBasketHeaderTail = 19;
 
 template <typename T>
 void byteswapAll(uint8_t* data, size_t count)
@@ -132,21 +130,12 @@ void BranchDecoder::forEachBasket(BranchInfo const& branch, bool parallel, F&& f
         bytes = storage;
       }
       Cursor c{bytes, mUrl, ref.seek};
-      auto nbytes = c.u32();
-      c.i16(); // key fVersion
-      auto objlen = c.u32();
-      c.u32(); // fDatime
-      auto keylen = c.u16();
-      if (nbytes != ref.nbytes || keylen < kBasketHeaderTail || keylen > nbytes) {
-        c.fail("basket header of branch " + branch.name + " does not match its TBranch");
-      }
-      c.seek(keylen - kBasketHeaderTail);
-      c.u16(); // fVersion
-      c.i32(); // fBufferSize
-      c.i32(); // fNevBufSize
-      auto nevbuf = c.i32();
-      auto last = c.i32();
-      if (nevbuf != ref.nEntries || last < keylen || (uint32_t)(last - keylen) > objlen) {
+      auto header = readBasketHeader(c);
+      auto nbytes = header.nbytes;
+      auto objlen = header.objlen;
+      auto keylen = header.keylen;
+      auto last = header.last;
+      if (nbytes != ref.nbytes || keylen > nbytes || header.nevbuf != ref.nEntries || last < keylen || (uint32_t)(last - keylen) > objlen) {
         c.fail("basket header of branch " + branch.name + " does not match its TBranch");
       }
       values.border = last - keylen;
@@ -278,6 +267,8 @@ std::vector<int32_t> BranchDecoder::decodeOffsets(BranchInfo const& sizeBranch)
   if (sizeBranch.leafClass != 'I' || sizeBranch.fixedLen != 1) {
     throw UnsupportedLayout("size branch " + sizeBranch.name + " is not a scalar 32-bit integer");
   }
+  // fEntries sizes the allocation: check it against the baskets first.
+  checkBaskets(sizeBranch);
   // Decode the sizes after the leading 0, then prefix-sum in place.
   std::vector<int32_t> offsets(sizeBranch.entries + 1);
   decodeValues(sizeBranch, (uint8_t*)(offsets.data() + 1), sizeBranch.entries * sizeof(int32_t));

Framework/AODNative/src/Objects.cxx

@@ -27,8 +27,6 @@ constexpr uint32_t kNewClassTag = 0xFFFFFFFF;
 constexpr uint16_t kStreamedMemberWise = 0x4000;
 constexpr uint32_t kIsReferenced = 1 << 4;
 constexpr int64_t kMapOffset = 2;
-// Fixed part of a TBasket key, after which come the class name, name and title.
-constexpr size_t kBasketHeaderTail = 19;
 
 struct LeafInfo {
   std::string name;
@@ -256,6 +254,9 @@ std::shared_ptr<LeafInfo> ObjectReader::tleaf(std::string const& className, int6
   leaf->className = className;
   leaf->name = tnamed();
   leaf->len = mC.i32();
+  if (leaf->len <= 0) {
+    mC.fail("leaf " + leaf->name + " with length " + std::to_string(leaf->len));
+  }
   mC.i32(); // fLenType
   mC.i32(); // fOffset
   mC.u8();  // fIsRange
@@ -277,23 +278,12 @@ std::shared_ptr<EmbeddedBasket> ObjectReader::tbasket(int64_t mapAt)
   // TBasket has no versioned header of its own: it starts with its TKey.
   auto basket = std::make_shared<EmbeddedBasket>();
   mObjects[mapAt] = basket;
-  auto start = mC.pos();
-  mC.i32(); // fNbytes
-  mC.i16(); // key fVersion
-  mC.i32(); // fObjlen
-  mC.u32(); // fDatime
-  auto keylen = mC.u16();
-  mC.u16(); // fCycle
-  if (keylen < kBasketHeaderTail) {
-    mC.fail("embedded basket with key length " + std::to_string(keylen));
-  }
-  mC.seek(start + keylen - kBasketHeaderTail);
-  checkVersion("TBasket", mC.u16() % 1000);
-  mC.i32(); // fBufferSize
-  mC.i32(); // fNevBufSize
-  auto nevbuf = mC.i32();
-  auto last = mC.i32();
-  int flag = mC.u8();
+  auto header = readBasketHeader(mC);
+  checkVersion("TBasket", header.version % 1000);
+  auto keylen = header.keylen;
+  auto nevbuf = header.nevbuf;
+  auto last = header.last;
+  int flag = header.flag;
   if (nevbuf < 0 || last < keylen) {
     mC.fail("inconsistent embedded basket header");
   }

Framework/AODNative/test/RootFraming.h

@@ -110,8 +110,10 @@ struct BigEndianWriter {
 };
 
 /// A TBasket key as written on disk: key header, TBasket fields, then the
-/// payload holding @a values followed by a fake entry-offset tail.
-inline std::vector<uint8_t> makeBasket(Algo algo, std::span<const uint8_t> values, int32_t nevbuf, size_t tailBytes = 0, size_t chunkSize = 0xffffff)
+/// payload holding @a values followed by a fake entry-offset tail. With
+/// @a ioBits, the header has the layout of TTree::kGenerateOffsetMap: a
+/// negative fNevBufSize followed by one fIOBits byte.
+inline std::vector<uint8_t> makeBasket(Algo algo, std::span<const uint8_t> values, int32_t nevbuf, size_t tailBytes = 0, size_t chunkSize = 0xffffff, bool ioBits = false)
 {
   std::vector<uint8_t> raw(values.begin(), values.end());
   // The tail is poisoned: a decoder that reads past the value region sees it.
@@ -132,9 +134,12 @@ inline std::vector<uint8_t> makeBasket(Algo algo, std::span<const uint8_t> value
     h.string("TBasket"); // class, name, title
     h.string("fX");
     h.string("O2track");
-    h.put<uint16_t>(3);                     // fVersion
-    h.put<int32_t>(32000);                  // fBufferSize
-    h.put<int32_t>(0);                      // fNevBufSize
+    h.put<uint16_t>(3);              // fVersion
+    h.put<int32_t>(32000);           // fBufferSize
+    h.put<int32_t>(ioBits ? -4 : 0); // fNevBufSize
+    if (ioBits) {
+      h.put<uint8_t>(0); // fIOBits
+    }
     h.put<int32_t>(nevbuf);                 // fNevBuf
     h.put<int32_t>(keylen + values.size()); // fLast
     h.put<uint8_t>(tailBytes ? 1 : 2);      // flag

Framework/AODNative/test/test_Decoder.cxx

@@ -45,9 +45,9 @@ struct BranchBuilder {
     branch.entries = branch.totBytes = branch.zipBytes = 0;
   }
 
-  void add(Algo algo, std::vector<uint8_t> const& values, int32_t nEntries, size_t tail = 0, size_t chunkSize = 0xffffff)
+  void add(Algo algo, std::vector<uint8_t> const& values, int32_t nEntries, size_t tail = 0, size_t chunkSize = 0xffffff, bool ioBits = false)
   {
-    auto basket = makeBasket(algo, values, nEntries, tail, chunkSize);
+    auto basket = makeBasket(algo, values, nEntries, tail, chunkSize, ioBits);
     branch.baskets.push_back(BasketRef{file.size(), (uint32_t)basket.size(), branch.entries, nEntries, false, {}});
     file.insert(file.end(), basket.begin(), basket.end());
     branch.entries += nEntries;
@@ -103,6 +103,25 @@ TEST_CASE("float values across codecs, uncompressed and embedded baskets")
   }
 }
 
+TEST_CASE("baskets with fIOBits in their header decode like the others")
+{
+  BranchBuilder b{'F'};
+  std::vector<float> expected;
+  for (auto [algo, ioBits] : {std::pair{Algo::ZSTD, true}, std::pair{Algo::None, false}, std::pair{Algo::LZ4, true}}) {
+    std::vector<float> values;
+    for (int i = 0; i < 100; ++i) {
+      values.push_back(expected.size() + i);
+    }
+    b.add(algo, bigEndian(values), values.size(), 404, 0xffffff, ioBits);
+    expected.insert(expected.end(), values.begin(), values.end());
+  }
+  MemoryByteSource source{b.file};
+  BranchDecoder decoder{source, Options{}};
+  std::vector<float> out(expected.size(), 42.f);
+  decoder.decodeValues(b.branch, (uint8_t*)out.data(), out.size() * sizeof(float));
+  CHECK(out == expected);
+}
+
 TEST_CASE("decoding stops exactly at the value region for every codec")
 {
   // A poisoned entry-offset tail after the values: nothing past fLast - fKeylen

Framework/AnalysisSupport/src/DataInputDirector.cxx

@@ -161,7 +161,7 @@ bool DataInputDescriptor::setFile(int counter, int wantedParentLevel, std::strin
   if (backend != "root") {
     try {
       mCurrentFile = openNativeAOD(filename, mFactory, mContext.nativeOptions);
-    } catch (o2::aod::native::FormatError const& e) {
+    } catch (std::exception const& e) {
       if (backend == "native") {
         throw AODReadError(AODReadError::Kind::FileOpen, fmt::format("Couldn't open file \"{}\" with the native reader: {}", filename, e.what()));
       }

Framework/AnalysisSupport/src/NativeAOD.cxx

@@ -206,7 +206,7 @@ void NativeAODFileSystem::addToPlan(std::string const& treeName)
   NativeTreeHandle* handle = nullptr;
   try {
     handle = findTree(mPlanDF, treeName);
-  } catch (FormatError const&) {
+  } catch (std::exception const&) {
     // Reported when the tree itself is asked for.
   }
   if (handle) {

Framework/AnalysisSupport/src/NativeAODPlugin.cxx

@@ -324,7 +324,9 @@ struct NativeAODObjectReadingCapability : o2::framework::RootObjectReadingCapabi
         }
         try {
           return nativeFS->getTreeHandle(path);
-        } catch (o2::aod::native::FormatError const& e) {
+        } catch (AODReadError const&) {
+          throw;
+        } catch (std::exception const& e) {
           // A DF which cannot be parsed: --aod-skip-invalid-reads skips it.
           throw AODReadError(AODReadError::Kind::InvalidTimeframe, e.what());
         }