Fix tf.raw_ops.AudioSummaryV2 vulnerability with invalid sample_rate.

Check that `sample_rate` has the required shape.

PiperOrigin-RevId: 460819553

tensorflow/core/kernels/summary_audio_op.cc

@@ -49,6 +49,11 @@ class SummaryAudioOp : public OpKernel {
     float sample_rate = sample_rate_attr_;
     if (!has_sample_rate_attr_) {
       const Tensor& sample_rate_tensor = c->input(2);
+      OP_REQUIRES(c,
+                  sample_rate_tensor.IsAligned() &&
+                      sample_rate_tensor.NumElements() == 1,
+                  errors::InvalidArgument(
+                      "sample_rate must be rank-0 or contain a single value"));
       sample_rate = sample_rate_tensor.scalar<float>()();
     }
     OP_REQUIRES(c, sample_rate > 0.0f,

tensorflow/python/summary/summary_test.py

@@ -23,6 +23,7 @@
 from tensorflow.core.framework import summary_pb2
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import meta_graph
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
@@ -183,6 +184,11 @@ def testAudioSummaryWithFamily(self):
         'family/outer/family/inner/audio/{}'.format(i) for i in range(3))
     self.assertEqual(tags, expected)
 
+  def testAudioSummaryWithInvalidSampleRate(self):
+    with self.assertRaises(errors.InvalidArgumentError):
+      invalid_sample_rate = [22000.0, 22000.0]
+      self.evaluate(summary_lib.audio('', [[1.0]], invalid_sample_rate))
+
   @test_util.run_deprecated_v1
   def testTextSummary(self):
     with self.cached_session():