#define _KERNTYPES

#include <sys/types.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/sysctl.h>
#if defined(__FreeBSD__)
#include <sys/user.h>
#endif
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <err.h>
#include <sched.h>
#include <inttypes.h>

#define ATF_REQUIRE_MSG(a,b,...) assert(a)
#define ATF_REQUIRE_EQ_MSG(a,b,c, ...) assert((a) == (b))
#define TWAIT_HAVE_STATUS
#define TWAIT_GENERIC(a,b,c) waitpid((a),(b),(c))
#define TEST_REQUIRE_EQ(a,b) assert((a) == (b))
#define TWAIT_REQUIRE_SUCCESS(a,b) TEST_REQUIRE_EQ((a), (b))
#define TWAIT_FNAME ""
#define ATF_REQUIRE(x) assert(x)
#define ATF_REQUIRE_EQ(a,b) assert((a) == (b))
#define TWAIT_REQUIRE_FAILURE(a,b) do { assert((b) == -1); assert(errno == (a)); } while(0)
#define DPRINTF(x, ...) printf(x, ##__VA_ARGS__)
#define atf_utils_fork fork

#define FORKEE_REQUIRE_SUCCESS(a,b)   FORKEE_ASSERT_EQ(a, b)

#ifndef __used
#define __used __attribute__((used))
#endif

#ifdef linux
size_t
strlcpy(char *dst, const char *src, size_t siz)
{
        char *d = dst;
        const char *s = src;
        size_t n = siz;

        /* Copy as many bytes as will fit */
        if (n != 0) {
                while (--n != 0) {
                        if ((*d++ = *s++) == '\0')
                                break;
                }
        }

        /* Not enough room in dst, add NUL and traverse rest of src */
        if (n == 0) {
                if (siz != 0)
                        *d = '\0';              /* NUL-terminate dst */
                while (*s++)
                        ;
        }
 
        return(s - src - 1);    /* count does not include NUL */
}
#endif

#define FORKEE_ASSERT_EQ(x, y)                                          \
do {                                                                    \
        uintmax_t vx = (x);                                             \
        uintmax_t vy = (y);                                             \
        int ret = vx == vy;                                             \
        if (!ret)                                                       \
                errx(EXIT_FAILURE, "%s:%d %s(): Assertion failed for: " \
                    "%s(%ju) == %s(%ju)", __FILE__, __LINE__, __func__, \
                    #x, vx, #y, vy);                                    \
} while (/*CONSTCOND*/0)


#define FORKEE_ASSERTX(x)                                               \
do {                                                                    \
        int ret = (x);                                                  \
        if (!ret)                                                       \
                errx(EXIT_FAILURE, "%s:%d %s(): Assertion failed for: %s",\
                     __FILE__, __LINE__, __func__, #x);                 \
} while (/*CONSTCOND*/0)

#define FORKEE_ASSERT(x)                                                \
do {                                                                    \
        int ret = (x);                                                  \
        if (!ret)                                                       \
                err(EXIT_FAILURE, "%s:%d %s(): Assertion failed for: %s",\
                     __FILE__, __LINE__, __func__, #x);                 \
} while (/*CONSTCOND*/0)   

static void __used
validate_status_stopped(int status, int expected)
{
        ATF_REQUIRE_MSG(!WIFEXITED(status), "Reported exited process");
        ATF_REQUIRE_MSG(!WIFCONTINUED(status), "Reported continued process");
        ATF_REQUIRE_MSG(!WIFSIGNALED(status), "Reported signaled process");
        ATF_REQUIRE_MSG(WIFSTOPPED(status), "Reported !stopped process");

        char st[128], ex[128];
        strlcpy(st, strsignal(WSTOPSIG(status)), sizeof(st));
        strlcpy(ex, strsignal(expected), sizeof(ex));

        ATF_REQUIRE_EQ_MSG(WSTOPSIG(status), expected,   
            "Unexpected stop signal received [%s] != [%s]", st, ex);
}
 
static void __used
forkee_status_stopped(int status, int expected)
{
        FORKEE_ASSERTX(!WIFEXITED(status));
        FORKEE_ASSERTX(!WIFCONTINUED(status));
        FORKEE_ASSERTX(!WIFSIGNALED(status));
        FORKEE_ASSERTX(WIFSTOPPED(status));

        FORKEE_ASSERT_EQ(WSTOPSIG(status), expected);
}

static void __used
forkee_status_exited(int status, int expected)
{
        FORKEE_ASSERTX(WIFEXITED(status));
        FORKEE_ASSERTX(!WIFCONTINUED(status));
        FORKEE_ASSERTX(!WIFSIGNALED(status));
        FORKEE_ASSERTX(!WIFSTOPPED(status));

        FORKEE_ASSERT_EQ(WEXITSTATUS(status), expected);
}   

static void __used
validate_status_exited(int status, int expected)
{
        ATF_REQUIRE_MSG(WIFEXITED(status), "Reported !exited process");
        ATF_REQUIRE_MSG(!WIFCONTINUED(status), "Reported continued process");
        ATF_REQUIRE_MSG(!WIFSIGNALED(status), "Reported signaled process");
        ATF_REQUIRE_MSG(!WIFSTOPPED(status), "Reported stopped process");
        
        ATF_REQUIRE_EQ_MSG(WEXITSTATUS(status), expected,
            "The process has exited with invalid value %d != %d",
            WEXITSTATUS(status), expected);
}


/* This function is currently designed to be run in the main/parent process */
static void __used
await_zombie(pid_t process)
{
#if defined(__NetBSD__)
        struct kinfo_proc2 p;
        size_t len = sizeof(p);

        const int name[] = {
                [0] = CTL_KERN,
                [1] = KERN_PROC2,
                [2] = KERN_PROC_PID,
                [3] = process,
                [4] = sizeof(p),
                [5] = 1
        };

        const size_t namelen = __arraycount(name);

        /* Await the process becoming a zombie */
        while(1) {
                int rv = sysctl(name, namelen, &p, &len, NULL, 0);
		if (rv != 0) {
			printf("rv=%d errno=%d\n", rv, errno);
			abort();
		}

                if (p.p_stat == LSZOMB)
                        break;

//		sched_yield();
//                ATF_REQUIRE(usleep(10) == 0);
        }
#elif defined(__FreeBSD__)
	/*
	 * Wait for a process to exit.  This is kind of gross, but
	 * there is not a better way.
	 *
	 * Prior to r325719, the kern.proc.pid.<pid> sysctl failed
	 * with ESRCH.  After that change, a valid struct kinfo_proc
	 * is returned for zombies with ki_stat set to SZOMB.
	 */
	for (;;) {
		struct kinfo_proc kp;
		size_t len;
		int mib[4];

		mib[0] = CTL_KERN;
		mib[1] = KERN_PROC;
		mib[2] = KERN_PROC_PID;
		mib[3] = process;
		len = sizeof(kp);
		if (sysctl(mib, nitems(mib), &kp, &len, NULL, 0) == -1) {
			ATF_REQUIRE(errno == ESRCH);
			break;
		}
		if (kp.ki_stat == SZOMB)
			break;
//		usleep(5000);
	}
#elif defined(__linux__)
        int iszombie = 0;    
        char pbuf[60];
        snprintf(pbuf, sizeof(pbuf), "/proc/%d/stat", (int)process);
        
        /* Detect that tracee is zombie */
        while(1) {
                FILE* fpstat = fopen(pbuf, "r");                                                                                      
                if (!fpstat) {
                        err(EXIT_FAILURE, "fopen");
                };
                int rpid =0; char rcmd[32]; char rstatc = 0;
                fscanf(fpstat, "%d %30s %c", &rpid, rcmd, &rstatc);
                iszombie = rstatc == 'Z';

                fclose(fpstat);
                if (iszombie == 1) {
                        break;
                }

//                usleep(1000);
        }
#else
#error fatal
#endif
}



/*	$NetBSD: msg.h,v 1.1 2017/04/02 21:44:00 kamil Exp $	*/

/*-
 * Copyright (c) 2016 The NetBSD Foundation, Inc.
 * All rights reserved.
 *
 * This code is derived from software contributed to The NetBSD Foundation
 * by Christos Zoulas.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

struct msg_fds {
	int pfd[2];
	int cfd[2];
};

#define CLOSEFD(fd) do { \
	if (fd != -1) { \
		close(fd); \
		fd = -1; \
	} \
} while (/*CONSTCOND*/ 0)

static int
msg_open(struct msg_fds *fds)
{
	if (pipe(fds->pfd) == -1)
		return -1;
	if (pipe(fds->cfd) == -1) {
		close(fds->pfd[0]);
		close(fds->pfd[1]);
		return -1;
	}
	return 0;
}

static void
msg_close(struct msg_fds *fds)
{
	CLOSEFD(fds->pfd[0]);
	CLOSEFD(fds->pfd[1]);
	CLOSEFD(fds->cfd[0]);
	CLOSEFD(fds->cfd[1]);
}

static int
msg_write_child(const char *info, struct msg_fds *fds, void *msg, size_t len)
{
	ssize_t rv;
	CLOSEFD(fds->cfd[1]);
	CLOSEFD(fds->pfd[0]);

	printf("Send %s\n", info);
	rv = write(fds->pfd[1], msg, len);
	if (rv != (ssize_t)len)
		return 1;
//	printf("Wait %s\n", info);
	rv = read(fds->cfd[0], msg, len);
	if (rv != (ssize_t)len)
		return 1;
	return 0;
}

static int
msg_write_parent(const char *info, struct msg_fds *fds, void *msg, size_t len)
{
	ssize_t rv;
	CLOSEFD(fds->pfd[1]);
	CLOSEFD(fds->cfd[0]);

	printf("Send %s\n", info);
	rv = write(fds->cfd[1], msg, len);
	if (rv != (ssize_t)len)
		return 1;
//	printf("Wait %s\n", info);
	rv = read(fds->pfd[0], msg, len);
	if (rv != (ssize_t)len)
		return 1;
	return 0;
}

static int
msg_read_parent(const char *info, struct msg_fds *fds, void *msg, size_t len)
{
	ssize_t rv;
	CLOSEFD(fds->pfd[1]);
	CLOSEFD(fds->cfd[0]);

	printf("Wait %s\n", info);
	rv = read(fds->pfd[0], msg, len);
	if (rv != (ssize_t)len)
		return 1;
//	printf("Send %s\n", info);
	rv = write(fds->cfd[1], msg, len);
	if (rv != (ssize_t)len)
		return 1;
	return 0;
}

static int
msg_read_child(const char *info, struct msg_fds *fds, void *msg, size_t len)
{
	ssize_t rv;
	CLOSEFD(fds->cfd[1]);
	CLOSEFD(fds->pfd[0]);

	printf("Wait %s\n", info);
	rv = read(fds->cfd[0], msg, len);
	if (rv != (ssize_t)len)
		return 1;
//	printf("Send %s\n", info);
	rv = write(fds->pfd[1], msg, len);
	if (rv != (ssize_t)len)
		return 1;
	return 0;
}

#define PARENT_TO_CHILD(info, fds, msg) \
    SYSCALL_REQUIRE(msg_write_child(info " to child " # fds, &fds, &msg, sizeof(msg)) == 0)

#define CHILD_FROM_PARENT(info, fds, msg) \
    FORKEE_ASSERT(msg_read_parent(info " from parent " # fds, &fds, &msg, sizeof(msg)) == 0)

#define CHILD_TO_PARENT(info, fds, msg) \
    FORKEE_ASSERT(msg_write_parent(info " to parent " # fds, &fds, &msg, sizeof(msg)) == 0)

#define PARENT_FROM_CHILD(info, fds, msg) \
    SYSCALL_REQUIRE(msg_read_child(info " from parent " # fds, &fds, &msg, sizeof(msg)) == 0)

#define SYSCALL_REQUIRE(expr) ATF_REQUIRE_MSG(expr, "%s: %s", # expr, \
    strerror(errno))
#define SYSCALL_REQUIRE_ERRNO(res, exp) ATF_REQUIRE_MSG(res == exp, \
    "%d(%s) != %d", res, strerror(res), exp)



/* Happy number sequence -- this function is used to just consume cpu cycles */
#define HAPPY_NUMBER    1

/* If n is not happy then its sequence ends in the cycle:
 * 4, 16, 37, 58, 89, 145, 42, 20, 4, ... */
#define SAD_NUMBER      4

/* Calculate the sum of the squares of the digits of n */
static unsigned __used
dsum(unsigned n)
{
        unsigned sum, x;
        for (sum = 0; n; n /= 10) {
                x = n % 10;
                sum += x * x;
        }
        return sum;
}

static int __used
#ifdef __clang__
__attribute__((__optnone__))
#else
__attribute__((__optimize__("O0")))
#endif
check_happy(unsigned n)
{
        for (;;) {
                unsigned total = dsum(n);

                if (total == HAPPY_NUMBER)
                        return 1;
                if (total == SAD_NUMBER)
                        return 0;

                n = total;
        }
}

union u {
        unsigned long raw;
        struct {
                unsigned long local_dr0_breakpoint : 1;         /* 0 */
                unsigned long global_dr0_breakpoint : 1;        /* 1 */
                unsigned long local_dr1_breakpoint : 1;         /* 2 */
                unsigned long global_dr1_breakpoint : 1;        /* 3 */
                unsigned long local_dr2_breakpoint : 1;         /* 4 */
                unsigned long global_dr2_breakpoint : 1;        /* 5 */
                unsigned long local_dr3_breakpoint : 1;         /* 6 */
                unsigned long global_dr3_breakpoint : 1;        /* 7 */
                unsigned long local_exact_breakpt : 1;          /* 8 */
                unsigned long global_exact_breakpt : 1;         /* 9 */
                unsigned long reserved_10 : 1;                  /* 10 */
                unsigned long rest_trans_memory : 1;            /* 11 */
                unsigned long reserved_12 : 1;                  /* 12 */
                unsigned long general_detect_enable : 1;        /* 13 */
                unsigned long reserved_14 : 1;                  /* 14 */
                unsigned long reserved_15 : 1;                  /* 15 */
                unsigned long condition_dr0 : 2;                /* 16-17 */
                unsigned long len_dr0 : 2;                      /* 18-19 */
                unsigned long condition_dr1 : 2;                /* 20-21 */
                unsigned long len_dr1 : 2;                      /* 22-23 */
                unsigned long condition_dr2 : 2;                /* 24-25 */
                unsigned long len_dr2 : 2;                      /* 26-27 */
                unsigned long condition_dr3 : 2;                /* 28-29 */
                unsigned long len_dr3 : 2;                      /* 30-31 */
        } bits;
};

unsigned N = 0;

void
loop(void)
{
        struct msg_fds parent_tracer, parent_tracee;
        const int exitval_tracee = 5;
        const int exitval_tracer1 = 10, exitval_tracer2 = 20;
        pid_t tracee, tracer, wpid;
        uint8_t msg = 0xde; /* dummy message for IPC based on pipe(2) */
#if defined(TWAIT_HAVE_STATUS)
        int status;
#endif

	printf("N=%d\n", N++);

        DPRINTF("Spawn tracee\n");
        SYSCALL_REQUIRE(msg_open(&parent_tracee) == 0);
        tracee = atf_utils_fork();
        if (tracee == 0) {
                /* Wait for message from the parent */
                CHILD_FROM_PARENT("Message 1", parent_tracee, msg);
                _exit(exitval_tracee);
        }

        DPRINTF("Spawn debugger\n");
        SYSCALL_REQUIRE(msg_open(&parent_tracer) == 0);
        tracer = atf_utils_fork();
        if (tracer == 0) {
                /* Fork again and drop parent to reattach to PID 1 */
                tracer = atf_utils_fork();
                if (tracer != 0)
                        _exit(exitval_tracer1);

                DPRINTF("Before calling PT_ATTACH from tracee %d\n", getpid());
                FORKEE_ASSERT(ptrace(PT_ATTACH, tracee, NULL, 0) != -1);

                /* Wait for tracee and assert that it was stopped w/ SIGSTOP */
                FORKEE_REQUIRE_SUCCESS(                                                                                               
                    wpid = TWAIT_GENERIC(tracee, &status, 0), tracee);

                forkee_status_stopped(status, SIGSTOP);

                /* Resume tracee with PT_CONTINUE */
                FORKEE_ASSERT(ptrace(PT_CONTINUE, tracee, (void *)1, 0) != -1);

                /* Inform parent that tracer has attached to tracee */
                CHILD_TO_PARENT("Message 1", parent_tracer, msg);
                CHILD_FROM_PARENT("Message 2", parent_tracer, msg);

                /* Wait for tracee and assert that it exited */
                FORKEE_REQUIRE_SUCCESS(
                    wpid = TWAIT_GENERIC(tracee, &status, 0), tracee);

                forkee_status_exited(status, exitval_tracee);
                                                                                                                                      
                DPRINTF("Before exiting of the tracer process\n");
                _exit(exitval_tracer2);
        }
        DPRINTF("Wait for the tracer process (direct child) to exit calling "
            "%s()\n", TWAIT_FNAME);
        TWAIT_REQUIRE_SUCCESS(
            wpid = TWAIT_GENERIC(tracer, &status, 0), tracer);

        validate_status_exited(status, exitval_tracer1);

        DPRINTF("Wait for the non-exited tracee process with %s()\n",
            TWAIT_FNAME);
        TWAIT_REQUIRE_SUCCESS(
            wpid = TWAIT_GENERIC(tracee, NULL, WNOHANG), 0);

        DPRINTF("Wait for the tracer to attach to the tracee\n");
        PARENT_FROM_CHILD("Message 1", parent_tracer, msg);                                                                           
        DPRINTF("Resume the tracee and let it exit\n");
        PARENT_TO_CHILD("Message 1", parent_tracee, msg);

        DPRINTF("Detect that tracee is zombie\n");
        await_zombie(tracee);

        DPRINTF("Assert that there is no status about tracee - "
            "Tracer must detect zombie first - calling %s()\n", TWAIT_FNAME);
        TWAIT_REQUIRE_SUCCESS(
            wpid = TWAIT_GENERIC(tracee, &status, WNOHANG), 0);

        DPRINTF("Resume the tracer and let it detect exited tracee\n");
        PARENT_TO_CHILD("Message 2", parent_tracer, msg);

        DPRINTF("Wait for tracee to finish its job and exit - calling %s()\n",
            TWAIT_FNAME);
        TWAIT_REQUIRE_SUCCESS(wpid = TWAIT_GENERIC(tracee, &status, 0),                                                         
            tracee); // tutaj

        validate_status_exited(status, exitval_tracee);

        msg_close(&parent_tracer);
        msg_close(&parent_tracee);
}

int
main(int argc, char **argv)
{
	while (1) {
		loop();
	}

	return 0;
}