path: root/tools/wdye/wdye.c

/*
 * wdye.c: what did you expect: Lua-based Expect tool
 *
 * Copyright (c) 2025, Přemysl Eric Janouch <p@janouch.name>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 */

#define LIBERTY_WANT_POLLER

#define PROGRAM_NAME "wdye"
#define PROGRAM_VERSION "1"
#include "../../liberty.c"

#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

#include <term.h>

static struct
{
	lua_State *L;                       ///< Lua state
	struct poller poller;               ///< Event loop
}
g;

static int
xlua_error_handler (lua_State *L)
{
	luaL_traceback (L, L, luaL_checkstring (L, 1), 1);
	return 1;
}

static bool
xlua_getfield (lua_State *L, int idx, const char *name,
	int expected, bool optional)
{
	int found = lua_getfield (L, idx, name);
	if (found == expected)
		return true;
	if (optional && found == LUA_TNIL)
		return false;

	const char *message = optional
		? "invalid field \"%s\" (found: %s, expected: %s or nil)"
		: "invalid or missing field \"%s\" (found: %s, expected: %s)";
	return luaL_error (L, message, name,
		lua_typename (L, found), lua_typename (L, expected));
}

// --- Process -----------------------------------------------------------------

#define XLUA_PROCESS_METATABLE "process"

struct process
{
	pid_t pid;                          ///< Process ID
	// TODO(p): File descriptors.
	// TODO(p): Either `TERMINAL *` or read-out values.
};

static struct process *
xlua_process_new (lua_State *L)
{
	struct process *self = lua_newuserdata (L, sizeof *self);
	luaL_setmetatable (L, XLUA_PROCESS_METATABLE);
	memset (self, 0, sizeof *self);
	return self;
}

static int
xlua_process_gc (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	// TODO(p): Clean up file descriptors.
	if (self->pid != -1)
		kill (self->pid, SIGKILL);
	return 0;
}

static int
xlua_process_expect (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	// TODO(p)
	return 0;
}

static int
xlua_process_send (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	// TODO(p)
	return 0;
}

static luaL_Reg xlua_process_table[] =
{
	{ "__gc",       xlua_process_gc       },
//	{ "__index",    xlua_process_index    },
//	{ "__newindex", xlua_process_newindex },
	{ "expect",     xlua_process_expect   },
	{ "send",       xlua_process_send     },
	{ NULL,         NULL                  }
};

// --- Terminal ----------------------------------------------------------------

// TODO(p): Probably create an object to interact with tinfo.

// --- Library -----------------------------------------------------------------

struct xlua_spawn_context
{
	struct str_map env;                 ///< Subprocess environment map
	struct strv envv;                   ///< Subprocess environment vector
	struct strv argv;                   ///< Subprocess argument vector

	pid_t child;                        ///< Child process ID
};

static struct xlua_spawn_context
xlua_spawn_context_make (void)
{
	struct xlua_spawn_context self = {};
	self.env = str_map_make (free);
	for (char **p = environ; *p; p++)
	{
		const char *equals = strchr (*p, '=');
		if (!equals)
			continue;

		char *key = xstrndup (*p, equals - *p);
		str_map_set (&self.env, key, xstrdup (equals + 1));
		free (key);
	}
	self.envv = strv_make ();
	self.argv = strv_make ();

	self.child = -1;
	return self;
}

static void
xlua_spawn_context_free (struct xlua_spawn_context *self)
{
	str_map_free (&self->env);
	strv_free (&self->envv);
	strv_free (&self->argv);

	if (self->child != -1)
		kill (self->child, SIGKILL);
}

// -0, +0, e
static void
environ_map_update (struct str_map *env, lua_State *L)
{
	lua_pushnil (L);
	while (lua_next (L, -2))
	{
		if (lua_type (L, -2) != LUA_TSTRING)
			luaL_error (L, "environment maps must be keyed by strings");

		str_map_set (env, lua_tostring (L, -2), xstrdup (lua_tostring (L, -1)));
		lua_pop (L, 1);
	}
}

// The environment will get pseudo-randomly reordered,
// which is fine by POSIX.
static void
environ_map_serialize (struct str_map *env, struct strv *envv)
{
	struct str_map_iter iter = str_map_iter_make (env);
	const char *value;
	while ((value = str_map_iter_next (&iter)))
		strv_append_owned (envv,
			xstrdup_printf ("%s=%s", iter.link->key, value));
}

// - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

static int
xlua_spawn_protected (lua_State *L)
{
	struct xlua_spawn_context *ctx = lua_touserdata (L, 1);

	// Step 1: Prepare a terminal object.
	(void) xlua_getfield (L, 1, "term", LUA_TSTRING, true);
	const char *term = lua_tostring (L, -1);
	// TODO(p): Process the terminal name,
	//   possibly by creating another indexable Lua object/table.
	//    - Remember to use a default when `term == NULL`.
	lua_pop (L, 1);

	// Step 2: Prepare process environment.
	if (xlua_getfield (L, 1, "environ", LUA_TTABLE, true))
	{
		environ_map_update (&ctx->env, L);
		lua_pop (L, 1);
	}
	environ_map_serialize (&ctx->env, &ctx->envv);

	// Step 3: Prepare process command line.
	size_t argc = lua_rawlen (L, 1);
	for (size_t i = 1; i <= argc; i++)
	{
		lua_pushinteger (L, i);
		lua_rawget (L, 1);
		const char *arg = lua_tostring (L, -1);
		if (!arg)
			return luaL_error (L, "spawn arguments must be strings");

		strv_append (&ctx->argv, arg);
		lua_pop (L, 1);
	}
	if (ctx->argv.len < 1)
		return luaL_error (L, "missing argument");

	// Step 4: Spawn the process.
	// Keeping the same process group for simplicity right now.

	// TODO(p): Allocate a pty, pipes in and out.
	//   See Advanced Programming in the UNIX® Environment.
	switch ((ctx->child = fork ()))
	{
	case -1:
		// TODO(p): Consider luaL_error().
		print_error ("failed to spawn %s: %s",
			ctx->argv.vector[0], strerror (errno));
		break;
	case 0:
		execve (ctx->argv.vector[0], ctx->argv.vector, ctx->envv.vector);
		print_error ("failed to spawn %s: %s",
			ctx->argv.vector[0], strerror (errno));
		_exit (EXIT_FAILURE);
	default:
		break;
	}

	// Step 5: Return a process object.
	struct process *process = xlua_process_new (L);
	process->pid = ctx->child;

	ctx->child = -1;
	return 1;
}

static int
xlua_spawn (lua_State *L)
{
	luaL_checktype (L, 1, LUA_TTABLE);

	lua_pushcfunction (L, xlua_error_handler);
	lua_pushcfunction (L, xlua_spawn_protected);

	// There are way too many opportunities for Lua to throw,
	// so maintain a context to clean up in one go.
	struct xlua_spawn_context ctx = xlua_spawn_context_make ();
	lua_pushlightuserdata (L, &ctx);
	int result = lua_pcall (L, 1, 1, -3);
	xlua_spawn_context_free (&ctx);
	if (result)
		lua_error (L);

	// Remove the error handler ("good programming practice").
	lua_remove (L, -2);
	return 1;
}

static luaL_Reg xlua_library[] =
{
	{ "spawn", xlua_spawn },
	{ NULL,    NULL       }
};

// --- Initialisation, event handling ------------------------------------------

static void *
xlua_alloc (void *ud, void *ptr, size_t o_size, size_t n_size)
{
	(void) ud;
	(void) o_size;

	if (n_size)
		return realloc (ptr, n_size);

	free (ptr);
	return NULL;
}

static int
xlua_panic (lua_State *L)
{
	print_fatal ("Lua panicked: %s", lua_tostring (L, -1));
	lua_close (L);
	exit (EXIT_FAILURE);
	return 0;
}

int
main (int argc, char *argv[])
{
	if (argc != 2)
	{
		fprintf (stderr, "Usage: %s program.lua\n", argv[0]);
		return 1;
	}

	if (!(g.L = lua_newstate (xlua_alloc, NULL)))
		exit_fatal ("Lua initialization failed");
	lua_atpanic (g.L, xlua_panic);
	luaL_openlibs (g.L);
	luaL_checkversion (g.L);

	// TODO(p): Analyse and set up SIGCHLD processing.
	//  - What would this even be useful for?
	//     - For having a global timeout while waiting on a process.

	luaL_newlib (g.L, xlua_library);
	lua_setglobal (g.L, PROGRAM_NAME);

	luaL_newmetatable (g.L, XLUA_PROCESS_METATABLE);
	luaL_setfuncs (g.L, xlua_process_table, 0);
	lua_pop (g.L, 1);

	poller_init (&g.poller);

	const char *path = argv[1];
	lua_pushcfunction (g.L, xlua_error_handler);
	if (luaL_loadfile (g.L, path)
	 || lua_pcall (g.L, 0, 0, -2))
	{
		print_error ("%s: %s", path, lua_tostring (g.L, -1));
		lua_pop (g.L, 1);
	}

	// TODO(p): Do we need to pop the error handler?
	lua_close (g.L);
	return 0;
}