p2pool/util/p2protocol.py

   1 '''
   2 Generic message-based protocol used by Bitcoin and P2Pool for P2P communication
   3 '''
   4
   5 import hashlib
   6 import struct
   7
   8 from twisted.internet import protocol
   9 from twisted.python import log
  10
  11 import p2pool
  12 from p2pool.util import datachunker, variable
  13
  14 class TooLong(Exception):
  15     pass
  16
  17 class Protocol(protocol.Protocol):
  18     def __init__(self, message_prefix, max_payload_length):
  19         self._message_prefix = message_prefix
  20         self._max_payload_length = max_payload_length
  21         self.dataReceived = datachunker.DataChunker(self.dataReceiver())
  22         self.paused_var = variable.Variable(False)
  23
  24     def connectionMade(self):
  25         self.transport.registerProducer(self, True)
  26
  27     def pauseProducing(self):
  28         self.paused_var.set(True)
  29
  30     def resumeProducing(self):
  31         self.paused_var.set(False)
  32
  33     def stopProducing(self):
  34         pass
  35
  36     def dataReceiver(self):
  37         while True:
  38             start = ''
  39             while start != self._message_prefix:
  40                 start = (start + (yield 1))[-len(self._message_prefix):]
  41
  42             command = (yield 12).rstrip('\0')
  43             length, = struct.unpack('<I', (yield 4))
  44             if length > self._max_payload_length:
  45                 print 'length too large'
  46                 continue
  47             checksum = yield 4
  48             payload = yield length
  49
  50             if hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4] != checksum:
  51                 print 'invalid hash for', self.transport.getPeer().host, repr(command), length, checksum.encode('hex'), hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4].encode('hex'), payload.encode('hex')
  52                 continue
  53
  54             type_ = getattr(self, 'message_' + command, None)
  55             if type_ is None:
  56                 if p2pool.DEBUG:
  57                     print 'no type for', repr(command)
  58                 continue
  59
  60             try:
  61                 self.packetReceived(command, type_.unpack(payload))
  62             except:
  63                 print 'RECV', command, payload[:100].encode('hex') + ('...' if len(payload) > 100 else '')
  64                 log.err(None, 'Error handling message: (see RECV line)')
  65                 self.badPeerHappened()
  66
  67     def packetReceived(self, command, payload2):
  68         handler = getattr(self, 'handle_' + command, None)
  69         if handler is None:
  70             if p2pool.DEBUG:
  71                 print 'no handler for', repr(command)
  72             return
  73
  74         handler(**payload2)
  75
  76     def badPeerHappened(self):
  77         self.transport.loseConnection()
  78
  79     def sendPacket(self, command, payload2):
  80         if len(command) >= 12:
  81             raise ValueError('command too long')
  82         type_ = getattr(self, 'message_' + command, None)
  83         if type_ is None:
  84             raise ValueError('invalid command')
  85         #print 'SEND', command, repr(payload2)[:500]
  86         payload = type_.pack(payload2)
  87         if len(payload) > self._max_payload_length:
  88             raise TooLong('payload too long')
  89         self.transport.write(self._message_prefix + struct.pack('<12sI', command, len(payload)) + hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4] + payload)
  90         return self.paused_var.get_when_satisfies(lambda paused: not paused)
  91
  92     def __getattr__(self, attr):
  93         prefix = 'send_'
  94         if attr.startswith(prefix):
  95             command = attr[len(prefix):]
  96             return lambda **payload2: self.sendPacket(command, payload2)
  97         #return protocol.Protocol.__getattr__(self, attr)
  98         raise AttributeError(attr)