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"))]
println!("B2B: Starting proxy loop :)");
let peer_address = stream.peer_addr().unwrap();
loop {
// first, forward all data we have.
while let Some((client_id, tag, flags, payload)) = local_receiver
.recv_buf_with_flags()
.expect("Error reading from local page!")
{
loop {
match local_receiver.recv_buf_with_flags() {
Ok(None) => break, // no more data to forward
Ok(Some((client_id, tag, flags, payload))) => {
if client_id == b2b_client_id {
println!(
"Ignored message we probably sent earlier (same id), TAG: {:x}",
@ -2116,7 +2118,7 @@ where
payload.len()
);
// We got a new message! Forward...
send_tcp_msg(
if let Err(e) = send_tcp_msg(
&mut stream,
&TcpRemoteNewMessage {
client_id,
@ -2124,8 +2126,17 @@ where
flags,
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.
@ -2134,7 +2145,8 @@ where
// Forwarding happens between each recv, too, as simplification.
// We ignore errors completely as they may be timeout, or stream closings.
// 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(
"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: somehow mangle in the other broker id? ClientId?
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.");
} 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"))]
println!("Received no input, timeout or closed. Looping back up :)");
}
}
}
});
let ret = recv.recv().map_err(|_| {