Fix TOB-EXECUTORCH-39, -42: validate tensor dimensions in XNNPACK compiler

Validate that dims array is non-null and num_dims matches the actual array
size in defineTensor to prevent heap buffer overflows. Change
flatbufferDimsToVector to return Result<> with null-check and per-dimension
validation against a 16M limit to prevent unbounded memory allocation from
malicious dimension values.

Authored-with: Claude

Author: Github Executorch

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -285,9 +285,22 @@ Error defineTensor(
   }
 
   ET_CHECK_OR_RETURN_ERROR(
-      tensor_value != nullptr,
-      Internal,
-      "Deserialized Tensor is Null, this should never happen");
+      tensor_value != nullptr, InvalidProgram, "Deserialized tensor is null");
+
+  ET_CHECK_OR_RETURN_ERROR(
+      tensor_value->num_dims() == 0 || tensor_value->dims() != nullptr,
+      InvalidProgram,
+      "Tensor dims is null but num_dims is %u",
+      tensor_value->num_dims());
+
+  if (tensor_value->dims() != nullptr) {
+    ET_CHECK_OR_RETURN_ERROR(
+        tensor_value->num_dims() == tensor_value->dims()->size(),
+        InvalidProgram,
+        "Tensor num_dims %u does not match dims array size %u",
+        tensor_value->num_dims(),
+        tensor_value->dims()->size());
+  }
 
   // Get tensor dims, here we need to use a vector in order
   // to properly convert the uint32_t* to size_t*
@@ -966,10 +979,14 @@ Error defineStaticTransposeNode(
   auto graph_node = node->xnode_union_as_XNNStaticTranspose();
 
   // Get tensor dims, we need to convert the uint32_t* to size_t*
+  ET_CHECK_OR_RETURN_ERROR(
+      graph_node->perm() != nullptr,
+      InvalidProgram,
+      "StaticTranspose: perm is null");
   std::vector<size_t> dims_data = flatbufferDimsToVector(graph_node->perm());
   xnn_status status = xnn_define_static_transpose(
       subgraph_ptr,
-      graph_node->num_dims(),
+      dims_data.size(),
       dims_data.data(),
       remapped_ids.at(graph_node->input_id()),
       remapped_ids.at(graph_node->output_id()),
@@ -1031,6 +1048,11 @@ Error defineStaticConstantPadNode(
   const fb_xnnpack::XNNStaticConstantPad* graph_node =
       node->xnode_union_as_XNNStaticConstantPad();
 
+  ET_CHECK_OR_RETURN_ERROR(
+      graph_node->pre_paddings() != nullptr &&
+          graph_node->post_paddings() != nullptr,
+      InvalidProgram,
+      "StaticConstantPad: pre_paddings or post_paddings is null");
   std::vector<size_t> pre_paddings_dims =
       flatbufferDimsToVector(graph_node->pre_paddings());
   std::vector<size_t> post_paddings_dims =
@@ -1111,11 +1133,15 @@ Error defineStaticReshapeNode(
   auto graph_node = node->xnode_union_as_XNNStaticReshape();
 
   // Get tensor dims, we need to convert the uint32_t* to size_t*
+  ET_CHECK_OR_RETURN_ERROR(
+      graph_node->new_shape() != nullptr,
+      InvalidProgram,
+      "StaticReshape: new_shape is null");
   std::vector<size_t> dims_data =
       flatbufferDimsToVector(graph_node->new_shape());
   xnn_status status = xnn_define_static_reshape(
       subgraph_ptr,
-      graph_node->num_dims(),
+      dims_data.size(),
       dims_data.data(),
       remapped_ids.at(graph_node->input_id()),
       remapped_ids.at(graph_node->output_id()),
@@ -1406,12 +1432,23 @@ Error defineStaticSliceNode(
 
   auto graph_node = node->xnode_union_as_XNNStaticSlice();
 
+  ET_CHECK_OR_RETURN_ERROR(
+      graph_node->offsets() != nullptr && graph_node->sizes() != nullptr,
+      InvalidProgram,
+      "StaticSlice: offsets or sizes is null");
   std::vector<size_t> offsets = flatbufferDimsToVector(graph_node->offsets());
   std::vector<size_t> sizes = flatbufferDimsToVector(graph_node->sizes());
 
+  ET_CHECK_OR_RETURN_ERROR(
+      offsets.size() == sizes.size(),
+      InvalidProgram,
+      "StaticSlice: offsets size %zu does not match sizes size %zu",
+      offsets.size(),
+      sizes.size());
+
   xnn_status status = xnn_define_static_slice(
       subgraph_ptr,
-      graph_node->num_dims(),
+      offsets.size(),
       offsets.data(),
       sizes.data(),
       remapped_ids.at(graph_node->input_id()),