1use std::{
29 io::{BufRead, BufReader, Read, Write},
30 net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs},
31 sync::{
32 Arc,
33 atomic::{AtomicBool, Ordering},
34 },
35 time::Duration,
36};
37
38#[derive(thiserror::Error, Debug)]
40pub enum PeerError {
41 #[error("peer request was not authorized")]
43 Unauthorized,
44 #[error("peer handle not found")]
46 NotFound,
47 #[error("peer returned HTTP {0}")]
49 Status(u16),
50 #[error("peer protocol error: {0}")]
52 Protocol(String),
53 #[error("{0}")]
55 Io(String),
56}
57
58impl From<std::io::Error> for PeerError {
59 fn from(error: std::io::Error) -> Self {
60 PeerError::Io(error.to_string())
61 }
62}
63
64pub trait ByteSource: Send + Sync {
70 fn len(&self) -> Option<u64>;
72
73 fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<usize>;
76
77 fn is_empty(&self) -> bool {
79 self.len() == Some(0)
80 }
81}
82
83pub struct BytesSource {
86 bytes: Vec<u8>,
87}
88
89impl BytesSource {
90 pub fn new(bytes: Vec<u8>) -> Self {
92 Self { bytes }
93 }
94}
95
96impl ByteSource for BytesSource {
97 fn len(&self) -> Option<u64> {
98 Some(self.bytes.len() as u64)
99 }
100
101 fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<usize> {
102 let offset = offset.min(self.bytes.len() as u64) as usize;
103 let available = &self.bytes[offset..];
104 let n = available.len().min(buf.len());
105 buf[..n].copy_from_slice(&available[..n]);
106 Ok(n)
107 }
108}
109
110pub type SourceResolver = Arc<dyn Fn(&str) -> Option<Arc<dyn ByteSource>> + Send + Sync>;
112
113pub struct PeerServer {
115 addr: SocketAddr,
116 running: Arc<AtomicBool>,
117}
118
119impl PeerServer {
120 pub fn start(
123 bind_addr: impl ToSocketAddrs,
124 token: impl Into<String>,
125 resolver: SourceResolver,
126 ) -> Result<PeerServer, PeerError> {
127 let listener = TcpListener::bind(bind_addr)?;
128 let addr = listener.local_addr()?;
129 let running = Arc::new(AtomicBool::new(true));
130 let token = token.into();
131
132 let loop_running = running.clone();
133 std::thread::Builder::new()
134 .name("cranpose-peer".to_string())
135 .spawn(move || {
136 for stream in listener.incoming() {
137 if !loop_running.load(Ordering::SeqCst) {
138 break;
139 }
140 let Ok(stream) = stream else { continue };
141 let token = token.clone();
142 let resolver = resolver.clone();
143 let _ = std::thread::Builder::new()
144 .name("cranpose-peer-conn".to_string())
145 .spawn(move || {
146 let _ = handle_connection(stream, &token, &resolver);
147 });
148 }
149 })
150 .map_err(|error| PeerError::Io(error.to_string()))?;
151
152 Ok(PeerServer { addr, running })
153 }
154
155 pub fn local_addr(&self) -> SocketAddr {
157 self.addr
158 }
159
160 pub fn port(&self) -> u16 {
162 self.addr.port()
163 }
164}
165
166impl Drop for PeerServer {
167 fn drop(&mut self) {
168 self.running.store(false, Ordering::SeqCst);
169 let _ = TcpStream::connect(self.addr);
170 }
171}
172
173fn handle_connection(
174 mut stream: TcpStream,
175 token: &str,
176 resolver: &SourceResolver,
177) -> Result<(), PeerError> {
178 stream.set_read_timeout(Some(Duration::from_secs(30)))?;
179 let mut reader = BufReader::new(stream.try_clone()?);
180
181 let mut request_line = String::new();
182 if reader.read_line(&mut request_line)? == 0 {
183 return Ok(());
184 }
185 let mut parts = request_line.split_whitespace();
186 let method = parts.next().unwrap_or("");
187 let path = parts.next().unwrap_or("");
188
189 let mut authorization = None;
190 let mut range = None;
191 loop {
192 let mut line = String::new();
193 if reader.read_line(&mut line)? == 0 {
194 break;
195 }
196 let line = line.trim_end();
197 if line.is_empty() {
198 break;
199 }
200 if let Some((name, value)) = line.split_once(':') {
201 let value = value.trim();
202 match name.trim().to_ascii_lowercase().as_str() {
203 "authorization" => authorization = Some(value.to_string()),
204 "range" => range = parse_range_header(value),
205 _ => {}
206 }
207 }
208 }
209
210 if method != "GET" {
211 return write_status(&mut stream, 405, "Method Not Allowed");
212 }
213 if authorization.as_deref() != Some(&format!("Bearer {token}")) {
214 return write_status(&mut stream, 401, "Unauthorized");
215 }
216 let Some(handle) = path.strip_prefix("/track/") else {
217 return write_status(&mut stream, 404, "Not Found");
218 };
219 let handle = crate::content::percent_decode_lossy(handle);
220 let Some(source) = resolver(&handle) else {
221 return write_status(&mut stream, 404, "Not Found");
222 };
223
224 serve_source(&mut stream, source.as_ref(), range)
225}
226
227fn serve_source(
228 stream: &mut TcpStream,
229 source: &dyn ByteSource,
230 range: Option<(u64, Option<u64>)>,
231) -> Result<(), PeerError> {
232 let total = source.len();
233
234 let (status, reason, start, length) = match (range, total) {
235 (Some((start, end)), Some(total)) if start < total => {
236 let last = end.unwrap_or(total - 1).min(total - 1);
237 if last < start {
238 return write_status(stream, 416, "Range Not Satisfiable");
239 }
240 (206, "Partial Content", start, last - start + 1)
241 }
242 (Some((start, _)), Some(total)) if start >= total => {
243 return write_status(stream, 416, "Range Not Satisfiable");
244 }
245 (_, Some(total)) => (200, "OK", 0, total),
246 (Some(_), None) => return write_status(stream, 416, "Range Not Satisfiable"),
247 (None, None) => {
248 return serve_unknown_length(stream, source);
249 }
250 };
251
252 let mut header = format!(
253 "HTTP/1.1 {status} {reason}\r\nContent-Length: {length}\r\nAccept-Ranges: bytes\r\nContent-Type: application/octet-stream\r\nConnection: close\r\n"
254 );
255 if status == 206
256 && let Some(total) = total
257 {
258 let end = start + length - 1;
259 header.push_str(&format!("Content-Range: bytes {start}-{end}/{total}\r\n"));
260 }
261 header.push_str("\r\n");
262 stream.write_all(header.as_bytes())?;
263
264 stream_bytes(stream, source, start, length)
265}
266
267fn serve_unknown_length(stream: &mut TcpStream, source: &dyn ByteSource) -> Result<(), PeerError> {
268 let header =
269 "HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nConnection: close\r\n\r\n";
270 stream.write_all(header.as_bytes())?;
271 let mut buf = vec![0u8; 64 * 1024];
272 let mut offset = 0u64;
273 loop {
274 let n = source.read_at(offset, &mut buf)?;
275 if n == 0 {
276 break;
277 }
278 stream.write_all(&buf[..n])?;
279 offset += n as u64;
280 }
281 Ok(())
282}
283
284fn stream_bytes(
285 stream: &mut TcpStream,
286 source: &dyn ByteSource,
287 start: u64,
288 length: u64,
289) -> Result<(), PeerError> {
290 let mut buf = vec![0u8; 64 * 1024];
291 let mut sent = 0u64;
292 while sent < length {
293 let want = ((length - sent) as usize).min(buf.len());
294 let n = source.read_at(start + sent, &mut buf[..want])?;
295 if n == 0 {
296 break;
297 }
298 stream.write_all(&buf[..n])?;
299 sent += n as u64;
300 }
301 Ok(())
302}
303
304fn write_status(stream: &mut TcpStream, code: u16, reason: &str) -> Result<(), PeerError> {
305 let response =
306 format!("HTTP/1.1 {code} {reason}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
307 stream.write_all(response.as_bytes())?;
308 Ok(())
309}
310
311fn parse_range_header(value: &str) -> Option<(u64, Option<u64>)> {
312 let spec = value.trim().strip_prefix("bytes=")?;
313 let (start, end) = spec.split_once('-')?;
314 let start = start.trim().parse::<u64>().ok()?;
315 let end = end.trim();
316 let end = if end.is_empty() {
317 None
318 } else {
319 Some(end.parse::<u64>().ok()?)
320 };
321 Some((start, end))
322}
323
324pub struct FetchResult {
326 pub total_len: Option<u64>,
328 pub bytes: Vec<u8>,
330}
331
332struct ResponseHead {
333 total_len: Option<u64>,
334 content_length: Option<u64>,
335 reader: BufReader<TcpStream>,
336}
337
338fn open_request(
339 base: &str,
340 token: &str,
341 handle: &str,
342 start: u64,
343 len: Option<u64>,
344) -> Result<ResponseHead, PeerError> {
345 let mut stream = TcpStream::connect(base)?;
346 stream.set_read_timeout(Some(Duration::from_secs(30)))?;
347
348 let range = match len {
349 Some(len) if len > 0 => format!("bytes={start}-{}", start + len - 1),
350 Some(_) => format!("bytes={start}-{start}"),
351 None => format!("bytes={start}-"),
352 };
353 let request = format!(
354 "GET /track/{} HTTP/1.1\r\nHost: {base}\r\nAuthorization: Bearer {token}\r\nRange: {range}\r\nConnection: close\r\n\r\n",
355 encode_handle(handle)
356 );
357 stream.write_all(request.as_bytes())?;
358
359 let mut reader = BufReader::new(stream);
360 let mut status_line = String::new();
361 reader.read_line(&mut status_line)?;
362 let status = parse_status(&status_line)?;
363
364 let mut total_len = None;
365 let mut content_length = None;
366 loop {
367 let mut line = String::new();
368 if reader.read_line(&mut line)? == 0 {
369 break;
370 }
371 let line = line.trim_end();
372 if line.is_empty() {
373 break;
374 }
375 if let Some((name, value)) = line.split_once(':') {
376 match name.trim().to_ascii_lowercase().as_str() {
377 "content-length" => content_length = value.trim().parse::<u64>().ok(),
378 "content-range" => total_len = parse_content_range_total(value.trim()),
379 _ => {}
380 }
381 }
382 }
383
384 match status {
385 401 => Err(PeerError::Unauthorized),
386 404 => Err(PeerError::NotFound),
387 200 | 206 => Ok(ResponseHead {
388 total_len,
389 content_length,
390 reader,
391 }),
392 other => Err(PeerError::Status(other)),
393 }
394}
395
396pub fn fetch_range(
399 base: &str,
400 token: &str,
401 handle: &str,
402 start: u64,
403 len: Option<u64>,
404) -> Result<FetchResult, PeerError> {
405 let mut head = open_request(base, token, handle, start, len)?;
406 let mut bytes = Vec::new();
407 match head.content_length {
408 Some(length) => {
409 bytes.resize(length as usize, 0);
410 head.reader.read_exact(&mut bytes)?;
411 }
412 None => {
413 head.reader.read_to_end(&mut bytes)?;
414 }
415 }
416 Ok(FetchResult {
417 total_len: head.total_len,
418 bytes,
419 })
420}
421
422pub fn fetch_to_writer(
426 base: &str,
427 token: &str,
428 handle: &str,
429 start: u64,
430 len: Option<u64>,
431 writer: &mut dyn Write,
432) -> Result<Option<u64>, PeerError> {
433 let mut head = open_request(base, token, handle, start, len)?;
434 let mut buf = vec![0u8; 64 * 1024];
435 let mut remaining = head.content_length;
436 loop {
437 let want = match remaining {
438 Some(0) => break,
439 Some(r) => (r as usize).min(buf.len()),
440 None => buf.len(),
441 };
442 let n = head.reader.read(&mut buf[..want])?;
443 if n == 0 {
444 break;
445 }
446 writer.write_all(&buf[..n])?;
447 if let Some(r) = remaining.as_mut() {
448 *r -= n as u64;
449 }
450 }
451 Ok(head.total_len)
452}
453
454pub fn content_length(base: &str, token: &str, handle: &str) -> Result<Option<u64>, PeerError> {
456 Ok(fetch_range(base, token, handle, 0, Some(1))?.total_len)
457}
458
459fn parse_status(line: &str) -> Result<u16, PeerError> {
460 line.split_whitespace()
461 .nth(1)
462 .and_then(|code| code.parse::<u16>().ok())
463 .ok_or_else(|| PeerError::Protocol(format!("bad status line: {line:?}")))
464}
465
466fn parse_content_range_total(value: &str) -> Option<u64> {
467 value.rsplit('/').next()?.trim().parse::<u64>().ok()
468}
469
470fn encode_handle(handle: &str) -> String {
471 let mut out = String::with_capacity(handle.len());
472 for byte in handle.as_bytes() {
473 match byte {
474 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
475 out.push(*byte as char);
476 }
477 other => out.push_str(&format!("%{other:02X}")),
478 }
479 }
480 out
481}
482
483#[cfg(test)]
484#[path = "tests/peer_tests.rs"]
485mod tests;