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 PROGRAM_NAME "wdye"
#define PROGRAM_VERSION "1"

#define LIBERTY_WANT_POLLER
#include "../../liberty.c"

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

#include <termios.h>
#if defined SOLARIS
#include <stropts.h>
#endif

#include <curses.h>
#include <term.h>

// --- Pseudoterminal ----------------------------------------------------------
// This is largely taken from Advanced Programming in the UNIX® Environment,
// just without a bunch of bugs.

static int
ptym_open (char **pts_name)
{
	int fdm = -1, err = 0;
	if ((fdm = posix_openpt (O_RDWR | O_NOCTTY)) < 0)
		return -1;
	if (grantpt (fdm) < 0
	 || unlockpt (fdm) < 0)
		goto errout;

	char *ptr = NULL;
	if ((ptr = ptsname (fdm)) == NULL)
		goto errout;

	cstr_set (pts_name, xstrdup (ptr));
	return fdm;

errout:
	err = errno;
	xclose (fdm);
	errno = err;
	return -1;
}

static int
ptys_open (const char *pts_name)
{
	int fds = -1;
#if defined SOLARIS
	int err = 0, setup = 0;
#endif
	if ((fds = open (pts_name, O_RDWR)) < 0)
		return -1;
#if defined SOLARIS
	if ((setup = ioctl (fds, I_FIND, "ldterm")) < 0)
		goto errout;
	if (setup == 0)
	{
		if (ioctl (fds, I_PUSH, "ptem") < 0
		 || ioctl (fds, I_PUSH, "ldterm") < 0)
			goto errout;

		if (ioctl (fds, I_PUSH, "ttcompat") < 0)
		{
errout:
			err = errno;
			xclose (fds);
			errno = err;
			return -1;
		}
	}
#endif
	return fds;
}

static pid_t
pty_fork (int *ptrfdm, char **slave_name,
	const struct termios *slave_termios, const struct winsize *slave_winsize,
	struct error **e)
{
	int fdm = -1, fds = -1;

	char *pts_name = NULL;
	if ((fdm = ptym_open (&pts_name)) < 0)
	{
		error_set (e, "can’t open master pty: %s", strerror (errno));
		return -1;
	}
	if (slave_name != NULL)
		cstr_set (slave_name, xstrdup (pts_name));

	pid_t pid = fork ();
	if (pid < 0)
	{
		error_set (e, "fork: %s", strerror (errno));
		xclose (fdm);
	}
	else if (pid != 0)
		*ptrfdm = fdm;
	else
	{
		if (setsid () < 0)
			exit_fatal ("setsid: %s", strerror (errno));
		if ((fds = ptys_open (pts_name)) < 0)
			exit_fatal ("can’t open slave pty: %s", strerror (errno));
		xclose (fdm);

#if defined BSD
		if (ioctl (fds, TIOCSCTTY, (char *) 0) < 0)
			exit_fatal ("TIOCSCTTY: %s", strerror (errno));
#endif

		if (slave_termios != NULL
		 && tcsetattr (fds, TCSANOW, slave_termios) < 0)
			exit_fatal ("tcsetattr error on slave pty: %s", strerror (errno));
		if (slave_winsize != NULL
		 && ioctl (fds, TIOCSWINSZ, slave_winsize) < 0)
			exit_fatal ("TIOCSWINSZ error on slave pty: %s", strerror (errno));

		if (dup2 (fds, STDIN_FILENO) != STDIN_FILENO)
			exit_fatal ("dup2 error to stdin");
		if (dup2 (fds, STDOUT_FILENO) != STDOUT_FILENO)
			exit_fatal ("dup2 error to stdout");
		if (dup2 (fds, STDERR_FILENO) != STDERR_FILENO)
			exit_fatal ("dup2 error to stderr");
		if (fds != STDIN_FILENO && fds != STDOUT_FILENO && fds != STDERR_FILENO)
			xclose (fds);
	}
	free (pts_name);
	return pid;
}

// --- Global state ------------------------------------------------------------

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
{
	int terminal_fd;                    ///< Process stdin/stdout/stderr
	pid_t pid;                          ///< Process ID
	int ref_term;                       ///< Terminal information
};

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);

	self->terminal_fd = -1;
	self->pid = -1;
	self->ref_term = LUA_NOREF;
	return self;
}

static int
xlua_process_gc (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	if (self->terminal_fd != -1)
		xclose (self->terminal_fd);
	if (self->pid != -1)
		kill (self->pid, SIGKILL);
	luaL_unref (L, LUA_REGISTRYINDEX, self->ref_term);
	return 0;
}

static int
xlua_process_expect (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	// TODO(p): Arguments.
	// TODO(p): Keep reading input into a buffer,
	//   until there's a match.
	//    - It seems we need to keep the buffer within the process object,
	//      because we can read too much.
	return 0;
}

static int
xlua_process_send (lua_State *L)
{
	struct process *self = luaL_checkudata (L, 1, XLUA_PROCESS_METATABLE);
	int nargs = lua_gettop (L);
	for (int i = 2; i <= nargs; i++)
		if (!lua_isstring (L, i))
			return luaL_error (L, "need string arguments");

	for (int i = 2; i <= nargs; i++)
	{
		size_t len = 0;
		const char *arg = lua_tolstring (L, i, &len);
		ssize_t written = write (self->terminal_fd, arg, len);
		if (written == -1)
			return luaL_error (L, "write failed: %s", strerror (errno));
		else if (written != len)
			return luaL_error (L, "write failed: %s", "short write");
	}
	lua_pushvalue (L, 1);
	return 1;
}

static luaL_Reg xlua_process_table[] =
{
	{ "__gc",       xlua_process_gc       },
	{ "expect",     xlua_process_expect   },
	{ "send",       xlua_process_send     },
	{ NULL,         NULL                  }
};

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

static bool
load_terminfo (const char *term, struct str_map *strings)
{
	// TODO(p): See if we can get away with using a bogus FD.
	int err = 0;
	TERMINAL *saved_term = set_curterm (NULL);
	if (setupterm ((char *) term, STDOUT_FILENO, &err) != OK)
	{
		set_curterm (saved_term);
		return false;
	}

	for (size_t i = 0; strfnames[i]; i++)
	{
		const char *value = tigetstr (strnames[i]);
		if (value && value != (char *) -1)
			str_map_set (strings, strfnames[i], xstrdup (value));
	}
	del_curterm (set_curterm (saved_term));
	return true;
}

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

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

	struct error *error;                ///< Error
};

static struct xlua_spawn_context
xlua_spawn_context_make (void)
{
	struct xlua_spawn_context self = {};
	self.env = str_map_make (free);
	self.term = str_map_make (free);

	// XXX: It might make sense to enable starting from an empty environment.
	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 ();
	return self;
}

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

	if (self->error)
		error_free (self->error);
}

// -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 process environment.
	if (xlua_getfield (L, 1, "environ", LUA_TTABLE, true))
	{
		environ_map_update (&ctx->env, L);
		lua_pop (L, 1);
	}
	char *term = str_map_find (&ctx->env, "TERM");
	if (!term)
	{
		print_debug ("setting a default TERM");
		str_map_set (&ctx->env, "TERM", (term = xstrdup ("dumb")));
	}
	environ_map_serialize (&ctx->env, &ctx->envv);

	// Step 2: Load terminal information.
	if (!load_terminfo (term, &ctx->term))
		luaL_error (L, "failed to initialize terminfo for %s", term);

	// 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: Create a process object.
	// This will get garbage collected as appropriate on failure.
	struct process *process = xlua_process_new (L);

	// This could be made into an object that can adjust winsize/termios.
	lua_createtable (L, 0, ctx->term.len);
	struct str_map_iter iter = str_map_iter_make (&ctx->term);
	const char *value = NULL;
	while ((value = str_map_iter_next (&iter)))
	{
		lua_pushstring (L, iter.link->key);
		lua_pushstring (L, value);
		lua_settable (L, -3);
	}
	process->ref_term = luaL_ref (L, LUA_REGISTRYINDEX);

	// Step 5: Spawn the process, which gets a new process group.
	process->pid =
		pty_fork (&process->terminal_fd, NULL, NULL, NULL, &ctx->error);
	if (process->pid < 0)
	{
		return luaL_error (L, "failed to spawn %s: %s",
			ctx->argv.vector[0], ctx->error->message);
	}
	if (!process->pid)
	{
		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);
	}
	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);
	}
	lua_pop (g.L, 1);
	lua_close (g.L);

	poller_free (&g.poller);
	return 0;
}