Optionally expose per-peer bitfields in PeerStats

crates/librqbit/src/lib.rs

@@ -99,6 +99,7 @@ pub use stream_connect::ConnectionOptions;
 pub use torrent_state::{
     ManagedTorrent, ManagedTorrentShared, ManagedTorrentState, TorrentMetadata, TorrentStats,
     TorrentStatsState,
+    live::peer::stats::snapshot::{PeerStats, PeerStatsFilter, PeerStatsFilterState},
 };
 pub use type_aliases::FileInfos;
 

crates/librqbit/src/torrent_state/live/mod.rs

@@ -104,7 +104,7 @@ use self::{
         PeerRx, PeerState, PeerTx, RemoveInflightRequestResult,
         stats::{
             atomic::PeerCountersAtomic as AtomicPeerCounters,
-            snapshot::{PeerStatsFilter, PeerStatsSnapshot},
+            snapshot::{PeerStats, PeerStatsFilter, PeerStatsSnapshot},
         },
     },
     peers::PeerStates,
@@ -122,6 +122,19 @@ fn make_piece_bitfield(lengths: &Lengths) -> BF {
     BF::from_boxed_slice(vec![0; lengths.piece_bitfield_bytes()].into_boxed_slice())
 }
 
+/// Builds a peer's bitfield from the bytes it sent, clearing the padding.
+///
+/// A bitfield is byte-padded, so the bits past `total_pieces` are spare. The
+/// spec says a peer zeroes them, but only the byte length is validated, so a
+/// peer can set them. Clearing them here rather than working around them at
+/// each use keeps `count_ones()` meaning "pieces this peer has", and stops a
+/// peer claiming a piece that does not exist.
+fn make_peer_bitfield(bytes: &[u8], lengths: &Lengths) -> BF {
+    let mut bf = BF::from_boxed_slice(bytes.to_vec().into_boxed_slice());
+    bf[lengths.total_pieces() as usize..].fill(false);
+    bf
+}
+
 pub(crate) struct TorrentStateLocked {
     // Coordinates piece state: what chunks we have, need, and what pieces are in-flight.
     // If this is None, the torrent was paused, and this live state is useless, and needs to be dropped.
@@ -752,7 +765,12 @@ impl TorrentStateLive {
                 .states
                 .iter()
                 .filter(|e| filter.state.matches(e.value().get_state()))
-                .map(|e| (e.key().to_string(), e.value().into()))
+                .map(|e| {
+                    (
+                        e.key().to_string(),
+                        PeerStats::from_peer(e.value(), filter.include_bitfield),
+                    )
+                })
                 .collect(),
         }
     }
@@ -1529,19 +1547,26 @@ impl PeerHandler {
                 if live.bitfield.is_empty() {
                     live.bitfield = make_piece_bitfield(&self.state.lengths);
                 }
-                match live.bitfield.get_mut(have as usize) {
-                    Some(mut v) => *v = true,
-                    None => {
-                        warn!(
-                            id = self.state.shared.id,
-                            info_hash = ?self.state.shared.info_hash,
-                            addr = ?self.addr,
-                            "received have {} out of range",
-                            have
-                        );
-                        return;
-                    }
-                };
+                // Bound by the piece count, not the bitfield length: the
+                // bitfield is byte-padded, so its length overshoots by up to 7
+                // and `get_mut` alone would let a peer set a padding bit.
+                let updated = have < self.state.lengths.total_pieces()
+                    && live
+                        .bitfield
+                        .get_mut(have as usize)
+                        .map(|mut v| *v = true)
+                        .is_some();
+                if !updated {
+                    warn!(
+                        id = self.state.shared.id,
+                        info_hash = ?self.state.shared.info_hash,
+                        addr = ?self.addr,
+                        "received have {} out of range",
+                        have
+                    );
+                    return;
+                }
+
                 trace!("updated bitfield with have={}", have);
                 if let Some(true) = live
                     .bitfield
@@ -1562,7 +1587,7 @@ impl PeerHandler {
                 self.state.lengths.piece_bitfield_bytes(),
             );
         }
-        let bf = BF::from_boxed_slice(bitfield.0.to_vec().into_boxed_slice());
+        let bf = make_peer_bitfield(bitfield.as_ref(), &self.state.lengths);
         if let Some(true) = bf
             .get(..self.state.lengths.total_pieces() as usize)
             .map(|s| s.all())
@@ -2092,3 +2117,58 @@ fn format_peer_client_name(value: &ByteBuf<'_>) -> Option<String> {
 
     Some(client_name)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::make_peer_bitfield;
+    use librqbit_core::lengths::Lengths;
+
+    /// 10 pieces needs 2 bytes, leaving 6 spare bits.
+    fn ten_pieces() -> Lengths {
+        let l = Lengths::new(10_000, 1_000).unwrap();
+        assert_eq!(l.total_pieces(), 10);
+        assert_eq!(l.piece_bitfield_bytes(), 2);
+        l
+    }
+
+    /// A peer that sets the spare trailing bits must not appear to hold more
+    /// than it does. Only the byte length is validated on ingest, so this is
+    /// reachable from the wire even though the spec says those bits are zero.
+    #[test]
+    fn padding_bits_are_cleared_on_ingest() {
+        // Claim 3 real pieces, then set every spare bit.
+        let bf = make_peer_bitfield(&[0b1110_0000, 0b0011_1111], &ten_pieces());
+        assert_eq!(bf.count_ones(), 3);
+    }
+
+    #[test]
+    fn a_seed_holds_every_piece() {
+        let bf = make_peer_bitfield(&[0xff, 0xff], &ten_pieces());
+        assert_eq!(bf.count_ones(), 10);
+    }
+
+    #[test]
+    fn an_empty_bitfield_holds_nothing() {
+        let bf = make_peer_bitfield(&[0x00, 0x00], &ten_pieces());
+        assert_eq!(bf.count_ones(), 0);
+    }
+
+    /// Real pieces are kept exactly as sent — the mask must not reach back
+    /// into the last byte's meaningful bits.
+    #[test]
+    fn pieces_inside_the_last_byte_survive() {
+        let bf = make_peer_bitfield(&[0b0000_0000, 0b1100_0000], &ten_pieces());
+        assert_eq!(bf.count_ones(), 2);
+        assert!(bf[8]);
+        assert!(bf[9]);
+    }
+
+    /// No padding to clear when the piece count is a multiple of 8.
+    #[test]
+    fn a_whole_number_of_bytes_is_untouched() {
+        let lengths = Lengths::new(8_000, 1_000).unwrap();
+        assert_eq!(lengths.total_pieces(), 8);
+        let bf = make_peer_bitfield(&[0xff], &lengths);
+        assert_eq!(bf.count_ones(), 8);
+    }
+}

crates/librqbit/src/torrent_state/live/peer/stats/snapshot.rs

@@ -29,6 +29,31 @@ pub struct PeerStats {
     pub state: &'static str,
     pub conn_kind: Option<ConnectionKind>,
     pub client_name: Option<String>,
+    /// How many pieces this peer has.
+    ///
+    /// `None` unless [`PeerStatsFilter::include_bitfield`] was set and the peer
+    /// is live. Provided alongside the bitfield to save every caller the same
+    /// `count_ones()`.
+    ///
+    /// Distinct from `counters.downloaded_and_checked_pieces`, which is how
+    /// many pieces this peer has sent *us*.
+    pub have_pieces: Option<u32>,
+    /// This peer's bitfield.
+    ///
+    /// `None` unless [`PeerStatsFilter::include_bitfield`] was set and the peer
+    /// is live. The trailing padding is cleared on ingest, so the bits are
+    /// exactly the pieces the peer holds.
+    ///
+    /// Raw bytes rather than the internal `BF`, because this struct is
+    /// `Serialize` and `bitvec` is built without its `serde` feature — and
+    /// naming `BitBox` here would put `bitvec` in the public API, where a
+    /// dependent would need a matching version to read the field.
+    ///
+    /// A count is not a substitute: computing piece availability across a
+    /// swarm — the rarest-piece copy count — needs to know *which* pieces each
+    /// peer holds. Two peers with 500 pieces each may overlap entirely or not
+    /// at all, and only one of those torrents can finish.
+    pub have_bitfield: Option<Vec<u8>>,
 }
 
 impl From<&super::atomic::PeerCountersAtomic> for PeerCounters {
@@ -54,8 +79,9 @@ impl From<&super::atomic::PeerCountersAtomic> for PeerCounters {
     }
 }
 
-impl From<&Peer> for PeerStats {
-    fn from(peer: &Peer) -> Self {
+impl PeerStats {
+    /// Builds a snapshot for one peer.
+    pub(crate) fn from_peer(peer: &Peer, include_bitfield: bool) -> Self {
         let state = peer.get_state();
         Self {
             counters: peer.stats.counters.as_ref().into(),
@@ -68,6 +94,18 @@ impl From<&Peer> for PeerStats {
                 PeerState::Live(l) => l.client_name.clone(),
                 _ => None,
             },
+            have_pieces: match state {
+                // The bitfield is sized from total_pieces, so the count fits a
+                // u32 by construction.
+                PeerState::Live(l) if include_bitfield => {
+                    Some(u32::try_from(l.bitfield.count_ones()).unwrap_or(u32::MAX))
+                }
+                _ => None,
+            },
+            have_bitfield: match state {
+                PeerState::Live(l) if include_bitfield => Some(l.bitfield.as_raw_slice().to_vec()),
+                _ => None,
+            },
         }
     }
 }
@@ -96,4 +134,22 @@ impl PeerStatsFilterState {
 pub struct PeerStatsFilter {
     #[serde(default)]
     pub state: PeerStatsFilterState,
+    /// Populate `have_pieces` and `have_bitfield`.
+    ///
+    /// Off by default, and nothing is computed unless it is set: the bitfield
+    /// costs a byte per eight pieces per peer to copy, and counting its bits
+    /// is a scan the existing callers have no use for.
+    #[serde(default)]
+    pub include_bitfield: bool,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::PeerStatsFilter;
+
+    /// The default filter computes nothing, so existing callers pay nothing.
+    #[test]
+    fn bitfields_are_off_by_default() {
+        assert!(!PeerStatsFilter::default().include_bitfield);
+    }
 }