diff --git a/src/OrcCommand/GetThis.h b/src/OrcCommand/GetThis.h index 01f9e7ea..3f9ba444 100644 --- a/src/OrcCommand/GetThis.h +++ b/src/OrcCommand/GetThis.h @@ -279,7 +279,7 @@ class ORCUTILS_API Main : public UtilitiesMain std::shared_ptr m_tableWriter; - HRESULT ConfigureSampleStreams(SampleRef& sampleRef); + HRESULT ConfigureSampleStreams(SampleRef& sample) const; HRESULT AddSampleRefToCSV(ITableOutput& output, const SampleRef& sampleRef) const; diff --git a/src/OrcCommand/GetThis_Run.cpp b/src/OrcCommand/GetThis_Run.cpp index 11275bbf..98a97f68 100644 --- a/src/OrcCommand/GetThis_Run.cpp +++ b/src/OrcCommand/GetThis_Run.cpp @@ -415,6 +415,34 @@ LimitStatus SampleLimitStatus(const Limits& globalLimits, const Limits& localLim return SampleWithinLimits; } +std::unique_ptr ConfigureStringStream( + const std::shared_ptr dataStream, + const ContentSpec& sampleSpec, + const ContentSpec& configSpec) +{ + size_t minChars, maxChars; + if (sampleSpec.MaxChars != 0 || sampleSpec.MinChars != 0) + { + minChars = sampleSpec.MinChars; + maxChars = sampleSpec.MaxChars; + } + else + { + maxChars = configSpec.MaxChars; + minChars = configSpec.MinChars; + } + + auto stream = std::make_unique(); + HRESULT hr = stream->OpenForStrings(dataStream, minChars, maxChars); + if (FAILED(hr)) + { + Log::Error(L"Failed to initialise strings stream"); + return {}; + } + + return stream; +} + HRESULT RegFlushKeys(logger& _L_) { bool bSuccess = true; @@ -452,88 +480,58 @@ Main::Main(logger pLog) { } -HRESULT Main::ConfigureSampleStreams(SampleRef& sampleRef) +HRESULT Main::ConfigureSampleStreams(SampleRef& sample) const { HRESULT hr = E_FAIL; - shared_ptr retval; + _ASSERT(sample.Matches.front()->MatchingAttributes[sample.AttributeIndex].DataStream->IsOpen() == S_OK); - _ASSERT(sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].DataStream->IsOpen() == S_OK); + auto& dataStream = sample.Matches.front()->MatchingAttributes[sample.AttributeIndex].DataStream; - if (sampleRef.SampleName.empty()) - return E_INVALIDARG; + // Stream are initially at eof std::shared_ptr stream; - - switch (sampleRef.Content.Type) + if (sample.Content.Type == ContentType::STRINGS) { - case ContentType::DATA: - stream = sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].DataStream; - break; - case ContentType::STRINGS: { - auto strings = std::make_shared(_L_); - if (sampleRef.Content.MaxChars == 0 && sampleRef.Content.MinChars == 0) - { - if (FAILED( - hr = strings->OpenForStrings( - sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].DataStream, - config.content.MinChars, - config.content.MaxChars))) - { - log::Error(_L_, hr, L"Failed to initialise strings stream\r\n"); - return hr; - } - } - else - { - if (FAILED( - hr = strings->OpenForStrings( - sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].DataStream, - sampleRef.Content.MinChars, - sampleRef.Content.MaxChars))) - { - log::Error(_L_, hr, L"Failed to initialise strings stream\r\n"); - return hr; - } - } - stream = strings; + stream = ::ConfigureStringStream(dataStream, sample.Content, config.content); + if (stream == nullptr) + { + return E_FAIL; } - break; - case ContentType::RAW: - stream = sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].RawStream; - break; - default: - stream = sampleRef.Matches.front()->MatchingAttributes[sampleRef.AttributeIndex].DataStream; - break; + } + else + { + stream = dataStream; } - std::shared_ptr upstream = stream; - - CryptoHashStream::Algorithm algs = config.CryptoHashAlgs; - + const auto algs = config.CryptoHashAlgs; if (algs != CryptoHashStream::Algorithm::Undefined) { - sampleRef.HashStream = make_shared(_L_); - if (FAILED(hr = sampleRef.HashStream->OpenToRead(algs, upstream))) + sample.HashStream = std::make_shared(); + hr = sample.HashStream->OpenToRead(algs, stream); + if (FAILED(hr)) + { return hr; - upstream = sampleRef.HashStream; - } - else - { - upstream = stream; + } + + stream = sample.HashStream; } - FuzzyHashStream::Algorithm fuzzy_algs = config.FuzzyHashAlgs; - if (fuzzy_algs != FuzzyHashStream::Algorithm::Undefined) + const auto fuzzyAlgs = config.FuzzyHashAlgs; + if (fuzzyAlgs != FuzzyHashStream::Algorithm::Undefined) { - sampleRef.FuzzyHashStream = make_shared(_L_); - if (FAILED(hr = sampleRef.FuzzyHashStream->OpenToRead(fuzzy_algs, upstream))) + sample.FuzzyHashStream = std::make_shared(); + hr = sample.FuzzyHashStream->OpenToRead(fuzzyAlgs, stream); + if (FAILED(hr)) + { return hr; - upstream = sampleRef.FuzzyHashStream; + } + + stream = sample.FuzzyHashStream; } - sampleRef.CopyStream = upstream; - sampleRef.SampleSize = sampleRef.CopyStream->GetSize(); + sample.CopyStream = stream; + sample.SampleSize = sample.CopyStream->GetSize(); return S_OK; }