Harden language model readers against malformed input

src/lm/ngram_model_trie.c

@@ -532,7 +532,7 @@ ngram_model_trie_read_dmp(ps_config_t * config,
         goto error_out;
     }
     if (strncmp(file_header, dmp_hdr, k) != 0) {
-        E_ERROR("Wrong header %s: %s is not a dump file\n", dmp_hdr);
+        E_ERROR("Wrong header: %s is not a dump file\n", file_name);
         goto error_out;
     }
     ckd_free(file_header);
@@ -675,9 +675,11 @@ ngram_model_trie_read_dmp(ps_config_t * config,
 
     /* Sentinel unigram and bigrams read before */
     ckd_free(unigram_next);
+    unigram_next = NULL;
 
     /* read ascii word strings */
-    read_word_str(base, fp, do_swap);
+    if (read_word_str(base, fp, do_swap) != 0)
+        goto error_out;
 
     fclose_comp(fp, is_pipe);
     return base;

src/lm/ngrams_raw.c

@@ -186,7 +186,9 @@ ngrams_raw_read_arpa(lineiter_t ** li, logmath_t * lmath, uint32 * counts,
         return NULL;
     } else {
         *li = lineiter_next(*li);
-	if (strcmp((*li)->buf, "\\end\\") != 0) {
+	if (*li == NULL) {
+	    E_WARN("Finished reading ARPA file without finding end mark\n");
+	} else if (strcmp((*li)->buf, "\\end\\") != 0) {
     	    E_WARN
         	("Finished reading ARPA file. Expecting end mark but found '%s'\n",
         	 (*li)->buf);
@@ -196,20 +198,44 @@ ngrams_raw_read_arpa(lineiter_t ** li, logmath_t * lmath, uint32 * counts,
     return raw_ngrams;
 }
 
-static void
+static int
 read_dmp_weight_array(FILE * fp, logmath_t * lmath, uint8 do_swap,
                       int32 counts, ngram_raw_t * raw_ngrams,
                       int weight_idx)
 {
     int32 i, k;
+    long filepos, remaining;
     dmp_weight_t *tmp_weight_arr;
 
-    fread(&k, sizeof(k), 1, fp);
+    if (fread(&k, sizeof(k), 1, fp) != 1) {
+        E_ERROR("Failed to read weight array size\n");
+        return -1;
+    }
     if (do_swap)
         SWAP_INT32(&k);
+    if (k <= 0) {
+        E_ERROR("Invalid weight array size %d\n", k);
+        return -1;
+    }
+    /* Reject a size that cannot fit in the remaining file, before
+     * allocating, so a corrupt count does not request a huge block.
+     * Skipped for non-seekable input, where the read below is the guard. */
+    filepos = ftell(fp);
+    if (filepos >= 0 && fseek(fp, 0, SEEK_END) == 0) {
+        remaining = ftell(fp) - filepos;
+        fseek(fp, filepos, SEEK_SET);
+        if (remaining < 0 || k > remaining / (long) sizeof(*tmp_weight_arr)) {
+            E_ERROR("Weight array size %d exceeds remaining file data\n", k);
+            return -1;
+        }
+    }
     tmp_weight_arr =
         (dmp_weight_t *) ckd_calloc(k, sizeof(*tmp_weight_arr));
-    fread(tmp_weight_arr, sizeof(*tmp_weight_arr), k, fp);
+    if (fread(tmp_weight_arr, sizeof(*tmp_weight_arr), k, fp) != (size_t) k) {
+        E_ERROR("Failed to read weight array\n");
+        ckd_free(tmp_weight_arr);
+        return -1;
+    }
     for (i = 0; i < k; i++) {
         if (do_swap)
             SWAP_INT32(&tmp_weight_arr[i].l);
@@ -219,15 +245,22 @@ read_dmp_weight_array(FILE * fp, logmath_t * lmath, uint8 do_swap,
     }
     /* replace indexes with real probs in raw bigrams */
     for (i = 0; i < counts; i++) {
+        int32 idx = (weight_idx == 0)
+            ? (int32) raw_ngrams[i].prob
+            : (int32) raw_ngrams[i].backoff;
+        if (idx < 0 || idx >= k) {
+            E_ERROR("Weight index %d out of range [0, %d)\n", idx, k);
+            ckd_free(tmp_weight_arr);
+            return -1;
+        }
 	if (weight_idx == 0) {
-	    raw_ngrams[i].prob =
-                tmp_weight_arr[(int) raw_ngrams[i].prob].f;
+	    raw_ngrams[i].prob = tmp_weight_arr[idx].f;
         } else {
-	    raw_ngrams[i].backoff =
-                tmp_weight_arr[(int) raw_ngrams[i].backoff].f;
+	    raw_ngrams[i].backoff = tmp_weight_arr[idx].f;
         }
     }
     ckd_free(tmp_weight_arr);
+    return 0;
 }
 
 #define BIGRAM_SEGMENT_SIZE 9
@@ -315,17 +348,29 @@ ngrams_raw_read_dmp(FILE * fp, logmath_t * lmath, uint32 * counts,
     }
 
     /* read prob2 */
-    read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[1],
-                          raw_ngrams[0], 0);
+    if (read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[1],
+                              raw_ngrams[0], 0) < 0) {
+        ckd_free(bigrams_next);
+        ngrams_raw_free(raw_ngrams, counts, order);
+        return NULL;
+    }
     /* read bo2 */
     if (order > 2) {
         int32 k;
         int32 *tseg_base;
-        read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[1],
-                              raw_ngrams[0], 1);
+        if (read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[1],
+                                  raw_ngrams[0], 1) < 0) {
+            ckd_free(bigrams_next);
+            ngrams_raw_free(raw_ngrams, counts, order);
+            return NULL;
+        }
         /* read prob3 */
-        read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[2],
-                              raw_ngrams[1], 0);
+        if (read_dmp_weight_array(fp, lmath, do_swap, (int32) counts[2],
+                                  raw_ngrams[1], 0) < 0) {
+            ckd_free(bigrams_next);
+            ngrams_raw_free(raw_ngrams, counts, order);
+            return NULL;
+        }
         /* Read tseg_base size and tseg_base to fill trigram's first words */
         fread(&k, sizeof(k), 1, fp);
         if (do_swap)
@@ -377,7 +422,11 @@ ngrams_raw_free(ngram_raw_t ** raw_ngrams, uint32 * counts, int order)
     uint32 num;
     int order_it;
 
+    if (raw_ngrams == NULL)
+        return;
     for (order_it = 0; order_it < order - 1; order_it++) {
+        if (raw_ngrams[order_it] == NULL)
+            continue;
         for (num = 0; num < counts[order_it + 1]; num++) {
             ckd_free(raw_ngrams[order_it][num].words);
         }

test/unit/CMakeLists.txt

@@ -24,6 +24,7 @@ set(TESTS
   test_keyphrase
   test_lattice
   test_lm_convert
+  test_lm_reader_hardening
   test_ngram_model_read
   test_log_shifted
   test_log_int8

test/unit/test_lm_reader_hardening.c

@@ -0,0 +1,222 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+#include <pocketsphinx.h>
+#include <pocketsphinx/model.h>
+#include <pocketsphinx/logmath.h>
+
+#include "test_macros.h"
+
+static char *
+write_temp(const void *data, size_t len)
+{
+    char template[] = "/tmp/ps_lmhardXXXXXX";
+    int fd;
+    FILE *fp;
+
+    fd = mkstemp(template);
+    TEST_ASSERT(fd >= 0);
+    fp = fdopen(fd, "wb");
+    TEST_ASSERT(fp != NULL);
+    if (len > 0)
+        TEST_EQUAL(len, fwrite(data, 1, len, fp));
+    fclose(fp);
+    return strdup(template);
+}
+
+static void
+put32(FILE *fp, int32 v)
+{
+    fwrite(&v, sizeof(v), 1, fp);
+}
+
+static void
+put16(FILE *fp, uint16 v)
+{
+    fwrite(&v, sizeof(v), 1, fp);
+}
+
+static void
+put_dmp_header(FILE *fp)
+{
+    put32(fp, 17);
+    fwrite("Darpa Trigram LM", 1, 16, fp);
+    fputc('\0', fp);
+}
+
+static char *
+build_dmp(int unigrams, int bigrams, int trigrams,
+          const uint16 *bigram_words, int32 weight_count,
+          int truncate_word_str)
+{
+    char template[] = "/tmp/ps_lmhardXXXXXX";
+    int fd = mkstemp(template);
+    FILE *fp;
+    int j;
+
+    TEST_ASSERT(fd >= 0);
+    fp = fdopen(fd, "wb");
+    TEST_ASSERT(fp != NULL);
+
+    put_dmp_header(fp);
+    put32(fp, 0);              /* filename length */
+    put32(fp, unigrams);       /* version > 0 means ucount, no extended header */
+    put32(fp, bigrams);        /* bcount */
+    put32(fp, trigrams);       /* tcount */
+
+    /* unigram table: ucount + 1 entries, 16 bytes each */
+    for (j = 0; j <= unigrams; j++) {
+        put32(fp, 0);          /* mapping id */
+        put32(fp, 0);          /* prob weight */
+        put32(fp, 0);          /* backoff weight */
+        put32(fp, 0);          /* bigram pointer */
+    }
+
+    if (bigrams > 0) {
+        /* bigram table: bcount + 1 entries, 8 bytes each */
+        for (j = 0; j <= bigrams; j++) {
+            uint16 prob_idx = 0;
+            if (bigram_words != NULL && j < bigrams)
+                prob_idx = bigram_words[j];
+            put16(fp, 0);              /* word id */
+            put16(fp, prob_idx);       /* prob index into weight array */
+            put16(fp, 0);              /* backoff index */
+            put16(fp, 0);              /* bigram next */
+        }
+        /* prob2 weight array size */
+        put32(fp, weight_count);
+    }
+
+    if (truncate_word_str) {
+        /* Declare a large word-string block but supply none of it, so
+         * read_word_str hits a short read and fails. */
+        put32(fp, 4096);
+    }
+
+    fclose(fp);
+    return strdup(template);
+}
+
+/* Defect class 1: an invalid (non-positive) DMP weight-array count in
+ * read_dmp_weight_array (ngrams_raw.c) must be rejected. */
+static void
+test_dmp_weight_array_count(ps_config_t *config, logmath_t *lmath)
+{
+    ngram_model_t *lm;
+    char *path = build_dmp(1, 1, 0, NULL, -1, 0);
+
+    lm = ngram_model_read(config, path, NGRAM_AUTO, lmath);
+    TEST_ASSERT(lm == NULL);
+    unlink(path);
+    free(path);
+}
+
+/* Defect class 1: a DMP weight-array count larger than the remaining file
+ * must be rejected in read_dmp_weight_array (ngrams_raw.c) before it is
+ * used to size an allocation. */
+static void
+test_dmp_weight_array_oversized(ps_config_t *config, logmath_t *lmath)
+{
+    ngram_model_t *lm;
+    char *path = build_dmp(1, 1, 0, NULL, 1000000, 0);
+
+    lm = ngram_model_read(config, path, NGRAM_AUTO, lmath);
+    TEST_ASSERT(lm == NULL);
+    unlink(path);
+    free(path);
+}
+
+/* Defect class 5: an unchecked read_word_str failure (ngram_model_trie.c)
+ * used to leave a partially built model; the reader must now fail cleanly. */
+static void
+test_dmp_truncated_word_str(ps_config_t *config, logmath_t *lmath)
+{
+    ngram_model_t *lm;
+    char *path = build_dmp(1, 0, 0, NULL, 0, 1);
+
+    lm = ngram_model_read(config, path, NGRAM_AUTO, lmath);
+    TEST_ASSERT(lm == NULL);
+    unlink(path);
+    free(path);
+}
+
+/* Defect class 2: an ARPA file that declares a trigram section it never
+ * provides makes ngrams_raw_read_arpa free a partially built array;
+ * ngrams_raw_free (ngrams_raw.c) must tolerate the NULL sub-array. */
+static void
+test_arpa_missing_section(ps_config_t *config, logmath_t *lmath)
+{
+    static const char arpa[] =
+        "\\data\\\n"
+        "ngram 1=2\n"
+        "ngram 2=1\n"
+        "ngram 3=1\n"
+        "\n"
+        "\\1-grams:\n"
+        "-1.0 <s> -0.5\n"
+        "-1.0 </s> -0.5\n"
+        "\n"
+        "\\2-grams:\n"
+        "-0.5 <s> </s> -0.3\n";
+    ngram_model_t *lm;
+    char *path = write_temp(arpa, sizeof(arpa) - 1);
+
+    lm = ngram_model_read(config, path, NGRAM_AUTO, lmath);
+    TEST_ASSERT(lm == NULL);
+    unlink(path);
+    free(path);
+}
+
+/* Defect class 3: a complete ARPA model with no end-mark used to
+ * dereference a NULL line iterator in ngrams_raw_read_arpa (ngrams_raw.c);
+ * the reader must now load it and only warn. */
+static void
+test_arpa_no_end_mark(ps_config_t *config, logmath_t *lmath)
+{
+    static const char arpa[] =
+        "\\data\\\n"
+        "ngram 1=2\n"
+        "ngram 2=1\n"
+        "\n"
+        "\\1-grams:\n"
+        "-1.0 <s> -0.5\n"
+        "-1.0 </s> -0.5\n"
+        "\n"
+        "\\2-grams:\n"
+        "-0.5 <s> </s>\n";
+    ngram_model_t *lm;
+    char *path = write_temp(arpa, sizeof(arpa) - 1);
+
+    lm = ngram_model_read(config, path, NGRAM_AUTO, lmath);
+    TEST_ASSERT(lm != NULL);
+    ngram_model_free(lm);
+    unlink(path);
+    free(path);
+}
+
+int
+main(int argc, char *argv[])
+{
+    ps_config_t *config;
+    logmath_t *lmath;
+
+    (void)argc;
+    (void)argv;
+
+    err_set_loglevel(ERR_FATAL);
+    TEST_ASSERT(config = ps_config_parse_json(NULL, "{}"));
+    TEST_ASSERT(lmath = logmath_init(1.0001, 0, 0));
+
+    test_dmp_weight_array_count(config, lmath);
+    test_dmp_weight_array_oversized(config, lmath);
+    test_dmp_truncated_word_str(config, lmath);
+    test_arpa_missing_section(config, lmath);
+    test_arpa_no_end_mark(config, lmath);
+
+    logmath_free(lmath);
+    ps_config_free(config);
+
+    return 0;
+}