Allow re-index of search index for larger accounts

- Issue #560
- Actual problem was in whoosh that closes a file object too early

1 files changed, 69 insertions, 9 deletions

diff --git a/service/pixelated/support/encrypted_file_storage.py b/service/pixelated/support/encrypted_file_storage.py
index 567a348a..a1dbffa8 100644
--- a/service/pixelated/support/encrypted_file_storage.py
+++ b/service/pixelated/support/encrypted_file_storage.py
@@ -14,35 +14,74 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with Pixelated. If not, see <http://www.gnu.org/licenses/>.
 
+from __future__ import with_statement
+
+import hmac
 import io
+import os
 from hashlib import sha256
 
-import os
-import hmac
-from whoosh.filedb.filestore import FileStorage
-from whoosh.filedb.structfile import StructFile, BufferFile
-from leap.soledad.client.crypto import encrypt_sym
 from leap.soledad.client.crypto import decrypt_sym
-from leap.soledad.common.crypto import EncryptionMethods
+from leap.soledad.client.crypto import encrypt_sym
+from whoosh.filedb.filestore import FileStorage
+from whoosh.filedb.structfile import BufferFile, StructFile
 from whoosh.util import random_name
 
 
+class DelayedCloseBytesIO(io.BytesIO):
+    def __init__(self, name):
+        super(DelayedCloseBytesIO, self).__init__()
+        self._name = name
+        self.shouldClose = False
+
+    def close(self):
+
+        self.shouldClose = True
+
+    def explicit_close(self):
+        super(DelayedCloseBytesIO, self).close()
+
+
+class DelayedCloseStructFile(StructFile):
+    def __init__(self, fileobj, name=None, onclose=None):
+        super(DelayedCloseStructFile, self).__init__(fileobj, name, onclose)
+
+    def close(self):
+        """Closes the wrapped file.
+        """
+
+        if self.is_closed:
+            raise Exception("This file is already closed")
+        if self.onclose:
+            self.onclose(self)
+        if hasattr(self.file, "explicit_close"):
+            self.file.explicit_close()
+        self.is_closed = True
+
+
 class EncryptedFileStorage(FileStorage):
     def __init__(self, path, masterkey=None):
+        FileStorage.__init__(self, path, supports_mmap=False)
         self.masterkey = masterkey[:32]
         self.signkey = masterkey[32:]
         self._tmp_storage = self.temp_storage
         self.length_cache = {}
-        FileStorage.__init__(self, path, supports_mmap=False)
+        self._open_files = {}
 
     def open_file(self, name, **kwargs):
-        return self._open_encrypted_file(name, onclose=self._encrypt_index_on_close(name))
+        return self._open_encrypted_file(name)
 
     def create_file(self, name, excl=False, mode="w+b", **kwargs):
-        f = StructFile(io.BytesIO(), name=name, onclose=self._encrypt_index_on_close(name))
+        f = DelayedCloseStructFile(DelayedCloseBytesIO(name), name=name, onclose=self._encrypt_index_on_close(name))
         f.is_real = False
+        self._open_files[name] = f
         return f
 
+    def delete_file(self, name):
+        super(EncryptedFileStorage, self).delete_file(name)
+        if name in self._open_files:
+            del self._open_files[name]
+
     def temp_storage(self, name=None):
         name = name or "%s.tmp" % random_name()
         path = os.path.join(self.folder, name)
@@ -78,10 +117,31 @@ class EncryptedFileStorage(FileStorage):
             encrypted_content = self.encrypt(content)
             with open(self._fpath(name), 'w+b') as f:
                 f.write(encrypted_content)
+
         return wrapper
 
     def _open_encrypted_file(self, name, onclose=lambda x: None):
+        if not self.file_exists(name):
+            if name in self._open_files:
+                f = self._open_files[name]
+                if not f.is_closed:
+                    state = 'closed' if f.file.shouldClose else 'open'
+                    if state == 'closed':
+                        self._store_file(name, f.file.getvalue())
+                        f.close()
+                        del self._open_files[name]
+            else:
+                raise NameError(name)
         file_content = open(self._fpath(name), "rb").read()
         decrypted = self.decrypt(file_content)
         self.length_cache[name] = (len(decrypted), sha256(decrypted).digest())
         return BufferFile(buffer(decrypted), name=name, onclose=onclose)
+
+    def _store_file(self, name, content):
+        try:
+            encrypted_content = self.encrypt(content)
+            with open(self._fpath(name), 'w+b') as f:
+                f.write(encrypted_content)
+        except Exception, e:
+            print e
+            raise