1use std::{
29 io::{BufRead, BufReader, Read, Write},
30 net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs},
31 sync::{
32 atomic::{AtomicBool, Ordering},
33 Arc,
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 {
91 Self { bytes }
92 }
93}
94
95impl ByteSource for BytesSource {
96 fn len(&self) -> Option<u64> {
97 Some(self.bytes.len() as u64)
98 }
99
100 fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<usize> {
101 let offset = offset.min(self.bytes.len() as u64) as usize;
102 let available = &self.bytes[offset..];
103 let n = available.len().min(buf.len());
104 buf[..n].copy_from_slice(&available[..n]);
105 Ok(n)
106 }
107}
108
109pub type SourceResolver = Arc<dyn Fn(&str) -> Option<Arc<dyn ByteSource>> + Send + Sync>;
111
112pub struct PeerServer {
114 addr: SocketAddr,
115 running: Arc<AtomicBool>,
116}
117
118impl PeerServer {
119 pub fn start(
122 bind_addr: impl ToSocketAddrs,
123 token: impl Into<String>,
124 resolver: SourceResolver,
125 ) -> Result<PeerServer, PeerError> {
126 let listener = TcpListener::bind(bind_addr)?;
127 let addr = listener.local_addr()?;
128 let running = Arc::new(AtomicBool::new(true));
129 let token = token.into();
130
131 let loop_running = running.clone();
132 std::thread::Builder::new()
133 .name("cranpose-peer".to_string())
134 .spawn(move || {
135 for stream in listener.incoming() {
136 if !loop_running.load(Ordering::SeqCst) {
137 break;
138 }
139 let Ok(stream) = stream else { continue };
140 let token = token.clone();
141 let resolver = resolver.clone();
142 let _ = std::thread::Builder::new()
145 .name("cranpose-peer-conn".to_string())
146 .spawn(move || {
147 let _ = handle_connection(stream, &token, &resolver);
148 });
149 }
150 })
151 .map_err(|error| PeerError::Io(error.to_string()))?;
152
153 Ok(PeerServer { addr, running })
154 }
155
156 pub fn local_addr(&self) -> SocketAddr {
158 self.addr
159 }
160
161 pub fn port(&self) -> u16 {
163 self.addr.port()
164 }
165}
166
167impl Drop for PeerServer {
168 fn drop(&mut self) {
169 self.running.store(false, Ordering::SeqCst);
170 let _ = TcpStream::connect(self.addr);
172 }
173}
174
175fn handle_connection(
176 mut stream: TcpStream,
177 token: &str,
178 resolver: &SourceResolver,
179) -> Result<(), PeerError> {
180 stream.set_read_timeout(Some(Duration::from_secs(30)))?;
181 let mut reader = BufReader::new(stream.try_clone()?);
182
183 let mut request_line = String::new();
184 if reader.read_line(&mut request_line)? == 0 {
185 return Ok(()); }
187 let mut parts = request_line.split_whitespace();
188 let method = parts.next().unwrap_or("");
189 let path = parts.next().unwrap_or("");
190
191 let mut authorization = None;
192 let mut range = None;
193 loop {
194 let mut line = String::new();
195 if reader.read_line(&mut line)? == 0 {
196 break;
197 }
198 let line = line.trim_end();
199 if line.is_empty() {
200 break;
201 }
202 if let Some((name, value)) = line.split_once(':') {
203 let value = value.trim();
204 match name.trim().to_ascii_lowercase().as_str() {
205 "authorization" => authorization = Some(value.to_string()),
206 "range" => range = parse_range_header(value),
207 _ => {}
208 }
209 }
210 }
211
212 if method != "GET" {
213 return write_status(&mut stream, 405, "Method Not Allowed");
214 }
215 if authorization.as_deref() != Some(&format!("Bearer {token}")) {
216 return write_status(&mut stream, 401, "Unauthorized");
217 }
218 let Some(handle) = path.strip_prefix("/track/") else {
219 return write_status(&mut stream, 404, "Not Found");
220 };
221 let handle = crate::content::percent_decode_lossy(handle);
222 let Some(source) = resolver(&handle) else {
223 return write_status(&mut stream, 404, "Not Found");
224 };
225
226 serve_source(&mut stream, source.as_ref(), range)
227}
228
229fn serve_source(
230 stream: &mut TcpStream,
231 source: &dyn ByteSource,
232 range: Option<(u64, Option<u64>)>,
233) -> Result<(), PeerError> {
234 let total = source.len();
235
236 let (status, reason, start, length) = match (range, total) {
237 (Some((start, end)), Some(total)) if start < total => {
238 let last = end.unwrap_or(total - 1).min(total - 1);
239 if last < start {
240 return write_status(stream, 416, "Range Not Satisfiable");
241 }
242 (206, "Partial Content", start, last - start + 1)
243 }
244 (Some((start, _)), Some(total)) if start >= total => {
245 return write_status(stream, 416, "Range Not Satisfiable");
246 }
247 (_, Some(total)) => (200, "OK", 0, total),
248 (Some(_), None) => return write_status(stream, 416, "Range Not Satisfiable"),
250 (None, None) => {
251 return serve_unknown_length(stream, source);
253 }
254 };
255
256 let mut header = format!(
257 "HTTP/1.1 {status} {reason}\r\nContent-Length: {length}\r\nAccept-Ranges: bytes\r\nContent-Type: application/octet-stream\r\nConnection: close\r\n"
258 );
259 if status == 206 {
260 if let Some(total) = total {
261 let end = start + length - 1;
262 header.push_str(&format!("Content-Range: bytes {start}-{end}/{total}\r\n"));
263 }
264 }
265 header.push_str("\r\n");
266 stream.write_all(header.as_bytes())?;
267
268 stream_bytes(stream, source, start, length)
269}
270
271fn serve_unknown_length(stream: &mut TcpStream, source: &dyn ByteSource) -> Result<(), PeerError> {
272 let header =
274 "HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nConnection: close\r\n\r\n";
275 stream.write_all(header.as_bytes())?;
276 let mut buf = vec![0u8; 64 * 1024];
277 let mut offset = 0u64;
278 loop {
279 let n = source.read_at(offset, &mut buf)?;
280 if n == 0 {
281 break;
282 }
283 stream.write_all(&buf[..n])?;
284 offset += n as u64;
285 }
286 Ok(())
287}
288
289fn stream_bytes(
290 stream: &mut TcpStream,
291 source: &dyn ByteSource,
292 start: u64,
293 length: u64,
294) -> Result<(), PeerError> {
295 let mut buf = vec![0u8; 64 * 1024];
296 let mut sent = 0u64;
297 while sent < length {
298 let want = ((length - sent) as usize).min(buf.len());
299 let n = source.read_at(start + sent, &mut buf[..want])?;
300 if n == 0 {
301 break;
302 }
303 stream.write_all(&buf[..n])?;
304 sent += n as u64;
305 }
306 Ok(())
307}
308
309fn write_status(stream: &mut TcpStream, code: u16, reason: &str) -> Result<(), PeerError> {
310 let response =
311 format!("HTTP/1.1 {code} {reason}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
312 stream.write_all(response.as_bytes())?;
313 Ok(())
314}
315
316fn parse_range_header(value: &str) -> Option<(u64, Option<u64>)> {
318 let spec = value.trim().strip_prefix("bytes=")?;
319 let (start, end) = spec.split_once('-')?;
320 let start = start.trim().parse::<u64>().ok()?;
321 let end = end.trim();
322 let end = if end.is_empty() {
323 None
324 } else {
325 Some(end.parse::<u64>().ok()?)
326 };
327 Some((start, end))
328}
329
330pub struct FetchResult {
336 pub total_len: Option<u64>,
338 pub bytes: Vec<u8>,
340}
341
342struct ResponseHead {
343 total_len: Option<u64>,
344 content_length: Option<u64>,
345 reader: BufReader<TcpStream>,
346}
347
348fn open_request(
352 base: &str,
353 token: &str,
354 handle: &str,
355 start: u64,
356 len: Option<u64>,
357) -> Result<ResponseHead, PeerError> {
358 let mut stream = TcpStream::connect(base)?;
359 stream.set_read_timeout(Some(Duration::from_secs(30)))?;
360
361 let range = match len {
362 Some(len) if len > 0 => format!("bytes={start}-{}", start + len - 1),
363 Some(_) => format!("bytes={start}-{start}"),
364 None => format!("bytes={start}-"),
365 };
366 let request = format!(
367 "GET /track/{} HTTP/1.1\r\nHost: {base}\r\nAuthorization: Bearer {token}\r\nRange: {range}\r\nConnection: close\r\n\r\n",
368 encode_handle(handle)
369 );
370 stream.write_all(request.as_bytes())?;
371
372 let mut reader = BufReader::new(stream);
373 let mut status_line = String::new();
374 reader.read_line(&mut status_line)?;
375 let status = parse_status(&status_line)?;
376
377 let mut total_len = None;
378 let mut content_length = None;
379 loop {
380 let mut line = String::new();
381 if reader.read_line(&mut line)? == 0 {
382 break;
383 }
384 let line = line.trim_end();
385 if line.is_empty() {
386 break;
387 }
388 if let Some((name, value)) = line.split_once(':') {
389 match name.trim().to_ascii_lowercase().as_str() {
390 "content-length" => content_length = value.trim().parse::<u64>().ok(),
391 "content-range" => total_len = parse_content_range_total(value.trim()),
392 _ => {}
393 }
394 }
395 }
396
397 match status {
398 401 => Err(PeerError::Unauthorized),
399 404 => Err(PeerError::NotFound),
400 200 | 206 => Ok(ResponseHead {
401 total_len,
402 content_length,
403 reader,
404 }),
405 other => Err(PeerError::Status(other)),
406 }
407}
408
409pub fn fetch_range(
412 base: &str,
413 token: &str,
414 handle: &str,
415 start: u64,
416 len: Option<u64>,
417) -> Result<FetchResult, PeerError> {
418 let mut head = open_request(base, token, handle, start, len)?;
419 let mut bytes = Vec::new();
420 match head.content_length {
421 Some(length) => {
422 bytes.resize(length as usize, 0);
423 head.reader.read_exact(&mut bytes)?;
424 }
425 None => {
426 head.reader.read_to_end(&mut bytes)?;
427 }
428 }
429 Ok(FetchResult {
430 total_len: head.total_len,
431 bytes,
432 })
433}
434
435pub fn fetch_to_writer(
439 base: &str,
440 token: &str,
441 handle: &str,
442 start: u64,
443 len: Option<u64>,
444 writer: &mut dyn Write,
445) -> Result<Option<u64>, PeerError> {
446 let mut head = open_request(base, token, handle, start, len)?;
447 let mut buf = vec![0u8; 64 * 1024];
448 let mut remaining = head.content_length;
449 loop {
450 let want = match remaining {
451 Some(0) => break,
452 Some(r) => (r as usize).min(buf.len()),
453 None => buf.len(),
454 };
455 let n = head.reader.read(&mut buf[..want])?;
456 if n == 0 {
457 break;
458 }
459 writer.write_all(&buf[..n])?;
460 if let Some(r) = remaining.as_mut() {
461 *r -= n as u64;
462 }
463 }
464 Ok(head.total_len)
465}
466
467pub fn content_length(base: &str, token: &str, handle: &str) -> Result<Option<u64>, PeerError> {
469 Ok(fetch_range(base, token, handle, 0, Some(1))?.total_len)
470}
471
472fn parse_status(line: &str) -> Result<u16, PeerError> {
473 line.split_whitespace()
474 .nth(1)
475 .and_then(|code| code.parse::<u16>().ok())
476 .ok_or_else(|| PeerError::Protocol(format!("bad status line: {line:?}")))
477}
478
479fn parse_content_range_total(value: &str) -> Option<u64> {
480 value.rsplit('/').next()?.trim().parse::<u64>().ok()
482}
483
484fn encode_handle(handle: &str) -> String {
485 let mut out = String::with_capacity(handle.len());
486 for byte in handle.as_bytes() {
487 match byte {
488 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
489 out.push(*byte as char)
490 }
491 other => out.push_str(&format!("%{other:02X}")),
492 }
493 }
494 out
495}
496
497#[cfg(test)]
498mod tests {
499 use super::*;
500
501 fn resolver_for(handle: &'static str, bytes: Vec<u8>) -> SourceResolver {
502 Arc::new(move |requested: &str| {
503 if requested == handle {
504 Some(Arc::new(BytesSource::new(bytes.clone())) as Arc<dyn ByteSource>)
505 } else {
506 None
507 }
508 })
509 }
510
511 #[test]
512 fn round_trips_full_and_partial() {
513 let data: Vec<u8> = (0..=255u8).cycle().take(5000).collect();
514 let server = PeerServer::start("127.0.0.1:0", "secret", resolver_for("song", data.clone()))
515 .expect("start");
516 let base = format!("127.0.0.1:{}", server.port());
517
518 let full = fetch_range(&base, "secret", "song", 0, None).expect("full");
519 assert_eq!(full.bytes, data);
520 assert_eq!(full.total_len, Some(5000));
521
522 let part = fetch_range(&base, "secret", "song", 1000, Some(256)).expect("part");
523 assert_eq!(part.bytes, data[1000..1256]);
524 assert_eq!(part.total_len, Some(5000));
525
526 assert_eq!(content_length(&base, "secret", "song").unwrap(), Some(5000));
527 }
528
529 #[test]
530 fn streams_to_writer_without_buffering() {
531 let data: Vec<u8> = (0..2000u32).map(|i| i as u8).collect();
532 let server =
533 PeerServer::start("127.0.0.1:0", "k", resolver_for("s", data.clone())).expect("start");
534 let base = format!("127.0.0.1:{}", server.port());
535 let mut out = Vec::new();
536 let total = fetch_to_writer(&base, "k", "s", 0, None, &mut out).expect("stream");
537 assert_eq!(out, data);
538 assert_eq!(total, Some(2000));
539 }
540
541 #[test]
542 fn rejects_wrong_token() {
543 let server =
544 PeerServer::start("127.0.0.1:0", "right", resolver_for("a", vec![1, 2, 3])).expect("s");
545 let base = format!("127.0.0.1:{}", server.port());
546 assert!(matches!(
547 fetch_range(&base, "wrong", "a", 0, None),
548 Err(PeerError::Unauthorized)
549 ));
550 }
551
552 #[test]
553 fn unknown_handle_is_not_found() {
554 let server =
555 PeerServer::start("127.0.0.1:0", "t", resolver_for("a", vec![1, 2, 3])).expect("s");
556 let base = format!("127.0.0.1:{}", server.port());
557 assert!(matches!(
558 fetch_range(&base, "t", "missing", 0, None),
559 Err(PeerError::NotFound)
560 ));
561 }
562
563 #[test]
564 fn handle_is_percent_encoded_round_trip() {
565 let server =
566 PeerServer::start("127.0.0.1:0", "t", resolver_for("a b/c.mp3", vec![9, 8, 7]))
567 .expect("s");
568 let base = format!("127.0.0.1:{}", server.port());
569 let got = fetch_range(&base, "t", "a b/c.mp3", 0, None).expect("fetch");
570 assert_eq!(got.bytes, vec![9, 8, 7]);
571 }
572}