Harden backup import and restore coverage

1 files changed, 392 insertions, 188 deletions

diff --git a/src/cli/import.c b/src/cli/import.c
index 3ddb9e9..f4688ab 100644
--- a/src/cli/import.c
+++ b/src/cli/import.c
@@ -9,8 +9,11 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <limits.h>
 
 #define IMPORT_MAX_LINE TINY_CLI_LINE_MAX
+#define IMPORT_MAX_RECORDS 100000U
+#define IMPORT_MAX_DECODED_BYTES (UINT64_C(4) * 1024U * 1024U * 1024U)
 
 typedef struct import_record_list
 {
@@ -64,27 +67,264 @@ import_get_string(const char *json, int len, const char *path, char *out, size_t
   return decoded > 0 && (size_t)decoded < out_size;
 }
 
+static int
+import_hex_digit(char value)
+{
+  if (value >= '0' && value <= '9')
+    {
+      return value - '0';
+    }
+  if (value >= 'a' && value <= 'f')
+    {
+      return value - 'a' + 10;
+    }
+  if (value >= 'A' && value <= 'F')
+    {
+      return value - 'A' + 10;
+    }
+  return -1;
+}
+
+static bool
+import_decode_json_string(const char *source, size_t length, char *output, size_t *output_length)
+{
+  size_t input = 0, used = 0;
+  while (input < length)
+    {
+      uint32_t codepoint;
+      unsigned count;
+      if (source[input] != '\\')
+        {
+          output[used++] = source[input++];
+          continue;
+        }
+      if (++input >= length)
+        {
+          return false;
+        }
+      switch (source[input++])
+        {
+        case '"':
+          output[used++] = '"';
+          continue;
+        case '\\':
+          output[used++] = '\\';
+          continue;
+        case '/':
+          output[used++] = '/';
+          continue;
+        case 'b':
+          output[used++] = '\b';
+          continue;
+        case 'f':
+          output[used++] = '\f';
+          continue;
+        case 'n':
+          output[used++] = '\n';
+          continue;
+        case 'r':
+          output[used++] = '\r';
+          continue;
+        case 't':
+          output[used++] = '\t';
+          continue;
+        case 'u':
+          break;
+        default:
+          return false;
+        }
+      if (input + 4 > length)
+        {
+          return false;
+        }
+      codepoint = 0;
+      for (count = 0; count < 4; ++count)
+        {
+          int digit = import_hex_digit(source[input++]);
+          if (digit < 0)
+            {
+              return false;
+            }
+          codepoint = (codepoint << 4) | (uint32_t)digit;
+        }
+      if (codepoint >= 0xd800 && codepoint <= 0xdbff)
+        {
+          uint32_t low = 0;
+          if (input + 6 > length || source[input] != '\\' || source[input + 1] != 'u')
+            {
+              return false;
+            }
+          input += 2;
+          for (count = 0; count < 4; ++count)
+            {
+              int digit = import_hex_digit(source[input++]);
+              if (digit < 0)
+                {
+                  return false;
+                }
+              low = (low << 4) | (uint32_t)digit;
+            }
+          if (low < 0xdc00 || low > 0xdfff)
+            {
+              return false;
+            }
+          codepoint = UINT32_C(0x10000) + ((codepoint - 0xd800U) << 10) + low - 0xdc00U;
+        }
+      else if (codepoint >= 0xdc00 && codepoint <= 0xdfff)
+        {
+          return false;
+        }
+      if (codepoint <= 0x7f)
+        {
+          output[used++] = (char)codepoint;
+        }
+      else if (codepoint <= 0x7ff)
+        {
+          output[used++] = (char)(0xc0U | (codepoint >> 6));
+          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
+        }
+      else if (codepoint <= 0xffff)
+        {
+          output[used++] = (char)(0xe0U | (codepoint >> 12));
+          output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU));
+          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
+        }
+      else
+        {
+          output[used++] = (char)(0xf0U | (codepoint >> 18));
+          output[used++] = (char)(0x80U | ((codepoint >> 12) & 0x3fU));
+          output[used++] = (char)(0x80U | ((codepoint >> 6) & 0x3fU));
+          output[used++] = (char)(0x80U | (codepoint & 0x3fU));
+        }
+    }
+  output[used] = '\0';
+  *output_length = used;
+  return true;
+}
+
+static bool
+import_get_allocated_string(const char *json,
+                            int len,
+                            const char *path,
+                            size_t max_length,
+                            bool allow_empty,
+                            char **out,
+                            size_t *decoded_length)
+{
+  const char *token = NULL;
+  int token_length = 0;
+  size_t decoded;
+  char *value;
+  *out = NULL;
+  if (decoded_length)
+    {
+      *decoded_length = 0;
+    }
+  if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_STRING || token_length < 2
+      || token_length == INT_MAX)
+    {
+      return false;
+    }
+  value = malloc((size_t)token_length);
+  if (!value)
+    {
+      return false;
+    }
+  if (!import_decode_json_string(token + 1, (size_t)token_length - 2U, value, &decoded)
+      || (!allow_empty && decoded == 0) || decoded > max_length || strlen(value) != decoded
+      || !tiny_utf8_valid(value, decoded))
+    {
+      free(value);
+      return false;
+    }
+  *out = value;
+  if (decoded_length)
+    {
+      *decoded_length = decoded;
+    }
+  return true;
+}
+
+static bool
+import_get_integer_token(const char *json,
+                         int len,
+                         const char *path,
+                         bool allow_negative,
+                         uint64_t maximum,
+                         uint64_t *magnitude,
+                         bool *negative)
+{
+  const char *token = NULL;
+  int token_length = 0;
+  int index = 0;
+  uint64_t value = 0;
+  if (mjson_find(json, len, path, &token, &token_length) != MJSON_TOK_NUMBER || token_length <= 0)
+    {
+      return false;
+    }
+  *negative = token[0] == '-';
+  if (*negative)
+    {
+      if (!allow_negative || ++index == token_length)
+        {
+          return false;
+        }
+    }
+  if (token[index] == '0' && index + 1 != token_length)
+    {
+      return false;
+    }
+  for (; index < token_length; ++index)
+    {
+      unsigned digit;
+      if (token[index] < '0' || token[index] > '9')
+        {
+          return false;
+        }
+      digit = (unsigned)(token[index] - '0');
+      if (value > (maximum - digit) / 10U)
+        {
+          return false;
+        }
+      value = value * 10U + digit;
+    }
+  *magnitude = value;
+  return true;
+}
+
 static bool
 import_get_i64(const char *json, int len, const char *path, int64_t *value)
 {
-  double number;
-  if (mjson_get_number(json, len, path, &number) != 1)
+  uint64_t magnitude;
+  bool negative;
+  if (!import_get_integer_token(
+          json, len, path, true, (uint64_t)INT64_MAX + 1U, &magnitude, &negative))
     {
       return false;
     }
-  *value = (int64_t)number;
+  if (negative)
+    {
+      *value = magnitude == (uint64_t)INT64_MAX + 1U ? INT64_MIN : -(int64_t)magnitude;
+    }
+  else
+    {
+      if (magnitude > INT64_MAX)
+        {
+          return false;
+        }
+      *value = (int64_t)magnitude;
+    }
   return true;
 }
 
 static bool
 import_get_u64(const char *json, int len, const char *path, uint64_t *value)
 {
-  double number;
-  if (mjson_get_number(json, len, path, &number) != 1 || number < 0)
+  bool negative;
+  if (!import_get_integer_token(json, len, path, false, UINT64_MAX, value, &negative))
     {
       return false;
     }
-  *value = (uint64_t)number;
   return true;
 }
 
@@ -126,23 +366,15 @@ import_parse_string_array(const char *json,
                               &value_type))
          != 0)
     {
-      char value[512];
-      int decoded;
+      char *value = NULL;
+      size_t decoded = 0;
       if (value_type != MJSON_TOK_STRING || index >= max_count
-          || (decoded = mjson_get_string(
-                  array_json + value_offset, value_length, "$", value, sizeof(value)))
-                 <= 0
-          || (size_t)decoded > max_length || decoded >= (int)sizeof(value))
-        {
-          break;
-        }
-      (*output)[index] = malloc((size_t)decoded + 1);
-      if (!(*output)[index])
+          || !import_get_allocated_string(
+              array_json + value_offset, value_length, "$", max_length, false, &value, &decoded))
         {
           break;
         }
-      memcpy((*output)[index], value, (size_t)decoded);
-      (*output)[index][decoded] = '\0';
+      (*output)[index] = value;
       index++;
     }
   *count = index;
@@ -183,31 +415,37 @@ import_parse_relations(const char *json, int len, const char *path, tiny_metadat
                               &value_type))
          != 0)
     {
-      char type_name[32];
-      char message_id[32];
-      int decoded;
+      char *type_name = NULL;
+      char *message_id = NULL;
+      size_t decoded = 0;
       tiny_relation_t *relation;
       if (value_type != MJSON_TOK_OBJECT || index >= TINY_RELATION_COUNT_MAX
-          || mjson_get_string(
-                 array_json + value_offset, value_length, "$.type", type_name, sizeof(type_name))
-                 <= 0
-          || (decoded = mjson_get_string(array_json + value_offset,
-                                         value_length,
-                                         "$.message_id",
-                                         message_id,
-                                         sizeof(message_id)))
-                 <= 0)
+          || !import_get_allocated_string(
+              array_json + value_offset, value_length, "$.type", 31, false, &type_name, NULL)
+          || !import_get_allocated_string(array_json + value_offset,
+                                          value_length,
+                                          "$.message_id",
+                                          TINY_MESSAGE_ID_LENGTH,
+                                          false,
+                                          &message_id,
+                                          &decoded))
         {
+          free(type_name);
+          free(message_id);
           return false;
         }
       relation = &metadata->relations[index];
       if (!tiny_relation_type_parse(type_name, &relation->type)
           || !tiny_message_id_valid(message_id))
         {
+          free(type_name);
+          free(message_id);
           return false;
         }
       memcpy(relation->message_id, message_id, (size_t)decoded);
       relation->message_id[decoded] = '\0';
+      free(type_name);
+      free(message_id);
       index++;
     }
   metadata->relation_count = index;
@@ -219,16 +457,30 @@ import_parse_metadata(const char *json, int len, tiny_message_t *message)
 {
   char tier_name[32] = { 0 };
   char kind_name[32] = { 0 };
-  char tool_call_id[TINY_TOOL_CALL_ID_MAX + 1] = { 0 };
-  char summary[TINY_SUMMARY_MAX + 1] = { 0 };
-  double importance_number = 0;
+  char *tool_call_id = NULL;
+  char *summary = NULL;
+  uint64_t importance_number = 0;
+  bool negative = false;
   bool has_tool_call_id, has_summary, has_importance;
-  has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0
-                     && import_get_string(
-                         json, len, "$.metadata.tool_call_id", tool_call_id, sizeof(tool_call_id));
-  has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0
-                && import_get_string(json, len, "$.metadata.summary", summary, sizeof(summary));
-  has_importance = mjson_get_number(json, len, "$.metadata.importance", &importance_number) == 1;
+  has_tool_call_id = mjson_find(json, len, "$.metadata.tool_call_id", NULL, NULL) != 0;
+  has_summary = mjson_find(json, len, "$.metadata.summary", NULL, NULL) != 0;
+  has_importance = mjson_find(json, len, "$.metadata.importance", NULL, NULL) != 0;
+  if ((has_tool_call_id
+       && !import_get_allocated_string(
+           json, len, "$.metadata.tool_call_id", TINY_TOOL_CALL_ID_MAX, false, &tool_call_id, NULL))
+      || (has_summary
+          && !import_get_allocated_string(
+              json, len, "$.metadata.summary", TINY_SUMMARY_MAX, false, &summary, NULL))
+      || (has_importance
+          && !import_get_integer_token(
+              json, len, "$.metadata.importance", false, 100, &importance_number, &negative)))
+    {
+      free(tool_call_id);
+      free(summary);
+      return false;
+    }
+  message->metadata.tool_call_id = tool_call_id;
+  message->metadata.summary = summary;
   if (!import_get_string(json, len, "$.metadata.memory_tier", tier_name, sizeof(tier_name))
       || !import_get_string(json, len, "$.metadata.memory_kind", kind_name, sizeof(kind_name))
       || !import_parse_string_array(json,
@@ -254,30 +506,8 @@ import_parse_metadata(const char *json, int len, tiny_message_t *message)
     {
       return false;
     }
-  if (has_tool_call_id)
-    {
-      message->metadata.tool_call_id = malloc(strlen(tool_call_id) + 1);
-      if (!message->metadata.tool_call_id)
-        {
-          return false;
-        }
-      memcpy(message->metadata.tool_call_id, tool_call_id, strlen(tool_call_id) + 1);
-    }
-  if (has_summary)
-    {
-      message->metadata.summary = malloc(strlen(summary) + 1);
-      if (!message->metadata.summary)
-        {
-          return false;
-        }
-      memcpy(message->metadata.summary, summary, strlen(summary) + 1);
-    }
   if (has_importance)
     {
-      if (importance_number < 0 || importance_number > 100)
-        {
-          return false;
-        }
       message->metadata.importance = (uint8_t)importance_number;
     }
   return true;
@@ -287,14 +517,14 @@ static bool
 import_manifest_valid(const char *line, int len)
 {
   char type_name[64];
-  double version = 0;
+  uint64_t version = 0;
   if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
       || strcmp(type_name, "tiny-convmem-manifest") != 0
-      || mjson_get_number(line, len, "$.format_version", &version) != 1)
+      || !import_get_u64(line, len, "$.format_version", &version))
     {
       return false;
     }
-  return (int)version == 2;
+  return version == 2;
 }
 
 static bool
@@ -318,18 +548,21 @@ static bool
 import_parse_message_record(const char *line, int len, tiny_import_record_t *record)
 {
   char role_name[32];
-  char session_id[TINY_SESSION_ID_MAX + 1];
-  char content_buffer[4096];
+  char *session_id = NULL;
+  char *content = NULL;
+  size_t content_length = 0;
   tiny_message_init(&record->message);
   record->deleted = false;
   if (!import_get_string(
           line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id))
-      || !import_get_string(line, len, "$.session_id", session_id, sizeof(session_id))
       || !import_get_string(line, len, "$.role", role_name, sizeof(role_name))
+      || !import_get_allocated_string(
+          line, len, "$.session_id", TINY_SESSION_ID_MAX, false, &session_id, NULL)
       || !import_get_i64(line, len, "$.timestamp", &record->message.timestamp_ms)
       || !import_get_i64(line, len, "$.updated_at", &record->message.updated_at_ms)
       || !import_get_u64(line, len, "$.revision", &record->message.revision))
     {
+      free(session_id);
       return false;
     }
   if (!tiny_message_id_valid(record->message.message_id)
@@ -337,57 +570,18 @@ import_parse_message_record(const char *line, int len, tiny_import_record_t *rec
       || record->message.updated_at_ms < 0 || record->message.revision == 0
       || record->message.revision > INT32_MAX)
     {
+      free(session_id);
       return false;
     }
-  if (mjson_find(line, len, "$.content", NULL, NULL) == 0)
-    {
-      return false;
-    }
-  if (mjson_get_string(line, len, "$.content", content_buffer, sizeof(content_buffer)) > 0)
-    {
-      size_t length = strlen(content_buffer);
-      if (length == 0 || length > TINY_CONTENT_MAX)
-        {
-          return false;
-        }
-      record->message.content = malloc(length + 1);
-      if (!record->message.content)
-        {
-          return false;
-        }
-      memcpy(record->message.content, content_buffer, length + 1);
-    }
-  else
-    {
-      const char *content = NULL;
-      int content_len = 0;
-      if (mjson_find(line, len, "$.content", &content, &content_len) == 0 || !content
-          || content_len < 2)
-        {
-          return false;
-        }
-      if ((size_t)(content_len - 2) == 0 || (size_t)(content_len - 2) > TINY_CONTENT_MAX)
-        {
-          return false;
-        }
-      record->message.content = malloc((size_t)content_len - 1);
-      if (!record->message.content)
-        {
-          return false;
-        }
-      memcpy(record->message.content, content + 1, (size_t)content_len - 2);
-      record->message.content[(size_t)content_len - 2] = '\0';
-    }
-  if (!tiny_utf8_valid(record->message.content, strlen(record->message.content)))
-    {
-      return false;
-    }
-  record->message.session_id = malloc(strlen(session_id) + 1);
-  if (!record->message.session_id)
+  if (!import_get_allocated_string(
+          line, len, "$.content", TINY_CONTENT_MAX, false, &content, &content_length))
     {
+      free(session_id);
       return false;
     }
-  memcpy(record->message.session_id, session_id, strlen(session_id) + 1);
+  (void)content_length;
+  record->message.content = content;
+  record->message.session_id = session_id;
   if (!import_parse_metadata(line, len, &record->message)
       || !tiny_metadata_valid(&record->message.metadata))
     {
@@ -398,60 +592,99 @@ import_parse_message_record(const char *line, int len, tiny_import_record_t *rec
   return true;
 }
 
+static uint64_t
+import_record_decoded_bytes(const tiny_import_record_t *record)
+{
+  uint64_t total = 0;
+  size_t index;
+  const tiny_message_t *message = &record->message;
+  if (record->deleted)
+    {
+      return TINY_MESSAGE_ID_LENGTH;
+    }
+  total = strlen(message->session_id) + strlen(message->content) + TINY_MESSAGE_ID_LENGTH;
+  total += message->metadata.tool_call_id ? strlen(message->metadata.tool_call_id) : 0;
+  total += message->metadata.summary ? strlen(message->metadata.summary) : 0;
+  for (index = 0; index < message->metadata.tag_count; ++index)
+    {
+      total += strlen(message->metadata.tags[index]);
+    }
+  total += message->metadata.relation_count * TINY_MESSAGE_ID_LENGTH;
+  total += message->metadata.source_count * TINY_MESSAGE_ID_LENGTH;
+  return total;
+}
+
 static bool
-import_parse_jsonl(const char *jsonl, import_record_list_t *records)
+import_parse_stream(FILE *file, import_record_list_t *records)
 {
-  const char *line_start = jsonl;
-  const char *line_end;
+  char *line = malloc((size_t)IMPORT_MAX_LINE + 2U);
   bool manifest_seen = false;
-  while (*line_start)
+  uint64_t decoded_total = 0;
+  if (!line)
     {
-      int length;
+      return false;
+    }
+  while (fgets(line, (int)IMPORT_MAX_LINE + 2, file) != NULL)
+    {
+      size_t length = strlen(line);
       tiny_import_record_t record;
-      line_end = strchr(line_start, '\n');
-      if (line_end == NULL)
+      uint64_t decoded;
+      if (length == 0 || length > IMPORT_MAX_LINE
+          || (length == IMPORT_MAX_LINE + 1U && line[length - 1] != '\n')
+          || (line[length - 1] != '\n' && !feof(file)))
         {
-          line_end = line_start + strlen(line_start);
+          free(line);
+          return false;
         }
-      length = (int)(line_end - line_start);
-      if (length > 0)
+      if (length > 0 && line[length - 1] == '\n')
         {
-          if (!manifest_seen)
-            {
-              if (!import_manifest_valid(line_start, length))
-                {
-                  return false;
-                }
-              manifest_seen = true;
-            }
-          else if (import_parse_del_record(line_start, length, &record))
-            {
-              if (!import_record_list_push(records))
-                {
-                  return false;
-                }
-              records->items[records->count++] = record;
-            }
-          else if (import_parse_message_record(line_start, length, &record))
-            {
-              if (!import_record_list_push(records))
-                {
-                  tiny_message_free(&record.message);
-                  return false;
-                }
-              records->items[records->count++] = record;
-            }
-          else
+          line[--length] = '\0';
+        }
+      if (length > 0 && line[length - 1] == '\r')
+        {
+          line[--length] = '\0';
+        }
+      if (length == 0)
+        {
+          continue;
+        }
+      if (!manifest_seen)
+        {
+          if (!import_manifest_valid(line, (int)length))
             {
+              free(line);
               return false;
             }
+          manifest_seen = true;
+          continue;
         }
-      if (*line_end == '\0')
+      if (records->count >= IMPORT_MAX_RECORDS)
         {
-          break;
+          free(line);
+          return false;
         }
-      line_start = line_end + 1;
+      if (!import_parse_del_record(line, (int)length, &record)
+          && !import_parse_message_record(line, (int)length, &record))
+        {
+          free(line);
+          return false;
+        }
+      decoded = import_record_decoded_bytes(&record);
+      if (decoded > IMPORT_MAX_DECODED_BYTES - decoded_total || !import_record_list_push(records))
+        {
+          tiny_message_free(&record.message);
+          free(line);
+          return false;
+        }
+      decoded_total += decoded;
+      records->items[records->count++] = record;
     }
+  if (ferror(file))
+    {
+      free(line);
+      return false;
+    }
+  free(line);
   return manifest_seen;
 }
 
@@ -462,9 +695,7 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
   size_t imported = 0, rejected = 0;
   tiny_store_status_t status;
   tiny_import_policy_t policy = TINY_IMPORT_OVERWRITE;
-  char *jsonl = NULL;
   FILE *file = NULL;
-  long file_size;
   int result = 1;
 
   if (!options || !options->import_path)
@@ -496,49 +727,22 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
       return 1;
     }
 
-  if (fseek(file, 0, SEEK_END) != 0 || (file_size = ftell(file)) < 0
-      || fseek(file, 0, SEEK_SET) != 0)
-    {
-      fprintf(stderr, "import: unable to seek '%s'\n", options->import_path);
-      fclose(file);
-      return 1;
-    }
-
-  if ((size_t)file_size > IMPORT_MAX_LINE * 1024ULL * 1024ULL)
+  if (!import_parse_stream(file, &records))
     {
-      fprintf(stderr, "import: file too large\n");
-      fclose(file);
-      return 1;
-    }
-
-  jsonl = malloc((size_t)file_size + 1);
-  if (!jsonl)
-    {
-      fprintf(stderr, "import: out of memory\n");
-      fclose(file);
-      return 1;
-    }
-
-  if (fread(jsonl, 1, (size_t)file_size, file) != (size_t)file_size)
-    {
-      fprintf(stderr, "import: unable to read '%s'\n", options->import_path);
-      free(jsonl);
+      fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path);
       fclose(file);
+      import_record_list_free(&records);
       return 1;
     }
-  jsonl[file_size] = '\0';
-  fclose(file);
-
-  if (!import_parse_jsonl(jsonl, &records))
+  if (fclose(file) != 0)
     {
-      fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path);
-      free(jsonl);
+      fprintf(stderr, "import: unable to close '%s'\n", options->import_path);
+      import_record_list_free(&records);
       return 1;
     }
 
   status = tiny_store_import(store, records.items, records.count, policy, &imported, &rejected);
   import_record_list_free(&records);
-  free(jsonl);
 
   if (status == TINY_STORE_OK)
     {