citizenaspirant/PyBitmessage
1
0
Fork 0
Code Issues 1 Pull Requests Releases Wiki Activity

Refactor packet header extraction and generation

Browse Source 
Demote payloadLength from class instance variable to processData local variable as no other function was using it
Improve processData:
 -Utilise shared.Header
 -Use a memoryview to reduce memory overhead
 -Clean up variables before a recursive call
 -Strip null bytes from command
Refactor sendData
Various functions:
 -Use shared.CreatePacket to generate packets
Fix typo in _checkIPv4Address
This commit is contained in:
bmng-dev 2014-05-22 15:57:48 +00:00
parent ccfbe8da47
commit 0f9625aac7
1 changed files with 66 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

148
src/class_receiveDataThread.py
View File

@ -47,7 +47,6 @@ class receiveDataThread(threading.Thread):
self.sock = sock self.sock = sock
self.peer = shared.Peer(HOST, port) self.peer = shared.Peer(HOST, port)
self.streamNumber = streamNumber self.streamNumber = streamNumber
self.payloadLength = 0 # This is the protocol payload length thus it doesn't include the 24 byte message header
self.objectsThatWeHaveYetToGetFromThisPeer = {} self.objectsThatWeHaveYetToGetFromThisPeer = {}
self.selfInitiatedConnections = selfInitiatedConnections self.selfInitiatedConnections = selfInitiatedConnections
self.sendDataThreadQueue = sendDataThreadQueue # used to send commands and data to the sendDataThread self.sendDataThreadQueue = sendDataThreadQueue # used to send commands and data to the sendDataThread
@ -113,28 +112,39 @@ class receiveDataThread(threading.Thread):
# with shared.printLock: # with shared.printLock:
# print 'self.data is currently ', repr(self.data) # print 'self.data is currently ', repr(self.data)
# #
if len(self.data) < 24: # if so little of the data has arrived that we can't even read the checksum then wait for more data. if len(self.data) < shared.Header.size: # if so little of the data has arrived that we can't even read the checksum then wait for more data.
return return
if self.data[0:4] != '\xe9\xbe\xb4\xd9': #Use a memoryview so we don't copy data unnecessarily
view = memoryview(self.data)
magic,command,payloadLength,checksum = shared.Header.unpack(view[:shared.Header.size])
view = view[shared.Header.size:]
if magic != 0xE9BEB4D9:
#if shared.verbose >= 1: #if shared.verbose >= 1:
# with shared.printLock: # with shared.printLock:
# print 'The magic bytes were not correct. First 40 bytes of data: ' + repr(self.data[0:40]) # print 'The magic bytes were not correct. First 40 bytes of data: ' + repr(self.data[0:40])
self.data = "" self.data = ""
return return
self.payloadLength, = unpack('>L', self.data[16:20]) if payloadLength > 20000000:
if self.payloadLength > 20000000: logger.info('The incoming message, which we have not yet download, is too large. Ignoring it. (unfortunately there is no way to tell the other node to stop sending it except to disconnect.) Message size: %s' % payloadLength)
logger.info('The incoming message, which we have not yet download, is too large. Ignoring it. (unfortunately there is no way to tell the other node to stop sending it except to disconnect.) Message size: %s' % self.payloadLength) self.data = view[payloadLength:].tobytes()
self.data = self.data[self.payloadLength + 24:] del view,magic,command,payloadLength,checksum #we don't need these anymore and better to clean them now before the recursive call rather than after
self.processData() self.processData()
return return
if len(self.data) < self.payloadLength + 24: # check if the whole message has arrived yet. if len(view) < payloadLength: # check if the whole message has arrived yet.
return return
if self.data[20:24] != hashlib.sha512(self.data[24:self.payloadLength + 24]).digest()[0:4]: # test the checksum in the message. If it is correct... payload = view[:payloadLength]
if checksum != hashlib.sha512(payload).digest()[0:4]: # test the checksum in the message. If it is correct...
print 'Checksum incorrect. Clearing this message.' print 'Checksum incorrect. Clearing this message.'
self.data = self.data[self.payloadLength + 24:] self.data = view[payloadLength:].tobytes()
del view,magic,command,payloadLength,checksum,payload #again better to clean up before the recursive call
self.processData() self.processData()
return return
#We can now revert back to bytestrings and take this message out
payload = payload.tobytes()
self.data = view[payloadLength:].tobytes()
del view,magic,payloadLength,checksum
# The time we've last seen this node is obviously right now since we # The time we've last seen this node is obviously right now since we
# just received valid data from it. So update the knownNodes list so # just received valid data from it. So update the knownNodes list so
# that other peers can be made aware of its existance. # that other peers can be made aware of its existance.
@ -143,37 +153,39 @@ class receiveDataThread(threading.Thread):
shared.knownNodes[self.streamNumber][self.peer] = int(time.time()) shared.knownNodes[self.streamNumber][self.peer] = int(time.time())
shared.knownNodesLock.release() shared.knownNodesLock.release()
remoteCommand = self.data[4:16] #Strip the nulls
command = command.rstrip('\x00')
with shared.printLock: with shared.printLock:
print 'remoteCommand', repr(remoteCommand.replace('\x00', '')), ' from', self.peer print 'remoteCommand', repr(command), ' from', self.peer
if remoteCommand == 'version\x00\x00\x00\x00\x00' and not self.connectionIsOrWasFullyEstablished: #TODO: Use a dispatcher here
self.recversion(self.data[24:self.payloadLength + 24]) if not self.connectionIsOrWasFullyEstablished:
elif remoteCommand == 'verack\x00\x00\x00\x00\x00\x00' and not self.connectionIsOrWasFullyEstablished: if command == 'version':
self.recverack() self.recversion(payload)
elif remoteCommand == 'addr\x00\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: elif command == 'verack':
self.recaddr(self.data[24:self.payloadLength + 24]) self.recverack()
elif remoteCommand == 'getpubkey\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: else:
shared.checkAndSharegetpubkeyWithPeers(self.data[24:self.payloadLength + 24]) if command == 'addr':
elif remoteCommand == 'pubkey\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recaddr(payload)
self.recpubkey(self.data[24:self.payloadLength + 24]) elif command == 'getpubkey':
elif remoteCommand == 'inv\x00\x00\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: shared.checkAndSharegetpubkeyWithPeers(payload)
self.recinv(self.data[24:self.payloadLength + 24]) elif command == 'pubkey':
elif remoteCommand == 'getdata\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recpubkey(payload)
self.recgetdata(self.data[24:self.payloadLength + 24]) elif command == 'inv':
elif remoteCommand == 'msg\x00\x00\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recinv(payload)
self.recmsg(self.data[24:self.payloadLength + 24]) elif command == 'getdata':
elif remoteCommand == 'broadcast\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recgetdata(payload)
self.recbroadcast(self.data[24:self.payloadLength + 24]) elif command == 'msg':
elif remoteCommand == 'ping\x00\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recmsg(payload)
self.sendpong() elif command == 'broadcast':
elif remoteCommand == 'pong\x00\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.recbroadcast(payload)
pass elif command == 'ping':
elif remoteCommand == 'alert\x00\x00\x00\x00\x00\x00\x00' and self.connectionIsOrWasFullyEstablished: self.sendpong(payload)
pass #elif command == 'pong':
# pass
#elif command == 'alert':
# pass
self.data = self.data[
self.payloadLength + 24:] # take this message out and then process the next message
if self.data == '': if self.data == '':
while len(self.objectsThatWeHaveYetToGetFromThisPeer) > 0: while len(self.objectsThatWeHaveYetToGetFromThisPeer) > 0:
shared.numberOfInventoryLookupsPerformed += 1 shared.numberOfInventoryLookupsPerformed += 1
@ -226,7 +238,7 @@ class receiveDataThread(threading.Thread):
def sendpong(self): def sendpong(self):
print 'Sending pong' print 'Sending pong'
self.sendDataThreadQueue.put((0, 'sendRawData', '\xE9\xBE\xB4\xD9\x70\x6F\x6E\x67\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xcf\x83\xe1\x35')) self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket('pong')))
def recverack(self): def recverack(self):
@ -297,7 +309,7 @@ class receiveDataThread(threading.Thread):
for hash, storedValue in bigInvList.items(): for hash, storedValue in bigInvList.items():
payload += hash payload += hash
numberOfObjectsInInvMessage += 1 numberOfObjectsInInvMessage += 1
if numberOfObjectsInInvMessage >= 50000: # We can only send a max of 50000 items per inv message but we may have more objects to advertise. They must be split up into multiple inv messages. if numberOfObjectsInInvMessage == 50000: # We can only send a max of 50000 items per inv message but we may have more objects to advertise. They must be split up into multiple inv messages.
self.sendinvMessageToJustThisOnePeer( self.sendinvMessageToJustThisOnePeer(
numberOfObjectsInInvMessage, payload) numberOfObjectsInInvMessage, payload)
payload = '' payload = ''
@ -311,13 +323,9 @@ class receiveDataThread(threading.Thread):
# function for broadcasting invs to everyone in our stream. # function for broadcasting invs to everyone in our stream.
def sendinvMessageToJustThisOnePeer(self, numberOfObjects, payload): def sendinvMessageToJustThisOnePeer(self, numberOfObjects, payload):
payload = encodeVarint(numberOfObjects) + payload payload = encodeVarint(numberOfObjects) + payload
headerData = '\xe9\xbe\xb4\xd9' # magic bits, slighly different from Bitcoin's magic bits.
headerData += 'inv\x00\x00\x00\x00\x00\x00\x00\x00\x00'
headerData += pack('>L', len(payload))
headerData += hashlib.sha512(payload).digest()[:4]
with shared.printLock: with shared.printLock:
print 'Sending huge inv message with', numberOfObjects, 'objects to just this one peer' print 'Sending huge inv message with', numberOfObjects, 'objects to just this one peer'
self.sendDataThreadQueue.put((0, 'sendRawData', headerData + payload)) self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket('inv', payload)))
def _sleepForTimingAttackMitigation(self, sleepTime): def _sleepForTimingAttackMitigation(self, sleepTime):
# We don't need to do the timing attack mitigation if we are # We don't need to do the timing attack mitigation if we are
@ -457,12 +465,7 @@ class receiveDataThread(threading.Thread):
print 'sending getdata to retrieve object with hash:', hash.encode('hex') print 'sending getdata to retrieve object with hash:', hash.encode('hex')
payload = '\x01' + hash payload = '\x01' + hash
headerData = '\xe9\xbe\xb4\xd9' # magic bits, slighly different from Bitcoin's magic bits. self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket('getdata', payload)))
headerData += 'getdata\x00\x00\x00\x00\x00'
headerData += pack('>L', len(
payload)) # payload length. Note that we add an extra 8 for the nonce.
headerData += hashlib.sha512(payload).digest()[:4]
self.sendDataThreadQueue.put((0, 'sendRawData', headerData + payload))
# We have received a getdata request from our peer # We have received a getdata request from our peer
@ -499,34 +502,17 @@ class receiveDataThread(threading.Thread):
# Our peer has requested (in a getdata message) that we send an object. # Our peer has requested (in a getdata message) that we send an object.
def sendData(self, objectType, payload): def sendData(self, objectType, payload):
headerData = '\xe9\xbe\xb4\xd9' # magic bits, slighly different from Bitcoin's magic bits. # pubkeyrequest??
if objectType == 'pubkey': if objectType == 'pubkeyrequest':
with shared.printLock: objectType = 'getpubkey'
print 'sending pubkey' elif objectType != 'pubkey' and objectType != 'getpubkey' and
objectType != 'msg' and objectType != 'broadcast':
headerData += 'pubkey\x00\x00\x00\x00\x00\x00'
elif objectType == 'getpubkey' or objectType == 'pubkeyrequest':
with shared.printLock:
print 'sending getpubkey'
headerData += 'getpubkey\x00\x00\x00'
elif objectType == 'msg':
with shared.printLock:
print 'sending msg'
headerData += 'msg\x00\x00\x00\x00\x00\x00\x00\x00\x00'
elif objectType == 'broadcast':
with shared.printLock:
print 'sending broadcast'
headerData += 'broadcast\x00\x00\x00'
else:
sys.stderr.write( sys.stderr.write(
'Error: sendData has been asked to send a strange objectType: %s\n' % str(objectType)) 'Error: sendData has been asked to send a strange objectType: %s\n' % str(objectType))
return return
headerData += pack('>L', len(payload)) # payload length. with shared.printLock:
headerData += hashlib.sha512(payload).digest()[:4] print 'sending', objectType
self.sendDataThreadQueue.put((0, 'sendRawData', headerData + payload)) self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket(objectType, payload)))
# Advertise this object to all of our peers # Advertise this object to all of our peers
@ -545,7 +531,7 @@ class receiveDataThread(threading.Thread):
if host[0] == '\x0A': if host[0] == '\x0A':
print 'Ignoring IP address in private range:', hostFromAddrMessage print 'Ignoring IP address in private range:', hostFromAddrMessage
return False return False
if host[0:2] == '\xC0A8': if host[0:2] == '\xC0\xA8':
print 'Ignoring IP address in private range:', hostFromAddrMessage print 'Ignoring IP address in private range:', hostFromAddrMessage
return False return False
return True return True
@ -745,11 +731,7 @@ class receiveDataThread(threading.Thread):
payload += pack('>H', PORT) # remote port payload += pack('>H', PORT) # remote port
payload = encodeVarint(numberOfAddressesInAddrMessage) + payload payload = encodeVarint(numberOfAddressesInAddrMessage) + payload
datatosend = '\xE9\xBE\xB4\xD9addr\x00\x00\x00\x00\x00\x00\x00\x00' self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket('addr', payload)))
datatosend = datatosend + pack('>L', len(payload)) # payload length
datatosend = datatosend + hashlib.sha512(payload).digest()[0:4]
datatosend = datatosend + payload
self.sendDataThreadQueue.put((0, 'sendRawData', datatosend))
# We have received a version message # We have received a version message
@ -832,7 +814,7 @@ class receiveDataThread(threading.Thread):
def sendverack(self): def sendverack(self):
with shared.printLock: with shared.printLock:
print 'Sending verack' print 'Sending verack'
self.sendDataThreadQueue.put((0, 'sendRawData', '\xE9\xBE\xB4\xD9\x76\x65\x72\x61\x63\x6B\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xcf\x83\xe1\x35')) self.sendDataThreadQueue.put((0, 'sendRawData', shared.CreatePacket('verack')))
self.verackSent = True self.verackSent = True
if self.verackReceived: if self.verackReceived:
self.connectionFullyEstablished() self.connectionFullyEstablished()