diff --git a/contrib/seekable_format/tests/seekable_tests.c b/contrib/seekable_format/tests/seekable_tests.c index 5fd9a51a80c..d878014d91c 100644 --- a/contrib/seekable_format/tests/seekable_tests.c +++ b/contrib/seekable_format/tests/seekable_tests.c @@ -7,6 +7,7 @@ #include "../zstd_seekable.h" #include "zstd_errors.h" +#include "xxhash.h" /* ZSTD_seekable_customFile implementation that reads/seeks a buffer while keeping track of total bytes read */ @@ -54,6 +55,44 @@ static int seekBuffWithTotal(void* opaque, long long offset, int origin) return 0; } +typedef struct { + buffWrapperWithTotal_t buff; + unsigned reads; + unsigned failAt; +} failingRead_t; + +static int readBuffFail(void* opaque, void* buffer, size_t n) +{ + failingRead_t* const src = (failingRead_t*)opaque; + if (++src->reads == src->failAt) return -1; + return readBuffWithTotal(&src->buff, buffer, n); +} + +static void* createSeekableBuffer(unsigned nbFrames, int checksumFlag, size_t* size) +{ + char const data[] = "a"; + char frame[100]; + size_t const cSize = ZSTD_compress(frame, sizeof(frame), data, sizeof(data), 1); + assert(!ZSTD_isError(cSize)); + size_t const capacity = nbFrames * (cSize + 12) + 17; + char* const buffer = malloc(capacity); + ZSTD_frameLog* const fl = ZSTD_seekable_createFrameLog(checksumFlag); + assert(buffer != NULL && fl != NULL); + unsigned i; + for (i = 0; i < nbFrames; ++i) { + memcpy(buffer + i * cSize, frame, cSize); + size_t const status = ZSTD_seekable_logFrame(fl, (unsigned)cSize, sizeof(data), + (unsigned)XXH64(data, sizeof(data), 0)); + assert(!ZSTD_isError(status)); + } + ZSTD_outBuffer out = {buffer, capacity, nbFrames * cSize}; + size_t const status = ZSTD_seekable_writeSeekTable(fl, &out); + assert(status == 0); + *size = out.pos; + ZSTD_seekable_freeFrameLog(fl); + return buffer; +} + /* Basic unit tests for zstd seekable format */ int main(int argc, const char** argv) { @@ -444,6 +483,54 @@ int main(int argc, const char** argv) } printf("Success!\n"); + printf("Test %u - reuse decoder after seek table load failures: ", testNb++); + { int checksumFlag; + for (checksumFlag = 0; checksumFlag <= 1; ++checksumFlag) { + unsigned const nbFrames = 40000; /* multiple seek table refills */ + size_t largeSize, smallSize; + void* const large = createSeekableBuffer(nbFrames, checksumFlag, &largeSize); + void* const small = createSeekableBuffer(1, checksumFlag, &smallSize); + ZSTD_seekable* const stream = ZSTD_seekable_create(); + assert(stream != NULL); + size_t const initStatus = ZSTD_seekable_initBuff(stream, small, smallSize); + assert(!ZSTD_isError(initStatus)); + ZSTD_seekTable* const copy = ZSTD_seekTable_create_fromSeekable(stream); + assert(copy != NULL); + failingRead_t src = {{large, largeSize, 0, 0}, 0, 0}; + ZSTD_seekable_customFile file = {&src, &readBuffFail, &seekBuffWithTotal}; + size_t const loadStatus = ZSTD_seekable_initAdvanced(stream, file); + assert(!ZSTD_isError(loadStatus)); + unsigned const reads = src.reads; + unsigned failAt; + int failed = 0; + for (failAt = 1; failAt <= reads; ++failAt) { + src.reads = 0; + src.failAt = failAt; + size_t const status = ZSTD_seekable_initAdvanced(stream, file); + if (ZSTD_getErrorCode(status) != ZSTD_error_seekableIO || + ZSTD_seekable_getNumFrames(stream) != nbFrames || + ZSTD_seekable_getFrameDecompressedOffset(stream, nbFrames - 1) != + (nbFrames - 1) * 2) { + failed = 1; + break; + } + } + size_t const retryStatus = ZSTD_seekable_initBuff(stream, small, smallSize); + char decoded[2]; + if (ZSTD_isError(retryStatus) || + ZSTD_seekable_decompress(stream, decoded, sizeof(decoded), 0) != sizeof(decoded) || + memcmp(decoded, "a", sizeof(decoded)) != 0) failed = 1; + ZSTD_seekable_free(stream); + if (ZSTD_seekTable_getNumFrames(copy) != 1 || + ZSTD_seekTable_getFrameDecompressedSize(copy, 0) != sizeof(decoded)) failed = 1; + ZSTD_seekTable_free(copy); + free(large); + free(small); + if (failed) goto _test_error; + } + } + printf("Success!\n"); + /* TODO: Add more tests */ printf("Finished tests\n"); return 0; diff --git a/contrib/seekable_format/zstdseek_decompress.c b/contrib/seekable_format/zstdseek_decompress.c index 352fbb5e863..7782f6f4cdd 100644 --- a/contrib/seekable_format/zstdseek_decompress.c +++ b/contrib/seekable_format/zstdseek_decompress.c @@ -435,7 +435,10 @@ static size_t ZSTD_seekable_loadSeekTable(ZSTD_seekable* zs) U32 const offset = SEEKABLE_BUFF_SIZE - pos; U32 const toRead = MIN(remaining, SEEKABLE_BUFF_SIZE - offset); memmove(zs->inBuff, zs->inBuff + pos, offset); /* move any data we haven't read yet */ - CHECK_IO(src.read(src.opaque, zs->inBuff+offset, toRead)); + if (src.read(src.opaque, zs->inBuff+offset, toRead) < 0) { + free(entries); + return ERROR(seekableIO); + } remaining -= toRead; pos = 0; } @@ -454,6 +457,7 @@ static size_t ZSTD_seekable_loadSeekTable(ZSTD_seekable* zs) entries[numFrames].cOffset = cOffset; entries[numFrames].dOffset = dOffset; + free(zs->seekTable.entries); zs->seekTable.entries = entries; zs->seekTable.tableLen = numFrames; zs->seekTable.checksumFlag = checksumFlag;