1 files changed, 252 insertions, 0 deletions
diff --git a/bin/tests/system/doth/stress_http_quota.py b/bin/tests/system/doth/stress_http_quota.py
new file mode 100755
index 0000000..12e29c8
--- /dev/null
+++ b/bin/tests/system/doth/stress_http_quota.py
@@ -0,0 +1,252 @@
+#!/usr/bin/env python
+
+# 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
+        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)
+
+    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())