33
30
class Terminate(Exception):
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"""
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"""
36
33
def __init__(self, final_response=None):
37
34
self.final_response = final_response
40
37
return repr(self.final_response)
42
class ProtocolError(Exception):
43
"""Exception thrown when client violates the the chat protocol"""
46
40
def start_server(port, magic, daemon_mode, handler, initializer = None):
47
41
# Attempt to open the socket.
59
53
os._exit(0) # kill off parent again.
63
MAXFD = os.sysconf("SC_OPEN_MAX")
67
# Close all file descriptors, except the socket.
68
for i in xrange(MAXFD):
76
si = os.open(os.devnull, os.O_RDONLY)
77
os.dup2(si, sys.stdin.fileno())
79
so = os.open(os.devnull, os.O_WRONLY)
80
os.dup2(so, sys.stdout.fileno())
82
se = os.open(os.devnull, os.O_WRONLY)
83
os.dup2(se, sys.stderr.fileno())
89
60
(conn, addr) = s.accept()
90
conn.settimeout(SOCKETTIMEOUT)
92
# Grab the input and try to decode
93
inp = recv_netstring(conn)
95
content = decode(inp, magic)
63
buf = cStringIO.StringIO()
68
blk = conn.recv(1024, socket.MSG_DONTWAIT)
70
# Exception thrown if it WOULD block (but we
71
# told it not to wait) - ie. we are done
74
env = cjson.decode(inp)
76
# Check that the message is
77
digest = md5.new(env['content'] + magic).digest().encode('hex')
78
if env['digest'] != digest:
82
content = cjson.decode(env['content'])
100
84
response = handler(content)
102
send_netstring(conn, cjson.encode(response))
86
conn.sendall(cjson.encode(response))
106
90
except Terminate, t:
107
91
# Try and send final response and then terminate
108
92
if t.final_response:
109
send_netstring(conn, cjson.encode(t.final_response))
93
conn.sendall(cjson.encode(t.final_response))
119
103
"value": str(e_val),
120
104
"traceback": tb_dump.getvalue()
122
send_netstring(conn, cjson.encode(json_exc))
106
conn.sendall(cjson.encode(json_exc))
126
110
def chat(host, port, msg, magic, decode = True):
127
111
sok = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
128
112
sok.connect((host, port))
129
sok.settimeout(SOCKETTIMEOUT)
131
json = encode(msg, magic)
133
send_netstring(sok, json)
134
inp = recv_netstring(sok)
113
content = cjson.encode(msg)
114
digest = md5.new(content + magic).digest().encode("hex")
115
env = {'digest':digest,'content':content}
116
sok.send(cjson.encode(env))
118
buf = cStringIO.StringIO()
123
blk = sok.recv(1024, socket.MSG_DONTWAIT)
124
except socket.error, e:
127
# Exception thrown if it WOULD block (but we
128
# told it not to wait) - ie. we are done
143
def encode(message, magic):
144
"""Converts a message into a JSON serialisation and uses a magic
145
string to attach a HMAC digest.
147
# XXX: Any reason that we double encode?
148
content = cjson.encode(message)
150
digest = hashlib.md5(content + magic).hexdigest()
151
env = {'digest':digest,'content':content}
152
json = cjson.encode(env)
157
def decode(message, magic):
158
"""Takes a message with an attached HMAC digest and validates the message.
160
msg = cjson.decode(message)
162
# Check that the message is valid
163
digest = hashlib.md5(msg['content'] + magic).hexdigest()
164
if msg['digest'] != digest:
165
raise ProtocolError("HMAC digest is invalid")
166
content = cjson.decode(msg['content'])
171
def send_netstring(sok, data):
172
""" Sends a netstring to a socket
174
netstring = "%d:%s,"%(len(data),data)
175
sok.sendall(netstring)
178
def recv_netstring(sok):
179
""" Attempts to recieve a Netstring from a socket.
180
Throws a ProtocolError if the received data violates the Netstring
187
# Limit the Netstring to less than 10^10 bytes (~1GB).
188
if len(size_buffer) >= 10:
190
"Could not read Netstring size in first 9 bytes: '%s'"%(
191
''.join(size_buffer)))
192
size_buffer.append(c)
195
# Message should be length plus ',' terminator
196
recv_expected = int(''.join(size_buffer)) + 1
197
except ValueError, e:
198
raise ProtocolError("Could not decode Netstring size as int: '%s'"%(
199
''.join(size_buffer)))
203
recv_data = sok.recv(min(recv_expected, BLOCKSIZE))
204
recv_size = len(recv_data)
205
while recv_size < recv_expected:
206
buf.append(recv_data)
207
recv_data = sok.recv(min(recv_expected-recv_size, 1024))
208
recv_size = recv_size + len(recv_data)
209
assert(recv_size == recv_expected)
211
# Did we receive the correct amount?
212
if recv_data[-1] != ',':
213
raise ProtocolError("Netstring did not end with ','")
214
buf.append(recv_data[:-1])