Protocol implementation, TCP, sending images through sockets - follow-up

Question

I would like to ask about your opinion about my code. The idea is simple: I designed my own protocol, where client asks server about the image, and the server sends the image, following the below steps (this is the actual protocol I wanted to implement):

CLIENT                                      SERVER
                      GET\r\n   
       ----------------------------------->
                      OK\r\n
       <----------------------------------- 
                   GET_SIZE\r\n
       ----------------------------------->
                    SIZE 1024\r\n
       <-----------------------------------
                   GET_IMG\r\n
       ----------------------------------->
    IMG_DATA an when image sending is over, EOIMG\r\n 
       <-----------------------------------
                   DONE\r\n
       ----------------------------------->
            server disconnects the client   

I improved the code as suggested by @TobErnack, but would still like to know your opinion of my improvements. Here's the link to the original question and code:

Network protocol using TCP, sending images through sockets

Everything works (it seems so) but I would like to ask about what else I could possibly improve here.

1. Does the implementation reflects in 100% what I wanted to achieve in my protocol flow?
2. Are the steps implemented as I wanted them to be in my protocol?

client.py

#!/usr/bin/env python

import socket
import sys

HOST = '127.0.0.1'
PORT = 6666

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_address = (HOST, PORT)
sock.connect(server_address)
fname = 'fromserver.png'


def recvall(sock, msgLen):
    msg = ""
    bytesRcvd = 0

    while bytesRcvd < msgLen:

        chunk = sock.recv(msgLen - bytesRcvd)

        if chunk == "": break

        bytesRcvd += len(chunk)
        msg += chunk

        if "\r\n" in msg: break
    return msg


try:

    sock.sendall("GET\r\n")
    data = recvall(sock, 4096)

    if data:
        txt = data.strip()
        print '--%s--' % txt

        if txt == 'OK':

            sock.sendall("GET_SIZE\r\n")
            data = recvall(sock, 4096)

            if data:
                txt = data.strip()
                print '--%s--' % txt

                if txt.startswith('SIZE'):

                    tmp = txt.split()
                    size = int(tmp[1])

                    print '--%s--' % size

                    sock.sendall("GET_IMG\r\n")

                    myfile = open(fname, 'wb')

                    amount_received = 0
                    while amount_received < size:
                        data = recvall(sock, 4096)
                        if not data:
                            break
                        amount_received += len(data)
                        print amount_received

                        txt = data.strip('\r\n')

                        if 'EOIMG' in str(txt):
                            print 'Image received successfully'
                            myfile.write(data)
                            myfile.close()
                            sock.sendall("DONE\r\n")
                        else:
                            myfile.write(data)
finally:
    sock.close()

server.py

# !/usr/bin/env python

import random
import socket, select
from time import gmtime, strftime

image = 'tux.png'

HOST = '127.0.0.1'
PORT = 6666

connected_clients_sockets = []

server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server_socket.bind((HOST, PORT))
server_socket.listen(10)

connected_clients_sockets.append(server_socket)


def recvall(sock, msgLen):
    msg = ""
    bytesRcvd = 0

    while bytesRcvd < msgLen:

        chunk = sock.recv(msgLen - bytesRcvd)

        if chunk == "": break

        bytesRcvd += len(chunk)
        msg += chunk

        if "\r\n" in msg: break
    return msg


while True:

    read_sockets, write_sockets, error_sockets = select.select(connected_clients_sockets, [], [])

    for sock in read_sockets:

        if sock == server_socket:

            sockfd, client_address = server_socket.accept()
            connected_clients_sockets.append(sockfd)

        else:
            try:
                data = recvall(sock, 4096)

                if data:

                    txt = data.strip()
                    print '--%s--' % txt

                    if txt == 'GET':
                        sock.sendall('OK\r\n')

                    elif txt == 'GET_SIZE':

                        with open('tux.png', 'rb') as f1:
                            file_size = len(f1.read())
                            f1.seek(0)

                        print '--%s--' % file_size

                        file_size = '%s' % file_size
                        sock.sendall('SIZE %s\r\n' % file_size)

                    elif txt == 'GET_IMG':
                        with open(image, 'rb') as fp:
                            image_data = fp.read()

                        msg = '%sEOIMG\r\n' % image_data
                        sock.sendall(msg)

                    elif txt == 'DONE':
                        sock.close()
                        connected_clients_sockets.remove(sock)

            except:
                sock.close()
                connected_clients_sockets.remove(sock)
                continue
server_socket.close()

Answer 1

You could simplify a lot the server code using the SocketServer module. You can declare your server using:

import SocketServer


class ImageServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    allow_reuse_address = True

You will then be able to wait for connections using:

server = ImageServer((HOST, PORT), <SomeHandler>)
server.serve_forever()

Doing so, each time a new client connect to the server, the accept is done by the server and the handle method of <SomeHandler> is called using the resulting socket.

Moreover, using the TheadingMixIn, each client will be processed in its own thread, so you can use blocking calls and sequencial flow instead of your kind of switch.

The handler can thus be defined allong the way of:

class ImageHandler(SocketServer.BaseRequestHandler):
    def read_message(self, expected_message):
        data = recvall(self.request, 4096)
        txt = data.strip()
        if txt != expected_message:
            raise ValueError

    def handle(self):
        self.read_message('GET')
        sock.sendall('OK\r\n')

        self.read_message('GET_SIZE')
        with open('tux.png', 'rb') as f1:
            file_size = len(f1.read())
        sock.sendall('SIZE %s\r\n' % file_size)

        self.get_message('GET_IMG')
        with open(image, 'rb') as fp:
            image_data = fp.read()
        sock.sendall('%sEOIMG\r\n' % image_data)

        self.get_message('DONE')

    def finish(self):
        self.request.close()

What is interesting here is that, whether or not handle is raising an exception, finish will always be called when the flow exits handle, ensuring all connections gets closed anyway.

I advise you, however, to use your own exception instead of a generic ValueError.

---

You are also using recvall in both files, so you could use a third one that just define this function (say common.py) that you import in both:

from common import recvall

---

New server.py could look like:

#!/usr/bin/env python2

import SocketServer

from common import recvall


IMAGE = 'tux.png'


class ImageServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
    allow_reuse_address = True


class ImageHandler(SocketServer.BaseRequestHandler):
    def read_message(self, expected_message):
        data = recvall(self.request, 4096)
        txt = data.strip()
        if txt != expected_message:
            raise ValueError

    def handle(self):
        self.read_message('GET')
        sock.sendall('OK\r\n')

        self.read_message('GET_SIZE')
        with open(IMAGE, 'rb') as f1:
            file_size = len(f1.read())
        sock.sendall('SIZE %s\r\n' % file_size)

        self.get_message('GET_IMG')
        with open(IMAGE, 'rb') as fp:
            image_data = fp.read()
        sock.sendall('%sEOIMG\r\n' % image_data)

        self.get_message('DONE')

    def finish(self):
        self.request.close()


if __name__ == '__main__':
    server = ImageServer(('', 6666), ImageHandler)
    server.serve_forever()

Answer 2

Please avoid polluting the global namespace with so many symbols. Just def main(): and put your top-level code in there.

This risks turning performance from O(n) linear into O(n^2) quadratic:

    msg += chunk

Some cPython implementation details may mask this, but the usual idiom is to build up a list of chunks and return ''.join(chunks).

Terminating upon CRLF, as well as awaiting msgLen bytes, is odd. Your test definitely suffers from quadratic cost. Scan just the chunk, rather than the whole accumulated msg, and deal with the corner case of msg ends with CR, chunk starts with LF.

                myfile = open(fname, 'wb')

It is usual to write this as:

                with open(fname, 'wb') as myfile:

Oh, wait. On re-reading it I see that msgLen is misnamed, you really meant max_msg_len, and CRLF is the usual way we terminate. So despite the sock.recv() binary approach, recvall() is actually line-oriented, and should be named recv_line() or read_line or get_line. It offers a hard-to-use public API to callers, since caller is responsible for locating first CRLF, processing what comes before, and retaining what comes after. A better API would pass back that pair with return msg, remainder. Relying on timing details is not robust, given that events like TCP retransmits may influence who wins racy access to the data.

Later you treat the API as a binary interface:

                while amount_received < size:
                    data = recvall(sock, 4096)
                    ...
                    txt = data.strip('\r\n')

But you apparently don't correctly handle binary JPEG images that happen to have a CRLF byte sequence within them. And you are incorrectly viewing "a buffer returned by recv" as "one record from a record stream". No, TCP does not offer record markers, it offers a byte stream, only. A sending TCP is free to segment the stream arbitrarily, just as recv could coalesce or break apart segments. What goes on the wire is not one-to-one with the chapter 2 API. If you want record markers you need to add them at the app layer, as you attempted to do here by obtaining length L then doing binary read of L bytes. To fix the bug, you would need to define binary_recvall() which does not pay attention to CRLFs in the binary bytestream.