xref: /netbsd-src/external/mpl/bind/dist/bin/tests/system/tcp/ans6/ans.py (revision b5c47949a45ac972130c38cf13dfd8afb1f09285) 
1############################################################################
2# Copyright (C) Internet Systems Consortium, Inc. ("ISC")
3#
4# This Source Code Form is subject to the terms of the Mozilla Public
5# License, v. 2.0. If a copy of the MPL was not distributed with this
6# file, you can obtain one at https://mozilla.org/MPL/2.0/.
7#
8# See the COPYRIGHT file distributed with this work for additional
9# information regarding copyright ownership.
10############################################################################
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
47def log(msg):
48    print(datetime.datetime.now().strftime('%d-%b-%Y %H:%M:%S.%f ') + msg)
49
50
51def open_connections(active_conns, count, host, port):
52    queued = []
53    errors = []
54
55    try:
56        socket.inet_aton(host)
57        family = socket.AF_INET
58    except socket.error:
59        family = socket.AF_INET6
60
61    log('Opening %d connections...' % count)
62
63    for _ in range(count):
64        sock = socket.socket(family, socket.SOCK_STREAM)
65        sock.setblocking(0)
66        err = sock.connect_ex((host, port))
67        if err not in (0, errno.EINPROGRESS):
68            log('%s on connect for socket %s' % (errno.errorcode[err], sock))
69            errors.append(sock)
70        else:
71            queued.append(sock)
72
73    start = time.time()
74    while queued:
75        now = time.time()
76        time_left = OPEN_TIMEOUT - (now - start)
77        if time_left <= 0:
78            break
79        _, wsocks, _ = select.select([], queued, [], time_left)
80        for sock in wsocks:
81            queued.remove(sock)
82            err = sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
83            if err:
84                log('%s for socket %s' % (errno.errorcode[err], sock))
85                errors.append(sock)
86            else:
87                sock.send(VERSION_QUERY)
88                active_conns.append(sock)
89
90    if errors:
91        log('result=FAIL: %d connection(s) failed' % len(errors))
92    elif queued:
93        log('result=FAIL: Timed out, aborting %d pending connections' % len(queued))
94        for sock in queued:
95            sock.close()
96    else:
97        log('result=OK: Successfully opened %d connections' % count)
98
99
100def close_connections(active_conns, count):
101    log('Closing %s connections...' % "all" if count == 0 else str(count))
102    if count == 0:
103        count = len(active_conns)
104    for _ in range(count):
105        sock = active_conns.pop(0)
106        sock.close()
107    log('result=OK: Successfully closed %d connections' % count)
108
109
110def sigterm(*_):
111    log('SIGTERM received, shutting down')
112    os.remove('ans.pid')
113    sys.exit(0)
114
115
116def main():
117    active_conns = []
118
119    signal.signal(signal.SIGTERM, sigterm)
120
121    with open('ans.pid', 'w') as pidfile:
122        print(os.getpid(), file=pidfile)
123
124    listenip = '10.53.0.6'
125    try:
126        port = int(os.environ['CONTROLPORT'])
127    except KeyError:
128        port = 5309
129
130    log('Listening on %s:%d' % (listenip, port))
131
132    ctlsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
133    ctlsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
134    ctlsock.bind((listenip, port))
135    ctlsock.listen(1)
136
137    while True:
138        (clientsock, _) = ctlsock.accept()
139        log('Accepted control connection from %s' % clientsock)
140        cmdline = clientsock.recv(512).decode('ascii').strip()
141        if cmdline:
142            log('Received command: %s' % cmdline)
143            cmd = cmdline.split()
144            if cmd[0] == 'open':
145                count, host, port = cmd[1:]
146                open_connections(active_conns, int(count), host, int(port))
147            elif cmd[0] == 'close':
148                (count, ) = cmd[1:]
149                close_connections(active_conns, int(count))
150            else:
151                log('result=FAIL: Unknown command')
152        clientsock.close()
153
154
155if __name__ == '__main__':
156    main()
157