codeql/python/extractor/semmle/cache.py

from semmle.util import makedirs
import os
from collections import deque
from functools import total_ordering

'''
Least Recently Written Disk-based Cache

Implements a LRW disk cache for trap files and similar.
This cache relies on the following properties which *must* hold.

Only one value can ever be associated with a key.
Keys should be ascii strings and cannot start with '$' or include any file or path separator characters.
Values should be byte strings (with any contents).

The cache is robust against arbitrary levels of concurrency.

'''


MAX_GENERATIONS = 50
MAX_FILES_PER_GENERATION = 200

def encode_keys(keys):
    'Convert a collection of keys to a byte string'
    return '\n'.join(keys).encode("ascii")

def decode_keys(data):
    'Convert a byte string into a set of keys'
    return set(data.decode("ascii").split('\n'))

@total_ordering
class Generation(object):

    def __init__(self, cachedir, age):
        self.cachedir = os.path.join(cachedir, str(age))
        self.age = age
        if not os.path.exists(self.cachedir):
            makedirs(self.cachedir)
        try:
            with open(os.path.join(self.cachedir, "$keys"), 'rb') as fd:
                self.keys = decode_keys(fd.read())
            self.full = True
        except Exception:
            self.keys = set()
            if os.path.isdir(self.cachedir):
                #Directory exists, but cannot read "$keys", so this is a non-full generation
                self.full = False
            else:
                self.full = True

    def get(self, key):
        if self.full and key not in self.keys:
            return None
        try:
            with open(os.path.join(self.cachedir, key), 'rb') as fd:
                return fd.read()
        except Exception:
            return None

    def set(self, key, value):
        '''Returns true if it should be able to store (key, value) even if in fact it can't.
        This means that this method will return True if the generation is not full.'''
        if self.full:
            return False
        if os.path.exists(os.path.join(self.cachedir, "$keys")):
            self.full = True
            try:
                with open(os.path.join(self.cachedir, "$keys"), 'rb') as fd:
                    self.keys = decode_keys(fd.read())
            except Exception:
                self.keys = set()
            return False
        self._try_atomic_write_file(key, value)
        if len(self._list_files()) >= MAX_FILES_PER_GENERATION:
            self.full = True
            self._write_keys()
        return True

    def _list_files(self):
        try:
            return os.listdir(self.cachedir)
        except Exception:
            #This probably means the directory has been deleted
            return []

    def _write_keys(self):
        keys = self._list_files()
        self._try_atomic_write_file("$keys", encode_keys(keys))
        self.keys = set(keys)

    def _try_atomic_write_file(self, name, contents):
        fullname = os.path.join(self.cachedir, name)
        tmpname = os.path.join(self.cachedir, '$%d%s' % (os.getpid(), name))
        try:
            with open(tmpname, 'wb') as tmp:
                tmp.write(contents)
            os.rename(tmpname, fullname)
        except Exception:
            #Failed for some reason. The folder may have been deleted, or on Windows, the file may already exist.
            #Attempt to tidy up
            if os.path.exists(tmpname):
                try:
                    os.remove(tmpname)
                except Exception:
                    #Give up :(
                    pass

    def clear(self):
        try:
            filenames = os.listdir(self.cachedir)
        except Exception:
            #Can't do anything
            return
        for filename in filenames:
            try:
                os.remove(os.path.join(self.cachedir, filename))
            except Exception:
                # Can't delete. Maybe another process has deleted it or it is open (on Windows)
                pass
        try:
            os.rmdir(self.cachedir)
        except Exception:
            # Can't delete
            pass

    def __lt__(self, other):
        #Smaller numbers are older
        return self.age > other.age

class Cache(object):

    cache_of_caches = {}

    def __init__(self, cachedir, verbose=False):
        self.cachedir = cachedir
        self.verbose = verbose
        self.generations = []
        if not os.path.exists(cachedir):
            makedirs(cachedir)
        generations = []
        for gen in os.listdir(self.cachedir):
            try:
                age = int(gen)
                generations.append(Generation(self.cachedir, age))
            except Exception:
                #gen might not be an int, or it may have been deleted
                pass
        if generations:
            generations.sort()
        else:
            generations = [Generation(self.cachedir, 1)]
        self.generations = deque(generations)
        while len(self.generations) > MAX_GENERATIONS:
            self.generations.pop().clear()

    def set(self, key, value):
        '''Add this (key, value) pair to the cache. keys should not start with '$' or include file or path separators.
        Either adds the (key, value) atomically or does nothing. Partial keys or values are never visible.
        '''
        try:
            while not self.generations[0].set(key, value):
                self.generations.appendleft(Generation(self.cachedir, self.generations[0].age+1))
                if len(self.generations) > MAX_GENERATIONS:
                    self.generations.pop().clear()
        except Exception as ex:
            #Its OK to fail but we must never raise
            if self.verbose:
                try:
                    print ("Exception setting cache key '%s': %s" % (key, ex))
                except Exception:
                    # Just in case
                    pass

    def get(self, key):
        if key is None:
            return None
        try:
            for gen in self.generations:
                res = gen.get(key)
                if res is not None:
                    return res
        except Exception as ex:
            if self.verbose:
                try:
                    print ("Exception getting cache key '%s': %s" % (key, ex))
                except Exception:
                    # Just in case
                    pass
        return None

    @staticmethod
    def for_directory(cachedir, verbose):
        '''Caches are relatively expensive objects, so we cache them.'''
        if (cachedir, verbose) not in Cache.cache_of_caches:
            Cache.cache_of_caches[(cachedir, verbose)] = Cache(cachedir, verbose)
        return Cache.cache_of_caches[(cachedir, verbose)]