From 478bb78ca41b29ba4a6c0d51304f5f3402f3166a Mon Sep 17 00:00:00 2001
From: Ryan Kuester <kuester@bdti.com>
Date: Tue, 17 Dec 2024 18:10:49 -0600
Subject: [PATCH] feat(compression): check metadata schema version (#3022)

Add a check in `micro_allocator.cc` to verify the compression
metadata schema version. If the schema version in the metadata is
greater than the expected version, log a schema version mismatch
error and return a `nullptr`. This prevents potential issues
arising from using a newer, unsupported schema version.

BUG=part of #2636
---
 tensorflow/lite/micro/micro_allocator.cc | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/lite/micro/micro_allocator.cc b/tensorflow/lite/micro/micro_allocator.cc
index 08203285f4c..eee24eb195d 100644
--- a/tensorflow/lite/micro/micro_allocator.cc
+++ b/tensorflow/lite/micro/micro_allocator.cc
@@ -409,6 +409,14 @@ const tflite::micro::compression::Metadata* GetCompressionMetadata(
         MicroPrintf("Compression: verification failure");
         return nullptr;
       } else {
+        tflite::micro::compression::MetadataT schema;
+        if (compression_metadata->schema_version() > schema.schema_version) {
+          MicroPrintf("Compression: schema version mismatch (using %d got %d)",
+                      schema.schema_version,
+                      compression_metadata->schema_version());
+          return nullptr;
+        }
+
         return compression_metadata;
       }
     }