/*
 * REP STOS single-stepping
 * Copyright (C) 2008  Vegard Nossum
 *
 * This program is distributed under the terms of the GNU GPL 2.
 *
 * Note on the observed behaviour: On newer machine, single-stepping a REP
 * instruction will generate a Debug Exception for each repetition. This is
 * not necessarily true for all CPUs. Therefore, on CPUs which only report
 * one trap per REP STOS, we would usually get a lower number of traps than
 * the number of iterations. This is because the instruction may still be
 * interrupted by for example the timer interrupt, and then we'll get one
 * trap just before the interrupt handler is entered.
 */

#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>

#include <sys/user.h>

#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/* Taken from the Linux Kernel v2.6.26. */
static inline void *__memset_generic(void *s, char c, size_t count)
{
	int d0, d1;
	asm volatile("rep\n\t"
		     "stosb"
		     : "=&c" (d0), "=&D" (d1)
		     : "a" (c), "1" (s), "0" (count)
		     : "memory");
	return s;
}

int
main(int argc, char *argv[])
{
	unsigned int size = 1000;

	if (argc == 2)
		size = atoi(argv[1]);

	char *buffer = malloc(size);
	if (!buffer) {
		fprintf(stderr, "error: malloc: %s\n", strerror(errno));
		exit(EXIT_FAILURE);
	}

	pid_t child = fork();
	if (child == -1) {
		fprintf(stderr, "error: fork: %s\n", strerror(errno));
		exit(EXIT_FAILURE);
	}

	if (child == 0) {
		/* We're the child */
		if (ptrace(PTRACE_TRACEME, 0, NULL, NULL) == -1) {
			fprintf(stderr, "error: ptrace: PTRACE_TRACEME: %s\n",
				strerror(errno));
			exit(EXIT_FAILURE);
		}

		/* Allow parent to single-step us */
		if (kill(getpid(), SIGTRAP) == -1) {
			fprintf(stderr, "error: kill: %s\n", strerror(errno));
			exit(EXIT_FAILURE);
		}

		__memset_generic(buffer, 0, size);

		exit(EXIT_SUCCESS);
	}

	unsigned int rep_insns = 0;

	/* We're the parent */
	while (1) {
		pid_t wchild;
		int status;

		wchild = wait(&status);
		if (wchild == -1) {
			if (errno != ECHILD) {
				fprintf(stderr, "error: wait: %s\n",
					strerror(errno));
				exit(EXIT_FAILURE);
			}

			break;
		}

		if (wchild != child) {
			fprintf(stderr, "warning: wait: got %d, expected %d\n",
				wchild, child);
			continue;
		}

		if (WIFEXITED(status))
			break;

		if (WIFSTOPPED(status) && WSTOPSIG(status) == SIGTRAP) {
			struct user_regs_struct regs;

			if (ptrace(PTRACE_GETREGS, wchild, NULL, &regs) == -1) {
				fprintf(stderr, "error: ptrace: "
					"PTRACE_GETREGS: %s\n",
					strerror(errno));
				exit(EXIT_FAILURE);
			}

			/* Might be unaligned... Oh well. */
			long ateip = ptrace(PTRACE_PEEKTEXT, wchild,
				(void *) regs.eip, NULL);

			/* Note: The instruction is reversed in memory */
			if ((ateip & 0xffff) == 0xaaf3)
				++rep_insns;

			/* Go on... */
			if (ptrace(PTRACE_SINGLESTEP, wchild,
				NULL, NULL) == -1)
			{
				fprintf(stderr, "error: ptrace: "
					"PTRACE_SINGLESTEP: %s\n",
					strerror(errno));
				exit(EXIT_FAILURE);
			}
		}
	}

	free(buffer);

	int status = system("grep '^\\(processor\\|cpu family\\|model\\)' "
		"/proc/cpuinfo | tail -n 4");
	if (status == -1)
	{
		fprintf(stderr, "warning: system: %s\n", strerror(errno));
	}

	if (WEXITSTATUS(status)) {
		fprintf(stderr, "warning: grep returned %d\n",
			WEXITSTATUS(status));
	}

	printf("\n");
	printf("Counted %d REP STOS instructions (expected %d).\n",
		rep_insns, size);

	return EXIT_SUCCESS;
}