#include "common.h"
#include "target.h"

#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <poll.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

enum {
	SHELL_UID = 2000,
};

static void set_root_env(void)
{
	setenv("PATH",
	       "/product/bin:/apex/com.android.runtime/bin:"
	       "/apex/com.android.art/bin:/system_ext/bin:/system/bin:"
	       "/system/xbin:/odm/bin:/vendor/bin:/vendor/xbin",
	       1);
	setenv("HOME", "/data/local/tmp", 1);
	setenv("USER", "root", 1);
	setenv("LOGNAME", "root", 1);
}

static int write_all(int fd, const void *data, size_t size)
{
	const unsigned char *p = data;

	while (size > 0) {
		ssize_t n = write(fd, p, size);

		if (n < 0 && errno == EINTR)
			continue;
		if (n <= 0)
			return -1;
		p += n;
		size -= (size_t)n;
	}
	return 0;
}

static int read_all(int fd, void *data, size_t size)
{
	unsigned char *p = data;

	while (size > 0) {
		ssize_t n = read(fd, p, size);

		if (n < 0 && errno == EINTR)
			continue;
		if (n <= 0)
			return 0;
		p += n;
		size -= (size_t)n;
	}
	return 1;
}

static int copy_fd(int source, int destination)
{
	unsigned char buffer[16384];

	for (;;) {
		ssize_t n = read(source, buffer, sizeof(buffer));

		if (n < 0 && errno == EINTR)
			continue;
		if (n < 0)
			return -1;
		if (n == 0)
			return 0;
		if (write_all(destination, buffer, (size_t)n) != 0)
			return -1;
	}
}

static int files_equal(const char *first, const char *second)
{
	unsigned char first_buffer[8192];
	unsigned char second_buffer[8192];
	int first_fd = open(first, O_RDONLY | O_CLOEXEC);
	int second_fd = open(second, O_RDONLY | O_CLOEXEC);
	int equal = 0;

	if (first_fd < 0 || second_fd < 0)
		goto out;
	for (;;) {
		ssize_t first_size;
		ssize_t second_size;

		do {
			first_size = read(first_fd, first_buffer,
					  sizeof(first_buffer));
		} while (first_size < 0 && errno == EINTR);
		do {
			second_size = read(second_fd, second_buffer,
					   sizeof(second_buffer));
		} while (second_size < 0 && errno == EINTR);
		if (first_size < 0 || second_size < 0 ||
		    first_size != second_size)
			goto out;
		if (first_size == 0) {
			equal = 1;
			goto out;
		}
		if (memcmp(first_buffer, second_buffer,
			   (size_t)first_size) != 0)
			goto out;
	}
out:
	if (first_fd >= 0)
		close(first_fd);
	if (second_fd >= 0)
		close(second_fd);
	return equal;
}

int su_install_self(const char *path)
{
	char temporary[PATH_MAX];
	int source = -1;
	int destination = -1;
	int saved_errno = EIO;
	int installed = 0;

	if (!path || snprintf(temporary, sizeof(temporary), "%s.new.%d",
			      path, getpid()) >= (int)sizeof(temporary)) {
		errno = ENAMETOOLONG;
		return -1;
	}
	unlink(temporary);
	source = open("/proc/self/exe", O_RDONLY | O_CLOEXEC);
	destination = open(temporary,
			   O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0755);
	if (source < 0 || destination < 0)
		goto out;
	if (copy_fd(source, destination) != 0 ||
	    fchmod(destination, 0755) != 0 || fsync(destination) != 0)
		goto out;
	if (close(source) != 0) {
		source = -1;
		goto out;
	}
	source = -1;
	if (close(destination) != 0) {
		destination = -1;
		goto out;
	}
	destination = -1;
	if (rename(temporary, path) != 0)
		goto out;
	if (!files_equal("/proc/self/exe", path)) {
		errno = EIO;
		unlink(path);
		goto out;
	}
	printf("SU_INSTALL path=%s source=/proc/self/exe verified=1\n", path);
	installed = 1;
out:
	saved_errno = errno;
	if (source >= 0)
		close(source);
	if (destination >= 0)
		close(destination);
	unlink(temporary);
	if (installed)
		return 0;
	errno = saved_errno;
	return -1;
}

static int connect_daemon(void)
{
	struct sockaddr_un address = {
		.sun_family = AF_UNIX,
	};
	struct ucred credential;
	char security_context[128];
	socklen_t credential_size = sizeof(credential);
	socklen_t context_size = sizeof(security_context) - 1;
	int fd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);

	if (fd < 0)
		return -1;
	snprintf(address.sun_path, sizeof(address.sun_path), "%s",
		 TARGET_SU_SOCKET);
	if (connect(fd, (struct sockaddr *)&address, sizeof(address)) != 0) {
		close(fd);
		return -1;
	}
	if (getsockopt(fd, SOL_SOCKET, SO_PEERCRED, &credential,
		       &credential_size) != 0 ||
	    credential.uid != 0 || credential.gid != 0 ||
	    getsockopt(fd, SOL_SOCKET, SO_PEERSEC, security_context,
		       &context_size) != 0) {
		close(fd);
		errno = EPERM;
		return -1;
	}
	security_context[context_size < sizeof(security_context) ?
			 context_size : sizeof(security_context) - 1] = 0;
	if (strcmp(security_context, TARGET_CARRIER_CONTEXT) != 0) {
		close(fd);
		errno = EPERM;
		return -1;
	}
	return fd;
}

static int pump_pair(int first, int second)
{
	char buffer[4096];
	int first_open = 1;
	int second_open = 1;

	while (first_open || second_open) {
		struct pollfd poll_fds[2];
		int count = 0;

		if (first_open) {
			poll_fds[count].fd = first;
			poll_fds[count].events = POLLIN;
			count++;
		}
		if (second_open) {
			poll_fds[count].fd = second;
			poll_fds[count].events = POLLIN;
			count++;
		}
		if (poll(poll_fds, (nfds_t)count, -1) < 0) {
			if (errno == EINTR)
				continue;
			return 1;
		}
		count = 0;
		if (first_open) {
			short events = poll_fds[count++].revents;

			if (events & POLLIN) {
				ssize_t n = read(first, buffer, sizeof(buffer));

				if (n > 0)
					(void)write_all(second, buffer, (size_t)n);
				else
					first_open = 0;
			} else if (events & (POLLHUP | POLLERR | POLLNVAL)) {
				first_open = 0;
			}
			if (!first_open)
				shutdown(second, SHUT_WR);
		}
		if (second_open) {
			short events = poll_fds[count].revents;

			if (events & POLLIN) {
				ssize_t n = read(second, buffer, sizeof(buffer));

				if (n > 0)
					(void)write_all(first, buffer, (size_t)n);
				else
					second_open = 0;
			} else if (events & (POLLHUP | POLLERR | POLLNVAL)) {
				second_open = 0;
			}
			if (!second_open)
				shutdown(first, SHUT_WR);
		}
	}
	return 0;
}

static char *read_socket_output(int fd)
{
	size_t capacity = 4096;
	size_t size = 0;
	char *output = malloc(capacity);

	if (!output)
		return NULL;
	for (;;) {
		ssize_t n;

		if (capacity - size < 2048) {
			char *grown;

			capacity *= 2;
			grown = realloc(output, capacity);
			if (!grown) {
				free(output);
				return NULL;
			}
			output = grown;
		}
		n = read(fd, output + size, capacity - size - 1);
		if (n < 0 && errno == EINTR)
			continue;
		if (n <= 0)
			break;
		size += (size_t)n;
	}
	output[size] = 0;
	return output;
}

char *su_run_command(const char *command, int *status_out)
{
	char mode = 'C';
	uint32_t length = (uint32_t)strlen(command);
	int fd = connect_daemon();
	char *output;

	if (status_out)
		*status_out = -1;
	if (fd < 0)
		return NULL;
	if (write_all(fd, &mode, sizeof(mode)) != 0 ||
	    write_all(fd, &length, sizeof(length)) != 0 ||
	    write_all(fd, command, length) != 0) {
		close(fd);
		return NULL;
	}
	shutdown(fd, SHUT_WR);
	output = read_socket_output(fd);
	close(fd);
	if (output && status_out)
		*status_out = 0;
	return output;
}

int su_client_main(int argc, char **argv)
{
	int status;

	if (argc == 3 && strcmp(argv[1], "-c") == 0) {
		char *output = su_run_command(argv[2], &status);

		if (!output)
			return 127;
		(void)write_all(STDOUT_FILENO, output, strlen(output));
		free(output);
		return status;
	}
	if (argc != 1) {
		fprintf(stderr, "usage: %s [-c command]\n", argv[0]);
		return 2;
	}

	{
		char mode = 'I';
		int fd = connect_daemon();

		if (fd < 0)
			return 127;
		if (write_all(fd, &mode, sizeof(mode)) != 0) {
			close(fd);
			return 127;
		}
		(void)pump_pair(STDIN_FILENO, fd);
		close(fd);
	}
	return 0;
}

static void run_command(int connection, const char *command)
{
	pid_t pid = fork();

	if (pid == 0) {
		dup2(connection, STDIN_FILENO);
		dup2(connection, STDOUT_FILENO);
		dup2(connection, STDERR_FILENO);
		close(connection);
		set_root_env();
		execl("/system/bin/sh", "sh", "-c", command, (char *)NULL);
		_exit(127);
	}
	if (pid > 0) {
		while (waitpid(pid, NULL, 0) < 0 && errno == EINTR)
			;
	}
}

static int open_pty(char *slave_name, size_t slave_name_size)
{
	int master = posix_openpt(O_RDWR | O_NOCTTY | O_CLOEXEC);

	if (master < 0)
		return -1;
	if (grantpt(master) != 0 || unlockpt(master) != 0 ||
	    ptsname_r(master, slave_name, slave_name_size) != 0) {
		close(master);
		return -1;
	}
	return master;
}

static void run_interactive(int connection)
{
	char slave_name[128];
	int master = open_pty(slave_name, sizeof(slave_name));
	pid_t pid;

	if (master < 0)
		return;
	pid = fork();
	if (pid == 0) {
		int slave;

		setsid();
		slave = open(slave_name, O_RDWR | O_NOCTTY);
		if (slave < 0)
			_exit(126);
		ioctl(slave, TIOCSCTTY, 0);
		dup2(slave, STDIN_FILENO);
		dup2(slave, STDOUT_FILENO);
		dup2(slave, STDERR_FILENO);
		if (slave > STDERR_FILENO)
			close(slave);
		close(master);
		close(connection);
		set_root_env();
		execl("/system/bin/sh", "sh", "-i", (char *)NULL);
		_exit(127);
	}
	if (pid > 0) {
		(void)pump_pair(connection, master);
		kill(pid, SIGHUP);
		while (waitpid(pid, NULL, 0) < 0 && errno == EINTR)
			;
	}
	close(master);
}

static int peer_allowed(int connection)
{
	struct ucred credential;
	socklen_t size = sizeof(credential);

	if (getsockopt(connection, SOL_SOCKET, SO_PEERCRED,
		       &credential, &size) != 0)
		return 0;
	return credential.uid == 0 || credential.uid == SHELL_UID;
}

static void serve_connection(int connection)
{
	char mode;

	if (!peer_allowed(connection) ||
	    !read_all(connection, &mode, sizeof(mode)))
		return;
	if (mode == 'C') {
		uint32_t length;
		char *command;

		if (!read_all(connection, &length, sizeof(length)) ||
		    length > 65536)
			return;
		command = calloc(1, (size_t)length + 1);
		if (!command)
			return;
		if (read_all(connection, command, length))
			run_command(connection, command);
		free(command);
	} else if (mode == 'I') {
		run_interactive(connection);
	}
}

int su_daemon_main(const char *self_path)
{
	struct sockaddr_un address = {
		.sun_family = AF_UNIX,
	};
	int fd;

	if (getuid() != 0 || geteuid() != 0)
		return 1;
	signal(SIGPIPE, SIG_IGN);
	signal(SIGCHLD, SIG_IGN);
	set_root_env();
	if (self_path && self_path[0] == '/') {
		if (chown(self_path, 0, 0) != 0 || chmod(self_path, 0755) != 0)
			return 1;
	}
	fd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0);
	if (fd < 0)
		return 1;
	unlink(TARGET_SU_SOCKET);
	snprintf(address.sun_path, sizeof(address.sun_path), "%s",
		 TARGET_SU_SOCKET);
	if (bind(fd, (struct sockaddr *)&address, sizeof(address)) != 0 ||
	    chown(TARGET_SU_SOCKET, 0, SHELL_UID) != 0 ||
	    chmod(TARGET_SU_SOCKET, 0660) != 0 || listen(fd, 16) != 0)
		return 1;
	fprintf(stderr, "cathash su ready pid=%d socket=%s uid=%d euid=%d\n",
		getpid(), TARGET_SU_SOCKET, getuid(), geteuid());

	for (;;) {
		int connection = accept4(fd, NULL, NULL, SOCK_CLOEXEC);
		pid_t pid;

		if (connection < 0) {
			if (errno == EINTR)
				continue;
			sleep(1);
			continue;
		}
		pid = fork();
		if (pid == 0) {
			signal(SIGCHLD, SIG_DFL);
			close(fd);
			serve_connection(connection);
			close(connection);
			_exit(0);
		}
		close(connection);
	}
}