SYS-OSS/FRET-LibAFL
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Handle broker-to-broker connection interruptions more gracefully (#921)

Browse Source 
* Handle broker-to-broker connection interruptions more gracefully

Exit gracefully instead of panicking or getting stuck in infinite loops

* Run cargo-fmt
This commit is contained in:
omergreen 2022-12-04 19:55:55 +02:00 committed by GitHub
parent c879a0a8d3
commit 3bad100cb7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 72 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
libafl/src/bolts/llmp.rs
View File

@ -2096,12 +2096,14 @@ where
#[cfg(all(feature = "llmp_debug", feature = "std"))] #[cfg(all(feature = "llmp_debug", feature = "std"))]
println!("B2B: Starting proxy loop :)"); println!("B2B: Starting proxy loop :)");
let peer_address = stream.peer_addr().unwrap();
loop { loop {
// first, forward all data we have. // first, forward all data we have.
while let Some((client_id, tag, flags, payload)) = local_receiver loop {
.recv_buf_with_flags() match local_receiver.recv_buf_with_flags() {
.expect("Error reading from local page!") Ok(None) => break, // no more data to forward
{ Ok(Some((client_id, tag, flags, payload))) => {
if client_id == b2b_client_id { if client_id == b2b_client_id {
println!( println!(
"Ignored message we probably sent earlier (same id), TAG: {:x}", "Ignored message we probably sent earlier (same id), TAG: {:x}",
@ -2116,7 +2118,7 @@ where
payload.len() payload.len()
); );
// We got a new message! Forward... // We got a new message! Forward...
send_tcp_msg( if let Err(e) = send_tcp_msg(
&mut stream, &mut stream,
&TcpRemoteNewMessage { &TcpRemoteNewMessage {
client_id, client_id,
@ -2124,8 +2126,17 @@ where
flags, flags,
payload: payload.to_vec(), payload: payload.to_vec(),
}, },
) ) {
.expect("Error sending message via broker 2 broker"); println!("Got error {} while trying to forward a message to broker {}, exiting thread", e, peer_address);
return;
}
}
Err(Error::ShuttingDown) => {
println!("Local broker is shutting down, exiting thread");
return;
}
Err(e) => panic!("Error reading from local page! {}", e),
}
} }
// Then, see if we can receive something. // Then, see if we can receive something.
@ -2134,7 +2145,8 @@ where
// Forwarding happens between each recv, too, as simplification. // Forwarding happens between each recv, too, as simplification.
// We ignore errors completely as they may be timeout, or stream closings. // We ignore errors completely as they may be timeout, or stream closings.
// Instead, we catch stream close when/if we next try to send. // Instead, we catch stream close when/if we next try to send.
if let Ok(val) = recv_tcp_msg(&mut stream) { match recv_tcp_msg(&mut stream) {
Ok(val) => {
let msg: TcpRemoteNewMessage = val.try_into().expect( let msg: TcpRemoteNewMessage = val.try_into().expect(
"Illegal message received from broker 2 broker connection - shutting down.", "Illegal message received from broker 2 broker connection - shutting down.",
); );
@ -2148,13 +2160,29 @@ where
// TODO: Could probably optimize this somehow to forward all queued messages between locks... oh well. // TODO: Could probably optimize this somehow to forward all queued messages between locks... oh well.
// Todo: somehow mangle in the other broker id? ClientId? // Todo: somehow mangle in the other broker id? ClientId?
new_sender new_sender
.send_buf_with_flags(msg.tag, msg.flags | LLMP_FLAG_FROM_B2B, &msg.payload) .send_buf_with_flags(
msg.tag,
msg.flags | LLMP_FLAG_FROM_B2B,
&msg.payload,
)
.expect("B2B: Error forwarding message. Exiting."); .expect("B2B: Error forwarding message. Exiting.");
} else { }
Err(e) => {
if let Error::File(e, _) = e {
if e.kind() == ErrorKind::UnexpectedEof {
println!(
"Broker {} seems to have disconnected, exiting",
peer_address
);
return;
}
}
#[cfg(all(feature = "llmp_debug", feature = "std"))] #[cfg(all(feature = "llmp_debug", feature = "std"))]
println!("Received no input, timeout or closed. Looping back up :)"); println!("Received no input, timeout or closed. Looping back up :)");
} }
} }
}
}); });
let ret = recv.recv().map_err(|_| { let ret = recv.recv().map_err(|_| {