python - How can I reliably read exactly n bytes from a TCP socket?

Question

Context:

It is common that a binary protocol defines frames of a given size. The struct module is good at parsing that, provided everything has been received in a single buffer.

Problem:

TCP sockets are streams. A read from a socket cannot give more bytes than requested but can return less. So this code is not reliable:

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

The naive workaround:

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

may not be efficient, because if we ask a large number of bytes, and the data if very fragmented, we will repeatedly re-allocate a new byte buffer.

Question:

How is it possible to reliably receive exactly n bytes from a stream socket with no risk of re-allocation?

References:

Those other questions are related, and do give hints, but none give a simple and clear answer:

• 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

The solution is to use recv_into and a memoryview. Python allows to pre-allocate a modifiable bytearray that can be passed to recv_into. But you cannot receive data into a slice of the bytearray, because the slice would be a copy. But a memoryview allows to recieve multiple fragments into the same 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