如何可靠地从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 个字节而无需重新分配的风险？

参考资料：

这些相关问题提供了一些提示，但没有提供简单明确的答案：

• 如何在套接字编程中接收特定大小的数据？
• 如何使用 ctypes 多次 recv_into C 缓冲区？

- Serge Ballesta

---

请参考以下答案：https://dev59.com/0mQo5IYBdhLWcg3wiv7A#15964489 - Jasha

Answer 1

解决方案是使用recv_into和一个memoryview。Python允许预先分配一个可修改的bytearray，并将其传递给recv_into。但是，您不能接收数据到bytearray的切片中，因为切片是一个副本。但是，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

- Serge Ballesta

Answer 2

你可以使用 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

- Mark Tolonen

---

clientfile.read(5) 读取确切的5个字节的文档在哪里可以找到？在 socket 文档中，它说 makefile 将返回一个文件对象。在 Python 词汇表中，它说 "file object" 是 "公开面向文件的 API 的对象（具有例如 read() 或 write() 的方法）"。它还说 "它们的接口在 io 模块中定义。" 在标准库 io 模块中，它说 "IOBase 不声明 read() 或 write()，因为它们的签名会有所不同。" 在 io 模块中描述的几个 IOBase 子类说 read(size=-1) 将读取 "最多 size 字节".... - Jasha

read 方法的说明文档中写道：“如果操作系统调用返回的字节数少于 size，则可能返回少于 size 字节。” 对于类文件对象，通常只有在 EOF（或套接字关闭/关闭）发生时才会出现这种情况，但如果担心这种情况，请检查返回值。 - Mark Tolonen

Answer 3

@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'

- Mateen Ulhaq

---

奖励：带缓冲的行读取器：https://dev59.com/-V4b5IYBdhLWcg3whB6V#65637828 - Mateen Ulhaq