Skip to content
Permalink
Browse files Browse the repository at this point in the history
[tfg] Fix out-of-bounds access due to mismatched integer type sizes i…
…n ValueMap::Manager::GetValueOrCreatePlaceholder

PiperOrigin-RevId: 497419488
  • Loading branch information
tlongeri authored and tensorflower-gardener committed Dec 23, 2022
1 parent 69c9be3 commit 760322a
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 10 deletions.
30 changes: 20 additions & 10 deletions tensorflow/core/ir/importexport/functiondef_import.cc
Expand Up @@ -86,11 +86,11 @@ class ValueMapManager {
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(':');
Expand All @@ -105,8 +105,16 @@ class ValueMapManager {
// 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);
}
}
Expand Down Expand Up @@ -171,8 +179,9 @@ Status ImportNodes(ValueMapManager value_manager,
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.
Expand Down Expand Up @@ -470,8 +479,9 @@ Status ImportGenericFunction(
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);
Expand All @@ -485,8 +495,8 @@ Status ImportGenericFunction(
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");
Expand Down
@@ -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
}
}
}
}
}

0 comments on commit 760322a

Please sign in to comment.