如何可靠地从 TCP 套接字准确读取 n 个字节？

Question

上下文：

二进制协议定义给定大小的帧是很常见的。 struct 模块擅长解析它，前提是所有内容都已在单个缓冲区中接收。

问题：

TCP 套接字是流。从套接字读取不能提供比请求更多的字节，但可以返回更少。所以这段代码不可靠：

def readnbytes(sock, n):
    return sock.recv(n)   # can return less than n bytes

天真的解决方法：

def readnbytes(sock, n):
    buff = b''
    while n > 0:
        b = sock.recv(n)
        buff += b
        if len(b) == 0:
            raise EOFError          # peer socket has received a SH_WR shutdown
        n -= len(b)
    return buff

可能效率不高，因为如果我们要求大量的字节，并且数据如果非常碎片化，我们会反复重新分配一个新的字节缓冲区。

问题：

如何才能可靠地从流套接字准确接收 n 个字节而没有重新分配的风险？

参考资料：

那些其他问题是相关的，确实给出了提示，但没有一个给出简单明了的答案：

• How to receive certain size of data in socket programming?
• How do I use ctypes to recv_into a C buffer multiple times?

Answer 1

@Serge 的答案的一个小补充，它返回一个IncompleteReadError（它是EOFError 的一个子类）。这包含一个包含部分读取数据的partial 属性。

import socket
from asyncio import IncompleteReadError
 
def readexactly(sock: socket.socket, num_bytes: int) -> bytes:
    buf = bytearray(num_bytes)
    pos = 0
    while pos < num_bytes:
        n = sock.recv_into(memoryview(buf)[pos:])
        if n == 0:
            raise IncompleteReadError(bytes(buf[:pos]), num_bytes)
        pos += n
    return bytes(buf)

用法：

try:
    print(readexactly(sock, 26))
except IncompleteReadError as e:
    print(f"Only read {len(e.partial)} out of {e.expected} bytes. :(")
    print(e.partial)

仅读取 5 个字节时的示例输出b"ABCDE"：

Only read 5 out of 26 bytes. :(
b'ABCDE'

Answer 2

解决方案是使用recv_into 和memoryview。 Python 允许预先分配一个可修改的bytearray，它可以传递给recv_into。但是您不能将数据接收到字节数组的切片中，因为切片将是一个副本。但是memoryview 允许将多个片段接收到同一个bytearray：

def readnbyte(sock, n):
    buff = bytearray(n)
    pos = 0
    while pos < n:
        cr = sock.recv_into(memoryview(buff)[pos:])
        if cr == 0:
            raise EOFError
        pos += cr
    return buff

Answer 3

您可以使用socket.makefile() 将套接字包装在类似文件的对象中。然后读取将准确返回请求的数量，除非套接字关闭，它可以返回剩余部分。这是一个例子：

server.py

from socket import *

sock = socket()
sock.bind(('',5000))
sock.listen(1)
with sock:
    client,addr = sock.accept()
    with client, client.makefile() as clientfile:
        while True:
            data = clientfile.read(5)
            if not data: break
            print(data)

client.py

from socket import *
import time

sock = socket()
sock.connect(('localhost',5000))
with sock:
    sock.sendall(b'123')
    time.sleep(.5)
    sock.sendall(b'451234')
    time.sleep(.5)
    sock.sendall(b'51234')

服务器输出

12345
12345
1234