2 """The KRPC communication protocol implementation.
4 @var UDP_PACKET_LIMIT: the maximum number of bytes that can be sent in a
5 UDP packet without fragmentation
7 @var KRPC_ERROR: the code for a generic error
8 @var KRPC_ERROR_SERVER_ERROR: the code for a server error
9 @var KRPC_ERROR_MALFORMED_PACKET: the code for a malformed packet error
10 @var KRPC_ERROR_METHOD_UNKNOWN: the code for a method unknown error
11 @var KRPC_ERROR_MALFORMED_REQUEST: the code for a malformed request error
12 @var KRPC_ERROR_INVALID_TOKEN: the code for an invalid token error
13 @var KRPC_ERROR_RESPONSE_TOO_LONG: the code for a response too long error
15 @var KRPC_ERROR_INTERNAL: the code for an internal error
16 @var KRPC_ERROR_RECEIVED_UNKNOWN: the code for an unknown message type error
17 @var KRPC_ERROR_TIMEOUT: the code for a timeout error
18 @var KRPC_ERROR_PROTOCOL_STOPPED: the code for a stopped protocol error
20 @var TID: the identifier for the transaction ID
21 @var REQ: the identifier for a request packet
22 @var RSP: the identifier for a response packet
23 @var TYP: the identifier for the type of packet
24 @var ARG: the identifier for the argument to the request
25 @var ERR: the identifier for an error packet
27 @group Remote node error codes: KRPC_ERROR, KRPC_ERROR_SERVER_ERROR,
28 KRPC_ERROR_MALFORMED_PACKET, KRPC_ERROR_METHOD_UNKNOWN,
29 KRPC_ERROR_MALFORMED_REQUEST, KRPC_ERROR_INVALID_TOKEN,
30 KRPC_ERROR_RESPONSE_TOO_LONG
31 @group Local node error codes: KRPC_ERROR_INTERNAL, KRPC_ERROR_RECEIVED_UNKNOWN,
32 KRPC_ERROR_TIMEOUT, KRPC_ERROR_PROTOCOL_STOPPED
33 @group Command identifiers: TID, REQ, RSP, TYP, ARG, ERR
37 from bencode import bencode, bdecode
38 from datetime import datetime, timedelta
41 from twisted.internet import defer
42 from twisted.internet import protocol, reactor
43 from twisted.python import log
44 from twisted.trial import unittest
46 from khash import newID
48 UDP_PACKET_LIMIT = 1472
52 KRPC_ERROR_SERVER_ERROR = 201
53 KRPC_ERROR_MALFORMED_PACKET = 202
54 KRPC_ERROR_METHOD_UNKNOWN = 203
55 KRPC_ERROR_MALFORMED_REQUEST = 204
56 KRPC_ERROR_INVALID_TOKEN = 205
57 KRPC_ERROR_RESPONSE_TOO_LONG = 206
60 KRPC_ERROR_INTERNAL = 100
61 KRPC_ERROR_RECEIVED_UNKNOWN = 101
62 KRPC_ERROR_TIMEOUT = 102
63 KRPC_ERROR_PROTOCOL_STOPPED = 103
73 class KrpcError(Exception):
74 """An error occurred in the KRPC protocol."""
77 def verifyMessage(msg):
78 """Check received message for corruption and errors.
80 @type msg: C{dictionary}
81 @param msg: the dictionary of information received on the connection
82 @raise KrpcError: if the message is corrupt
86 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "not a dictionary")
88 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "no message type")
91 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "request type not specified")
92 if type(msg[REQ]) != str:
93 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "request type is not a string")
95 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "no arguments for request")
96 if type(msg[ARG]) != dict:
97 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "arguments for request are not in a dictionary")
100 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "response not specified")
101 if type(msg[RSP]) != dict:
102 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "response is not a dictionary")
103 elif msg[TYP] == ERR:
105 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "error not specified")
106 if type(msg[ERR]) != list:
107 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "error is not a list")
108 if len(msg[ERR]) != 2:
109 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "error is not a 2-element list")
110 if type(msg[ERR][0]) not in (int, long):
111 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "error number is not a number")
112 if type(msg[ERR][1]) != str:
113 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "error string is not a string")
115 # raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "unknown message type")
117 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "no transaction ID specified")
118 if type(msg[TID]) != str:
119 raise KrpcError, (KRPC_ERROR_MALFORMED_PACKET, "transaction id is not a string")
121 class hostbroker(protocol.DatagramProtocol):
122 """The factory for the KRPC protocol.
124 @type server: L{khashmir.Khashmir}
125 @ivar server: the main Khashmir program
126 @type stats: L{stats.StatsLogger}
127 @ivar stats: the statistics logger to save transport info
128 @type config: C{dictionary}
129 @ivar config: the configuration parameters for the DHT
130 @type connections: C{dictionary}
131 @ivar connections: all the connections that have ever been made to the
132 protocol, keys are IP address and port pairs, values are L{KRPC}
133 protocols for the addresses
134 @ivar protocol: the protocol to use to handle incoming connections
136 @type addr: (C{string}, C{int})
137 @ivar addr: the IP address and port of this node
140 def __init__(self, server, stats, config):
141 """Initialize the factory.
143 @type server: L{khashmir.Khashmir}
144 @param server: the main DHT program
145 @type stats: L{stats.StatsLogger}
146 @param stats: the statistics logger to save transport info
147 @type config: C{dictionary}
148 @param config: the configuration parameters for the DHT
153 # this should be changed to storage that drops old entries
154 self.connections = {}
156 def datagramReceived(self, datagram, addr):
157 """Optionally create a new protocol object, and handle the new datagram.
159 @type datagram: C{string}
160 @param datagram: the data received from the transport.
161 @type addr: (C{string}, C{int})
162 @param addr: source IP address and port of datagram.
164 c = self.connectionForAddr(addr)
165 c.datagramReceived(datagram, addr)
167 # del self.connections[addr]
169 def connectionForAddr(self, addr):
170 """Get a protocol object for the source.
172 @type addr: (C{string}, C{int})
173 @param addr: source IP address and port of datagram.
175 # Don't connect to ourself
176 if addr == self.addr:
179 # Create a new protocol object if necessary
180 if not self.connections.has_key(addr):
181 conn = self.protocol(addr, self.server, self.stats, self.transport, self.config)
182 self.connections[addr] = conn
184 conn = self.connections[addr]
187 def makeConnection(self, transport):
188 """Make a connection to a transport and save our address."""
189 protocol.DatagramProtocol.makeConnection(self, transport)
190 tup = transport.getHost()
191 self.addr = (tup.host, tup.port)
193 def stopProtocol(self):
194 """Stop all the open connections."""
195 for conn in self.connections.values():
197 protocol.DatagramProtocol.stopProtocol(self)
199 class KrpcRequest(defer.Deferred):
200 """An outstanding request to a remote node.
202 @type protocol: L{KRPC}
203 @ivar protocol: the protocol to send data with
204 @ivar tid: the transaction ID of the request
205 @type method: C{string}
206 @ivar method: the name of the method to call on the remote node
207 @type data: C{string}
208 @ivar data: the message to send to the remote node
209 @type config: C{dictionary}
210 @ivar config: the configuration parameters for the DHT
212 @ivar delay: the last timeout delay sent
213 @type start: C{datetime}
214 @ivar start: the time to request was started at
215 @type later: L{twisted.internet.interfaces.IDelayedCall}
216 @ivar later: the pending call to timeout the last sent request
219 def __init__(self, protocol, newTID, method, data, config):
220 """Initialize the request, and send it out.
222 @type protocol: L{KRPC}
223 @param protocol: the protocol to send data with
224 @param newTID: the transaction ID of the request
225 @type method: C{string}
226 @param method: the name of the method to call on the remote node
227 @type data: C{string}
228 @param data: the message to send to the remote node
229 @type config: C{dictionary}
230 @param config: the configuration parameters for the DHT
232 defer.Deferred.__init__(self)
233 self.protocol = protocol
238 self.delay = self.config.get('KRPC_INITIAL_DELAY', 2)
239 self.start = datetime.now()
241 reactor.callLater(0, self.send)
244 """Send the request to the remote node."""
245 assert not self.later, 'There is already a pending request'
246 self.later = reactor.callLater(self.delay, self.timeOut)
248 self.protocol.sendData(self.method, self.data)
253 """Check for a unrecoverable timeout, otherwise resend."""
255 delay = datetime.now() - self.start
256 if delay > timedelta(seconds = self.config.get('KRPC_TIMEOUT', 9)):
257 log.msg('%r timed out after %0.2f sec' %
258 (self.tid, delay.seconds + delay.microseconds/1000000.0))
259 self.protocol.timeOut(self.tid, self.method)
260 elif self.protocol.stopped:
261 log.msg('Timeout but can not resend %r, protocol has been stopped' % self.tid)
264 log.msg('Trying to resend %r now with delay %d sec' % (self.tid, self.delay))
265 reactor.callLater(0, self.send)
267 def callback(self, resp):
269 defer.Deferred.callback(self, resp)
271 def errback(self, resp):
273 defer.Deferred.errback(self, resp)
275 def dropTimeOut(self):
276 """Cancel the timeout call when a response is received."""
277 if self.later and self.later.active():
282 """The KRPC protocol implementation.
284 @ivar transport: the transport to use for the protocol
285 @type factory: L{khashmir.Khashmir}
286 @ivar factory: the main Khashmir program
287 @type stats: L{stats.StatsLogger}
288 @ivar stats: the statistics logger to save transport info
289 @type addr: (C{string}, C{int})
290 @ivar addr: the IP address and port of the source node
291 @type config: C{dictionary}
292 @ivar config: the configuration parameters for the DHT
293 @type tids: C{dictionary}
294 @ivar tids: the transaction IDs outstanding for requests, keys are the
295 transaction ID of the request, values are the deferreds to call with
297 @type stopped: C{boolean}
298 @ivar stopped: whether the protocol has been stopped
301 def __init__(self, addr, server, stats, transport, config = {}):
302 """Initialize the protocol.
304 @type addr: (C{string}, C{int})
305 @param addr: the IP address and port of the source node
306 @type server: L{khashmir.Khashmir}
307 @param server: the main Khashmir program
308 @type stats: L{stats.StatsLogger}
309 @param stats: the statistics logger to save transport info
310 @param transport: the transport to use for the protocol
311 @type config: C{dictionary}
312 @param config: the configuration parameters for the DHT
313 (optional, defaults to using defaults)
315 self.transport = transport
316 self.factory = server
323 def datagramReceived(self, data, addr):
324 """Process the new datagram.
326 @type data: C{string}
327 @param data: the data received from the transport.
328 @type addr: (C{string}, C{int})
329 @param addr: source IP address and port of datagram.
331 self.stats.receivedBytes(len(data))
333 if self.config.get('SPEW', False):
334 log.msg("stopped, dropping message from %r: %s" % (addr, data))
336 # Bdecode the message
340 if self.config.get('SPEW', False):
341 log.msg("krpc bdecode error: ")
345 # Make sure the remote node isn't trying anything funny
349 log.msg("krpc message verification error: ")
353 if self.config.get('SPEW', False):
354 log.msg("%d received from %r: %s" % (self.factory.port, addr, msg))
356 # Process it based on its type
360 # Requests are handled by the factory
361 f = getattr(self.factory ,"krpc_" + msg[REQ], None)
362 msg[ARG]['_krpc_sender'] = self.addr
363 if f and callable(f):
364 self.stats.receivedAction(msg[REQ])
366 ret = f(*(), **msg[ARG])
368 log.msg('Got a Krpc error while running: krpc_%s' % msg[REQ])
369 if e[0] != KRPC_ERROR_INVALID_TOKEN:
371 self.stats.errorAction(msg[REQ])
372 olen = self._sendResponse(msg[REQ], addr, msg[TID], ERR,
375 log.msg('Got a malformed request for: krpc_%s' % msg[REQ])
377 self.stats.errorAction(msg[REQ])
378 olen = self._sendResponse(msg[REQ], addr, msg[TID], ERR,
379 [KRPC_ERROR_MALFORMED_REQUEST, str(e)])
381 log.msg('Got an unknown error while running: krpc_%s' % msg[REQ])
383 self.stats.errorAction(msg[REQ])
384 olen = self._sendResponse(msg[REQ], addr, msg[TID], ERR,
385 [KRPC_ERROR_SERVER_ERROR, str(e)])
387 olen = self._sendResponse(msg[REQ], addr, msg[TID], RSP, ret)
389 # Request for unknown method
390 log.msg("ERROR: don't know about method %s" % msg[REQ])
391 self.stats.receivedAction('unknown')
392 olen = self._sendResponse(msg[REQ], addr, msg[TID], ERR,
393 [KRPC_ERROR_METHOD_UNKNOWN, "unknown method "+str(msg[REQ])])
395 log.msg("%s >>> %s %s %s" % (addr, ilen, msg[REQ], olen))
396 elif msg[TYP] == RSP:
397 # Responses get processed by their TID's deferred
398 if self.tids.has_key(msg[TID]):
399 req = self.tids[msg[TID]]
401 del(self.tids[msg[TID]])
402 msg[RSP]['_krpc_sender'] = addr
403 req.callback(msg[RSP])
405 # no tid, this transaction was finished already...
406 if self.config.get('SPEW', False):
407 log.msg('received response from %r for completed request: %r' %
408 (msg[RSP]['id'], msg[TID]))
409 elif msg[TYP] == ERR:
410 # Errors get processed by their TID's deferred's errback
411 if self.tids.has_key(msg[TID]):
412 req = self.tids[msg[TID]]
413 del(self.tids[msg[TID]])
415 req.errback(KrpcError(*msg[ERR]))
417 # no tid, this transaction was finished already...
418 log.msg('received an error %r from %r for completed request: %r' %
419 (msg[ERR], msg[RSP]['id'], msg[TID]))
421 # Received an unknown message type
422 if self.config.get('SPEW', False):
423 log.msg("unknown message type: %r" % msg)
424 if msg[TID] in self.tids:
425 req = self.tids[msg[TID]]
426 del(self.tids[msg[TID]])
428 req.errback(KrpcError(KRPC_ERROR_RECEIVED_UNKNOWN,
429 "Received an unknown message type: %r" % msg[TYP]))
431 def _sendResponse(self, request, addr, tid, msgType, response):
432 """Helper function for sending responses to nodes.
434 @param request: the name of the requested method
435 @type addr: (C{string}, C{int})
436 @param addr: source IP address and port of datagram.
437 @param tid: the transaction ID of the request
438 @param msgType: the type of message to respond with
439 @param response: the arguments for the response
445 # Create the response message
446 msg = {TID : tid, TYP : msgType, msgType : response}
448 if self.config.get('SPEW', False):
449 log.msg("%d responding to %r: %s" % (self.factory.port, addr, msg))
453 # Make sure its not too long
454 if len(out) > UDP_PACKET_LIMIT:
455 # Can we remove some values to shorten it?
456 if 'values' in response:
457 # Save the original list of values
458 orig_values = response['values']
459 len_orig_values = len(bencode(orig_values))
461 # Caclulate the maximum value length possible
462 max_len_values = len_orig_values - (len(out) - UDP_PACKET_LIMIT)
463 assert max_len_values > 0
465 # Start with a calculation of how many values should be included
466 # (assumes all values are the same length)
467 per_value = (float(len_orig_values) - 2.0) / float(len(orig_values))
468 num_values = len(orig_values) - int(ceil(float(len(out) - UDP_PACKET_LIMIT) / per_value))
470 # Do a linear search for the actual maximum number possible
471 bencoded_values = len(bencode(orig_values[:num_values]))
472 while bencoded_values < max_len_values and num_values + 1 < len(orig_values):
473 bencoded_values += len(bencode(orig_values[num_values]))
475 while bencoded_values > max_len_values and num_values > 0:
477 bencoded_values -= len(bencode(orig_values[num_values]))
478 assert num_values > 0
481 response['values'] = orig_values[:num_values]
483 assert len(out) < UDP_PACKET_LIMIT
484 log.msg('Shortened a long packet from %d to %d values, new packet length: %d' %
485 (len(orig_values), num_values, len(out)))
487 # Too long a response, send an error
488 log.msg('Could not send response, too long: %d bytes' % len(out))
489 self.stats.errorAction(request)
490 msg = {TID : tid, TYP : ERR, ERR : [KRPC_ERROR_RESPONSE_TOO_LONG, "response was %d bytes" % len(out)]}
494 # Unknown error, send an error message
495 self.stats.errorAction(request)
496 msg = {TID : tid, TYP : ERR, ERR : [KRPC_ERROR_SERVER_ERROR, "unknown error sending response: %s" % str(e)]}
499 self.stats.sentBytes(len(out))
500 self.transport.write(out, addr)
503 def sendRequest(self, method, args):
504 """Send a request to the remote node.
506 @type method: C{string}
507 @param method: the method name to call on the remote node
508 @param args: the arguments to send to the remote node's method
511 return defer.fail(KrpcError(KRPC_ERROR_PROTOCOL_STOPPED,
512 "cannot send, connection has been stopped"))
514 # Create the request message
516 msg = {TID : newTID, TYP : REQ, REQ : method, ARG : args}
517 if self.config.get('SPEW', False):
518 log.msg("%d sending to %r: %s" % (self.factory.port, self.addr, msg))
521 # Create the request object and save it with the TID
522 req = KrpcRequest(self, newTID, method, data, self.config)
523 self.tids[newTID] = req
525 # Save the conclusion of the action
526 req.addCallbacks(self.stats.responseAction, self.stats.failedAction,
527 callbackArgs = (method, datetime.now()),
528 errbackArgs = (method, datetime.now()))
532 def sendData(self, method, data):
533 """Write a request to the transport and save the stats.
535 @type method: C{string}
536 @param method: the name of the method to call on the remote node
537 @type data: C{string}
538 @param data: the message to send to the remote node
540 self.transport.write(data, self.addr)
541 self.stats.sentAction(method)
542 self.stats.sentBytes(len(data))
544 def timeOut(self, badTID, method):
545 """Call the deferred's errback if a timeout occurs.
547 @param badTID: the transaction ID of the request
548 @type method: C{string}
549 @param method: the name of the method that timed out on the remote node
551 if badTID in self.tids:
552 req = self.tids[badTID]
553 del(self.tids[badTID])
554 req.errback(KrpcError(KRPC_ERROR_TIMEOUT, "timeout waiting for '%s' from %r" %
555 (method, self.addr)))
557 log.msg('Received a timeout for an unknown request for %s from %r' % (method, self.addr))
560 """Cancel all pending requests."""
561 for req in self.tids.values():
562 req.errback(KrpcError(KRPC_ERROR_PROTOCOL_STOPPED,
563 'connection has been stopped while waiting for response'))
567 #{ For testing the KRPC protocol
568 def connectionForAddr(host, port):
571 class Receiver(protocol.Factory):
575 def krpc_store(self, msg, _krpc_sender):
578 def krpc_echo(self, msg, _krpc_sender):
580 def krpc_values(self, length, num, _krpc_sender):
581 return {'values': ['1'*length]*num}
584 from stats import StatsLogger
586 a = hostbroker(af, StatsLogger(None, None),
587 {'KRPC_TIMEOUT': 9, 'KRPC_INITIAL_DELAY': 2, 'SPEW': False})
589 p = reactor.listenUDP(port, a)
592 class KRPCTests(unittest.TestCase):
596 self.af, self.a, self.ap = make(1180)
597 self.bf, self.b, self.bp = make(1181)
600 self.ap.stopListening()
601 self.bp.stopListening()
603 def bufEquals(self, result, value):
604 self.failUnlessEqual(self.bf.buf, value)
606 def testSimpleMessage(self):
607 d = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('store', {'msg' : "This is a test."})
608 d.addCallback(self.bufEquals, ["This is a test."])
611 def testMessageBlast(self):
613 d = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('store', {'msg' : "This is a test."})
614 d.addCallback(self.bufEquals, ["This is a test."] * 100)
618 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
619 df.addCallback(self.gotMsg, "This is a test.")
622 def gotMsg(self, dict, should_be):
623 _krpc_sender = dict['_krpc_sender']
624 self.failUnlessEqual(dict['msg'], should_be)
626 def testManyEcho(self):
627 for i in xrange(100):
628 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
629 df.addCallback(self.gotMsg, "This is a test.")
632 def testMultiEcho(self):
633 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
634 df.addCallback(self.gotMsg, "This is a test.")
636 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is another test."})
637 df.addCallback(self.gotMsg, "This is another test.")
639 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is yet another test."})
640 df.addCallback(self.gotMsg, "This is yet another test.")
644 def testEchoReset(self):
645 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test."})
646 df.addCallback(self.gotMsg, "This is a test.")
648 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is another test."})
649 df.addCallback(self.gotMsg, "This is another test.")
650 df.addCallback(self.echoReset)
653 def echoReset(self, dict):
654 del(self.a.connections[('127.0.0.1', 1181)])
655 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is yet another test."})
656 df.addCallback(self.gotMsg, "This is yet another test.")
659 def testUnknownMeth(self):
660 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('blahblah', {'msg' : "This is a test."})
661 df = self.failUnlessFailure(df, KrpcError)
662 df.addBoth(self.gotErr, KRPC_ERROR_METHOD_UNKNOWN)
665 def testMalformedRequest(self):
666 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('echo', {'msg' : "This is a test.", 'foo': 'bar'})
667 df = self.failUnlessFailure(df, KrpcError)
668 df.addBoth(self.gotErr, KRPC_ERROR_MALFORMED_REQUEST, TypeError)
671 def gotErr(self, value, should_be, *errorTypes):
672 self.failUnlessEqual(value[0], should_be)
674 self.flushLoggedErrors(*errorTypes)
676 def testLongPackets(self):
677 df = self.a.connectionForAddr(('127.0.0.1', 1181)).sendRequest('values', {'length' : 1, 'num': 2000})
678 df.addCallback(self.gotLongRsp)
681 def gotLongRsp(self, dict):
682 # Not quite accurate, but good enough
683 self.failUnless(len(bencode(dict))-10 < UDP_PACKET_LIMIT)