diff options
author | Victor Shyba <victor1984@riseup.net> | 2017-03-22 21:43:07 -0300 |
---|---|---|
committer | drebs <drebs@leap.se> | 2017-04-04 18:27:38 +0200 |
commit | 45b73d58930a2a66394a6797c94a50c34e8f96e7 (patch) | |
tree | 395e8085ebfb7773ccba8a5c292523b4018b1b0b | |
parent | 4cf662b0b043579badd231f30fc4eb58f9e6f09c (diff) |
[refactor] adds a PreamblePipe for preamble download
Downloading until there is a space then splitting the content was a
mess. Extracted this behaviour out of DecrypterBuffer into a new
component so it eases testing by introducing a single responsibility
class.
-rw-r--r-- | client/src/leap/soledad/client/_blobs.py | 27 | ||||
-rw-r--r-- | client/src/leap/soledad/client/_pipes.py | 29 | ||||
-rw-r--r-- | testing/tests/pipes/test_pipes.py | 21 |
3 files changed, 56 insertions, 21 deletions
diff --git a/client/src/leap/soledad/client/_blobs.py b/client/src/leap/soledad/client/_blobs.py index c64c2bb2..7ecd4a4d 100644 --- a/client/src/leap/soledad/client/_blobs.py +++ b/client/src/leap/soledad/client/_blobs.py @@ -36,7 +36,7 @@ import treq from leap.soledad.client.sqlcipher import SQLCipherOptions from leap.soledad.client import pragmas -from leap.soledad.client._pipes import TruncatedTailPipe +from leap.soledad.client._pipes import TruncatedTailPipe, PreamblePipe from leap.soledad.common.errors import SoledadError from _crypto import DocInfo, BlobEncryptor, BlobDecryptor @@ -106,33 +106,22 @@ def check_http_status(code): class DecrypterBuffer(object): def __init__(self, doc_id, rev, secret, tag): - self.decrypter = None - self.preamble = BytesIO() self.doc_info = DocInfo(doc_id, rev) self.secret = secret self.tag = tag - self.d = None + self.preamble_pipe = PreamblePipe(self._make_decryptor) - def write(self, data): - if not self.decrypter: - return self.prepare_decrypter(data) - - self.output.write(data) - - def prepare_decrypter(self, data): - if ' ' not in data: - self.preamble.write(data) - return - preamble_chunk, remaining = data.split(' ', 1) - self.preamble.write(preamble_chunk) + def _make_decryptor(self, preamble): self.decrypter = BlobDecryptor( - self.doc_info, BytesIO(self.preamble.getvalue()), + self.doc_info, preamble, secret=self.secret, armor=False, start_stream=False, tag=self.tag) - self.output = TruncatedTailPipe(self.decrypter, tail_size=16) - self.output.write(remaining) + return TruncatedTailPipe(self.decrypter, tail_size=len(self.tag)) + + def write(self, data): + self.preamble_pipe.write(data) def close(self): return self.decrypter._end_stream(), self.decrypter.size diff --git a/client/src/leap/soledad/client/_pipes.py b/client/src/leap/soledad/client/_pipes.py index ed89e14d..443f4a70 100644 --- a/client/src/leap/soledad/client/_pipes.py +++ b/client/src/leap/soledad/client/_pipes.py @@ -20,7 +20,7 @@ Components for piping data on streams. from io import BytesIO -__all__ = ['TruncatedTailPipe'] +__all__ = ['TruncatedTailPipe', 'PreamblePipe'] class TruncatedTailPipe(object): @@ -49,3 +49,30 @@ class TruncatedTailPipe(object): def close(self): return self.output + + +class PreamblePipe(object): + """ + Consumes data until a space is found, then calls a callback with it and + starts forwarding data to consumer returned by this callback. + """ + + def __init__(self, callback): + self.callback = callback + self.preamble = BytesIO() + self.output = None + + def write(self, data): + if not self.output: + self._write_preamble(data) + else: + self.output.write(data) + + def _write_preamble(self, data): + if ' ' not in data: + self.preamble.write(data) + return + preamble_chunk, remaining = data.split(' ', 1) + self.preamble.write(preamble_chunk) + self.output = self.callback(self.preamble) + self.output.write(remaining) diff --git a/testing/tests/pipes/test_pipes.py b/testing/tests/pipes/test_pipes.py index d7db2716..42ed81ac 100644 --- a/testing/tests/pipes/test_pipes.py +++ b/testing/tests/pipes/test_pipes.py @@ -19,9 +19,11 @@ Tests for streaming components. """ from twisted.trial import unittest from leap.soledad.client._pipes import TruncatedTailPipe +from leap.soledad.client._pipes import PreamblePipe +from io import BytesIO -class TruncatedTailTestCase(unittest.TestCase): +class TruncatedTailPipeTestCase(unittest.TestCase): def test_tail_truncating_pipe(self): pipe = TruncatedTailPipe(tail_size=20) @@ -30,3 +32,20 @@ class TruncatedTailTestCase(unittest.TestCase): pipe.write(data) result = pipe.close() assert result.getvalue() == 'A' * 100 + + +class PreamblePipeTestCase(unittest.TestCase): + + def test_preamble_pipe(self): + remaining = BytesIO() + preamble = BytesIO() + + def callback(dataBuffer): + preamble.write(dataBuffer.getvalue()) + return remaining + pipe = PreamblePipe(callback) + payload = 'A' * 100 + ' ' + 'B' * 20 + for data in payload: + pipe.write(data) + assert remaining.getvalue() == 'B' * 20 + assert preamble.getvalue() == 'A' * 100 |