gpacache.c

Go to the documentation of this file.

/*
 * Copyright (c) 2020 Bitdefender
 * SPDX-License-Identifier: Apache-2.0
 */
#include "gpacache.h"
#include "glue.h"
#include "introcrt.h"


static __inline DWORD
IntGpaCacheHashLine(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa
    )
{
    return (Gpa >> 12) & (Cache->LinesCount - 1);
}


static INTSTATUS
IntGpaCacheAddVictim(
    _In_ PGPA_CACHE Cache,
    _In_ PGPA_CACHE_ENTRY Entry
    )
{
    GPA_CACHE_VICTIM *pVictim = HpAllocWithTag(sizeof(*pVictim), IC_TAG_GPCV);
    if (NULL == pVictim)
    {
        return INT_STATUS_INSUFFICIENT_RESOURCES;
    }

    memcpy(&pVictim->Entry, Entry, sizeof(GPA_CACHE_ENTRY));

    InsertTailList(&Cache->Victims, &pVictim->Link);

    return INT_STATUS_SUCCESS;
}


void
IntGpaCacheDump(
    _In_ PGPA_CACHE Cache
    )
{
    DWORD i, j;
    QWORD totalHitCount = 0;

    LOG("Gpa Cache:\n");

    NLOG("Number of lines:  %d\n", Cache->LinesCount);
    NLOG("Entries per line: %d\n", Cache->EntriesCount);

    // Dump the entries inside the Cache
    for (i = 0; i < Cache->LinesCount; i++)
    {
        QWORD lineHitCount = 0;

        for (j = 0; j < Cache->EntriesCount; j++)
        {
            if (Cache->Lines[i].Entries[j].Valid)
            {
                NLOG("%04d:%02d: 0x%016llx:%p, %04d %04d %d\n",
                     i, j,
                     Cache->Lines[i].Entries[j].Gpa,
                     Cache->Lines[i].Entries[j].Hva,
                     Cache->Lines[i].Entries[j].HitCount,
                     Cache->Lines[i].Entries[j].UseCount,
                     Cache->Lines[i].Entries[j].Valid);

                totalHitCount += Cache->Lines[i].Entries[j].HitCount;
                lineHitCount += Cache->Lines[i].Entries[j].HitCount;
            }
        }

        if (lineHitCount > 0)
        {
            NLOG("    HitCount: %lld\n", lineHitCount);
        }
    }

    NLOG("TOTALHC: %lld\n", totalHitCount);
}


INTSTATUS
IntGpaCacheInit(
    _Inout_ PGPA_CACHE *Cache,
    _In_ DWORD LinesCount,
    _In_ DWORD EntriesCount
    )
{
    INTSTATUS status;
    GPA_CACHE *pCache;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    // LinesCount must be power of two.
    if (0 != (LinesCount & (LinesCount - 1)))
    {
        return INT_STATUS_INVALID_PARAMETER_3;
    }

    pCache = HpAllocWithTag(sizeof(*pCache), IC_TAG_GPCA);
    if (NULL == pCache)
    {
        return INT_STATUS_INSUFFICIENT_RESOURCES;
    }

    pCache->LinesCount = LinesCount;
    pCache->EntriesCount = EntriesCount;

    pCache->Lines = HpAllocWithTag(sizeof(*pCache->Lines) * LinesCount, IC_TAG_GPCA);
    if (NULL == pCache->Lines)
    {
        status = INT_STATUS_INSUFFICIENT_RESOURCES;
        goto cleanup_and_exit;
    }

    for (DWORD i = 0; i < LinesCount; i++)
    {
        pCache->Lines[i].Entries = HpAllocWithTag(sizeof(*pCache->Lines[i].Entries) * EntriesCount, IC_TAG_GPCA);
        if (NULL == pCache->Lines[i].Entries)
        {
            status = INT_STATUS_INSUFFICIENT_RESOURCES;
            goto cleanup_and_exit;
        }
    }

    InitializeListHead(&pCache->Victims);

    status = INT_STATUS_SUCCESS;

cleanup_and_exit:
    if (!INT_SUCCESS(status))
    {
        if (NULL != pCache)
        {
            if (NULL != pCache->Lines)
            {
                for (DWORD i = 0; i < LinesCount; i++)
                {
                    if (NULL != pCache->Lines[i].Entries)
                    {
                        HpFreeAndNullWithTag(&pCache->Lines[i].Entries, IC_TAG_GPCA);
                    }
                }

                HpFreeAndNullWithTag(&pCache->Lines, IC_TAG_GPCA);
            }

            HpFreeAndNullWithTag(&pCache, IC_TAG_GPCA);
        }
    }

    *Cache = pCache;

    return status;
}


INTSTATUS
IntGpaCacheUnInit(
    _In_ PGPA_CACHE *Cache
    )
{
    DWORD i, j;
    PGPA_CACHE pCache;
    LIST_ENTRY *list;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    pCache = *Cache;
    if (NULL == pCache)
    {
        return INT_STATUS_NOT_INITIALIZED_HINT;
    }

    for (i = 0; i < pCache->LinesCount; i++)
    {
        for (j = 0; j < pCache->EntriesCount; j++)
        {
            if (pCache->Lines[i].Entries[j].Valid)
            {
                IntPhysMemUnmap(&pCache->Lines[i].Entries[j].Hva);
            }
        }

        HpFreeAndNullWithTag(&pCache->Lines[i].Entries, IC_TAG_GPCA);
    }

    HpFreeAndNullWithTag(&pCache->Lines, IC_TAG_GPCA);

    list = pCache->Victims.Flink;

    while (list != &pCache->Victims)
    {
        PGPA_CACHE_VICTIM pVictim = CONTAINING_RECORD(list, GPA_CACHE_VICTIM, Link);

        list = list->Flink;

        if (pVictim->Entry.Valid && (pVictim->Entry.UseCount > 0) && (NULL != pVictim->Entry.Hva))
        {
            IntPhysMemUnmap(&pVictim->Entry.Hva);
        }

        RemoveEntryList(&pVictim->Link);

        HpFreeAndNullWithTag(&pVictim, IC_TAG_GPCV);
    }

    HpFreeAndNullWithTag(Cache, IC_TAG_GPCA);

    return INT_STATUS_SUCCESS;
}


static INTSTATUS
IntGpaCacheAddEntry(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa,
    _Out_ PGPA_CACHE_ENTRY *Entry
    )
{
    INTSTATUS status;
    DWORD line, entry, minHits;

    entry = 0;
    line = IntGpaCacheHashLine(Cache, Gpa);
    minHits = 0xFFFFFFFF;

    // GPA must be 4K aligned.
    Gpa &= PHYS_PAGE_MASK;

    for (DWORD i = 0; i < Cache->EntriesCount; i++)
    {
        if (!Cache->Lines[line].Entries[i].Valid)
        {
            entry = i;
            break;
        }
        else
        {
            // Select the entry with the lowest LFU score, as long as it is not in use.
            if ((minHits > Cache->Lines[line].Entries[i].HitCount) && (0 == Cache->Lines[line].Entries[i].UseCount))
            {
                minHits = Cache->Lines[line].Entries[i].HitCount;
                entry = i;
            }
        }
    }

    // Couldn't decide upon selecting a proper entry to evict, select a random one.
    if (minHits == 0xFFFFFFFF)
    {
        entry = __rdtsc() % Cache->EntriesCount;
    }

    // If the entry is in use, add it to the victim cache, and wait for it to be released.
    if (Cache->Lines[line].Entries[entry].Valid && (Cache->Lines[line].Entries[entry].UseCount > 0))
    {
        TRACE("[GPACACHE] Adding victim entry for GPA %llx, use count %d...\n",
              Cache->Lines[line].Entries[entry].Gpa, Cache->Lines[line].Entries[entry].UseCount);

        status = IntGpaCacheAddVictim(Cache, &Cache->Lines[line].Entries[entry]);
        if (!INT_SUCCESS(status))
        {
            ERROR("[ERROR] IntGpaCacheAddVictim failed: 0x%08x\n", status);
            goto cleanup_and_exit;
        }
    }
    else if (Cache->Lines[line].Entries[entry].Valid && (Cache->Lines[line].Entries[entry].UseCount == 0))
    {
        IntPhysMemUnmap(&Cache->Lines[line].Entries[entry].Hva);
    }

    // Map it and leave
    status = IntPhysMemMap(Gpa, PAGE_SIZE, PHYS_MAP_FLG_NO_FASTMAP, &Cache->Lines[line].Entries[entry].Hva);

    if (INT_SUCCESS(status))
    {
        Cache->Lines[line].Entries[entry].Valid = TRUE;
        Cache->Lines[line].Entries[entry].HitCount = 0;
        Cache->Lines[line].Entries[entry].UseCount = 0;
        Cache->Lines[line].Entries[entry].Gpa = Gpa;
    }
    else
    {
        Cache->Lines[line].Entries[entry].Valid = FALSE;
        Cache->Lines[line].Entries[entry].Gpa = 0;
    }

    *Entry = &Cache->Lines[line].Entries[entry];

cleanup_and_exit:
    return status;
}


static INTSTATUS
IntGpaCacheLookupEntry(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa,
    _Out_ PGPA_CACHE_ENTRY *Entry
    )
{
    DWORD line, i;
    BOOLEAN bFound;
    LIST_ENTRY *list;

    bFound = FALSE;

    line = IntGpaCacheHashLine(Cache, Gpa);

    // Gpa must be page-aligned.
    Gpa &= PHYS_PAGE_MASK;

    for (i = 0; i < Cache->EntriesCount; i++)
    {
        if (Cache->Lines[line].Entries[i].Valid && (Cache->Lines[line].Entries[i].Gpa == Gpa))
        {
            *Entry = &Cache->Lines[line].Entries[i];
            bFound = TRUE;
            break;
        }
    }

    if (!bFound)
    {
        // Search the victim cache. If an entry is found in the victim cache, it could be brought back to the
        // main cache.
        list = Cache->Victims.Flink;

        while (list != &Cache->Victims)
        {
            PGPA_CACHE_VICTIM pVictim = CONTAINING_RECORD(list, GPA_CACHE_VICTIM, Link);

            list = list->Flink;

            if (pVictim->Entry.Valid && pVictim->Entry.Gpa == Gpa)
            {
                TRACE("[GPACACHE] Entry %llx found in victim cache, will not add it back...\n", Gpa);
                *Entry = &pVictim->Entry;
                bFound = TRUE;
                break;
            }
        }
    }

    if (bFound)
    {
        return INT_STATUS_SUCCESS;
    }

    return INT_STATUS_NOT_FOUND;
}


INTSTATUS
IntGpaCacheFindAndAdd(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa,
    _Out_ void **Hva
    )
{
    INTSTATUS status;
    PGPA_CACHE_ENTRY entry;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    if (NULL == Hva)
    {
        return INT_STATUS_INVALID_PARAMETER_3;
    }

    entry = NULL;

    status = IntGpaCacheLookupEntry(Cache, Gpa, &entry);
    if (!INT_SUCCESS(status))
    {
        status = IntGpaCacheAddEntry(Cache, Gpa, &entry);
        if (!INT_SUCCESS(status))
        {
            return status;
        }
    }

    entry->HitCount++;
    entry->UseCount++;

    *Hva = entry->Hva;

    return INT_STATUS_SUCCESS;
}


INTSTATUS
IntGpaCacheFetchAndAdd(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa,
    _In_ DWORD Size,
    _Out_ PBYTE Buffer
    )
{
    INTSTATUS status;
    PGPA_CACHE_ENTRY entry;
    PBYTE data;
    DWORD offset;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    if (NULL == Buffer)
    {
        return INT_STATUS_INVALID_PARAMETER_4;
    }

    entry = NULL;
    offset = Gpa & PAGE_OFFSET;

    if (offset + Size > PAGE_SIZE)
    {
        return INT_STATUS_NOT_SUPPORTED;
    }

    status = IntGpaCacheLookupEntry(Cache, Gpa, &entry);
    if (!INT_SUCCESS(status))
    {
        status = IntGpaCacheAddEntry(Cache, Gpa, &entry);
        if (!INT_SUCCESS(status))
        {
            return status;
        }
    }

    entry->HitCount++;

    data = (PBYTE)entry->Hva + offset;

    switch (Size)
    {
    case 8:
        *(PQWORD)Buffer = *(PQWORD)data;
        break;
    case 4:
        *(PDWORD)Buffer = *(PDWORD)data;
        break;
    case 2:
        *(PWORD)Buffer = *(PWORD)data;
        break;
    case 1:
        *Buffer = *data;
        break;
    default:
        memcpy(Buffer, data, Size);
        break;
    }

    return INT_STATUS_SUCCESS;
}


INTSTATUS
IntGpaCachePatchAndAdd(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa,
    _In_ DWORD Size,
    _In_ PBYTE Buffer
    )
{
    INTSTATUS status;
    PGPA_CACHE_ENTRY entry;
    PBYTE data;
    DWORD offset;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    if (NULL == Buffer)
    {
        return INT_STATUS_INVALID_PARAMETER_4;
    }

    entry = NULL;
    offset = Gpa & PAGE_OFFSET;

    if (offset + Size > PAGE_SIZE)
    {
        return INT_STATUS_NOT_SUPPORTED;
    }

    status = IntGpaCacheLookupEntry(Cache, Gpa, &entry);
    if (!INT_SUCCESS(status))
    {
        status = IntGpaCacheAddEntry(Cache, Gpa, &entry);
        if (!INT_SUCCESS(status))
        {
            return status;
        }
    }

    entry->HitCount++;

    data = (PBYTE)entry->Hva + offset;

    switch (Size)
    {
    case 8:
        *(PQWORD)data = *(PQWORD)Buffer;
        break;
    case 4:
        *(PDWORD)data = *(PDWORD)Buffer;
        break;
    case 2:
        *(PWORD)data = *(PWORD)Buffer;
        break;
    case 1:
        *data = *Buffer;
        break;
    default:
        memcpy(data, Buffer, Size);
        break;
    }

    return INT_STATUS_SUCCESS;
}


INTSTATUS
IntGpaCacheRelease(
    _In_ PGPA_CACHE Cache,
    _In_ QWORD Gpa
    )
{
    DWORD line, i;
    BOOLEAN bFound;
    LIST_ENTRY *list;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    bFound = FALSE;

    line = IntGpaCacheHashLine(Cache, Gpa);

    // Gpa must be page-aligned.
    Gpa &= PHYS_PAGE_MASK;

    for (i = 0; i < Cache->EntriesCount; i++)
    {
        if (Cache->Lines[line].Entries[i].Valid && (Cache->Lines[line].Entries[i].Gpa == Gpa))
        {
            if (Cache->Lines[line].Entries[i].UseCount > 0)
            {
                Cache->Lines[line].Entries[i].UseCount--;
            }

            bFound = TRUE;
            break;
        }
    }

    if (!bFound)
    {
        TRACE("[GPACACHE] Entry %llx not found in main cache, searching victim cache...\n", Gpa);

        list = Cache->Victims.Flink;

        while (list != &Cache->Victims)
        {
            PGPA_CACHE_VICTIM pVictim = CONTAINING_RECORD(list, GPA_CACHE_VICTIM, Link);

            list = list->Flink;

            if (pVictim->Entry.Gpa == Gpa)
            {
                if (pVictim->Entry.UseCount > 0)
                {
                    pVictim->Entry.UseCount--;
                }

                TRACE("[GPACACHE] Entry %llx is victim, use = %d\n", Gpa, pVictim->Entry.UseCount);

                if (0 == pVictim->Entry.UseCount)
                {
                    IntPhysMemUnmap(&pVictim->Entry.Hva);

                    RemoveEntryList(&pVictim->Link);

                    HpFreeAndNullWithTag(&pVictim, IC_TAG_GPCV);
                }

                bFound = TRUE;

                break;
            }
        }
    }

    if (bFound)
    {
        return INT_STATUS_SUCCESS;
    }

    return INT_STATUS_NOT_FOUND;
}


INTSTATUS
IntGpaCacheFlush(
    _In_ PGPA_CACHE Cache
    )
{
    INTSTATUS status;
    DWORD i, j;

    if (NULL == Cache)
    {
        return INT_STATUS_INVALID_PARAMETER_1;
    }

    for (i = 0; i < Cache->LinesCount; i++)
    {
        for (j = 0; j < Cache->EntriesCount; j++)
        {
            if (Cache->Lines[i].Entries[j].Valid)
            {
                if (Cache->Lines[i].Entries[j].UseCount == 0)
                {
                    IntPhysMemUnmap(&Cache->Lines[i].Entries[j].Hva);
                }
                else
                {
                    status = IntGpaCacheAddVictim(Cache, &Cache->Lines[i].Entries[j]);
                    if (!INT_SUCCESS(status))
                    {
                        ERROR("[ERROR] IntGpaCacheAddVictim failed: 0x%08x\n", status);
                    }
                }

                Cache->Lines[i].Entries[j].Valid = FALSE;
                Cache->Lines[i].Entries[j].Gpa = 0;
                Cache->Lines[i].Entries[j].HitCount = 0;
                Cache->Lines[i].Entries[j].UseCount = 0;
            }
        }
    }

    return INT_STATUS_SUCCESS;
}