From 19601eea1326237507a939ce7fd26c496654f5ea Mon Sep 17 00:00:00 2001
From: Piotr Sikora <piotrsikora@google.com>
Date: Mon, 26 Oct 2020 07:39:47 +0000
Subject: [PATCH] review: fix getRootContext().

Signed-off-by: Piotr Sikora <piotrsikora@google.com>
---
 include/proxy-wasm/wasm.h | 4 +---
 src/wasm.cc               | 8 ++++++++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/include/proxy-wasm/wasm.h b/include/proxy-wasm/wasm.h
index b8199ace..3e0fdb58 100644
--- a/include/proxy-wasm/wasm.h
+++ b/include/proxy-wasm/wasm.h
@@ -60,9 +60,7 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   WasmVm *wasm_vm() const { return wasm_vm_.get(); }
   ContextBase *vm_context() const { return vm_context_.get(); }
   ContextBase *createAndSaveRootContext(const std::shared_ptr<PluginBase> &plugin);
-  ContextBase *getRootContext(std::string_view plugin_key) {
-    return root_contexts_[std::string(plugin_key)].get();
-  }
+  ContextBase *getRootContext(std::string_view plugin_key);
   ContextBase *getOrCreateRootContext(const std::shared_ptr<PluginBase> &plugin);
   ContextBase *getContext(uint32_t id) {
     auto it = contexts_.find(id);
diff --git a/src/wasm.cc b/src/wasm.cc
index dd99b21c..de1e8e72 100644
--- a/src/wasm.cc
+++ b/src/wasm.cc
@@ -321,6 +321,14 @@ ContextBase *WasmBase::createAndSaveRootContext(const std::shared_ptr<PluginBase
   return root_context;
 }
 
+ContextBase *WasmBase::getRootContext(std::string_view plugin_key) {
+  auto it = root_contexts_.find(std::string(plugin_key));
+  if (it == root_contexts_.end()) {
+    return nullptr;
+  }
+  return it->second.get();
+}
+
 ContextBase *WasmBase::getOrCreateRootContext(const std::shared_ptr<PluginBase> &plugin) {
   auto root_context = getRootContext(plugin->key());
   if (!root_context) {