From 6dee94a429e456888e692910f7bb0d82b1c06d60 Mon Sep 17 00:00:00 2001
From: Rob Browning <rlb@defaultvalue.org>
Date: Fri, 24 Jun 2022 16:36:04 -0500
Subject: [PATCH] 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>
---
 lib/bup/hlinkdb.py | 96 ++++++++++++++++------------------------------
 1 file changed, 34 insertions(+), 62 deletions(-)

diff --git a/lib/bup/hlinkdb.py b/lib/bup/hlinkdb.py
index 102ee3b..f7e5d72 100644
--- 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
-- 
2.39.2