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, traffic_happened=variable.Event(), ignore_trailing_payload=False):
  19         self._message_prefix = message_prefix
  20         self._max_payload_length = max_payload_length
  21         self.dataReceived2 = datachunker.DataChunker(self.dataReceiver())
  22         self.traffic_happened = traffic_happened
  23         self.ignore_trailing_payload = ignore_trailing_payload
  24
  25     def dataReceived(self, data):
  26         self.traffic_happened.happened('p2p/in', len(data))
  27         self.dataReceived2(data)
  28
  29     def dataReceiver(self):
  30         while True:
  31             start = ''
  32             while start != self._message_prefix:
  33                 start = (start + (yield 1))[-len(self._message_prefix):]
  34
  35             command = (yield 12).rstrip('\0')
  36             length, = struct.unpack('<I', (yield 4))
  37             if length > self._max_payload_length:
  38                 print 'length too large'
  39                 continue
  40             checksum = yield 4
  41             payload = yield length
  42
  43             if hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4] != checksum:
  44                 print 'invalid hash for', self.transport.getPeer().host, repr(command), length, checksum.encode('hex')
  45                 if p2pool.DEBUG:
  46                     print hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4].encode('hex'), payload.encode('hex')
  47                 self.badPeerHappened()
  48                 continue
  49
  50             type_ = getattr(self, 'message_' + command, None)
  51             if type_ is None:
  52                 if p2pool.DEBUG:
  53                     print 'no type for', repr(command)
  54                 continue
  55
  56             try:
  57                 self.packetReceived(command, type_.unpack(payload, self.ignore_trailing_payload))
  58             except:
  59                 print 'RECV', command, payload[:100].encode('hex') + ('...' if len(payload) > 100 else '')
  60                 log.err(None, 'Error handling message: (see RECV line)')
  61                 self.disconnect()
  62
  63     def packetReceived(self, command, payload2):
  64         handler = getattr(self, 'handle_' + command, None)
  65         if handler is None:
  66             if p2pool.DEBUG:
  67                 print 'no handler for', repr(command)
  68             return
  69
  70         if getattr(self, 'connected', True) and not getattr(self, 'disconnecting', False):
  71             handler(**payload2)
  72
  73     def disconnect(self):
  74         if hasattr(self.transport, 'abortConnection'):
  75             # Available since Twisted 11.1
  76             self.transport.abortConnection()
  77         else:
  78             # This doesn't always close timed out connections! warned about in main
  79             self.transport.loseConnection()
  80
  81     def badPeerHappened(self):
  82         self.disconnect()
  83
  84     def sendPacket(self, command, payload2):
  85         if len(command) >= 12:
  86             raise ValueError('command too long')
  87         type_ = getattr(self, 'message_' + command, None)
  88         if type_ is None:
  89             raise ValueError('invalid command')
  90         #print 'SEND', command, repr(payload2)[:500]
  91         payload = type_.pack(payload2)
  92         if len(payload) > self._max_payload_length:
  93             raise TooLong('payload too long')
  94         data = self._message_prefix + struct.pack('<12sI', command, len(payload)) + hashlib.sha256(hashlib.sha256(payload).digest()).digest()[:4] + payload
  95         self.traffic_happened.happened('p2p/out', len(data))
  96         self.transport.write(data)
  97
  98     def __getattr__(self, attr):
  99         prefix = 'send_'
 100         if attr.startswith(prefix):
 101             command = attr[len(prefix):]
 102             return lambda **payload2: self.sendPacket(command, payload2)
 103         #return protocol.Protocol.__getattr__(self, attr)
 104         raise AttributeError(attr)