/*

 * linux_vdso_my_sched_getcpu.c - use Linux VDSO

 * shortcut for |sched_getcpu()|, even if userland (g)libc

 * does not have support for VDSO

 *

 * Compile with:

 * $ clang -std=c99 -g -m32 -Wall \

 * 	-Wextra linux_vdso_my_sched_getcpu.c -lpthread \

 *	-o linux_vdso_my_sched_getcpu

 *

 * Written by Roland Mainz <roland.mainz@nrubsig.org>

 *

 */

#define _XOPEN_SOURCE	700

#define _GNU_SOURCE	1

#include <stdio.h>

#include <stddef.h>

#include <stdlib.h>

#include <string.h>

#include <elf.h>

#include <time.h>

#include <pthread.h>

#include <sched.h>

#include <sys/auxv.h>

/* Catch cursed idiots playing around with XPG config */

#if (_XOPEN_SOURCE < 600)

#error _XOPEN_SOURCE wrong version

#endif

/* Build test code ? */

#define MY_SCHED_GETCPU_TESTS 1

/* debug */

#if 1

#define D(x) x

#else

#define D(x)

#endif

/* 32bit vs. 64bit */

typedef Elf32_Ehdr	ElfXX_Ehdr;

typedef Elf32_Shdr	ElfXX_Shdr;

typedef Elf32_Sym	ElfXX_Sym;

typedef Elf64_Ehdr	ElfXX_Ehdr;

typedef Elf64_Shdr	ElfXX_Shdr;

typedef Elf64_Sym	ElfXX_Sym;

#else

#endif

/*

 * |get_linux_vdso_sym()| - Look up symbol in AUX AT_SYSINFO_EHDR

 */

static

void *get_linux_vdso_sym(const char *funcname)

{

	const char	*dynstr;

	void		*ret = NULL;

	unsigned char	*vdso_addr;

	ssize_t		i;

	size_t		si;

	const char	*name;

	ElfXX_Ehdr	*elf_header;

	ElfXX_Shdr	*section_header;

	ElfXX_Shdr	*s;

	ElfXX_Shdr	*s2;

	ElfXX_Sym	*sym;

	/*

	 * valgrind can disable |AT_SYSINFO_EHDR| - see

	 * https://www.mail-archive.com/kde-bugs-dist@kde.org/msg550927.html

	 */

	vdso_addr = (unsigned char *)getauxval(AT_SYSINFO_EHDR);

	if (!vdso_addr) {

		D((void)fprintf(stderr,

			"get_linux_vdso_sym: getauxval(AT_SYSINFO_EHDR) failed.\n"));

		return NULL;

	}

	elf_header = (ElfXX_Ehdr *)vdso_addr;

	section_header = (ElfXX_Shdr *)(vdso_addr + elf_header->e_shoff);

	for (dynstr = NULL, i=0L; i < elf_header->e_shnum; i++) {

		s    = &section_header[i];

		s2   = &section_header[elf_header->e_shstrndx];

		name = (const char*)(vdso_addr + s2->sh_offset + s->sh_name);

		if (strcmp(name, ".dynstr") == 0) {

			dynstr = (const char *)(vdso_addr + s->sh_offset);

			break;

		}

	}

	if (!dynstr)

		return NULL;

	for (i=0L; i < elf_header->e_shnum; i++) {

		s    = &section_header[i];

		s2   = &section_header[elf_header->e_shstrndx];

		name = (const char*)(vdso_addr + s2->sh_offset + s->sh_name);

		/* Not .dynsym ? Then look at next section... */

		if (strcmp(name, ".dynsym"))

			continue;

		for (si=0L; si < (s->sh_size/s->sh_entsize); si++) {

			sym = &(((ElfXX_Sym*)(vdso_addr + s->sh_offset))[si]);

			name = dynstr + sym->st_name;

			D((void)fprintf(stderr, "get_linux_vdso_sym: "

				"symbol=|%s| vs. |%s|\n",

				funcname, name));

			if (!strcmp(name, funcname)) {

				ret = (void *)(vdso_addr + sym->st_value);

				break;

			}

		}

		if (ret)

			break;

	}

	return ret;

}

/*

 * Initalisation code

 */

/*

 * Note that the |getcpu()| prototype here differs between

 * Linux versions, the third argument (|tcache|) has been

 * removed, and the prototype was moved from

 * <linux/getcpu.h> to <bits/sched.h>.

 */

typedef int (getcpu_func_t)(unsigned *cpu, unsigned *node, void *tcache);

static pthread_once_t my_getcpu_func_is_initialized = PTHREAD_ONCE_INIT;

static getcpu_func_t *my_getcpu_func = NULL;

static

void my_getcpu_func_init(void)

{

	my_getcpu_func = (getcpu_func_t *)get_linux_vdso_sym("__vdso_getcpu");

	if (!my_getcpu_func)

		my_getcpu_func = (getcpu_func_t *)get_linux_vdso_sym("getcpu");

	D((void)fprintf(stderr, "my_getcpu_func_init: "

		"my_getcpu_func = 0x%lx (%s)\n",

		(long)my_getcpu_func,

		(my_getcpu_func?

			"Using fast codepath":

			"Using slow codepath")));

	if (!my_getcpu_func) {

		(void)fprintf(stderr, "my_getcpu_func_init: "

			"WARNING: aux symbol for getcpu not found"

			", using slow libc codepath.\n");

	}

}

/*

 * |my_sched_getcpu()| - our "public" interface

 *

 * Notes:

 * - The first call might be slower because the call to

 * |my_getcpu_func_init()|

 */

int my_sched_getcpu(void)

{

	unsigned int	cpu;

	int		s;

	(void)pthread_once(&my_getcpu_func_is_initialized,

		my_getcpu_func_init);

	if (my_getcpu_func) {

		/* our MY_getcpu_TESTS own "fast" codepath */

		s = (*my_getcpu_func)(&cpu, NULL, NULL);

	}

	else

	{

		/*

		 * libc codepath, whatever they use - might be VDSO,

		 * might be syscall

		 *

		 * Note that the prototype here differs between Linux

		 * versions, the third argument (|tcache|) has been

		 * removed, and the prototype was moved from

		 * <linux/getcpu.h> to <bits/sched.h>.

		 * Depending on glibc version this might cause

		 * crashes.

		 */

		s = getcpu(&cpu, NULL/*, NULL*/);

	}

	if (s == -1)

		return -1;

	return (int)cpu;

}

#ifdef MY_SCHED_GETCPU_TESTS

/*

 * test code

 */

static

void test_my_sched_getcpu1(void)

{

	int cpu_num;

	cpu_num = my_sched_getcpu();

	(void)printf("my_sched_getcpu() = %d\n", cpu_num);

	cpu_num = sched_getcpu();

	(void)printf("sched_getcpu()    = %d\n", cpu_num);

}

int main(int ac, char *av[])

{

	(void)ac;

	(void)av;

	test_my_sched_getcpu1();

	return EXIT_SUCCESS;

}

#endif /* MY_SCHED_GETCPU_TESTS */