python-werkzeug/ephemeral_port_reserve.py
2023-05-10 14:46:54 +08:00

62 lines
1.9 KiB
Python
Executable File

#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import unicode_literals
import contextlib
import errno
from socket import error as SocketError
from socket import SO_REUSEADDR
from socket import socket
from socket import SOL_SOCKET
LOCALHOST = '127.0.0.1'
def reserve(ip=LOCALHOST, port=0):
"""Bind to an ephemeral port, force it into the TIME_WAIT state, and unbind it.
This means that further ephemeral port alloctions won't pick this "reserved" port,
but subprocesses can still bind to it explicitly, given that they use SO_REUSEADDR.
By default on linux you have a grace period of 60 seconds to reuse this port.
To check your own particular value:
$ cat /proc/sys/net/ipv4/tcp_fin_timeout
60
By default, the port will be reserved for localhost (aka 127.0.0.1).
To reserve a port for a different ip, provide the ip as the first argument.
Note that IP 0.0.0.0 is interpreted as localhost.
"""
port = int(port)
with contextlib.closing(socket()) as s:
s.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
try:
s.bind((ip, port))
except SocketError as e:
# socket.error: EADDRINUSE Address already in use
if e.errno == errno.EADDRINUSE and port != 0:
s.bind((ip, 0))
else:
raise
# the connect below deadlocks on kernel >= 4.4.0 unless this arg is greater than zero
s.listen(1)
sockname = s.getsockname()
# these three are necessary just to get the port into a TIME_WAIT state
with contextlib.closing(socket()) as s2:
s2.connect(sockname)
sock, _ = s.accept()
with contextlib.closing(sock):
return sockname[1]
def main(): # pragma: no cover
from sys import argv
port = reserve(*argv[1:])
print(port)
if __name__ == '__main__':
exit(main())