diff options
| -rw-r--r-- | client/src/leap/soledad/client/_blobs.py | 23 | ||||
| -rw-r--r-- | client/src/leap/soledad/client/_pipes.py | 51 | ||||
| -rw-r--r-- | testing/tests/pipes/test_pipes.py | 32 |
3 files changed, 91 insertions, 15 deletions
diff --git a/client/src/leap/soledad/client/_blobs.py b/client/src/leap/soledad/client/_blobs.py index 0d25702c..c64c2bb2 100644 --- a/client/src/leap/soledad/client/_blobs.py +++ b/client/src/leap/soledad/client/_blobs.py @@ -36,6 +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.common.errors import SoledadError from _crypto import DocInfo, BlobEncryptor, BlobDecryptor @@ -106,7 +107,7 @@ class DecrypterBuffer(object): def __init__(self, doc_id, rev, secret, tag): self.decrypter = None - self.buffer = BytesIO() + self.preamble = BytesIO() self.doc_info = DocInfo(doc_id, rev) self.secret = secret self.tag = tag @@ -116,30 +117,22 @@ class DecrypterBuffer(object): if not self.decrypter: return self.prepare_decrypter(data) - self.buffer.write(data) - if self.buffer.tell() > 16: - overflow_size = self.buffer.tell() - 16 - self.buffer.seek(0) - self.decrypter.write(self.buffer.read(overflow_size)) - remaining = self.buffer.read() - self.buffer.seek(0) - self.buffer.write(remaining) - self.buffer.truncate() + self.output.write(data) def prepare_decrypter(self, data): if ' ' not in data: - self.buffer.write(data) + self.preamble.write(data) return preamble_chunk, remaining = data.split(' ', 1) - self.buffer.write(preamble_chunk) + self.preamble.write(preamble_chunk) self.decrypter = BlobDecryptor( - self.doc_info, BytesIO(self.buffer.getvalue()), + self.doc_info, BytesIO(self.preamble.getvalue()), secret=self.secret, armor=False, start_stream=False, tag=self.tag) - self.buffer = BytesIO() - self.write(remaining) + self.output = TruncatedTailPipe(self.decrypter, tail_size=16) + self.output.write(remaining) 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 new file mode 100644 index 00000000..ed89e14d --- /dev/null +++ b/client/src/leap/soledad/client/_pipes.py @@ -0,0 +1,51 @@ +# -*- coding: utf-8 -*- +# _pipes.py +# Copyright (C) 2017 LEAP +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see <http://www.gnu.org/licenses/>. +""" +Components for piping data on streams. +""" +from io import BytesIO + + +__all__ = ['TruncatedTailPipe'] + + +class TruncatedTailPipe(object): + """ + Truncate the last `tail_size` bytes from the stream. + """ + + def __init__(self, output=None, tail_size=16): + self.tail_size = tail_size + self.output = output or BytesIO() + self.buffer = BytesIO() + + def write(self, data): + self.buffer.write(data) + if self.buffer.tell() > self.tail_size: + self._truncate_tail() + + def _truncate_tail(self): + overflow_size = self.buffer.tell() - self.tail_size + self.buffer.seek(0) + self.output.write(self.buffer.read(overflow_size)) + remaining = self.buffer.read() + self.buffer.seek(0) + self.buffer.write(remaining) + self.buffer.truncate() + + def close(self): + return self.output diff --git a/testing/tests/pipes/test_pipes.py b/testing/tests/pipes/test_pipes.py new file mode 100644 index 00000000..d7db2716 --- /dev/null +++ b/testing/tests/pipes/test_pipes.py @@ -0,0 +1,32 @@ +# -*- coding: utf-8 -*- +# test_pipes.py +# Copyright (C) 2017 LEAP +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see <http://www.gnu.org/licenses/>. +""" +Tests for streaming components. +""" +from twisted.trial import unittest +from leap.soledad.client._pipes import TruncatedTailPipe + + +class TruncatedTailTestCase(unittest.TestCase): + + def test_tail_truncating_pipe(self): + pipe = TruncatedTailPipe(tail_size=20) + payload = 'A' * 100 + 'B' * 20 + for data in payload: + pipe.write(data) + result = pipe.close() + assert result.getvalue() == 'A' * 100 |