Skip to content

Commit

Permalink
Don't track peers' id
Browse files Browse the repository at this point in the history
  • Loading branch information
@ivanjermakov
ivanjermakov committed Oct 27, 2023
1 parent 04325f8 commit 51b74cd
Show file tree
Hide file tree
Showing 4 changed files with 15 additions and 39 deletions.
23 changes: 10 additions & 13 deletions src/peer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -187,15 +187,12 @@ pub async fn handshake(peer: &PeerInfo, state: Arc<Mutex<State>>) -> Result<TcpS
debug!("peer response: {}", hex(&msg));
if let Message::Handshake {
info_hash: h_info_hash,
peer_id: h_peer_id,
..
} = Message::try_from(msg)
.map_err(Error::msg)
.context("handshake parse error")?
{
ensure!(h_info_hash == *info_hash, "response `info_hash` differ");
if h_peer_id != peer.peer_id {
debug!("peer id differ")
}
Ok(stream)
} else {
Err(Error::msg("unexpected message"))
Expand Down Expand Up @@ -325,15 +322,15 @@ pub async fn handle_peer(peer: PeerInfo, state: Arc<Mutex<State>>) -> Result<()>
{
debug!("connecting to peer: {:?}", peer);
let mut state = state.lock().await;
match state.peers.get_mut(&peer.peer_id) {
match state.peers.get_mut(&peer) {
Some(p) if p.status == PeerStatus::Connected => {
return Err(Error::msg("peer is already connected"))
}
Some(p) => p.status = PeerStatus::Connected,
None => {
let mut p = Peer::new(peer.clone());
p.status = PeerStatus::Connected;
state.peers.insert(peer.peer_id.clone(), p);
state.peers.insert(peer.clone(), p);
}
};
};
Expand All @@ -345,7 +342,7 @@ pub async fn handle_peer(peer: PeerInfo, state: Arc<Mutex<State>>) -> Result<()>
.lock()
.await
.peers
.get_mut(&peer.peer_id)
.get_mut(&peer)
.context("no peer")?
.status = if res.is_err() {
PeerStatus::Disconnected
Expand All @@ -362,7 +359,7 @@ pub async fn do_handle_peer(peer: PeerInfo, state: Arc<Mutex<State>>) -> Result<
.context("handshake error")?;
info!("successfull handshake with peer {:?}", peer);

if let Some(p) = state.lock().await.peers.get_mut(&peer.peer_id) {
if let Some(p) = state.lock().await.peers.get_mut(&peer) {
p.status = PeerStatus::Connected;
}

Expand Down Expand Up @@ -394,7 +391,7 @@ pub async fn write_loop(
{
let state = state.lock().await;
if state.config.respect_choke {
let p = state.peers.get(&peer.peer_id).cloned();
let p = state.peers.get(&peer).cloned();
if let Some(p) = p {
if p.choked {
info!("peer is choked, waiting");
Expand Down Expand Up @@ -442,13 +439,13 @@ async fn read_loop(
) -> Result<()> {
loop {
match read_message(&mut stream).await {
Ok(Message::Choke) => match state.lock().await.peers.get_mut(&peer.peer_id) {
Ok(Message::Choke) => match state.lock().await.peers.get_mut(&peer) {
Some(p) => p.choked = true,
_ => debug!("no peer with id {:?}", peer.peer_id),
_ => debug!("no peer {:?}", peer),
},
Ok(Message::Unchoke) => match state.lock().await.peers.get_mut(&peer.peer_id) {
Ok(Message::Unchoke) => match state.lock().await.peers.get_mut(&peer) {
Some(p) => p.choked = false,
_ => debug!("no peer with id {:?}", peer.peer_id),
_ => debug!("no peer {:?}", peer),
},
Ok(Message::Piece {
piece_index,
Expand Down
21 changes: 2 additions & 19 deletions src/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ pub struct State {
pub info_hash: Vec<u8>,
pub peer_id: Vec<u8>,
pub pieces: BTreeMap<u32, Piece>,
pub peers: BTreeMap<ByteString, Peer>,
pub peers: BTreeMap<PeerInfo, Peer>,
pub status: TorrentStatus,
}

Expand Down Expand Up @@ -114,29 +114,12 @@ pub enum PeerStatus {
Done,
}

#[derive(Clone, PartialEq, PartialOrd, Hash)]
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PeerInfo {
pub peer_id: ByteString,
pub ip: String,
pub port: i64,
}

impl fmt::Debug for PeerInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
f.debug_struct("TrackerPeer")
.field(
"peer_id",
&match String::from_utf8(self.peer_id.clone()) {
Ok(str) => str,
_ => "<non-utf>".into(),
},
)
.field("ip", &self.ip)
.field("port", &self.port)
.finish()
}
}

pub fn init_pieces(info: &Info) -> BTreeMap<u32, Piece> {
let total_len = info.file_info.total_length() as u32;
if info.pieces.len() != (total_len as f64 / info.piece_length as f64).ceil() as usize {
Expand Down
2 changes: 1 addition & 1 deletion src/torrent.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ pub async fn download_torrent(path: &Path, peer_id: &ByteString, config: &Config
peers: resp
.peers
.into_iter()
.map(|p| (p.peer_id.clone(), Peer::new(p)))
.map(|p| (p.clone(), Peer::new(p)))
.collect(),
status: TorrentStatus::Started,
}));
Expand Down
8 changes: 2 additions & 6 deletions src/tracker.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,6 @@ impl TryFrom<BencodeValue> for TrackerResponse {
.iter()
.map(|p| match p {
BencodeValue::Dict(p_dict) => Ok(PeerInfo {
peer_id: match p_dict.get("peer id") {
Some(BencodeValue::String(v)) => v.clone(),
_ => return Err("'peer id' missing".into()),
},
ip: match p_dict.get("ip") {
Some(BencodeValue::String(i)) => {
String::from_utf8(i.clone()).map_err(|e| e.to_string())?
Expand Down Expand Up @@ -256,12 +252,12 @@ pub async fn tracker_loop(state: Arc<Mutex<State>>) {
let new_peers: Vec<_> = resp
.peers
.into_iter()
.filter(|p| !state.peers.contains_key(&p.peer_id))
.filter(|p| !state.peers.contains_key(p))
.map(Peer::new)
.collect();
info!("received {} new peers", new_peers.len());
for p in new_peers {
state.peers.insert(p.info.peer_id.clone(), p);
state.peers.insert(p.info.clone(), p);
}
info!(
"total {} peers, {} connected",
Expand Down

0 comments on commit 51b74cd

Please sign in to comment.