MojoLPM: Refactor to remove fuzzer-specific Context.

This change cleans up the CodeCacheHost fuzzer, moving the per-testcase
code into the CodeCacheHostTestcase object and removing
CodeCacheHostFuzzerContext since this was redundant.

Bug: 1076336
Change-Id: I8736106edadb4a88c9d46a007ed5404fec6d4d4c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2329571
Commit-Queue: Mark Brand <markbrand@google.com>
Reviewed-by: Scott Violet <sky@chromium.org>
Cr-Commit-Position: refs/heads/master@{#799623}

content/test/fuzzer/code_cache_host_mojolpm_fuzzer.cc

@@ -33,10 +33,21 @@
 
 using url::Origin;
 
-namespace content {
-
 const char* cmdline[] = {"code_cache_host_mojolpm_fuzzer", nullptr};
 
+// Global environment needed to run the interface being tested.
+//
+// This will be created once, before fuzzing starts, and will be shared between
+// all testcases. It is created on the main thread.
+//
+// At a minimum, we should always be able to set up the command line, i18n and
+// mojo, and create the thread on which the fuzzer will be run. We want to avoid
+// (as much as is reasonable) any state being preserved between testcases.
+//
+// For CodeCacheHost, we can also safely re-use a single BrowserTaskEnvironment
+// and the TestContentClientInitializer between testcases. We try to create an
+// environment that matches the real browser process as much as possible, so we
+// use real platform threads in the task environment.
 class ContentFuzzerEnvironment {
  public:
   ContentFuzzerEnvironment()
@@ -60,154 +71,189 @@ class ContentFuzzerEnvironment {
 
  private:
   base::AtExitManager at_exit_manager_;
-  std::unique_ptr<base::FieldTrialList> field_trial_list_;
-  base::test::ScopedFeatureList scoped_feature_list_;
   base::Thread fuzzer_thread_;
-  BrowserTaskEnvironment task_environment_;
-  TestContentClientInitializer content_client_initializer_;
+  content::BrowserTaskEnvironment task_environment_;
+  content::TestContentClientInitializer content_client_initializer_;
 };
 
-ContentFuzzerEnvironment g_environment;
-
-ContentFuzzerEnvironment& SingletonEnvironment() {
-  return g_environment;
+ContentFuzzerEnvironment& GetEnvironment() {
+  static base::NoDestructor<ContentFuzzerEnvironment> environment;
+  return *environment;
 }
 
 scoped_refptr<base::SequencedTaskRunner> GetFuzzerTaskRunner() {
-  return SingletonEnvironment().fuzzer_task_runner();
+  return GetEnvironment().fuzzer_task_runner();
 }
 
-class CodeCacheHostFuzzerContext : public mojolpm::Context {
-  const Origin kOriginA;
-  const Origin kOriginB;
-  const Origin kOriginOpaque;
-  const Origin kOriginEmpty;
-
+// Per-testcase state needed to run the interface being tested.
+//
+// The lifetime of this is scoped to a single testcase, and it is created and
+// destroyed from the fuzzer sequence.
+//
+// For CodeCacheHost, this needs the basic common Browser process state provided
+// by TestBrowserContext, and to set up the cache storage that will provide the
+// storage backing for the code cache.
+//
+// Since the Browser process will host one CodeCacheHostImpl per
+// RenderProcessHost, we emulate this by allowing the fuzzer to create (and
+// destroy) multiple CodeCacheHostImpl instances.
+class CodeCacheHostTestcase {
  public:
-  CodeCacheHostFuzzerContext()
-      : kOriginA(url::Origin::Create(GURL("http://aaa.com/"))),
-        kOriginB(url::Origin::Create(GURL("http://bbb.com/"))),
-        kOriginOpaque(url::Origin::Create(GURL("opaque"))),
-        kOriginEmpty(url::Origin::Create(GURL("file://this_becomes_empty"))),
-        browser_context_() {
-    base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
-    base::PostTaskAndReply(
-        FROM_HERE, {BrowserThread::UI},
-        base::BindOnce(&CodeCacheHostFuzzerContext::InitializeOnUIThread,
-                       base::Unretained(this)),
-        run_loop.QuitClosure());
-    run_loop.Run();
-  }
+  explicit CodeCacheHostTestcase(
+      const content::fuzzing::code_cache_host::proto::Testcase& testcase);
 
-  void InitializeOnUIThread() {
-    cache_storage_context_ = base::MakeRefCounted<CacheStorageContextImpl>();
-    cache_storage_context_->Init(browser_context_.GetPath(),
-                                 browser_context_.GetSpecialStoragePolicy(),
-                                 nullptr);
+  // Returns true once either all of the actions in the testcase have been
+  // performed, or the per-testcase action limit has been exceeded.
+  //
+  // This should only be called from the fuzzer sequence.
+  bool IsFinished();
 
-    generated_code_cache_context_ =
-        base::MakeRefCounted<GeneratedCodeCacheContext>();
-    generated_code_cache_context_->Initialize(browser_context_.GetPath(),
-                                              65536);
-  }
+  // If there are still actions remaining in the testcase, this will perform the
+  // next sequence of actions before returning.
+  //
+  // If IsFinished() would return true, then calling this function is a no-op.
+  //
+  // This should only be called from the fuzzer sequence.
+  void NextAction();
 
+  void SetUp();
+  void TearDown();
+
+ private:
+  using Action = content::fuzzing::code_cache_host::proto::Action;
+
+  void SetUpOnUIThread();
+  void TearDownOnUIThread();
+
+  // Used by AddCodeCacheHost to create and bind CodeCacheHostImpl on the UI
+  // thread.
   void AddCodeCacheHostImpl(
       uint32_t id,
       int renderer_id,
       const Origin& origin,
-      mojo::PendingReceiver<::blink::mojom::CodeCacheHost>&& receiver) {
-    code_cache_hosts_[renderer_id] = std::make_unique<CodeCacheHostImpl>(
-        renderer_id, cache_storage_context_, generated_code_cache_context_,
-        std::move(receiver));
-  }
+      mojo::PendingReceiver<::blink::mojom::CodeCacheHost>&& receiver);
 
+  // Create and bind a new instance for fuzzing. This needs to  make sure that
+  // the new instance has been created and bound on the correct sequence before
+  // returning.
   void AddCodeCacheHost(
       uint32_t id,
       int renderer_id,
       content::fuzzing::code_cache_host::proto::NewCodeCacheHostAction::OriginId
-          origin_id) {
-    mojo::Remote<::blink::mojom::CodeCacheHost> remote;
-    auto receiver = remote.BindNewPipeAndPassReceiver();
-
-    const Origin* origin = &kOriginA;
-    if (origin_id == 1) {
-      origin = &kOriginB;
-    } else if (origin_id == 2) {
-      origin = &kOriginOpaque;
-    } else if (origin_id == 3) {
-      origin = &kOriginEmpty;
-    }
+          origin_id);
 
-    base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
-    base::PostTaskAndReply(
-        FROM_HERE, {BrowserThread::UI},
-        base::BindOnce(&CodeCacheHostFuzzerContext::AddCodeCacheHostImpl,
-                       base::Unretained(this), id, renderer_id, *origin,
-                       std::move(receiver)),
-        run_loop.QuitClosure());
-    run_loop.Run();
+  // The proto message describing the test actions to perform.
+  const content::fuzzing::code_cache_host::proto::Testcase& testcase_;
 
-    AddInstance(id, std::move(remote));
-  }
+  // This set of origins should cover all of the origin types which have special
+  // handling in CodeCacheHostImpl, and give us two distinct "normal" origins,
+  // which should be enough to exercise all of the code.
+  const Origin origin_a_;
+  const Origin origin_b_;
+  const Origin origin_opaque_;
+  const Origin origin_empty_;
 
- private:
-  TestBrowserContext browser_context_;
+  // Apply a reasonable upper-bound on testcase complexity to avoid timeouts.
+  const int max_action_count_ = 512;
 
-  scoped_refptr<CacheStorageContextImpl> cache_storage_context_;
-  scoped_refptr<GeneratedCodeCacheContext> generated_code_cache_context_;
+  // Count of total actions performed in this testcase.
+  int action_count_ = 0;
 
-  std::map<int, std::unique_ptr<CodeCacheHostImpl>> code_cache_hosts_;
-};
-}  // namespace content
+  // The index of the next sequence of actions to execute.
+  int next_sequence_idx_ = 0;
 
-class CodeCacheHostTestcase : public mojolpm::TestcaseBase {
-  content::CodeCacheHostFuzzerContext& cch_context_;
-  const content::fuzzing::code_cache_host::proto::Testcase& testcase_;
-  int next_idx_ = 0;
-  int action_count_ = 0;
-  const int MAX_ACTION_COUNT = 512;
+  // Prerequisite state.
+  std::unique_ptr<content::TestBrowserContext> browser_context_;
+  scoped_refptr<content::CacheStorageContextImpl> cache_storage_context_;
+  scoped_refptr<content::GeneratedCodeCacheContext>
+      generated_code_cache_context_;
 
- public:
-  CodeCacheHostTestcase(
-      content::CodeCacheHostFuzzerContext& cch_context,
-      const content::fuzzing::code_cache_host::proto::Testcase& testcase);
-  ~CodeCacheHostTestcase() override;
-  bool IsFinished() override;
-  void NextAction() override;
+  // Mapping from renderer id to CodeCacheHostImpl instances being fuzzed.
+  // Access only from UI thread.
+  std::map<int, std::unique_ptr<content::CodeCacheHostImpl>> code_cache_hosts_;
+
+  SEQUENCE_CHECKER(sequence_checker_);
 };
 
 CodeCacheHostTestcase::CodeCacheHostTestcase(
-    content::CodeCacheHostFuzzerContext& cch_context,
     const content::fuzzing::code_cache_host::proto::Testcase& testcase)
-    : cch_context_(cch_context), testcase_(testcase) {}
+    : testcase_(testcase),
+      origin_a_(url::Origin::Create(GURL("http://aaa.com/"))),
+      origin_b_(url::Origin::Create(GURL("http://bbb.com/"))),
+      origin_opaque_(url::Origin::Create(GURL("opaque"))),
+      origin_empty_(url::Origin::Create(GURL("file://this_becomes_empty"))) {
+  // CodeCacheHostTestcase is created on the main thread, but the actions that
+  // we want to validate the sequencing of take place on the fuzzer sequence.
+  DETACH_FROM_SEQUENCE(sequence_checker_);
+}
+
+void CodeCacheHostTestcase::SetUp() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
+  base::PostTaskAndReply(FROM_HERE, {content::BrowserThread::UI},
+                         base::BindOnce(&CodeCacheHostTestcase::SetUpOnUIThread,
+                                        base::Unretained(this)),
+                         run_loop.QuitClosure());
+  run_loop.Run();
+}
 
-CodeCacheHostTestcase::~CodeCacheHostTestcase() {}
+void CodeCacheHostTestcase::SetUpOnUIThread() {
+  browser_context_ = std::make_unique<content::TestBrowserContext>();
+
+  cache_storage_context_ =
+      base::MakeRefCounted<content::CacheStorageContextImpl>();
+  cache_storage_context_->Init(browser_context_->GetPath(),
+                               browser_context_->GetSpecialStoragePolicy(),
+                               nullptr);
+
+  generated_code_cache_context_ =
+      base::MakeRefCounted<content::GeneratedCodeCacheContext>();
+  generated_code_cache_context_->Initialize(browser_context_->GetPath(), 65536);
+}
+
+void CodeCacheHostTestcase::TearDown() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
+  base::PostTaskAndReply(
+      FROM_HERE, {content::BrowserThread::UI},
+      base::BindOnce(&CodeCacheHostTestcase::TearDownOnUIThread,
+                     base::Unretained(this)),
+      run_loop.QuitClosure());
+  run_loop.Run();
+}
+
+void CodeCacheHostTestcase::TearDownOnUIThread() {
+  code_cache_hosts_.clear();
+  generated_code_cache_context_.reset();
+  cache_storage_context_.reset();
+  browser_context_.reset();
+}
 
 bool CodeCacheHostTestcase::IsFinished() {
-  return next_idx_ >= testcase_.sequence_indexes_size();
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  return next_sequence_idx_ >= testcase_.sequence_indexes_size();
 }
 
 void CodeCacheHostTestcase::NextAction() {
-  if (next_idx_ < testcase_.sequence_indexes_size()) {
-    auto sequence_idx = testcase_.sequence_indexes(next_idx_++);
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  if (next_sequence_idx_ < testcase_.sequence_indexes_size()) {
+    auto sequence_idx = testcase_.sequence_indexes(next_sequence_idx_++);
     const auto& sequence =
         testcase_.sequences(sequence_idx % testcase_.sequences_size());
     for (auto action_idx : sequence.action_indexes()) {
-      if (!testcase_.actions_size() || ++action_count_ > MAX_ACTION_COUNT) {
+      if (!testcase_.actions_size() || ++action_count_ > max_action_count_) {
         return;
       }
       const auto& action =
           testcase_.actions(action_idx % testcase_.actions_size());
       switch (action.action_case()) {
-        case content::fuzzing::code_cache_host::proto::Action::
-            kNewCodeCacheHost: {
-          cch_context_.AddCodeCacheHost(
-              action.new_code_cache_host().id(),
-              action.new_code_cache_host().render_process_id(),
-              action.new_code_cache_host().origin_id());
+        case Action::kNewCodeCacheHost: {
+          AddCodeCacheHost(action.new_code_cache_host().id(),
+                           action.new_code_cache_host().render_process_id(),
+                           action.new_code_cache_host().origin_id());
         } break;
 
-        case content::fuzzing::code_cache_host::proto::Action::kRunThread: {
+        case Action::kRunThread: {
           if (action.run_thread().id()) {
             base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
             base::PostTask(FROM_HERE, {content::BrowserThread::UI},
@@ -221,8 +267,7 @@ void CodeCacheHostTestcase::NextAction() {
           }
         } break;
 
-        case content::fuzzing::code_cache_host::proto::Action::
-            kCodeCacheHostRemoteAction: {
+        case Action::kCodeCacheHostRemoteAction: {
           mojolpm::HandleRemoteAction(action.code_cache_host_remote_action());
         } break;
 
@@ -233,53 +278,100 @@ void CodeCacheHostTestcase::NextAction() {
   }
 }
 
-void NextAction(content::CodeCacheHostFuzzerContext* context,
+void CodeCacheHostTestcase::AddCodeCacheHostImpl(
+    uint32_t id,
+    int renderer_id,
+    const Origin& origin,
+    mojo::PendingReceiver<::blink::mojom::CodeCacheHost>&& receiver) {
+  code_cache_hosts_[renderer_id] = std::make_unique<content::CodeCacheHostImpl>(
+      renderer_id, cache_storage_context_, generated_code_cache_context_,
+      std::move(receiver));
+}
+
+void CodeCacheHostTestcase::AddCodeCacheHost(
+    uint32_t id,
+    int renderer_id,
+    content::fuzzing::code_cache_host::proto::NewCodeCacheHostAction::OriginId
+        origin_id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  mojo::Remote<::blink::mojom::CodeCacheHost> remote;
+  auto receiver = remote.BindNewPipeAndPassReceiver();
+
+  const Origin* origin = &origin_a_;
+  if (origin_id == 1) {
+    origin = &origin_b_;
+  } else if (origin_id == 2) {
+    origin = &origin_opaque_;
+  } else if (origin_id == 3) {
+    origin = &origin_empty_;
+  }
+
+  base::RunLoop run_loop(base::RunLoop::Type::kNestableTasksAllowed);
+  base::PostTaskAndReply(
+      FROM_HERE, {content::BrowserThread::UI},
+      base::BindOnce(&CodeCacheHostTestcase::AddCodeCacheHostImpl,
+                     base::Unretained(this), id, renderer_id, *origin,
+                     std::move(receiver)),
+      run_loop.QuitClosure());
+  run_loop.Run();
+
+  mojolpm::GetContext()->AddInstance(id, std::move(remote));
+}
+
+// Helper function to keep scheduling fuzzer actions on the current runloop
+// until the testcase has completed, and then quit the runloop.
+void NextAction(CodeCacheHostTestcase* testcase,
                 base::RepeatingClosure quit_closure) {
-  if (!context->IsFinished()) {
-    context->NextAction();
-    content::GetFuzzerTaskRunner()->PostTask(
-        FROM_HERE, base::BindOnce(NextAction, base::Unretained(context),
+  if (!testcase->IsFinished()) {
+    testcase->NextAction();
+    GetFuzzerTaskRunner()->PostTask(
+        FROM_HERE, base::BindOnce(NextAction, base::Unretained(testcase),
                                   std::move(quit_closure)));
   } else {
-    content::GetFuzzerTaskRunner()->PostTask(FROM_HERE,
-                                             std::move(quit_closure));
+    GetFuzzerTaskRunner()->PostTask(FROM_HERE, std::move(quit_closure));
   }
 }
 
-void RunTestcase(
-    content::CodeCacheHostFuzzerContext* context,
-    const content::fuzzing::code_cache_host::proto::Testcase* testcase) {
+// Helper function to setup and run the testcase, since we need to do that from
+// the fuzzer sequence rather than the main thread.
+void RunTestcase(CodeCacheHostTestcase* testcase) {
   mojo::Message message;
   auto dispatch_context =
       std::make_unique<mojo::internal::MessageDispatchContext>(&message);
 
-  CodeCacheHostTestcase cch_testcase(*context, *testcase);
-  context->StartTestcase(&cch_testcase, content::GetFuzzerTaskRunner());
+  testcase->SetUp();
+
+  mojolpm::GetContext()->StartTestcase();
 
   base::RunLoop fuzzer_run_loop(base::RunLoop::Type::kNestableTasksAllowed);
-  content::GetFuzzerTaskRunner()->PostTask(
-      FROM_HERE, base::BindOnce(NextAction, base::Unretained(context),
+  GetFuzzerTaskRunner()->PostTask(
+      FROM_HERE, base::BindOnce(NextAction, base::Unretained(testcase),
                                 fuzzer_run_loop.QuitClosure()));
   fuzzer_run_loop.Run();
 
-  context->EndTestcase();
+  mojolpm::GetContext()->EndTestcase();
+
+  testcase->TearDown();
 }
 
 DEFINE_BINARY_PROTO_FUZZER(
-    const content::fuzzing::code_cache_host::proto::Testcase& testcase) {
-  if (!testcase.actions_size() || !testcase.sequences_size() ||
-      !testcase.sequence_indexes_size()) {
+    const content::fuzzing::code_cache_host::proto::Testcase& proto_testcase) {
+  if (!proto_testcase.actions_size() || !proto_testcase.sequences_size() ||
+      !proto_testcase.sequence_indexes_size()) {
     return;
   }
 
-  content::CodeCacheHostFuzzerContext context;
-  mojolpm::SetContext(&context);
+  // Make sure that the environment is initialized before we do anything else.
+  GetEnvironment();
+
+  CodeCacheHostTestcase testcase(proto_testcase);
 
   base::RunLoop ui_run_loop(base::RunLoop::Type::kNestableTasksAllowed);
-  content::GetFuzzerTaskRunner()->PostTaskAndReply(
-      FROM_HERE,
-      base::BindOnce(RunTestcase, base::Unretained(&context),
-                     base::Unretained(&testcase)),
+
+  // Unretained is safe here, because ui_run_loop has to finish before testcase
+  // goes out of scope.
+  GetFuzzerTaskRunner()->PostTaskAndReply(
+      FROM_HERE, base::BindOnce(RunTestcase, base::Unretained(&testcase)),
       ui_run_loop.QuitClosure());
 
   ui_run_loop.Run();

mojo/public/tools/fuzzers/mojolpm.cc

@@ -22,14 +22,16 @@ Context::Storage::Storage() = default;
 
 Context::Storage::~Storage() = default;
 
-void Context::StartTestcase(
-    TestcaseBase* testcase,
-    scoped_refptr<base::SequencedTaskRunner> task_runner) {
-  testcase_ = testcase;
-  task_runner_ = task_runner;
+void Context::StartTestcase() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
 }
 
 void Context::EndTestcase() {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  // Some fuzzers need to destroy the fuzzer thread along with their testcase,
+  // so we need to detach the sequence checker here so that it will be attached
+  // to the new sequence for the next testcase.
+  DETACH_FROM_SEQUENCE(sequence_checker_);
   // We need to destroy all Remotes/Receivers before we start destroying other
   // objects (like callbacks).
   for (const TypeId& interface_type_id : interface_type_ids_) {
@@ -38,41 +40,13 @@ void Context::EndTestcase() {
       instances_iter->second.clear();
     }
   }
+  interface_type_ids_.clear();
   instances_.clear();
-  testcase_ = nullptr;
 }
 
-bool Context::IsFinished() {
-  if (testcase_) {
-    return testcase_->IsFinished();
-  }
-  return true;
-}
-
-void Context::NextAction() {
-  // fprintf(stderr, "NextAction\n");
-  CHECK(task_runner_->RunsTasksInCurrentSequence());
-  if (testcase_) {
-    testcase_->NextAction();
-  }
-}
-
-void Context::PostNextAction() {
-  if (task_runner_) {
-    task_runner_->PostTask(FROM_HERE, base::BindOnce(&Context::NextAction,
-                                                     base::Unretained(this)));
-  }
-}
-
-Context* g_context = nullptr;
-
 Context* GetContext() {
-  DCHECK(g_context);
-  return g_context;
-}
-
-void SetContext(Context* context) {
-  g_context = context;
+  static base::NoDestructor<Context> context;
+  return context.get();
 }
 
 bool FromProto(const bool& input, bool& output) {

mojo/public/tools/fuzzers/mojolpm.h

@@ -11,6 +11,7 @@
 #include "base/check.h"
 #include "base/containers/flat_map.h"
 #include "base/optional.h"
+#include "base/sequence_checker.h"
 #include "mojo/public/cpp/bindings/associated_receiver.h"
 #include "mojo/public/cpp/bindings/associated_remote.h"
 #include "mojo/public/cpp/bindings/receiver.h"
@@ -31,6 +32,7 @@ namespace mojolpm {
 
 typedef void* TypeId;
 
+// Returns a unique TypeId for every different typename T.
 template <typename T>
 TypeId type_id() {
   static std::remove_reference<T>* ptr = nullptr;
@@ -45,14 +47,89 @@ std::string type_name() {
 }
 #endif
 
-class TestcaseBase {
+// Common state shared between generated MojoLPM code and fuzzer-specific
+// testcases.
+//
+// The main state is (various) object instances used as parameters and response
+// parameters to Mojo method calls, along with the instances of Mojo Remote and
+// Receiver objects.
+//
+// At code-generation time we don't have access to all of the types used, we
+// can't know ahead of time all of the types that might be stored. This is
+// further complicated by the move-only and sequence-bound semantics of some
+// Mojo objects.
+//
+// The interface
+//   AddInstance<T>(), GetInstance<T>(), RemoveInstance<T>(), etc.
+// implements a per-type mapping from integer id to instance of type T. Adding a
+// new instance can be performed with or without providing an id; in the case
+// that a fuzzer attempts to store two instances of the same type with the same
+// id at once, the second instance will be immediately destroyed.
+//
+// We want most calls to GetInstance to return a valid object, even if there
+// wasn't an object previously stored with the given id (lookup ids are
+// randomly generated during fuzzing), so GetInstance does a very fuzzy lookup.
+//
+// The details of the algorithms used to generate a new id, and match the id
+// when there is no exact match should not be depended on. If a fuzzer needs to
+// know the id that will be used to store an object, but does not have a
+// testcase-supplied id, then the function NextId<T>() should be used to choose
+// an id.
+//
+// All methods should only be called from the fuzzer sequence.
+class Context {
  public:
-  virtual ~TestcaseBase() = default;
-  virtual bool IsFinished() = 0;
-  virtual void NextAction() = 0;
-};
+  Context();
+  Context(const Context&) = delete;
+  ~Context();
 
-class Context {
+  // Lookup the previously stored instance of type T using a fuzzy-match on the
+  // provided id. Returns nullptr if there's no matching instance.
+  template <typename T>
+  T* GetInstance(uint32_t id);
+
+  // Lookup the previously stored instance of type T using a fuzzy-match on
+  // the provided id, and remove that instance from the object storage, passing
+  // ownership of that instance to the caller. Returns nullptr if there's no
+  // matching instance.
+  template <typename T>
+  std::unique_ptr<T> GetAndRemoveInstance(uint32_t id);
+
+  // Lookup the previously stored instance of type T using a fuzzy-match on
+  // the provided id, and remove that instance from the object storage.
+  template <typename T>
+  void RemoveInstance(uint32_t id);
+
+  // Adds an instance of type T to the object storage using an automatically
+  // chosen id, which it returns. Equivalent to calling
+  //   AddInstance<T>(NextId<T>(), instance);
+  template <typename T>
+  uint32_t AddInstance(T instance);
+
+  // Adds an instance of type T to the object storage using the provided id. If
+  // the provided id already exists in the object storage, the existing instance
+  // is not modified, and the implementation will assign a new id. Returns the
+  // id that can be used to lookup the instance.
+  template <typename T>
+  uint32_t AddInstance(uint32_t id, T instance);
+
+  // Returns an instance id for the given type T that is guaranteed to be
+  // available for storing an instance of type T at the time of calling.
+  //
+  // NB: This does NOT reserve the id; so the following snippet is not correct.
+  //
+  // uint32_t id = NextId<T>();
+  // AddInstance<T>(some_t);
+  // CHECK_EQ(id, AddInstance<T>(id, some_t));
+  template <typename T>
+  uint32_t NextId();
+
+  void StartTestcase();
+  void EndTestcase();
+
+  mojo::Message& message() { return message_; }
+
+ private:
   // mojolpm::Context::Storage implements generic storage for all possible
   // object types that might be created during fuzzing. This allows the fuzzer
   // to reference objects by id, even when the possible types of those objects
@@ -97,8 +174,6 @@ class Context {
     std::unique_ptr<StorageWrapperBase> wrapper_;
   };
 
-  std::map<TypeId, std::map<uint32_t, Storage>> instances_;
-
   // mojolpm::Context::StorageTraits implements type-specific details in the
   // handling of stored instances. In particular, we need to guarantee that
   // certain types are destroyed before other types - all fuzzer-owned
@@ -152,51 +227,16 @@ class Context {
     Context* context_;
   };
 
-  std::set<TypeId> interface_type_ids_;
+  std::map<TypeId, std::map<uint32_t, Storage>> instances_
+      GUARDED_BY_CONTEXT(sequence_checker_);
+  std::set<TypeId> interface_type_ids_ GUARDED_BY_CONTEXT(sequence_checker_);
 
-  TestcaseBase* testcase_;
-  scoped_refptr<base::SequencedTaskRunner> task_runner_;
   mojo::Message message_;
 
- public:
-  explicit Context();
-  Context(const Context&) = delete;
-  ~Context();
-
-  template <typename T>
-  T* GetInstance(uint32_t id);
-
-  template <typename T>
-  std::unique_ptr<T> GetAndRemoveInstance(uint32_t id);
-
-  template <typename T>
-  void RemoveInstance(uint32_t id);
-
-  template <typename T>
-  uint32_t AddInstance(T&& instance);
-
-  template <typename T>
-  uint32_t AddInstance(uint32_t id, T&& instance);
-
-  template <typename T>
-  uint32_t NextId();
-
-  void StartTestcase(TestcaseBase* testcase,
-                     scoped_refptr<base::SequencedTaskRunner> task_runner);
-  void EndTestcase();
-  bool IsFinished();
-  void NextAction();
-  void PostNextAction();
-
-  scoped_refptr<base::SequencedTaskRunner> task_runner() const {
-    return task_runner_;
-  }
-
-  mojo::Message& message() { return message_; }
+  SEQUENCE_CHECKER(sequence_checker_);
 };
 
 Context* GetContext();
-void SetContext(Context* context);
 
 template <typename T>
 Context::Storage::Storage(T&& value) {
@@ -205,21 +245,21 @@ Context::Storage::Storage(T&& value) {
 
 template <typename T>
 T& Context::Storage::Storage::get() {
-  // DCHECK(wrapper_->type() == type_id<T>());
+  DCHECK(wrapper_->type() == type_id<T>());
   DCHECK(static_cast<StorageWrapper<T>*>(wrapper_.get()));
   return static_cast<StorageWrapper<T>*>(wrapper_.get())->value();
 }
 
 template <typename T>
 const T& Context::Storage::Storage::get() const {
-  // DCHECK(wrapper_->type() == type_id<T>());
+  DCHECK(wrapper_->type() == type_id<T>());
   DCHECK(static_cast<StorageWrapper<T>*>(wrapper_.get()));
   return static_cast<StorageWrapper<T>*>(wrapper_.get())->value();
 }
 
 template <typename T>
 std::unique_ptr<T> Context::Storage::Storage::release() {
-  // DCHECK(wrapper_->type() == type_id<T>());
+  DCHECK(wrapper_->type() == type_id<T>());
   DCHECK(static_cast<StorageWrapper<T>*>(wrapper_.get()));
   return std::make_unique<T>(
       std::move(static_cast<StorageWrapper<T>*>(wrapper_.get())->value()));
@@ -246,9 +286,11 @@ const T& Context::Storage::StorageWrapper<T>::value() const {
 
 template <typename T>
 void Context::StorageTraits<::mojo::Remote<T>>::OnInstanceAdded(uint32_t id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(context_->sequence_checker_);
   context_->interface_type_ids_.insert(type_id<::mojo::Remote<T>>());
   auto instance = context_->GetInstance<::mojo::Remote<T>>(id);
   CHECK(instance);
+  // Unretained is safe here since context_ owns instance.
   instance->set_disconnect_handler(
       base::BindOnce(&Context::RemoveInstance<::mojo::Remote<T>>,
                      base::Unretained(context_), id));
@@ -257,9 +299,11 @@ void Context::StorageTraits<::mojo::Remote<T>>::OnInstanceAdded(uint32_t id) {
 template <typename T>
 void Context::StorageTraits<::mojo::AssociatedRemote<T>>::OnInstanceAdded(
     uint32_t id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(context_->sequence_checker_);
   context_->interface_type_ids_.insert(type_id<::mojo::AssociatedRemote<T>>());
   auto instance = context_->GetInstance<::mojo::AssociatedRemote<T>>(id);
   CHECK(instance);
+  // Unretained is safe here since context_ owns instance.
   instance->set_disconnect_handler(
       base::BindOnce(&Context::RemoveInstance<::mojo::AssociatedRemote<T>>,
                      base::Unretained(context_), id));
@@ -268,11 +312,13 @@ void Context::StorageTraits<::mojo::AssociatedRemote<T>>::OnInstanceAdded(
 template <typename T>
 void Context::StorageTraits<
     std::unique_ptr<::mojo::Receiver<T>>>::OnInstanceAdded(uint32_t id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(context_->sequence_checker_);
   context_->interface_type_ids_.insert(
       type_id<std::unique_ptr<::mojo::Receiver<T>>>());
   auto instance =
       context_->GetInstance<std::unique_ptr<::mojo::Receiver<T>>>(id);
   CHECK(instance);
+  // Unretained is safe here since context_ owns instance.
   (*instance)->set_disconnect_handler(base::BindOnce(
       &Context::RemoveInstance<std::unique_ptr<::mojo::Receiver<T>>>,
       base::Unretained(context_), id));
@@ -281,11 +327,13 @@ void Context::StorageTraits<
 template <typename T>
 void Context::StorageTraits<std::unique_ptr<::mojo::AssociatedReceiver<T>>>::
     OnInstanceAdded(uint32_t id) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(context_->sequence_checker_);
   context_->interface_type_ids_.insert(
       type_id<std::unique_ptr<::mojo::AssociatedReceiver<T>>>());
   auto instance =
       context_->GetInstance<std::unique_ptr<::mojo::AssociatedReceiver<T>>>(id);
   CHECK(instance);
+  // Unretained is safe here since context_ owns instance.
   (*instance)->set_disconnect_handler(base::BindOnce(
       &Context::RemoveInstance<std::unique_ptr<::mojo::AssociatedReceiver<T>>>,
       base::Unretained(context_), id));
@@ -293,7 +341,7 @@ void Context::StorageTraits<std::unique_ptr<::mojo::AssociatedReceiver<T>>>::
 
 template <typename T>
 T* Context::GetInstance(uint32_t id) {
-  DCHECK(task_runner_->RunsTasksInCurrentSequence());
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   mojolpmdbg("getInstance(%s, %i) = ", type_name<T>().c_str(), id);
   auto instances_iter = instances_.find(type_id<T>());
   if (instances_iter == instances_.end()) {
@@ -321,7 +369,7 @@ T* Context::GetInstance(uint32_t id) {
 
 template <typename T>
 std::unique_ptr<T> Context::GetAndRemoveInstance(uint32_t id) {
-  DCHECK(task_runner_->RunsTasksInCurrentSequence());
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   mojolpmdbg("getAndRemoveInstance(%s, %i) = ", type_name<T>().c_str(), id);
   auto instances_iter = instances_.find(type_id<T>());
   if (instances_iter == instances_.end()) {
@@ -351,7 +399,7 @@ std::unique_ptr<T> Context::GetAndRemoveInstance(uint32_t id) {
 
 template <typename T>
 void Context::RemoveInstance(uint32_t id) {
-  DCHECK(task_runner_->RunsTasksInCurrentSequence());
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   mojolpmdbg("RemoveInstance(%s, %u) = ", type_name<T>().c_str(), id);
   auto instances_iter = instances_.find(type_id<T>());
   if (instances_iter != instances_.end()) {
@@ -376,19 +424,21 @@ void Context::RemoveInstance(uint32_t id) {
 }
 
 template <typename T>
-uint32_t Context::AddInstance(T&& instance) {
+uint32_t Context::AddInstance(T instance) {
   return AddInstance(1, std::move(instance));
 }
 
 template <typename T>
-uint32_t Context::AddInstance(uint32_t id, T&& instance) {
-  DCHECK(task_runner_->RunsTasksInCurrentSequence());
+uint32_t Context::AddInstance(uint32_t id, T instance) {
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   auto instances_iter = instances_.find(type_id<T>());
   if (instances_iter == instances_.end()) {
     instances_[type_id<T>()].emplace(id, std::move(instance));
   } else {
     auto& instance_map = instances_iter->second;
     auto instance_map_iter = instance_map.find(id);
+    // if this id is a collision with an existing entry, loop until we find a
+    // free id.
     while (instance_map_iter != instance_map.end()) {
       id = instance_map_iter->first + 1;
       instance_map_iter = instance_map.find(id);
@@ -403,13 +453,17 @@ uint32_t Context::AddInstance(uint32_t id, T&& instance) {
 
 template <typename T>
 uint32_t Context::NextId() {
-  DCHECK(task_runner_->RunsTasksInCurrentSequence());
+  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
   uint32_t id = 1;
   auto instances_iter = instances_.find(type_id<T>());
   if (instances_iter != instances_.end()) {
     auto& instance_map = instances_iter->second;
-    if (instance_map.size() > 0) {
-      id = instance_map.rbegin()->first + 1;
+    auto instance_map_iter = instance_map.find(id);
+    // if this id is a collision with an existing entry, loop until we find a
+    // free id.
+    while (instance_map_iter != instance_map.end()) {
+      id = instance_map_iter->first + 1;
+      instance_map_iter = instance_map.find(id);
     }
   }
   return id;