atomic_ofstream: fix follow_symlinks fallback and default file modeHEADmaster

Handle OSError from mkstemp for (default) follow_symlinks mode,
not following the symlink if necessary (the target's parent may
not exist or may be readonly). This restores the fallback
behavior that existed before the introduction of mkstemp in
commit de19f3a7215d.

Handle missing _file and _tmp_name attributes during close.
Also set the default file mode respecting umask if a previous
file does not exist, which fixes the mode of CONTENTS files
since mkstemp.

Fixes: de19f3a7215d ("atomic_ofstream: Use mkstemp rather than getpid (pid namespace safety)")
Signed-off-by: Zac Medico <zmedico@gentoo.org>

3 files changed, 116 insertions, 15 deletions

diff --git a/lib/portage/tests/util/meson.build b/lib/portage/tests/util/meson.build
index 010dfa784..7f4db871f 100644
--- a/lib/portage/tests/util/meson.build
+++ b/lib/portage/tests/util/meson.build
@@ -1,5 +1,6 @@
 py.install_sources(
     [
+        'test_atomic_ofstream.py',
         'test_checksum.py',
         'test_digraph.py',
         'test_file_copier.py',
diff --git a/lib/portage/tests/util/test_atomic_ofstream.py b/lib/portage/tests/util/test_atomic_ofstream.py
new file mode 100644
index 000000000..bbaf0f1b0
--- /dev/null
+++ b/lib/portage/tests/util/test_atomic_ofstream.py
@@ -0,0 +1,85 @@
+# Copyright 2024 Gentoo Authors
+# Distributed under the terms of the GNU General Public License v2
+
+import errno
+import os
+import stat
+import tempfile
+
+from portage.tests import TestCase
+from portage.util import atomic_ofstream
+
+
+class AtomicOFStreamTestCase(TestCase):
+    def test_enospc_rollback(self):
+        file_name = "foo"
+        start_dir = os.getcwd()
+        with tempfile.TemporaryDirectory() as tempdir:
+            try:
+                os.chdir(tempdir)
+                with self.assertRaises(OSError):
+                    with atomic_ofstream(file_name) as f:
+                        f.write("hello")
+                        raise OSError(errno.ENOSPC, "No space left on device")
+                self.assertFalse(os.path.exists(file_name))
+                self.assertEqual(os.listdir(tempdir), [])
+            finally:
+                os.chdir(start_dir)
+
+    def test_open_failure(self):
+        file_name = "bad/path"
+        start_dir = os.getcwd()
+        with tempfile.TemporaryDirectory() as tempdir:
+            try:
+                os.chdir(tempdir)
+                with self.assertRaises(OSError):
+                    with atomic_ofstream(file_name):
+                        pass
+                self.assertEqual(os.listdir(tempdir), [])
+            finally:
+                os.chdir(start_dir)
+
+    def test_broken_symlink(self):
+        content = "foo"
+        broken_symlink = "symlink"
+        symlink_targets = (("foo/bar/baz", False), ("baz", True))
+        start_dir = os.getcwd()
+        for symlink_target, can_follow in symlink_targets:
+            with tempfile.TemporaryDirectory() as tempdir:
+                try:
+                    os.chdir(tempdir)
+                    with open(broken_symlink, "w") as f:
+                        default_file_mode = stat.S_IMODE(os.fstat(f.fileno()).st_mode)
+                        os.unlink(broken_symlink)
+                    os.symlink(symlink_target, broken_symlink)
+                    with atomic_ofstream(broken_symlink) as f:
+                        f.write(content)
+                    with open(broken_symlink) as f:
+                        self.assertEqual(f.read(), content)
+                    self.assertEqual(os.path.islink(broken_symlink), can_follow)
+                    self.assertEqual(
+                        stat.S_IMODE(os.stat(broken_symlink).st_mode), default_file_mode
+                    )
+                finally:
+                    os.chdir(start_dir)
+
+    def test_preserves_mode(self):
+        file_name = "foo"
+        file_mode = 0o604
+        start_dir = os.getcwd()
+        with tempfile.TemporaryDirectory() as tempdir:
+            try:
+                os.chdir(tempdir)
+                with open(file_name, "wb"):
+                    pass
+                self.assertNotEqual(stat.S_IMODE(os.stat(file_name).st_mode), file_mode)
+                os.chmod(file_name, file_mode)
+                st_before = os.stat(file_name)
+                self.assertEqual(stat.S_IMODE(st_before.st_mode), file_mode)
+                with atomic_ofstream(file_name):
+                    pass
+                st_after = os.stat(file_name)
+                self.assertNotEqual(st_before.st_ino, st_after.st_ino)
+                self.assertEqual(stat.S_IMODE(st_after.st_mode), file_mode)
+            finally:
+                os.chdir(start_dir)
diff --git a/lib/portage/util/__init__.py b/lib/portage/util/__init__.py
index f338f274a..50fa1680f 100644
--- a/lib/portage/util/__init__.py
+++ b/lib/portage/util/__init__.py
@@ -1434,7 +1434,7 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
     error occurs)."""
 
     def __init__(self, filename, mode="w", follow_links=True, **kargs):
-        """Opens a temporary filename.pid in the same directory as filename."""
+        """Opens a temporary file in the same directory as filename."""
         ObjectProxy.__init__(self)
         object.__setattr__(self, "_aborted", False)
         if "b" in mode:
@@ -1447,10 +1447,11 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
         if follow_links:
             canonical_path = os.path.realpath(filename)
             object.__setattr__(self, "_real_name", canonical_path)
-            parent, basename = os.path.split(canonical_path)
-            fd, tmp_name = tempfile.mkstemp(prefix=basename, dir=parent)
-            object.__setattr__(self, "_tmp_name", tmp_name)
+            fd = None
             try:
+                parent, basename = os.path.split(canonical_path)
+                fd, tmp_name = tempfile.mkstemp(prefix=basename, dir=parent)
+                object.__setattr__(self, "_tmp_name", tmp_name)
                 object.__setattr__(
                     self,
                     "_file",
@@ -1462,7 +1463,8 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
                 )
                 return
             except OSError:
-                os.close(fd)
+                if fd is not None:
+                    os.close(fd)
                 if canonical_path == filename:
                     raise
                 # Ignore this error, since it's irrelevant
@@ -1470,10 +1472,11 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
                 # new error if necessary.
 
         object.__setattr__(self, "_real_name", filename)
-        parent, basename = os.path.split(filename)
-        fd, tmp_name = tempfile.mkstemp(prefix=basename, dir=parent)
-        object.__setattr__(self, "_tmp_name", tmp_name)
+        fd = None
         try:
+            parent, basename = os.path.split(filename)
+            fd, tmp_name = tempfile.mkstemp(prefix=basename, dir=parent)
+            object.__setattr__(self, "_tmp_name", tmp_name)
             object.__setattr__(
                 self,
                 "_file",
@@ -1484,7 +1487,8 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
                 ),
             )
         except OSError:
-            os.close(fd)
+            if fd is not None:
+                os.close(fd)
             raise
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1505,15 +1509,26 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
         """Closes the temporary file, copies permissions (if possible),
         and performs the atomic replacement via os.rename().  If the abort()
         method has been called, then the temp file is closed and removed."""
-        f = object.__getattribute__(self, "_file")
-        tmp_name = object.__getattribute__(self, "_tmp_name")
-        real_name = object.__getattribute__(self, "_real_name")
-        if not f.closed:
+        try:
+            f = object.__getattribute__(self, "_file")
+        except AttributeError:
+            f = None
+        if f is not None and not f.closed:
+            real_name = object.__getattribute__(self, "_real_name")
+            tmp_name = object.__getattribute__(self, "_tmp_name")
             try:
                 f.close()
                 if not object.__getattribute__(self, "_aborted"):
                     try:
-                        apply_stat_permissions(tmp_name, os.stat(real_name))
+                        try:
+                            st = os.stat(real_name)
+                        except FileNotFoundError:
+                            umask_test_file = f"{tmp_name}_umask_test"
+                            with open(umask_test_file, "w") as f:
+                                st = os.fstat(f.fileno())
+                                os.unlink(umask_test_file)
+
+                        apply_stat_permissions(tmp_name, st)
                     except OperationNotPermitted:
                         pass
                     except FileNotFound:
@@ -1529,7 +1544,7 @@ class atomic_ofstream(AbstractContextManager, ObjectProxy):
                 # even if an exception is raised.
                 try:
                     os.unlink(tmp_name)
-                except OSError as oe:
+                except OSError:
                     pass
 
     def abort(self):