Implementing TLS (Thread Local storage) for x86_64

#include <threads.h>
#include <stdio.h>
#include <stdlib.h>

thread_local int counter = 0;

int thread_func(void *arg) {
    int id = *(int*)arg;
    counter++;  // Each thread increments its own copy
    printf("Thread %d: counter = %d\n", id, counter);
    return 0;
}

int main() {
    thrd_t threads[4];
    int ids[4] = {1, 2, 3, 4};

    for (int i = 0; i < 4; i++) {
        thrd_create(&threads[i], thread_func, &ids[i]);
    }

    for (int i = 0; i < 4; i++) {
        thrd_join(threads[i], NULL);
    }

    printf("Main thread counter: %d\n", counter);  // Main's own copy
    return 0;
}

/* int thread_func(void *arg) */
thread_func:
        /* Push new stack frame */
        push    rbp
        mov     rbp, rsp
        mov     qword ptr [rbp - 8], rdi /* store arg on the stack frame */


        /* Read the ID value */
        /* int id = *(int*)arg; */
        mov     rax, qword ptr [rbp - 8]
        mov     eax, dword ptr [rax]
        mov     dword ptr [rbp - 12], eax


        /* counter++; */
        mov     rax, qword ptr fs:[0] /* ?????????? */
        lea     rax, [rax + counter@TPOFF]
        mov     ecx, dword ptr [rax]
        add     ecx, 1 /* do the ++ */
        mov     dword ptr [rax], ecx

        /* return 0; */
        xor     eax, eax
        pop     rbp
        ret

/* The rest is irrelevant here... */

counter:
        .long   0

movl NUMBER_OF_MSR, %ecx
movl VALUE_BITS_LOW, %eax
movl VALUE_BITS_HIGH, %edx
wrmsr

movl NUMBER_OF_MSR, %ecx
rdmsr
/* now %eax contains high bits and %edx low bits. These two shall be concatinated into a 64 bit value */

#define MSR_FS_BASE        _AC (0xc0000100, U) /* 64bit FS base */

void do_sched (struct proc* proc, spin_lock_t* cpu_lock, spin_lock_ctx_t* ctxcpu) {
  spin_lock_ctx_t ctxpr;

  spin_lock (&proc->lock, &ctxpr);

  thiscpu->tss.rsp0 = proc->pdata.kernel_stack; /* set TSS kernel stack */
  thiscpu->syscall_kernel_stack = proc->pdata.kernel_stack; /* set syscall entry stack */
  amd64_wrmsr (MSR_FS_BASE, proc->pdata.fs_base); /* switch to proc's fs base */

  spin_unlock (&proc->lock, &ctxpr);
  spin_unlock (cpu_lock, ctxcpu);

  amd64_do_sched ((void*)&proc->pdata.regs, (void*)proc->procgroup->pd.cr3_paddr);
}

/// Read a model-specific register
uint64_t amd64_rdmsr (uint32_t msr) {
  uint32_t low, high;
  __asm__ volatile ("rdmsr" : "=a"(low), "=d"(high) : "c"(msr));
  return ((uint64_t)high << 32 | (uint64_t)low);
}

/// Write a model-specific register
void amd64_wrmsr (uint32_t msr, uint64_t value) {
  uint32_t low = (uint32_t)(value & 0xFFFFFFFF);
  uint32_t high = (uint32_t)(value >> 32);
  __asm__ volatile ("wrmsr" ::"c"(msr), "a"(low), "d"(high));
}

The TLS memory:

  Var 1   Var 2    Var 3    Var 4 ....                                The pointer
+-------------------------------------------------------------------------------+
|       |        |        |        |        |        |        |        |        |  <---+
+-------------------------------------------------------------------------------+      |
                                                                                       |
                                                                            ^          |
                                                                            |          |
                                                                       TLS (fs base)   |
                                                                                       |
                                                                  %fs:0  --------------+

__thread int a = 123;

void my_thread (void) {
  printf ("a = %d\n", a);

  a = 456;

  printf ("a = %d\n", a);
}

    case PT_TLS: {
#if defined(__x86_64__)
      if (phdr->p_memsz > 0) {
        /* What is the aligment we need to use? */
        size_t tls_align = phdr->p_align ? phdr->p_align : sizeof (uintptr_t);
        /* Size of the TLS memory block (variables go here) */
        size_t tls_size = align_up (phdr->p_memsz, tls_align);
        /* Size needed - TLS block size + 8 bytes (64 bits) for back pointer */
        size_t tls_total_needed = tls_size + sizeof (uintptr_t);
        /* amount of pages to allocate */
        size_t blks = div_align_up (tls_total_needed, PAGE_SIZE);
        /* Initialize TLS template in the procgroup. This will be copied into individual TLSes */
        proc->procgroup->tls.tls_tmpl_pages = blks;
        proc->procgroup->tls.tls_tmpl_size = tls_size;
        proc->procgroup->tls.tls_tmpl_total_size = tls_total_needed;

        /* malloc () and zero out */
        proc->procgroup->tls.tls_tmpl = malloc (blks * PAGE_SIZE);
        memset (proc->procgroup->tls.tls_tmpl, 0, blks * PAGE_SIZE);

        /* copy initialized stuff */
        memcpy (proc->procgroup->tls.tls_tmpl, (void*)((uintptr_t)elf + phdr->p_offset),
                phdr->p_filesz);

        proc_init_tls (proc);
      }
#endif
    } break;

void proc_init_tls (struct proc* proc) {
  struct limine_hhdm_response* hhdm = limine_hhdm_request.response;

  /* This application doesn't use TLS */
  if (proc->procgroup->tls.tls_tmpl == NULL)
    return;

  size_t tls_size = proc->procgroup->tls.tls_tmpl_size;
  size_t pages = proc->procgroup->tls.tls_tmpl_pages;

  uintptr_t tls_paddr;
  uint32_t flags = MM_PG_USER | MM_PG_PRESENT | MM_PG_RW;

  /* allocate a new TLS memory space and map it into the procgroup's address space */
  uintptr_t tls_vaddr = procgroup_map (proc->procgroup, 0, pages, flags, &tls_paddr);

  uintptr_t k_tls_addr = (uintptr_t)hhdm->offset + tls_paddr;

  /* zero and copy the template contents */
  memset ((void*)k_tls_addr, 0, pages * PAGE_SIZE);
  memcpy ((void*)k_tls_addr, (void*)proc->procgroup->tls.tls_tmpl, tls_size);

  /* kernel address and user address + size will point to the tls pointer */
  uintptr_t ktcb = k_tls_addr + tls_size;
  uintptr_t utcb = tls_vaddr + tls_size;

  /* write the pointer value, which makes the TLS point to itself */
  *(uintptr_t*)ktcb = utcb;

  /* store as fs_base for switching during scheduling */
  proc->pdata.fs_base = utcb;
  /* save allocation address to later free it when not needed */
  proc->pdata.tls_vaddr = tls_vaddr;
}

#define MUTEX 2000

LOCAL volatile char letter = 'c';

void app_proc1 (void) {
  letter = 'a';

  for (;;) {
    mutex_lock (MUTEX);

    for (int i = 0; i < 3; i++)
      test (letter);

    mutex_unlock (MUTEX);
  }

  process_quit ();
}

void app_proc2 (void) {
  letter = 'b';

  for (;;) {
    mutex_lock (MUTEX);

    for (int i = 0; i < 3; i++)
      test (letter);

    mutex_unlock (MUTEX);
  }

  process_quit ();
}

void app_proc3 (void) {
  letter = 'c';

  for (;;) {
    mutex_lock (MUTEX);

    for (int i = 0; i < 3; i++)
      test (letter);

    mutex_unlock (MUTEX);
  }

  process_quit ();
}

void app_main (void) {
  mutex_create (MUTEX);

  letter = 'a';

  process_spawn (&app_proc1, NULL);
  process_spawn (&app_proc2, NULL);
  process_spawn (&app_proc3, NULL);

  for (;;) {
    mutex_lock (MUTEX);

    for (int i = 0; i < 3; i++)
      test (letter);

    mutex_unlock (MUTEX);
  }
}