/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2015 Ryan Zezeski
 * Copyright 2004 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

/*
 * pbind - bind a process to a processor (non-exclusively)
 */

#include <sys/types.h>
#include <sys/procset.h>
#include <sys/processor.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <procfs.h>
#include <fcntl.h>
#include <errno.h>
#include <dirent.h>
#include <locale.h>
#include <libproc.h>
#include <stdarg.h>

#if !defined(TEXT_DOMAIN)		/* should be defined by cc -D */
#define	TEXT_DOMAIN 	"SYS_TEST"	/* Use this only if it weren't */
#endif

#define	ERR_OK		0		/* exit status for success */
#define	ERR_FAIL	1		/* exit status for errors */
#define	ERR_USAGE	2		/* exit status for usage errors */

static char	*progname;
static char	bflag;
static char	eflag;
static char	qflag;
static char	Qflag;
static char	uflag;
static char	Uflag;
static int	errors;

#define	MAX_PROCFS_PATH	80

/*PRINTFLIKE1*/
static void
warn(char *format, ...)
{
	int err = errno;
	va_list alist;

	(void) fprintf(stderr, "%s: ", progname);
	va_start(alist, format);
	(void) vfprintf(stderr, format, alist);
	va_end(alist);
	if (strchr(format, '\n') == NULL)
		(void) fprintf(stderr, ": %s\n", strerror(err));
}

/*PRINTFLIKE1*/
static void
die(char *format, ...)
{
	int err = errno;
	va_list alist;

	(void) fprintf(stderr, "%s: ", progname);
	va_start(alist, format);
	(void) vfprintf(stderr, format, alist);
	va_end(alist);
	if (strchr(format, '\n') == NULL)
		(void) fprintf(stderr, ": %s\n", strerror(err));
	exit(ERR_FAIL);
}

/*
 * Output for query.
 */
static void
query_out(id_t pid, id_t lwpid, processorid_t cpu)
{
	char *proclwp;
	char pidstr[21];

	if (lwpid == -1) {
		(void) snprintf(pidstr, 20, "%d", (int)pid);
		proclwp = "process";
	} else {
		(void) snprintf(pidstr, 20, "%d/%d", (int)pid, (int)lwpid);
		proclwp = "lwp";
	}

	if (cpu == PBIND_NONE)
		(void) printf(gettext("%s id %s: not bound\n"),
		    proclwp, pidstr);
	else
		(void) printf(gettext("%s id %s: %d\n"),
		    proclwp, pidstr, cpu);
}

/*
 * Binding error.
 */
static void
bind_err(processorid_t cpu, id_t pid, id_t lwpid, int err)
{
	char *msg;

	switch (cpu) {
	case PBIND_NONE:
		msg = gettext("unbind");
		break;
	case PBIND_QUERY:
		msg = gettext("query");
		break;
	default:
		msg = gettext("bind");
		break;
	}
	if (lwpid == -1)
		warn(gettext("cannot %s pid %d: %s\n"), msg,
		    (int)pid, strerror(err));
	else
		warn(gettext("cannot %s lwpid %d/%d: %s\n"), msg,
		    (int)pid, (int)lwpid, strerror(err));
}

/*
 * Output for bind.
 */
static void
bind_out(id_t pid, id_t lwpid, processorid_t old, processorid_t new)
{
	char *proclwp;
	char pidstr[21];

	if (lwpid == -1) {
		(void) snprintf(pidstr, 20, "%d", (int)pid);
		proclwp = "process";
	} else {
		(void) snprintf(pidstr, 20, "%d/%d", (int)pid, (int)lwpid);
		proclwp = "lwp";
	}

	if (old == PBIND_NONE) {
		if (new == PBIND_NONE)
			(void) printf(gettext("%s id %s: was not bound, "
			    "now not bound\n"), proclwp, pidstr);
		else
			(void) printf(gettext("%s id %s: was not bound, "
			    "now %d\n"), proclwp, pidstr, new);
	} else {
		if (new == PBIND_NONE)
			(void) printf(gettext("%s id %s: was %d, "
			    "now not bound\n"), proclwp, pidstr, old);
		else
			(void) printf(gettext("%s id %s: was %d, "
			    "now %d\n"), proclwp, pidstr, old, new);
	}
}

static struct ps_prochandle *
grab_proc(id_t pid)
{
	int ret;
	struct ps_prochandle *Pr;

	if ((Pr = Pgrab(pid, 0, &ret)) == NULL) {
		warn(gettext("cannot control process %d: %s\n"),
		    (int)pid, Pgrab_error(ret));
		errors = ERR_FAIL;
		return (NULL);
	}

	/*
	 * Set run-on-last-close flag so the controlled process
	 * runs even if we die on a signal, and create an agent LWP.
	 */
	if (Psetflags(Pr, PR_RLC) != 0 || Pcreate_agent(Pr) != 0) {
		warn(gettext("cannot control process %d\n"), (int)pid);
		errors = ERR_FAIL;
		Prelease(Pr, 0);
		return (NULL);
	}
	return (Pr);
}

static void
rele_proc(struct ps_prochandle *Pr)
{
	if (Pr == NULL)
		return;
	Pdestroy_agent(Pr);
	Prelease(Pr, 0);
}

static void
bind_lwp(struct ps_prochandle *Pr, id_t pid, id_t lwpid, processorid_t cpu)
{
	processorid_t old_cpu;

	if (pr_processor_bind(Pr, P_LWPID, lwpid, cpu, &old_cpu) < 0) {
		bind_err(cpu, pid, lwpid, errno);
		errors = ERR_FAIL;
	} else {
		if (qflag)
			query_out(pid, lwpid, old_cpu);
		else
			bind_out(pid, lwpid, old_cpu, cpu);
	}
}

/*
 * Query, set, or clear bindings for the range of LWPs in the given process.
 */
static int
do_lwps(id_t pid, const char *range, processorid_t cpu)
{
	char procfile[MAX_PROCFS_PATH];
	struct ps_prochandle *Pr;
	struct prheader header;
	processorid_t binding;
	struct lwpsinfo *lwp;
	char *lpsinfo, *ptr;
	int nent, size;
	int i, fd, found;

	/*
	 * Report bindings for LWPs in process 'pid'.
	 */
	(void) snprintf(procfile, MAX_PROCFS_PATH,
	    "/proc/%d/lpsinfo", (int)pid);
	if ((fd = open(procfile, O_RDONLY)) < 0) {
		if (errno == ENOENT)
			errno = ESRCH;
		bind_err(cpu, pid, -1, errno);
		return (ERR_FAIL);
	}
	if (pread(fd, &header, sizeof (header), 0) != sizeof (header)) {
		(void) close(fd);
		bind_err(cpu, pid, -1, errno);
		return (ERR_FAIL);
	}
	nent = header.pr_nent;
	size = header.pr_entsize * nent;
	ptr = lpsinfo = malloc(size);
	if (lpsinfo == NULL) {
		bind_err(cpu, pid, -1, errno);
		return (ERR_FAIL);
	}
	if (pread(fd, lpsinfo, size, sizeof (header)) != size) {
		bind_err(cpu, pid, -1, errno);
		free(lpsinfo);
		(void) close(fd);
		return (ERR_FAIL);
	}

	if ((bflag || uflag) && (Pr = grab_proc(pid)) == NULL) {
		free(lpsinfo);
		(void) close(fd);
		return (ERR_FAIL);
	}
	found = 0;
	for (i = 0; i < nent; i++, ptr += header.pr_entsize) {
		/*LINTED ALIGNMENT*/
		lwp = (lwpsinfo_t *)ptr;
		binding = lwp->pr_bindpro;
		if (!proc_lwp_in_set(range, lwp->pr_lwpid))
			continue;
		found++;
		if (bflag || uflag)
			bind_lwp(Pr, pid, lwp->pr_lwpid, cpu);
		else if (binding != PBIND_NONE)
			query_out(pid, lwp->pr_lwpid, binding);
	}
	if (bflag || uflag)
		rele_proc(Pr);
	free(lpsinfo);
	(void) close(fd);
	if (found == 0) {
		warn(gettext("cannot %s lwpid %d/%s: "
		    "No matching LWPs found\n"),
		    bflag ? "bind" : "query", pid, range);
		return (ERR_FAIL);
	}
	return (ERR_OK);
}

/*ARGSUSED*/
static int
query_all_proc(psinfo_t *psinfo, lwpsinfo_t *lwpsinfo, void *arg)
{
	id_t pid = psinfo->pr_pid;
	processorid_t binding;

	if (processor_bind(P_PID, pid, PBIND_QUERY, &binding) < 0) {
		/*
		 * Ignore search errors.  The process may have exited
		 * since we read the directory.
		 */
		if (errno == ESRCH)
			return (0);
		bind_err(PBIND_QUERY, pid, -1, errno);
		errors = ERR_FAIL;
		return (0);
	}
	if (binding != PBIND_NONE)
		query_out(pid, -1, binding);
	return (0);
}

static int
query_all_lwp(psinfo_t *psinfo, lwpsinfo_t *lwpsinfo, void *arg)
{
	id_t pid = psinfo->pr_pid;
	id_t lwpid = lwpsinfo->pr_lwpid;
	processorid_t *cpuid = arg;
	processorid_t binding = lwpsinfo->pr_bindpro;

	if (psinfo->pr_nlwp == 1)
		lwpid = -1;	/* report process bindings if only 1 lwp */
	if ((cpuid != NULL && *cpuid == binding) ||
	    (cpuid == NULL && binding != PBIND_NONE))
		query_out(pid, lwpid, binding);
	return (0);
}

/*
 * Execute the cmd with args while bound to cpu. Does not return:
 * either executes cmd successfully or dies trying.
 */
static void
exec_cmd(processorid_t cpu, char *cmd, char **args)
{
	if (processor_bind(P_PID, P_MYID, cpu, NULL) == -1) {
		bind_err(cpu, getpid(), -1, errno);
		exit(ERR_FAIL);
	}

	if (execvp(cmd, args) == -1)
		die(gettext("failed to exec %s\n"), cmd);
}

/*
 * Attempt to parse str as a CPU identifier. Return the identifier or
 * die.
 */
static processorid_t
parse_cpu(char *str)
{
	processorid_t cpu;
	char *endstr;

	cpu = strtol(str, &endstr, 10);
	if (endstr != NULL && *endstr != '\0' || cpu < 0)
		die(gettext("invalid processor ID %s\n"), optarg);

	return (cpu);
}

static int
usage(void)
{
	(void) fprintf(stderr,
	    gettext("usage: \n\t%1$s -b processor_id pid[/lwpids] ...\n"
	    "\t%1$s -e processor_id cmd [args...]\n"
	    "\t%1$s -U [processor_id] ...\n"
	    "\t%1$s -Q [processor_id] ...\n"
	    "\t%1$s -u pid[/lwpids] ...\n"
	    "\t%1$s [-q] [pid[/lwpids] ...]\n"),
	    progname);
	return (ERR_USAGE);
}

int
main(int argc, char *argv[])
{
	int c;
	int ret;
	id_t pid;
	processorid_t cpu, old_cpu;
	char *endstr;

	progname = argv[0];	/* put actual command name in messages */

	(void) setlocale(LC_ALL, "");	/* setup localization */
	(void) textdomain(TEXT_DOMAIN);

	while ((c = getopt(argc, argv, "b:e:qQuU")) != EOF) {
		switch (c) {

		case 'b':
			bflag = 1;
			cpu = parse_cpu(optarg);
			break;

		case 'e':
			eflag = 1;
			cpu = parse_cpu(optarg);
			break;

		case 'q':
			qflag = 1;
			cpu = PBIND_QUERY;
			break;

		case 'Q':
			Qflag = 1;
			cpu = PBIND_QUERY;
			break;

		case 'u':
			uflag = 1;
			cpu = PBIND_NONE;
			break;

		case 'U':
			Uflag = 1;
			break;

		default:
			return (usage());
		}
	}


	/*
	 * Make sure that at most one of the options b, e, q, Q, u, or
	 * U was specified.
	 */
	c = bflag + eflag + qflag + Qflag + uflag + Uflag;
	if (c < 1) {				/* nothing specified */
		qflag = 1;			/* default to query */
		cpu = PBIND_QUERY;
	} else if (c > 1) {
		warn(gettext("options -b, -e, -q, -Q, -u and -U "
		    "are mutually exclusive\n"));
		return (usage());
	}

	errors = 0;
	argc -= optind;
	argv += optind;

	/*
	 * Handle query of all processes.
	 */
	if (argc == 0) {
		if (bflag || uflag) {
			warn(gettext("must specify at least one pid\n"));
			return (usage());
		}
		if (eflag) {
			warn(gettext("must specify command\n"));
			return (usage());
		}
		if (Uflag) {
			if (processor_bind(P_ALL, 0, PBIND_NONE, &old_cpu) != 0)
				die(gettext("failed to unbind some LWPs"));
		}
		if (Qflag) {
			(void) proc_walk(query_all_lwp, NULL, PR_WALK_LWP);
			return (errors);
		} else {
			(void) proc_walk(query_all_proc, NULL, PR_WALK_PROC);
			return (errors);
		}
	}

	if (eflag)
		exec_cmd(cpu, argv[0], argv);

	if (Qflag || Uflag) {
		/*
		 * Go through listed processor IDs.
		 */
		for (; argc > 0; argv++, argc--) {
			errno = 0;
			cpu = (id_t)strtol(*argv, &endstr, 10);
			if (errno != 0 || (endstr != NULL && *endstr != '\0') ||
			    p_online(cpu, P_STATUS) == -1) {
				warn(gettext("invalid processor ID\n"));
				continue;
			}
			if (Qflag) {
				(void) proc_walk(query_all_lwp,
				    &cpu, PR_WALK_LWP);
				continue;
			}
			if (Uflag) {
				if (processor_bind(P_CPUID, cpu,
				    PBIND_NONE, &old_cpu) != 0) {
					warn(gettext("failed to unbind from "
					    "processor %d"), (int)cpu);
					errors = ERR_FAIL;
				}
				continue;
			}
		}
		return (errors);
	}

	/*
	 * Go through listed process[/lwp_ranges].
	 */
	for (; argc > 0; argv++, argc--) {
		errno = 0;
		pid = (id_t)strtol(*argv, &endstr, 10);
		if (errno != 0 ||
		    (endstr != NULL && *endstr != '\0' && *endstr != '/')) {
			warn(gettext("invalid process ID: %s\n"), *argv);
			continue;
		}
		if (endstr != NULL && *endstr == '/') {
			/*
			 * Handle lwp range case
			 */
			const char *lwps = (const char *)(++endstr);
			if (*lwps == '\0' ||
			    proc_lwp_range_valid(lwps) != 0) {
				warn(gettext("invalid lwp range "
				    "for pid %d\n"), (int)pid);
				errors = ERR_FAIL;
				continue;
			}
			if (!qflag)
				(void) proc_initstdio();
			ret = do_lwps(pid, lwps, qflag ? PBIND_QUERY : cpu);
			if (!qflag)
				(void) proc_finistdio();
			if (ret != ERR_OK)
				errors = ret;
		} else {
			/*
			 * Handle whole process case.
			 */
			if (processor_bind(P_PID, pid, cpu, &old_cpu) < 0) {
				bind_err(cpu, pid, -1, errno);
				errors = ERR_FAIL;
				continue;
			}
			if (qflag)
				query_out(pid, -1, old_cpu);
			else
				bind_out(pid, -1, old_cpu, cpu);
		}
	}
	return (errors);
}