From e6a2bdfacdb584bc77d51076bb49d852febfdeb6 Mon Sep 17 00:00:00 2001
From: Hariharan Seshadri <shariharan91@gmail.com>
Date: Fri, 29 Mar 2019 10:42:34 -0700
Subject: [PATCH] Handle incorrect perm data in Transpose op gracefully  (#739)

---
 onnxruntime/core/providers/cpu/tensor/transpose.cc  |  4 +++-
 onnxruntime/core/providers/cpu/tensor/transpose.h   | 13 ++++++++++++-
 onnxruntime/core/providers/cuda/tensor/transpose.cc |  4 +++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.cc b/onnxruntime/core/providers/cpu/tensor/transpose.cc
index 16a516532b..ff3380e4e7 100644
--- a/onnxruntime/core/providers/cpu/tensor/transpose.cc
+++ b/onnxruntime/core/providers/cpu/tensor/transpose.cc
@@ -267,7 +267,9 @@ Status Transpose::Compute(OpKernelContext* ctx) const {
   std::vector<int64_t> output_dims(rank);
   const std::vector<int64_t>* p_perm;
   std::vector<int64_t> default_perm(rank);
-  ComputeOutputShape(X, output_dims, default_perm, p_perm);
+  const auto& status = ComputeOutputShape(X, output_dims, default_perm, p_perm);
+  if (!status.IsOK())
+    return status;
 
   TensorShape output_shape{output_dims};
   Tensor& Y = *ctx->Output(0, output_shape);
diff --git a/onnxruntime/core/providers/cpu/tensor/transpose.h b/onnxruntime/core/providers/cpu/tensor/transpose.h
index 8be880a712..50432c7875 100644
--- a/onnxruntime/core/providers/cpu/tensor/transpose.h
+++ b/onnxruntime/core/providers/cpu/tensor/transpose.h
@@ -6,6 +6,7 @@
 #include "gsl/gsl_util"
 #include "core/common/common.h"
 #include "core/framework/op_kernel.h"
+#include <sstream>
 
 namespace onnxruntime {
 
@@ -36,7 +37,7 @@ class TransposeBase {
     }
   }
 
-  void ComputeOutputShape(const Tensor& X, std::vector<int64_t>& output_dims,
+  Status ComputeOutputShape(const Tensor& X, std::vector<int64_t>& output_dims,
                           std::vector<int64_t>& default_perm, const std::vector<int64_t>*& p_perm) const {
     size_t rank = X.Shape().NumDimensions();
     const auto& input_dims = X.Shape().GetDims();
@@ -57,8 +58,18 @@ class TransposeBase {
     output_dims.resize(rank);
     for (int i = 0; i < rank; i++) {
       size_t inpdim = (*p_perm)[i];
+      if (inpdim >= rank) {
+        std::ostringstream ss;
+        ss << "[ ";
+        for (const auto& p : *p_perm)
+          ss << p << " ";
+        ss << "]";
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "perm: ", ss.str(), " does not align with rank of input data: ", std::to_string(rank));
+      }
       output_dims[i] = input_dims[inpdim];
     }
+    return Status::OK();
   }
 
   bool perm_specified_ = false;
diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.cc b/onnxruntime/core/providers/cuda/tensor/transpose.cc
index d443336deb..53bdc71c61 100644
--- a/onnxruntime/core/providers/cuda/tensor/transpose.cc
+++ b/onnxruntime/core/providers/cuda/tensor/transpose.cc
@@ -31,7 +31,9 @@ Status Transpose<T>::ComputeInternal(OpKernelContext* ctx) const {
   std::vector<int64_t> output_dims(rank);
   std::vector<int64_t> default_perm(rank);
   const std::vector<int64_t>* p_perm = nullptr;
-  ComputeOutputShape(X, output_dims, default_perm, p_perm);
+  const auto& status = ComputeOutputShape(X, output_dims, default_perm, p_perm);
+  if (!status.IsOK())
+    return status;
 
   TensorShape output_shape{output_dims};
   Tensor* Y = ctx->Output(0, output_shape);