← Back to branch summary

~azzar1/unity/add-show-desktop-key

« back to all changes in this revision

Viewing changes to ivle/chat.py

  • Committer: David Coles
  • Date: 2010-02-17 23:24:30 UTC
  • Revision ID: coles.david@gmail.com-20100217232430-wjk8ybvvori3q80e
https://launchpad.net/bugs/494872
Modify chat protocol to use Netstrings to track message boundaries rather than 
relying on the unsound "read until it would block" method. This should fix 
random corruptions of very long chat messages. Fixes lp:494872.

Show diffs side-by-side

added added

removed removed

Lines of Context:
ivle/chat.py
27
27
import socket
28
28
import traceback
29
29
 
 
30
SOCKETTIMEOUT = 60
 
31
BLOCKSIZE = 1024
 
32
 
30
33
class Terminate(Exception):
31
 
    """Exception thrown when server is to be shut down. It will attempt to sned 
32
 
    the final_response to the client and then exits"""
 
34
    """Exception thrown when server is to be shut down. It will attempt to
 
35
    send the final_response to the client and then exits"""
33
36
    def __init__(self, final_response=None):
34
37
        self.final_response = final_response
35
38
 
81
84
 
82
85
    while True:
83
86
        (conn, addr) = s.accept()
 
87
        conn.settimeout(SOCKETTIMEOUT)
84
88
        try:
85
89
            # Grab the input
86
 
            buf = cStringIO.StringIO()
87
 
            blk = conn.recv(1024)
88
 
            while blk:
89
 
                buf.write(blk)
90
 
                try:
91
 
                    blk = conn.recv(1024, socket.MSG_DONTWAIT)
92
 
                except:
93
 
                    # Exception thrown if it WOULD block (but we
94
 
                    # told it not to wait) - ie. we are done
95
 
                    blk = None
96
 
            inp = buf.getvalue()
 
90
            inp = recv_netstring(conn)
97
91
            env = cjson.decode(inp)
98
 
            
 
92
 
99
93
            # Check that the message is 
100
94
            digest = hashlib.md5(env['content'] + magic).hexdigest()
101
95
            if env['digest'] != digest:
106
100
 
107
101
            response = handler(content)
108
102
 
109
 
            conn.sendall(cjson.encode(response))
 
103
            send_netstring(conn, cjson.encode(response))
110
104
 
111
105
            conn.close()
112
106
 
113
107
        except Terminate, t:
114
108
            # Try and send final response and then terminate
115
109
            if t.final_response:
116
 
                conn.sendall(cjson.encode(t.final_response))
 
110
                send_netstring(conn, cjson.encode(t.final_response))
117
111
            conn.close()
118
112
            sys.exit(0)
119
113
        except Exception:
126
120
                "value": str(e_val),
127
121
                "traceback": tb_dump.getvalue()
128
122
            }
129
 
            conn.sendall(cjson.encode(json_exc))
 
123
            send_netstring(conn, cjson.encode(json_exc))
130
124
            conn.close()
131
125
 
132
126
 
133
127
def chat(host, port, msg, magic, decode = True):
134
128
    sok = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
135
129
    sok.connect((host, port))
 
130
    sok.settimeout(SOCKETTIMEOUT)
136
131
    content = cjson.encode(msg)
137
132
    digest = hashlib.md5(content + magic).hexdigest()
138
133
    env = {'digest':digest,'content':content}
139
 
    sok.send(cjson.encode(env))
140
 
 
141
 
    buf = cStringIO.StringIO()
142
 
    blk = sok.recv(1024)
143
 
    while blk:
144
 
        buf.write(blk)
145
 
        try:
146
 
            blk = sok.recv(1024, socket.MSG_DONTWAIT)
147
 
        except socket.error, e:
148
 
            if e[0] != 11:
149
 
                raise
150
 
            # Exception thrown if it WOULD block (but we
151
 
            # told it not to wait) - ie. we are done
152
 
            blk = None
153
 
    inp = buf.getvalue()
 
134
    json = cjson.encode(env)
 
135
 
 
136
    send_netstring(sok, json)
 
137
    inp = recv_netstring(sok)
 
138
 
154
139
    sok.close()
155
140
 
156
141
    if decode:
158
143
    else:
159
144
        return inp
160
145
 
 
146
 
 
147
def send_netstring(sok, data):
 
148
    netstring = "%d:%s,"%(len(data),data)
 
149
    sok.sendall(netstring)
 
150
 
 
151
 
 
152
def recv_netstring(sok):
 
153
    # Decode netstring
 
154
    size_buffer = []
 
155
    c = sok.recv(1)
 
156
    while c != ':':
 
157
        # Limit the Netstring to less than 10^10 bytes (~1GB).
 
158
        if len(size_buffer) >= 10:
 
159
            raise ValueError("Not valid Netstring: More than 10^10 bytes")
 
160
        size_buffer.append(c)
 
161
        c = sok.recv(1)
 
162
    try:
 
163
        # Message should be length plus ',' terminator
 
164
        recv_expected = int(''.join(size_buffer)) + 1
 
165
    except ValueError, e:
 
166
        raise ValueError("Not valid Netstring: '%s'"%blk)
 
167
 
 
168
    # Read data
 
169
    buf = []
 
170
    recv_data = sok.recv(min(recv_expected, BLOCKSIZE))
 
171
    recv_size = len(recv_data)
 
172
    while recv_size < recv_expected:
 
173
        buf.append(recv_data)
 
174
        recv_data = sok.recv(min(recv_expected-recv_size, 1024))
 
175
        recv_size = recv_size + len(recv_data)
 
176
    assert(recv_size == recv_expected)
 
177
 
 
178
    # Did we receive the correct amount?
 
179
    if recv_data[-1] != ',':
 
180
        raise ValueError("Not valid Netstring: Did not end with ','")
 
181
    buf.append(recv_data[:-1])
 
182
 
 
183
    return ''.join(buf)