hlinkdb: respect umask/sgid/etc. when creating new db

Use atomically_replaced_file to ensure that the new db respects the
current umask, any directory sgid bit, etc., and use an ExitStack to
try to make sure all the corner cases are covered.

Signed-off-by: Rob Browning <rlb@defaultvalue.org>
Tested-by: Rob Browning <rlb@defaultvalue.org>

diff --git a/lib/bup/hlinkdb.py b/lib/bup/hlinkdb.py
index 102ee3bd26096d150d4b18312e186e014bc710e9..f7e5d72153b6d72e7e5934296eb9169bfe4a12d8 100644 (file)
--- a/lib/bup/hlinkdb.py
+++ b/lib/bup/hlinkdb.py
@@ -1,11 +1,17 @@
 
-import errno, os, pickle, tempfile
+from contextlib import ExitStack
+import os, pickle
 
-from bup.compat import pending_raise
+from bup.helpers import atomically_replaced_file, unlink
 
 
-def pickle_load(f):
-    return pickle.load(f, encoding='bytes')
+def pickle_load(filename):
+    try:
+        f = open(filename, 'rb')
+    except FileNotFoundError:
+        return None
+    with f:
+        return pickle.load(f, encoding='bytes')
 
 
 class Error(Exception):
@@ -14,28 +20,13 @@ class Error(Exception):
 class HLinkDB:
     def __init__(self, filename):
         self.closed = False
+        self._cleanup = ExitStack()
+        self._filename = filename
+        self._pending_save = None
         # Map a "dev:ino" node to a list of paths associated with that node.
-        self._node_paths = {}
-        # Map a path to a "dev:ino" node.
+        self._node_paths = pickle_load(filename) or {}
+        # Map a path to a "dev:ino" node (a reverse hard link index).
         self._path_node = {}
-        self._filename = filename
-        self._save_prepared = None
-        self._tmpname = None
-        f = None
-        try:
-            f = open(filename, 'rb')
-        except IOError as e:
-            if e.errno == errno.ENOENT:
-                pass
-            else:
-                raise
-        if f:
-            try:
-                self._node_paths = pickle_load(f)
-            finally:
-                f.close()
-                f = None
-        # Set up the reverse hard link index.
         for node, paths in self._node_paths.items():
             for path in paths:
                 self._path_node[path] = node
@@ -43,59 +34,40 @@ class HLinkDB:
     def prepare_save(self):
         """ Commit all of the relevant data to disk.  Do as much work
         as possible without actually making the changes visible."""
-        if self._save_prepared:
+        if self._pending_save:
             raise Error('save of %r already in progress' % self._filename)
-        if self._node_paths:
-            (dir, name) = os.path.split(self._filename)
-            (ffd, self._tmpname) = tempfile.mkstemp(b'.tmp', name, dir)
-            try:
-                try:
-                    f = os.fdopen(ffd, 'wb', 65536)
-                except:
-                    os.close(ffd)
-                    raise
-                try:
+        with self._cleanup:
+            if self._node_paths:
+                dir, name = os.path.split(self._filename)
+                self._pending_save = atomically_replaced_file(self._filename,
+                                                              mode='wb',
+                                                              buffering=65536)
+                with self._cleanup.enter_context(self._pending_save) as f:
                     pickle.dump(self._node_paths, f, 2)
-                finally:
-                    f.close()
-                    f = None
-            except:
-                tmpname = self._tmpname
-                self._tmpname = None
-                os.unlink(tmpname)
-                raise
-        self._save_prepared = True
+            else: # No data
+                self._cleanup.callback(lambda: unlink(self._filename))
+            self._cleanup = self._cleanup.pop_all()
 
     def commit_save(self):
         self.closed = True
-        if not self._save_prepared:
+        if self._node_paths and not self._pending_save:
             raise Error('cannot commit save of %r; no save prepared'
                         % self._filename)
-        if self._tmpname:
-            os.rename(self._tmpname, self._filename)
-            self._tmpname = None
-        else: # No data -- delete _filename if it exists.
-            try:
-                os.unlink(self._filename)
-            except OSError as e:
-                if e.errno == errno.ENOENT:
-                    pass
-                else:
-                    raise
-        self._save_prepared = None
+        self._cleanup.close()
+        self._pending_save = None
 
     def abort_save(self):
         self.closed = True
-        if self._tmpname:
-            os.unlink(self._tmpname)
-            self._tmpname = None
+        with self._cleanup:
+            if self._pending_save:
+                self._pending_save.cancel()
+        self._pending_save = None
 
     def __enter__(self):
         return self
 
     def __exit__(self, type, value, traceback):
-        with pending_raise(value, rethrow=True):
-            self.abort_save()
+        self.abort_save()
 
     def __del__(self):
         assert self.closed