From fd3d8762dc8952249696fde6f7da1c3a26a29c63 Mon Sep 17 00:00:00 2001
From: A-zero-dev <agbozero@gmail.com>
Date: Mon, 31 Aug 2026 09:52:04 +0000
Subject: [PATCH] Fix #417: Allow atomic operator replacement in set_operator

Issue #417: set_operator couldn't replace an existing operator - it returned
Error::OperatorAlreadySet instead of allowing the replacement. This required
two transactions to rotate a compromised key: revoke_operator then set_operator,
leaving a window where the stream has no operator.

Solution: Modify set_operator to allow replacement while maintaining idempotent
behavior for same-address calls. When a different operator is provided:
- Emit operator_revoked event for the old address
- Set the new operator and emit operator_set event
- This provides atomic key rotation in a single transaction

Changes:
- Modify set_operator() in contracts/stream/src/lib.rs to allow replacement
- Keep idempotent early-return for same-address calls
- Emit both operator_revoked and operator_set events on replacement
- Update tests to verify atomic replacement behavior
- Add test for idempotent same-address calls
- Add missing error variants to Error enum (InvalidRecipient, BackdatedStream,
  StreamUnderfunded) that were blocking compilation
---
 contracts/stream/src/errors.rs |  3 +++
 contracts/stream/src/lib.rs    | 11 ++++++++---
 contracts/stream/src/tests.rs  | 33 +++++++++++++++++++++++++++++----
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/contracts/stream/src/errors.rs b/contracts/stream/src/errors.rs
index c84f785e..6ac45f90 100644
--- a/contracts/stream/src/errors.rs
+++ b/contracts/stream/src/errors.rs
@@ -22,4 +22,7 @@ pub enum Error {
     ReentrancyForbidden = 16,
     OperatorAlreadySet = 17,
     NotInitialized = 18,
+    InvalidRecipient = 19,
+    BackdatedStream = 20,
+    StreamUnderfunded = 21,
 }
diff --git a/contracts/stream/src/lib.rs b/contracts/stream/src/lib.rs
index 81d5d290..ce0faa2d 100644
--- a/contracts/stream/src/lib.rs
+++ b/contracts/stream/src/lib.rs
@@ -708,12 +708,17 @@ impl DripStream {
             .instance()
             .get::<_, Address>(&DataKey::Operator)
         {
-            if existing != operator {
-                return Err(Error::OperatorAlreadySet);
+            // Issue #417: If the new operator is the same, it's idempotent — return early.
+            // If different, replace atomically: revoke the old one, then set the new one.
+            // This allows key rotation in a single transaction without a no-operator gap.
+            if existing == operator {
+                return Ok(());
             }
-            return Ok(());
+            // Revoke the old operator
+            events::operator_revoked(&env, &caller);
         }
 
+        // Set the new operator
         env.storage().instance().set(&DataKey::Operator, &operator);
         events::operator_set(&env, &caller, &operator);
         Ok(())
diff --git a/contracts/stream/src/tests.rs b/contracts/stream/src/tests.rs
index f31834c5..2cda1e78 100644
--- a/contracts/stream/src/tests.rs
+++ b/contracts/stream/src/tests.rs
@@ -1169,15 +1169,40 @@ fn set_operator_rejects_on_cancelled_stream() {
 }
 
 #[test]
-fn set_operator_requires_revoke_before_replacement() {
+fn set_operator_replaces_existing_operator_atomically() {
+    // Issue #417: set_operator should replace an existing operator in one call,
+    // allowing atomic key rotation without a no-operator gap.
     let s = Setup::new(100, 3600, false);
     let op1 = Address::generate(&s.env);
     let op2 = Address::generate(&s.env);
+    
+    // Set initial operator
     s.client.set_operator(&s.sender, &op1);
     assert_eq!(s.client.operator(), Some(op1.clone()));
-    let result = s.client.try_set_operator(&s.sender, &op2);
-    assert_eq!(result, Err(Ok(Error::OperatorAlreadySet)));
-    assert_eq!(s.client.operator(), Some(op1));
+    
+    // Replace with a different operator - should succeed now
+    s.client.set_operator(&s.sender, &op2);
+    assert_eq!(s.client.operator(), Some(op2.clone()));
+    
+    // Verify we can replace again
+    let op3 = Address::generate(&s.env);
+    s.client.set_operator(&s.sender, &op3);
+    assert_eq!(s.client.operator(), Some(op3.clone()));
+}
+
+#[test]
+fn set_operator_is_idempotent_for_same_address() {
+    // Issue #417: Setting the same operator twice should be idempotent.
+    // The early return keeps this operation idempotent.
+    let s = Setup::new(100, 3600, false);
+    let operator = Address::generate(&s.env);
+    
+    s.client.set_operator(&s.sender, &operator);
+    assert_eq!(s.client.operator(), Some(operator.clone()));
+    
+    // Setting the same operator again should succeed without error
+    s.client.set_operator(&s.sender, &operator);
+    assert_eq!(s.client.operator(), Some(operator));
 }
 
 #[test]