From a4422e188e2043d9cc643804be559b735881cee6 Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Wed, 18 Dec 2019 09:25:42 +0100
Subject: [PATCH] Import source code.

---
 build.sh |   3 ++
 cpuset.c | 147 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 150 insertions(+)
 create mode 100755 build.sh
 create mode 100644 cpuset.c

diff --git a/build.sh b/build.sh
new file mode 100755
index 0000000..29c3ef0
--- /dev/null
+++ b/build.sh
@@ -0,0 +1,3 @@
+#!/bin/sh
+
+gcc -Wall cpuset.c -o cpuset
diff --git a/cpuset.c b/cpuset.c
new file mode 100644
index 0000000..906ed4a
--- /dev/null
+++ b/cpuset.c
@@ -0,0 +1,147 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <sys/stat.h>
+
+#define CPUSET(x) "/sys/fs/cgroup/cpuset/"x
+
+#define ACTION_ENABLE 1
+#define ACTION_DISABLE 2
+
+char *readLine(FILE *fp, char *buffer) {
+	int ch;
+	int i = 0;
+	size_t buff_len = 0;
+	buffer = malloc(buff_len + 1);
+	if(!buffer) {
+		return NULL;
+	}
+	while((ch = fgetc(fp)) != '\n' && ch != EOF) {
+	buff_len++;
+	void *tmp = realloc(buffer, buff_len + 1);
+	if(tmp == NULL) {
+		free(buffer);
+		return NULL;
+	}
+	buffer = tmp;
+	buffer[i] = (char) ch;
+	i++;
+	}
+	buffer[i] = '\0';
+	if(ch == EOF && (i == 0 || ferror(fp))) {
+		free(buffer);
+		return NULL;
+	}
+	return buffer;
+}
+
+int moveTasks(char *src, char *dst) {
+	FILE *fd1, *fd2;
+	char *pid;
+	fd1 = fopen(src, "rb");
+	if(fd1 == NULL) {
+		return 255;
+	}
+	while((pid = readLine(fd1, 0)) != NULL) {
+		if(strtol(pid, (char **)NULL, 10) != getppid()) {
+			fd2 = fopen(dst, "wb");
+			if(fd2 == NULL) {
+				return 255;
+			}
+			fprintf(fd2, "%s\n", pid);
+			fclose(fd2);
+			free(pid);
+		}
+	}
+	fclose(fd1);
+	return 0;
+}
+
+int setGroup(char *dest, char *data) {
+	FILE *fd;
+	fd = fopen(dest, "wb");
+	if(fd == NULL) {
+		return 255;
+	}
+	fprintf(fd, "%s", data);
+	fclose(fd);
+	return 0;
+}
+
+int disableGroup() {
+	struct stat sb = {0};
+	printf("Disabling CGROUPS ... ");
+	if(stat(CPUSET("host"), &sb) == 0) {
+		if(moveTasks(CPUSET("host/tasks"), CPUSET("tasks")) != 0) {
+			return 255;
+		}
+		rmdir(CPUSET("host"));
+		printf("OK!\n");
+	} else {
+		printf("SKIPPED!\n");
+	}
+	return 0;
+}
+
+int enableGroup(char *cpu, char *mem) {
+	struct stat sb = {0};
+	printf("Enabling CGROUPS ... ");
+	if(stat(CPUSET("host"), &sb) == -1) {
+		mkdir(CPUSET("host"), 0755);
+	}
+	if(setGroup(CPUSET("host/cpuset.cpus"), cpu) != 0) {
+		printf("FAIL!\n");
+		return 255;
+	}
+	if(setGroup(CPUSET("host/cpuset.mems"), mem) != 0) {
+		printf("FAIL!\n");
+		return 255;
+	}
+	if(moveTasks(CPUSET("tasks"), CPUSET("host/tasks")) != 0) {
+		printf("FAIL!\n");
+		return 255;
+	}
+	printf("OK!\n");
+	return 0;
+}
+
+int main(int argc, char *argv[]) {
+	int opt;
+	int action = 0;
+	char *cpu = {0};
+	char *mem = {0};
+	while((opt = getopt(argc, argv, ":edc:m:")) != -1) {
+		switch(opt) {
+			case 'e':
+				action = ACTION_ENABLE;
+				break;
+			case 'd':
+				action = ACTION_DISABLE;
+				break;
+			case 'c':
+				cpu = optarg;
+				break;
+			case 'm':
+				mem = optarg;
+				break;
+			case ':':
+				printf("Option needs a value!\n");
+				return 1;
+			default:
+				printf("Invalid option: -%c!\n", optopt);
+				return 1;
+		}
+	}
+	if(action == ACTION_ENABLE) {
+		if(!cpu || !mem) {
+			printf("No CPU or memory pinning specified!\n");
+			return 1;
+		}
+		return enableGroup(cpu, mem);
+	} else if(action == ACTION_DISABLE) {
+		return disableGroup();
+	} else {
+		printf("No action specified!\n");
+	}
+	return 0;
+}