From f760f88b4267d981e13f4b302c437ae800445968 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Fri, 24 Apr 2020 10:49:08 -0700
Subject: [PATCH] Properly handle negative shape dimensions from improper saved
 models.

PiperOrigin-RevId: 308283636
Change-Id: Ib10849425de7d541d8dacfe4d0c709fbac9180b6
---
 tensorflow/cc/saved_model/loader.cc           |  26 ++++++++++++++++++
 .../cc/saved_model/saved_model_bundle_test.cc |  14 ++++++++++
 .../testdata/negative_shape/fuzz_generated    | Bin 0 -> 10820 bytes
 3 files changed, 40 insertions(+)
 create mode 100644 tensorflow/cc/saved_model/testdata/negative_shape/fuzz_generated

diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc
index 3bb4660e449..d193679ec19 100644
--- a/tensorflow/cc/saved_model/loader.cc
+++ b/tensorflow/cc/saved_model/loader.cc
@@ -19,12 +19,16 @@ limitations under the License.
 
 #include "tensorflow/cc/saved_model/constants.h"
 #include "tensorflow/cc/saved_model/reader.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/monitoring/counter.h"
 #include "tensorflow/core/lib/monitoring/sampler.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/protobuf_internal.h"
 #include "tensorflow/core/protobuf/graph_debug_info.pb.h"
 #include "tensorflow/core/protobuf/saver.pb.h"
@@ -65,12 +69,34 @@ uint64 GetLatencyMicroseconds(const uint64 start_microseconds) {
   return end_microseconds - start_microseconds;
 }
 
+// Ensure that constant tensors loaded from the saved model have valid shape.
+// TODO(b/154763635): this is temporary and will be replaced with a better audit
+static Status ValidateSavedTensors(const GraphDef& graph_def) {
+  for (const auto& node : graph_def.node()) {
+    const auto node_iterator = node.attr().find("value");
+    if (node_iterator != node.attr().end()) {
+      AttrValue node_value = node_iterator->second;
+      if (node_value.has_tensor()) {
+        const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
+        if (node_shape.num_elements() < 0) {
+          return errors::FailedPrecondition(
+              "Saved model contains node \"", node.name(), "\" (op \"",
+              node.op(), "\") which initializes from a tensor with ",
+              node_shape.num_elements(), " elements");
+        }
+      }
+    }
+  }
+  return Status::OK();
+}
+
 Status LoadMetaGraphIntoSession(const MetaGraphDef& meta_graph_def,
                                 const SessionOptions& session_options,
                                 std::unique_ptr<Session>* session) {
   Session* session_p = nullptr;
   TF_RETURN_IF_ERROR(NewSession(session_options, &session_p));
   session->reset(session_p);
+  TF_RETURN_IF_ERROR(ValidateSavedTensors(meta_graph_def.graph_def()));
   return (*session)->Create(meta_graph_def.graph_def());
 }
 
diff --git a/tensorflow/cc/saved_model/saved_model_bundle_test.cc b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
index 9fc71552d6f..46f365613f1 100644
--- a/tensorflow/cc/saved_model/saved_model_bundle_test.cc
+++ b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
@@ -40,6 +40,8 @@ constexpr char kTestDataInitOpV2[] =
     "cc/saved_model/testdata/half_plus_two_v2/00000123";
 constexpr char kTestDataV2DebugInfo[] =
     "cc/saved_model/testdata/x_plus_y_v2_debuginfo";
+constexpr char kTestNegativeShapeFuzzGenerated[] =
+    "cc/saved_model/testdata/negative_shape/fuzz_generated";
 
 class LoaderTest : public ::testing::Test {
  protected:
@@ -256,5 +258,17 @@ TEST_F(LoaderTest, SavedModelV2DebugInfo) {
   EXPECT_NE(bundle.debug_info.get(), nullptr);
 }
 
+TEST_F(LoaderTest, NegativeShapeDimension) {
+  SavedModelBundle bundle;
+  RunOptions run_options;
+  SessionOptions session_options;
+
+  const string export_dir = io::JoinPath(testing::TensorFlowSrcRoot(),
+                                         kTestNegativeShapeFuzzGenerated);
+  Status st = LoadSavedModel(session_options, run_options, export_dir,
+                             {kSavedModelTagServe}, &bundle);
+  EXPECT_FALSE(st.ok());
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/cc/saved_model/testdata/negative_shape/fuzz_generated b/tensorflow/cc/saved_model/testdata/negative_shape/fuzz_generated
new file mode 100644
index 0000000000000000000000000000000000000000..5ee5c360ce03537c5cf8ec7d8513eec1ada74ae8
GIT binary patch
literal 10820
zcmd5CU60#XHBK`5IFp&o+?h_(JMDIrrKRg?n~b}Q+Ndnu6e>XKP-y9PwP-9S_O#Yy
z;t+eLnLYpyJb;AI3WQh*A;BYpKfntTzldifBslk8U&p?V)1>VKt=jnG`h1^r&&MVI
zd<<W%k$=~q1BFj)o0LGCG_}X<_Y8k2{=UTD<}TpFVsJKcIvcQV)YiAOe6h4rs8&cB
zCQ$ms_uXehvI>Q<bBG`OOw$O!%BeLtaTpxD@lJx#Nnl3ro7D(rlfpVKb!{tfO#itx
za!9e~c>|iirZvyQ>WS}|1FwJN4xde5e5q-03s&xVLqDK(C<*w@6`N7ntcS=+DR73q
zH*T(hcq6!g8&LU<?F<7qIAesmL!^xnx*7pOuBQ}&1LT1-e&#$e?mc(<M<dT220l3t
zTm9(P93gM+#PJ(tvIE3+{J<Srf$I&!H;ppNv2Hs92if!nwrRU#f6CA=kfE<b;lasZ
zNi}S8J$!+mm!P=seQyL0p!m@0ACYw^aTWN;ny&39xvPP;PYTG#HLb11ctzP2t%>XN
zcQoI~!N2lQduWY)=UWr&cr<VjF8I#aMU}mDY-G0zD?{s;bKZbef5g_BN6s0FLND)&
zdz1;F#HTMx`}~$PU|l@7okQzn5co#vF@HdJVLgQ94sExO@T(CxdI{FTmx4v37Tz-b
z`;<nnxj`j9-GQ3G7=9VzYi>Y9tdI+eLQ)w%xB>|@1O~1jTw^HTzovD{uwoRnYPoqC
zLYP?#E=dvub!3&p;Wr2f?!oHAfz@}OqnSA4w0S43xw#E1!fnJ4Eq3w^4W7bI49KT$
z@cf%X!l=K{;)PGV5<aev!-vnI`~yckL5B%F3?l@39e>&^J?6hjgPXb&I()c7xFgNC
z3mr6GBVVY32UP?fcw=-n^v#lRE3~o=s0mc2G`4Nn;CM`HXq*0kn@E{%LRqXsOXAVM
zZ#E#jO6sGnF{c|q>5+B%n%Ky)*uD+*N6)RX?b!F-fx|c<HK_C~pI-}SLZ-%cqrm)Z
zcyi1_;@I{Z1&Rv~!;cze*avtNj4>#F7o!nbgA$Ko%;(V6QE(z6V$?}``gy4O4!Wp?
zA|@rDL!fz*10Udtg<jnope)cPcq|tdg7Lhm7=aQ|F-}+e-Y~G-p_2gHynqn_8E|Z_
zwm3!<escu_$*H5m-i<xuM*D+tzipIHhDSs1#gP2{W9UQOYPmx<z+hu?cO@d^H`iLW
zbBY@8+<W-LPFp`G1AeUw6^-QdVlJ2ahQ2Dpe!i{gJFsqs+YnNOQGyg9-^oD$nnhx~
ztFRXVqHiYIM1kt{xJcLo1jG>&hV^$ERIFCGbA~%8)hwzqtWr~E`zB*&JGE=#eu+;@
zW`AJ$KKf3DlyKkd+_e~W@m-WH%#z0EP_kNM$FfPZMIu=$6|>P~z|<wpDC&A|A$?wz
z`n=TLMSX6P)yjY6AC2TI8dvb$g%#YJr{FBIm(^*48m{18O2Iv<N+2y-H5<y(Yi|j?
zPIZ6^?xhsmOH+}yqh6NLOo1Omy}yux7jml0(h7FE+tRD=-Zd)d+x!o6-TpkyE`()S
zHFM4OQ=08lmG!ht7u9V4bv4^hX|}IYA)Rq%sOV>@7!Ck}L(>UUz<7S5BEalKggKsv
z#GRt7W45QnVBGh4x?%%~Cz&mt*=I!iO=%aEZVBxojIDgL=b<qJ3vO~}u$91qS*kfP
zGSIvy(a>%K(G;V&lj}t$7-vt@v64~Pt3av4Kg9i&CQyedR7=Qe_O`YoDOl^SRtkcG
zj;6(t7Um-TjRxDZuo9WU(OL4Kb0}4{U5j!#%;M(JG7S=Ea5+I!;)W-8849@|#j(*{
zuapUp94V1v9b-XUBAHf(`hNje6W9_rqAFm8i`mYgfJh3e-7->;>uQrE>WsYvUM`4u
zF9UHbl9jW3G?@i4>hcrgWe)z6hYJZt<2NKx+tHYl(^qMUE-1?d*X_|65rED#U)2cT
zQR=C6s0ptV|Lbeh(he(H9c?zOKJ$&d-bp#_HMrVIryMgL=~$G+>rPDaHlo^(S2S2)
zxY$*RB~#g|P`UbQim-FAHH!)v(&33{Ah0e14bwP^EgdCUULNjK(H~+>A}WB36ynVx
zKf+d1OgId~n*)6cD<}bTkGw(<Qf~U95`iDV3R1`*Gf@aEv?O+2mG0%CKY{$25tD|%
z_2zH~fPo+ecQ+mYTJf9Co{DWz(M-epXF~c$Qu;{2n}ZI;kDw#*8*f+qH)CSr;(rTN
zyJ4R)4DFfW`$M_8T<^ZmH<$NX!SSf&jooM1N*uJD$;cZAEiCSwz;7LT-i=@qB(@)H
zR9=!T0#kHBW@h7<hl%bisdv{ns`4ywS(IWfQ&-T$v+%6XmM)T$y!}Ii9R;hJgsW81
z`O=U!Th1t^n^7^{fE}{Ui)n>|B@%Twg>B_)HCtEdpAxtTDOtJ4W2lKfLo41wc~cws
zBQ~6s_JvA5P2hc-Le6E{SQvH7a0iO4b3}MgoK_EHvjs0@S&o;oj#FY2e9Zc3Zh(>g
z^}fg6JOADg-@|#R^LuMBFay?O!|oGr)*DcNHug?NrtOUU7JC$8EFTB}bVFQBzl1&5
zRKAJm48KytIL=C7`bmgk<6!!TW2QuXV(x(swdt}L9-(`>8v6u3gDoF>sV!;d7<_QD
zGL3s-&|%dKOHWshJ)DOmygW;<FvO0295l)3Zqr#5FceYiZ83Zf4mlV;QHN$BOVX4@
zd^URcxdsi6XWGjZLi177q9|O5Uj2@Y=r_B2T)5Y912{&rYTtJD4lUy&XYclh_Sf3C
zZd+~ZX7A=Lqqlc>%Q0FbYaGA{)D<d7th-EOqM(`NW^#`%xldT8jUCZx@0s<SzaOHV
zwz(zGi9)MDa;d5}E7%<Y6`Uly?5xZl+aYwt+2aB(WO9+FC<nj7@_H6FqBi-dDqcE=
z%5d>iC@Q^KyuXmP#OE2U__!p4^NZ4Iwz`{WHLO!n=xSD1>vgq$TrTA~rF{Wsz!!xW
z1Y+)Lv8hZ3a;2OX$a7u%yEx`wYuI>ae)@=`CgCiH&Qy!E-L!b<T>4pGUh7dU)l1uq
z21{c0ht=MnG&s)*V&jmpH(+ty8DeXpAFp#}lUF;4v7#znZfsbr$kyy7*d-TuRwBMj
zntYZe!{0Qx&Ob=cv|BR<+JsM&vm`1=PBW{LKc_XOLpX9~wRyy57&SEJWJ028a><Ob
zJY9yjcw|#4K9`=>V<#Sdsf9aHbb??yC`rUgA(=GMw6UnnB3WcP{7cy7ii{?vi}-w^
zm6WZPf<8Amz)#?^g8MDI{iRg<V->2prtzLm0sf-l`2gnm5%Xj^Igaa<&XW;+6%(pe
zyqltmHJfnYF8^UgAg9eyQ24!!QcA*%aYf-lMK^oUw)i|x`cuZd5dI*oSx|;aM>LD<
z(T${h)6Id<Kr?{PoG>giv`SHDOsVUqs7qDquR+~1gBSP|2p7;;oL3No`Y^W+LwT4a
zuM3ZtGN#FUm2fOwgybb)O#MT!TLO@L!IK)rWSaTPF1>?&N#+|od+oG|NIzu6>a;`Q
zNjQk$-IZ(FN4dKf@88Nd|ExjvG&~F7lizCO3k{D<6)Xf4&8#V2u`zL{ZEPwt`cC^6
zU*BWvqgo5E<a%WWEZ6thdRdTByprpcks@2){Hq3A!5GgQuwn|&28fdj|F6`(g9jN{
z+s0&-(Z~v(`*>Qe3gu8g`U$L$ooA>H-<+6%XP%Khyd#fvkZnBr63yN^?Ts?Uctfs0
zSsZ8Z35CKLKG$fvYfY(00}lt@ixxX=fhR^#upF2EQ`m^m8XWIUas*c*#Jf<O7#V=c
z)hT{lcXl&l-CP2Gj0J$83G*a={)-&IW!5l9mL?ex2`)HWnrA#ngr~>?{S-&k?0|%O
zh@-znW&IT>`F#)dGZ*dd3{gv5`h7TOn#?I}lb5|FFYmn(`G8!{dPS18H&eU*0_u^@
aIQa&#Uui3fuCjDGPgllv%wUP^7yb`2@I3wi

literal 0
HcmV?d00001