Use DataBlock value instead of pointer in reassembly map

This commit is contained in:
Jon Siwek 2019-09-13 14:17:41 -07:00
parent e1e779e90b
commit 69d1620374
6 changed files with 169 additions and 129 deletions

View file

@ -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 ( it != block_list.Begin() )
{
const auto& prev = std::prev(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 )
if ( prev.upper < b.seq )
break;
}
if ( b->upper > n )
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;

View file

@ -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);
}
}

View file

@ -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<uint64_t, DataBlock*>;
using DataBlockMap = std::map<uint64_t, DataBlock>;
// 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;

View file

@ -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;

View file

@ -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;

View file

@ -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);