linux_programming_interface/rtsched_cap.c at master · anthonyc-r/linux_programming_interface · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
/**
* Linux-specific
* Run a program with a specific realtime scheduler and priority
* This version still assumes the file is set-uid-root, and doesn't assume use of
* file capabilities, however it does make use of capabilities via libcap rather than
* changing the UID.
* $./rtsched policy(r|f) priority command args...
*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <sched.h>
#include <string.h>
#include <sys/prctl.h>
#include <sys/capability.h>
#include <linux/securebits.h>
#include <linux/capability.h>


static void exit_usage() {
	setuid(getuid());
	puts("Usage: ./rtsched policy priority command arg");
	exit(EXIT_FAILURE);
}

static void exit_error(char *tag) {
	setuid(getuid());
	printf("rtsched (%s): %s\n", tag, strerror(errno));
	exit(errno);
}


// Want to drop root permenantly, but retain required caps.
static void drop_root() {
	cap_t caps;
	cap_value_t cap;

	if (prctl(PR_SET_SECUREBITS, SECBIT_NO_SETUID_FIXUP | SECBIT_NOROOT |
		SECBIT_NO_SETUID_FIXUP_LOCKED | SECBIT_NOROOT_LOCKED) == -1)
		exit_error("prctl");
	if ((caps = cap_init()) == NULL)
		exit_error("cap_init");

	cap = CAP_SYS_NICE;
	if (cap_set_flag(caps, CAP_PERMITTED, 1, &cap, 1) == -1)
		exit_error("cap_set_flag");
	if (cap_set_proc(caps) == -1)
		exit_error("cap_set_proc");
	if (cap_free(caps) == -1)
		exit_error("cap_free");
	if (setuid(getuid()) == -1)
		exit_error("setuid");
}

static void drop_caps(int forever) {
	cap_flag_t flags;
	cap_t caps;
	cap_value_t cap;

	if ((caps = cap_init()) == NULL)
		exit_error("cap_init");
	cap = CAP_SYS_NICE;
	if (!forever) {
		if (cap_set_flag(caps, CAP_PERMITTED, 1, &cap, 1) == -1)
			exit_error("cap_set_flag");
	}
	if (cap_set_proc(caps) == -1)
		exit_error("cap_set_proc");
	if (cap_free(caps) == -1)
		exit_error("cap_free");
}

static void raise_caps() {
	cap_t caps;
	cap_value_t cap;

	if ((caps = cap_init()) == NULL)
		exit_error("cap_init");
	cap = CAP_SYS_NICE;
	if (cap_set_flag(caps, CAP_PERMITTED, 1, &cap, 1) == -1)
		exit_error("cap_set_flag");
	if (cap_set_flag(caps, CAP_EFFECTIVE, 1, &cap, 1) == -1)
		exit_error("cap_set_flag");
	if (cap_set_proc(caps) == -1)
		exit_error("cap_set_proc");
	if (cap_free(caps) == -1)
		exit_error("cap_free");
}

int main(int argc, char **argv) {
	int policy;
	struct sched_param param;
	char *ep;

	drop_root();
	if (argc < 4)
		exit_usage();
	param.sched_priority = strtol(argv[2], &ep, 10);
	if (*ep != '\0')
		exit_usage();
	if (*argv[1] == 'f')
		policy = SCHED_FIFO;
	else if (*argv[1] == 'r')
		policy = SCHED_RR;
	else
		exit_usage();

	raise_caps();
	if (sched_setscheduler(0, policy, &param) == -1)
		exit_error("sched_setscheduler");
	drop_caps(1);
	if (setuid(getuid()))
		exit_error("setuid");
	execvp(argv[3], argv + 3);
	exit_error("execve");
}