jail.c

/*
 * AUTHOR: Jiahua Guo <gjhdgm@gmail.com>
 *
 * run this to compile: gcc -o jail jail.c -lcap
 */
#define _GNU_SOURCE
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/mount.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/prctl.h>
#include <sys/capability.h>
#include <linux/capability.h>
#include <linux/securebits.h>
#include <fcntl.h>
#include <dirent.h>
#include <sched.h>
#include <grp.h>
#include <time.h>
#include <signal.h>
#include <pwd.h>
#include <grp.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <errno.h>
#include <limits.h>

enum {
	CP_WAIT     = 1,
	CP_ROOT     = 2,
	CP_NCHROOT  = 4,
	CP_NCDHOME  = 8,
};

#define STACK_SIZE (4*1024*1024)
#define PATH_BUF_SIZE 8192


int checkerr_noexit = 0;

static inline void _check_perror(int err, char *msg, char *file, int lineno)
{
	if(!err) return;
	fprintf(stderr, "\nFILE: %s\nLINE: %d\n", file, lineno);
	perror(msg);
	if( !checkerr_noexit )
		exit(EXIT_FAILURE);
}
static inline void _check_error(int err, char *msg, char *file, int lineno)
{
	if(!err) return;
	fprintf(stderr, "\nFILE: %s\nLINE: %d\n%s\n", file, lineno, msg);
	if( !checkerr_noexit )
		exit(EXIT_FAILURE);
}
#define check_perror(err, msg) _check_perror( (err), (msg), __FILE__, __LINE__ )
#define check_error(err, msg)  _check_error( (err), (msg), __FILE__, __LINE__ )
#define massert(cond)  check_error( !(cond), "ASSERT FAILED: " #cond )

#define DEBUG(...) (fprintf( stderr, "!%s~%d: ", __FILE__, __LINE__ ), fprintf( stderr, __VA_ARGS__ ))


static int close_nonstdin_opened_files(void)
{
	int dirfd;
	DIR *dir;
	struct dirent *dirent;
	char *name, *end;
	int fd;
	long int ret;
	int count = 0;

	dirfd = open("/proc/self/fd", O_RDONLY);
	check_perror( dirfd == -1, "open" );

	dir = fdopendir( dirfd );
	check_perror(dir == NULL, "opendir");

	while(1) {
		errno = 0;
		dirent = readdir(dir);
		if(dirent == NULL) {
			check_perror( errno, "readdir" );
			break;
		}
		name = dirent->d_name;
		if( name[0] == '.' )
			continue;

		errno = 0;
		ret  = strtol( name, &end, 10 );
		check_perror( errno, "strtol" );
		check_error( *end != 0, "bad fd in /proc/self/fd/" );
		check_error( ret < 0 || ret > INT_MAX, "bad fd, out of range" );

		fd = ret;

		if(fd == dirfd || fd == 0 || fd == 1 || fd == 2)
			continue;
		ret = close(fd);
		check_perror(ret, "close");
	}
	closedir(dir);
	return count;
}

static uid_t user_uid, target_uid;
static gid_t user_gid, target_gid;
static char *chroot_dir;
static char *mount_file;
static char *init_file;
static char *shell = "/bin/bash";
static char **program_argv;

#define HOME      "/tmp/home"
#define INIT_FILE "/tmp/init_file"


static const char *chrooted_path(const char *path)
{
	static char spath[PATH_BUF_SIZE];
	const int size = sizeof(spath);
	int len1, len2;
	massert( chroot_dir[0] == '/' );
	massert( path[0] == '/' );

	len1 = strnlen( chroot_dir, size );
	len2 = strnlen( path, size );

	massert( len1 + len2 < size );

	strcpy( spath, chroot_dir );
	strcat( spath, path );

	return spath;
}

static void do_chroot(void)
{
	check_perror( chroot( chroot_dir ), "chroot" );
	check_perror( chdir( "/" ), "chdir" );
}

static void drop_privileges(uid_t target_uid, gid_t target_gid)
{
	gid_t gids[1] = {-1};
	cap_t cap;
	int i;
	int ret;

	check_perror( setgroups( 0, gids ), "setgroups" );
	check_perror( setregid( target_gid, target_gid ), "setregid" );

	check_perror(
		prctl(PR_SET_SECUREBITS,
		      SECBIT_NO_SETUID_FIXUP |
		      SECBIT_NO_SETUID_FIXUP_LOCKED |
		      SECBIT_NOROOT |
		      SECBIT_NOROOT_LOCKED),
		"prctl");

	for( i = 0; i <= CAP_LAST_CAP; ++i ) {
		if( cap_valid( i ) ) {
			ret = prctl( PR_CAPBSET_DROP, i );
			check_perror( ret == -1 && errno != EINVAL, "prctl pr_capbset_drop" );
		}
	}

	check_perror( setreuid( target_uid, target_uid ), "setreuid" );

	cap = cap_init();
	check_perror( cap == NULL, "cap_init" );
	check_perror( cap_set_proc( cap ), "cap_set_proc" );
	check_perror( cap_free( cap ), "cap_free" );
}

static void safe_exec(char **argv, int flags, uid_t uid, gid_t gid)
{
	const char *home;
	if( !(flags & CP_NCHROOT) ) {
		do_chroot();
		home = HOME;
	} else {
		home = chrooted_path( HOME );
	}
	setenv( "HOME", home, 1 );
	if( !(flags & CP_NCDHOME) )
		check_perror( chdir( home ), "chdir" );
	if( !(flags & CP_ROOT) )
		drop_privileges( uid, gid );
	umask(0022);
	check_perror( execv(argv[0], argv), "exec" );
}

static int create_process(char **argv, int flags, uid_t uid, gid_t gid)
{
	pid_t pid;
	pid = fork();
	check_perror( pid == -1, "fork" );

	if( pid == 0 ) {
		safe_exec( argv, flags, uid, gid );
	} else {
		if( flags & CP_WAIT ) {
			int status;
			int ret;
			ret = waitpid( pid, &status, 0 );
			check_perror( ret == -1, "waitpid" );
			return WEXITSTATUS(status);
		}
	}
	return pid;
}

static void initial_devfs(void)
{
	check_perror( mkdir( "/dev/shm", 01755 ), "mkdir /dev/shm" );
	check_perror( mount( NULL, "/dev/shm", "tmpfs",  0, NULL), "mount /dev/shm" );

	check_perror( mknod( "/dev/null",    S_IFCHR | 0666, makedev( 1, 3 ) ), "mknod" );
	check_perror( mknod( "/dev/zero",    S_IFCHR | 0666, makedev( 1, 5 ) ), "mknod" );
	check_perror( mknod( "/dev/random",  S_IFCHR | 0666, makedev( 1, 8 ) ), "mknod" );
	check_perror( mknod( "/dev/urandom", S_IFCHR | 0666, makedev( 1, 9 ) ), "mknod" );

	check_perror( mknod( "/dev/tty",     S_IFCHR | 0666, makedev( 5, 0 ) ), "mknod" );

	check_perror( symlink( "/proc/self/fd", "/dev/fd" ), "symlink" );
	check_perror( symlink( "fd/0", "/dev/stdin"  ), "symlink" );
	check_perror( symlink( "fd/1", "/dev/stdout" ), "symlink" );
	check_perror( symlink( "fd/2", "/dev/stderr" ), "symlink" );
}

static void prepare_filesystem(void)
{
	if( mount_file ) {
		char *mount_file_argv[] = {
			"/bin/mount",
			"-r",
			"--",
			mount_file,
			chroot_dir,
			NULL
		};
		int ret = create_process( mount_file_argv, CP_WAIT | CP_ROOT | CP_NCHROOT | CP_NCDHOME, 0, 0 );
		check_error( ret, "exec /bin/mount failed" );
		check_perror( chdir( "/" ), "chdir" );
	}

	umask( 0 );

	do_chroot();

	check_perror( mount( NULL, "/proc", "proc",  0, NULL), "mount /proc" );
	check_perror( mount( NULL, "/sys",  "sysfs", 0, NULL), "mount /sys" );
	check_perror( mount( NULL, "/dev",  "tmpfs", 0, NULL), "mount /dev" );
	check_perror( mount( NULL, "/tmp",  "tmpfs", 0, NULL), "mount /tmp" );

	initial_devfs();

	check_perror( mkdir( HOME, 0755 ), "mkdir " HOME );
	check_perror( chown( HOME, target_uid, target_gid ), "chown " HOME );
}

static int fork_run(int (*fn)(void *arg), void *arg)
{
	pid_t pid;
	int status;

	pid = fork();
	if( pid == 0 )
		exit( fn(arg) );
	
	check_perror( pid == -1, "fork" );
	check_perror( waitpid( pid, &status, 0 ) == -1, "waitpid" );
	check_error( !WIFEXITED(status), "process not exited" );
	return WEXITSTATUS( status );
}

static void copy_file(const char *from, const char *to) 
{
	int fdr, fdw;
	char buf[4096];
	ssize_t rs, ws;

	fdr = open( from, O_RDONLY );
	check_perror( fdr < 0, "open" );

	fdw = open( to, O_WRONLY | O_CREAT | O_EXCL, 0700 );
	check_perror( fdw < 0, "open" );

	while(1) {
		rs = read( fdr, buf, sizeof(buf) );
		check_perror( rs < 0, "read" );

		if(!rs) break;

		ws = write( fdw, buf, rs );
		check_perror( ws < 0, "write" );
		massert( ws == rs );
	}

	close( fdr );
	close( fdw );
}

static int init(void *arg)
{
	char **argv, *basic_argv[2];

	check_error( fork_run( (int(*)(void*))prepare_filesystem, NULL ),
		     "prepare_filesystem failed"
		);

	setenv( "SHELL", shell, 1 );

	argv = program_argv;
	if( *argv == NULL ) {
		basic_argv[0] = shell;
		basic_argv[1] = NULL;
		argv = basic_argv;
	}

	if( init_file ) {
		char *init_file_abspath = strdupa( chrooted_path( INIT_FILE ) );
		char *iargv[] = {
			shell,
			"--",
			init_file_abspath,
			NULL
		};
		copy_file( init_file, init_file_abspath );
		create_process( iargv, CP_WAIT | CP_NCHROOT, 0, 0 );
	}

	create_process( argv, CP_WAIT, target_uid, target_gid );

	exit( EXIT_SUCCESS );
}

static void usage(const char *program)
{
	printf( "Usage:\n\n\
\t%s [-r chroot_dir] [-m mount_file] [-i init_file] [-u user] [-g group] [-U uid] [-G gid] [program [args...]]\n\
\n\
Example:\n\
\n\
\t%s -r /tmp -m /path/to/disk.squashfs -i /path/to/initial.sh -U 1000 -G 1000 /bin/bash\n\
", program, program );
	
}


int main(int argc, char **argv)
{
	int ret;
	int opt;
	struct passwd *passwd;
	struct group  *group;

	char *stack;

	pid_t pid;
	int status;

	user_uid = getuid();
	user_gid = getgid();

	target_uid = user_uid;
	target_gid = user_gid;
	chroot_dir = "/tmp";

	while( (opt = getopt(argc, argv, "r:m:i:u:g:U:G:")) != -1 ) {
		switch( opt ) {
		case 'r':
			chroot_dir = optarg;
			massert( chroot_dir[0] == '/' );
			break;
		case 'm':
			mount_file = optarg;
			massert( mount_file[0] == '/' );
			break;
		case 'i':
			init_file = optarg;
			massert( init_file[0] == '/' );
			break;
		case 'u':
			passwd = getpwnam( optarg );
			if( passwd ) {
				target_uid = passwd->pw_uid;
				target_gid = passwd->pw_gid;
			} else {
				fprintf( stderr, "can't find user %s", optarg );
				exit( EXIT_FAILURE );
			}
			break;
		case 'g':
			group = getgrnam( optarg );
			if( group ) {
				target_gid = group->gr_gid;
			} else {
				fprintf( stderr, "can't find group %s", optarg );
				exit( EXIT_FAILURE );
			}
			break;
		case 'U':
			target_uid = atoi(optarg);
			break;
		case 'G':
			target_gid = atoi(optarg);
			break;
		case '?':
		default:
			usage(argv[0]);
			exit( EXIT_FAILURE );
		}
	}
	program_argv = argv + optind;

	massert( user_uid == 0 || user_uid == target_uid );
	massert( user_gid == 0 || user_gid == target_gid );


	close_nonstdin_opened_files();

	stack = malloc( STACK_SIZE );
	check_perror( !stack, "malloc" );

	check_perror( setreuid( 0, 0 ), "setreuid" );

	pid = clone(
		init,
		stack + STACK_SIZE - 4,
		CLONE_NEWIPC | CLONE_NEWNET | CLONE_NEWNS | CLONE_NEWUTS | CLONE_NEWPID,
		NULL
		);
	check_perror( pid == -1, "clone" );

	ret = waitpid( pid, &status, __WALL );
	check_perror( ret != pid, "waitpid" );

	return 0;
}