diff options
Diffstat (limited to 'bin/tests/system/tcp/ans6/ans.py')
-rw-r--r-- | bin/tests/system/tcp/ans6/ans.py | 157 |
1 files changed, 157 insertions, 0 deletions
diff --git a/bin/tests/system/tcp/ans6/ans.py b/bin/tests/system/tcp/ans6/ans.py new file mode 100644 index 0000000..4595ddc --- /dev/null +++ b/bin/tests/system/tcp/ans6/ans.py @@ -0,0 +1,157 @@ +# 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. + +############################################################################ +# +# This tool allows an arbitrary number of TCP connections to be made to the +# specified service and to keep them open until told otherwise. It is +# controlled by writing text commands to a TCP socket (default port: 5309). +# +# Currently supported commands: +# +# - open <COUNT> <HOST> <PORT> +# +# Opens <COUNT> TCP connections to <HOST>:<PORT> and keeps them open. +# <HOST> must be an IP address (IPv4 or IPv6). +# +# - close <COUNT> +# +# Close the oldest <COUNT> previously established connections. +# +############################################################################ + +from __future__ import print_function + +import datetime +import errno +import os +import select +import signal +import socket +import sys +import time + + +# Timeout for establishing all connections requested by a single 'open' command. +OPEN_TIMEOUT = 2 +VERSION_QUERY = b"\x00\x1e\xaf\xb8\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x07version\x04bind\x00\x00\x10\x00\x03" + + +def log(msg): + print(datetime.datetime.now().strftime("%d-%b-%Y %H:%M:%S.%f ") + msg) + + +def open_connections(active_conns, count, host, port): + queued = [] + errors = [] + + try: + socket.inet_aton(host) + family = socket.AF_INET + except socket.error: + family = socket.AF_INET6 + + log("Opening %d connections..." % count) + + for _ in range(count): + sock = socket.socket(family, socket.SOCK_STREAM) + sock.setblocking(0) + err = sock.connect_ex((host, port)) + if err not in (0, errno.EINPROGRESS): + log("%s on connect for socket %s" % (errno.errorcode[err], sock)) + errors.append(sock) + else: + queued.append(sock) + + start = time.time() + while queued: + now = time.time() + time_left = OPEN_TIMEOUT - (now - start) + if time_left <= 0: + break + _, wsocks, _ = select.select([], queued, [], time_left) + for sock in wsocks: + queued.remove(sock) + err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR) + if err: + log("%s for socket %s" % (errno.errorcode[err], sock)) + errors.append(sock) + else: + sock.send(VERSION_QUERY) + active_conns.append(sock) + + if errors: + log("result=FAIL: %d connection(s) failed" % len(errors)) + elif queued: + log("result=FAIL: Timed out, aborting %d pending connections" % len(queued)) + for sock in queued: + sock.close() + else: + log("result=OK: Successfully opened %d connections" % count) + + +def close_connections(active_conns, count): + log("Closing %s connections..." % "all" if count == 0 else str(count)) + if count == 0: + count = len(active_conns) + for _ in range(count): + sock = active_conns.pop(0) + sock.close() + log("result=OK: Successfully closed %d connections" % count) + + +def sigterm(*_): + log("SIGTERM received, shutting down") + os.remove("ans.pid") + sys.exit(0) + + +def main(): + active_conns = [] + + signal.signal(signal.SIGTERM, sigterm) + + with open("ans.pid", "w") as pidfile: + print(os.getpid(), file=pidfile) + + listenip = "10.53.0.6" + try: + port = int(os.environ["CONTROLPORT"]) + except KeyError: + port = 5309 + + log("Listening on %s:%d" % (listenip, port)) + + ctlsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + ctlsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + ctlsock.bind((listenip, port)) + ctlsock.listen(1) + + while True: + (clientsock, _) = ctlsock.accept() + log("Accepted control connection from %s" % clientsock) + cmdline = clientsock.recv(512).decode("ascii").strip() + if cmdline: + log("Received command: %s" % cmdline) + cmd = cmdline.split() + if cmd[0] == "open": + count, host, port = cmd[1:] + open_connections(active_conns, int(count), host, int(port)) + elif cmd[0] == "close": + (count,) = cmd[1:] + close_connections(active_conns, int(count)) + else: + log("result=FAIL: Unknown command") + clientsock.close() + + +if __name__ == "__main__": + main() |