From d32015377e3b9235623b328585f86941bef3fab2 Mon Sep 17 00:00:00 2001 From: RJ Rybarczyk Date: Wed, 31 Jul 2019 02:57:08 +0000 Subject: [PATCH] Make SocketDescriptor::send_data take slice instead of vec and offset --- fuzz/fuzz_targets/full_stack_target.rs | 5 ++--- net-tokio/src/lib.rs | 12 ++++++------ src/ln/peer_handler.rs | 15 ++++++++------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/fuzz/fuzz_targets/full_stack_target.rs b/fuzz/fuzz_targets/full_stack_target.rs index 76beac59..6145d003 100644 --- a/fuzz/fuzz_targets/full_stack_target.rs +++ b/fuzz/fuzz_targets/full_stack_target.rs @@ -124,9 +124,8 @@ struct Peer<'a> { peers_connected: &'a RefCell<[bool; 256]>, } impl<'a> SocketDescriptor for Peer<'a> { - fn send_data(&mut self, data: &Vec, write_offset: usize, _resume_read: bool) -> usize { - assert!(write_offset < data.len()); - data.len() - write_offset + fn send_data(&mut self, data: &[u8], _resume_read: bool) -> usize { + data.len() } fn disconnect_socket(&mut self) { assert!(self.peers_connected.borrow()[self.id as usize]); diff --git a/net-tokio/src/lib.rs b/net-tokio/src/lib.rs index 54752fdb..0bc36b28 100644 --- a/net-tokio/src/lib.rs +++ b/net-tokio/src/lib.rs @@ -128,7 +128,7 @@ impl Connection { let (reader, us) = Self::new(event_notify, stream); if let Ok(initial_send) = peer_manager.new_outbound_connection(their_node_id, SocketDescriptor::new(us.clone(), peer_manager.clone())) { - if SocketDescriptor::new(us.clone(), peer_manager.clone()).send_data(&initial_send, 0, true) == initial_send.len() { + if SocketDescriptor::new(us.clone(), peer_manager.clone()).send_data(&initial_send, true) == initial_send.len() { Self::schedule_read(peer_manager, us, reader); } else { println!("Failed to write first full message to socket!"); @@ -170,7 +170,7 @@ impl SocketDescriptor { } } impl peer_handler::SocketDescriptor for SocketDescriptor { - fn send_data(&mut self, data: &Vec, write_offset: usize, resume_read: bool) -> usize { + fn send_data(&mut self, data: &[u8], resume_read: bool) -> usize { macro_rules! schedule_read { ($us_ref: expr) => { tokio::spawn(future::lazy(move || -> Result<(), ()> { @@ -211,20 +211,20 @@ impl peer_handler::SocketDescriptor for SocketDescriptor { let us_ref = self.clone(); schedule_read!(us_ref); } - if data.len() == write_offset { return 0; } + if data.is_empty() { return 0; } if us.writer.is_none() { us.read_paused = true; return 0; } - let mut bytes = bytes::BytesMut::with_capacity(data.len() - write_offset); - bytes.put(&data[write_offset..]); + let mut bytes = bytes::BytesMut::with_capacity(data.len()); + bytes.put(data); let write_res = us.writer.as_mut().unwrap().start_send(bytes.freeze()); match write_res { Ok(res) => { match res { AsyncSink::Ready => { - data.len() - write_offset + data.len() }, AsyncSink::NotReady(_) => { us.read_paused = true; diff --git a/src/ln/peer_handler.rs b/src/ln/peer_handler.rs index 03aeedcd..a13705f4 100644 --- a/src/ln/peer_handler.rs +++ b/src/ln/peer_handler.rs @@ -45,12 +45,13 @@ pub struct MessageHandler { /// careful to ensure you don't have races whereby you might register a new connection with an fd /// the same as a yet-to-be-disconnect_event()-ed. pub trait SocketDescriptor : cmp::Eq + hash::Hash + Clone { - /// Attempts to send some data from the given Vec starting at the given offset to the peer. + /// Attempts to send some data from the given slice to the peer. + /// /// Returns the amount of data which was sent, possibly 0 if the socket has since disconnected. /// Note that in the disconnected case, a disconnect_event must still fire and further write /// attempts may occur until that time. /// - /// If the returned size is smaller than data.len() - write_offset, a write_available event must + /// If the returned size is smaller than data.len(), a write_available event must /// trigger the next time more data can be written. Additionally, until the a send_data event /// completes fully, no further read_events should trigger on the same peer! /// @@ -58,7 +59,7 @@ pub trait SocketDescriptor : cmp::Eq + hash::Hash + Clone { /// events should be paused to prevent DoS in the send buffer), resume_read may be set /// indicating that read events on this descriptor should resume. A resume_read of false does /// *not* imply that further read events should be paused. - fn send_data(&mut self, data: &Vec, write_offset: usize, resume_read: bool) -> usize; + fn send_data(&mut self, data: &[u8], resume_read: bool) -> usize; /// Disconnect the socket pointed to by this SocketDescriptor. Once this function returns, no /// more calls to write_event, read_event or disconnect_event may be made with this descriptor. /// No disconnect_event should be generated as a result of this call, though obviously races @@ -387,7 +388,8 @@ impl PeerManager { }; let should_be_reading = peer.pending_outbound_buffer.len() < MSG_BUFF_SIZE; - let data_sent = descriptor.send_data(next_buff, peer.pending_outbound_buffer_first_msg_offset, should_be_reading); + let pending = &next_buff[peer.pending_outbound_buffer_first_msg_offset..]; + let data_sent = descriptor.send_data(pending, should_be_reading); peer.pending_outbound_buffer_first_msg_offset += data_sent; if peer.pending_outbound_buffer_first_msg_offset == next_buff.len() { true } else { false } } { @@ -1122,9 +1124,8 @@ mod tests { } impl SocketDescriptor for FileDescriptor { - fn send_data(&mut self, data: &Vec, write_offset: usize, _resume_read: bool) -> usize { - assert!(write_offset < data.len()); - data.len() - write_offset + fn send_data(&mut self, data: &[u8], _resume_read: bool) -> usize { + data.len() } fn disconnect_socket(&mut self) {} -- 2.30.2