Add safety checks when rendering kernel key strings

Differential Revision: D69324821

Pull Request resolved: pytorch#8327

Author: dbort

runtime/kernel/operator_registry.cpp

@@ -114,44 +114,106 @@ Error register_kernels(const Span<const Kernel> kernels) {
 }
 
 namespace {
-int copy_char_as_number_to_buf(char num, char* buf) {
-  if ((char)num < 10) {
+/**
+ * Writes `num` as a decimal string to `buf` and returns the number of bytes
+ * written. Returns -1 if `buf` is too small or if `num` is not supported.
+ */
+int copy_char_as_number_to_buf(int num, char* buf, size_t buf_size) {
+  if (num < 0) {
+    return -1;
+  }
+  if (num < 10) {
+    if (buf_size < 1) {
+      return -1;
+    }
     *buf = '0' + (char)num;
-    buf += 1;
     return 1;
-  } else {
-    *buf = '0' + ((char)num) / 10;
-    buf += 1;
+  }
+  if (num < 100) {
+    if (buf_size < 2) {
+      return -1;
+    }
+    *buf++ = '0' + ((char)num) / 10;
     *buf = '0' + ((char)num) % 10;
-    buf += 1;
     return 2;
   }
+  return -1;
 }
 } // namespace
 
 namespace internal {
-void make_kernel_key_string(Span<const TensorMeta> key, char* buf) {
+Error make_kernel_key_string(
+    Span<const TensorMeta> key,
+    char* buf,
+    size_t buf_size) {
   if (key.empty()) {
-    // If no tensor is present in an op, kernel key does not apply
-    return;
+    // If no tensor is present in an op, kernel key does not apply.
+    if (buf_size > 0) {
+      buf[0] = '\0';
+    }
+    return Error::Ok;
   }
-  strncpy(buf, "v1/", 3);
+
+  // Reserve one byte for null terminator.
+  if (buf_size < 1) {
+    return Error::InvalidArgument;
+  }
+  buf_size -= 1;
+
+  // Add prefix.
+  if (buf_size < 3) {
+    return Error::InvalidArgument;
+  }
+  memcpy(buf, "v1/", 3);
   buf += 3;
+  buf_size -= 3;
+
+  // Add tensor meta.
   for (size_t i = 0; i < key.size(); i++) {
     auto& meta = key[i];
-    buf += copy_char_as_number_to_buf((char)meta.dtype_, buf);
-    *buf = ';';
-    buf += 1;
+
+    // Add dtype.
+    int n = copy_char_as_number_to_buf((int)meta.dtype_, buf, buf_size);
+    if (n < 0) {
+      return Error::InvalidArgument;
+    }
+    buf += n;
+    buf_size -= n;
+
+    // Add separator between dtype and dim order.
+    if (buf_size < 1) {
+      return Error::InvalidArgument;
+    }
+    *buf++ = ';';
+    buf_size -= 1;
+
+    // Add dim order.
     for (int j = 0; j < meta.dim_order_.size(); j++) {
-      buf += copy_char_as_number_to_buf((char)meta.dim_order_[j], buf);
-      if (j != meta.dim_order_.size() - 1) {
-        *buf = ',';
-        buf += 1;
+      n = copy_char_as_number_to_buf((int)meta.dim_order_[j], buf, buf_size);
+      if (n < 0) {
+        return Error::InvalidArgument;
+      }
+      buf += n;
+      buf_size -= n;
+
+      if (j < meta.dim_order_.size() - 1) {
+        if (buf_size < 1) {
+          return Error::InvalidArgument;
+        }
+        *buf++ = ',';
+        buf_size -= 1;
+      }
+    }
+    if (i < key.size() - 1) {
+      if (buf_size < 1) {
+        return Error::InvalidArgument;
       }
+      *buf++ = '|';
+      buf_size -= 1;
     }
-    *buf = (i < (key.size() - 1)) ? '|' : 0x00;
-    buf += 1;
   }
+  *buf = '\0'; // Space for this was reserved above.
+  return Error::Ok;
 }
 } // namespace internal
 
@@ -164,10 +226,14 @@ bool registry_has_op_function(
 Result<OpFunction> get_op_function_from_registry(
     const char* name,
     Span<const TensorMeta> meta_list) {
-  // @lint-ignore CLANGTIDY facebook-hte-CArray
-  char buf[KernelKey::MAX_SIZE] = {0};
-  internal::make_kernel_key_string(meta_list, buf);
-  KernelKey kernel_key = KernelKey(buf);
+  std::array<char, internal::kKernelKeyBufSize> key_string;
+  Error err = internal::make_kernel_key_string(
+      meta_list, key_string.data(), key_string.size());
+  if (err != Error::Ok) {
+    ET_LOG(Error, "Failed to make kernel key string");
+    return err;
+  }
+  KernelKey kernel_key = KernelKey(key_string.data());
 
   int32_t fallback_idx = -1;
   for (size_t idx = 0; idx < num_registered_kernels; idx++) {

runtime/kernel/operator_registry.h

@@ -96,39 +96,43 @@ struct TensorMeta {
 
 /**
  * Describes which dtype & dim order specialized kernel to be bound to an
- * operator. If `is_fallback_` is true, it means this kernel can be used as a
- * fallback, if false, it means this kernel can only be used if all the
- * `TensorMeta` are matched. Fallback means this kernel will be used for
- * all input tensor dtypes and dim orders, if the specialized kernel is not
- * registered.
+ * operator.
  *
- * The format of a kernel key data is a string:
- *                              "v<version>/<tensor_meta>|<tensor_meta>..."
- * Size: Up to 691               1    1    1     (42     +1) * 16
- *           Assuming max number of tensors is 16               ^
- * Kernel key version is v1 for now. If the kernel key format changes,
- * update the version to avoid breaking pre-existing kernel keys.
- * Example: v1/7;0,1,2,3
- * The kernel key has only one tensor: a double tensor with dimension 0, 1, 2, 3
+ * Kernel key data is a string with the format:
+ *
+ *     "v<version>/<tensor_meta>|<tensor_meta>..."
+ *
+ * The version is v1 for now. If the kernel key format changes, update the
+ * version to avoid breaking pre-existing kernel keys.
  *
  * Each tensor_meta has the following format: "<dtype>;<dim_order,...>"
- * Size: Up to 42                               1-2   1    24 (1 byte for 0-9; 2
- * for 10-15) + 15 commas Assuming that the max number of dims is 16 ^ Example:
- * 7;0,1,2,3 for [double; 0, 1, 2, 3]
+ *
+ * Example kernel key data: "v1/7;0,1,2,3|1;0,1,2,3,4,5,6,7"
+ *
+ * This has two tensors: the first with dtype=7 and dim order 0,1,2,3, and the
+ * second with dtype=1 and dim order 0,1,2,3,4,5,6,7.
  *
  * IMPORTANT:
  * Users should not construct a kernel key manually. Instead, it should be
  * generated from kernel yaml.
  */
 struct KernelKey {
  public:
+  /**
+   * Creates a fallback (non-specialized) kernel key: this kernel can be used
+   * for all input tensor dtypes and dim orders if the specialized kernel is not
+   * registered.
+   */
   KernelKey() : is_fallback_(true) {}
 
+  /**
+   * Creates a specialized (non-fallback) kernel key that matches a specific
+   * set of input tensor dtypes and dim orders. See the class comment for the
+   * expected format of `kernel_key_data`.
+   */
   /* implicit */ KernelKey(const char* kernel_key_data)
       : kernel_key_data_(kernel_key_data), is_fallback_(false) {}
 
-  constexpr static int MAX_SIZE = 691;
-
   bool operator==(const KernelKey& other) const {
     return this->equals(other);
   }
@@ -144,7 +148,7 @@ struct KernelKey {
     if (is_fallback_) {
       return true;
     }
-    return strncmp(kernel_key_data_, other.kernel_key_data_, MAX_SIZE) == 0;
+    return strcmp(kernel_key_data_, other.kernel_key_data_) == 0;
   }
 
   bool is_fallback() const {
@@ -194,7 +198,23 @@ struct Kernel {
 };
 
 namespace internal {
-void make_kernel_key_string(Span<const TensorMeta> key, char* buf);
+
+/**
+ * A make_kernel_key_string buffer size that is large enough to hold a kernel
+ * key string with 16 tensors of 16 dimensions, plus the trailing NUL byte.
+ */
+constexpr size_t kKernelKeyBufSize = 659;
+
+/**
+ * Given the list of input tensor dtypes + dim orders, writes the kernel key
+ * string into the buffer. Returns an error if the buffer is too small or if the
+ * tensors cannot be represented as a valid key string.
+ */
+Error make_kernel_key_string(
+    Span<const TensorMeta> key,
+    char* buf,
+    size_t buf_size);
+
 } // namespace internal
 
 /**