forked from LBRYCommunity/lbry-sdk
346 lines
14 KiB
Python
346 lines
14 KiB
Python
import json
|
|
import logging
|
|
|
|
from twisted.internet.error import ConnectionRefusedError
|
|
from twisted.protocols.basic import FileSender
|
|
from twisted.internet.protocol import Protocol, ClientFactory
|
|
from twisted.internet import defer, error
|
|
|
|
from lbrynet.reflector.common import IncompleteResponse, ReflectorRequestError
|
|
from lbrynet.reflector.common import REFLECTOR_V1, REFLECTOR_V2
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class EncryptedFileReflectorClient(Protocol):
|
|
# Protocol stuff
|
|
def connectionMade(self):
|
|
log.debug("Connected to reflector")
|
|
self.response_buff = b''
|
|
self.outgoing_buff = b''
|
|
self.blob_hashes_to_send = []
|
|
self.failed_blob_hashes = []
|
|
self.next_blob_to_send = None
|
|
self.read_handle = None
|
|
self.sent_stream_info = False
|
|
self.received_descriptor_response = False
|
|
self.received_server_version = False
|
|
self.server_version = None
|
|
self.stream_descriptor = None
|
|
self.descriptor_needed = None
|
|
self.needed_blobs = []
|
|
self.reflected_blobs = []
|
|
self.file_sender = None
|
|
self.producer = None
|
|
self.streaming = False
|
|
|
|
self.blob_manager = self.factory.blob_manager
|
|
self.protocol_version = self.factory.protocol_version
|
|
self.stream_hash = self.factory.stream_hash
|
|
self.sd_hash = self.factory.sd_hash
|
|
|
|
d = self.load_descriptor()
|
|
d.addCallback(lambda _: self.send_handshake())
|
|
d.addErrback(lambda err: log.warning("An error occurred immediately: %s", err.getTraceback()))
|
|
|
|
def dataReceived(self, data):
|
|
self.response_buff += data
|
|
try:
|
|
msg = self.parse_response(self.response_buff)
|
|
except IncompleteResponse:
|
|
pass
|
|
else:
|
|
self.response_buff = b''
|
|
d = self.handle_response(msg)
|
|
d.addCallback(lambda _: self.send_next_request())
|
|
d.addErrback(self.response_failure_handler)
|
|
|
|
def store_result(self, result):
|
|
if not self.needed_blobs or len(self.reflected_blobs) == len(self.needed_blobs):
|
|
reflected = True
|
|
else:
|
|
reflected = False
|
|
|
|
d = self.blob_manager.storage.update_reflected_stream(self.sd_hash, self.transport.getPeer().host, reflected)
|
|
d.addCallback(lambda _: result)
|
|
return d
|
|
|
|
def connectionLost(self, reason):
|
|
# make sure blob file readers get closed
|
|
self.set_not_uploading()
|
|
|
|
if reason.check(error.ConnectionDone):
|
|
if not self.needed_blobs:
|
|
log.info("Reflector has all blobs for %s", self.stream_descriptor)
|
|
elif not self.reflected_blobs:
|
|
log.info("No more completed blobs for %s to reflect, %i are still needed",
|
|
self.stream_descriptor, len(self.needed_blobs))
|
|
else:
|
|
log.info('Finished sending reflector %i blobs for %s',
|
|
len(self.reflected_blobs), self.stream_descriptor)
|
|
result = self.reflected_blobs
|
|
elif reason.check(error.ConnectionLost):
|
|
log.warning("Stopped reflecting %s after sending %i blobs",
|
|
self.stream_descriptor, len(self.reflected_blobs))
|
|
result = self.reflected_blobs
|
|
else:
|
|
log.info('Reflector finished for %s: %s', self.stream_descriptor,
|
|
reason)
|
|
result = reason
|
|
self.factory.finished_deferred.addCallback(self.store_result)
|
|
self.factory.finished_deferred.callback(result)
|
|
|
|
# IConsumer stuff
|
|
|
|
def registerProducer(self, producer, streaming):
|
|
self.producer = producer
|
|
self.streaming = streaming
|
|
if self.streaming is False:
|
|
from twisted.internet import reactor
|
|
reactor.callLater(0, self.producer.resumeProducing)
|
|
|
|
def unregisterProducer(self):
|
|
self.producer = None
|
|
|
|
def write(self, data):
|
|
self.transport.write(data)
|
|
if self.producer is not None and self.streaming is False:
|
|
from twisted.internet import reactor
|
|
reactor.callLater(0, self.producer.resumeProducing)
|
|
|
|
def get_validated_blobs(self, blobs_in_stream):
|
|
def get_blobs(blobs):
|
|
for crypt_blob in blobs:
|
|
if crypt_blob.blob_hash and crypt_blob.length:
|
|
yield self.blob_manager.get_blob(crypt_blob.blob_hash, crypt_blob.length)
|
|
|
|
dl = defer.DeferredList(list(get_blobs(blobs_in_stream)), consumeErrors=True)
|
|
dl.addCallback(lambda blobs: [blob for r, blob in blobs if r and blob.get_is_verified()])
|
|
return dl
|
|
|
|
def set_blobs_to_send(self, blobs_to_send):
|
|
for blob in blobs_to_send:
|
|
if blob.blob_hash not in self.blob_hashes_to_send:
|
|
self.blob_hashes_to_send.append(blob.blob_hash)
|
|
|
|
def get_blobs_to_send(self):
|
|
def _show_missing_blobs(filtered):
|
|
if filtered:
|
|
needs_desc = "" if not self.descriptor_needed else "descriptor and "
|
|
log.info("Reflector needs %s%i blobs for stream",
|
|
needs_desc,
|
|
len(filtered))
|
|
return filtered
|
|
|
|
d = self.factory.blob_manager.storage.get_blobs_for_stream(self.stream_hash)
|
|
d.addCallback(self.get_validated_blobs)
|
|
if not self.descriptor_needed:
|
|
d.addCallback(lambda filtered:
|
|
[blob for blob in filtered if blob.blob_hash in self.needed_blobs])
|
|
d.addCallback(_show_missing_blobs)
|
|
d.addCallback(self.set_blobs_to_send)
|
|
d.addCallback(lambda _: None if self.descriptor_needed else self.set_not_uploading())
|
|
return d
|
|
|
|
def send_request(self, request_dict):
|
|
self.write(json.dumps(request_dict).encode())
|
|
|
|
def send_handshake(self):
|
|
self.send_request({'version': self.protocol_version})
|
|
|
|
@defer.inlineCallbacks
|
|
def load_descriptor(self):
|
|
if self.sd_hash:
|
|
self.stream_descriptor = yield self.factory.blob_manager.get_blob(self.sd_hash)
|
|
else:
|
|
raise ValueError("no sd hash for stream %s" % self.stream_hash)
|
|
|
|
def parse_response(self, buff):
|
|
try:
|
|
return json.loads(buff)
|
|
except ValueError:
|
|
raise IncompleteResponse()
|
|
|
|
def response_failure_handler(self, err):
|
|
log.warning("An error occurred handling the response: %s", err.getTraceback())
|
|
|
|
def handle_response(self, response_dict):
|
|
if not self.received_server_version:
|
|
return self.handle_handshake_response(response_dict)
|
|
elif not self.received_descriptor_response and self.server_version == REFLECTOR_V2:
|
|
return self.handle_descriptor_response(response_dict)
|
|
else:
|
|
return self.handle_normal_response(response_dict)
|
|
|
|
def set_not_uploading(self):
|
|
if self.next_blob_to_send is not None:
|
|
log.debug("Close %s", self.next_blob_to_send)
|
|
self.read_handle.close()
|
|
self.read_handle = None
|
|
self.next_blob_to_send = None
|
|
if self.file_sender is not None:
|
|
self.file_sender.stopProducing()
|
|
self.file_sender = None
|
|
return defer.succeed(None)
|
|
|
|
def start_transfer(self):
|
|
assert self.read_handle is not None, \
|
|
"self.read_handle was None when trying to start the transfer"
|
|
d = self.file_sender.beginFileTransfer(self.read_handle, self)
|
|
d.addCallback(lambda _: self.read_handle.close())
|
|
return d
|
|
|
|
def handle_handshake_response(self, response_dict):
|
|
if 'version' not in response_dict:
|
|
raise ValueError("Need protocol version number!")
|
|
self.server_version = int(response_dict['version'])
|
|
if self.server_version not in [REFLECTOR_V1, REFLECTOR_V2]:
|
|
raise ValueError("I can't handle protocol version {}!".format(self.server_version))
|
|
self.received_server_version = True
|
|
return defer.succeed(True)
|
|
|
|
def handle_descriptor_response(self, response_dict):
|
|
if self.file_sender is None: # Expecting Server Info Response
|
|
if 'send_sd_blob' not in response_dict:
|
|
raise ReflectorRequestError("I don't know whether to send the sd blob or not!")
|
|
if response_dict['send_sd_blob'] is True:
|
|
self.file_sender = FileSender()
|
|
else:
|
|
self.received_descriptor_response = True
|
|
self.descriptor_needed = response_dict['send_sd_blob']
|
|
self.needed_blobs = response_dict.get('needed_blobs', [])
|
|
return self.get_blobs_to_send()
|
|
else: # Expecting Server Blob Response
|
|
if 'received_sd_blob' not in response_dict:
|
|
raise ValueError("I don't know if the sd blob made it to the intended destination!")
|
|
else:
|
|
self.received_descriptor_response = True
|
|
disconnect = False
|
|
if response_dict['received_sd_blob']:
|
|
self.reflected_blobs.append(self.next_blob_to_send.blob_hash)
|
|
log.info("Sent reflector descriptor %s", self.next_blob_to_send)
|
|
else:
|
|
log.warning("Reflector failed to receive descriptor %s",
|
|
self.next_blob_to_send)
|
|
disconnect = True
|
|
d = self.set_not_uploading()
|
|
if disconnect:
|
|
d.addCallback(lambda _: self.transport.loseConnection())
|
|
return d
|
|
|
|
def handle_normal_response(self, response_dict):
|
|
if self.file_sender is None: # Expecting Server Info Response
|
|
if 'send_blob' not in response_dict:
|
|
raise ValueError("I don't know whether to send the blob or not!")
|
|
if response_dict['send_blob'] is True:
|
|
self.file_sender = FileSender()
|
|
return defer.succeed(True)
|
|
else:
|
|
log.info("Reflector already has %s", self.next_blob_to_send)
|
|
return self.set_not_uploading()
|
|
else: # Expecting Server Blob Response
|
|
if 'received_blob' not in response_dict:
|
|
raise ValueError("I don't know if the blob made it to the intended destination!")
|
|
else:
|
|
if response_dict['received_blob']:
|
|
self.reflected_blobs.append(self.next_blob_to_send.blob_hash)
|
|
log.debug("Sent reflector blob %s", self.next_blob_to_send)
|
|
else:
|
|
log.warning("Reflector failed to receive blob %s", self.next_blob_to_send)
|
|
return self.set_not_uploading()
|
|
|
|
def open_blob_for_reading(self, blob):
|
|
if blob.get_is_verified():
|
|
read_handle = blob.open_for_reading()
|
|
if read_handle is not None:
|
|
log.debug('Getting ready to send %s', blob.blob_hash)
|
|
self.next_blob_to_send = blob
|
|
self.read_handle = read_handle
|
|
return defer.succeed(None)
|
|
return defer.fail(ValueError(
|
|
"Couldn't open that blob for some reason. blob_hash: {}".format(blob.blob_hash)))
|
|
|
|
def send_blob_info(self):
|
|
assert self.next_blob_to_send is not None, "need to have a next blob to send at this point"
|
|
r = {
|
|
'blob_hash': self.next_blob_to_send.blob_hash,
|
|
'blob_size': self.next_blob_to_send.length
|
|
}
|
|
self.send_request(r)
|
|
|
|
def send_descriptor_info(self):
|
|
assert self.stream_descriptor is not None, "need to have a sd blob to send at this point"
|
|
r = {
|
|
'sd_blob_hash': self.stream_descriptor.blob_hash,
|
|
'sd_blob_size': self.stream_descriptor.length
|
|
}
|
|
self.sent_stream_info = True
|
|
self.send_request(r)
|
|
|
|
def skip_missing_blob(self, err, blob_hash):
|
|
err.trap(ValueError)
|
|
if blob_hash not in self.failed_blob_hashes:
|
|
log.warning("Failed to reflect blob %s, reason: %s",
|
|
str(blob_hash)[:16], err.getTraceback())
|
|
self.blob_hashes_to_send.append(blob_hash)
|
|
self.failed_blob_hashes.append(blob_hash)
|
|
else:
|
|
log.warning("Failed second try reflecting blob %s, giving up, reason: %s",
|
|
str(blob_hash)[:16], err.getTraceback())
|
|
|
|
def send_next_request(self):
|
|
if self.file_sender is not None:
|
|
# send the blob
|
|
return self.start_transfer()
|
|
elif not self.sent_stream_info:
|
|
# open the sd blob to send
|
|
blob = self.stream_descriptor
|
|
d = self.open_blob_for_reading(blob)
|
|
d.addCallbacks(lambda _: self.send_descriptor_info(),
|
|
lambda err: self.skip_missing_blob(err, blob.blob_hash))
|
|
return d
|
|
elif self.blob_hashes_to_send:
|
|
# open the next blob to send
|
|
blob_hash = self.blob_hashes_to_send[0]
|
|
self.blob_hashes_to_send = self.blob_hashes_to_send[1:]
|
|
d = self.blob_manager.get_blob(blob_hash)
|
|
d.addCallback(self.open_blob_for_reading)
|
|
d.addCallbacks(lambda _: self.send_blob_info(),
|
|
lambda err: self.skip_missing_blob(err, blob.blob_hash))
|
|
return d
|
|
# close connection
|
|
self.transport.loseConnection()
|
|
|
|
|
|
class EncryptedFileReflectorClientFactory(ClientFactory):
|
|
protocol = EncryptedFileReflectorClient
|
|
protocol_version = REFLECTOR_V2
|
|
|
|
def __init__(self, blob_manager, stream_hash, sd_hash):
|
|
self.blob_manager = blob_manager
|
|
self.stream_hash = stream_hash
|
|
self.sd_hash = sd_hash
|
|
self.p = None
|
|
self.finished_deferred = defer.Deferred()
|
|
|
|
def buildProtocol(self, addr):
|
|
p = self.protocol()
|
|
p.factory = self
|
|
self.p = p
|
|
return p
|
|
|
|
def startFactory(self):
|
|
log.debug('Starting reflector factory')
|
|
ClientFactory.startFactory(self)
|
|
|
|
def startedConnecting(self, connector):
|
|
log.debug('Connecting to reflector')
|
|
|
|
def clientConnectionLost(self, connector, reason):
|
|
"""If we get disconnected, reconnect to server."""
|
|
|
|
def clientConnectionFailed(self, connector, reason):
|
|
if reason.check(ConnectionRefusedError):
|
|
log.warning("Could not connect to reflector server")
|
|
else:
|
|
log.error("Reflector connection failed: %s", reason)
|