252 lines
7.7 KiB
Python
Executable file
252 lines
7.7 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
# Copyright (C) Internet Systems Consortium, Inc. ("ISC")
|
|
#
|
|
# SPDX-License-Identifier: MPL-2.0
|
|
#
|
|
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, you can obtain one at https://mozilla.org/MPL/2.0/.
|
|
#
|
|
# See the COPYRIGHT file distributed with this work for additional
|
|
# information regarding copyright ownership.
|
|
|
|
import os
|
|
import sys
|
|
import socket
|
|
import subprocess
|
|
import random
|
|
import time
|
|
|
|
from functools import reduce
|
|
from resource import getrlimit
|
|
from resource import setrlimit
|
|
from resource import RLIMIT_NOFILE
|
|
|
|
MULTIDIG_INSTANCES = 10
|
|
CONNECT_TRIES = 5
|
|
|
|
random.seed()
|
|
|
|
# Ensure we have enough file desriptors to work
|
|
rlimit_nofile = getrlimit(RLIMIT_NOFILE)
|
|
if rlimit_nofile[0] < 1024:
|
|
setrlimit(RLIMIT_NOFILE, (1024, rlimit_nofile[1]))
|
|
|
|
|
|
# Introduce some random delay
|
|
def jitter():
|
|
time.sleep((500 + random.randint(0, 250)) / 1000000.0)
|
|
|
|
|
|
# A set of simple procedures to get the test's configuration options
|
|
def get_http_port(http_secure=False):
|
|
http_port_env = None
|
|
if http_secure:
|
|
http_port_env = os.getenv("HTTPSPORT")
|
|
else:
|
|
http_port_env = os.getenv("HTTPPORT")
|
|
if http_port_env:
|
|
return int(http_port_env)
|
|
return 443
|
|
|
|
|
|
def get_http_host():
|
|
bind_host = os.getenv("BINDHOST")
|
|
if bind_host:
|
|
return bind_host
|
|
return "localhost"
|
|
|
|
|
|
def get_dig_path():
|
|
dig_path = os.getenv("DIG")
|
|
if dig_path:
|
|
return dig_path
|
|
return "dig"
|
|
|
|
|
|
# A simple class which creates the given number of TCP connections to
|
|
# the given host in order to stress the BIND's quota facility
|
|
class TCPConnector:
|
|
def __init__(self, host, port):
|
|
self.host = host
|
|
self.port = port
|
|
self.connections = []
|
|
|
|
def connect_one(self):
|
|
tries = CONNECT_TRIES
|
|
while tries > 0:
|
|
try:
|
|
sock = socket.create_connection(
|
|
address=(self.host, self.port), timeout=None
|
|
)
|
|
self.connections.append(sock)
|
|
break
|
|
except ConnectionResetError:
|
|
# some jitter for BSDs
|
|
jitter()
|
|
continue
|
|
except TimeoutError:
|
|
jitter()
|
|
continue
|
|
finally:
|
|
tries -= 1
|
|
|
|
# Close an established connection (randomly)
|
|
def disconnect_random(self):
|
|
pos = random.randint(0, len(self.connections) - 1)
|
|
conn = self.connections[pos]
|
|
try:
|
|
conn.shutdown(socket.SHUT_RDWR)
|
|
conn.close()
|
|
except OSError:
|
|
conn.close()
|
|
finally:
|
|
self.connections.remove(conn)
|
|
|
|
def disconnect_all(self):
|
|
while len(self.connections) != 0:
|
|
self.disconnect_random()
|
|
|
|
|
|
# A simple class which allows running a dig instance under control of
|
|
# the process
|
|
class SubDIG:
|
|
def __init__(self, http_secure=None, extra_args=None):
|
|
self.sub_process = None
|
|
self.dig_path = get_dig_path()
|
|
self.host = get_http_host()
|
|
self.port = get_http_port(http_secure=http_secure)
|
|
if http_secure:
|
|
self.http_secure = True
|
|
else:
|
|
self.http_secure = False
|
|
self.extra_args = extra_args
|
|
|
|
# This method constructs a command string
|
|
def get_command(self):
|
|
command = self.dig_path + " -p " + str(self.port) + " "
|
|
command = command + "+noadd +nosea +nostat +noquest +nocmd +time=30 "
|
|
if self.http_secure:
|
|
command = command + "+https "
|
|
else:
|
|
command = command + "+http-plain "
|
|
command = command + "@" + self.host + " "
|
|
if self.extra_args:
|
|
command = command + self.extra_args
|
|
return command
|
|
|
|
def run(self):
|
|
# pylint: disable=consider-using-with
|
|
with open(os.devnull, "w", encoding="utf-8") as devnull:
|
|
self.sub_process = subprocess.Popen(
|
|
self.get_command(), shell=True, stdout=devnull
|
|
)
|
|
|
|
def wait(self, timeout=None):
|
|
res = None
|
|
if timeout is None:
|
|
return self.sub_process.wait()
|
|
try:
|
|
res = self.sub_process.wait(timeout=timeout)
|
|
except subprocess.TimeoutExpired:
|
|
return None
|
|
return res
|
|
|
|
def alive(self):
|
|
return self.sub_process.poll() is None
|
|
|
|
|
|
# A simple wrapper class which allows running multiple dig instances
|
|
# and examining their statuses in one logical operation.
|
|
class MultiDIG:
|
|
def __init__(self, numdigs, http_secure=None, extra_args=None):
|
|
assert int(numdigs) > 0, f"numdigs={numdigs}"
|
|
digs = []
|
|
for _ in range(1, int(numdigs) + 1):
|
|
digs.append(SubDIG(http_secure=http_secure, extra_args=extra_args))
|
|
self.digs = digs
|
|
assert len(self.digs) == int(numdigs), f"len={len(self.digs)} numdigs={numdigs}"
|
|
|
|
def run(self):
|
|
for p in self.digs:
|
|
p.run()
|
|
|
|
def wait(self):
|
|
return map(lambda p: (p.wait()), self.digs)
|
|
|
|
# Wait for the all instances to terminate with expected given
|
|
# status. Returns true or false.
|
|
def wait_for_result(self, result):
|
|
return reduce(
|
|
lambda a, b: ((a == result or a is True) and b == result), self.wait()
|
|
)
|
|
|
|
def alive(self):
|
|
return reduce(lambda a, b: (a and b), map(lambda p: (p.alive()), self.digs))
|
|
|
|
def completed(self):
|
|
total = 0
|
|
for p in self.digs:
|
|
if not p.alive():
|
|
total += 1
|
|
return total
|
|
|
|
|
|
# The test's main logic
|
|
def run_test(http_secure=True):
|
|
query_args = "SOA ."
|
|
# Let's try to make a successful query
|
|
subdig = SubDIG(http_secure=http_secure, extra_args=query_args)
|
|
subdig.run()
|
|
assert subdig.wait() == 0, "DIG was expected to succeed"
|
|
# Let's create a lot of TCP connections to the server stress the
|
|
# HTTP quota
|
|
connector = TCPConnector(get_http_host(), get_http_port(http_secure=http_secure))
|
|
# Let's make queries until the quota kicks in
|
|
subdig = SubDIG(http_secure=http_secure, extra_args=query_args)
|
|
subdig.run()
|
|
while True:
|
|
connector.connect_one()
|
|
subdig = SubDIG(http_secure=http_secure, extra_args=query_args)
|
|
subdig.run()
|
|
if subdig.wait(timeout=5) is None:
|
|
break
|
|
|
|
# At this point quota has kicked in. Additionally, let's create a
|
|
# bunch of dig processes all trying to make a query against the
|
|
# server with exceeded quota
|
|
multidig = MultiDIG(
|
|
MULTIDIG_INSTANCES, http_secure=http_secure, extra_args=query_args
|
|
)
|
|
multidig.run()
|
|
# Wait for the dig instance to complete. Not a single instance has
|
|
# a chance to complete successfully because of the exceeded quota
|
|
assert (
|
|
subdig.wait(timeout=5) is None
|
|
), "The single DIG instance has stopped prematurely"
|
|
assert subdig.alive(), "The single DIG instance is expected to be alive"
|
|
assert multidig.alive(), (
|
|
"The DIG instances from the set are all expected to "
|
|
"be alive, but {} of them have completed"
|
|
).format(multidig.completed())
|
|
# Let's close opened connections (in random order) to let all dig
|
|
# processes to complete
|
|
connector.disconnect_all()
|
|
# Wait for all processes to complete successfully
|
|
assert subdig.wait() == 0, "Single DIG instance failed"
|
|
assert (
|
|
multidig.wait_for_result(0) is True
|
|
), "One or more of DIG instances returned unexpected results"
|
|
|
|
|
|
def main():
|
|
run_test(http_secure=True)
|
|
run_test(http_secure=False)
|
|
# If we have reached this point we could safely return 0
|
|
# (success). If the test fails because of an assert, the whole
|
|
# program will return non-zero exit code and produce the backtrace
|
|
return 0
|
|
|
|
|
|
sys.exit(main())
|