grin/p2p/src/store.rs

// Copyright 2016 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Storage implementation for peer data.

use std::net::SocketAddr;
use num::FromPrimitive;

use core::ser::{self, Readable, Writeable, Reader, Writer};
use grin_store::{self, Error, to_key, option_to_not_found};
use msg::SockAddr;
use types::Capabilities;

use tokio_io::codec::{Encoder, Decoder};
use peer_codec::PeerCodec;

const STORE_SUBPATH: &'static str = "peers";
const PEER_PREFIX: u8 = 'p' as u8;

/// Types of messages
enum_from_primitive! {
  #[derive(Debug, Clone, Copy, PartialEq)]
  pub enum State {
    Healthy,
    Banned,
    Dead,
  }
}

/// Data stored for any given peer we've encountered.
#[derive(Debug, Clone, PartialEq)]
pub struct PeerData {
	/// Network address of the peer.
	pub addr: SocketAddr,
	/// What capabilities the peer advertises. Unknown until a successful
	/// connection.
	pub capabilities: Capabilities,
	/// The peer user agent.
	pub user_agent: String,
	/// State the peer has been detected with.
	pub flags: State,
}

/// Storage facility for peer data.
pub struct PeerStore {
	db: grin_store::Store,
}

impl PeerStore {
	/// Instantiates a new peer store under the provided root path.
	pub fn new(root_path: String) -> Result<PeerStore, Error> {
		let db = grin_store::Store::open(format!("{}/{}", root_path, STORE_SUBPATH).as_str())?;
		Ok(PeerStore { db: db })
	}

	pub fn save_peer(&self, p: &PeerData) -> Result<(), Error> {
		let key = to_key(PEER_PREFIX, &mut format!("{}", p.addr).into_bytes());
		self.db.put_enc(&mut PeerCodec, &key, p.clone())
	}

	fn get_peer(&self, peer_addr: SocketAddr) -> Result<PeerData, Error> {
		let key = peer_key(peer_addr);
		option_to_not_found(self.db.get_dec(&mut PeerCodec, &key))
	}

	pub fn exists_peer(&self, peer_addr: SocketAddr) -> Result<bool, Error> {
		let key = peer_key(peer_addr);
		self.db.exists(&key)
	}

	pub fn delete_peer(&self, peer_addr: SocketAddr) -> Result<(), Error> {
		let key = peer_key(peer_addr);
		self.db.delete(&key)
	}

	pub fn find_peers(&self, state: State, cap: Capabilities, count: usize) -> Vec<PeerData> {
		
		let key = to_key(PEER_PREFIX, &mut "".to_string().into_bytes());
		let peers_iter = self.db.iter_dec(PeerCodec::default(), &key);
		
		let mut peers = Vec::with_capacity(count);
		for p in peers_iter {
			if p.flags == state && p.capabilities.contains(cap) {
				peers.push(p);
			}
			if peers.len() >= count {
				break;
			}
		}
		peers
	}

	/// Convenience method to load a peer data, update its status and save it
	/// back.
	pub fn update_state(&self, peer_addr: SocketAddr, new_state: State) -> Result<(), Error> {
		let mut peer = self.get_peer(peer_addr)?;
		peer.flags = new_state;
		self.save_peer(&peer)
	}
}

fn peer_key(peer_addr: SocketAddr) -> Vec<u8> {
	to_key(PEER_PREFIX, &mut format!("{}", peer_addr).into_bytes())
}