Skip to content

Commit e08f3fb

Browse files
committed
Authenticate with getsocketname and getpeername
1 parent d079ce3 commit e08f3fb

File tree

2 files changed

+58
-88
lines changed

2 files changed

+58
-88
lines changed

Lib/socket.py

+28-73
Original file line numberDiff line numberDiff line change
@@ -628,81 +628,36 @@ def socketpair(family=AF_INET, type=SOCK_STREAM, proto=0):
628628
if proto != 0:
629629
raise ValueError("Only protocol zero is supported")
630630

631-
import secrets # Delay import until actually needed.
632-
auth_len = secrets.DEFAULT_ENTROPY
633-
reason = "unknown"
634-
for _ in range(5):
635-
# We do not try forever, that'd just provide another process
636-
# the ability to make us waste CPU retrying. In all normal
637-
# circumstances the first connection auth attempt succeeds.
638-
s_auth = secrets.token_bytes()
639-
c_auth = secrets.token_bytes()
640-
641-
# We create a connected TCP socket. Note the trick with
642-
# setblocking(False) prevents us from having to create a thread.
643-
csock = None
644-
ssock = None
645-
reason = "unknown"
646-
lsock = socket(family, type, proto)
631+
# We create a connected TCP socket. Note the trick with
632+
# setblocking(False) that prevents us from having to create a thread.
633+
lsock = socket(family, type, proto)
634+
try:
635+
lsock.bind((host, 0))
636+
lsock.listen()
637+
# On IPv6, ignore flow_info and scope_id
638+
addr, port = lsock.getsockname()[:2]
639+
csock = socket(family, type, proto)
647640
try:
648-
lsock.bind((host, 0))
649-
lsock.listen()
650-
# On IPv6, ignore flow_info and scope_id
651-
addr, port = lsock.getsockname()[:2]
652-
csock = socket(family, type, proto)
653-
try:
654-
csock.setblocking(False)
655-
try:
656-
csock.connect((addr, port))
657-
except (BlockingIOError, InterruptedError):
658-
pass
659-
csock.setblocking(True)
660-
ssock, _ = lsock.accept()
661-
except Exception:
662-
csock.close()
663-
raise
664-
665-
def authenticate_socket_conn(send_sock, recv_sock, auth_bytes):
666-
nonlocal auth_len
667-
data_buf = bytearray(auth_len)
668-
data_mem = memoryview(data_buf)
669-
data_len = 0
670-
671-
# Send the authentication bytes.
672-
if send_sock.send(auth_bytes) != auth_len:
673-
raise ConnectionError("send() sent too few auth bytes.")
674-
675-
# Attempt to read the authentication bytes from the socket.
676-
max_auth_time = time.monotonic() + 3.0
677-
while time.monotonic() < max_auth_time and data_len < auth_len:
678-
bytes_received = recv_sock.recv_into(data_mem, auth_len - data_len)
679-
if bytes_received == 0:
680-
break # Connection closed.
681-
data_len += bytes_received
682-
data_mem = data_mem[bytes_received:]
683-
684-
# Check that the authentication bytes match.
685-
if len(data_buf) != auth_len:
686-
raise ConnectionError("recv() got too few auth bytes.")
687-
if bytes(data_buf) != auth_bytes:
688-
raise ConnectionError(f"Mismatched auth token.")
689-
690-
# Authenticating avoids using a connection from something else
691-
# able to connect to {host}:{port} instead of us.
641+
csock.setblocking(False)
692642
try:
693-
authenticate_socket_conn(ssock, csock, s_auth)
694-
authenticate_socket_conn(csock, ssock, c_auth)
695-
except OSError as exc:
696-
csock.close()
697-
ssock.close()
698-
reason = str(exc)
699-
continue
700-
# authentication successful, both sockets are our process.
701-
break
702-
finally:
703-
lsock.close()
704-
else:
705-
raise ConnectionError(f"socketpair authentication failed: {reason}")
643+
csock.connect((addr, port))
644+
except (BlockingIOError, InterruptedError):
645+
pass
646+
csock.setblocking(True)
647+
ssock, _ = lsock.accept()
648+
except:
649+
csock.close()
650+
raise
651+
finally:
652+
lsock.close()
653+
654+
# Authenticating avoids using a connection from something else
655+
# able to connect to {host}:{port} instead of us.
656+
if (
657+
ssock.getsockname()[:2] != csock.getpeername()[:2]
658+
or csock.getsockname()[:2] != ssock.getpeername()[:2]
659+
):
660+
raise ConnectionError("Unexpected peer connection")
706661
return (ssock, csock)
707662
__all__.append("socketpair")
708663

Lib/test/test_socket.py

+30-15
Original file line numberDiff line numberDiff line change
@@ -4921,21 +4921,36 @@ def _test_injected_authentication_failure(self):
49214921
pass
49224922

49234923
def test_injected_authentication_failure(self):
4924-
import secrets
4925-
orig_token_bytes = secrets.token_bytes
4926-
fake_n = secrets.DEFAULT_ENTROPY - 1
4927-
from unittest import mock
4928-
with mock.patch.object(
4929-
secrets, 'token_bytes',
4930-
new=lambda nbytes=None: orig_token_bytes(fake_n)):
4931-
s1, s2 = None, None
4932-
try:
4933-
with self.assertRaisesRegex(ConnectionError,
4934-
"authentication fail"):
4935-
s1, s2 = socket.socketpair()
4936-
finally:
4937-
if s1: s1.close()
4938-
if s2: s2.close()
4924+
orig_getsockname = socket.socket.getsockname
4925+
inject_sock = None
4926+
4927+
def inject_getsocketname(self):
4928+
nonlocal inject_sock
4929+
sockname = orig_getsockname(self)
4930+
# Connect to the listening socket ahead of the
4931+
# client socket.
4932+
if inject_sock is None:
4933+
inject_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
4934+
inject_sock.setblocking(False)
4935+
try:
4936+
inject_sock.connect(sockname[:2])
4937+
except (BlockingIOError, InterruptedError):
4938+
pass
4939+
inject_sock.setblocking(True)
4940+
return sockname
4941+
4942+
sock1 = sock2 = None
4943+
try:
4944+
socket.socket.getsockname = inject_getsocketname
4945+
with self.assertRaises(OSError):
4946+
sock1, sock2 = socket.socketpair()
4947+
finally:
4948+
socket.socket.getsockname = orig_getsockname
4949+
if inject_sock:
4950+
inject_sock.close()
4951+
if sock1: # This cleanup isn't needed on a successful test.
4952+
sock1.close()
4953+
sock2.close()
49394954

49404955

49414956
class NonBlockingTCPTests(ThreadedTCPSocketTest):

0 commit comments

Comments
 (0)