Mirror/zeek
1
0
Fork 0
mirror of https://github.com/zeek/zeek.git synced 2025-10-02 14:48:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Similar map changes for async requests

Browse source
This commit is contained in:
Tim Wojtulewicz 2022-02-07 14:49:55 -07:00
parent e8f833b8a6
commit b531ec97ef
2 changed files with 42 additions and 55 deletions
Download patch file Download diff file Expand all files Collapse all files

78
src/DNS_Mgr.cc
View file

@ -281,8 +281,8 @@ void DNS_Request::ProcessAsyncResult(bool timed_out, DNS_Mgr* mgr)
mgr->CheckAsyncHostRequest(host, timed_out); mgr->CheckAsyncHostRequest(host, timed_out);
else if ( request_type == T_PTR ) else if ( request_type == T_PTR )
mgr->CheckAsyncAddrRequest(addr, timed_out); mgr->CheckAsyncAddrRequest(addr, timed_out);
else if ( request_type == T_TXT ) else
mgr->CheckAsyncTextRequest(host, timed_out); mgr->CheckAsyncOtherRequest(host, timed_out, request_type);
} }
/** /**
@ -842,16 +842,16 @@ void DNS_Mgr::LookupHost(const std::string& name, LookupCallback* callback)
// If we already have a request waiting for this host, we don't need to make // If we already have a request waiting for this host, we don't need to make
// another one. We can just add the callback to it and it'll get handled // another one. We can just add the callback to it and it'll get handled
// when the first request comes back. // when the first request comes back.
AsyncRequestNameMap::iterator i = asyncs_names.find(name); auto key = std::make_pair(T_A, name);
if ( i != asyncs_names.end() ) auto i = asyncs.find(key);
if ( i != asyncs.end() )
req = i->second; req = i->second;
else else
{ {
// A new one. // A new one.
req = new AsyncRequest{}; req = new AsyncRequest{name, T_A};
req->host = name;
asyncs_queued.push_back(req); asyncs_queued.push_back(req);
asyncs_names.emplace_hint(i, name, req); asyncs.emplace_hint(i, std::move(key), req);
} }
req->callbacks.push_back(callback); req->callbacks.push_back(callback);
@ -862,16 +862,16 @@ void DNS_Mgr::LookupHost(const std::string& name, LookupCallback* callback)
IssueAsyncRequests(); IssueAsyncRequests();
} }
void DNS_Mgr::LookupAddr(const IPAddr& host, LookupCallback* callback) void DNS_Mgr::LookupAddr(const IPAddr& addr, LookupCallback* callback)
{ {
if ( mode == DNS_FAKE ) if ( mode == DNS_FAKE )
{ {
resolve_lookup_cb(callback, fake_addr_lookup_result(host)); resolve_lookup_cb(callback, fake_addr_lookup_result(addr));
return; return;
} }
// Do we already know the answer? // Do we already know the answer?
if ( auto name = LookupAddrInCache(host, true, false) ) if ( auto name = LookupAddrInCache(addr, true, false) )
{ {
resolve_lookup_cb(callback, name->CheckString()); resolve_lookup_cb(callback, name->CheckString());
return; return;
@ -882,16 +882,15 @@ void DNS_Mgr::LookupAddr(const IPAddr& host, LookupCallback* callback)
// If we already have a request waiting for this host, we don't need to make // If we already have a request waiting for this host, we don't need to make
// another one. We can just add the callback to it and it'll get handled // another one. We can just add the callback to it and it'll get handled
// when the first request comes back. // when the first request comes back.
AsyncRequestAddrMap::iterator i = asyncs_addrs.find(host); auto i = asyncs.find(addr);
if ( i != asyncs_addrs.end() ) if ( i != asyncs.end() )
req = i->second; req = i->second;
else else
{ {
// A new one. // A new one.
req = new AsyncRequest{}; req = new AsyncRequest{addr};
req->addr = host;
asyncs_queued.push_back(req); asyncs_queued.push_back(req);
asyncs_addrs.emplace_hint(i, host, req); asyncs.emplace_hint(i, addr, req);
} }
req->callbacks.push_back(callback); req->callbacks.push_back(callback);
@ -904,12 +903,6 @@ void DNS_Mgr::LookupAddr(const IPAddr& host, LookupCallback* callback)
void DNS_Mgr::Lookup(const std::string& name, int request_type, LookupCallback* callback) void DNS_Mgr::Lookup(const std::string& name, int request_type, LookupCallback* callback)
{ {
if ( request_type == T_A || request_type == T_AAAA )
{
LookupHost(name, callback);
return;
}
if ( mode == DNS_FAKE ) if ( mode == DNS_FAKE )
{ {
resolve_lookup_cb(callback, fake_lookup_result(name, request_type)); resolve_lookup_cb(callback, fake_lookup_result(name, request_type));
@ -928,17 +921,16 @@ void DNS_Mgr::Lookup(const std::string& name, int request_type, LookupCallback*
// If we already have a request waiting for this host, we don't need to make // If we already have a request waiting for this host, we don't need to make
// another one. We can just add the callback to it and it'll get handled // another one. We can just add the callback to it and it'll get handled
// when the first request comes back. // when the first request comes back.
AsyncRequestTextMap::iterator i = asyncs_texts.find(name); auto key = std::make_pair(request_type, name);
if ( i != asyncs_texts.end() ) auto i = asyncs.find(key);
if ( i != asyncs.end() )
req = i->second; req = i->second;
else else
{ {
// A new one. // A new one.
req = new AsyncRequest{}; req = new AsyncRequest{name, request_type};
req->host = name;
req->is_txt = true;
asyncs_queued.push_back(req); asyncs_queued.push_back(req);
asyncs_texts.emplace_hint(i, name, req); asyncs.emplace_hint(i, std::move(key), req);
} }
req->callbacks.push_back(callback); req->callbacks.push_back(callback);
@ -1281,14 +1273,14 @@ void DNS_Mgr::IssueAsyncRequests()
++num_requests; ++num_requests;
req->time = util::current_time(); req->time = util::current_time();
if ( req->IsAddrReq() ) if ( req->type == T_PTR )
dns_req = new DNS_Request(req->addr, true); dns_req = new DNS_Request(req->addr, true);
else if ( req->is_txt ) else if ( req->type == T_A || req->type == T_AAAA )
dns_req = new DNS_Request(req->host.c_str(), T_TXT, true);
else
// We pass T_A here, but DNSRequest::MakeRequest() will special-case that in // We pass T_A here, but DNSRequest::MakeRequest() will special-case that in
// a request that gets both T_A and T_AAAA results at one time. // a request that gets both T_A and T_AAAA results at one time.
dns_req = new DNS_Request(req->host.c_str(), T_A, true); dns_req = new DNS_Request(req->host.c_str(), T_A, true);
else
dns_req = new DNS_Request(req->host.c_str(), req->type, true);
dns_req->MakeRequest(channel, this); dns_req->MakeRequest(channel, this);
@ -1300,10 +1292,9 @@ void DNS_Mgr::IssueAsyncRequests()
void DNS_Mgr::CheckAsyncHostRequest(const std::string& host, bool timeout) void DNS_Mgr::CheckAsyncHostRequest(const std::string& host, bool timeout)
{ {
// Note that this code is a mirror of that for CheckAsyncAddrRequest. // Note that this code is a mirror of that for CheckAsyncAddrRequest.
auto i = asyncs.find(std::make_pair(T_A, host));
AsyncRequestNameMap::iterator i = asyncs_names.find(host); if ( i != asyncs.end() )
if ( i != asyncs_names.end() )
{ {
if ( timeout ) if ( timeout )
{ {
@ -1319,7 +1310,7 @@ void DNS_Mgr::CheckAsyncHostRequest(const std::string& host, bool timeout)
return; return;
delete i->second; delete i->second;
asyncs_names.erase(i); asyncs.erase(i);
--asyncs_pending; --asyncs_pending;
} }
} }
@ -1330,9 +1321,9 @@ void DNS_Mgr::CheckAsyncAddrRequest(const IPAddr& addr, bool timeout)
// In the following, if it's not in the respective map anymore, we've // In the following, if it's not in the respective map anymore, we've
// already finished it earlier and don't have anything to do. // already finished it earlier and don't have anything to do.
AsyncRequestAddrMap::iterator i = asyncs_addrs.find(addr); auto i = asyncs.find(addr);
if ( i != asyncs_addrs.end() ) if ( i != asyncs.end() )
{ {
if ( timeout ) if ( timeout )
{ {
@ -1348,25 +1339,24 @@ void DNS_Mgr::CheckAsyncAddrRequest(const IPAddr& addr, bool timeout)
return; return;
delete i->second; delete i->second;
asyncs_addrs.erase(i); asyncs.erase(i);
--asyncs_pending; --asyncs_pending;
} }
} }
void DNS_Mgr::CheckAsyncTextRequest(const std::string& host, bool timeout) void DNS_Mgr::CheckAsyncOtherRequest(const std::string& host, bool timeout, int request_type)
{ {
// Note that this code is a mirror of that for CheckAsyncAddrRequest. // Note that this code is a mirror of that for CheckAsyncAddrRequest.
AsyncRequestTextMap::iterator i = asyncs_texts.find(host); auto i = asyncs.find(std::make_pair(request_type, host));
if ( i != asyncs_texts.end() ) if ( i != asyncs.end() )
{ {
if ( timeout ) if ( timeout )
{ {
AsyncRequestTextMap::iterator it = asyncs_texts.begin();
++failed; ++failed;
i->second->Timeout(); i->second->Timeout();
} }
else if ( auto name = LookupOtherInCache(host, T_TXT, true) ) else if ( auto name = LookupOtherInCache(host, request_type, true) )
{ {
++successful; ++successful;
i->second->Resolved(name->CheckString()); i->second->Resolved(name->CheckString());
@ -1375,7 +1365,7 @@ void DNS_Mgr::CheckAsyncTextRequest(const std::string& host, bool timeout)
return; return;
delete i->second; delete i->second;
asyncs_texts.erase(i); asyncs.erase(i);
--asyncs_pending; --asyncs_pending;
} }
} }

19
src/DNS_Mgr.h
View file

@ -256,7 +256,7 @@ protected:
// requested. // requested.
void CheckAsyncAddrRequest(const IPAddr& addr, bool timeout); void CheckAsyncAddrRequest(const IPAddr& addr, bool timeout);
void CheckAsyncHostRequest(const std::string& host, bool timeout); void CheckAsyncHostRequest(const std::string& host, bool timeout);
void CheckAsyncTextRequest(const std::string& host, bool timeout); void CheckAsyncOtherRequest(const std::string& host, bool timeout, int request_type);
void Event(EventHandlerPtr e, DNS_Mapping* dm); void Event(EventHandlerPtr e, DNS_Mapping* dm);
void Event(EventHandlerPtr e, DNS_Mapping* dm, ListValPtr l1, ListValPtr l2); void Event(EventHandlerPtr e, DNS_Mapping* dm, ListValPtr l1, ListValPtr l2);
@ -303,10 +303,13 @@ protected:
IPAddr addr; IPAddr addr;
std::string host; std::string host;
CallbackList callbacks; CallbackList callbacks;
bool is_txt = false; int type = 0;
bool processed = false; bool processed = false;
bool IsAddrReq() const { return host.empty(); } AsyncRequest(std::string host, int request_type) : host(std::move(host)), type(request_type)
{
}
AsyncRequest(const IPAddr& addr) : addr(addr), type(T_PTR) { }
void Resolved(const std::string& name); void Resolved(const std::string& name);
void Resolved(TableValPtr addrs); void Resolved(TableValPtr addrs);
@ -318,14 +321,8 @@ protected:
bool operator()(const AsyncRequest* a, const AsyncRequest* b) { return a->time > b->time; } bool operator()(const AsyncRequest* a, const AsyncRequest* b) { return a->time > b->time; }
}; };
using AsyncRequestAddrMap = std::map<IPAddr, AsyncRequest*>; using AsyncRequestMap = std::map<MappingKey, AsyncRequest*>;
AsyncRequestAddrMap asyncs_addrs; AsyncRequestMap asyncs;
using AsyncRequestNameMap = std::map<std::string, AsyncRequest*>;
AsyncRequestNameMap asyncs_names;
using AsyncRequestTextMap = std::map<std::string, AsyncRequest*>;
AsyncRequestTextMap asyncs_texts;
using QueuedList = std::list<AsyncRequest*>; using QueuedList = std::list<AsyncRequest*>;
QueuedList asyncs_queued; QueuedList asyncs_queued;