ZSTD_DCtx_refPrefix(): whether the prefix survives ZSTD_DCtx_reset(ZSTD_reset_session_only) after an error depends on where the frame failed #4843

/* ZSTD_DCtx_refPrefix(): a prefix referenced for a frame whose header is rejected survives
 * ZSTD_DCtx_reset(ZSTD_reset_session_only) and is used by the next frame.
 *
 * Frame P is compressed against a prefix (it cannot be decoded without it).
 * Each case: reference the prefix, feed a first input, reset the session (as documented after an error),
 * then decode frame P WITHOUT referencing the prefix again.
 *
 *   cc -O1 -g -fsanitize=address,undefined prefix_after_reset.c -Ilib lib/libzstd.a -o prefix_after_reset
 *   ./prefix_after_reset          functional cases (no memory is freed)
 *   ./prefix_after_reset free     the application frees the prefix buffer after the reset (run under ASan)
 */
#define ZSTD_STATIC_LINKING_ONLY   /* only for ZSTD_frameHeaderSize(), used to build the corrupted test input */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <zstd.h>

#define N 8192
static unsigned char src[N], comp[2 * N], out[8 * N];

static size_t decode(ZSTD_DCtx *d, const void *in, size_t n, size_t *outlen) {
    ZSTD_inBuffer ib = { in, n, 0 };
    ZSTD_outBuffer ob = { out, sizeof out, 0 };
    size_t r;
    do { r = ZSTD_decompressStream(d, &ob, &ib); } while (!ZSTD_isError(r) && r != 0 && ib.pos < ib.size);
    *outlen = ob.pos;
    return r;
}

static const char *res(size_t r, size_t outlen) {
    static char buf[2][64];
    static int k;
    if (ZSTD_isError(r)) return ZSTD_getErrorName(r);
    k ^= 1;
    snprintf(buf[k], sizeof buf[k], "ok, %zu bytes", outlen);
    return buf[k];
}

int main(int argc, char **argv) {
    int free_prefix = argc > 1 && !strcmp(argv[1], "free");
    unsigned char *prefix = malloc(N);
    unsigned long long x = 88172645463325252ULL;
    for (size_t i = 0; i < N; i++) { x ^= x << 13; x ^= x >> 7; x ^= x << 17; src[i] = prefix[i] = (unsigned char)x; }

    /* frame P: payload equal to the prefix, so it is one long match into the prefix */
    ZSTD_CCtx *c = ZSTD_createCCtx();
    ZSTD_CCtx_refPrefix(c, prefix, N);
    size_t n = ZSTD_compress2(c, comp, sizeof comp, src, N);
    ZSTD_freeCCtx(c);
    if (ZSTD_isError(n)) { printf("compress: %s\n", ZSTD_getErrorName(n)); return 1; }

    unsigned char bad_header[2 * N], bad_block[2 * N];
    memcpy(bad_header, comp, n);
    bad_header[4] |= 0x08;                          /* frame header descriptor: reserved bit -> rejected header */
    memcpy(bad_block, comp, n);
    bad_block[ZSTD_frameHeaderSize(comp, n)] |= 0x06; /* first block header: reserved block type */

    printf("zstd %s, frame P = %zu bytes\n", ZSTD_versionString(), n);
    {   ZSTD_DCtx *d = ZSTD_createDCtx(); size_t ol, r = decode(d, comp, n, &ol);
        printf("fresh DCtx, no prefix, frame P:                         %s\n", res(r, ol));
        ZSTD_freeDCtx(d); }

    struct { const char *name; const unsigned char *first; ZSTD_ResetDirective reset; int ref_null; } cases[] = {
        { "first frame OK,             reset session_only",            comp,       ZSTD_reset_session_only, 0 },
        { "first frame bad block,      reset session_only",            bad_block,  ZSTD_reset_session_only, 0 },
        { "first frame bad header,     reset session_only",            bad_header, ZSTD_reset_session_only, 0 },
        { "first frame bad header,     reset session_and_parameters",  bad_header, ZSTD_reset_session_and_parameters, 0 },
        { "first frame bad header,     reset session_only + refPrefix(NULL)", bad_header, ZSTD_reset_session_only, 1 },
    };
    for (size_t i = 0; i < sizeof cases / sizeof cases[0]; i++) {
        if (free_prefix && i != 2) continue;
        ZSTD_DCtx *d = ZSTD_createDCtx();
        ZSTD_DCtx_refPrefix(d, prefix, N);
        size_t ol1, ol, r1 = decode(d, cases[i].first, n, &ol1);
        ZSTD_DCtx_reset(d, cases[i].reset);
        if (cases[i].ref_null) ZSTD_DCtx_refPrefix(d, NULL, 0);
        if (free_prefix) { free(prefix); prefix = NULL; }   /* the failed frame is over; release its prefix */
        size_t r2 = decode(d, comp, n, &ol);                 /* next frame, no prefix referenced for it */
        printf("%-62s first: %-28s then frame P: %s\n", cases[i].name, res(r1, ol1), res(r2, ol));
        ZSTD_freeDCtx(d);
    }
    free(prefix);
    return 0;
}