TensorFlow vulnerable to null-dereference in `mlir::tfg::TFOp::nameAttr`

Impact

When mlir::tfg::TFOp::nameAttr receives null type list attributes, it crashes.

StatusOr<unsigned> GraphDefImporter::ArgNumType(const NamedAttrList &attrs,
                                                const OpDef::ArgDef &arg_def,
                                                SmallVectorImpl<Type> &types) {
  // Check whether a type list attribute is specified.
  if (!arg_def.type_list_attr().empty()) {
    if (auto v = attrs.get(arg_def.type_list_attr()).dyn_cast<ArrayAttr>()) {
      for (Attribute attr : v) {
        if (auto dtype = attr.dyn_cast<TypeAttr>()) {
          types.push_back(UnrankedTensorType::get(dtype.getValue()));
        } else {
          return InvalidArgument("Expected '", arg_def.type_list_attr(),
                                 "' to be a list of types");
        }
      }
      return v.size();
    }
    return NotFound("Type attr not found: ", arg_def.type_list_attr());
  }

  unsigned num = 1;
  // Check whether a number attribute is specified.
  if (!arg_def.number_attr().empty()) {
    if (auto v = attrs.get(arg_def.number_attr()).dyn_cast<IntegerAttr>()) {
      num = v.getValue().getZExtValue();
    } else {
      return NotFound("Type attr not found: ", arg_def.number_attr());
    }
  }

  // Check for a type or type attribute.
  Type dtype;
  if (arg_def.type() != DataType::DT_INVALID) {
    TF_RETURN_IF_ERROR(ConvertDataType(arg_def.type(), b_, &dtype));
  } else if (arg_def.type_attr().empty()) {
    return InvalidArgument("Arg '", arg_def.name(),
                           "' has invalid type and no type attribute");
  } else {
    if (auto v = attrs.get(arg_def.type_attr()).dyn_cast<TypeAttr>()) {
      dtype = v.getValue();
    } else {
      return NotFound("Type attr not found: ", arg_def.type_attr());
    }
  }
  types.append(num, UnrankedTensorType::get(dtype));
  return num;
}

Patches

We have patched the issue in GitHub commits 3a754740d5414e362512ee981eefba41561a63a6 and .