Move initiateKeyUpdate to event queue

Summary: InitiateKeyUpdate was implemented in a dangerous way, bypassing the event queue. To safely invoke the state machine in middle of a connection, we instead need to create a new pending event type, and insert an event into the queue. Reviewed By: mingtaoy Differential Revision: D50860258 fbshipit-source-id: 28bd4dbc4f4bffe6031700f871f481ca68398e42
facebookincubator · Nov 6, 2023 · b6bd537 · b6bd537
1 parent 7d37de5
commit b6bd537
Show file tree

Hide file tree

Showing 7 changed files with 56 additions and 29 deletions.
diff --git a/fizz/client/FizzClient-inl.h b/fizz/client/FizzClient-inl.h
@@ -57,13 +57,6 @@ Buf FizzClient<ActionMoveVisitor, SM>::getEarlyEkm(
       length);
 }
 
-template <typename ActionMoveVisitor, typename SM>
-void FizzClient<ActionMoveVisitor, SM>::initiateKeyUpdate(
-    KeyUpdateInitiation keyUpdateInitiation) {
-  this->addProcessingActions(this->machine_.processKeyUpdateInitiation(
-      this->state_, std::move(keyUpdateInitiation)));
-}
-
 template <typename ActionMoveVisitor, typename SM>
 void FizzClient<ActionMoveVisitor, SM>::startActions(Actions actions) {
   this->processActions(actions);

diff --git a/fizz/client/FizzClient.h b/fizz/client/FizzClient.h
@@ -51,11 +51,6 @@ class FizzClient : public FizzBase<
       const Buf& context,
       uint16_t length) const;
 
-  /**
-   * Initialize a key update.
-   */
-  void initiateKeyUpdate(KeyUpdateInitiation keyUpdateInitiation);
-
  protected:
   void visitActions(typename SM::CompletedActions& actions) override;
 

diff --git a/fizz/protocol/FizzBase-inl.h b/fizz/protocol/FizzBase-inl.h
@@ -16,6 +16,13 @@ void FizzBase<Derived, ActionMoveVisitor, StateMachine>::writeNewSessionTicket(
   processPendingEvents();
 }
 
+template <typename Derived, typename ActionMoveVisitor, typename StateMachine>
+void FizzBase<Derived, ActionMoveVisitor, StateMachine>::initiateKeyUpdate(
+    KeyUpdateInitiation kui) {
+  pendingEvents_.push_back(std::move(kui));
+  processPendingEvents();
+}
+
 template <typename Derived, typename ActionMoveVisitor, typename StateMachine>
 void FizzBase<Derived, ActionMoveVisitor, StateMachine>::appWrite(AppWrite w) {
   pendingEvents_.push_back(std::move(w));
@@ -81,6 +88,7 @@ void FizzBase<Derived, ActionMoveVisitor, StateMachine>::moveToErrorState(
         break;
       case detail::PendingEvent::Type::AppClose_E:
       case detail::PendingEvent::Type::WriteNewSessionTicket_E:
+      case detail::PendingEvent::Type::KeyUpdateInitiation_E:
         break;
     }
   }
@@ -182,6 +190,10 @@ void FizzBase<Derived, ActionMoveVisitor, StateMachine>::
             actions.emplace(machine_.processAppCloseImmediate(state_));
           }
           break;
+        case detail::PendingEvent::Type::KeyUpdateInitiation_E:
+          actions.emplace(machine_.processKeyUpdateInitiation(
+              state_, std::move(*event.asKeyUpdateInitiation())));
+          break;
       }
     } else {
       actionGuard_.reset();

diff --git a/fizz/protocol/FizzBase.h b/fizz/protocol/FizzBase.h
@@ -16,11 +16,12 @@ namespace fizz {
 
 namespace detail {
 
-#define FIZZ_PENDING_EVENT(F, ...) \
-  F(AppWrite, __VA_ARGS__)         \
-  F(EarlyAppWrite, __VA_ARGS__)    \
-  F(AppClose, __VA_ARGS__)         \
-  F(WriteNewSessionTicket, __VA_ARGS__)
+#define FIZZ_PENDING_EVENT(F, ...)      \
+  F(AppWrite, __VA_ARGS__)              \
+  F(EarlyAppWrite, __VA_ARGS__)         \
+  F(AppClose, __VA_ARGS__)              \
+  F(WriteNewSessionTicket, __VA_ARGS__) \
+  F(KeyUpdateInitiation, __VA_ARGS__)
 
 // `fizz::detail::PendingEvent` is declared here, rather than
 // being declared as a private struct within `fizz::FizzBase` in order to
@@ -63,6 +64,11 @@ class FizzBase {
    */
   void writeNewSessionTicket(WriteNewSessionTicket writeNewSessionTicket);
 
+  /**
+   * Called to initiate a key update.
+   */
+  void initiateKeyUpdate(KeyUpdateInitiation keyUpdateInitiation);
+
   /**
    * Called to write application data.
    */

diff --git a/fizz/protocol/test/FizzBaseTest.cpp b/fizz/protocol/test/FizzBaseTest.cpp
@@ -97,6 +97,15 @@ class TestStateMachine {
       WriteNewSessionTicket ticket) {
     return processWriteNewSessionTicket_(state, ticket);
   }
+  MOCK_METHOD(
+      Future<Actions>,
+      processKeyUpdateInitiation_,
+      (const State&, KeyUpdateInitiation&));
+  Future<Actions> processKeyUpdateInitiation(
+      const State& state,
+      KeyUpdateInitiation kui) {
+    return processKeyUpdateInitiation_(state, kui);
+  }
   MOCK_METHOD(Future<Actions>, processAppWrite_, (const State&, AppWrite&));
   MOCK_METHOD(
       Future<Actions>,
@@ -292,6 +301,30 @@ TEST_F(FizzBaseTest, TestWriteMulti) {
   testFizz_->appWrite(appWrite("write2"));
 }
 
+TEST_F(FizzBaseTest, TestInitiateKeyUpdate) {
+  EXPECT_CALL(*TestStateMachine::instance, processAppWrite_(_, _))
+      .WillOnce(InvokeWithoutArgs([]() {
+        Actions actions;
+        actions.push_back(A1());
+        return actions;
+      }));
+  bool inCallback = false;
+  EXPECT_CALL(testFizz_->visitor_, a1())
+      .WillOnce(InvokeWithoutArgs([this, &inCallback] {
+        inCallback = true;
+        SCOPE_EXIT {
+          inCallback = false;
+        };
+        testFizz_->initiateKeyUpdate(KeyUpdateInitiation());
+      }));
+  EXPECT_CALL(*TestStateMachine::instance, processKeyUpdateInitiation_(_, _))
+      .WillOnce(InvokeWithoutArgs([&inCallback]() {
+        EXPECT_FALSE(inCallback);
+        return Actions();
+      }));
+  testFizz_->appWrite(AppWrite());
+}
+
 TEST_F(FizzBaseTest, TestAppClose) {
   EXPECT_CALL(*TestStateMachine::instance, processAppClose(_))
       .WillOnce(InvokeWithoutArgs([]() {

diff --git a/fizz/server/FizzServer-inl.h b/fizz/server/FizzServer-inl.h
@@ -64,13 +64,6 @@ Buf FizzServer<ActionMoveVisitor, SM>::getEarlyEkm(
       length);
 }
 
-template <typename ActionMoveVisitor, typename SM>
-void FizzServer<ActionMoveVisitor, SM>::initiateKeyUpdate(
-    KeyUpdateInitiation keyUpdateInitiation) {
-  this->addProcessingActions(this->machine_.processKeyUpdateInitiation(
-      this->state_, std::move(keyUpdateInitiation)));
-}
-
 template <typename ActionMoveVisitor, typename SM>
 void FizzServer<ActionMoveVisitor, SM>::startActions(AsyncActions actions) {
   folly::variant_match(

diff --git a/fizz/server/FizzServer.h b/fizz/server/FizzServer.h
@@ -43,11 +43,6 @@ class FizzServer : public FizzBase<
       const Buf& context,
       uint16_t length) const;
 
-  /**
-   * Initialize a key update.
-   */
-  void initiateKeyUpdate(KeyUpdateInitiation keyUpdateInitiation);
-
  protected:
   void visitActions(typename SM::CompletedActions& actions) override;