SERVER-54732 Added validate cachedir tool

3 files changed, 207 insertions, 3 deletions

diff --git a/SConstruct b/SConstruct
index 9e4cad24503..638d28ceb73 100644
--- a/SConstruct
+++ b/SConstruct
@@ -469,6 +469,12 @@ add_option('cache-dir',
     help='Specify the directory to use for caching objects if --cache is in use',
 )
 
+add_option('cache-signature-mode',
+    choices=['none', 'validate'],
+    default="none",
+    help='Extra check to validate integrity of cache files after pulling from cache',
+)
+
 add_option("cxx-std",
     choices=["17"],
     default="17",
@@ -1202,6 +1208,13 @@ if get_option('build-tools') == 'next':
 env = Environment(variables=env_vars, **envDict)
 del envDict
 
+if get_option('cache-signature-mode') == 'validate':
+    validate_cache_dir = Tool('validate_cache_dir')
+    if validate_cache_dir.exists(env):
+        validate_cache_dir(env)
+    else:
+        env.FatalError("Failed to enable validate_cache_dir tool.")
+
 # Only print the spinner if stdout is a tty
 if sys.stdout.isatty():
     Progress(['-\r', '\\\r', '|\r', '/\r'], interval=50)
diff --git a/buildscripts/scons_cache_prune.py b/buildscripts/scons_cache_prune.py
index f2a28b2c5a4..ce327f35da0 100644
--- a/buildscripts/scons_cache_prune.py
+++ b/buildscripts/scons_cache_prune.py
@@ -24,6 +24,19 @@ GIGBYTES = 1024 * 1024 * 1024
 CacheItem = collections.namedtuple("CacheContents", ["path", "time", "size"])
 
 
+def get_cachefile_size(file_path):
+    """Get the size of the cachefile."""
+
+    if file_path.endswith('.cksum'):
+        size = 0
+        for cksum_path in os.listdir(file_path):
+            cksum_path = os.path.join(file_path, cksum_path)
+            size += os.stat(cksum_path).st_size
+    else:
+        size = os.stat(file_path).st_size
+    return size
+
+
 def collect_cache_contents(cache_path):
     """Collect the cache contents."""
     # map folder names to timestamps
@@ -37,15 +50,19 @@ def collect_cache_contents(cache_path):
         if os.path.isdir(path):
             for file_name in os.listdir(path):
                 file_path = os.path.join(path, file_name)
-                if os.path.isdir(file_path):
+                # Cache prune script is allowing only directories with this extension
+                # which comes from the validate_cache_dir.py tool in scons, it must match
+                # the extension set in that file.
+                if os.path.isdir(file_path) and not file_path.endswith('.cksum'):
                     LOGGER.warning(
                         "cache item %s is a directory and not a file. "
                         "The cache may be corrupt.", file_path)
                     continue
 
                 try:
+
                     item = CacheItem(path=file_path, time=os.stat(file_path).st_atime,
-                                     size=os.stat(file_path).st_size)
+                                     size=get_cachefile_size(file_path))
 
                     total += item.size
 
@@ -90,7 +107,10 @@ def prune_cache(cache_path, cache_size_gb, clean_ratio):
                 LOGGER.warning("Unable to rename %s : %s", cache_item, err)
             else:
                 try:
-                    os.remove(to_remove)
+                    if os.path.isdir(to_remove):
+                        shutil.rmtree(to_remove)
+                    else:
+                        os.remove(to_remove)
                     total_size -= cache_item.size
                 except Exception as err:  # pylint: disable=broad-except
                     # this should not happen, but who knows?
diff --git a/site_scons/site_tools/validate_cache_dir.py b/site_scons/site_tools/validate_cache_dir.py
new file mode 100644
index 00000000000..1f13ca6befb
--- /dev/null
+++ b/site_scons/site_tools/validate_cache_dir.py
@@ -0,0 +1,171 @@
+# Copyright 2021 MongoDB Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files (the
+# "Software"), to deal in the Software without restriction, including
+# without limitation the rights to use, copy, modify, merge, publish,
+# distribute, sublicense, and/or sell copies of the Software, and to
+# permit persons to whom the Software is furnished to do so, subject to
+# the following conditions:
+#
+# The above copyright notice and this permission notice shall be included
+# in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
+# KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+
+import os
+import pathlib
+import shutil
+
+import SCons
+
+cache_debug_suffix = " (target: %s, cachefile: %s) "
+
+class InvalidChecksum(SCons.Errors.BuildError):
+    def __init__(self, src, dst, reason):
+        self.message = f"ERROR: md5 checksum {reason} for {src} ({dst})"
+
+    def __str__(self):
+        return self.message
+
+class CacheTransferFailed(SCons.Errors.BuildError):
+    def __init__(self, src, dst, reason):
+        self.message = f"ERROR: cachedir transfer {reason} while transfering {src} to {dst}"
+
+    def __str__(self):
+        return self.message
+
+class UnsupportedError(SCons.Errors.BuildError):
+    def __init__(self, class_name, feature):
+        self.message = f"{class_name} does not support {feature}"
+
+    def __str__(self):
+        return self.message
+
+class CacheDirValidate(SCons.CacheDir.CacheDir):
+
+    @staticmethod
+    def get_ext():
+        # Cache prune script is allowing only directories with this extension
+        # if this is changed, cache prune script should also be updated.
+        return '.cksum'
+
+    @staticmethod
+    def get_file_contents_path(path):
+        return str(pathlib.Path(path + CacheDirValidate.get_ext()) / pathlib.Path(path).name)
+
+    @staticmethod
+    def get_hash_path(path):
+        return str(pathlib.Path(path).parent / 'content_hash')
+
+    @classmethod
+    def copy_from_cache(cls, env, src, dst):
+
+        if not str(pathlib.Path(src).parent).endswith(cls.get_ext()):
+            return super().copy_from_cache(env, src, dst)
+
+        if env.cache_timestamp_newer:
+            raise UnsupportedError(cls.__name__, "timestamp-newer")
+
+        csig = None
+        invalid_cksum = InvalidChecksum(cls.get_hash_path(src), dst, "failed to read hash file")
+        try:
+            with open(cls.get_hash_path(src), 'rb') as f_out:
+                csig = f_out.read().decode().strip()
+        except OSError as ex:
+            raise invalid_cksum from ex
+        finally:
+            if not csig:
+                raise invalid_cksum from ex
+
+        try:
+            shutil.copy2(src, dst)
+        except OSError as ex:
+            raise CacheTransferFailed(src, dst, "failed to copy from cache") from ex
+
+        new_csig = SCons.Util.MD5filesignature(dst,
+            chunksize=SCons.Node.FS.File.md5_chunksize*1024)
+
+        if csig != new_csig:
+            raise InvalidChecksum(
+                cls.get_hash_path(src), dst, f"checksums don't match {csig} != {new_csig}")
+
+    @classmethod
+    def copy_to_cache(cls, env, src, dst):
+
+        # dst is bsig/file from cachepath method, so
+        # we make sure to make the bsig dir first
+        os.makedirs(pathlib.Path(dst).parent, exist_ok=True)
+
+        try:
+            shutil.copy2(src, dst)
+        except OSError as ex:
+            raise CacheTransferFailed(src, dst, "failed to copy to cache") from ex
+
+        try:
+            with open(cls.get_hash_path(dst), 'w') as f_out:
+                f_out.write(env.File(src).get_content_hash())
+        except OSError as ex:
+            raise CacheTransferFailed(src, dst, "failed to create hash file") from ex
+
+    def retrieve(self, node):
+        try:
+            return super().retrieve(node)
+        except InvalidChecksum as ex:
+            self.print_cache_issue(node, str(ex))
+            self.clean_bad_cachefile(node)
+            return False
+        except (UnsupportedError, CacheTransferFailed) as ex:
+            self.print_cache_issue(node, str(ex))
+            return False
+
+    def push(self, node):
+        try:
+            return super().push(node)
+        except CacheTransferFailed as ex:
+            self.print_cache_issue(node, str(ex))
+            return False
+
+    def print_cache_issue(self, node, msg):
+
+        cksum_dir = pathlib.Path(self.cachepath(node)[1]).parent
+        print(msg)
+        self.CacheDebug(msg + cache_debug_suffix, node, cksum_dir)
+
+    def clean_bad_cachefile(self, node):
+
+        cksum_dir = pathlib.Path(self.cachepath(node)[1]).parent
+        if cksum_dir.is_dir():
+            rm_path = f"{cksum_dir}.{SCons.CacheDir.cache_tmp_uuid}.del"
+            cksum_dir.replace(rm_path)
+            shutil.rmtree(rm_path)
+
+            clean_msg = f"Removed bad cachefile {cksum_dir} from cache."
+            print(clean_msg)
+            self.CacheDebug(clean_msg + cache_debug_suffix, node, cksum_dir)
+
+    def get_cachedir_csig(self, node):
+        cachedir, cachefile = self.cachepath(node)
+        if cachefile and os.path.exists(cachefile):
+            with open(self.get_hash_path(cachefile), 'rb') as f_out:
+                return f_out.read().decode()
+
+    def cachepath(self, node):
+        dir, path = super().cachepath(node)
+        if node.fs.exists(path):
+            return dir, path
+        return dir, self.get_file_contents_path(path)
+
+def exists(env):
+    return True
+
+def generate(env):
+    if not env.get('CACHEDIR_CLASS'):
+        env['CACHEDIR_CLASS'] = CacheDirValidate