diff --git a/tests-extra/tests/modules/cookies/test.py b/tests-extra/tests/modules/cookies/test.py index 3be0096879e796a71042c78156c05ffd5163f0e7..73acd9d1d36305fd7790578901f8e275acb99296 100644 --- a/tests-extra/tests/modules/cookies/test.py +++ b/tests-extra/tests/modules/cookies/test.py @@ -30,19 +30,12 @@ def reconfigure(server, zone, secret_lifetime, badcookie_slip): server.reload() server.zone_wait(zone) -def check_rcode(server, query, rcode, msg): +def check_rcode(server, query, rcode, msg, tcp=False): try: - response = dns.query.udp(query, server.addr, port=server.port, timeout=1) - except dns.exception.Timeout: - response = None - if response is None: - return None - compare(response.rcode(), rcode, msg) - return response - -def check_rcode_tcp(server, query, rcode, msg): - try: - response = dns.query.tcp(query, server.addr, port=server.port, timeout=1) + if tcp: + response = dns.query.tcp(query, server.addr, port=server.port, timeout=1) + else: + response = dns.query.udp(query, server.addr, port=server.port, timeout=1) except dns.exception.Timeout: response = None if response is None: @@ -74,7 +67,12 @@ check_rcode(knot, query, rcodeNoerror, "NO COOKIE OPT") # Try a query without a server cookie over TCP cookieOpt = dns.edns.option_from_wire(cookieOpcode, clientCookie, 0, clientCookieLen) query = dns.message.make_query("dns1.example.com", "A", use_edns=True, options=[cookieOpt]) -response = check_rcode_tcp(knot, query, rcodeNoerror, "ONLY CLIENT COOKIE [TCP]") +response = check_rcode(knot, query, rcodeNoerror, "ONLY CLIENT COOKIE [TCP]", tcp=True) + +# Try a query with the received cookie over TCP +cookieOpt = response.options[0] +query = dns.message.make_query("dns1.example.com", "A", use_edns=True, options=[cookieOpt]) +check_rcode(knot, query, rcodeNoerror, "CORRECT COOKIE", tcp=True) # Try a query without a server cookie cookieOpt = dns.edns.option_from_wire(cookieOpcode, clientCookie, 0, clientCookieLen)