Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent null dereference read in SpecializeType()
For some adversarial protos, the attribute for a key might not exist.

PiperOrigin-RevId: 408382090
Change-Id: Ie7eabe532c9ff280fce5dce1f6cdb93c76c2e040
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Nov 8, 2021
1 parent 258112d commit 8a513ce
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion tensorflow/core/framework/full_type_util.cc
Expand Up @@ -22,6 +22,7 @@ limitations under the License.
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"

namespace tensorflow {

Expand Down Expand Up @@ -102,7 +103,11 @@ StatusOr<FullTypeDef> SpecializeType(const AttrSlice& attrs,
auto* arg = t->mutable_args(i);
if (arg->type_id() == TFT_VAR) {
const auto* attr = attrs.Find(arg->s());
DCHECK(attr != nullptr);
if (attr == nullptr) {
return Status(
error::INVALID_ARGUMENT,
absl::StrCat("Could not find an attribute for key ", arg->s()));
}
if (attr->value_case() == AttrValue::kList) {
const auto& attr_list = attr->list();
arg->set_type_id(TFT_PRODUCT);
Expand Down

0 comments on commit 8a513ce

Please sign in to comment.