/*-
 * Copyright (c) 2023 Klara, Inc.
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <errno.h>
#include <stdio.h>

#include <atf-c.h>

#define BUFSIZE 16

static const char seq[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

struct stream {
        char buf[BUFSIZE];
        unsigned int len;
        unsigned int pos;
};

static int
writefn(void *cookie, const char *buf, int len)
{
        struct stream *s = cookie;
        int written = 0;

        if (len <= 0)
                return (0);
        while (len > 0 && s->pos < s->len) {
                s->buf[s->pos++] = *buf++;
                written++;
                len--;
        }
        if (written > 0)
                return (written);
        errno = EAGAIN;
        return (-1);
}

ATF_TC_WITHOUT_HEAD(flushlbuf_partial);
ATF_TC_BODY(flushlbuf_partial, tc)
{
        static struct stream s;
        static char buf[BUFSIZE + 1];
        FILE *f;
        unsigned int i = 0;
        int ret = 0;

        /*
         * Create the stream and its buffer, print just enough characters
         * to the stream to fill the buffer without triggering a flush,
         * then check the state.
         */
        s.len = BUFSIZE / 2; // write will fail after this amount
        ATF_REQUIRE((f = fwopen(&s, writefn)) != NULL);
        ATF_REQUIRE(setvbuf(f, buf, _IOLBF, BUFSIZE) == 0);
        while (i < BUFSIZE)
                if ((ret = fprintf(f, "%c", seq[i++])) < 0)
                        break;
        ATF_CHECK_EQ(BUFSIZE, i);
        ATF_CHECK_EQ(seq[i - 1], buf[BUFSIZE - 1]);
        ATF_CHECK_EQ(1, ret);
        ATF_CHECK_EQ(0, s.pos);

        /*
         * At this point, the buffer is full but writefn() has not yet
         * been called.  The next fprintf() call will trigger a preemptive
         * fflush(), and writefn() will consume s.len characters before
         * returning EAGAIN, causing fprintf() to fail without having
         * written anything (which is why we don't increment i here).
         */
        ret = fprintf(f, "%c", seq[i]);
        ATF_CHECK_ERRNO(EAGAIN, ret < 0);
        ATF_CHECK_EQ(s.len, s.pos);

        /*
         * We have consumed s.len characters from the buffer, so continue
         * printing until it is full again and check that no overflow has
         * occurred yet.
         */
        while (i < BUFSIZE + s.len)
                fprintf(f, "%c", seq[i++]);
        ATF_CHECK_EQ(BUFSIZE + s.len, i);
        ATF_CHECK_EQ(seq[i - 1], buf[BUFSIZE - 1]);
        ATF_CHECK_EQ(0, buf[BUFSIZE]);

        /*
         * The straw that breaks the camel's back: libc fails to recognize
         * that the buffer is full and continues to write beyond its end.
         */
        fprintf(f, "%c", seq[i++]);
        ATF_CHECK_EQ(0, buf[BUFSIZE]);
}

ATF_TC_WITHOUT_HEAD(flushlbuf_full);
ATF_TC_BODY(flushlbuf_full, tc)
{
        static struct stream s;
        static char buf[BUFSIZE];
        FILE *f;
        unsigned int i = 0;
        int ret = 0;

        /*
         * Create the stream and its buffer, print just enough characters
         * to the stream to fill the buffer without triggering a flush,
         * then check the state.
         */
        s.len = 0; // any attempt to write will fail
        ATF_REQUIRE((f = fwopen(&s, writefn)) != NULL);
        ATF_REQUIRE(setvbuf(f, buf, _IOLBF, BUFSIZE) == 0);
        while (i < BUFSIZE)
                if ((ret = fprintf(f, "%c", seq[i++])) < 0)
                        break;
        ATF_CHECK_EQ(BUFSIZE, i);
        ATF_CHECK_EQ(seq[i - 1], buf[BUFSIZE - 1]);
        ATF_CHECK_EQ(1, ret);
        ATF_CHECK_EQ(0, s.pos);

        /*
         * At this point, the buffer is full but writefn() has not yet
         * been called.  The next fprintf() call will trigger a preemptive
         * fflush(), and writefn() will immediately return EAGAIN, causing
         * fprintf() to fail without having written anything (which is why
         * we don't increment i here).
         */
        ret = fprintf(f, "%c", seq[i]);
        ATF_CHECK_ERRNO(EAGAIN, ret < 0);
        ATF_CHECK_EQ(s.len, s.pos);

        /*
         * Now make our stream writeable.
         */
        s.len = sizeof(s.buf);

        /*
         * Flush the stream again.  The data we failed to write previously
         * should still be in the buffer and will now be written to the
         * stream.
         */
        ATF_CHECK_EQ(0, fflush(f));
        ATF_CHECK_EQ(seq[0], s.buf[0]);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, flushlbuf_partial);
        ATF_TP_ADD_TC(tp, flushlbuf_full);

        return (atf_no_error());
}