Making LD_PRELOAD/LD_AUDIT modules with Python

Makefile

.PHONY: all clean test

PYTHON_CONFIG = python3-config

CC = gcc -std=gnu11

CFLAGS   = `$(PYTHON_CONFIG) --cflags`
INCLUDES = `$(PYTHON_CONFIG) --includes`
LIBS     = `$(PYTHON_CONFIG) --libs`
LD_PATH  = `$(PYTHON_CONFIG) --prefix`/lib

CYTHON_TARGETS = malloc
TARGETS = preload
SHARED_OBJS = $(addsuffix .so,$(TARGETS)) $(addsuffix .so,$(CYTHON_TARGETS))

all: $(SHARED_OBJS)

#preload.so: preload.c
	#$(CC) -Wl,--export-dynamic -fPIC -shared -o $@ $^ -ldl

%.so: %.c
	$(CC) $(CFLAGS) $(INCLUDES) -L$(LD_PATH) -Wl,--export-dynamic -fPIC -shared -o $@ $^ -ldl $(LIBS)

%.c: %.pyx
	cython $^

prog: prog.c
	$(CC) -o $@ $^ -lm

test: $(SHARED_OBJS)
	./preload malloc ls

clean:
	rm -f $(SHARED_OBJS) $(addsuffix .c,$(CYTHON_TARGETS)) $(addsuffix .h,$(CYTHON_TARGETS))

malloc.pyx

from libc.stdlib cimport malloc as libc_malloc

cdef public void* malloc(size_t x):
    print('malloc({})'.format(x))
    return libc_malloc(x)

preload

#/bin/sh

PY_PRELOAD="$1"
shift 1
env LD_AUDIT=./preload.so PY_PRELOAD="$PY_PRELOAD" $@

preload.c

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

#define __USE_GNU
#define _GNU_SOURCE

#include "Python.h"
#include <link.h>
#include <dlfcn.h>

static void __attribute__ ((destructor)) fini(void) {
    Py_Finalize();
}

static void __attribute__ ((constructor)) init(void) {
    Py_Initialize();
    PyRun_SimpleStringFlags("from sys import path; from os import getcwd; path.insert(0, getcwd())", NULL);
    atexit(fini);
}

unsigned int la_version(unsigned int version) {
    return version;
}

uintptr_t la_symbind64(Elf64_Sym *sym, unsigned int ndx, uintptr_t *refcook,
                       uintptr_t *defcook, unsigned int *flags, const char *symname) {
    if (0 == strcmp(symname, "malloc")) {
        const char *modname = getenv("PY_PRELOAD");
        if (!modname) goto error;

        PyObject *mod = PyImport_ImportModule(modname);
        if (!mod) goto error;
        PyObject *modfile = PyObject_GetAttrString(mod, "__file__");
        if (!modfile) goto error;
        char *filename = PyUnicode_AsUTF8(modfile);
        if (!filename) goto error;
        void* obj = dlopen(filename, RTLD_LAZY | RTLD_LOCAL);
        if (!obj) goto error;
        uintptr_t sym = (uintptr_t)dlsym(obj, symname);
        if (!sym) goto error;

        return (uintptr_t)sym;
    }
error:
    return sym->st_value;
}

unsigned int la_objopen(struct link_map *map, Lmid_t lmid, uintptr_t *cookie) {
   return LA_FLG_BINDTO | LA_FLG_BINDFROM;
}

prog.c

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

int main(int argc, char *argv[]) {
  for (int i = 1; i < argc; ++i) {
    int arg = atoi(argv[i]);
    void *rv = malloc(arg);
    printf("%s(%d) -> %p\n", argv[0], arg, rv);
  }

  exit(EXIT_SUCCESS);
  return 0;
}