CVE-2023-25671: Out-of-bounds Write (CWE-787)
value: TensorFlow is an open source platform for machine learning. There is out-of-bounds access due to mismatched integer type sizes. A fix is included in TensorFlow version 2.12.0 and version 2.11.1.

CWE description: The product writes data past the end, or before the beginning, of the intended buffer.
CVSS v3 base score: 7.5
Published: 2023-03-25T00:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Other
Files changed: tensorflow/core/ir/importexport/functiondef_import.cc, tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_named_edge_index.pbtxt
Commit date: 2022-12-23 15:00:24 -0800
Commit message: [tfg] Fix out-of-bounds access due to mismatched integer type sizes in ValueMap::Manager::GetValueOrCreatePlaceholder

PiperOrigin-RevId: 497419488

Diff: diff --git a/tensorflow/core/ir/importexport/functiondef_import.cc b/tensorflow/core/ir/importexport/functiondef_import.cc
index 96cdb5767b5e6..d1d54b2478833 100644
--- a/tensorflow/core/ir/importexport/functiondef_import.cc
+++ b/tensorflow/core/ir/importexport/functiondef_import.cc
@@ -84,15 +84,15 @@ class ValueMapManager {
     base_operation.push_back(op->getResult(1));
     base_operation.push_back(op->getResult(0));
     return ::tensorflow::OkStatus();
   }
 
-  Value GetValueOrCreatePlaceholder(StringRef full_name) {
+  tensorflow::StatusOr<Value> GetValueOrCreatePlaceholder(StringRef full_name) {
     StringRef node_name;
     StringRef output_name = "";
     bool is_control_dep = full_name[0] == '^';
-    int output_num = 0;
+    size_t output_num = 0;
     if (is_control_dep) full_name = full_name.drop_front();
     {
       size_t colon_sep = full_name.find_first_of(':');
       if (colon_sep == StringRef::npos) {
         node_name = full_name;
@@ -103,12 +103,20 @@ class ValueMapManager {
       colon_sep = output_name.find_last_of(':');
       if (colon_sep != StringRef::npos) {
         // NOLINTNEXTLINE: type matching the API taking a reference.
         unsigned long long value;
         if (!llvm::getAsUnsignedInteger(output_name.drop_front(colon_sep + 1),
-                                        10, value))
-          output_num = value;
+                                        10, value)) {
+          if (LLVM_LIKELY(
+                  value <=
+                  std::numeric_limits<llvm::SmallVectorSizeType<Value>>::max() -
+                      1))
+            output_num = value;
+          else
+            return InvalidArgument("Output index ", value,
+                                   " is invalid (too large)");
+        }
         output_name = output_name.take_front(colon_sep);
       }
     }
 
     llvm::StringMap<SmallVector<Value, 1>>& op_info = values_map_[node_name];
@@ -169,12 +177,13 @@ Status ImportNodes(ValueMapManager value_manager,
     OperationState state(unknown_loc, absl::StrCat("tfg.", node.op()));
     // Fetch the inputs, creating placeholder if an input hasn't been visited.
     for (const std::string& input : node.input()) {
       if (input.empty())
         return InvalidArgument("Node '", node.name(), "' has an empty input");
-      state.operands.push_back(
-          value_manager.GetValueOrCreatePlaceholder(input));
+      TF_ASSIGN_OR_RETURN(Value value,
+                          value_manager.GetValueOrCreatePlaceholder(input));
+      state.operands.push_back(value);
     }
     // Retrieve the entry in the nodes_map for this node and infer the result
     // count from what was inferred during the first traversal above.
     state.types.push_back(placeholder_ty);
     state.types.push_back(control_ty);
@@ -468,12 +477,13 @@ Status ImportGenericFunction(
     }
     if (ret_val.second.empty()) {
       return InvalidArgument("Function '", func.signature().name(),
                              "' has empty result name");
     }
-    ret_vals[position->second] =
-        value_manager.GetValueOrCreatePlaceholder(ret_val.second);
+    TF_ASSIGN_OR_RETURN(
+        ret_vals[position->second],
+        value_manager.GetValueOrCreatePlaceholder(ret_val.second));
   }
   for (const auto& ret_val : func.control_ret()) {
     auto position = control_output_to_position.find(ret_val.first);
     if (position == control_output_to_position.end()) {
       return InvalidArgument(
@@ -483,12 +493,12 @@ Status ImportGenericFunction(
     }
     if (ret_val.second.empty()) {
       return InvalidArgument("Function '", func.signature().name(),
                              "' has empty control result name");
     }
-    Value result = value_manager.GetValueOrCreatePlaceholder(
-        (Twine("^") + ret_val.second).str());
+    TF_ASSIGN_OR_RETURN(Value result, value_manager.GetValueOrCreatePlaceholder(
+                                          (Twine("^") + ret_val.second).str()));
     if (!result.getType().isa<ControlType>())
       return InvalidArgument("failed to map returned value ", ret_val.second,
                              ", isn't a control output");
     ret_vals[func.ret_size() + position->second] = result;
   }
diff --git a/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_named_edge_index.pbtxt b/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_named_edge_index.pbtxt
new file mode 100644
index 0000000000000..7275ac6f1e6ea
--- /dev/null
+++ b/tensorflow/core/ir/importexport/tests/graphdef_to_mlir/invalid_generic_function_named_edge_index.pbtxt
@@ -0,0 +1,52 @@
+# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
+
+# CHECK: Output index 18446744073709551615 is invalid (too large)
+
+library {
+  function {
+    signature {
+      name: "foo"
+      attr {
+        name: "T"
+        type: "type"
+      }
+    }
+    node_def {
+      name: "two"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value {
+          type: DT_INT64
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT64
+            tensor_shape {}
+            int64_val: 2
+          }
+        }
+      }
+    }
+    node_def {
+      name: "a"
+      op: "Cast"
+      input: "two:output:18446744073709551615"
+      attr {
+        key: "DstT"
+        value {
+          placeholder: "T"
+        }
+      }
+      attr {
+        key: "SrcT"
+        value {
+          type: DT_INT64
+        }
+      }
+    }
+  }
+}

Vulnerable code:   Value GetValueOrCreatePlaceholder(StringRef full_name) {
    int output_num = 0;
                                        10, value))
          output_num = value;
      state.operands.push_back(
          value_manager.GetValueOrCreatePlaceholder(input));
    ret_vals[position->second] =
        value_manager.GetValueOrCreatePlaceholder(ret_val.second);
    Value result = value_manager.GetValueOrCreatePlaceholder(
        (Twine("^") + ret_val.second).str());
Fixed code:   tensorflow::StatusOr<Value> GetValueOrCreatePlaceholder(StringRef full_name) {
    size_t output_num = 0;
                                        10, value)) {
          if (LLVM_LIKELY(
                  value <=
                  std::numeric_limits<llvm::SmallVectorSizeType<Value>>::max() -
                      1))
            output_num = value;
          else
            return InvalidArgument("Output index ", value,
                                   " is invalid (too large)");
        }
      TF_ASSIGN_OR_RETURN(Value value,
                          value_manager.GetValueOrCreatePlaceholder(input));
      state.operands.push_back(value);
    TF_ASSIGN_OR_RETURN(
        ret_vals[position->second],
        value_manager.GetValueOrCreatePlaceholder(ret_val.second));
    TF_ASSIGN_OR_RETURN(Value result, value_manager.GetValueOrCreatePlaceholder(
                                          (Twine("^") + ret_val.second).str()));
# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s

# CHECK: Output index 18446744073709551615 is invalid (too large)

library {
  function {
    signature {
      name: "foo"
      attr {
        name: "T"
        type: "type"
      }
    }
    node_def {
      name: "two"
      op: "Const"
      attr {
        key: "dtype"
        value {
          type: DT_INT64
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_INT64
            tensor_shape {}
            int64_val: 2
          }
        }
      }
    }
    node_def {
      name: "a"
      op: "Cast"
      input: "two:output:18446744073709551615"
      attr {
        key: "DstT"
        value {
          placeholder: "T"
        }
      }
      attr {
        key: "SrcT"
        value {
          type: DT_INT64
        }
      }
    }
  }
}