Prevent a crash due to heap OOB write in grappler.

mihaimaruseac · tensorflower-gardener · commit 97282c6d0d34 · 2021-11-08T05:54:06.000-08:00
PiperOrigin-RevId: 408318417
Change-Id: If095feb8c001e3a8ac4a85b7387b81e8309df47d
diff --git a/tensorflow/core/grappler/costs/graph_properties.cc b/tensorflow/core/grappler/costs/graph_properties.cc
@@ -1134,7 +1134,12 @@ class SymbolicShapeRefiner {
         GetUnknownOutputShape(node, output_port);
     InferenceContext* ctx = GetContext(node);
     if (ctx == nullptr) {
-      return errors::InvalidArgument("Missing context");
+      return errors::InvalidArgument("SetUnknownShape: Missing context");
+    }
+    if (output_port < 0 || output_port >= ctx->num_outputs()) {
+      return errors::InvalidArgument(
+          "SetUnknownShape: output_port must be in [0, ", ctx->num_outputs(),
+          ") but was ", output_port);
     }
     ctx->set_output(output_port, shape);
     return Status::OK();
