1 // Copyright 2017 NSONE, Inc
2
3 #include <iostream>
4 #include <iterator>
5 #include <map>
6 #include <queue>
7 #include <sstream>
8 #include <string>
9 #include <vector>
10
11 #include "config.h"
12 #include "docopt.h"
13 #include "http.h"
14 #include "metrics.h"
15 #include "query.h"
16 #include "trafgen.h"
17 #include "utils.h"
18
19 #include <uvw.hpp>
20
21 #include "version.h"
22
23 static const char USAGE[] =
24 R"(Flamethrower.
25 Usage:
26 flame [-b BIND_IP] [-q QCOUNT] [-c TCOUNT] [-p PORT] [-d DELAY_MS] [-r RECORD] [-T QTYPE]
27 [-o FILE] [-l LIMIT_SECS] [-t TIMEOUT] [-F FAMILY] [-f FILE] [-n LOOP] [-P PROTOCOL] [-M HTTPMETHOD]
28 [-Q QPS] [-g GENERATOR] [-v VERBOSITY] [-R] [--class CLASS] [--qps-flow SPEC]
29 [--dnssec] [--targets FILE]
30 TARGET [GENOPTS]...
31 flame (-h | --help)
32 flame --version
33
34 TARGET may be a hostname, an IP address, or a comma separated list of either. If multiple targets are specified,
35 they will be sent queries in a strict round robin fashion across all concurrent generators. All targets must
36 share the same port, protocol, and internet family.
37
38 TARGET may also be the special value "file", in which case the --targets option needs to also be specified.
39
40 Options:
41 -h --help Show this screen
42 --version Show version
43 --class CLASS Default query class, defaults to IN. May also be CH [default: IN]
44 -b BIND_IP IP address to bind to [defaults: 0.0.0.0 for inet, ::0 for inet6]
45 -c TCOUNT Number of concurrent traffic generators per process [default: 10]
46 -d DELAY_MS ms delay between each traffic generator's query [default: 1]
47 -q QCOUNT Number of queries to send every DELAY ms [default: 10]
48 -l LIMIT_SECS Limit traffic generation to N seconds, 0 is unlimited [default: 0]
49 -t TIMEOUT_SECS Query timeout in seconds [default: 3]
50 -n LOOP Loop LOOP times through record list, 0 is unlimited [default: 0]
51 -Q QPS Rate limit to a maximum of QPS, 0 is no limit [default: 0]
52 --qps-flow SPEC Change rate limit over time, format: QPS,MS;QPS,MS;...
53 -r RECORD The base record to use as the DNS query for generators [default: test.com]
54 -T QTYPE The query type to use for generators [default: A]
55 -f FILE Read records from FILE, one per row, QNAME TYPE
56 -p PORT Which port to flame [defaults: 53, 443 for DoH, 853 for DoT]
57 -F FAMILY Internet family (inet/inet6) [default: inet]
58 -P PROTOCOL Protocol to use (udp/tcp/dot/doh) [default: udp]
59 -M HTTPMETHOD HTTP method to use (POST/GET) when DoH is used [default: GET]
60 -g GENERATOR Generate queries with the given generator [default: static]
61 -o FILE Metrics output file, JSON format
62 -v VERBOSITY How verbose output should be, 0 is silent [default: 1]
63 -R Randomize the query list before sending [default: false]
64 --targets FILE Get the list of TARGETs from the given file, one line per host or IP
65 --dnssec Set DO flag in EDNS
66
67 Generators:
68
69 Using generator modules you can craft the type of packet or query which is sent.
70
71 Specify generator arguments by passing in KEY=VAL pairs, where the KEY is a specific configuration
72 key interpreted by the generator as specified below in caps (although keys are not case sensitive).
73
74 static The basic static generator, used by default, has a single qname/qtype
75 which you can set with -r and -T. There are no KEYs for this generator.
76
77 file The basic file generator, used with -f, reads in one qname/qtype pair
78 per line in the file. There are no KEYs for this generator.
79
80 numberqname Synthesize qnames with random numbers, between [LOW, HIGH], at zone specified with -r
81
82 LOW An integer representing the lowest number queried, default 0
83 HIGH An integer representing the highest number queried, default 100000
84
85 randompkt Generate COUNT randomly generated packets, of random size [1,SIZE]
86
87 COUNT An integer representing the number of packets to generate, default 1000
88 SIZE An integer representing the maximum size of the random packet, default 600
89
90 randomqname Generate COUNT queries of randomly generated QNAME's (including nulls) of random length
91 [1,SIZE], at base zone specified with -r
92
93 COUNT An integer representing the number of queries to generate, default 1000
94 SIZE An integer representing the maximum length of the random qname, default 255
95
96 randomlabel Generate COUNT queries in base zone, each with LBLCOUNT random labels of size [1,LBLSIZE]
97 Use -r to set the base zone to create the labels in. Queries will have a random QTYPE
98 from the most popular set.
99
100 COUNT An integer representing the number of queries to generate, default 1000
101 LBLSIZE An integer representing the maximum length of a single label, default 10
102 LBLCOUNT An integer representing the maximum number of labels in the qname, default 5
103
104
105 Generator Example:
106 flame target.test.com -T ANY -g randomlabel lblsize=10 lblcount=4 count=1000
107
108 )";
109
parse_flowspec(std::string spec,std::queue<std::pair<uint64_t,uint64_t>> & result,int verbosity,long c_count)110 void parse_flowspec(std::string spec, std::queue<std::pair<uint64_t, uint64_t>> &result, int verbosity, long c_count)
111 {
112
113 std::vector<std::string> groups = split(spec, ';');
114 for (unsigned i = 0; i < groups.size(); i++) {
115 std::vector<std::string> nums = split(groups[i], ',');
116 if (verbosity > 1) {
117 std::cout << "adding QPS flow: " << nums[0] << "qps, " << nums[1] << "ms" << std::endl;
118 }
119 long want_r = std::stol(nums[0]);
120 if (want_r < c_count) {
121 std::cerr << "WARNING: QPS flow limit is less than concurrent senders, changing limit to " << c_count << std::endl;
122 want_r = c_count;
123 }
124 result.push(std::make_pair(want_r, std::stol(nums[1])));
125 }
126 }
127
flow_change(std::queue<std::pair<uint64_t,uint64_t>> qps_flow,std::vector<std::shared_ptr<TokenBucket>> rl_list,int verbosity,long c_count)128 void flow_change(std::queue<std::pair<uint64_t, uint64_t>> qps_flow,
129 std::vector<std::shared_ptr<TokenBucket>> rl_list,
130 int verbosity,
131 long c_count)
132 {
133 auto flow = qps_flow.front();
134 qps_flow.pop();
135 if (verbosity) {
136 if (qps_flow.size()) {
137 std::cout << "QPS flow now " << flow.first << " for " << flow.second << "ms, flows left: "
138 << qps_flow.size() << std::endl;
139 } else {
140 std::cout << "QPS flow now " << flow.first << " until completion" << std::endl;
141 }
142 }
143 for (auto &rl : rl_list) {
144 *rl = TokenBucket(flow.first / c_count);
145 }
146 if (qps_flow.size() == 0)
147 return;
148 auto loop = uvw::Loop::getDefault();
149 auto qps_timer = loop->resource<uvw::TimerHandle>();
150 qps_timer->on<uvw::TimerEvent>([qps_flow, rl_list, verbosity, c_count](const auto& event, auto& handle) {
151 handle.stop();
152 flow_change(qps_flow, rl_list, verbosity, c_count);
153 });
154 qps_timer->start(uvw::TimerHandle::Time{flow.second}, uvw::TimerHandle::Time{0});
155 }
156
arg_exists(const char * needle,int argc,char * argv[])157 bool arg_exists(const char *needle, int argc, char *argv[])
158 {
159 for (int i = 0; i < argc; i++) {
160 if (std::string(needle) == std::string(argv[i])) {
161 return true;
162 }
163 }
164 return false;
165 }
166
main(int argc,char * argv[])167 int main(int argc, char *argv[])
168 {
169
170 std::map<std::string, docopt::value> args = docopt::docopt(USAGE,
171 {argv + 1, argv + argc},
172 true, // show help if requested
173 FLAME_VERSION); // version string
174
175 if (args["-v"].asLong() > 3) {
176 for (auto const &arg : args) {
177 std::cout << arg.first << ": " << arg.second << std::endl;
178 }
179 }
180
181 auto loop = uvw::Loop::getDefault();
182
183 auto sigint = loop->resource<uvw::SignalHandle>();
184 auto sigterm = loop->resource<uvw::SignalHandle>();
185 std::shared_ptr<uvw::TimerHandle> run_timer;
186 std::shared_ptr<uvw::TimerHandle> qgen_loop_timer;
187
188 std::string output_file;
189 if (args["-o"]) {
190 output_file = args["-o"].asString();
191 }
192
193 // these defaults change based on protocol
194 long s_delay = args["-d"].asLong();
195 long b_count = args["-q"].asLong();
196 long c_count = args["-c"].asLong();
197
198 Protocol proto{Protocol::UDP};
199 // note: tcptls is available as a deprecated alternative to dot
200 if (args["-P"].asString() == "tcp" || args["-P"].asString() == "dot" || args["-P"].asString() == "tcptls" || args["-P"].asString() == "doh") {
201 if (args["-P"].asString() == "dot" || args["-P"].asString() == "tcptls") {
202 proto = Protocol::DOT;
203 } else if (args["-P"].asString() == "doh") {
204 #ifdef DOH_ENABLE
205 proto = Protocol::DOH;
206 #else
207 std::cerr << "DNS over HTTPS (DoH) support is not enabled" << std::endl;
208 return 1;
209 #endif
210 } else {
211 proto = Protocol::TCP;
212 }
213 if (!arg_exists("-d", argc, argv))
214 s_delay = 1000;
215 if (!arg_exists("-q", argc, argv))
216 b_count = 100;
217 if (!arg_exists("-c", argc, argv))
218 c_count = 30;
219 } else if (args["-P"].asString() == "udp") {
220 proto = Protocol::UDP;
221 } else {
222 std::cerr << "protocol must be 'udp', 'tcp', dot' or 'doh'" << std::endl;
223 return 1;
224 }
225
226 if (!args["-p"]) {
227 if (proto == Protocol::DOT)
228 args["-p"] = std::string("853");
229 #ifdef DOH_ENABLE
230 else if (proto == Protocol::DOH)
231 args["-p"] = std::string("443");
232 #endif
233 else
234 args["-p"] = std::string("53");
235 }
236
237 #ifdef DOH_ENABLE
238 HTTPMethod method{HTTPMethod::GET};
239 if(args["-M"].asString() == "POST") {
240 method = HTTPMethod::POST;
241 }
242 #endif
243
244 auto runtime_limit = args["-l"].asLong();
245
246 auto family_s = args["-F"].asString();
247 int family{0};
248 if (family_s == "inet") {
249 family = AF_INET;
250 } else if (family_s == "inet6") {
251 family = AF_INET6;
252 } else {
253 std::cerr << "internet family must be 'inet' or 'inet6'" << std::endl;
254 return 1;
255 }
256
257 if (!args["-b"]) {
258 if (family == AF_INET)
259 args["-b"] = std::string("0.0.0.0");
260 else
261 args["-b"] = std::string("::0");
262 }
263 auto bind_ip = args["-b"].asString();
264 auto bind_ip_request = loop->resource<uvw::GetAddrInfoReq>();
265 auto bind_ip_resolved = bind_ip_request->addrInfoSync(bind_ip, "0");
266 if (!bind_ip_resolved.first) {
267 std::cerr << "unable to resolve bind ip address: " << bind_ip << std::endl;
268 return 1;
269 }
270
271 std::vector<std::string> raw_target_list;
272 if (args["TARGET"].asString() == "file" && args["--targets"]) {
273 std::ifstream inFile(args["--targets"].asString());
274 if (!inFile.is_open()) {
275 std::cerr << "couldn't open targets file: " << args["--targets"].asString() << std::endl;
276 return 1;
277 }
278
279 std::string line;
280 while (getline(inFile, line)) {
281 raw_target_list.push_back(line);
282 }
283 inFile.close();
284 } else {
285 raw_target_list = split(args["TARGET"].asString(), ',');
286 }
287
288 std::vector<Target> target_list;
289 auto request = loop->resource<uvw::GetAddrInfoReq>();
290 for (uint i = 0; i < raw_target_list.size(); i++) {
291 uvw::Addr addr;
292 struct http_parser_url parsed = {};
293 std::string url = raw_target_list[i];
294 if(url.rfind("https://", 0) != 0) {
295 url.insert(0, "https://");
296 }
297 int ret = http_parser_parse_url(url.c_str(), strlen(url.c_str()), 0, &parsed);
298 if(ret != 0) {
299 std::cerr << "could not parse url: " << url << std::endl;
300 return 1;
301 }
302 std::string authority(&url[parsed.field_data[UF_HOST].off], parsed.field_data[UF_HOST].len);
303
304 auto target_resolved = request->addrInfoSync(authority, args["-p"].asString());
305 if (!target_resolved.first) {
306 std::cerr << "unable to resolve target address: " << authority << std::endl;
307 if (raw_target_list[i] == "file") {
308 std::cerr << "(did you mean to include --targets?)" << std::endl;
309 }
310 return 1;
311 }
312 addrinfo *node{target_resolved.second.get()};
313 while (node && node->ai_family != family) {
314 node = node->ai_next;
315 }
316 if (!node) {
317 std::cerr << "name did not resolve to valid IP address for this inet family: " << raw_target_list[i] << std::endl;
318 return 1;
319 }
320
321 if (family == AF_INET) {
322 addr = uvw::details::address<uvw::IPv4>((struct sockaddr_in *)node->ai_addr);
323 } else if (family == AF_INET6) {
324 addr = uvw::details::address<uvw::IPv6>((struct sockaddr_in6 *)node->ai_addr);
325 }
326 target_list.push_back({&parsed, addr.ip, url});
327 }
328
329 long want_r_limit = args["-Q"].asLong();
330 if (want_r_limit && want_r_limit < c_count) {
331 std::cerr << "WARNING: QPS limit is less than concurrent senders, changing limit to " << c_count << std::endl;
332 want_r_limit = c_count;
333 }
334 auto config = std::make_shared<Config>(
335 args["-v"].asLong(),
336 output_file,
337 want_r_limit);
338
339 std::shared_ptr<QueryGenerator> qgen;
340 try {
341 if (args["-f"]) {
342 qgen = std::make_shared<FileQueryGenerator>(config, args["-f"].asString());
343 } else if (args["-g"] && args["-g"].asString() == "numberqname") {
344 qgen = std::make_shared<NumberNameQueryGenerator>(config);
345 } else if (args["-g"] && args["-g"].asString() == "randompkt") {
346 qgen = std::make_shared<RandomPktQueryGenerator>(config);
347 } else if (args["-g"] && args["-g"].asString() == "randomqname") {
348 qgen = std::make_shared<RandomQNameQueryGenerator>(config);
349 } else if (args["-g"] && args["-g"].asString() == "randomlabel") {
350 qgen = std::make_shared<RandomLabelQueryGenerator>(config);
351 } else {
352 qgen = std::make_shared<StaticQueryGenerator>(config);
353 }
354 qgen->set_args(args["GENOPTS"].asStringList());
355 qgen->set_qclass(args["--class"].asString());
356 qgen->set_loops(args["-n"].asLong());
357 qgen->set_dnssec(args["--dnssec"].asBool());
358 qgen->set_qname(args["-r"].asString());
359 qgen->set_qtype(args["-T"].asString());
360 qgen->init();
361 if (!qgen->synthesizedQueries() && qgen->size() == 0) {
362 throw std::runtime_error("no queries were generated");
363 }
364 } catch (const std::exception &e) {
365 std::cerr << "generator error: " << e.what() << std::endl;
366 return 1;
367 }
368
369 if (args["-R"].asBool()) {
370 qgen->randomize();
371 }
372
373 std::string cmdline{};
374 for (int i = 0; i < argc; i++) {
375 cmdline.append(argv[i]);
376 if (i != argc - 1) {
377 cmdline.push_back(' ');
378 }
379 }
380 auto metrics_mgr = std::make_shared<MetricsMgr>(loop, config, cmdline);
381
382 std::queue<std::pair<uint64_t, uint64_t>> qps_flow;
383 std::vector<std::shared_ptr<TokenBucket>> rl_list;
384 if (args["--qps-flow"]) {
385 parse_flowspec(args["--qps-flow"].asString(), qps_flow, config->verbosity(), c_count);
386 }
387
388 auto traf_config = std::make_shared<TrafGenConfig>();
389 traf_config->batch_count = b_count;
390 traf_config->family = family;
391 traf_config->bind_ip = bind_ip;
392 traf_config->target_list = target_list;
393 traf_config->port = static_cast<unsigned int>(args["-p"].asLong());
394 traf_config->s_delay = s_delay;
395 traf_config->protocol = proto;
396 #ifdef DOH_ENABLE
397 traf_config->method = method;
398 #endif
399 traf_config->r_timeout = args["-t"].asLong();
400
401 std::vector<std::shared_ptr<TrafGen>> throwers;
402 for (auto i = 0; i < c_count; i++) {
403 std::shared_ptr<TokenBucket> rl;
404 if (config->rate_limit()) {
405 rl = std::make_shared<TokenBucket>(config->rate_limit() / static_cast<double>(c_count));
406 } else if (args["--qps-flow"]) {
407 rl = std::make_shared<TokenBucket>();
408 rl_list.push_back(rl);
409 }
410 throwers.push_back(std::make_shared<TrafGen>(loop,
411 metrics_mgr->create_trafgen_metrics(),
412 config,
413 traf_config,
414 qgen,
415 rl));
416 throwers[i]->start();
417 }
418 if (args["--qps-flow"]) {
419 flow_change(qps_flow, rl_list, config->verbosity(), c_count);
420 }
421
422 auto have_in_flight = [&throwers]() {
423 for (const auto &i : throwers) {
424 if (i->in_flight_cnt()) {
425 return true;
426 }
427 }
428 return false;
429 };
430
431 auto shutdown = [&]() {
432 sigint->stop();
433 sigterm->stop();
434 if (run_timer.get())
435 run_timer->stop();
436 if (qgen_loop_timer.get())
437 qgen_loop_timer->stop();
438 for (auto &t : throwers) {
439 t->stop();
440 }
441 metrics_mgr->stop();
442 if (have_in_flight() && config->verbosity()) {
443 std::cout << "stopping, waiting up to " << traf_config->r_timeout << "s for in flight to finish..." << std::endl;
444 }
445 };
446
447 auto stop_traffic = [&](uvw::SignalEvent &, uvw::SignalHandle &) {
448 shutdown();
449 };
450
451 if (runtime_limit != 0) {
452 run_timer = loop->resource<uvw::TimerHandle>();
453 run_timer->on<uvw::TimerEvent>([&shutdown](const auto &, auto &) { shutdown(); });
454 run_timer->start(uvw::TimerHandle::Time{runtime_limit * 1000}, uvw::TimerHandle::Time{0});
455 }
456
457 if (qgen->loops()) {
458 qgen_loop_timer = loop->resource<uvw::TimerHandle>();
459 qgen_loop_timer->on<uvw::TimerEvent>([&qgen, &shutdown](const auto &, auto &) {
460 if (qgen->finished()) {
461 shutdown();
462 } });
463 qgen_loop_timer->start(uvw::TimerHandle::Time{500}, uvw::TimerHandle::Time{500});
464 }
465
466 sigint->on<uvw::SignalEvent>(stop_traffic);
467 sigint->start(SIGINT);
468
469 sigterm->on<uvw::SignalEvent>(stop_traffic);
470 sigterm->start(SIGTERM);
471
472 if (config->verbosity()) {
473 std::cout << "binding to " << traf_config->bind_ip << std::endl;
474 std::cout << "flaming target(s) [";
475 for (uint i = 0; i < 3; i++) {
476 std::cout << traf_config->target_list[i].address;
477 if (i == traf_config->target_list.size()-1) {
478 break;
479 }
480 else {
481 std::cout << ", ";
482 }
483 }
484 if (traf_config->target_list.size() > 3) {
485 std::cout << "and " << traf_config->target_list.size()-3 << " more";
486 }
487 std::cout << "] on port "
488 << args["-p"].asLong()
489 << " with " << c_count << " concurrent generators, each sending " << b_count
490 << " queries every " << s_delay << "ms on protocol " << args["-P"].asString()
491 << std::endl;
492 std::cout << "query generator [" << qgen->name() << "] contains " << qgen->size() << " record(s)" << std::endl;
493 if (args["-R"].asBool()) {
494 std::cout << "query list randomized" << std::endl;
495 }
496 if (config->rate_limit()) {
497 std::cout << "rate limit @ " << config->rate_limit() << " QPS (" << config->rate_limit() / static_cast<double>(c_count) <<
498 " QPS per concurrent sender)" << std::endl;
499 }
500 }
501
502 metrics_mgr->start();
503 loop->run();
504
505 // break from loop with ^C or timer
506 loop = nullptr;
507
508 // when loop is complete, finalize metrics
509 metrics_mgr->finalize();
510
511 return 0;
512 }
513