/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 */

/*
 * Copyright 2020 Joyent, Inc.
 * Copyright 2022 Jason King
 */

#include <sys/debug.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/avl.h>
#include <sys/fcntl.h>
#include <sys/sysmacros.h>
#include <dirent.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <gelf.h>
#include <libcustr.h>
#include <libelf.h>
#include <libgen.h>
#include <limits.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

static inline bool
is_same(const struct stat *s1, const struct stat *s2)
{
	if ((s1->st_ino == s2->st_ino) && (s1->st_dev == s2->st_dev))
		return (true);
	return (false);
}

typedef enum dir_flags {
	DF_NONE =	0,
	DF_IS_SELF =	1 << 0,
	DF_IS_SYMLINK =	1 << 2,
} dir_flags_t;

typedef struct path {
	custr_t	*p_name;
	size_t	p_pfxidx;
} path_t;

typedef struct name {
	char	*n_name;
	bool	n_is_symlink;
} name_t;

typedef struct names {
	name_t	*ns_names;
	uint_t	ns_num;
	uint_t	ns_alloc;
} names_t;

typedef struct elfinfo {
	int		ei_class;
	uint16_t	ei_type;
	bool		ei_hasverdef;
} elfinfo_t;

typedef struct obj {
	avl_node_t	obj_avlid;
	avl_node_t	obj_avlname;
	dev_t		obj_dev;
	ino_t		obj_inode;
	names_t		obj_names;
	elfinfo_t	obj_elfinfo;
} obj_t;

static void path_init(path_t *, const char *, bool);
static size_t path_append(path_t *, const char *);
static const char *path_name(const path_t *);
static const char *path_fullpath(const path_t *);
static void path_pop(path_t *, size_t);

static bool maybe_obj(const char *, mode_t);
static bool get_elfinfo(const path_t *, int, elfinfo_t *);
static void add_name(obj_t *, const path_t *, bool);
static int cmp_id(const void *, const void *);
static int cmp_objname(const void *, const void *);
static int cmp_names(const void *, const void *);

static void process_dir(path_t *, int, const struct stat *, dir_flags_t);
static void process_file(path_t *, int, const struct stat *, bool);
static void process_arg(char *);

static void sort_names(void *, void *);
static void print_entry(void *, void *);

static void foreach_avl(avl_tree_t *, void (*)(void *, void *), void *);

static void nomem(void);
static char *xstrdup(const char *);
static void *xcalloc(size_t, size_t);

static avl_tree_t avl_byid;
static avl_tree_t avl_byname;

static const char *special_paths[] = {
	"usr/bin/alias",
	"usr/lib/isaexec",
};

static int rootfd = -1;

/* By default, we process aliases */
static bool do_alias = true;

/* By default, we treat certain well known names (e.g. isaexec) as special */
static bool do_special = true;

/* fast_mode, relpath, and so_only are all false by default */
static bool fast_mode;
static bool relpath;
static bool so_only;

static void __NORETURN
usage(const char *name)
{
	(void) fprintf(stderr,
	    "Usage: %s [-afnrs] file | dir\n"
	    "\t[-a]\texpand symlink aliases\n"
	    "\t[-f]\tuse file name at mode to speed search\n"
	    "\t[-h]\tshow this help\n"
	    "\t[-n]\tdon\'t treat well known paths as special in sorting\n"
	    "\t[-r]\treport relative paths\n"
	    "\t[-s]\tonly remote shareable (ET_DYN) objects\n", name);
	exit(EXIT_FAILURE);
}

int
main(int argc, char **argv)
{
	int c;

	if (elf_version(EV_CURRENT) == EV_NONE)
		errx(EXIT_FAILURE, "elf library is out of date");

	while ((c = getopt(argc, argv, "ahfnrs")) != -1) {
		switch (c) {
		case 'a':
			do_alias = false;
			break;
		case 'f':
			fast_mode = true;
			break;
		case 'n':
			do_special = false;
			break;
		case 'r':
			relpath = true;
			break;
		case 's':
			so_only = true;
			break;
		case 'h':
			usage(argv[0]);
		case '?':
			(void) fprintf(stderr, "Unknown option -%c\n", optopt);
			usage(argv[0]);
		}
	}

	if (optind == argc) {
		(void) fprintf(stderr, "Missing file or dir parameter\n");
		usage(argv[0]);
	}

	if (argv[optind][0] == '\0')
		errx(EXIT_FAILURE, "Invalid file or dir value");

	avl_create(&avl_byid, cmp_id, sizeof (obj_t),
	    offsetof(obj_t, obj_avlid));
	avl_create(&avl_byname, cmp_objname, sizeof (obj_t),
	    offsetof(obj_t, obj_avlname));

	process_arg(argv[optind]);
	foreach_avl(&avl_byid, sort_names, &avl_byname);
	foreach_avl(&avl_byname, print_entry, NULL);

	return (EXIT_SUCCESS);
}

static void
process_arg(char *arg)
{
	path_t path;
	struct stat sb;
	int fd;

	if ((fd = open(arg, O_RDONLY)) == -1) {
		err(EXIT_FAILURE, "could not open %s", arg);
	}

	if (fstat(fd, &sb) < 0) {
		err(EXIT_FAILURE, "failed to stat %s", arg);
	}

	if (S_ISDIR(sb.st_mode)) {
		path_init(&path, arg, relpath);
		if (relpath) {
			(void) printf("PREFIX %s\n", arg);
		}

		rootfd = fd;
		/* process_dir() will close fd */
		process_dir(&path, fd, &sb, DF_NONE);
		return;
	}

	char *argcopy = xstrdup(arg);
	char *dir = dirname(argcopy);

	if (!S_ISREG(sb.st_mode)) {
		err(EXIT_FAILURE, "not a file or directory: %s", arg);
	}

#ifndef O_DIRECTORY
	struct stat tsb;
	if (stat(dir, &tsb) == -1) {
		err(EXIT_FAILURE, "failed to stat %s", dir);
	}
	if (!S_ISDIR(tsb.st_mode)) {
		errx(EXIT_FAILURE, "not a directory: %s", dir);
	}
	rootfd = open(dir, O_RDONLY);
#else
	rootfd = open(dir, O_RDONLY|O_DIRECTORY);
#endif
	if (rootfd < 0) {
		err(EXIT_FAILURE, "%s", dir);
	}

	path_init(&path, dir, relpath);
	if (relpath) {
		(void) printf("PREFIX %s\n", dir);
	}
	free(argcopy);

	process_file(&path, fd, &sb, DF_NONE);
}

/*
 * Processing a directory has some subtleties. When we encounter a
 * subdirectory that's a symlink, we want any files we encounter when
 * we traverse it to be treated as aliases. We may also have a symlink that's
 * a link back to ourself (e.g. 32 -> .). In this special case, we still want
 * to reprocess the directory to generate alias names, but we use the
 * DFLAG_SELF flag to avoid recursing forever.
 *
 * A limitation of this approach is that we cannot handle a loop over multiple
 * levels (e.g. foo -> ../..). Nothing currently in illumos-gate creates any
 * such symlinks in the proto area, so we've opted to avoid complicating
 * processing further to handle scenarios that aren't realistically expected
 * to occur.
 *
 * Note that dirfd is always closed upon return from process_dir().
 */
static void
process_dir(path_t *p, int dirfd, const struct stat *dirsb, dir_flags_t dflags)
{
	DIR *d;
	struct dirent *de;

	d = fdopendir(dirfd);
	if (d == NULL) {
		warn("opendir(%s)", path_fullpath(p));
		VERIFY0(close(dirfd));
		return;
	}

	while ((de = readdir(d)) != NULL) {
		struct stat sb;
		int fd;
		bool is_link = false;
		size_t plen;

		if (strcmp(de->d_name, ".") == 0 ||
		    strcmp(de->d_name, "..") == 0) {
			continue;
		}

		plen = path_append(p, de->d_name);

		if (fstatat(rootfd, path_name(p), &sb,
		    AT_SYMLINK_NOFOLLOW) < 0) {
			warn("%s", path_fullpath(p));
			path_pop(p, plen);
			continue;
		}

		fd = openat(dirfd, de->d_name, O_RDONLY);
		if (fd < 0) {
			/*
			 * Symlinks in the proto area may point to a path
			 * that's not present (e.g. /dev/stdin -> fd/0).
			 * Silently skip such entries.
			 */
			if (errno != ENOENT || !S_ISLNK(sb.st_mode)) {
				warn("%s", path_fullpath(p));
			}
			path_pop(p, plen);
			continue;
		}

		if (S_ISLNK(sb.st_mode)) {
			is_link = true;

			if (fstat(fd, &sb) < 0) {
				warn("stat %s", path_fullpath(p));
				path_pop(p, plen);
				continue;
			}
		}

		if (S_ISDIR(sb.st_mode)) {
			if ((dflags & DF_IS_SELF) != 0) {
				VERIFY0(close(fd));
				path_pop(p, plen);
				continue;
			}

			dir_flags_t newflags = dflags;

			/* The recursive process_dir() call closes fd */
			if (is_link)
				newflags |= DF_IS_SYMLINK;
			if (is_same(&sb, dirsb))
				newflags |= DF_IS_SELF;
			process_dir(p, fd, &sb, newflags);
		} else if (S_ISREG(sb.st_mode)) {
			if (!maybe_obj(de->d_name, sb.st_mode)) {
				VERIFY0(close(fd));
				path_pop(p, plen);
				continue;
			}

			if ((dflags & (DF_IS_SELF | DF_IS_SYMLINK)) != 0)
				is_link = true;

			/* process_file() closes fd */
			process_file(p, fd, &sb, is_link);
		}

		path_pop(p, plen);
	}

	/* Closes dirfd */
	VERIFY0(closedir(d));
}

/* Note that fd is always closed upon return */
static void
process_file(path_t *p, int fd, const struct stat *sb, bool is_link)
{
	avl_index_t where = { 0 };
	obj_t templ = {
		.obj_dev = sb->st_dev,
		.obj_inode = sb->st_ino,
	};
	obj_t *obj = avl_find(&avl_byid, &templ, &where);
	elfinfo_t elfinfo = { 0 };

	if (obj != NULL)
		goto done;

	if (!get_elfinfo(p, fd, &elfinfo)) {
		VERIFY0(close(fd));
		return;
	}

	obj = xcalloc(1, sizeof (*obj));
	obj->obj_dev = sb->st_dev;
	obj->obj_inode = sb->st_ino;
	obj->obj_elfinfo = elfinfo;
	avl_add(&avl_byid, obj);

done:
	add_name(obj, p, is_link);
	VERIFY0(close(fd));
}

static void
print_entry(void *a, void *arg __unused)
{
	obj_t *obj = a;
	const char *objname = obj->obj_names.ns_names[0].n_name;
	const char *bits = "";
	const char *type = "";
	const char *verdef = obj->obj_elfinfo.ei_hasverdef ?
	    "VERDEF" : "NOVERDEF";

	switch (obj->obj_elfinfo.ei_class) {
	case ELFCLASS32:
		bits = "32";
		break;
	case ELFCLASS64:
		bits = "64";
		break;
	default:
		errx(EXIT_FAILURE, "unknown elfclass value %x for %s",
		    obj->obj_elfinfo.ei_class, objname);
	}

	switch (obj->obj_elfinfo.ei_type) {
	case ET_REL:
		type = "REL";
		break;
	case ET_DYN:
		type = "DYN";
		break;
	case ET_EXEC:
		type = "EXEC";
		break;
	default:
		errx(EXIT_FAILURE, "unexpected elf type %x for %s",
		    obj->obj_elfinfo.ei_type, objname);
	}

	names_t *names = &obj->obj_names;

	VERIFY3U(names->ns_num, >, 0);
	VERIFY(!names->ns_names[0].n_is_symlink);

	(void) printf("OBJECT %2s %-4s %-8s %s\n", bits, type, verdef,
	    objname);

	for (uint_t i = 1; i < names->ns_num; i++) {
		if (do_alias) {
			(void) printf("%-23s %s\t%s\n",
			    "ALIAS", objname, names->ns_names[i].n_name);
		} else {
			(void) printf("OBJECT %2s %-4s %-8s %s\n", bits, type,
			    verdef, names->ns_names[i].n_name);
		}
	}
}

/*
 * Returns true and eip populated if name was an elf object, otherwise
 * returns false.
 */
static bool
get_elfinfo(const path_t *p, int fd, elfinfo_t *eip)
{
	Elf *elf = NULL;
	Elf_Scn *scn = NULL;
	GElf_Ehdr ehdr = { 0 };
	int eval;

	if ((elf = elf_begin(fd, ELF_C_READ, NULL)) == NULL)
		goto fail_noend;

	if ((eip->ei_class = gelf_getclass(elf)) == ELFCLASSNONE) {
		VERIFY0(elf_end(elf));
		return (false);
	}

	if (gelf_getehdr(elf, &ehdr) == NULL)
		goto fail;

	eip->ei_type = ehdr.e_type;
	eip->ei_hasverdef = false;

	while ((scn = elf_nextscn(elf, scn)) != NULL) {
		Elf_Data *data = NULL;
		GElf_Shdr shdr = { 0 };

		if (gelf_getshdr(scn, &shdr) == NULL)
			goto fail;

		if (shdr.sh_type != SHT_DYNAMIC)
			continue;

		if ((data = elf_getdata(scn, NULL)) == NULL)
			continue;

		size_t nent = shdr.sh_size / shdr.sh_entsize;

		for (size_t i = 0; i < nent; i++) {
			GElf_Dyn dyn = { 0 };

			if (gelf_getdyn(data, i, &dyn) == NULL) {
				goto fail;
			}

			if (dyn.d_tag == DT_VERDEF) {
				eip->ei_hasverdef = true;
				break;
			}
		}
	}

	VERIFY0(elf_end(elf));
	return (true);

fail:
	VERIFY0(elf_end(elf));

fail_noend:
	eval = elf_errno();

	warnx("%s: %s", path_fullpath(p), elf_errmsg(eval));
	return (false);
}

static bool
is_special(const char *name)
{
	for (uint_t i = 0; i < ARRAY_SIZE(special_paths); i++) {
		if (strcmp(special_paths[i], name) == 0)
			return (true);
	}
	return (false);
}

static void
sort_names(void *a, void *b)
{
	obj_t *obj = a;
	avl_tree_t *name_avl = b;
	names_t *names = &obj->obj_names;

	/* We should always have at least one name */
	VERIFY3U(names->ns_num, >, 0);

	/* If there is only one, they get the prize and we're done */
	if (names->ns_num == 1)
		return;

	name_t *first = NULL;

	/*
	 * Find the first (in sorted order) pathname for this object
	 * that isn't a symlink.
	 */
	for (uint_t i = 0; i < names->ns_num; i++) {
		name_t *n = &names->ns_names[i];

		if (n->n_is_symlink)
			continue;

		if (first == NULL) {
			first = n;
			continue;
		}

		/*
		 * If we're treating isaexec as special, we always
		 * want it to be the first entry. Otherwise, pick
		 * the 'lowest' sorted pathname.
		 */
		if (do_special) {
			if (is_special(n->n_name)) {
				first = n;
				break;
			}

			if (strcmp(n->n_name, first->n_name) < 0)
				first = n;
		}
	}

	/*
	 * If the first pathname (in sorted order) isn't the first
	 * name entry, we swap it into first place (while also updating
	 * the names AVL tree).
	 */
	if (first != NULL && first != &names->ns_names[0]) {
		name_t tmp = names->ns_names[0];

		avl_remove(name_avl, obj);
		(void) memcpy(&names->ns_names[0], first, sizeof (name_t));
		(void) memcpy(first, &tmp, sizeof (name_t));
		avl_add(name_avl, obj);
	}

	/*
	 * The remaining names (symlinks or not) are sorted by
	 * their pathname.
	 */
	qsort(&names->ns_names[1], names->ns_num - 1, sizeof (name_t),
	    cmp_names);
}

/*
 * We grow the names array by NAME_CHUNK entries every time we need to
 * extend it.
 */
#define	NAME_CHUNK	4

static name_t *
name_new(names_t *names)
{
	if (names->ns_num < names->ns_alloc)
		return (&names->ns_names[names->ns_num++]);

	name_t *newn = NULL;
	uint_t newamt = names->ns_alloc + NAME_CHUNK;

	/*
	 * Since we may be building on a machine that doesn't
	 * have reallocarray or the like, we avoid their use here.
	 * If we ever officially designate an illumos-gate build
	 * as the minimum required to build master that contains
	 * reallocarray and such, we can replace most of this logic
	 * with it.
	 *
	 * Also use xcalloc so we get the unused entries zeroed already.
	 * Not strictly necessary, but useful for debugging.
	 */
	newn = xcalloc(newamt, sizeof (name_t));

	/* Move over existing entries */
	(void) memcpy(newn, names->ns_names, names->ns_num * sizeof (name_t));

	free(names->ns_names);

	names->ns_names = newn;
	names->ns_alloc = newamt;
	return (&names->ns_names[names->ns_num++]);
}

static void
add_name(obj_t *obj, const path_t *p, bool is_symlink)
{
	names_t *ns = &obj->obj_names;
	const char *srcname = path_name(p);

	/* We should never have duplicates */
	for (size_t i = 0; i < ns->ns_num; i++)
		VERIFY3S(strcmp(ns->ns_names[i].n_name, srcname), !=, 0);

	name_t *n = name_new(ns);

	n->n_name = xstrdup(srcname);
	n->n_is_symlink = is_symlink;

	if (is_symlink)
		return;

	/*
	 * If the name was not a symlink, first determine if there are
	 * existing (hard) links to this name already, and if so, which entry
	 * is the first one. Typically this will be the name we just added
	 * unless the file does have multiple hard links (e.g. isaexec).
	 */
	uint_t nhlink = 0;
	uint_t firsthlink = UINT_MAX;

	for (uint_t i = 0; i < ns->ns_num; i++) {
		if (ns->ns_names[i].n_is_symlink)
			continue;
		if (nhlink == 0)
			firsthlink = i;
		nhlink++;
	}

	if (nhlink > 1)
		return;

	/*
	 * Put the first non-symlink name as the very first
	 * entry.
	 */
	VERIFY3U(firsthlink, !=, UINT_MAX);

	if (firsthlink != 0) {
		name_t tmp = {
			.n_name = ns->ns_names[0].n_name,
			.n_is_symlink = ns->ns_names[0].n_is_symlink,
		};

		(void) memcpy(&ns->ns_names[0], &ns->ns_names[firsthlink],
		    sizeof (name_t));
		(void) memcpy(&ns->ns_names[firsthlink], &tmp, sizeof (name_t));
	}

	avl_add(&avl_byname, obj);
}

/*
 * This is an arbitrary initial value -- basically we can nest 16 directories
 * deep before having to grow p->p_idx (which seems like a nice power of 2).
 */
#define	PATH_INIT	16

static void
path_init(path_t *p, const char *name, bool relpath)
{
	(void) memset(p, '\0', sizeof (*p));

	if (custr_alloc(&p->p_name) < 0) {
		nomem();
	}

	if (name != NULL && custr_append(p->p_name, name) < 0)
		nomem();

	size_t len = custr_len(p->p_name);

	/* Trim off any trailing /'s, but don't trim '/' to an empty path */
	while (len > 1 && custr_cstr(p->p_name)[len - 1] == '/') {
		VERIFY0(custr_rtrunc(p->p_name, 0));
		len--;
	}

	p->p_pfxidx = relpath ? len + 1 : 0;
}

static size_t
path_append(path_t *p, const char *name)
{
	size_t clen = custr_len(p->p_name);

	if (clen > 0)
		VERIFY0(custr_appendc(p->p_name, '/'));
	VERIFY0(custr_append(p->p_name, name));

	return (clen);
}

static const char *
path_name(const path_t *p)
{
	return (custr_cstr(p->p_name) + p->p_pfxidx);
}

static const char *
path_fullpath(const path_t *p)
{
	return (custr_cstr(p->p_name));
}

static void
path_pop(path_t *p, size_t idx)
{
	VERIFY0(custr_trunc(p->p_name, idx));
}

static int
cmp_id(const void *a, const void *b)
{
	const obj_t *l = a;
	const obj_t *r = b;

	if (l->obj_dev < r->obj_dev)
		return (-1);
	if (l->obj_dev > r->obj_dev)
		return (1);
	if (l->obj_inode < r->obj_inode)
		return (-1);
	if (l->obj_inode > r->obj_inode)
		return (1);
	return (0);
}

static int
cmp_objname(const void *a, const void *b)
{
	const obj_t *l = a;
	const obj_t *r = b;
	const name_t *ln = &l->obj_names.ns_names[0];
	const name_t *rn = &r->obj_names.ns_names[0];

	return (cmp_names(ln, rn));
}

static int
cmp_names(const void *a, const void *b)
{
	const name_t *l = a;
	const name_t *r = b;
	int cmp = strcmp(l->n_name, r->n_name);

	if (cmp < 0)
		return (-1);
	if (cmp > 0)
		return (1);
	return (0);
}

/*
 * Use the filename and permission bits to determine if this file could
 * possibly be an executable.
 */
static bool
maybe_obj(const char *name, mode_t mode)
{
	/* If not in fast mode, check everything, so always return true */
	if (!fast_mode)
		return (true);

	size_t len = strlen(name);

	/* If the file name ends in .so, we check */
	if (len >= 3 && strcmp(&name[len - 3], ".so") == 0) {
		return (true);
	}

	/* If the file name contains '.so', we check */
	if (len >= 4 && strstr(name, ".so.") != NULL) {
		return (true);
	}

	/* If the above checks fail, we assume it's not a shared library */
	if (so_only)
		return (false);

	/*
	 * The original perl script used a -x test which only looked at
	 * file permissions. This may return slightly different results
	 * than using access(2) or faccessat(2).
	 */
	if ((mode & (S_IXUSR|S_IXGRP|S_IXOTH)) == 0)
		return (false);

	return (true);
}

static void
foreach_avl(avl_tree_t *avl, void (*cb)(void *, void *), void *arg)
{
	void *obj;

	for (obj = avl_first(avl); obj != NULL; obj = AVL_NEXT(avl, obj))
		cb(obj, arg);
}

static char *
xstrdup(const char *s)
{
	char *news = strdup(s);

	if (news == NULL) {
		nomem();
	}
	return (news);
}

static void *
xcalloc(size_t nelem, size_t elsize)
{
	void *p = calloc(nelem, elsize);

	if (p == NULL) {
		nomem();
	}

	return (p);
}

#define	NOMEM_MSG "out of memory\n"
static void __NORETURN
nomem(void)
{
	/* Try to get the error out before aborting */
	(void) write(STDERR_FILENO, NOMEM_MSG, sizeof (NOMEM_MSG));
	abort();
}