OrcLib: CacheStream: use reference for underlying stream

fabienfl-orc · fabienfl-orc · commit d7b534997ce5 · 2024-04-25T16:42:48.000+02:00
diff --git a/src/OrcLib/Authenticode.cpp b/src/OrcLib/Authenticode.cpp
@@ -665,7 +665,8 @@ HRESULT Authenticode::Verify(LPCWSTR szFileName, const std::shared_ptr<ByteStrea
     HRESULT hr = E_FAIL;
     std::error_code ec;
 
-    PeParser pe(std::make_shared<CacheStream>(pStream), ec);
+    CacheStream cache(*pStream, 1048576);
+    PeParser pe(cache, ec);
     if (ec)
     {
         Log::Debug(L"Failed to parse pe file (filename: {}) [{}]", szFileName, ec);
diff --git a/src/OrcLib/CacheStream.cpp b/src/OrcLib/CacheStream.cpp
@@ -14,9 +14,9 @@ using namespace std::string_view_literals;
 
 namespace Orc {
 
-CacheStream::CacheStream(std::shared_ptr<ByteStream> stream, size_t cacheSize)
+CacheStream::CacheStream(ByteStream& stream, size_t cacheSize)
     : ByteStream()
-    , m_stream(std::move(stream))
+    , m_stream(stream)
     , m_streamOffset(0)
     , m_offset(0)
     , m_cache()
@@ -27,14 +27,14 @@ CacheStream::CacheStream(std::shared_ptr<ByteStream> stream, size_t cacheSize)
     m_cache.resize(cacheSize);
 }
 
-STDMETHODIMP Orc::CacheStream::Clone(std::shared_ptr<ByteStream>& clone)
+STDMETHODIMP Orc::CacheStream::Clone(ByteStream& clone)
 {
     return E_NOTIMPL;
 }
 
 HRESULT CacheStream::Close()
 {
-    HRESULT hr = m_stream->Close();
+    HRESULT hr = m_stream.Close();
     if (FAILED(hr))
     {
         Log::Debug("Failed to close underlying cached stream [{}]", SystemError(hr));
@@ -59,7 +59,7 @@ HRESULT CacheStream::Duplicate(const CacheStream& other)
 
 HRESULT CacheStream::Open()
 {
-    HRESULT hr = m_stream->IsOpen();
+    HRESULT hr = m_stream.IsOpen();
     if (hr != S_OK)
     {
         Log::Debug("Failed to open CacheStream: underlying stream is closed [{}]", SystemError(hr));
@@ -107,7 +107,7 @@ HRESULT CacheStream::Read_(
         else
         {
             ULONGLONG streamRead = 0;
-            hr = m_stream->Read(m_cache.data(), m_cache.size(), &streamRead);
+            hr = m_stream.Read(m_cache.data(), m_cache.size(), &streamRead);
             if (FAILED(hr))
             {
                 return hr;
@@ -149,7 +149,7 @@ CacheStream::SetFilePointer(__in LONGLONG DistanceToMove, __in DWORD dwMoveMetho
         pCurrPointer = &newOffset;
     }
 
-    HRESULT hr = m_stream->SetFilePointer(DistanceToMove, dwMoveMethod, pCurrPointer);
+    HRESULT hr = m_stream.SetFilePointer(DistanceToMove, dwMoveMethod, pCurrPointer);
     if (FAILED(hr))
     {
         return hr;
@@ -162,7 +162,7 @@ CacheStream::SetFilePointer(__in LONGLONG DistanceToMove, __in DWORD dwMoveMetho
 
 ULONG64 CacheStream::GetSize()
 {
-    return m_stream->GetSize();
+    return m_stream.GetSize();
 }
 
 HRESULT CacheStream::SetSize(ULONG64 ullNewSize)
diff --git a/src/OrcLib/CacheStream.h b/src/OrcLib/CacheStream.h
@@ -10,7 +10,8 @@
 
 #include "OrcLib.h"
 
-#include "MemoryStream.h"
+#include "ByteStream.h"
+#include "Utils/MetaPtr.h"
 
 #pragma managed(push, off)
 
@@ -21,17 +22,17 @@ class CBinaryBuffer;
 class CacheStream : public ByteStream
 {
 public:
-    CacheStream(std::shared_ptr<ByteStream> stream, size_t cacheSize = 1048576);
+    CacheStream(ByteStream& stream, size_t cacheSize);
     ~CacheStream();
     HRESULT Open();
 
     void Accept(ByteStreamVisitor& visitor) override { return visitor.Visit(*this); };
 
-    STDMETHOD(IsOpen)() { return m_stream->IsOpen(); };
+    STDMETHOD(IsOpen)() { return m_stream.IsOpen(); };
 
-    STDMETHOD(CanRead)() { return m_stream->CanRead(); };
+    STDMETHOD(CanRead)() { return m_stream.CanRead(); };
     STDMETHOD(CanWrite)() { return S_FALSE; };
-    STDMETHOD(CanSeek)() { return m_stream->CanSeek(); };
+    STDMETHOD(CanSeek)() { return m_stream.CanSeek(); };
 
     //
     // CByteStream implementation
@@ -52,13 +53,13 @@ class CacheStream : public ByteStream
     STDMETHOD_(ULONG64, GetSize)();
     STDMETHOD(SetSize)(ULONG64 ullSize);
 
-    STDMETHOD(Clone)(std::shared_ptr<ByteStream>& clone);
+    STDMETHOD(Clone)(ByteStream& clone);
     STDMETHOD(Close)();
 
     HRESULT Duplicate(const CacheStream& other);
 
 private:
-    std::shared_ptr<ByteStream> m_stream;
+    ByteStream& m_stream;
     uint64_t m_streamOffset;
     uint64_t m_offset;
     std::vector<uint8_t> m_cache;
diff --git a/src/OrcLib/FileFormat/PeParser.cpp b/src/OrcLib/FileFormat/PeParser.cpp
@@ -193,23 +193,23 @@ void ParseImageSections(
 
 namespace Orc {
 
-PeParser::PeParser(std::shared_ptr<ByteStream> stream, std::error_code& ec)
-    : m_stream(std::move(stream))
+PeParser::PeParser(ByteStream& stream, std::error_code& ec)
+    : m_stream(stream)
 {
-    ::ParseImageDosHeader(*m_stream, m_imageDosHeader, ec);
+    ::ParseImageDosHeader(m_stream, m_imageDosHeader, ec);
     if (ec)
     {
         return;
     }
 
     ::ParseImageNtHeader(
-        *m_stream, m_imageDosHeader, m_imageNtHeader, m_imageOptionalHeaders32, m_imageOptionalHeaders64, ec);
+        m_stream, m_imageDosHeader, m_imageNtHeader, m_imageOptionalHeaders32, m_imageOptionalHeaders64, ec);
     if (ec)
     {
         return;
     }
 
-    ::ParseImageSections(*m_stream, m_imageNtHeader, m_imageSectionsHeaders, ec);
+    ::ParseImageSections(m_stream, m_imageNtHeader, m_imageSectionsHeaders, ec);
     if (ec)
     {
         return;
@@ -313,7 +313,7 @@ void PeParser::ReadDirectory(uint8_t index, std::vector<uint8_t>& buffer, std::e
     }
 
     buffer.resize(directory.Size);
-    ReadChunkAt(*m_stream, directory.VirtualAddress, buffer, ec);
+    ReadChunkAt(m_stream, directory.VirtualAddress, buffer, ec);
     if (ec)
     {
         Log::Debug("Failed to read directory (index: {}) [{}]", index, ec);
@@ -403,7 +403,7 @@ void PeParser::GetHashedChunks(PeChunks& chunks, std::error_code& ec) const
     }
 
     chunks[3].offset = GetSizeOfOptionalHeaders();
-    chunks[3].length = m_stream->GetSize() - chunks[3].offset - secdir.Size;
+    chunks[3].length = m_stream.GetSize() - chunks[3].offset - secdir.Size;
 }
 
 void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks, PeHash& output, std::error_code& ec)
@@ -418,7 +418,7 @@ void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks
         return;
     }
 
-    const auto written = ::CopyChunks(chunks, *m_stream, *hashstream, ec);
+    const auto written = ::CopyChunks(chunks, m_stream, *hashstream, ec);
     if (ec)
     {
         Log::Debug("Failed to hash chunks [{}]", ec);
@@ -427,7 +427,7 @@ void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks
 
     // MS does zero padding for PEs that are not 8 modulo (often with catalogs)
     const uint8_t padding[8] = {0, 0, 0, 0, 0, 0, 0, 0};
-    const auto alignment = m_stream->GetSize() % sizeof(padding);
+    const auto alignment = m_stream.GetSize() % sizeof(padding);
     if (alignment != 0)
     {
         const auto paddingLength = sizeof(padding) - alignment;
diff --git a/src/OrcLib/FileFormat/PeParser.h b/src/OrcLib/FileFormat/PeParser.h
@@ -44,7 +44,7 @@ class PeParser
     using PeChunks = std::array<PeChunk, 4>;
     using SectionHeaders = fmt::basic_memory_buffer<IMAGE_SECTION_HEADER, 16>;
 
-    PeParser(std::shared_ptr<ByteStream> stream, std::error_code& ec);
+    PeParser(ByteStream& stream, std::error_code& ec);
 
     bool HasDebugDirectory() const;
     void ReadDebugDirectory(std::vector<uint8_t>& buffer, std::error_code& ec) const;
@@ -66,7 +66,7 @@ class PeParser
     void Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks, PeHash& output, std::error_code& ec) const;
 
 private:
-    mutable std::shared_ptr<ByteStream> m_stream;
+    ByteStream& m_stream;
     IMAGE_DOS_HEADER m_imageDosHeader;
     ImageNtHeader m_imageNtHeader;
     std::optional<IMAGE_OPTIONAL_HEADER32> m_imageOptionalHeaders32;
diff --git a/src/OrcLib/PEInfo.cpp b/src/OrcLib/PEInfo.cpp
@@ -307,7 +307,7 @@ HRESULT PEInfo::OpenVersionInformation()
         return E_POINTER;
     }
 
-    auto stream = std::make_shared<CacheStream>(directStream, 4096);
+    auto stream = std::make_shared<CacheStream>(*directStream, 4096);
 
     ULONGLONG ullBytesRead;
     size_t rsrc_rsrc_offset = 0;
@@ -648,7 +648,7 @@ HRESULT PEInfo::OpenAllHash(Intentions localIntentions)
     // Pe Hash
     {
         std::error_code ec;
-        PeParser pe(memstream, ec);
+        PeParser pe(*memstream, ec);
         if (ec)
         {
             Log::Debug(L"Failed to parse pe hash '{}' [{}]", m_FileInfo.m_szFullName, ec);
@@ -723,7 +723,8 @@ HRESULT PEInfo::OpenPeHash(Intentions localIntentions)
         return E_POINTER;
 
     std::error_code ec;
-    PeParser pe(std::make_shared<CacheStream>(std::move(stream)), ec);
+    CacheStream cache(*stream, 1048576);
+    PeParser pe(cache, ec);
     if (ec)
     {
         Log::Error(L"Failed to parse PE '{}' [{}]", m_FileInfo.m_szFullName, ec);

Original file line number	Diff line number	Diff line change
`@@ -665,7 +665,8 @@ HRESULT Authenticode::Verify(LPCWSTR szFileName, const std::shared_ptr<ByteStrea`
`665`	`665`	`HRESULT hr = E_FAIL;`
`666`	`666`	`std::error_code ec;`
`667`	`667`
`668`		`- PeParser pe(std::make_shared<CacheStream>(pStream), ec);`
	`668`	`+ CacheStream cache(*pStream, 1048576);`
	`669`	`+ PeParser pe(cache, ec);`
`669`	`670`	`if (ec)`
`670`	`671`	`{`
`671`	`672`	`Log::Debug(L"Failed to parse pe file (filename: {}) [{}]", szFileName, ec);`
Original file line number	Diff line number	Diff line change
`@@ -14,9 +14,9 @@ using namespace std::string_view_literals;`
`14`	`14`
`15`	`15`	`namespace Orc {`
`16`	`16`
`17`		`-CacheStream::CacheStream(std::shared_ptr<ByteStream> stream, size_t cacheSize)`
	`17`	`+CacheStream::CacheStream(ByteStream& stream, size_t cacheSize)`
`18`	`18`	`: ByteStream()`
`19`		`- , m_stream(std::move(stream))`
	`19`	`+ , m_stream(stream)`
`20`	`20`	`, m_streamOffset(0)`
`21`	`21`	`, m_offset(0)`
`22`	`22`	`, m_cache()`
`@@ -27,14 +27,14 @@ CacheStream::CacheStream(std::shared_ptr<ByteStream> stream, size_t cacheSize)`
`27`	`27`	`m_cache.resize(cacheSize);`
`28`	`28`	`}`
`29`	`29`
`30`		`-STDMETHODIMP Orc::CacheStream::Clone(std::shared_ptr<ByteStream>& clone)`
	`30`	`+STDMETHODIMP Orc::CacheStream::Clone(ByteStream& clone)`
`31`	`31`	`{`
`32`	`32`	`return E_NOTIMPL;`
`33`	`33`	`}`
`34`	`34`
`35`	`35`	`HRESULT CacheStream::Close()`
`36`	`36`	`{`
`37`		`- HRESULT hr = m_stream->Close();`
	`37`	`+ HRESULT hr = m_stream.Close();`
`38`	`38`	`if (FAILED(hr))`
`39`	`39`	`{`
`40`	`40`	`Log::Debug("Failed to close underlying cached stream [{}]", SystemError(hr));`
`@@ -59,7 +59,7 @@ HRESULT CacheStream::Duplicate(const CacheStream& other)`
`59`	`59`
`60`	`60`	`HRESULT CacheStream::Open()`
`61`	`61`	`{`
`62`		`- HRESULT hr = m_stream->IsOpen();`
	`62`	`+ HRESULT hr = m_stream.IsOpen();`
`63`	`63`	`if (hr != S_OK)`
`64`	`64`	`{`
`65`	`65`	`Log::Debug("Failed to open CacheStream: underlying stream is closed [{}]", SystemError(hr));`
`@@ -107,7 +107,7 @@ HRESULT CacheStream::Read_(`
`107`	`107`	`else`
`108`	`108`	`{`
`109`	`109`	`ULONGLONG streamRead = 0;`
`110`		`- hr = m_stream->Read(m_cache.data(), m_cache.size(), &streamRead);`
	`110`	`+ hr = m_stream.Read(m_cache.data(), m_cache.size(), &streamRead);`
`111`	`111`	`if (FAILED(hr))`
`112`	`112`	`{`
`113`	`113`	`return hr;`
`@@ -149,7 +149,7 @@ CacheStream::SetFilePointer(__in LONGLONG DistanceToMove, __in DWORD dwMoveMetho`
`149`	`149`	`pCurrPointer = &newOffset;`
`150`	`150`	`}`
`151`	`151`
`152`		`- HRESULT hr = m_stream->SetFilePointer(DistanceToMove, dwMoveMethod, pCurrPointer);`
	`152`	`+ HRESULT hr = m_stream.SetFilePointer(DistanceToMove, dwMoveMethod, pCurrPointer);`
`153`	`153`	`if (FAILED(hr))`
`154`	`154`	`{`
`155`	`155`	`return hr;`
`@@ -162,7 +162,7 @@ CacheStream::SetFilePointer(__in LONGLONG DistanceToMove, __in DWORD dwMoveMetho`
`162`	`162`
`163`	`163`	`ULONG64 CacheStream::GetSize()`
`164`	`164`	`{`
`165`		`- return m_stream->GetSize();`
	`165`	`+ return m_stream.GetSize();`
`166`	`166`	`}`
`167`	`167`
`168`	`168`	`HRESULT CacheStream::SetSize(ULONG64 ullNewSize)`
Original file line number	Diff line number	Diff line change
`@@ -193,23 +193,23 @@ void ParseImageSections(`
`193`	`193`
`194`	`194`	`namespace Orc {`
`195`	`195`
`196`		`-PeParser::PeParser(std::shared_ptr<ByteStream> stream, std::error_code& ec)`
`197`		`- : m_stream(std::move(stream))`
	`196`	`+PeParser::PeParser(ByteStream& stream, std::error_code& ec)`
	`197`	`+ : m_stream(stream)`
`198`	`198`	`{`
`199`		`- ::ParseImageDosHeader(*m_stream, m_imageDosHeader, ec);`
	`199`	`+ ::ParseImageDosHeader(m_stream, m_imageDosHeader, ec);`
`200`	`200`	`if (ec)`
`201`	`201`	`{`
`202`	`202`	`return;`
`203`	`203`	`}`
`204`	`204`
`205`	`205`	`::ParseImageNtHeader(`
`206`		`- *m_stream, m_imageDosHeader, m_imageNtHeader, m_imageOptionalHeaders32, m_imageOptionalHeaders64, ec);`
	`206`	`+ m_stream, m_imageDosHeader, m_imageNtHeader, m_imageOptionalHeaders32, m_imageOptionalHeaders64, ec);`
`207`	`207`	`if (ec)`
`208`	`208`	`{`
`209`	`209`	`return;`
`210`	`210`	`}`
`211`	`211`
`212`		`- ::ParseImageSections(*m_stream, m_imageNtHeader, m_imageSectionsHeaders, ec);`
	`212`	`+ ::ParseImageSections(m_stream, m_imageNtHeader, m_imageSectionsHeaders, ec);`
`213`	`213`	`if (ec)`
`214`	`214`	`{`
`215`	`215`	`return;`
`@@ -313,7 +313,7 @@ void PeParser::ReadDirectory(uint8_t index, std::vector<uint8_t>& buffer, std::e`
`313`	`313`	`}`
`314`	`314`
`315`	`315`	`buffer.resize(directory.Size);`
`316`		`- ReadChunkAt(*m_stream, directory.VirtualAddress, buffer, ec);`
	`316`	`+ ReadChunkAt(m_stream, directory.VirtualAddress, buffer, ec);`
`317`	`317`	`if (ec)`
`318`	`318`	`{`
`319`	`319`	`Log::Debug("Failed to read directory (index: {}) [{}]", index, ec);`
`@@ -403,7 +403,7 @@ void PeParser::GetHashedChunks(PeChunks& chunks, std::error_code& ec) const`
`403`	`403`	`}`
`404`	`404`
`405`	`405`	`chunks[3].offset = GetSizeOfOptionalHeaders();`
`406`		`- chunks[3].length = m_stream->GetSize() - chunks[3].offset - secdir.Size;`
	`406`	`+ chunks[3].length = m_stream.GetSize() - chunks[3].offset - secdir.Size;`
`407`	`407`	`}`
`408`	`408`
`409`	`409`	`void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks, PeHash& output, std::error_code& ec)`
`@@ -418,7 +418,7 @@ void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks`
`418`	`418`	`return;`
`419`	`419`	`}`
`420`	`420`
`421`		`- const auto written = ::CopyChunks(chunks, m_stream, hashstream, ec);`
	`421`	`+ const auto written = ::CopyChunks(chunks, m_stream, *hashstream, ec);`
`422`	`422`	`if (ec)`
`423`	`423`	`{`
`424`	`424`	`Log::Debug("Failed to hash chunks [{}]", ec);`
`@@ -427,7 +427,7 @@ void PeParser::Hash(CryptoHashStreamAlgorithm algorithms, const PeChunks& chunks`
`427`	`427`
`428`	`428`	`// MS does zero padding for PEs that are not 8 modulo (often with catalogs)`
`429`	`429`	`const uint8_t padding[8] = {0, 0, 0, 0, 0, 0, 0, 0};`
`430`		`- const auto alignment = m_stream->GetSize() % sizeof(padding);`
	`430`	`+ const auto alignment = m_stream.GetSize() % sizeof(padding);`
`431`	`431`	`if (alignment != 0)`
`432`	`432`	`{`
`433`	`433`	`const auto paddingLength = sizeof(padding) - alignment;`
Original file line number	Diff line number	Diff line change
`@@ -307,7 +307,7 @@ HRESULT PEInfo::OpenVersionInformation()`
`307`	`307`	`return E_POINTER;`
`308`	`308`	`}`
`309`	`309`
`310`		`- auto stream = std::make_shared<CacheStream>(directStream, 4096);`
	`310`	`+ auto stream = std::make_shared<CacheStream>(*directStream, 4096);`
`311`	`311`
`312`	`312`	`ULONGLONG ullBytesRead;`
`313`	`313`	`size_t rsrc_rsrc_offset = 0;`
`@@ -648,7 +648,7 @@ HRESULT PEInfo::OpenAllHash(Intentions localIntentions)`
`648`	`648`	`// Pe Hash`
`649`	`649`	`{`
`650`	`650`	`std::error_code ec;`
`651`		`- PeParser pe(memstream, ec);`
	`651`	`+ PeParser pe(*memstream, ec);`
`652`	`652`	`if (ec)`
`653`	`653`	`{`
`654`	`654`	`Log::Debug(L"Failed to parse pe hash '{}' [{}]", m_FileInfo.m_szFullName, ec);`
`@@ -723,7 +723,8 @@ HRESULT PEInfo::OpenPeHash(Intentions localIntentions)`
`723`	`723`	`return E_POINTER;`
`724`	`724`
`725`	`725`	`std::error_code ec;`
`726`		`- PeParser pe(std::make_shared<CacheStream>(std::move(stream)), ec);`
	`726`	`+ CacheStream cache(*stream, 1048576);`
	`727`	`+ PeParser pe(cache, ec);`
`727`	`728`	`if (ec)`
`728`	`729`	`{`
`729`	`730`	`Log::Error(L"Failed to parse PE '{}' [{}]", m_FileInfo.m_szFullName, ec);`