Keep the same read buffer unless the last message was overly large

[rust-lightning] / lightning / src / ln / peer_handler.rs

diff --git a/lightning/src/ln/peer_handler.rs b/lightning/src/ln/peer_handler.rs
index bb7b697e420ea5603e18dca088808b5150938ecd..f771cf4ff2a4b36d25720f526ad577a182e8b1e0 100644 (file)
--- a/lightning/src/ln/peer_handler.rs
+++ b/lightning/src/ln/peer_handler.rs
@@ -33,7 +33,7 @@ use routing::network_graph::{NetworkGraph, NetGraphMsgHandler};
 use prelude::*;
 use io;
 use alloc::collections::LinkedList;
-use sync::{Arc, Mutex, MutexGuard, RwLock};
+use sync::{Arc, Mutex, MutexGuard, FairRwLock};
 use core::sync::atomic::{AtomicBool, Ordering};
 use core::{cmp, hash, fmt, mem};
 use core::ops::Deref;
@@ -428,7 +428,7 @@ pub struct PeerManager<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: De
                L::Target: Logger,
                CMH::Target: CustomMessageHandler {
        message_handler: MessageHandler<CM, RM>,
-       peers: RwLock<PeerHolder<Descriptor>>,
+       peers: FairRwLock<PeerHolder<Descriptor>>,
        /// Only add to this set when noise completes.
        /// Locked *after* peers. When an item is removed, it must be removed with the `peers` write
        /// lock held. Entries may be added with only the `peers` read lock held (though the
@@ -570,7 +570,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
 
                PeerManager {
                        message_handler,
-                       peers: RwLock::new(PeerHolder {
+                       peers: FairRwLock::new(PeerHolder {
                                peers: HashMap::new(),
                        }),
                        node_id_to_descriptor: Mutex::new(HashMap::new()),
@@ -972,7 +972,7 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                                                                if peer.pending_read_is_header {
                                                                        let msg_len = try_potential_handleerror!(peer,
                                                                                peer.channel_encryptor.decrypt_length_header(&peer.pending_read_buffer[..]));
-                                                                       peer.pending_read_buffer = Vec::with_capacity(msg_len as usize + 16);
+                                                                       if peer.pending_read_buffer.capacity() > 8192 { peer.pending_read_buffer = Vec::new(); }
                                                                        peer.pending_read_buffer.resize(msg_len as usize + 16, 0);
                                                                        if msg_len < 2 { // Need at least the message type tag
                                                                                return Err(PeerHandleError{ no_connection_possible: false });
@@ -984,7 +984,8 @@ impl<Descriptor: SocketDescriptor, CM: Deref, RM: Deref, L: Deref, CMH: Deref> P
                                                                        assert!(msg_data.len() >= 2);
 
                                                                        // Reset read buffer
-                                                                       peer.pending_read_buffer = [0; 18].to_vec();
+                                                                       if peer.pending_read_buffer.capacity() > 8192 { peer.pending_read_buffer = Vec::new(); }
+                                                                       peer.pending_read_buffer.resize(18, 0);
                                                                        peer.pending_read_is_header = true;
 
                                                                        let mut reader = io::Cursor::new(&msg_data[..]);