XFESystem

#include <stdint.h>
#include <stddef.h>

// ELF x86_64 TLS variant II. This is kernel code (ring 0), not a host libc.
// A scheduler must save/restore IA32_FS_BASE on every thread switch. No managed
// object is stored in this allocation; the Microsoft runtime owns its TLS roots.
extern unsigned char __tls_start[], __tls_init_end[], __tls_end[];
extern unsigned char __tls_align[];
extern void* xfe_runtime_alloc_pages(uint64_t pages);
extern int xfe_runtime_free_pages(void* address, uint64_t pages);
extern void xfe_runtime_halt(const char* reason);

static __thread int xfe_thread_errno;
static int xfe_boot_errno;
struct XfeTlsControl {
    void* self;
    void* dynamic_vector;
    void* tls_base;
    uint64_t tls_bytes;
    void* allocation;
    uint64_t allocation_pages;
    uint64_t signature;
};
static uint64_t get_fs_base(void)
{
    uint32_t low,high;
    __asm__ volatile("rdmsr" : "=a"(low),"=d"(high) : "c"(0xc0000100));
    return ((uint64_t)high<<32)|low;
}
static void set_fs_base(uint64_t address)
{
    __asm__ volatile("wrmsr" :: "c"(0xc0000100),"a"((uint32_t)address),"d"((uint32_t)(address>>32)) : "memory");
}

void* XfeRuntimeTlsCreate(void)
{
    uint64_t initialized=(uintptr_t)__tls_init_end-(uintptr_t)__tls_start;
    uint64_t bytes=(uintptr_t)__tls_end-(uintptr_t)__tls_start;
    uint64_t alignment=(uintptr_t)__tls_align;
    if(alignment<16)alignment=16;
    if(initialized>bytes || bytes>16*1024*1024 || (alignment&(alignment-1)) || alignment>4096)return NULL;
    bytes=(bytes+alignment-1)&~(alignment-1);
    uint64_t total=bytes+sizeof(struct XfeTlsControl)+alignment;
    uint64_t pages=(total+4095)/4096;
    unsigned char* raw=xfe_runtime_alloc_pages(pages);
    if(!raw)return NULL;
    unsigned char* base=(unsigned char*)(((uintptr_t)raw+alignment-1)&~(alignment-1));
    for(uint64_t i=0;i<bytes;i++)base[i]=i<initialized?__tls_start[i]:0;
    struct XfeTlsControl* control=(struct XfeTlsControl*)(base+bytes);
    *control=(struct XfeTlsControl){control,NULL,base,bytes,raw,pages,0x584645544c533031ULL};
    return control;
}
int XfeRuntimeTlsSet(void* base)
{
    struct XfeTlsControl* control=(struct XfeTlsControl*)base;
    if(!control || control->self!=control || control->signature!=0x584645544c533031ULL)return 0;
    set_fs_base((uintptr_t)base);
    return 1;
}
void* XfeRuntimeTlsCurrent(void) { return (void*)get_fs_base(); }
int XfeRuntimeTlsRelease(void* base)
{
    struct XfeTlsControl* control=(struct XfeTlsControl*)base;
    if(!control || (void*)get_fs_base()==base || control->signature!=0x584645544c533031ULL)return 0;
    return xfe_runtime_free_pages(control->allocation,control->allocation_pages);
}
int XfeRuntimeTlsInitialize(void)
{
    if(get_fs_base())return 0; // Never silently replace firmware or another thread's TLS.
    void* control=XfeRuntimeTlsCreate();
    return control?XfeRuntimeTlsSet(control):0;
}
__attribute__((noinline)) static int* runtime_errno_address(void)
{
    // Keep this load behind the bootstrap check. A ternary address expression
    // can otherwise be lowered to an unconditional FS:0 read plus cmov.
    __asm__ volatile("" ::: "memory");
    return &xfe_thread_errno;
}
int* __errno_location(void)
{
    // During the pre-runtime kernel selftest no threads exist and FS is zero.
    // Once a runtime thread is installed the address follows that thread's TLS.
    if(!get_fs_base())return &xfe_boot_errno;
    return runtime_errno_address();
}
struct XfeTlsIndex { uint64_t module,offset; };
void* __tls_get_addr(struct XfeTlsIndex* index)
{
    struct XfeTlsControl* control=(struct XfeTlsControl*)get_fs_base();
    if(!control || !index || index->module!=1 || index->offset>=control->tls_bytes || control->signature!=0x584645544c533031ULL) {
        xfe_runtime_halt("Invalid NativeAOT TLS access");
        for(;;)__asm__ volatile("cli; hlt");
    }
    return (unsigned char*)control->tls_base+index->offset;
}