From efad146c17f0f933111eafdccd65c990e51df6fc Mon Sep 17 00:00:00 2001
From: Jason Simmons <jsimmons@google.com>
Date: Thu, 29 Feb 2024 23:18:53 +0000
Subject: [PATCH] [Impeller] Fix a buffer overrun in ImpellerC reflector
 resource offsets

---
 impeller/compiler/reflector.cc | 23 ++++++++++++++---------
 impeller/compiler/reflector.h  |  3 +++
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/impeller/compiler/reflector.cc b/impeller/compiler/reflector.cc
index dd085b25de0c3..1fd6ba2e0d204 100644
--- a/impeller/compiler/reflector.cc
+++ b/impeller/compiler/reflector.cc
@@ -421,9 +421,7 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
     const auto inputs = compiler_->get_shader_resources().stage_inputs;
     auto input_offsets = ComputeOffsets(inputs);
     for (const auto& input : inputs) {
-      auto location = compiler_->get_decoration(
-          input.id, spv::Decoration::DecorationLocation);
-      std::optional<size_t> offset = input_offsets[location];
+      std::optional<size_t> offset = GetOffset(input.id, input_offsets);
 
       const auto type = compiler_->get_type(input.type_id);
 
@@ -518,9 +516,7 @@ std::shared_ptr<ShaderBundleData> Reflector::GenerateShaderBundleData() const {
     const auto inputs = compiler_->get_shader_resources().stage_inputs;
     auto input_offsets = ComputeOffsets(inputs);
     for (const auto& input : inputs) {
-      auto location = compiler_->get_decoration(
-          input.id, spv::Decoration::DecorationLocation);
-      std::optional<size_t> offset = input_offsets[location];
+      std::optional<size_t> offset = GetOffset(input.id, input_offsets);
 
       const auto type = compiler_->get_type(input.type_id);
 
@@ -624,6 +620,17 @@ std::vector<size_t> Reflector::ComputeOffsets(
   return offsets;
 }
 
+std::optional<size_t> Reflector::GetOffset(
+    spirv_cross::ID id,
+    const std::vector<size_t>& offsets) const {
+  uint32_t location =
+      compiler_->get_decoration(id, spv::Decoration::DecorationLocation);
+  if (location >= offsets.size()) {
+    return std::nullopt;
+  }
+  return offsets[location];
+}
+
 std::optional<nlohmann::json::object_t> Reflector::ReflectResource(
     const spirv_cross::Resource& resource,
     std::optional<size_t> offset) const {
@@ -698,9 +705,7 @@ std::optional<nlohmann::json::array_t> Reflector::ReflectResources(
   for (const auto& resource : resources) {
     std::optional<size_t> maybe_offset = std::nullopt;
     if (compute_offsets) {
-      auto location = compiler_->get_decoration(
-          resource.id, spv::Decoration::DecorationLocation);
-      maybe_offset = offsets[location];
+      maybe_offset = GetOffset(resource.id, offsets);
     }
     if (auto reflected = ReflectResource(resource, maybe_offset);
         reflected.has_value()) {
diff --git a/impeller/compiler/reflector.h b/impeller/compiler/reflector.h
index 1448e412d4738..b246ada5b4399 100644
--- a/impeller/compiler/reflector.h
+++ b/impeller/compiler/reflector.h
@@ -225,6 +225,9 @@ class Reflector {
   std::vector<size_t> ComputeOffsets(
       const spirv_cross::SmallVector<spirv_cross::Resource>& resources) const;
 
+  std::optional<size_t> GetOffset(spirv_cross::ID id,
+                                  const std::vector<size_t>& offsets) const;
+
   std::optional<nlohmann::json::object_t> ReflectType(
       const spirv_cross::TypeID& type_id) const;