1 | #include "DNSResolver.h" |
2 | #include <common/SimpleCache.h> |
3 | #include <Common/Exception.h> |
4 | #include <Common/ProfileEvents.h> |
5 | #include <Core/Names.h> |
6 | #include <Core/Types.h> |
7 | #include <Poco/Net/DNS.h> |
8 | #include <Poco/Net/NetException.h> |
9 | #include <Poco/NumberParser.h> |
10 | #include <Poco/Logger.h> |
11 | #include <common/logger_useful.h> |
12 | #include <arpa/inet.h> |
13 | #include <atomic> |
14 | #include <optional> |
15 | |
16 | namespace ProfileEvents |
17 | { |
18 | extern Event NetworkErrors; |
19 | } |
20 | |
21 | |
22 | namespace DB |
23 | { |
24 | |
25 | namespace ErrorCodes |
26 | { |
27 | extern const int BAD_ARGUMENTS; |
28 | } |
29 | |
30 | |
31 | /// Slightly altered implementation from https://github.com/pocoproject/poco/blob/poco-1.6.1/Net/src/SocketAddress.cpp#L86 |
32 | static void splitHostAndPort(const std::string & host_and_port, std::string & out_host, UInt16 & out_port) |
33 | { |
34 | String port_str; |
35 | out_host.clear(); |
36 | |
37 | auto it = host_and_port.begin(); |
38 | auto end = host_and_port.end(); |
39 | |
40 | if (*it == '[') /// Try parse case '[<IPv6 or something else>]:<port>' |
41 | { |
42 | ++it; |
43 | while (it != end && *it != ']') |
44 | out_host += *it++; |
45 | if (it == end) |
46 | throw Exception("Malformed IPv6 address" , ErrorCodes::BAD_ARGUMENTS); |
47 | ++it; |
48 | } |
49 | else /// Case '<IPv4 or domain name or something else>:<port>' |
50 | { |
51 | while (it != end && *it != ':') |
52 | out_host += *it++; |
53 | } |
54 | |
55 | if (it != end && *it == ':') |
56 | { |
57 | ++it; |
58 | while (it != end) |
59 | port_str += *it++; |
60 | } |
61 | else |
62 | throw Exception("Missing port number" , ErrorCodes::BAD_ARGUMENTS); |
63 | |
64 | unsigned port; |
65 | if (Poco::NumberParser::tryParseUnsigned(port_str, port) && port <= 0xFFFF) |
66 | { |
67 | out_port = static_cast<UInt16>(port); |
68 | } |
69 | else |
70 | { |
71 | struct servent * se = getservbyname(port_str.c_str(), nullptr); |
72 | if (se) |
73 | out_port = ntohs(static_cast<UInt16>(se->s_port)); |
74 | else |
75 | throw Exception("Service not found" , ErrorCodes::BAD_ARGUMENTS); |
76 | } |
77 | } |
78 | |
79 | static Poco::Net::IPAddress resolveIPAddressImpl(const std::string & host) |
80 | { |
81 | /// NOTE: Poco::Net::DNS::resolveOne(host) doesn't work for IP addresses like 127.0.0.2 |
82 | /// Therefore we use SocketAddress constructor with dummy port to resolve IP |
83 | return Poco::Net::SocketAddress(host, 0U).host(); |
84 | } |
85 | |
86 | struct DNSResolver::Impl |
87 | { |
88 | SimpleCache<decltype(resolveIPAddressImpl), &resolveIPAddressImpl> cache_host; |
89 | |
90 | std::mutex drop_mutex; |
91 | std::mutex update_mutex; |
92 | |
93 | /// Cached server host name |
94 | std::optional<String> host_name; |
95 | |
96 | /// Store hosts, which was asked to resolve from last update of DNS cache. |
97 | NameSet new_hosts; |
98 | |
99 | /// Store all hosts, which was whenever asked to resolve |
100 | NameSet known_hosts; |
101 | |
102 | /// If disabled, will not make cache lookups, will resolve addresses manually on each call |
103 | std::atomic<bool> disable_cache{false}; |
104 | }; |
105 | |
106 | |
107 | DNSResolver::DNSResolver() : impl(std::make_unique<DNSResolver::Impl>()) {} |
108 | |
109 | Poco::Net::IPAddress DNSResolver::resolveHost(const std::string & host) |
110 | { |
111 | if (impl->disable_cache) |
112 | return resolveIPAddressImpl(host); |
113 | |
114 | addToNewHosts(host); |
115 | return impl->cache_host(host); |
116 | } |
117 | |
118 | Poco::Net::SocketAddress DNSResolver::resolveAddress(const std::string & host_and_port) |
119 | { |
120 | if (impl->disable_cache) |
121 | return Poco::Net::SocketAddress(host_and_port); |
122 | |
123 | String host; |
124 | UInt16 port; |
125 | splitHostAndPort(host_and_port, host, port); |
126 | |
127 | addToNewHosts(host); |
128 | return Poco::Net::SocketAddress(impl->cache_host(host), port); |
129 | } |
130 | |
131 | Poco::Net::SocketAddress DNSResolver::resolveAddress(const std::string & host, UInt16 port) |
132 | { |
133 | if (impl->disable_cache) |
134 | return Poco::Net::SocketAddress(host, port); |
135 | |
136 | addToNewHosts(host); |
137 | return Poco::Net::SocketAddress(impl->cache_host(host), port); |
138 | } |
139 | |
140 | void DNSResolver::dropCache() |
141 | { |
142 | impl->cache_host.drop(); |
143 | |
144 | std::scoped_lock lock(impl->update_mutex, impl->drop_mutex); |
145 | |
146 | impl->known_hosts.clear(); |
147 | impl->new_hosts.clear(); |
148 | impl->host_name.reset(); |
149 | } |
150 | |
151 | void DNSResolver::setDisableCacheFlag(bool is_disabled) |
152 | { |
153 | impl->disable_cache = is_disabled; |
154 | } |
155 | |
156 | String DNSResolver::getHostName() |
157 | { |
158 | if (impl->disable_cache) |
159 | return Poco::Net::DNS::hostName(); |
160 | |
161 | std::lock_guard lock(impl->drop_mutex); |
162 | |
163 | if (!impl->host_name.has_value()) |
164 | impl->host_name.emplace(Poco::Net::DNS::hostName()); |
165 | |
166 | return *impl->host_name; |
167 | } |
168 | |
169 | bool DNSResolver::updateCache() |
170 | { |
171 | { |
172 | std::lock_guard lock(impl->drop_mutex); |
173 | for (auto & host : impl->new_hosts) |
174 | impl->known_hosts.insert(std::move(host)); |
175 | impl->new_hosts.clear(); |
176 | |
177 | impl->host_name.emplace(Poco::Net::DNS::hostName()); |
178 | } |
179 | |
180 | std::lock_guard lock(impl->update_mutex); |
181 | |
182 | bool updated = false; |
183 | String lost_hosts; |
184 | for (const auto & host : impl->known_hosts) |
185 | { |
186 | try |
187 | { |
188 | updated |= updateHost(host); |
189 | } |
190 | catch (const Poco::Net::NetException &) |
191 | { |
192 | ProfileEvents::increment(ProfileEvents::NetworkErrors); |
193 | |
194 | if (!lost_hosts.empty()) |
195 | lost_hosts += ", " ; |
196 | lost_hosts += host; |
197 | } |
198 | catch (...) |
199 | { |
200 | tryLogCurrentException(__PRETTY_FUNCTION__); |
201 | } |
202 | } |
203 | |
204 | if (!lost_hosts.empty()) |
205 | LOG_INFO(&Logger::get("DNSResolver" ), "Cached hosts not found: " << lost_hosts); |
206 | |
207 | return updated; |
208 | } |
209 | |
210 | bool DNSResolver::updateHost(const String & host) |
211 | { |
212 | /// Usage of updateHost implies that host is already in cache and there is no extra computations |
213 | auto old_value = impl->cache_host(host); |
214 | impl->cache_host.update(host); |
215 | return old_value != impl->cache_host(host); |
216 | } |
217 | |
218 | void DNSResolver::addToNewHosts(const String & host) |
219 | { |
220 | std::lock_guard lock(impl->drop_mutex); |
221 | impl->new_hosts.insert(host); |
222 | } |
223 | |
224 | DNSResolver::~DNSResolver() = default; |
225 | |
226 | DNSResolver & DNSResolver::instance() |
227 | { |
228 | static DNSResolver ret; |
229 | return ret; |
230 | } |
231 | |
232 | } |
233 | |