kitty-mirror/kitty/mime-cache.c

/*
 * mime-cache.c
 * Copyright (C) 2026 Kovid Goyal <kovid at kovidgoyal.net>
 *
 * Distributed under terms of the GPL3 license.
 */

// Fast lookup of MIME types for file names using the binary mime.cache files
// generated by update-mime-database from shared-mime-info. The lookup
// algorithm mirrors the one used by xdgmime (GLib/GIO). See
// https://specifications.freedesktop.org/shared-mime-info-spec/latest/ar01s02.html

#include "data-types.h"
#include "monotonic.h"

#ifdef __APPLE__

static PyObject *
mime_type_for_filename(PyObject *self UNUSED, PyObject *name UNUSED) {
    Py_RETURN_NONE;
}

static PyObject *
mime_type_aliases(PyObject *self UNUSED, PyObject *mime) {
    if (!PyUnicode_Check(mime)) {
        PyErr_SetString(PyExc_TypeError, "mime must be a string");
        return NULL;
    }
    return PyTuple_Pack(1, mime);
}

static PyObject *
clear_mime_cache_data(PyObject *self UNUSED, PyObject *args UNUSED) {
    Py_RETURN_NONE;
}

#else

#include <fcntl.h>
#include <fnmatch.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include "safe-wrappers.h"

#define CASE_SENSITIVE_FLAG 0x100
#define WEIGHT_MASK 0xff
#define MAX_MATCHES 16
#define MAX_SUFFIX_TREE_DEPTH 256
#define HEADER_SIZE 40
#define ALIAS_LIST_OFFSET_POS 4
#define LITERAL_LIST_OFFSET_POS 12
#define REVERSE_SUFFIX_TREE_OFFSET_POS 16
#define GLOB_LIST_OFFSET_POS 20

typedef struct {
    char *path;
    const uint8_t *data;
    size_t sz;
    bool exists;
    dev_t dev;
    ino_t ino;
    off_t size;
    struct timespec mtime;
} MimeCache;

typedef struct {
    const char *mime;
    uint32_t weight;
} Match;

typedef struct {
    Match items[MAX_MATCHES];
    unsigned count;
} Matches;

static struct {
    MimeCache *items;
    size_t count;
    bool loaded;
    monotonic_t last_check_at;
} caches = {0};

// Reading data {{{
static inline bool
read_u32(const MimeCache *c, size_t offset, uint32_t *ans) {
    if (offset > c->sz || c->sz - offset < 4) return false;
    const uint8_t *p = c->data + offset;
    *ans = ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | (uint32_t)p[3];
    return true;
}

static inline const char *
read_string(const MimeCache *c, uint32_t offset) {
    if (offset >= c->sz || !memchr(c->data + offset, 0, c->sz - offset)) return NULL;
    return (const char *)c->data + offset;
}

static inline bool
read_entry(const MimeCache *c, size_t offset, uint32_t *a, uint32_t *b, uint32_t *weight) {
    return read_u32(c, offset, a) && read_u32(c, offset + 4, b) && read_u32(c, offset + 8, weight);
}

static void
add_match(const MimeCache *c, Matches *m, uint32_t mime_offset, uint32_t weight) {
    if (m->count >= MAX_MATCHES) return;
    const char *mime = read_string(c, mime_offset);
    if (!mime) return;
    m->items[m->count].mime = mime;
    m->items[m->count++].weight = weight & WEIGHT_MASK;
}
// }}}

// Lookups {{{
static bool
lookup_literal(const MimeCache *c, const char *name, bool case_sensitive_check, Matches *ans) {
    uint32_t list_offset, n;
    if (!read_u32(c, LITERAL_LIST_OFFSET_POS, &list_offset) || !read_u32(c, list_offset, &n)) return false;
    size_t lo = 0, hi = n;
    while (lo < hi) {
        size_t mid = lo + (hi - lo) / 2;
        uint32_t literal_offset, mime_offset, weight;
        if (!read_entry(c, (size_t)list_offset + 4 + 12 * mid, &literal_offset, &mime_offset, &weight)) return false;
        const char *literal = read_string(c, literal_offset);
        if (!literal) return false;
        int cmp = strcmp(literal, name);
        if (cmp < 0) lo = mid + 1;
        else if (cmp > 0) hi = mid;
        else {
            if (!case_sensitive_check && (weight & CASE_SENSITIVE_FLAG)) return false;
            unsigned before = ans->count;
            add_match(c, ans, mime_offset, weight);
            return ans->count > before;
        }
    }
    return false;
}

static void
lookup_suffix_node(
    const MimeCache *c, uint32_t num_nodes, uint32_t offset, const char_type *name, size_t len, bool case_sensitive_check, Matches *ans, unsigned depth) {
    if (!len || depth > MAX_SUFFIX_TREE_DEPTH) return;
    const char_type ch = name[len - 1];
    size_t lo = 0, hi = num_nodes;
    while (lo < hi) {
        size_t mid = lo + (hi - lo) / 2;
        uint32_t node_char, num_children, first_child_offset;
        if (!read_entry(c, (size_t)offset + 12 * mid, &node_char, &num_children, &first_child_offset)) return;
        if (ch < node_char) hi = mid;
        else if (ch > node_char) lo = mid + 1;
        else {
            const unsigned before = ans->count;
            lookup_suffix_node(c, num_children, first_child_offset, name, len - 1, case_sensitive_check, ans, depth + 1);
            if (ans->count == before) {
                // No longer suffix matched, use the leaf nodes of this node,
                // which are sorted before all other children.
                for (uint32_t i = 0; i < num_children; i++) {
                    uint32_t leaf_char, mime_offset, weight;
                    if (!read_entry(c, (size_t)first_child_offset + 12 * (size_t)i, &leaf_char, &mime_offset, &weight) || leaf_char) break;
                    if (case_sensitive_check || !(weight & CASE_SENSITIVE_FLAG)) add_match(c, ans, mime_offset, weight);
                }
            }
            return;
        }
    }
}

static void
lookup_suffix(const MimeCache *c, const char_type *name, size_t len, bool case_sensitive_check, Matches *ans) {
    uint32_t tree_offset, num_roots, first_root_offset;
    if (!read_u32(c, REVERSE_SUFFIX_TREE_OFFSET_POS, &tree_offset) || !read_u32(c, tree_offset, &num_roots) ||
        !read_u32(c, (size_t)tree_offset + 4, &first_root_offset))
        return;
    lookup_suffix_node(c, num_roots, first_root_offset, name, len, case_sensitive_check, ans, 0);
}

static void
lookup_glob(const MimeCache *c, const char *name, bool case_sensitive_check, Matches *ans) {
    uint32_t list_offset, n;
    if (!read_u32(c, GLOB_LIST_OFFSET_POS, &list_offset) || !read_u32(c, list_offset, &n)) return;
    for (size_t i = 0; i < n && ans->count < MAX_MATCHES; i++) {
        uint32_t glob_offset, mime_offset, weight;
        if (!read_entry(c, (size_t)list_offset + 4 + 12 * i, &glob_offset, &mime_offset, &weight)) return;
        if (!case_sensitive_check && (weight & CASE_SENSITIVE_FLAG)) continue;
        const char *glob = read_string(c, glob_offset);
        if (glob && fnmatch(glob, name, 0) == 0) add_match(c, ans, mime_offset, weight);
    }
}

typedef struct {
    const char *utf8;
    const char_type *ucs4;
    size_t ucs4_len;
} FileName;

static const char *
lookup(const FileName *name, const FileName *lowercased) {
    Matches m = {0};
    for (size_t i = 0; i < caches.count; i++) {
        if (caches.items[i].data && lookup_literal(caches.items + i, name->utf8, true, &m)) return m.items[0].mime;
    }
    for (size_t i = 0; i < caches.count; i++) {
        if (caches.items[i].data && lookup_literal(caches.items + i, lowercased->utf8, false, &m)) return m.items[0].mime;
    }
    for (size_t i = 0; i < caches.count; i++) {
        if (caches.items[i].data) lookup_suffix(caches.items + i, lowercased->ucs4, lowercased->ucs4_len, false, &m);
    }
    if (m.count < 2) {
        for (size_t i = 0; i < caches.count; i++) {
            if (caches.items[i].data) lookup_suffix(caches.items + i, name->ucs4, name->ucs4_len, true, &m);
        }
    }
    if (!m.count) {
        for (size_t i = 0; i < caches.count; i++) {
            if (caches.items[i].data) lookup_glob(caches.items + i, name->utf8, true, &m);
        }
    }
    if (!m.count) {
        for (size_t i = 0; i < caches.count; i++) {
            if (caches.items[i].data) lookup_glob(caches.items + i, lowercased->utf8, false, &m);
        }
    }
    if (!m.count) return NULL;
    // Highest weight wins, ties go to the earliest match, which means the
    // longest suffix and the highest priority cache file.
    const Match *best = m.items;
    for (unsigned i = 1; i < m.count; i++) {
        if (m.items[i].weight > best->weight) best = m.items + i;
    }
    return best->mime;
}

static bool
read_alias_list(const MimeCache *c, uint32_t *list_offset, uint32_t *n) {
    return read_u32(c, ALIAS_LIST_OFFSET_POS, list_offset) && read_u32(c, *list_offset, n);
}

static bool
read_alias_entry(const MimeCache *c, uint32_t list_offset, size_t idx, const char **alias, const char **mime) {
    uint32_t alias_offset, mime_offset;
    const size_t offset = (size_t)list_offset + 4 + 8 * idx;
    if (!read_u32(c, offset, &alias_offset) || !read_u32(c, offset + 4, &mime_offset)) return false;
    *alias = read_string(c, alias_offset);
    *mime = read_string(c, mime_offset);
    return *alias && *mime;
}

// Returns the canonical MIME type for the specified alias or NULL if mime is
// not an alias. The alias list is sorted by alias.
static const char *
canonical_mime_type(const MimeCache *c, const char *mime) {
    uint32_t list_offset, n;
    if (!read_alias_list(c, &list_offset, &n)) return NULL;
    size_t lo = 0, hi = n;
    while (lo < hi) {
        size_t mid = lo + (hi - lo) / 2;
        const char *alias, *canonical;
        if (!read_alias_entry(c, list_offset, mid, &alias, &canonical)) return NULL;
        int cmp = strcmp(alias, mime);
        if (cmp < 0) lo = mid + 1;
        else if (cmp > 0) hi = mid;
        else return canonical;
    }
    return NULL;
}

// Corrupt cache files can contain strings that are not valid UTF-8. These are
// treated as missing: NULL is returned without setting an exception.
static PyObject *
string_from_cache(const char *val) {
    PyObject *ans = PyUnicode_DecodeUTF8(val, (Py_ssize_t)strlen(val), NULL);
    if (!ans && PyErr_ExceptionMatches(PyExc_UnicodeDecodeError)) PyErr_Clear();
    return ans;
}

static bool
add_unique_string(PyObject *list, const char *val) {
    RAII_PyObject(s, string_from_cache(val));
    if (!s) return !PyErr_Occurred();
    int ret = PySequence_Contains(list, s);
    if (ret < 0) return false;
    return ret == 1 || PyList_Append(list, s) == 0;
}

static bool
add_aliases(const MimeCache *c, const char *canonical, PyObject *ans) {
    uint32_t list_offset, n;
    if (!read_alias_list(c, &list_offset, &n)) return true;
    for (size_t i = 0; i < n; i++) {
        const char *alias, *mime;
        if (!read_alias_entry(c, list_offset, i, &alias, &mime)) break;
        if (strcmp(mime, canonical) == 0 && !add_unique_string(ans, alias)) return false;
    }
    return true;
}
// }}}

// Loading {{{
static void
unmap_cache(MimeCache *c) {
    if (c->data) munmap((void *)c->data, c->sz);
    c->data = NULL;
    c->sz = 0;
}

static bool
stat_cache(const MimeCache *c, struct stat *st) {
    return stat(c->path, st) == 0 && S_ISREG(st->st_mode);
}

static bool
cache_changed(const MimeCache *c) {
    struct stat st;
    if (!stat_cache(c, &st)) return c->exists;
    return !c->exists || st.st_dev != c->dev || st.st_ino != c->ino || st.st_size != c->size || st.st_mtim.tv_sec != c->mtime.tv_sec ||
           st.st_mtim.tv_nsec != c->mtime.tv_nsec;
}

// Note for security reviews: if a mapped cache file is truncated in place,
// accessing the mapped pages beyond the new end of file raises SIGBUS. This
// is deliberately not handled. update-mime-database replaces cache files
// atomically via rename(), so the mapped data stays valid, and truncating them
// in place requires write access to the XDG data directories. GLib (xdgmime)
// and Qt (QMimeBinaryProvider) mmap these files in exactly the same way,
// without handling SIGBUS. All reads from the mapped data are bounds checked,
// so corrupt or malicious cache file contents are safe.
static void
map_cache(MimeCache *c) {
    unmap_cache(c);
    c->exists = false;
    int fd = safe_open(c->path, O_RDONLY | O_CLOEXEC, 0);
    if (fd < 0) {
        if (errno != ENOENT && errno != ENOTDIR) log_error("Failed to open MIME cache file %s with error: %s", c->path, strerror(errno));
        return;
    }
    struct stat st;
    if (fstat(fd, &st) != 0 || !S_ISREG(st.st_mode)) {
        safe_close(fd, __FILE__, __LINE__);
        return;
    }
    c->exists = true;
    c->dev = st.st_dev;
    c->ino = st.st_ino;
    c->size = st.st_size;
    c->mtime = st.st_mtim;
    if (st.st_size < HEADER_SIZE) {
        safe_close(fd, __FILE__, __LINE__);
        log_error("The MIME cache file %s is too small, ignoring it", c->path);
        return;
    }
    void *addr = mmap(NULL, (size_t)st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    safe_close(fd, __FILE__, __LINE__);
    if (addr == MAP_FAILED) {
        log_error("Failed to mmap MIME cache file %s with error: %s", c->path, strerror(errno));
        return;
    }
    const uint8_t *p = addr;
    const unsigned major = (p[0] << 8) | p[1], minor = (p[2] << 8) | p[3];
    if (major != 1 || minor < 1) {
        log_error("The MIME cache file %s has unsupported version %u.%u, ignoring it", c->path, major, minor);
        munmap(addr, (size_t)st.st_size);
        return;
    }
    c->data = p;
    c->sz = (size_t)st.st_size;
}

static void
free_caches(void) {
    for (size_t i = 0; i < caches.count; i++) {
        unmap_cache(caches.items + i);
        free(caches.items[i].path);
    }
    free(caches.items);
    zero_at_ptr(&caches);
}

static bool
add_cache_path(const char *data_dir, size_t len) {
    if (!len || data_dir[0] != '/') return true; // relative paths are ignored as per the XDG spec
    static const char suffix[] = "/mime/mime.cache";
    char *path = malloc(len + sizeof(suffix));
    if (!path) return false;
    memcpy(path, data_dir, len);
    memcpy(path + len, suffix, sizeof(suffix));
    for (size_t i = 0; i < caches.count; i++) {
        if (strcmp(caches.items[i].path, path) == 0) {
            free(path);
            return true;
        }
    }
    MimeCache *items = realloc(caches.items, sizeof(caches.items[0]) * (caches.count + 1));
    if (!items) {
        free(path);
        return false;
    }
    caches.items = items;
    zero_at_ptr(caches.items + caches.count);
    caches.items[caches.count++].path = path;
    return true;
}

static bool
find_cache_paths(void) {
    // Cache files are in order of decreasing priority
    const char *data_home = getenv("XDG_DATA_HOME");
    if (data_home && data_home[0]) {
        if (!add_cache_path(data_home, strlen(data_home))) return false;
    } else {
        const char *home = getenv("HOME");
        if (home && home[0]) {
            static const char rel[] = "/.local/share";
            const size_t hl = strlen(home);
            RAII_ALLOC(char, buf, malloc(hl + sizeof(rel)));
            if (!buf) return false;
            memcpy(buf, home, hl);
            memcpy(buf + hl, rel, sizeof(rel));
            if (!add_cache_path(buf, hl + sizeof(rel) - 1)) return false;
        }
    }
    const char *data_dirs = getenv("XDG_DATA_DIRS");
    if (!data_dirs || !data_dirs[0]) data_dirs = "/usr/local/share/:/usr/share/";
    for (const char *p = data_dirs; *p;) {
        const char *end = strchr(p, ':');
        size_t len = end ? (size_t)(end - p) : strlen(p);
        while (len > 1 && p[len - 1] == '/') len--;
        if (!add_cache_path(p, len)) return false;
        if (!end) break;
        p = end + 1;
    }
    return true;
}

static bool
ensure_caches_loaded(void) {
    const monotonic_t now = monotonic();
    if (caches.loaded) {
        // Check for updated cache files at most once every few seconds, like xdgmime
        if (now - caches.last_check_at < s_to_monotonic_t(5)) return true;
        caches.last_check_at = now;
        for (size_t i = 0; i < caches.count; i++) {
            if (cache_changed(caches.items + i)) map_cache(caches.items + i);
        }
        return true;
    }
    free_caches();
    if (!find_cache_paths()) {
        free_caches();
        PyErr_NoMemory();
        return false;
    }
    for (size_t i = 0; i < caches.count; i++) map_cache(caches.items + i);
    caches.loaded = true;
    caches.last_check_at = now;
    return true;
}
// }}}

// Python API {{{
static bool
filename_from_python(PyObject *name, FileName *ans, PyObject **utf8_holder, Py_UCS4 **ucs4_holder) {
    // Use surrogateescape so that file names that are not valid UTF-8 round trip to their original bytes
    // Lone surrogates outside the surrogateescape range cannot be encoded and
    // cannot come from a real file name, so they are treated as not matching:
    // false is returned without setting an exception.
    *utf8_holder = PyUnicode_AsEncodedString(name, "utf-8", "surrogateescape");
    if (!*utf8_holder) {
        if (PyErr_ExceptionMatches(PyExc_UnicodeEncodeError)) PyErr_Clear();
        return false;
    }
    *ucs4_holder = PyUnicode_AsUCS4Copy(name);
    if (!*ucs4_holder) return false;
    ans->utf8 = PyBytes_AS_STRING(*utf8_holder);
    ans->ucs4 = *ucs4_holder;
    ans->ucs4_len = (size_t)PyUnicode_GET_LENGTH(name);
    return true;
}

static void
cleanup_ucs4(Py_UCS4 **p) {
    PyMem_Free(*p);
    *p = NULL;
}
#define RAII_UCS4(name) __attribute__((cleanup(cleanup_ucs4))) Py_UCS4 *name = NULL

static PyObject *
mime_type_for_filename(PyObject *self UNUSED, PyObject *path) {
    if (!PyUnicode_Check(path)) {
        PyErr_SetString(PyExc_TypeError, "path must be a string");
        return NULL;
    }
    const Py_ssize_t path_len = PyUnicode_GET_LENGTH(path);
    const Py_ssize_t slash = PyUnicode_FindChar(path, '/', 0, path_len, -1);
    if (slash == -2) return NULL;
    RAII_PyObject(name, PyUnicode_Substring(path, slash + 1, path_len));
    if (!name) return NULL;
    const Py_ssize_t name_len = PyUnicode_GET_LENGTH(name);
    if (!name_len) Py_RETURN_NONE;
    const Py_ssize_t nul = PyUnicode_FindChar(name, 0, 0, name_len, 1);
    if (nul == -2) return NULL;
    if (nul != -1) Py_RETURN_NONE; // embedded NUL cannot be part of a valid file name
    RAII_PyObject(lowercased, PyObject_CallMethod(name, "lower", NULL));
    if (!lowercased) return NULL;
    if (!ensure_caches_loaded()) return NULL;
    if (!caches.count) Py_RETURN_NONE;
    RAII_PyObject(utf8, NULL);
    RAII_PyObject(lutf8, NULL);
    RAII_UCS4(ucs4);
    RAII_UCS4(lucs4);
    FileName fname, lname;
    if (!filename_from_python(name, &fname, &utf8, &ucs4) || !filename_from_python(lowercased, &lname, &lutf8, &lucs4)) {
        if (PyErr_Occurred()) return NULL;
        Py_RETURN_NONE;
    }
    const char *ans = lookup(&fname, &lname);
    if (!ans) Py_RETURN_NONE;
    PyObject *ret = string_from_cache(ans);
    if (!ret && !PyErr_Occurred()) Py_RETURN_NONE;
    return ret;
}

static PyObject *
mime_type_aliases(PyObject *self UNUSED, PyObject *mime) {
    if (!PyUnicode_Check(mime)) {
        PyErr_SetString(PyExc_TypeError, "mime must be a string");
        return NULL;
    }
    const char *q = PyUnicode_AsUTF8(mime);
    if (!q) {
        // Lone surrogates cannot be encoded and so cannot match anything in the cache
        if (!PyErr_ExceptionMatches(PyExc_UnicodeEncodeError)) return NULL;
        PyErr_Clear();
        return PyTuple_Pack(1, mime);
    }
    if (!ensure_caches_loaded()) return NULL;
    const char *canonical = NULL;
    for (size_t i = 0; i < caches.count && !canonical; i++) {
        if (caches.items[i].data) canonical = canonical_mime_type(caches.items + i, q);
    }
    if (!canonical) canonical = q;
    RAII_PyObject(ans, PyList_New(0));
    if (!ans || !add_unique_string(ans, canonical)) return NULL;
    for (size_t i = 0; i < caches.count; i++) {
        if (caches.items[i].data && !add_aliases(caches.items + i, canonical, ans)) return NULL;
    }
    // Ensure the queried name is always present, even if a cache entry is corrupt
    if (!add_unique_string(ans, q)) return NULL;
    return PyList_AsTuple(ans);
}

static PyObject *
clear_mime_cache_data(PyObject *self UNUSED, PyObject *args UNUSED) {
    free_caches();
    Py_RETURN_NONE;
}
// }}}

#endif

static PyMethodDef methods[] = {
    {"mime_type_for_filename",
     mime_type_for_filename,
     METH_O,
     "Return the MIME type for the specified file name based on the shared-mime-info mime.cache files or None if no match is found"},
    {"mime_type_aliases",
     mime_type_aliases,
     METH_O,
     "Return all names for the specified MIME type, starting with its canonical name, followed by all its aliases"},
    {"clear_mime_cache_data", clear_mime_cache_data, METH_NOARGS, "Release all loaded MIME cache data, it will be re-loaded on next use"},
    {NULL, NULL, 0, NULL},
};

bool
init_mime_cache(PyObject *module) {
    return PyModule_AddFunctions(module, methods) == 0;
}