xref: /netbsd-src/external/mpl/bind/dist/bin/tests/system/tcp/ans6/ans.py (revision f281902de12281841521aa31ef834ad944d725e2) 
1# Copyright (C) Internet Systems Consortium, Inc. ("ISC")
2#
3# SPDX-License-Identifier: MPL-2.0
4#
5# This Source Code Form is subject to the terms of the Mozilla Public
6# License, v. 2.0.  If a copy of the MPL was not distributed with this
7# file, you can obtain one at https://mozilla.org/MPL/2.0/.
8#
9# See the COPYRIGHT file distributed with this work for additional
10# information regarding copyright ownership.
11
12############################################################################
13#
14# This tool allows an arbitrary number of TCP connections to be made to the
15# specified service and to keep them open until told otherwise.  It is
16# controlled by writing text commands to a TCP socket (default port: 5309).
17#
18# Currently supported commands:
19#
20#   - open <COUNT> <HOST> <PORT>
21#
22#     Opens <COUNT> TCP connections to <HOST>:<PORT> and keeps them open.
23#     <HOST> must be an IP address (IPv4 or IPv6).
24#
25#   - close <COUNT>
26#
27#     Close the oldest <COUNT> previously established connections.
28#
29############################################################################
30
31from __future__ import print_function
32
33import datetime
34import errno
35import os
36import select
37import signal
38import socket
39import sys
40import time
41
42
43# Timeout for establishing all connections requested by a single 'open' command.
44OPEN_TIMEOUT = 2
45VERSION_QUERY = b"\x00\x1e\xaf\xb8\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x07version\x04bind\x00\x00\x10\x00\x03"
46
47
48def log(msg):
49    print(datetime.datetime.now().strftime("%d-%b-%Y %H:%M:%S.%f ") + msg)
50
51
52def open_connections(active_conns, count, host, port):
53    queued = []
54    errors = []
55
56    try:
57        socket.inet_aton(host)
58        family = socket.AF_INET
59    except socket.error:
60        family = socket.AF_INET6
61
62    log("Opening %d connections..." % count)
63
64    for _ in range(count):
65        sock = socket.socket(family, socket.SOCK_STREAM)
66        sock.setblocking(0)
67        err = sock.connect_ex((host, port))
68        if err not in (0, errno.EINPROGRESS):
69            log("%s on connect for socket %s" % (errno.errorcode[err], sock))
70            errors.append(sock)
71        else:
72            queued.append(sock)
73
74    start = time.time()
75    while queued:
76        now = time.time()
77        time_left = OPEN_TIMEOUT - (now - start)
78        if time_left <= 0:
79            break
80        _, wsocks, _ = select.select([], queued, [], time_left)
81        for sock in wsocks:
82            queued.remove(sock)
83            err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
84            if err:
85                log("%s for socket %s" % (errno.errorcode[err], sock))
86                errors.append(sock)
87            else:
88                sock.send(VERSION_QUERY)
89                active_conns.append(sock)
90
91    if errors:
92        log("result=FAIL: %d connection(s) failed" % len(errors))
93    elif queued:
94        log("result=FAIL: Timed out, aborting %d pending connections" % len(queued))
95        for sock in queued:
96            sock.close()
97    else:
98        log("result=OK: Successfully opened %d connections" % count)
99
100
101def close_connections(active_conns, count):
102    log("Closing %s connections..." % "all" if count == 0 else str(count))
103    if count == 0:
104        count = len(active_conns)
105    for _ in range(count):
106        sock = active_conns.pop(0)
107        sock.close()
108    log("result=OK: Successfully closed %d connections" % count)
109
110
111def sigterm(*_):
112    log("SIGTERM received, shutting down")
113    os.remove("ans.pid")
114    sys.exit(0)
115
116
117def main():
118    active_conns = []
119
120    signal.signal(signal.SIGTERM, sigterm)
121
122    with open("ans.pid", "w") as pidfile:
123        print(os.getpid(), file=pidfile)
124
125    listenip = "10.53.0.6"
126    try:
127        port = int(os.environ["CONTROLPORT"])
128    except KeyError:
129        port = 5309
130
131    log("Listening on %s:%d" % (listenip, port))
132
133    ctlsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
134    ctlsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
135    ctlsock.bind((listenip, port))
136    ctlsock.listen(1)
137
138    while True:
139        (clientsock, _) = ctlsock.accept()
140        log("Accepted control connection from %s" % clientsock)
141        cmdline = clientsock.recv(512).decode("ascii").strip()
142        if cmdline:
143            log("Received command: %s" % cmdline)
144            cmd = cmdline.split()
145            if cmd[0] == "open":
146                count, host, port = cmd[1:]
147                open_connections(active_conns, int(count), host, int(port))
148            elif cmd[0] == "close":
149                (count,) = cmd[1:]
150                close_connections(active_conns, int(count))
151            else:
152                log("result=FAIL: Unknown command")
153        clientsock.close()
154
155
156if __name__ == "__main__":
157    main()
158