feat: schedule automatic undelegation on AML check failure#1300
feat: schedule automatic undelegation on AML check failure#1300Dodecahedr0x wants to merge 3 commits into
Conversation
📝 WalkthroughWalkthroughThis PR implements AML-triggered undelegation scheduling. When FetchCloner detects high-risk addresses during post-delegation action validation, it can trigger undelegation scheduling through a new Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@magicblock-api/src/magic_validator.rs`:
- Around line 115-120: Wrap the await on self.0.schedule_undelegation(pubkey,
request.account).await in a tokio timeout (e.g.,
tokio::time::timeout(Duration::from_secs(...), ...)) so the call cannot hang
indefinitely; after awaiting the timeout, map a timeout error to the same
chainlink/committor error variant currently produced (the "committor response
channel closed" mapping) so both channel-closure and timeout produce the same
error path, and then continue to .and_then(|result| ...) as before; update
imports if necessary to include tokio::time::timeout and Duration and ensure the
timeout mapping converts into a String consistent with the existing
.map_err(|message| ...) flow.
In `@magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs`:
- Around line 251-255: The mock server's blocking accept loop inside the
tokio::task::spawn_blocking worker (which calls listener.accept() in a for
0..expected_calls loop) can hang CI if fewer connections arrive; modify the
worker to enforce a hard timeout (e.g., 5s) for each accept or for the entire
loop: set the TcpListener to non-blocking (listener.set_nonblocking(true)) or
use Instant::now()/elapsed to track time, repeatedly try accept() with a short
sleep, and if the deadline is exceeded before receiving all expected_calls,
break and return an error (or panic) so the test fails fast instead of hanging
while the test awaits the worker join. Ensure the change is applied to the same
worker block that spawns the mock server and to the other similar instance
mentioned (around the other accept loop).
In `@magicblock-chainlink/tests/10_aml_undelegation.rs`:
- Around line 70-74: The mock risk-server worker can block forever on
listener.accept() inside the spawned closure (the worker) causing join() to
hang; change the worker loop to avoid blocking by calling
listener.set_nonblocking(true) (or use try_accept if available) and perform a
bounded retry loop with a short sleep, counting successful accepts until
expected_calls or until a configurable timeout elapses, then break and return an
error/result so the task completes; update the code that awaits the worker (the
join() site) to handle the early-return error case. Ensure you reference the
spawned closure (worker), listener.accept, expected_calls and the join()
awaiting the worker when making the change.
In `@magicblock-committor-service/src/committor_processor.rs`:
- Around line 302-309: The current intent_id() uses wall-clock nanoseconds which
can collide; replace it with a collision-free generator such as a monotonic
atomic counter or a UUID-based ID instead of SystemTime::now(). Specifically,
update the intent_id() function to read from a shared AtomicU64 (e.g.,
NEXT_INTENT_ID.fetch_add(1, Ordering::SeqCst)) or generate a UUID and
convert/encode it for ScheduledIntentBundle.id so each undelegation gets a
unique, monotonic or globally-unique ID; ensure the chosen generator is
initialized and accessible where intent_id() is called and that persisted
keys/status lookups use the same ID format.
In `@magicblock-committor-service/src/service.rs`:
- Around line 68-72: The ScheduleUndelegation variant embeds AccountSharedData
which causes CommittorService::try_send() to log full account bytes on
TrySendError::{Full,Closed}; change the logging to avoid dumping the account
payload by either redacting ScheduleUndelegation's Debug output or logging only
the pubkey and message type. Concretely, update the CommittorService::try_send()
error paths that currently log the CommittorMessage to instead pattern-match the
message and log a concise string for ScheduleUndelegation (e.g.,
"ScheduleUndelegation { pubkey: <pubkey> }") or implement a custom Debug/Display
for that enum variant that omits AccountSharedData; ensure references to
ScheduleUndelegation, CommittorMessage, AccountSharedData, Pubkey, and
TrySendError::Full/Closed are used to locate and change the relevant code.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: dd15f542-18ab-4785-a422-c9f5e2e1d1a1
⛔ Files ignored due to path filters (2)
Cargo.lockis excluded by!**/*.locktest-integration/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (15)
magicblock-api/Cargo.tomlmagicblock-api/src/magic_validator.rsmagicblock-chainlink/Cargo.tomlmagicblock-chainlink/src/chainlink/errors.rsmagicblock-chainlink/src/chainlink/fetch_cloner/mod.rsmagicblock-chainlink/src/chainlink/fetch_cloner/tests.rsmagicblock-chainlink/src/chainlink/mod.rsmagicblock-chainlink/tests/10_aml_undelegation.rsmagicblock-chainlink/tests/utils/test_context.rsmagicblock-committor-service/src/committor_processor.rsmagicblock-committor-service/src/service.rsmagicblock-committor-service/src/service_ext.rsmagicblock-committor-service/src/stubs/changeset_committor_stub.rstest-integration/test-chainlink/src/ixtest_context.rstest-integration/test-chainlink/src/test_context.rs
| self.0 | ||
| .schedule_undelegation(pubkey, request.account) | ||
| .await | ||
| .map_err(|err| format!("committor response channel closed: {err}")) | ||
| .and_then(|result| result.map_err(|err| err.to_string())) | ||
| .map_err(|message| { |
There was a problem hiding this comment.
Add a timeout around committor scheduling await.
Line 115 awaits the committor response without a deadline. If that actor stalls, this task can hang indefinitely and back up AML-rejection handling. Bound this external call with a timeout and map timeout to the same chainlink error variant.
Suggested fix
let pubkey = request.pubkey;
- self.0
- .schedule_undelegation(pubkey, request.account)
- .await
+ tokio::time::timeout(
+ Duration::from_secs(5),
+ self.0.schedule_undelegation(pubkey, request.account),
+ )
+ .await
+ .map_err(|_| "committor schedule_undelegation timed out".to_string())?
.map_err(|err| format!("committor response channel closed: {err}"))
.and_then(|result| result.map_err(|err| err.to_string()))
.map_err(|message| {
magicblock_chainlink::errors::ChainlinkError::FailedToScheduleUndelegationAfterAmlRejection(
pubkey, message,
)
})🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-api/src/magic_validator.rs` around lines 115 - 120, Wrap the await
on self.0.schedule_undelegation(pubkey, request.account).await in a tokio
timeout (e.g., tokio::time::timeout(Duration::from_secs(...), ...)) so the call
cannot hang indefinitely; after awaiting the timeout, map a timeout error to the
same chainlink/committor error variant currently produced (the "committor
response channel closed" mapping) so both channel-closure and timeout produce
the same error path, and then continue to .and_then(|result| ...) as before;
update imports if necessary to include tokio::time::timeout and Duration and
ensure the timeout mapping converts into a String consistent with the existing
.map_err(|message| ...) flow.
| let worker = tokio::task::spawn_blocking(move || { | ||
| for _ in 0..expected_calls { | ||
| let (mut stream, _) = | ||
| listener.accept().expect("accept mock risk request"); | ||
| stream |
There was a problem hiding this comment.
Bound the mock server accept loop to prevent CI hangs.
Line 254 blocks on listener.accept() for each expected request, and Line 289 waits on that worker. If AML calls are fewer than expected_calls, the test can hang instead of failing fast.
Proposed hard-timeout pattern for the mock server worker
impl MockRiskServer {
async fn start(
address_scores: Vec<(String, u64)>,
expected_calls: usize,
) -> Self {
let listener =
TcpListener::bind("127.0.0.1:0").expect("bind mock risk server");
+ listener
+ .set_nonblocking(true)
+ .expect("set mock risk listener nonblocking");
let addr = listener.local_addr().expect("mock risk server address");
let score_by_address: HashMap<String, u64> =
address_scores.into_iter().collect();
let worker = tokio::task::spawn_blocking(move || {
- for _ in 0..expected_calls {
- let (mut stream, _) =
- listener.accept().expect("accept mock risk request");
+ let deadline = std::time::Instant::now() + Duration::from_secs(3);
+ let mut served = 0usize;
+ while served < expected_calls {
+ let (mut stream, _) = match listener.accept() {
+ Ok(conn) => conn,
+ Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "timed out waiting for mock risk requests: served {served}/{expected_calls}"
+ );
+ std::thread::sleep(Duration::from_millis(10));
+ continue;
+ }
+ Err(err) => panic!("accept mock risk request: {err}"),
+ };
+ served += 1;
stream
.set_read_timeout(Some(Duration::from_secs(2)))
.expect("set mock risk read timeout");
...
}
});Also applies to: 288-290
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-chainlink/src/chainlink/fetch_cloner/tests.rs` around lines 251 -
255, The mock server's blocking accept loop inside the
tokio::task::spawn_blocking worker (which calls listener.accept() in a for
0..expected_calls loop) can hang CI if fewer connections arrive; modify the
worker to enforce a hard timeout (e.g., 5s) for each accept or for the entire
loop: set the TcpListener to non-blocking (listener.set_nonblocking(true)) or
use Instant::now()/elapsed to track time, repeatedly try accept() with a short
sleep, and if the deadline is exceeded before receiving all expected_calls,
break and return an error (or panic) so the test fails fast instead of hanging
while the test awaits the worker join. Ensure the change is applied to the same
worker block that spawns the mock server and to the other similar instance
mentioned (around the other accept loop).
| let worker = tokio::task::spawn_blocking(move || { | ||
| for _ in 0..expected_calls { | ||
| let (mut stream, _) = | ||
| listener.accept().expect("accept mock risk request"); | ||
| stream |
There was a problem hiding this comment.
Prevent indefinite blocking in the mock risk server worker.
Line 73 can block forever on accept() when the expected call count is not reached, and Line 108 then waits forever on join(). This can stall the whole test run.
Suggested fail-fast listener loop
async fn start(
address_scores: Vec<(String, u64)>,
expected_calls: usize,
) -> Self {
let listener =
TcpListener::bind("127.0.0.1:0").expect("bind mock risk server");
+ listener
+ .set_nonblocking(true)
+ .expect("set mock risk listener nonblocking");
let addr = listener.local_addr().expect("mock risk server address");
let score_by_address: HashMap<String, u64> =
address_scores.into_iter().collect();
let worker = tokio::task::spawn_blocking(move || {
- for _ in 0..expected_calls {
- let (mut stream, _) =
- listener.accept().expect("accept mock risk request");
+ let deadline = std::time::Instant::now() + Duration::from_secs(3);
+ let mut served = 0usize;
+ while served < expected_calls {
+ let (mut stream, _) = match listener.accept() {
+ Ok(conn) => conn,
+ Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
+ assert!(
+ std::time::Instant::now() < deadline,
+ "timed out waiting for mock risk requests: served {served}/{expected_calls}"
+ );
+ std::thread::sleep(Duration::from_millis(10));
+ continue;
+ }
+ Err(err) => panic!("accept mock risk request: {err}"),
+ };
+ served += 1;
...
}
});Also applies to: 107-109
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-chainlink/tests/10_aml_undelegation.rs` around lines 70 - 74, The
mock risk-server worker can block forever on listener.accept() inside the
spawned closure (the worker) causing join() to hang; change the worker loop to
avoid blocking by calling listener.set_nonblocking(true) (or use try_accept if
available) and perform a bounded retry loop with a short sleep, counting
successful accepts until expected_calls or until a configurable timeout elapses,
then break and return an error/result so the task completes; update the code
that awaits the worker (the join() site) to handle the early-return error case.
Ensure you reference the spawned closure (worker), listener.accept,
expected_calls and the join() awaiting the worker when making the change.
| fn intent_id() -> u64 { | ||
| SystemTime::now() | ||
| .duration_since(UNIX_EPOCH) | ||
| .unwrap_or_default() | ||
| .as_nanos() | ||
| .try_into() | ||
| .unwrap_or(u64::MAX) | ||
| } |
There was a problem hiding this comment.
Generate undelegation intent IDs with a collision-free source.
intent_id() is using SystemTime::now().as_nanos() as the sole message ID. Two undelegations scheduled in the same clock tick can get the same ScheduledIntentBundle.id, and that ID is later used as the key for persisted rows / status lookup. The result is unrelated undelegations being merged under one message ID.
Use a monotonic counter or another collision-free generator here instead of raw wall-clock time.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-committor-service/src/committor_processor.rs` around lines 302 -
309, The current intent_id() uses wall-clock nanoseconds which can collide;
replace it with a collision-free generator such as a monotonic atomic counter or
a UUID-based ID instead of SystemTime::now(). Specifically, update the
intent_id() function to read from a shared AtomicU64 (e.g.,
NEXT_INTENT_ID.fetch_add(1, Ordering::SeqCst)) or generate a UUID and
convert/encode it for ScheduledIntentBundle.id so each undelegation gets a
unique, monotonic or globally-unique ID; ensure the chosen generator is
initialized and accessible where intent_id() is called and that persisted
keys/status lookups use the same ID format.
| ScheduleUndelegation { | ||
| pubkey: Pubkey, | ||
| account: AccountSharedData, | ||
| respond_to: oneshot::Sender<CommittorServiceResult<()>>, | ||
| }, |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial | ⚡ Quick win
Avoid enqueue-failure logs that dump full account payloads.
CommittorService::try_send() still logs the full CommittorMessage on TrySendError::{Full,Closed}. With AccountSharedData embedded here, a saturated/closed queue will now print raw account bytes for every failed undelegation request. Please log just the pubkey / message type, or redact this variant’s Debug output.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@magicblock-committor-service/src/service.rs` around lines 68 - 72, The
ScheduleUndelegation variant embeds AccountSharedData which causes
CommittorService::try_send() to log full account bytes on
TrySendError::{Full,Closed}; change the logging to avoid dumping the account
payload by either redacting ScheduleUndelegation's Debug output or logging only
the pubkey and message type. Concretely, update the CommittorService::try_send()
error paths that currently log the CommittorMessage to instead pattern-match the
message and log a concise string for ScheduleUndelegation (e.g.,
"ScheduleUndelegation { pubkey: <pubkey> }") or implement a custom Debug/Display
for that enum variant that omits AccountSharedData; ensure references to
ScheduleUndelegation, CommittorMessage, AccountSharedData, Pubkey, and
TrySendError::Full/Closed are used to locate and change the relevant code.
|
@Dodecahedr0x : can you please add some description as to when/why undelegation should be scheduled and automated. |
There was a problem hiding this comment.
I think this PR digresses from current committor flow, and imo there's no need in this deviation.
Delegation is an account state property, therefore it has to go through builtin program and modify delegated flags.
This makes flow similar to regular Intent scheduling one, in fact you do schedule an Intent youself.
- Validator detect account is malicious
- Validator sends TX in magic-program(new entrypoint, or an old one with commit)
- magic-program mutates account states correctly and schedules undelegation(
mark_account_as_undelegated) - CommittorService picks it up and executes as in current flow
This gets rid of new schedule_undelegation functionality completely as it will seamlessly integrates into current Committor logic. All is needed is new entrypoint in magic-program and TX sender
Current state bypasses all of that flow to which it belongs. This additionally lead to an amigious state where in AccountsDB account is delegated but on Base it is not. Maybe cloner will fix this(which I'm not sure about) after some updates, but this leaves window for an attack where in meantime it could be used by user.
| struct CommittorUndelegationScheduler(Arc<CommittorService>); | ||
|
|
||
| #[async_trait::async_trait] | ||
| impl UndelegationScheduler for CommittorUndelegationScheduler { |
There was a problem hiding this comment.
magic_validator.rs is already massive, let's extract this into separate file
| } | ||
| } | ||
|
|
||
| pub(crate) fn undelegation_intent_bundle( |
There was a problem hiding this comment.
This has to be created by magic-program which additionally would call mark_account_as_undelegated
| } | ||
| } | ||
|
|
||
| fn intent_id() -> u64 { |
There was a problem hiding this comment.
Intent ids are assigned by magic-program
snawaz
changed the title
3 participants
Summary
Automatically schedules the undelegation of AML-flagged accounts. Currently, if an account fails the AML check, it is just not cloned, meaning it stays owned by the DLP on mainnet but can't be used in the ER. This PR automatically undelegates it.
Breaking Changes
Summary by CodeRabbit
Release Notes
New Features
Tests
Chores