Merge pull request #11 from opsengine/develop

OS X support

8 files changed, 198 insertions, 71 deletions

diff --git a/README b/README
index 2ad5570..7b9ddd3 100644
--- a/README
+++ b/README
@@ -25,7 +25,7 @@ https://github.com/opsengine/cpulimit
  Install instructions
 ======================
 
-On Linux:
+On Linux/OS X:
 
 $ make
 # cp src/cpulimit /usr/bin
diff --git a/src/cpulimit.c b/src/cpulimit.c
index e7d5309..002205d 100644
--- a/src/cpulimit.c
+++ b/src/cpulimit.c
@@ -46,6 +46,10 @@
 #include "process_group.h"
 #include "list.h"
 
+#ifdef __APPLE__
+#include "memrchr.c"
+#endif
+
 //some useful macro
 #ifndef MIN
 #define MIN(a,b) (((a)<(b))?(a):(b))
@@ -147,6 +151,26 @@ static int get_ncpu() {
 	return ncpu;
 }
 
+int get_pid_max()
+{
+#ifdef __linux__
+	//read /proc/sys/kernel/pid_max
+	static char buffer[1024];
+	FILE *fd = fopen("/proc/sys/kernel/pid_max", "r");
+	if (fd==NULL) return -1;
+	if (fgets(buffer, sizeof(buffer), fd)==NULL) {
+		fclose(fd);
+		return -1;
+	}
+	fclose(fd);
+	return atoi(buffer);
+#elif defined __FreeBSD__
+	return 99998;
+#elif defined __APPLE__
+	return 99998;
+#endif
+}
+
 void limit_process(pid_t pid, double limit, int include_children)
 {
 	//slice of the slot in which the process is allowed to run
@@ -347,7 +371,7 @@ int main(int argc, char **argv) {
 		}
 	} while(next_option != -1);
 
-	if (pid_ok && (pid <= 1 || pid >= 65536)) {
+	if (pid_ok && (pid <= 1 || pid >= get_pid_max())) {
 		fprintf(stderr,"Error: Invalid value for argument PID\n");
 		print_usage(stderr, 1);
 		exit(1);
diff --git a/src/memrchr.c b/src/memrchr.c
new file mode 100644
index 0000000..ba788f5
--- /dev/null
+++ b/src/memrchr.c
@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2007 Todd C. Miller <Todd.Miller@courtesan.com>
+ *
+ * Permission to use, copy, modify, and distribute this software for any
+ * purpose with or without fee is hereby granted, provided that the above
+ * copyright notice and this permission notice appear in all copies.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+ * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+ * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+ * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+ * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+ */
+
+#include <sys/types.h>
+
+/*
+ * Reverse memchr()
+ * Find the last occurrence of 'c' in the buffer 's' of size 'n'.
+ */
+void *
+memrchr(s, c, n)
+    const void *s;
+    int c;
+    size_t n;
+{
+    const unsigned char *cp;
+
+    if (n != 0) {
+	cp = (unsigned char *)s + n;
+	do {
+	    if (*(--cp) == (unsigned char)c)
+		return((void *)cp);
+	} while (--n != 0);
+    }
+    return((void *)0);
+}
diff --git a/src/process_iterator.c b/src/process_iterator.c
index 165f420..8b4019d 100644
--- a/src/process_iterator.c
+++ b/src/process_iterator.c
@@ -22,7 +22,9 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#ifndef __APPLE__
 #include <sys/procfs.h>
+#endif
 #include <time.h>
 #include "process_iterator.h"
 
diff --git a/src/process_iterator.h b/src/process_iterator.h
index ecbe995..70520b6 100644
--- a/src/process_iterator.h
+++ b/src/process_iterator.h
@@ -81,7 +81,9 @@ struct process_iterator {
 	int count;
 	int i;
 #elif defined __APPLE__
-	struct kinfo_proc *proclist;
+	int i;
+	int count;
+	int *pidlist;
 #endif
 	struct process_filter *filter;
 };
diff --git a/src/process_iterator_apple.c b/src/process_iterator_apple.c
index d7aa5f3..b878ed8 100644
--- a/src/process_iterator_apple.c
+++ b/src/process_iterator_apple.c
@@ -17,82 +17,132 @@
  * You should have received a copy of the GNU General Public License
  * along with this program; if not, write to the Free Software
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ * 
+ * Author: Simon Sigurdhsson
+ *
  */
 
-int init_process_iterator(struct process_iterator *it) {
+#include <errno.h>
+#include <stdio.h>
+#include <libproc.h>
+
+int unique_nonzero_ints(int* arr_in, int len_in, int* arr_out) {
+	int* source = arr_in;
+	if (arr_out == NULL) return -1;
+	if (arr_in == arr_out) {
+		source = malloc(sizeof(int)*len_in);
+		memcpy(source, arr_in, sizeof(int)*len_in);
+		memset(arr_out, -1, sizeof(int)*len_in);
+	}
+	int len_out = 0;
+        int i, j;
+	for (i=0; i<len_in; i++) {
+		int found = 0;
+		if (source[i] == 0) continue;
+		for (j=0; !found && j<len_out; j++) {
+			found = (source[i] == arr_out[j]) ? 1 : 0;
+		}
+		if (!found) {
+			arr_out[len_out++] = source[i];
+		}
+	}
+	if (arr_in == arr_out) {
+		free(source);
+	}
+	return len_out-1;
+}
+
+int init_process_iterator(struct process_iterator *it, struct process_filter *filter) {
+	it->i = 0;
+	/* Find out how much to allocate for it->pidlist */
+	if ((it->count = proc_listpids(PROC_ALL_PIDS, 0, NULL, 0)) <= 0) {
+		fprintf(stderr, "proc_listpids: %s\n", strerror(errno));
+		return -1;
+	}
+	/* Allocate and populate it->pidlist */
+	if ((it->pidlist = (int *)malloc((it->count)*sizeof(int))) == NULL) {
+		fprintf(stderr, "malloc: %s\n", strerror(errno));
+	}
+	if ((it->count = proc_listpids(PROC_ALL_PIDS, 0, it->pidlist, it->count)) <= 0) {
+		fprintf(stderr, "proc_listpids: %s\n", strerror(errno));
+		return -1;
+	}
+	it->count = unique_nonzero_ints(it->pidlist, it->count, it->pidlist);
+	it->filter = filter;
+	return 0;
+}
+
+static int pti2proc(struct proc_taskallinfo *ti, struct process *process) {
+	int bytes;
+	process->pid = ti->pbsd.pbi_pid;
+	process->ppid = ti->pbsd.pbi_ppid;
+	process->starttime = ti->pbsd.pbi_start_tvsec;
+	process->cputime = (ti->ptinfo.pti_total_user + ti->ptinfo.pti_total_system) / 1000000;
+	bytes = strlen(ti->pbsd.pbi_comm);
+	memcpy(process->command, ti->pbsd.pbi_comm, (bytes < PATH_MAX ? bytes : PATH_MAX) + 1);
+	return 0;
+}
+
+static int get_process_pti(pid_t pid, struct proc_taskallinfo *ti) {
+	int bytes;
+	bytes = proc_pidinfo(pid, PROC_PIDTASKALLINFO, 0, ti, sizeof(*ti));
+	if (bytes <= 0) {
+		if (!(errno & (EPERM | ESRCH))) {
+			fprintf(stderr, "proc_pidinfo: %s\n", strerror(errno));
+		}
+		return -1;
+	} else if (bytes < sizeof(ti)) {
+		fprintf(stderr, "proc_pidinfo: too few bytes; expected %ld, got %d\n", sizeof(ti), bytes);
+		return -1;
+	}
 	return 0;
 }
 
 int get_next_process(struct process_iterator *it, struct process *p) {
+	if (it->i == it->count) return -1;
+	if (it->filter->pid != 0 && !it->filter->include_children) {
+		struct proc_taskallinfo ti;
+		if (get_process_pti(it->filter->pid, &ti) != 0) {
+			it->i = it->count = 0;
+			return -1;
+		}
+		it->i = it->count = 1;
+		return pti2proc(&ti, p);
+	}
+	while (it->i < it->count) {
+		struct proc_taskallinfo ti;
+		if (get_process_pti(it->pidlist[it->i], &ti) != 0) {
+			it->i++;
+			continue;
+		}
+		if (ti.pbsd.pbi_flags & PROC_FLAG_SYSTEM) {
+			it->i++;
+			continue;
+		}
+		if (it->filter->pid != 0 && it->filter->include_children) {
+			pti2proc(&ti, p);
+			it->i++;
+			if (p->pid != it->pidlist[it->i - 1]) // I don't know why this can happen
+				continue;
+			if (p->pid != it->filter->pid && p->ppid != it->filter->pid)
+				continue;
+			return 0;
+		}
+		else if (it->filter->pid == 0)
+		{
+			pti2proc(&ti, p);
+			it->i++;
+			return 0;
+		}
+	}
 	return -1;
 }
 
 int close_process_iterator(struct process_iterator *it) {
+	free(it->pidlist);
+	it->pidlist = NULL;
+	it->filter = NULL;
+	it->count = 0;
+	it->i = 0;
 	return 0;
 }
-
-	// int err;
-	// struct kinfo_proc *result = NULL;
-	// size_t length;
-	// int mib[] = {CTL_KERN, KERN_PROC, KERN_PROC_ALL, 0};
-
-	// /* We start by calling sysctl with result == NULL and length == 0.
-	//    That will succeed, and set length to the appropriate length.
-	//    We then allocate a buffer of that size and call sysctl again
-	//    with that buffer.
-	// */
-	// length = 0;
-	// err = sysctl(mib, 4, NULL, &length, NULL, 0);
-	// if (err == -1) {
-	// 	err = errno;
-	// }
-	// if (err == 0) {
-	// 	result = malloc(length);
-	// 	err = sysctl(mib, 4, result, &length, NULL, 0);
-	// 	if (err == -1)
-	// 		err = errno;
-	// 	if (err == ENOMEM) {
-	// 		free(result); /* clean up */
-	// 		result = NULL;
-	// 	}
-	// }
-
-	// i->proclist = result;
-	// i->count = err == 0 ? length / sizeof *result : 0;
-	// i->c = 0;
-
-// int get_proc_info(struct process *p, pid_t pid) {
-// 	int err;
-// 	struct kinfo_proc *result = NULL;
-// 	size_t length;
-// 	int mib[] = {CTL_KERN, KERN_PROC, KERN_PROC_PID, pid};
-
-// 	/* We start by calling sysctl with result == NULL and length == 0.
-// 	   That will succeed, and set length to the appropriate length.
-// 	   We then allocate a buffer of that size and call sysctl again
-// 	   with that buffer.
-// 	*/
-// 	length = 0;
-// 	err = sysctl(mib, 4, NULL, &length, NULL, 0);
-// 	if (err == -1) {
-// 		err = errno;
-// 	}
-// 	if (err == 0) {
-// 		result = malloc(length);
-// 		err = sysctl(mib, 4, result, &length, NULL, 0);
-// 		if (err == -1)
-// 			err = errno;
-// 		if (err == ENOMEM) {
-// 			free(result); /* clean up */
-// 			result = NULL;
-// 		}
-// 	}
-
-// 	p->pid = result->kp_proc.p_pid;
-// 	p->ppid = result->kp_eproc.e_ppid;
-// 	p->starttime = result->kp_proc.p_starttime.tv_sec;
-// 	p->last_jiffies = result->kp_proc.p_cpticks;
-// 	//p_pctcpu
-
-// 	return 0;
-// }
diff --git a/tests/Makefile b/tests/Makefile
index ffc80fa..764db87 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -12,7 +12,7 @@ endif
 
 all::	$(TARGETS)
 
-busy:	busy.c $(SYSLIBS)
+busy:	busy.c
 	$(CC) -o busy busy.c $(SYSLIBS) $(CFLAGS)
 
 process_iterator_test: process_iterator_test.c $(LIBS)
diff --git a/tests/process_iterator_test.c b/tests/process_iterator_test.c
index 3d5e885..20ef479 100644
--- a/tests/process_iterator_test.c
+++ b/tests/process_iterator_test.c
@@ -28,6 +28,10 @@
 #include <signal.h>
 #include <string.h>
 
+#ifdef __APPLE__
+#include <libgen.h>
+#endif
+
 #include <process_iterator.h>
 #include <process_group.h>
 
@@ -198,7 +202,13 @@ void test_process_name(const char * command)
 	assert(get_next_process(&it, &process) == 0);
 	assert(process.pid == getpid());
 	assert(process.ppid == getppid());
+	#ifdef __APPLE__
+	// proc_pidinfo only gives us the first 15 chars
+	// of the basename of the command on OSX.
+	assert(strncmp(basename((char*)command), process.command, 15) == 0);
+	#else
 	assert(strncmp(command, process.command, strlen(process.command)) == 0);
+	#endif
 	assert(get_next_process(&it, &process) != 0);
 	close_process_iterator(&it);
 }