Merge pull request #7011 from paritytech/whisper-testing

test framework and basic test for whisper
This commit is contained in:
Robert Habermeier 2017-12-18 11:48:20 +01:00 committed by GitHub
commit f99b977399
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 214 additions and 15 deletions

View File

@ -30,6 +30,9 @@ use rlp::{DecoderError, RlpStream, UntrustedRlp};
use message::{Message, Error as MessageError}; use message::{Message, Error as MessageError};
#[cfg(test)]
mod tests;
// how often periodic relays are. when messages are imported // how often periodic relays are. when messages are imported
// we directly broadcast. // we directly broadcast.
const RALLY_TOKEN: TimerToken = 1; const RALLY_TOKEN: TimerToken = 1;
@ -341,8 +344,8 @@ impl Peer {
.map_or(true, |filter| &(filter & message.bloom()) == message.bloom()) .map_or(true, |filter| &(filter & message.bloom()) == message.bloom())
} }
// note a message as known. returns true if it was already // note a message as known. returns false if it was already
// known, false otherwise. // known, true otherwise.
fn note_known(&mut self, message: &Message) -> bool { fn note_known(&mut self, message: &Message) -> bool {
self.known_messages.insert(message.hash().clone()) self.known_messages.insert(message.hash().clone())
} }
@ -657,6 +660,22 @@ impl<T: MessageHandler> Network<T> {
io.send(*peer, packet::STATUS, ::rlp::EMPTY_LIST_RLP.to_vec()); io.send(*peer, packet::STATUS, ::rlp::EMPTY_LIST_RLP.to_vec());
} }
fn on_packet<C: Context>(&self, io: &C, peer: &PeerId, packet_id: u8, data: &[u8]) {
let rlp = UntrustedRlp::new(data);
let res = match packet_id {
packet::STATUS => self.on_status(peer, rlp),
packet::MESSAGES => self.on_messages(peer, rlp),
packet::POW_REQUIREMENT => self.on_pow_requirement(peer, rlp),
packet::TOPIC_FILTER => self.on_topic_filter(peer, rlp),
_ => Ok(()), // ignore unknown packets.
};
if let Err(e) = res {
trace!(target: "whisper", "Disabling peer due to misbehavior: {}", e);
io.disable_peer(*peer);
}
}
fn on_disconnect(&self, peer: &PeerId) { fn on_disconnect(&self, peer: &PeerId) {
trace!(target: "whisper", "Disconnecting peer {}", peer); trace!(target: "whisper", "Disconnecting peer {}", peer);
let _ = self.peers.write().remove(peer); let _ = self.peers.write().remove(peer);
@ -673,19 +692,7 @@ impl<T: MessageHandler> ::network::NetworkProtocolHandler for Network<T> {
} }
fn read(&self, io: &NetworkContext, peer: &PeerId, packet_id: u8, data: &[u8]) { fn read(&self, io: &NetworkContext, peer: &PeerId, packet_id: u8, data: &[u8]) {
let rlp = UntrustedRlp::new(data); self.on_packet(io, peer, packet_id, data)
let res = match packet_id {
packet::STATUS => self.on_status(peer, rlp),
packet::MESSAGES => self.on_messages(peer, rlp),
packet::POW_REQUIREMENT => self.on_pow_requirement(peer, rlp),
packet::TOPIC_FILTER => self.on_topic_filter(peer, rlp),
_ => Ok(()), // ignore unknown packets.
};
if let Err(e) = res {
trace!(target: "whisper", "Disabling peer due to misbehavior: {}", e);
io.disable_peer(*peer);
}
} }
fn connected(&self, io: &NetworkContext, peer: &PeerId) { fn connected(&self, io: &NetworkContext, peer: &PeerId) {

192
whisper/src/net/tests.rs Normal file
View File

@ -0,0 +1,192 @@
// Copyright 2015-2017 Parity Technologies (UK) Ltd.
// This file is part of Parity.
// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
// You should have received a copy of the GNU General Public License
// along with Parity. If not, see <http://www.gnu.org/licenses/>.
//! Tests for the whisper network module.
use std::collections::HashSet;
use std::sync::mpsc;
use parking_lot::Mutex;
use network::{NodeId, PeerId};
use message::{CreateParams, Message};
use super::*;
struct TestHandler(Mutex<mpsc::Sender<Message>>);
impl MessageHandler for TestHandler {
fn handle_messages(&self, messages: &[Message]) {
let tx = self.0.lock();
for message in messages {
let _ = tx.send(message.clone());
}
}
}
struct TestPeer {
network: Network<TestHandler>,
recv: mpsc::Receiver<Message>,
disconnected: Mutex<HashSet<PeerId>>,
}
impl TestPeer {
fn create() -> Self {
let (tx, rx) = mpsc::channel();
TestPeer {
network: Network::new(10 * 1024 * 1024, TestHandler(Mutex::new(tx))),
recv: rx,
disconnected: Mutex::new(HashSet::new()),
}
}
}
struct TestNetwork {
peers: Vec<TestPeer>,
}
impl TestNetwork {
fn new(n_peers: usize) -> Self {
let unconnected_peers: Vec<_> = (0..n_peers).map(|_| TestPeer::create()).collect();
for i in 0..n_peers {
for j in (i + 1)..n_peers {
let (peer1, peer2) = (&unconnected_peers[i], &unconnected_peers[j]);
let ctx1 = TestContext::new(&unconnected_peers, i);
let ctx2 = TestContext::new(&unconnected_peers, j);
peer1.network.on_connect(&ctx1, &j);
peer2.network.on_connect(&ctx2, &i);
}
}
TestNetwork {
peers: unconnected_peers,
}
}
fn post_message_from(&self, id: PeerId, msg: Message) {
self.peers[id].network.post_message(msg, &TestContext::new(&self.peers, id));
}
}
enum Event {
Disconnect(PeerId, PeerId),
Send(PeerId, PeerId, u8, Vec<u8>),
}
struct TestContext<'a> {
peers: &'a [TestPeer],
local_id: PeerId,
events: Mutex<Vec<Event>>,
}
impl<'a> TestContext<'a> {
fn new(peers: &'a [TestPeer], local_id: PeerId) -> Self {
TestContext {
peers,
local_id,
events: Mutex::new(Vec::new()),
}
}
}
impl<'a> Context for TestContext<'a> {
fn disconnect_peer(&self, id: PeerId) {
self.events.lock().push(Event::Disconnect(self.local_id, id));
}
fn disable_peer(&self, id: PeerId) {
self.events.lock().push(Event::Disconnect(self.local_id, id));
}
fn node_key(&self, peer: PeerId) -> Option<NodeId> {
let mut id = NodeId::default();
id[0] = peer as _;
Some(id)
}
fn protocol_version(&self, id: ::network::ProtocolId, _peer: PeerId) -> Option<u8> {
if &id == b"shh" || &id == b"pwh" {
Some(PROTOCOL_VERSION as _)
} else {
None
}
}
fn send(&self, peer: PeerId, packet: u8, data: Vec<u8>) {
self.events.lock().push(Event::Send(self.local_id, peer, packet, data));
}
}
impl<'a> Drop for TestContext<'a> {
fn drop(&mut self) {
let events = self.events.get_mut();
while !events.is_empty() {
let mut deferred = Vec::new();
for event in events.drain(..) {
match event {
Event::Disconnect(from, target) => {
self.peers[from].network.on_disconnect(&target);
self.peers[target].network.on_disconnect(&from);
self.peers[from].disconnected.lock().insert(target);
self.peers[target].disconnected.lock().insert(from);
}
Event::Send(from, target, packet, data) => {
if self.peers[from].disconnected.lock().contains(&target) {
continue;
}
let mut inner_ctx = TestContext::new(self.peers, target);
self.peers[target].network.on_packet(
&inner_ctx,
&from,
packet,
&data[..]
);
// don't recursively apply disconnects or new messages
// from the receiver's actions yet.
let inner_events = ::std::mem::replace(inner_ctx.events.get_mut(), Vec::new());
deferred.extend(inner_events);
}
}
}
events.extend(deferred);
}
}
}
#[test]
fn message_gets_relayed() {
let network = TestNetwork::new(5);
let message = Message::create(CreateParams {
ttl: 500,
payload: b"this is my payload, pal".to_vec(),
topics: vec![[0, 1, 2, 3].into()],
work: 25,
}).unwrap();
network.post_message_from(0, message.clone());
assert!(network.peers[0].recv.try_recv().is_err());
for i in 1..5 {
assert_eq!(network.peers[i].recv.try_recv().unwrap(), message);
}
}