Skip to main content

pfx_steam/
blob.rs

1use std::collections::BTreeMap;
2
3use crate::{Delivery, Error, SteamApi, User, checksum};
4
5pub const BLOB_CHUNK: usize = 32 * 1024;
6pub const BLOB_MAX: usize = 16 * 1024 * 1024;
7pub const BLOB_OPEN_MAX: usize = 4;
8const TAG: u8 = 0xb1;
9const HEADER: usize = 29;
10const BATCH: usize = 64;
11
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub enum Blob {
14    Complete { peer: User, id: u32, bytes: Vec<u8> },
15    Corrupt { peer: User, id: u32 },
16}
17
18#[derive(Clone, Debug)]
19struct Partial {
20    count: u32,
21    len: u64,
22    sum: u64,
23    chunks: BTreeMap<u32, Vec<u8>>,
24}
25
26#[derive(Clone, Debug)]
27pub struct Blobs {
28    channel: u32,
29    open: BTreeMap<(User, u32), Partial>,
30}
31
32fn chunks_for(len: usize) -> u32 {
33    len.div_ceil(BLOB_CHUNK).max(1) as u32
34}
35
36fn chunk_len(len: u64, count: u32, index: u32) -> u64 {
37    if index + 1 < count {
38        BLOB_CHUNK as u64
39    } else {
40        len - BLOB_CHUNK as u64 * u64::from(count - 1)
41    }
42}
43
44pub fn send_blob(
45    steam: &mut impl SteamApi,
46    peer: User,
47    channel: u32,
48    id: u32,
49    bytes: &[u8],
50) -> Result<u32, Error> {
51    if bytes.len() > BLOB_MAX {
52        return Err(Error::TooBig);
53    }
54    let count = chunks_for(bytes.len());
55    let sum = checksum(bytes);
56    for index in 0..count {
57        let start = index as usize * BLOB_CHUNK;
58        let end = (start + BLOB_CHUNK).min(bytes.len());
59        let mut packet = Vec::with_capacity(HEADER + end - start);
60        packet.push(TAG);
61        packet.extend_from_slice(&id.to_le_bytes());
62        packet.extend_from_slice(&index.to_le_bytes());
63        packet.extend_from_slice(&count.to_le_bytes());
64        packet.extend_from_slice(&(bytes.len() as u64).to_le_bytes());
65        packet.extend_from_slice(&sum.to_le_bytes());
66        packet.extend_from_slice(&bytes[start..end]);
67        steam.send_message(peer, channel, &packet, Delivery::Reliable)?;
68    }
69    Ok(count)
70}
71
72fn word<const N: usize>(bytes: &[u8], at: usize) -> [u8; N] {
73    let mut out = [0u8; N];
74    out.copy_from_slice(&bytes[at..at + N]);
75    out
76}
77
78impl Blobs {
79    pub fn new(channel: u32) -> Self {
80        Self {
81            channel,
82            open: BTreeMap::new(),
83        }
84    }
85
86    pub fn pending(&self) -> usize {
87        self.open.len()
88    }
89
90    pub fn receive(&mut self, steam: &mut impl SteamApi) -> Result<Vec<Blob>, Error> {
91        let mut done = Vec::new();
92        loop {
93            let batch = steam.receive_messages(self.channel, BATCH)?;
94            if batch.is_empty() {
95                return Ok(done);
96            }
97            for message in batch {
98                done.extend(self.accept(message.peer, &message.bytes));
99            }
100        }
101    }
102
103    pub fn accept(&mut self, peer: User, bytes: &[u8]) -> Option<Blob> {
104        if bytes.len() < HEADER || bytes[0] != TAG {
105            return None;
106        }
107        let id = u32::from_le_bytes(word(bytes, 1));
108        let index = u32::from_le_bytes(word(bytes, 5));
109        let count = u32::from_le_bytes(word(bytes, 9));
110        let len = u64::from_le_bytes(word(bytes, 13));
111        let sum = u64::from_le_bytes(word(bytes, 21));
112        let payload = &bytes[HEADER..];
113        let key = (peer, id);
114        let shaped = len <= BLOB_MAX as u64
115            && count == chunks_for(len as usize)
116            && index < count
117            && payload.len() as u64 == chunk_len(len, count, index);
118        if !shaped {
119            return self.open.remove(&key).map(|_| Blob::Corrupt { peer, id });
120        }
121        if !self.open.contains_key(&key) {
122            let mine = self.open.keys().filter(|(who, _)| *who == peer).count();
123            if mine >= BLOB_OPEN_MAX {
124                return None;
125            }
126            self.open.insert(
127                key,
128                Partial {
129                    count,
130                    len,
131                    sum,
132                    chunks: BTreeMap::new(),
133                },
134            );
135        }
136        let partial = self.open.get_mut(&key)?;
137        if (partial.count, partial.len, partial.sum) != (count, len, sum) {
138            self.open.remove(&key);
139            return Some(Blob::Corrupt { peer, id });
140        }
141        partial
142            .chunks
143            .entry(index)
144            .or_insert_with(|| payload.to_vec());
145        if partial.chunks.len() < count as usize {
146            return None;
147        }
148        let partial = self.open.remove(&key)?;
149        let whole: Vec<u8> = partial.chunks.into_values().flatten().collect();
150        if checksum(&whole) == partial.sum {
151            Some(Blob::Complete {
152                peer,
153                id,
154                bytes: whole,
155            })
156        } else {
157            Some(Blob::Corrupt { peer, id })
158        }
159    }
160}