Prevent use after free.

mihaimaruseac · tensorflower-gardener · commit 5ecec9c6fbdb · 2021-07-30T19:27:22.000-07:00
A very old version of the code used `result` as a simple pointer to a resource. Two years later, the pointer got changed to a `unique_ptr` but author forgot to remove the call to `Unref`. Three years after that, we finally uncover the UAF.

PiperOrigin-RevId: 387924872
Change-Id: I70fb6f199164de49fac20c168132a07b84903f9b
diff --git a/tensorflow/core/kernels/boosted_trees/resource_ops.cc b/tensorflow/core/kernels/boosted_trees/resource_ops.cc
@@ -53,6 +53,7 @@ class BoostedTreesCreateEnsembleOp : public OpKernel {
     if (!result->InitFromSerialized(
             tree_ensemble_serialized_t->scalar<tstring>()(), stamp_token)) {
       result->Unref();
+      result.release();  // Needed due to the `->Unref` above, to prevent UAF
       OP_REQUIRES(
           context, false,
           errors::InvalidArgument("Unable to parse tree ensemble proto."));
