diff --git a/Framework/AODNative/src/Container.cxx b/Framework/AODNative/src/Container.cxx index cf47a59f1da71..5136c8557afcb 100644 --- a/Framework/AODNative/src/Container.cxx +++ b/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 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(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; diff --git a/Framework/AODNative/src/Cursor.h b/Framework/AODNative/src/Cursor.h index fd7e097ad1d65..965cc031537d5 100644 --- a/Framework/AODNative/src/Cursor.h +++ b/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_ diff --git a/Framework/AODNative/src/Decoder.cxx b/Framework/AODNative/src/Decoder.cxx index a98e1267f6efe..cb30eeef7a5db 100644 --- a/Framework/AODNative/src/Decoder.cxx +++ b/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 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 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 offsets(sizeBranch.entries + 1); decodeValues(sizeBranch, (uint8_t*)(offsets.data() + 1), sizeBranch.entries * sizeof(int32_t)); diff --git a/Framework/AODNative/src/Objects.cxx b/Framework/AODNative/src/Objects.cxx index 03048dffbda21..6cca05ebcbfd2 100644 --- a/Framework/AODNative/src/Objects.cxx +++ b/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 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 ObjectReader::tbasket(int64_t mapAt) // TBasket has no versioned header of its own: it starts with its TKey. auto basket = std::make_shared(); 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"); } diff --git a/Framework/AODNative/test/RootFraming.h b/Framework/AODNative/test/RootFraming.h index 24492e61706ba..04b7b240f5546 100644 --- a/Framework/AODNative/test/RootFraming.h +++ b/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 makeBasket(Algo algo, std::span 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 makeBasket(Algo algo, std::span values, int32_t nevbuf, size_t tailBytes = 0, size_t chunkSize = 0xffffff, bool ioBits = false) { std::vector 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 makeBasket(Algo algo, std::span value h.string("TBasket"); // class, name, title h.string("fX"); h.string("O2track"); - h.put(3); // fVersion - h.put(32000); // fBufferSize - h.put(0); // fNevBufSize + h.put(3); // fVersion + h.put(32000); // fBufferSize + h.put(ioBits ? -4 : 0); // fNevBufSize + if (ioBits) { + h.put(0); // fIOBits + } h.put(nevbuf); // fNevBuf h.put(keylen + values.size()); // fLast h.put(tailBytes ? 1 : 2); // flag diff --git a/Framework/AODNative/test/test_Decoder.cxx b/Framework/AODNative/test/test_Decoder.cxx index cfb207ac9e851..17a7b3b3684e6 100644 --- a/Framework/AODNative/test/test_Decoder.cxx +++ b/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 const& values, int32_t nEntries, size_t tail = 0, size_t chunkSize = 0xffffff) + void add(Algo algo, std::vector 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 expected; + for (auto [algo, ioBits] : {std::pair{Algo::ZSTD, true}, std::pair{Algo::None, false}, std::pair{Algo::LZ4, true}}) { + std::vector 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 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 diff --git a/Framework/AnalysisSupport/src/DataInputDirector.cxx b/Framework/AnalysisSupport/src/DataInputDirector.cxx index 48bc6682408de..2c5d362fd1ffc 100644 --- a/Framework/AnalysisSupport/src/DataInputDirector.cxx +++ b/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())); } diff --git a/Framework/AnalysisSupport/src/NativeAOD.cxx b/Framework/AnalysisSupport/src/NativeAOD.cxx index 3f62f7bf03970..b5d92015b718f 100644 --- a/Framework/AnalysisSupport/src/NativeAOD.cxx +++ b/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) { diff --git a/Framework/AnalysisSupport/src/NativeAODPlugin.cxx b/Framework/AnalysisSupport/src/NativeAODPlugin.cxx index bfb054b2ff046..03a0fb0e88dc7 100644 --- a/Framework/AnalysisSupport/src/NativeAODPlugin.cxx +++ b/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()); }