Commit message: Prevent format string vulnerability in `tf.strings.as_string`.

The `printf` format specifier only allows `#`, `0`, `-`, `+` and space as flag characters. Others are interpreted as width/precision/length modifier or conversion specifiers. If a character does not fit into any of these sets `printf` just displays it.

Also add a test suite for `tf.strings.as_string`. Also fix the issue where the flag character was used only if width was specified.

PiperOrigin-RevId: 332553548
Change-Id: Ie57cf2a7c14d1a36097642794c14329db669bbba
CVE: CVE-2020-15203
CWE: CWE-703
Vulnerability description: In Tensorflow before versions 1.15.4, 2.0.3, 2.1.2, 2.2.1 and 2.3.1, by controlling the `fill` argument of tf.strings.as_string, a malicious attacker is able to trigger a format string vulnerability due to the way the internal format use in a `printf` call is constructed. This may result in segmentation fault. The issue is patched in commit 33be22c65d86256e6826666662e40dbdfe70ee83, and is released in TensorFlow versions 1.15.4, 2.0.3, 2.1.2, 2.2.1, or 2.3.1.

Project: tensorflow

  Status Init(DataType input_type, const string& fill = "", int width = -1,
              int precision = -1, bool scientific = false,
              bool shortest = false) {
    TF_CHECK_OK(NodeDefBuilder("op", "AsString")
                    .Input(FakeInput(input_type))
                    .Attr("fill", fill)
                    .Attr("precision", precision)
                    .Attr("scientific", scientific)
                    .Attr("shortest", shortest)
                    .Attr("width", width)
                    .Finalize(node_def()));
    return InitOp();
  }

Project: tensorflow

TEST_F(AsStringGraphTest, FloatPrecisionOnly) {
  TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/2));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(&expected, {"-42.00", "0.00", "3.14", "42.00"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FillWithChar1) {
  TF_ASSERT_OK(Init(DT_INT64, /*fill=*/"-", /*width=*/4));

  AddInputFromArray<int64>(TensorShape({3}), {-42, 0, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(&expected, {"-42 ", "0   ", "42  "});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, NoScientificForNonFloat) {
  Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
                  /*scientific=*/true);
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(absl::StrContains(
      s.error_message(),
      "scientific and shortest format not supported for datatype"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, Float_5_2_Format) {
  TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/5, /*precision=*/2));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(&expected, {"-42.00", " 0.00", " 3.14", "42.00"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

  void Compute(OpKernelContext* context) override {
    const Tensor* input_tensor;
    OP_REQUIRES_OK(context, context->input("input", &input_tensor));
    const DataType& dtype = input_tensor->dtype();

    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context,
                   context->allocate_output("output", input_tensor->shape(),
                                            &output_tensor));
    auto output_flat = output_tensor->flat<tstring>();

#define ENCODE_TYPE(type, T, enc_str)                                     \
  case (type): {                                                          \
    const auto& input_flat = input_tensor->flat<T>();                     \
    for (int i = 0; i < input_flat.size(); ++i) {                         \
      output_flat(i) = strings::Printf((enc_str.c_str()), input_flat(i)); \
    }                                                                     \
  } break

    switch (dtype) {
      ENCODE_TYPE(DT_INT32, int32, format_);
      ENCODE_TYPE(DT_INT64, int64, format_);
      ENCODE_TYPE(DT_FLOAT, float, format_);
      ENCODE_TYPE(DT_DOUBLE, double, format_);
      ENCODE_TYPE(DT_INT8, int8, format_);
      ENCODE_TYPE(DT_INT16, int16, format_);
      case (DT_BOOL): {
        const auto& input_flat = input_tensor->flat<bool>();
        for (int i = 0; i < input_flat.size(); ++i) {
          output_flat(i) = (input_flat(i)) ? "true" : "false";
        }
      } break;
      case (DT_COMPLEX64): {
        const auto& input_flat = input_tensor->flat<complex64>();
        for (int i = 0; i < input_flat.size(); ++i) {
          output_flat(i) = strings::Printf(
              format_.c_str(), input_flat(i).real(), input_flat(i).imag());
        }
      } break;
      case (DT_COMPLEX128): {
        const auto& input_flat = input_tensor->flat<complex128>();
        for (int i = 0; i < input_flat.size(); ++i) {
          output_flat(i) = strings::Printf(
              format_.c_str(), input_flat(i).real(), input_flat(i).imag());
        }
      } break;
      default:
        bool can_encode_type = false;
        OP_REQUIRES(context, can_encode_type,
                    errors::InvalidArgument("Cannot encode input of type ",
                                            DataTypeString(dtype)));
    }

#undef ENCODE_TYPE
  }

Project: tensorflow

TEST_F(AsStringGraphTest, FloatDefault) {
  TF_ASSERT_OK(Init(DT_FLOAT));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(
      &expected, {"-42.000000", "0.000000", "3.141590", "42.000000"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, Int8) {
  TF_ASSERT_OK(Init(DT_INT8));

  AddInputFromArray<int8>(TensorShape({3}), {-42, 0, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(&expected, {"-42", "0", "42"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FloatScientific) {
  TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
                    /*scientific=*/true));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(&expected, {"-4.200000e+01", "0.000000e+00",
                                        "3.141590e+00", "4.200000e+01"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FillWithSpace) {
  TF_ASSERT_OK(Init(DT_INT64, /*fill=*/" ", /*width=*/4));

  AddInputFromArray<int64>(TensorShape({3}), {-42, 0, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(&expected, {" -42", "   0", "  42"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FloatWidthOnly) {
  TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/5));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(
      &expected, {"-42.000000", "0.000000", "3.141590", "42.000000"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, LongFill) {
  Status s = Init(DT_INT32, /*fill=*/"asdf");
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(absl::StrContains(s.error_message(),
                                "Fill string must be one or fewer characters"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, NoShortestForNonFloat) {
  Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
                  /*scientific=*/false, /*shortest=*/true);
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(absl::StrContains(
      s.error_message(),
      "scientific and shortest format not supported for datatype"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FillWithChar3) {
  Status s = Init(DT_INT32, /*fill=*/"s");
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(
      absl::StrContains(s.error_message(), "Fill argument not supported"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FillWithChar4) {
  Status s = Init(DT_INT32, /*fill=*/"n");
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(
      absl::StrContains(s.error_message(), "Fill argument not supported"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, NoPrecisionForNonFloat) {
  Status s = Init(DT_INT32, /*fill=*/"", /*width=*/-1, /*precision=*/5);
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(absl::StrContains(s.error_message(),
                                "precision not supported for datatype"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, OnlyOneOfScientificAndShortest) {
  Status s = Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
                  /*scientific=*/true, /*shortest=*/true);
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(
      absl::StrContains(s.error_message(),
                        "Cannot select both scientific and shortest notation"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, Bool) {
  TF_ASSERT_OK(Init(DT_BOOL));

  AddInputFromArray<bool>(TensorShape({2}), {true, false});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({2}));
  test::FillValues<tstring>(&expected, {"true", "false"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FloatShortest) {
  TF_ASSERT_OK(Init(DT_FLOAT, /*fill=*/"", /*width=*/-1, /*precision=*/-1,
                    /*scientific=*/false, /*shortest=*/true));

  AddInputFromArray<float>(TensorShape({4}), {-42, 0, 3.14159, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({4}));
  test::FillValues<tstring>(&expected, {"-42", "0", "3.14159", "42"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

  explicit AsStringOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
    int32 precision;
    bool scientific;
    bool shortest;
    int32 width;
    string fill_string;
    DataType dtype;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("precision", &precision));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("scientific", &scientific));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("shortest", &shortest));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("width", &width));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("fill", &fill_string));
    switch (dtype) {
      case DT_FLOAT:
      case DT_DOUBLE:
      case DT_COMPLEX64:
      case DT_COMPLEX128:
        break;
      default:
        OP_REQUIRES(ctx, !(scientific || shortest),
                    errors::InvalidArgument("scientific and shortest format "
                                            "not supported for datatype ",
                                            DataTypeString(dtype)));
        OP_REQUIRES(ctx, precision < 0,
                    errors::InvalidArgument("precision not supported "
                                            "for datatype ",
                                            DataTypeString(dtype)));
    }
    OP_REQUIRES(
        ctx, fill_string.size() <= 1,
        errors::InvalidArgument("Fill string must be one or fewer characters"));
    OP_REQUIRES(ctx, !(scientific && shortest),
                errors::InvalidArgument(
                    "Cannot select both scientific and shortest notation"));

    format_ = "%";
    if (!fill_string.empty()) {
      switch (fill_string[0]) {
        case ' ':
        case '+':
        case '-':
        case '0':
        case '#':
          strings::Appendf(&format_, "%s", fill_string.c_str());
          break;
        default:
          bool fill_not_supported = true;
          OP_REQUIRES(ctx, !fill_not_supported,
                      errors::InvalidArgument("Fill argument not supported: \"",
                                              fill_string, "\""));
      }
    }
    if (width > -1) {
      strings::Appendf(&format_, "%d", width);
    }
    if (precision > -1) {
      strings::Appendf(&format_, ".%d", precision);
    }
    switch (dtype) {
      case DT_INT8:
      case DT_INT16:
      case DT_INT32:
        strings::Appendf(&format_, "d");
        break;
      case DT_INT64:
        strings::Appendf(&format_, "lld");
        break;
      case DT_FLOAT:
      case DT_DOUBLE:
      case DT_COMPLEX64:
      case DT_COMPLEX128:
        if (shortest) {
          strings::Appendf(&format_, "g");
        } else if (scientific) {
          strings::Appendf(&format_, "e");
        } else {
          strings::Appendf(&format_, "f");
        }
        break;
      case DT_BOOL:
        break;
      default:
        bool type_not_supported = true;
        OP_REQUIRES(ctx, !type_not_supported,
                    errors::InvalidArgument("Type not supported: ",
                                            DataTypeString(dtype)));
    }

    if (dtype == DT_COMPLEX64 || dtype == DT_COMPLEX128) {
      format_ = strings::Printf("(%s,%s)", format_.c_str(), format_.c_str());
    }
  }

Project: tensorflow

TEST_F(AsStringGraphTest, String) {
  Status s = Init(DT_STRING);
  ASSERT_EQ(error::INVALID_ARGUMENT, s.code());
  ASSERT_TRUE(absl::StrContains(
      s.error_message(),
      "Value for attr 'T' of string is not in the list of allowed values"));
}

Project: tensorflow

TEST_F(AsStringGraphTest, Complex) {
  TF_ASSERT_OK(Init(DT_COMPLEX64, /*fill=*/"", /*width=*/5, /*precision=*/2));

  AddInputFromArray<complex64>(TensorShape({3}), {{-4, 2}, {0}, {3.14159, -1}});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(
      &expected, {"(-4.00, 2.00)", "( 0.00, 0.00)", "( 3.14,-1.00)"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, Int64) {
  TF_ASSERT_OK(Init(DT_INT64));

  AddInputFromArray<int64>(TensorShape({3}), {-42, 0, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(&expected, {"-42", "0", "42"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}

Project: tensorflow

TEST_F(AsStringGraphTest, FillWithZero) {
  TF_ASSERT_OK(Init(DT_INT64, /*fill=*/"0", /*width=*/4));

  AddInputFromArray<int64>(TensorShape({3}), {-42, 0, 42});
  TF_ASSERT_OK(RunOpKernel());
  Tensor expected(allocator(), DT_STRING, TensorShape({3}));
  test::FillValues<tstring>(&expected, {"-042", "0000", "0042"});
  test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
}