diff --git a/src/Frag.cc b/src/Frag.cc index 3e3300efdd..4d1f7ace47 100644 --- a/src/Frag.cc +++ b/src/Frag.cc @@ -189,7 +189,7 @@ void FragReassembler::BlockInserted(DataBlockMap::const_iterator /* it */) { auto it = block_list.Begin(); - if ( it->second->seq > 0 || ! frag_size ) + if ( it->second.seq > 0 || ! frag_size ) // For sure don't have it all yet. return; @@ -198,19 +198,19 @@ void FragReassembler::BlockInserted(DataBlockMap::const_iterator /* it */) // We might have it all - look for contiguous all the way. while ( next != block_list.End() ) { - if ( it->second->upper != next->second->seq ) + if ( it->second.upper != next->second.seq ) break; ++it; ++next; } - auto last = std::prev(block_list.End())->second; + const auto& last = std::prev(block_list.End())->second; if ( next != block_list.End() ) { // We have a hole. - if ( it->second->upper >= frag_size ) + if ( it->second.upper >= frag_size ) { // We're stuck. The point where we stopped is // contiguous up through the expected end of @@ -223,19 +223,19 @@ void FragReassembler::BlockInserted(DataBlockMap::const_iterator /* it */) // We decide to analyze the contiguous portion now. // Extend the fragment up through the end of what // we have. - frag_size = it->second->upper; + frag_size = it->second.upper; } else return; } - else if ( last->upper > frag_size ) + else if ( last.upper > frag_size ) { Weird("fragment_size_inconsistency"); - frag_size = last->upper; + frag_size = last.upper; } - else if ( last->upper < frag_size ) + else if ( last.upper < frag_size ) // Missing the tail. return; @@ -258,16 +258,20 @@ void FragReassembler::BlockInserted(DataBlockMap::const_iterator /* it */) for ( it = block_list.Begin(); it != block_list.End(); ++it ) { - auto b = it->second; - DataBlock* prev = it == block_list.Begin() ? nullptr : std::prev(it)->second; + const auto& b = it->second; - // If we're above a hole, stop. This can happen because - // the logic above regarding a hole that's above the - // expected fragment size. - if ( prev && prev->upper < b->seq ) - break; + if ( it != block_list.Begin() ) + { + const auto& prev = std::prev(it)->second; - if ( b->upper > n ) + // If we're above a hole, stop. This can happen because + // the logic above regarding a hole that's above the + // expected fragment size. + if ( prev.upper < b.seq ) + break; + } + + if ( b.upper > n ) { reporter->InternalWarning("bad fragment reassembly"); DeleteTimer(); @@ -276,7 +280,7 @@ void FragReassembler::BlockInserted(DataBlockMap::const_iterator /* it */) return; } - memcpy(&pkt[b->seq], b->block, b->upper - b->seq); + memcpy(&pkt[b.seq], b.block, b.upper - b.seq); } delete reassembled_pkt; diff --git a/src/Reassem.cc b/src/Reassem.cc index 3f6dce1659..02545072ca 100644 --- a/src/Reassem.cc +++ b/src/Reassem.cc @@ -23,32 +23,31 @@ void DataBlockList::DataSize(uint64_t seq_cutoff, uint64_t* below, uint64_t* abo for ( const auto& e : block_map ) { - auto b = e.second; + const auto& b = e.second; - if ( b->seq <= seq_cutoff ) - *above += b->Size(); + if ( b.seq <= seq_cutoff ) + *above += b.Size(); else - *below += b->Size(); + *below += b.Size(); } } void DataBlockList::Delete(DataBlockMap::const_iterator it) { - auto b = it->second; - auto size = b->Size(); + const auto& b = it->second; + auto size = b.Size(); block_map.erase(it); total_data_size -= size; - delete b; Reassembler::total_size -= size + sizeof(DataBlock); Reassembler::sizes[reassembler->rtype] -= size + sizeof(DataBlock); } -DataBlock* DataBlockList::Remove(DataBlockMap::const_iterator it) +DataBlock DataBlockList::Remove(DataBlockMap::const_iterator it) { - auto b = it->second; - auto size = b->Size(); + auto b = std::move(it->second); + auto size = b.Size(); block_map.erase(it); total_data_size -= size; @@ -58,18 +57,19 @@ DataBlock* DataBlockList::Remove(DataBlockMap::const_iterator it) void DataBlockList::Clear() { - // TODO: maybe can just use clear() - while ( ! block_map.empty() ) - Delete(block_map.begin()); - + auto total_db_size = sizeof(DataBlock) * block_map.size(); + auto total = total_data_size + total_db_size; + Reassembler::total_size -= total; + Reassembler::sizes[reassembler->rtype] -= total; + total_data_size = 0; block_map.clear(); } -void DataBlockList::Append(DataBlock* block, uint64_t limit) +void DataBlockList::Append(DataBlock block, uint64_t limit) { - total_data_size += block->Size(); + total_data_size += block.Size(); - block_map.emplace_hint(block_map.end(), block->seq, block); + block_map.emplace_hint(block_map.end(), block.seq, std::move(block)); while ( block_map.size() > limit ) Delete(block_map.begin()); @@ -94,9 +94,7 @@ DataBlockList::Insert(uint64_t seq, uint64_t upper, const u_char* data, DataBlockMap::const_iterator hint) { auto size = upper - seq; - auto db = new DataBlock(data, size, seq); - - auto rval = block_map.emplace_hint(hint, seq, db); + auto rval = block_map.emplace_hint(hint, seq, DataBlock(data, size, seq)); total_data_size += size; Reassembler::sizes[reassembler->rtype] += size + sizeof(DataBlock); @@ -113,10 +111,10 @@ DataBlockList::Insert(uint64_t seq, uint64_t upper, const u_char* data, if ( block_map.empty() ) return Insert(seq, upper, data, block_map.end()); - auto last = block_map.rbegin()->second; + const auto& last = block_map.rbegin()->second; // Special check for the common case of appending to the end. - if ( seq == last->upper ) + if ( seq == last.upper ) return Insert(seq, upper, data, block_map.end()); // Find the first block that doesn't come completely before the new data. @@ -132,44 +130,39 @@ DataBlockList::Insert(uint64_t seq, uint64_t upper, const u_char* data, it = block_map.begin(); } - while ( std::next(it) != block_map.end() && it->second->upper <= seq ) + while ( std::next(it) != block_map.end() && it->second.upper <= seq ) ++it; - DataBlock* b = it->second; + const auto& b = it->second; - if ( b->upper <= seq ) + if ( b.upper <= seq ) // b is the last block, and it comes completely before the new block. return Insert(seq, upper, data, block_map.end()); - if ( upper <= b->seq ) + if ( upper <= b.seq ) // The new block comes completely before b. return Insert(seq, upper, data, it); - DataBlock* new_b; DataBlockMap::const_iterator rval; // The blocks overlap. - if ( seq < b->seq ) + if ( seq < b.seq ) { // The new block has a prefix that comes before b. - uint64_t prefix_len = b->seq - seq; + uint64_t prefix_len = b.seq - seq; rval = Insert(seq, seq + prefix_len, data, it); - new_b = rval->second; data += prefix_len; seq += prefix_len; } else - { rval = it; - new_b = b; - } uint64_t overlap_start = seq; - uint64_t overlap_offset = overlap_start - b->seq; + uint64_t overlap_offset = overlap_start - b.seq; uint64_t new_b_len = upper - seq; - uint64_t b_len = b->upper - overlap_start; + uint64_t b_len = b.upper - overlap_start; uint64_t overlap_len = min(new_b_len, b_len); if ( overlap_len < new_b_len ) @@ -180,7 +173,7 @@ DataBlockList::Insert(uint64_t seq, uint64_t upper, const u_char* data, auto r = Insert(seq, upper, data, &it); - if ( new_b == b ) + if ( rval == it ) rval = r; } @@ -197,11 +190,11 @@ uint64_t DataBlockList::Trim(uint64_t seq, uint64_t max_old, if ( ! block_map.empty() ) { - auto first = block_map.begin()->second; + const auto& first = block_map.begin()->second; - if ( first->seq > reassembler->LastReassemSeq() ) + if ( first.seq > reassembler->LastReassemSeq() ) // An initial hole. - num_missing += first->seq - reassembler->LastReassemSeq(); + num_missing += first.seq - reassembler->LastReassemSeq(); } else if ( seq > reassembler->LastReassemSeq() ) { @@ -219,26 +212,25 @@ uint64_t DataBlockList::Trim(uint64_t seq, uint64_t max_old, while ( ! block_map.empty() ) { auto first_it = block_map.begin(); - auto first = first_it->second; + const auto& first = first_it->second; - if ( first->upper > seq ) + if ( first.upper > seq ) break; - auto next_it = std::next(first_it); - auto next = next_it == block_map.end() ? nullptr : next_it->second; + auto next = std::next(first_it); - if ( next && next->seq <= seq ) + if ( next != block_map.end() && next->second.seq <= seq ) { - if ( first->upper != next->seq ) - num_missing += next->seq - first->upper; + if ( first.upper != next->second.seq ) + num_missing += next->second.seq - first.upper; } else { // No more blocks - did this one make it to seq? // Second half of test is for acks of FINs, which // don't get entered into the sequence space. - if ( first->upper != seq && first->upper != seq - 1 ) - num_missing += seq - first->upper; + if ( first.upper != seq && first.upper != seq - 1 ) + num_missing += seq - first.upper; } if ( max_old ) @@ -250,12 +242,12 @@ uint64_t DataBlockList::Trim(uint64_t seq, uint64_t max_old, if ( ! block_map.empty() ) { auto first_it = block_map.begin(); - auto first = first_it->second; + const auto& first = first_it->second; // If we skipped over some undeliverable data, then // it's possible that this block is now deliverable. // Give it a try. - if ( first->seq == reassembler->LastReassemSeq() ) + if ( first.seq == reassembler->LastReassemSeq() ) reassembler->BlockInserted(first_it); } @@ -277,9 +269,9 @@ void Reassembler::CheckOverlap(const DataBlockList& list, if ( list.Empty() ) return; - auto last = list.LastBlock(); + const auto& last = list.LastBlock(); - if ( seq == last->upper ) + if ( seq == last.upper ) // Special case check for common case of appending to the end. return; @@ -292,31 +284,31 @@ void Reassembler::CheckOverlap(const DataBlockList& list, for ( ; it != list.End(); ++it ) { - auto b = it->second; + const auto& b = it->second; uint64_t nseq = seq; uint64_t nupper = upper; const u_char* ndata = data; - if ( nupper <= b->seq ) + if ( nupper <= b.seq ) break; - if ( nseq >= b->upper ) + if ( nseq >= b.upper ) continue; - if ( nseq < b->seq ) + if ( nseq < b.seq ) { - ndata += (b->seq - seq); - nseq = b->seq; + ndata += (b.seq - seq); + nseq = b.seq; } - if ( nupper > b->upper ) - nupper = b->upper; + if ( nupper > b.upper ) + nupper = b.upper; - uint64_t overlap_offset = (nseq - b->seq); + uint64_t overlap_offset = (nseq - b.seq); uint64_t overlap_len = (nupper - nseq); if ( overlap_len ) - Overlap(&b->block[overlap_offset], ndata, overlap_len); + Overlap(&b.block[overlap_offset], ndata, overlap_len); } } diff --git a/src/Reassem.h b/src/Reassem.h index 31e9c02599..2c5535966b 100644 --- a/src/Reassem.h +++ b/src/Reassem.h @@ -26,6 +26,50 @@ class DataBlock { public: DataBlock(const u_char* data, uint64_t size, uint64_t seq); + DataBlock(const DataBlock& other) + { + seq = other.seq; + upper = other.upper; + auto size = other.Size(); + block = new u_char[size]; + memcpy(block, other.block, size); + } + + DataBlock(DataBlock&& other) + { + seq = other.seq; + upper = other.upper; + block = other.block; + other.block = nullptr; + } + + DataBlock& operator=(const DataBlock& other) + { + if ( this == &other ) + return *this; + + seq = other.seq; + upper = other.upper; + auto size = other.Size(); + delete [] block; + block = new u_char[size]; + memcpy(block, other.block, size); + return *this; + } + + DataBlock& operator=(DataBlock&& other) + { + if ( this == &other ) + return *this; + + seq = other.seq; + upper = other.upper; + delete [] block; + block = other.block; + other.block = nullptr; + return *this; + } + ~DataBlock() { delete [] block; } @@ -37,7 +81,7 @@ public: u_char* block; }; -using DataBlockMap = std::map; +using DataBlockMap = std::map; // TODO: add comments class DataBlockList { @@ -58,10 +102,10 @@ public: DataBlockMap::const_iterator End() const { return block_map.end(); } - DataBlock* FirstBlock() const + const DataBlock& FirstBlock() const { return block_map.begin()->second; } - DataBlock* LastBlock() const + const DataBlock& LastBlock() const { return block_map.rbegin()->second; } bool Empty() const @@ -81,7 +125,7 @@ public: Insert(uint64_t seq, uint64_t upper, const u_char* data, DataBlockMap::const_iterator* hint = nullptr); - void Append(DataBlock* block, uint64_t limit); + void Append(DataBlock block, uint64_t limit); uint64_t Trim(uint64_t seq, uint64_t max_old, DataBlockList* old_list); @@ -95,7 +139,7 @@ private: void Delete(DataBlockMap::const_iterator it); - DataBlock* Remove(DataBlockMap::const_iterator it); + DataBlock Remove(DataBlockMap::const_iterator it); Reassembler* reassembler = nullptr; size_t total_data_size = 0; diff --git a/src/analyzer/protocol/tcp/TCP_Reassembler.cc b/src/analyzer/protocol/tcp/TCP_Reassembler.cc index d38ee3bf64..e5b1250a05 100644 --- a/src/analyzer/protocol/tcp/TCP_Reassembler.cc +++ b/src/analyzer/protocol/tcp/TCP_Reassembler.cc @@ -70,8 +70,8 @@ void TCP_Reassembler::Done() { auto last_block = std::prev(block_list.End())->second; - if ( last_reassem_seq < last_block->upper ) - RecordToSeq(last_reassem_seq, last_block->upper, + if ( last_reassem_seq < last_block.upper ) + RecordToSeq(last_reassem_seq, last_block.upper, record_contents_file); } @@ -92,7 +92,7 @@ uint64_t TCP_Reassembler::NumUndeliveredBytes() const return 0; auto last_block = std::prev(block_list.End())->second; - return last_block->upper - last_reassem_seq; + return last_block.upper - last_reassem_seq; } void TCP_Reassembler::SetContentsFile(BroFile* f) @@ -109,7 +109,7 @@ void TCP_Reassembler::SetContentsFile(BroFile* f) else { if ( ! block_list.Empty() ) - RecordToSeq(block_list.Begin()->second->seq, last_reassem_seq, f); + RecordToSeq(block_list.Begin()->second.seq, last_reassem_seq, f); } Ref(f); @@ -241,21 +241,21 @@ void TCP_Reassembler::Undelivered(uint64_t up_to_seq) while ( it != block_list.End() ) { - auto b = it->second; + const auto& b = it->second; - if ( b->seq < last_reassem_seq ) + if ( b.seq < last_reassem_seq ) { // Already delivered this block. ++it; continue; } - if ( b->seq >= up_to_seq ) + if ( b.seq >= up_to_seq ) // Block is beyond what we need to process at this point. break; uint64_t gap_at_seq = last_reassem_seq; - uint64_t gap_len = b->seq - last_reassem_seq; + uint64_t gap_len = b.seq - last_reassem_seq; Gap(gap_at_seq, gap_len); last_reassem_seq += gap_len; @@ -289,10 +289,10 @@ void TCP_Reassembler::MatchUndelivered(uint64_t up_to_seq, bool use_last_upper) if ( block_list.Empty() || ! rule_matcher ) return; - auto last_block = std::prev(block_list.End())->second; + const auto& last_block = std::prev(block_list.End())->second; if ( use_last_upper ) - up_to_seq = last_block->upper; + up_to_seq = last_block.upper; // ### Note: the original code did not check whether blocks have // already been delivered, but not ACK'ed, and therefore still @@ -302,7 +302,7 @@ void TCP_Reassembler::MatchUndelivered(uint64_t up_to_seq, bool use_last_upper) // min(last_block->upper, up_to_seq). // Is there such data? if ( up_to_seq <= last_reassem_seq || - last_block->upper <= last_reassem_seq ) + last_block.upper <= last_reassem_seq ) return; // Skip blocks that are already delivered (but not ACK'ed). @@ -311,12 +311,12 @@ void TCP_Reassembler::MatchUndelivered(uint64_t up_to_seq, bool use_last_upper) for ( auto it = block_list.Begin(); it != block_list.End(); ++it ) { - auto b = it->second; + const auto& b = it->second; - if ( b->upper > last_reassem_seq ) + if ( b.upper > last_reassem_seq ) break; - tcp_analyzer->Conn()->Match(Rule::PAYLOAD, b->block, b->Size(), + tcp_analyzer->Conn()->Match(Rule::PAYLOAD, b.block, b.Size(), false, false, IsOrig(), false); } } @@ -326,7 +326,7 @@ void TCP_Reassembler::RecordToSeq(uint64_t start_seq, uint64_t stop_seq, BroFile auto it = block_list.Begin(); // Skip over blocks up to the start seq. - while ( it != block_list.End() && it->second->upper <= start_seq ) + while ( it != block_list.End() && it->second.upper <= start_seq ) ++it; if ( it == block_list.End() ) @@ -334,15 +334,15 @@ void TCP_Reassembler::RecordToSeq(uint64_t start_seq, uint64_t stop_seq, BroFile uint64_t last_seq = start_seq; - while ( it != block_list.End() && it->second->upper <= stop_seq ) + while ( it != block_list.End() && it->second.upper <= stop_seq ) { - auto b = it->second; + const auto& b = it->second; - if ( b->seq > last_seq ) - RecordGap(last_seq, b->seq, f); + if ( b.seq > last_seq ) + RecordGap(last_seq, b.seq, f); RecordBlock(b, f); - last_seq = b->upper; + last_seq = b.upper; ++it; } @@ -352,9 +352,9 @@ void TCP_Reassembler::RecordToSeq(uint64_t start_seq, uint64_t stop_seq, BroFile RecordGap(last_seq, stop_seq, f); } -void TCP_Reassembler::RecordBlock(const DataBlock* b, BroFile* f) +void TCP_Reassembler::RecordBlock(const DataBlock& b, BroFile* f) { - if ( f->Write((const char*) b->block, b->Size()) ) + if ( f->Write((const char*) b.block, b.Size()) ) return; reporter->Error("TCP_Reassembler contents write failed"); @@ -388,10 +388,10 @@ void TCP_Reassembler::RecordGap(uint64_t start_seq, uint64_t upper_seq, BroFile* void TCP_Reassembler::BlockInserted(DataBlockMap::const_iterator it) { - auto start_block = it->second; + const auto& start_block = it->second; - if ( start_block->seq > last_reassem_seq || - start_block->upper <= last_reassem_seq ) + if ( start_block.seq > last_reassem_seq || + start_block.upper <= last_reassem_seq ) return; // We've filled a leading hole. Deliver as much as possible. @@ -402,21 +402,21 @@ void TCP_Reassembler::BlockInserted(DataBlockMap::const_iterator it) // data. while ( it != block_list.End() ) { - auto b = it->second; + const auto& b = it->second; - if ( b->seq > last_reassem_seq ) + if ( b.seq > last_reassem_seq ) break; - if ( b->seq == last_reassem_seq ) + if ( b.seq == last_reassem_seq ) { // New stuff. - uint64_t len = b->Size(); + uint64_t len = b.Size(); uint64_t seq = last_reassem_seq; last_reassem_seq += len; if ( record_contents_file ) RecordBlock(b, record_contents_file); - DeliverBlock(seq, len, b->block); + DeliverBlock(seq, len, b.block); } ++it; diff --git a/src/analyzer/protocol/tcp/TCP_Reassembler.h b/src/analyzer/protocol/tcp/TCP_Reassembler.h index 4f93ef73fb..a934dcbebc 100644 --- a/src/analyzer/protocol/tcp/TCP_Reassembler.h +++ b/src/analyzer/protocol/tcp/TCP_Reassembler.h @@ -87,7 +87,7 @@ private: void Gap(uint64_t seq, uint64_t len); void RecordToSeq(uint64_t start_seq, uint64_t stop_seq, BroFile* f); - void RecordBlock(const DataBlock* b, BroFile* f); + void RecordBlock(const DataBlock& b, BroFile* f); void RecordGap(uint64_t start_seq, uint64_t upper_seq, BroFile* f); void BlockInserted(DataBlockMap::const_iterator it) override; diff --git a/src/file_analysis/FileReassembler.cc b/src/file_analysis/FileReassembler.cc index 5949c6adc3..6eb5453beb 100644 --- a/src/file_analysis/FileReassembler.cc +++ b/src/file_analysis/FileReassembler.cc @@ -29,11 +29,11 @@ uint64_t FileReassembler::Flush() if ( block_list.Empty() ) return 0; - auto last_block = std::prev(block_list.End())->second; + const auto& last_block = std::prev(block_list.End())->second; // This is expected to call back into FileReassembler::Undelivered(). flushing = true; - uint64_t rval = TrimToSeq(last_block->upper); + uint64_t rval = TrimToSeq(last_block.upper); flushing = false; return rval; } @@ -52,24 +52,24 @@ uint64_t FileReassembler::FlushTo(uint64_t sequence) void FileReassembler::BlockInserted(DataBlockMap::const_iterator it) { - auto start_block = it->second; + const auto& start_block = it->second; - if ( start_block->seq > last_reassem_seq || - start_block->upper <= last_reassem_seq ) + if ( start_block.seq > last_reassem_seq || + start_block.upper <= last_reassem_seq ) return; while ( it != block_list.End() ) { - auto b = it->second; + const auto& b = it->second; - if ( b->seq > last_reassem_seq ) + if ( b.seq > last_reassem_seq ) break; - if ( b->seq == last_reassem_seq ) + if ( b.seq == last_reassem_seq ) { // New stuff. - uint64_t len = b->Size(); + uint64_t len = b.Size(); last_reassem_seq += len; - the_file->DeliverStream(b->block, len); + the_file->DeliverStream(b.block, len); } ++it; @@ -86,21 +86,21 @@ void FileReassembler::Undelivered(uint64_t up_to_seq) while ( it != block_list.End() ) { - auto b = it->second; + const auto& b = it->second; - if ( b->seq < last_reassem_seq ) + if ( b.seq < last_reassem_seq ) { // Already delivered this block. ++it; continue; } - if ( b->seq >= up_to_seq ) + if ( b.seq >= up_to_seq ) // Block is beyond what we need to process at this point. break; uint64_t gap_at_seq = last_reassem_seq; - uint64_t gap_len = b->seq - last_reassem_seq; + uint64_t gap_len = b.seq - last_reassem_seq; the_file->Gap(gap_at_seq, gap_len); last_reassem_seq += gap_len; BlockInserted(it);