nfextract: Check number of available entries

3 files changed, 17 insertions, 11 deletions

diff --git a/bin/nfextract.c b/bin/nfextract.c
index d7fc9ae..11fe608 100644
--- a/bin/nfextract.c
+++ b/bin/nfextract.c
@@ -63,13 +63,14 @@ static void extract_each(const char *storage_dir, const char *filename) {
     sprintf(fullpath, "%s/%s", storage_dir, filename);
 
     debug("Extracting storage file: %s", fullpath);
-    if (nfl_extract_worker(fullpath, &trunk) < 0)
+    int entries = nfl_extract_worker(fullpath, &trunk);
+    if (entries < 0)
         return;
 
     free(fullpath);
 
     char output[1024];
-    for (int i = 0; i < trunk.header->n_entries; ++i) {
+    for (int i = 0; i < entries; ++i) {
         nfl_format_output(output, &trunk.store[i]);
         puts((char *)output);
     }
diff --git a/lib/common.c b/lib/common.c
index 4380f27..3cbf19b 100644
--- a/lib/common.c
+++ b/lib/common.c
@@ -102,8 +102,7 @@ void nfl_cal_trunk(uint32_t total_size, uint32_t *trunk_cnt,
 
 void nfl_cal_entries(uint32_t trunk_size, uint32_t *entries_cnt) {
     assert(entries_cnt);
-    *entries_cnt =
-        (trunk_size - sizeof(nfl_header_t)) / sizeof(nfl_entry_t);
+    *entries_cnt = (trunk_size - sizeof(nfl_header_t)) / sizeof(nfl_entry_t);
 }
 
 void nfl_format_output(char *output, nfl_entry_t *entry) {
diff --git a/lib/extract.c b/lib/extract.c
index e5fa681..11381d7 100644
--- a/lib/extract.c
+++ b/lib/extract.c
@@ -29,9 +29,10 @@ static int nfl_verify_header(nfl_header_t *header) {
 }
 
 static int nfl_extract_default(FILE *f, nfl_state_t *state) {
-    fread(state->store, state->header->n_entries, sizeof(nfl_entry_t), f);
+    int entries =
+        fread(state->store, sizeof(nfl_entry_t), state->header->n_entries, f);
     WARN_RETURN(ferror(f), "%s", strerror(errno));
-    return 0;
+    return entries;
 }
 
 static int nfl_extract_zstd(FILE *f, nfl_state_t *state) {
@@ -40,6 +41,9 @@ static int nfl_extract_zstd(FILE *f, nfl_state_t *state) {
                  expected_decom_size =
                      state->header->n_entries * sizeof(nfl_entry_t);
 
+    // It's possible that data or header is not written due to broken commit
+    WARN_RETURN(compressed_size <= 0, "%s", "zstd: no data in this trunk");
+
     ERR(!(buf = malloc(compressed_size)), "zstd: cannot malloc");
     fread(buf, compressed_size, 1, f);
     WARN_RETURN(ferror(f), "%s", strerror(errno));
@@ -61,7 +65,7 @@ static int nfl_extract_zstd(FILE *f, nfl_state_t *state) {
     }
 
     free(buf);
-    return 0;
+    return actual_decom_size / sizeof(nfl_entry_t);
 }
 
 static int nfl_extract_lz4(FILE *f, nfl_state_t *state) {
@@ -94,16 +98,18 @@ int nfl_extract_worker(const char *filename, nfl_state_t *state) {
         "extract malloc store");
     switch (h->compression_opt) {
     case COMPRESS_NONE:
-        debug("Extract worker #%u: extract without compression\n", h->id)
-            nfl_extract_default(f, state);
+        debug("Extract worker #%u: extract without compression\n", h->id);
+        ret = nfl_extract_default(f, state);
         break;
     case COMPRESS_LZ4:
         debug("Extract worker #%u: extract with compression algorithm: lz4",
-              h->id) nfl_extract_lz4(f, state);
+              h->id);
+        ret = nfl_extract_lz4(f, state);
         break;
     case COMPRESS_ZSTD:
         debug("Extract worker #%u: extract with compression algorithm: zstd",
-              h->id) nfl_extract_zstd(f, state);
+              h->id);
+        ret = nfl_extract_zstd(f, state);
         break;
     // Must not reach here ...
     default: