Commit 06ef4950 authored by Alexandre's avatar Alexandre
Browse files

Merge branch 'streams-async' into streams

parents 61353d89 b50f9227
......@@ -62,6 +62,7 @@ Possible options, besides `--dot`:
* --file INPUT_FILE: provide an input file with a list of domain name to query
(read the first line only, use --repeat N to read up to N lines of the file)
* --repeat N: repeat a test N times or read up to N lines of a file
* --silent: suppress response output (with --multistreams only)
### Check
......
......@@ -53,6 +53,8 @@ forceIPv4 = False
forceIPv6 = False
connectTo = None
multistreams = False
sync = False
silent = False
check = False
mandatory_level = None
check_additional = True
......@@ -313,8 +315,8 @@ class Connection:
def __str__(self):
return self.server
def do_test(self, qname, qtype=rtype, synchronous=True):
# Routine doing one actual test. Returns a Request object
def do_test(self, request):
# Routine doing one actual test. Returns nothing
pass
......@@ -427,45 +429,13 @@ class ConnectionDoT(Connection):
self.send_data(request.data)
self.receive_data(request)
def do_test(self, qname, qtype=rtype, synchronous=True):
request = RequestDoT(qname, qtype, want_dnssec=dnssec, use_edns=edns)
request.to_wire()
def do_test(self, request, synchronous=True):
self.send_and_receive(request)
request.check_response()
return request
class CurlHandle():
def __init__(self, connect=None, forceIPv4=False, forceIPv6=False,
verbose=verbose, insecure=insecure):
self.handle = pycurl.Curl()
# Does not work if pycurl was not compiled with nghttp2 (recent Debian
# packages are OK) https://github.com/pycurl/pycurl/issues/477
self.handle.setopt(pycurl.HTTP_VERSION, pycurl.CURL_HTTP_VERSION_2)
if verbose:
self.handle.setopt(pycurl.VERBOSE, True)
if insecure:
self.handle.setopt(pycurl.SSL_VERIFYPEER, False)
self.handle.setopt(pycurl.SSL_VERIFYHOST, False)
if forceIPv4:
self.handle.setopt(pycurl.IPRESOLVE, pycurl.IPRESOLVE_V4)
if forceIPv6:
self.handle.setopt(pycurl.IPRESOLVE, pycurl.IPRESOLVE_V6)
if connect is not None:
family, repraddress = check_ip_address(connect, dot=False)
self.handle.setopt(pycurl.CONNECT_TO, [f'::{repraddress}:443',])
self.handle.setopt(pycurl.HTTPHEADER, ["Accept: application/dns-message", "Content-type: application/dns-message"])
def close(self):
self.handle.close()
def setopt(self, key, value):
self.handle.setopt(key, value)
def perform(self):
self.handle.perform()
def reset_opt_default(self):
def create_handle(connection):
def reset_opt_default(handle):
opts = {
pycurl.NOBODY: False,
pycurl.POST: False,
......@@ -473,35 +443,46 @@ class CurlHandle():
pycurl.URL: ''
}
for opt, value in opts.items():
self.handle.setopt(opt, value)
handle.setopt(opt, value)
def prepare(self, connection, request):
def prepare(handle, connection, request):
if not connection.multistreams:
self.reset_opt_default()
handle.reset_opt_default(handle)
if request.post:
self.prepare_post(connection, request)
elif request.head:
self.prepare_head(connection, request)
handle.setopt(pycurl.POST, True)
handle.setopt(pycurl.POSTFIELDS, request.data)
handle.setopt(pycurl.URL, connection.server)
else:
self.prepare_get(connection, request)
request.buffer = io.BytesIO()
self.handle.setopt(pycurl.WRITEDATA, request.buffer)
def prepare_get(self, connection, request):
self.handle.setopt(pycurl.HTTPGET, True)
dns_req = base64.urlsafe_b64encode(request.data).decode('UTF8').rstrip('=')
self.handle.setopt(pycurl.URL, connection.server + ("?dns=%s" % dns_req))
def prepare_post(self, connection, request):
request.post = True
self.handle.setopt(pycurl.POST, True)
self.handle.setopt(pycurl.POSTFIELDS, request.data)
self.handle.setopt(pycurl.URL, connection.server)
def prepare_head(self, connection, request):
request.head = True
self.prepare_get(connection, request)
self.handle.setopt(pycurl.NOBODY, True)
handle.setopt(pycurl.HTTPGET, True) # automatically sets CURLOPT_NOBODY to 0
if request.head:
handle.setopt(pycurl.NOBODY, True)
dns_req = base64.urlsafe_b64encode(request.data).decode('UTF8').rstrip('=')
handle.setopt(pycurl.URL, connection.server + ("?dns=%s" % dns_req))
handle.buffer = io.BytesIO()
handle.setopt(pycurl.WRITEDATA, handle.buffer)
handle.request = request
handle = pycurl.Curl()
# Does not work if pycurl was not compiled with nghttp2 (recent Debian
# packages are OK) https://github.com/pycurl/pycurl/issues/477
handle.setopt(pycurl.HTTP_VERSION, pycurl.CURL_HTTP_VERSION_2)
if connection.verbose:
handle.setopt(pycurl.VERBOSE, True)
if connection.insecure:
handle.setopt(pycurl.SSL_VERIFYPEER, False)
handle.setopt(pycurl.SSL_VERIFYHOST, False)
if connection.forceIPv4:
handle.setopt(pycurl.IPRESOLVE, pycurl.IPRESOLVE_V4)
if connection.forceIPv6:
handle.setopt(pycurl.IPRESOLVE, pycurl.IPRESOLVE_V6)
if connection.connect is not None:
family, repraddress = check_ip_address(connection.connect, dot=False)
handle.setopt(pycurl.CONNECT_TO, [f'::{repraddress}:443',])
handle.setopt(pycurl.HTTPHEADER,
["Accept: application/dns-message", "Content-type: application/dns-message"])
handle.reset_opt_default = reset_opt_default
handle.prepare = prepare
return handle
class ConnectionDoH(Connection):
......@@ -515,11 +496,9 @@ class ConnectionDoH(Connection):
self.multistreams = multistreams
if self.multistreams:
self.multi = self.create_multi()
self.all_handles = []
else:
self.curl_handle = self.create_handle()
def create_handle(self):
return CurlHandle(self.connect, self.forceIPv4, self.forceIPv6, self.verbose, self.insecure)
self.curl_handle = create_handle(self)
def create_multi(self):
multi = pycurl.CurlMulti()
......@@ -540,14 +519,6 @@ class ConnectionDoH(Connection):
h.close()
self.multi.remove_handle(h)
def perform(self, request):
request.buffer = io.BytesIO()
self.curl_handle.setopt(pycurl.WRITEDATA, request.buffer)
try:
self.curl_handle.perform()
except pycurl.error as e:
error(e.args[1])
def perform_multi(self):
while 1:
ret, num_handles = self.multi.perform()
......@@ -559,15 +530,24 @@ class ConnectionDoH(Connection):
continue
while 1:
ret, num_handles = self.multi.perform()
if not sync:
n, handle_pass, handle_fail = self.multi.info_read()
for handle in handle_pass:
self.read_result_handle(handle)
if ret != pycurl.E_CALL_MULTI_PERFORM:
break
def send(self, request):
self.perform(request)
def send(self, handle):
handle.buffer = io.BytesIO()
handle.setopt(pycurl.WRITEDATA, handle.buffer)
try:
handle.perform()
except pycurl.error as e:
error(e.args[1])
def receive(self, request):
handle = request.handle.handle
body = request.buffer.getvalue()
def receive(self, handle):
request = handle.request
body = handle.buffer.getvalue()
body_size = len(body)
http_code = handle.getinfo(pycurl.RESPONSE_CODE)
try:
......@@ -578,28 +558,34 @@ class ConnectionDoH(Connection):
request.response_size = body_size
request.rcode = http_code
request.ctype = content_type
request.buffer.close()
handle.buffer.close()
def send_and_receive(self, request):
self.send(request)
self.receive(request)
def send_and_receive(self, handle):
self.send(handle)
self.receive(handle)
def do_test(self, qname, qtype=rtype, synchronous=True):
request = RequestDoH(qname, qtype, want_dnssec=dnssec, use_edns=edns)
request.head = head
request.post = post
request.to_wire()
def read_result_handle(self, handle):
self.receive(handle)
handle.request.check_response()
if not silent:
print("Return code %s:\n%s\n" % (handle.request.rcode, handle.request.response))
def read_results(self):
for handle in self.all_handles:
self.read_result_handle(handle)
def do_test(self, request, synchronous=True):
if synchronous:
request.handle = self.curl_handle
handle = self.curl_handle
else:
request.handle = self.create_handle()
request.handle.prepare(self, request)
handle = create_handle(self)
self.all_handles.append(handle)
handle.prepare(handle, self, request)
if synchronous:
self.send_and_receive(request)
self.send_and_receive(handle)
request.check_response()
else:
self.multi.add_handle(request.handle.handle)
return request
self.multi.add_handle(handle)
def get_next_domain(input_file):
......@@ -657,19 +643,11 @@ def print_result(connection, request, prefix=None, display_err=True):
ok = False
return ok
# pending_requests must be an array
def read_results(connection, pending_requests):
for i in range(0, len(pending_requests)):
request = pending_requests[i]
connection.receive(request)
request.check_response()
return pending_requests
def create_request(dot=dot, trunc=False, **req_args):
def create_request(qname, qtype=rtype, use_edns=edns, want_dnssec=dnssec, dot=dot, trunc=False):
if dot:
request = RequestDoT(**req_args)
request = RequestDoT(qname, rtype, use_edns, want_dnssec)
else:
request = RequestDoH(**req_args)
request = RequestDoH(qname, rtype, use_edns, want_dnssec)
if trunc:
request.trunc_data()
else:
......@@ -712,10 +690,13 @@ def run_check_default(connection):
request.post = True
elif method == DOH_HEAD:
request.head = True
request.handle = connection.curl_handle
request.handle.prepare(connection, request)
handle = connection.curl_handle
handle.prepare(handle, connection, request)
bundle = handle
else:
bundle = request
try:
connection.send_and_receive(request)
connection.send_and_receive(bundle)
except CustomException as e:
ok = False
error(e)
......@@ -735,11 +716,11 @@ def run_check_mime(connection, accept="application/dns-message", content_type="a
header = [f"Accept: {accept}", f"Content-type: {content_type}"]
req_args = { 'qname': name, 'qtype': rtype, 'use_edns': edns, 'want_dnssec': dnssec }
request = create_request(**req_args)
connection.curl_handle.setopt(pycurl.HTTPHEADER, header)
request.handle = connection.curl_handle
request.handle.prepare(connection, request)
handle = connection.curl_handle
handle.setopt(pycurl.HTTPHEADER, header)
handle.prepare(handle, connection, request)
try:
connection.send_and_receive(request)
connection.send_and_receive(handle)
except CustomException as e:
ok = False
error(e)
......@@ -748,7 +729,7 @@ def run_check_mime(connection, accept="application/dns-message", content_type="a
ok = False
default = "application/dns-message"
default_header = [f"Accept: {default}", f"Content-type: {default}"]
connection.curl_handle.setopt(pycurl.HTTPHEADER, default_header)
handle.setopt(pycurl.HTTPHEADER, default_header)
return ok
def run_check_trunc(connection):
......@@ -759,14 +740,16 @@ def run_check_trunc(connection):
req_args = { 'qname': name, 'qtype': rtype, 'use_edns': edns, 'want_dnssec': dnssec }
if dot:
request = create_request(dot=dot, trunc=True, **req_args)
bundle = request
else:
request = create_request(trunc=True, **req_args)
request.post = True
request.handle = connection.curl_handle
request.handle.prepare(connection, request)
handle = connection.curl_handle
handle.prepare(handle, connection, request)
bundle = handle
try:
# 8.8.8.8 replies FORMERR but most DoT servers violently shut down the connection (which is legal)
connection.send_and_receive(request)
connection.send_and_receive(bundle)
except CustomException as e:
ok = False
error(e)
......@@ -812,6 +795,7 @@ if not monitoring:
optlist, args = getopt.getopt (sys.argv[1:], "hvPkeV:r:f:d:t46",
["help", "verbose", "dot", "head",
"insecure", "POST", "vhost=", "multistreams",
"sync", "silent",
"dnssec", "noedns", "ecs", "repeat=", "file=", "delay=",
"key=", "nosni",
"v4only", "v6only",
......@@ -834,6 +818,10 @@ if not monitoring:
insecure = True
elif option == "--multistreams":
multistreams = True
elif option == "--sync":
sync = True
elif option == "--silent":
silent = True
elif option == "--dnssec":
dnssec = True
elif option == "--nosni":
......@@ -887,6 +875,12 @@ if not monitoring:
if multistreams and ifile is None:
usage("Multi-streams requires an input file")
sys.exit(1)
if sync and not multistreams:
usage("--sync cannot be used without --multistreams")
sys.exit(1)
if silent and not multistreams:
usage("--silent cannot be used without --multistreams")
sys.exit(1)
if not edns and not no_ecs:
usage("ECS requires EDNS")
sys.exit(1)
......@@ -911,8 +905,6 @@ if not monitoring:
name = args[1]
if len(args) == 3:
rtype = args[2]
if multistreams:
pending = {}
else: # Monitoring plugin
dot = (me == "check_dot")
name = None
......@@ -1049,12 +1041,16 @@ for connectTo in ip_set:
input = open(ifile)
if not check:
for i in range (0, tests):
if tests > 1:
if tests > 1 and (verbose or not silent):
print("\nTest %i" % i)
if ifile is not None:
name, rtype = get_next_domain(input)
request = create_request(name, qtype=rtype, use_edns=edns, want_dnssec=dnssec, dot=dot)
if not dot:
request.head = head
request.post = post
try:
request = conn.do_test(name, rtype, synchronous = not multistreams)
conn.do_test(request, synchronous = not multistreams)
except (OpenSSL.SSL.Error, CustomException) as e:
ok = False
error(e)
......@@ -1062,18 +1058,15 @@ for connectTo in ip_set:
if not multistreams:
if not print_result(conn, request):
ok = False
else: # We do multistreams
pending[i] = request # No result yet
if tests > 1 and i == 0:
start2 = time.time()
if delay is not None:
time.sleep(delay)
if multistreams:
conn.perform_multi()
print("")
result = read_results(conn, pending)
for j in result:
print("Return code %s: %s\n" % (result[j].rcode, result[j].response))
if sync:
print()
conn.read_results()
else:
ok = run_check(conn) and ok # need to run run_check first
stop = time.time()
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment