mirror of
https://gitlab.isc.org/isc-projects/bind9
synced 2025-10-11 13:47:54 +00:00
For duration measurements, i.e. deadlines and timeouts, it's more suitable to use monotonic time as it's guaranteed to only go forward, unlike time.time() which can be affected by local clock settings.
158 lines
4.6 KiB
Python
158 lines
4.6 KiB
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.
|
|
|
|
############################################################################
|
|
#
|
|
# 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.monotonic()
|
|
while queued:
|
|
now = time.monotonic()
|
|
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()
|