From f9e292577be29a9490f62420c6bbed44dfc30a09 Mon Sep 17 00:00:00 2001
From: mbearup <mbearup@microsoft.com>
Date: Tue, 26 Feb 2019 14:14:24 -0800
Subject: [PATCH] Add fixes for handling swap file

---
 .../daemon/resourcedisk/default.py            | 35 ++++++++++++++-----
 .../daemon/resourcedisk/freebsd.py            |  6 ++--
 tests/distro/test_resourceDisk.py             | 35 +++++++++++++++++--
 3 files changed, 61 insertions(+), 15 deletions(-)

diff --git a/azurelinuxagent/daemon/resourcedisk/default.py b/azurelinuxagent/daemon/resourcedisk/default.py
index 0f0925d..321c7bc 100644
--- a/azurelinuxagent/daemon/resourcedisk/default.py
+++ b/azurelinuxagent/daemon/resourcedisk/default.py
@@ -16,6 +16,7 @@
 #
 
 import os
+import stat
 import re
 import sys
 import threading
@@ -245,16 +246,28 @@ def get_mount_string(mount_options, partition, mount_point):
         else:
             return 'mount {0} {1}'.format(partition, mount_point)
 
+    @staticmethod
+    def check_existing_swap_file(swapfile, size):
+        swaplist = shellutil.run_get_output("swapon -s")[1]
+
+        if swapfile in swaplist and os.path.isfile(swapfile) and os.path.getsize(swapfile) == size:
+            logger.info("Swap already enabled")
+            # restrict access to owner (remove all access from group, others)
+            swapfile_mode = os.stat(swapfile).st_mode
+            if swapfile_mode & (stat.S_IRWXG | stat.S_IRWXO):
+                swapfile_mode = swapfile_mode & ~(stat.S_IRWXG | stat.S_IRWXO)
+                logger.info("Changing mode of {0} to {1:o}".format(swapfile, swapfile_mode))
+                os.chmod(swapfile, swapfile_mode)
+            return True
+
+        return False
+
     def create_swap_space(self, mount_point, size_mb):
         size_kb = size_mb * 1024
         size = size_kb * 1024
         swapfile = os.path.join(mount_point, 'swapfile')
-        swaplist = shellutil.run_get_output("swapon -s")[1]
 
-        if swapfile in swaplist \
-                and os.path.isfile(swapfile) \
-                and os.path.getsize(swapfile) == size:
-            logger.info("Swap already enabled")
+        if self.check_existing_swap_file(swapfile, size):
             return
 
         if os.path.isfile(swapfile) and os.path.getsize(swapfile) != size:
@@ -305,13 +318,18 @@ def mkfile(self, filename, nbytes):
                 # Probable errors:
                 #  - OSError: Seen on Cygwin, libc notimpl?
                 #  - AttributeError: What if someone runs this under...
+                fd = None
+
                 try:
-                    with open(filename, 'w') as f:
-                        os.posix_fallocate(f.fileno(), 0, nbytes)
-                        return 0
+                    fd = os.open(filename, os.O_CREAT | os.O_WRONLY, stat.S_IRUSR | stat.S_IWUSR)
+                    os.posix_fallocate(fd, 0, nbytes)
+                    return 0
                 except:
                     # Not confident with this thing, just keep trying...
                     pass
+                finally:
+                    if fd is not None:
+                        os.close(fd)
 
             # fallocate command
             ret = shellutil.run(
@@ -340,3 +358,4 @@ def mkfile(self, filename, nbytes):
             logger.error("dd unsuccessful")
 
         return ret
+
diff --git a/azurelinuxagent/daemon/resourcedisk/freebsd.py b/azurelinuxagent/daemon/resourcedisk/freebsd.py
index ece166b..f1a5d91 100644
--- a/azurelinuxagent/daemon/resourcedisk/freebsd.py
+++ b/azurelinuxagent/daemon/resourcedisk/freebsd.py
@@ -130,10 +130,7 @@ def create_swap_space(self, mount_point, size_mb):
         swapfile = os.path.join(mount_point, 'swapfile')
         swaplist = shellutil.run_get_output("swapctl -l")[1]
 
-        if swapfile in swaplist \
-                and os.path.isfile(swapfile) \
-                and os.path.getsize(swapfile) == size:
-            logger.info("Swap already enabled")
+        if self.check_existing_swap_file(swapfile, size):
             return
 
         if os.path.isfile(swapfile) and os.path.getsize(swapfile) != size:
@@ -161,3 +158,4 @@ def create_swap_space(self, mount_point, size_mb):
             if shellutil.run("swapon /dev/{0}".format(mddevice)):
                 raise ResourceDiskError("/dev/{0}".format(mddevice))
             logger.info("Enabled {0}KB of swap at /dev/{1} ({2})".format(size_kb, mddevice, swapfile))
+
diff --git a/tests/distro/test_resourceDisk.py b/tests/distro/test_resourceDisk.py
index d2ce6e1..9b27ade 100644
--- a/tests/distro/test_resourceDisk.py
+++ b/tests/distro/test_resourceDisk.py
@@ -18,12 +18,11 @@
 # http://msdn.microsoft.com/en-us/library/cc227282%28PROT.10%29.aspx
 # http://msdn.microsoft.com/en-us/library/cc227259%28PROT.13%29.aspx
 
-import sys
+import stat
 from azurelinuxagent.common.utils import shellutil
 from azurelinuxagent.daemon.resourcedisk import get_resourcedisk_handler
 from tests.tools import *
 
-
 class TestResourceDisk(AgentTestCase):
     def test_mkfile(self):
         # setup
@@ -38,6 +37,10 @@ def test_mkfile(self):
         # assert
         assert os.path.exists(test_file)
 
+        # only the owner should have access
+        mode = os.stat(test_file).st_mode & (stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
+        assert mode == stat.S_IRUSR | stat.S_IWUSR
+
         # cleanup
         os.remove(test_file)
 
@@ -83,7 +86,6 @@ def test_mkfile_xfs_fs(self):
             assert run_patch.call_count == 1
             assert "dd if" in run_patch.call_args_list[0][0][0]
 
-
     def test_change_partition_type(self):
         resource_handler = get_resourcedisk_handler()
         # test when sfdisk --part-type does not exist
@@ -105,6 +107,33 @@ def test_change_partition_type(self):
             assert run_patch.call_count == 1
             assert "sfdisk --part-type" in run_patch.call_args_list[0][0][0]
 
+    def test_check_existing_swap_file(self):
+        test_file = os.path.join(self.tmp_dir, 'test_swap_file')
+        file_size = 1024 * 128
+        if os.path.exists(test_file):
+            os.remove(test_file)
+
+        with open(test_file, "wb") as file:
+            file.write(bytes(file_size))
+
+        os.chmod(test_file,  stat.S_ISUID | stat.S_ISGID | stat.S_IRUSR | stat.S_IWUSR | stat.S_IRWXG | stat.S_IRWXO)  # 0o6677
+
+        def swap_on(_):   # mimic the output of "swapon -s"
+            return [
+                "Filename   Type        Size      Used  Priority",
+                "{0}        partition	16498684  0     -2".format(test_file)
+            ]
+
+        with patch.object(shellutil, "run_get_output", side_effect=swap_on):
+            get_resourcedisk_handler().check_existing_swap_file(test_file, file_size)
+
+        # it should remove access from group, others
+        mode = os.stat(test_file).st_mode & (stat.S_ISUID | stat.S_ISGID | stat.S_IRWXU | stat.S_IWUSR | stat.S_IRWXG | stat.S_IRWXO)  # 0o6777
+        assert mode == stat.S_ISUID | stat.S_ISGID | stat.S_IRUSR | stat.S_IWUSR  # 0o6600
+
+        os.remove(test_file)
+
 
 if __name__ == '__main__':
     unittest.main()
+
