Простой способ передать большие файлы через socket на Python

Python предлагает очень удобную обертку над сокетами (socket), однако, как известно, сокеты являются достаточно низкоуровневыми и не гарантируют доставку сообщения целиком. То есть, отправив большой объем данных, вы их получите, но, скорее всего, по частям. В этой заметке будет показан простой способ сделать обертку над сокетами для получения примитивного протокола уровня пакетов и передавать таким образом большие файлы.

Идея

Фрагментация сообщений при передаче через сокеты приводит к тому, что невозможно заранее знать размер получаемого сообщения, а прием/передача с применением заведомо большего буфера также не гарантирует передачу за один вызов socket.send.

Простейшим решением этой проблемы является добавление в начало посылки ее размера в байтах. Таким образом, первым вызовом recvall можно получать только размер (4 байта), который практически гарантированно будет передан в одном TCP пакете (фрагментация, как правило, возникает при передаче больше, чем 1 кБ), после чего считать из буфера оставшееся сообщение, длинна которого уже известна.

Передача пакета

Для передачи большого сообщения (например, картинки) к нему, сформированным привычным образом (как для socket) добавим в начало длину в байтах (в сетевом порядке следования байт) и отправим целиком через вызов socket.send.

Реализуется это простой функцией:

def send_msg(self, sock, msg):
    # Каждое сообщение будет иметь префикс в 4 байта блинной(network byte order)
    msg = struct.pack('>I', len(msg)) + msg
    sock.send(msg)

Получение пакета

Получение пакета менее тривиально: первоначально напишем вспомогательную функцию для получения заданного числа байт:

def recvall(self, sock, n):
    # Функция для получения n байт или возврата None если получен EOF
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            return None
        data += packet
    return data

Так, пока не будет достигнут заданный размер, в byteArray добавляются принятые данные. По достижению заданного размера возвращается весь принятый пакет.

Теперь можно написать функцию для приема пакета целиком:

def recv_msg(self, sock):
    # Получение длины сообщения и распаковка в integer
    raw_msglen = self.recvall(sock, 4)
    if not raw_msglen:
        return None
    msglen = struct.unpack('>I', raw_msglen)[0]
    # Получение данных
    return self.recvall(sock, msglen)

Тут первоначально из сокета считывается размер сообщения (4 байта), распаковывается в целое число и, узнав таким образом размер сообщения, получается остаток данных с помощью вспомогательной функции

Простой класс

Упростим использование: чтобы не передавать сокет каждый раз явно, создадим класс.

Фактически он ничем не отличается от предложенных ранее функций, но объект сокета передается один раз при создании, а не каждый раз при вызове функций:

import socket
import struct

class SuperSocket():
    def __init__(self, sock):
        self._sock = sock

    def send_msg(self, msg):
        # Каждое сообщение будет иметь префикс в 4 байта блинной(network byte order)
        msg = struct.pack('>I', len(msg)) + msg
        self._sock.send(msg)

    def recv_msg(self):
        # Получение длины сообщения и распаковка в integer
        raw_msglen = self.recvall(4)
        if not raw_msglen:
            return None
        msglen = struct.unpack('>I', raw_msglen)[0]
        # Получение данных
        return self.recvall(msglen)

    def recvall(self, n):
        # Функция для получения n байт или возврата None если получен EOF
        data = b''
        while len(data) < n:
            packet = self._sock.recv(n - len(data))
            if not packet:
                return None
            data += packet
        return data