Why this Lua program that limits memory usage crashes?

I'm working on Exercise 27.4 from the 4th edition of Programming in Lua. I am using Lua 5.4.4, and my solution is the following.

/*
Exercise 27.4: Write a library that allows a script to limit the total amount of memory used
by its Lua state. It may offer a single function, `setlimit`, to set that limit.

The library should set its own allocation function. This function, before calling the
original allocator, checks the total memory in use and returns NULL if the requested memory
exceeds the limit.

(Hint: the library can use the user data of the allocation function to keep its state:
the byte count, the current memory limit, etc.; remember to use the original user data
when calling the original allocation function.)
*/

#include "lauxlib.h"
#include <stdlib.h>

typedef struct {
  void *o_ud;
  lua_Alloc o_alloc;
  size_t mem;
  size_t limit;
} memlimit_ud;

static void *memlimit_alloc(void *ud, void *ptr, size_t osize, size_t nsize) {
  memlimit_ud *m_ud = ud;
  size_t mem        = m_ud->mem;
  if (mem + nsize - osize > m_ud->limit) {
    return NULL;
  }
  void *res = m_ud->o_alloc(m_ud->o_ud, ptr, osize, nsize);
  if (res) m_ud->mem += nsize - osize;
  return res;
}

static int l_setlimit(lua_State *L) {
  int n = luaL_checkinteger(L, 1) * 1024;

  memlimit_ud *ud;
  lua_getallocf(L, (void **) &ud);

  ud->limit = n;
  return 0;
}

static const struct luaL_Reg memlimit[] = {
    {"setlimit", l_setlimit},
    {      NULL,       NULL},
};

int luaopen_memlimit(lua_State *L) {
  memlimit_ud *ud = lua_newuserdata(L, sizeof(memlimit_ud));
  luaL_ref(L, LUA_REGISTRYINDEX);

  ud->o_alloc = lua_getallocf(L, &ud->o_ud);
  ud->mem     = 0;
  ud->limit   = -1;

  lua_setallocf(L, memlimit_alloc, ud);

  luaL_newlib(L, memlimit);
  return 1;
}

However, it throws a segment fault when the script ends. I found this question, and have tried to change lua_newuserdata to malloc, but it still doesn't work. While debugging, I noticed that it was trying to call memlimit_alloc before it crashed, but the function seemed like didn't exist (the former calls to this function were successful).

How can I fix this?