diff options
| -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 |