1use std::io;
21
22#[cfg(target_os = "linux")]
23use std::os::unix::io::AsRawFd;
24#[cfg(target_os = "linux")]
25use tokio::io::{AsyncReadExt, AsyncWriteExt};
26#[cfg(target_os = "linux")]
27use tracing::debug;
28
29#[cfg(target_os = "linux")]
30use crate::util::zero_copy::splice_transfer;
31
32#[cfg(target_os = "linux")]
36const MAX_HEADER_SIZE: usize = 8 * 1024;
37
38#[cfg(target_os = "linux")]
64pub async fn try_splice_download(
65 url: &str,
66 offset: u64,
67 length: u64,
68 file: &std::fs::File,
69 file_offset: u64,
70) -> io::Result<u64> {
71 if length == 0 {
72 return Ok(0);
73 }
74
75 let parsed = url::Url::parse(url)
77 .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, format!("invalid URL: {e}")))?;
78
79 if parsed.scheme() != "http" {
80 return Err(io::Error::new(
81 io::ErrorKind::Unsupported,
82 "splice requires plain HTTP (no HTTPS)",
83 ));
84 }
85
86 let host = parsed
87 .host_str()
88 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "URL has no host"))?;
89 let port = parsed.port_or_known_default().unwrap_or(80);
90 let path = if parsed.path().is_empty() {
91 "/"
92 } else {
93 parsed.path()
94 };
95 let query = parsed.query().map(|q| format!("?{q}")).unwrap_or_default();
96 let path_query = format!("{path}{query}");
97
98 let addr = tokio::net::lookup_host((host, port))
100 .await
101 .map_err(|e| {
102 io::Error::new(
103 io::ErrorKind::AddrNotAvailable,
104 format!("DNS resolution failed: {e}"),
105 )
106 })?
107 .next()
108 .ok_or_else(|| io::Error::new(io::ErrorKind::AddrNotAvailable, "no addresses resolved"))?;
109
110 debug!(host, port, %addr, "splice_download: connecting");
111
112 let mut stream = tokio::net::TcpStream::connect(addr).await.map_err(|e| {
114 io::Error::new(
115 io::ErrorKind::ConnectionRefused,
116 format!("TCP connect failed: {e}"),
117 )
118 })?;
119 let _ = stream.set_nodelay(true);
122
123 let host_header = if port == 80 {
125 host.to_string()
126 } else {
127 format!("{host}:{port}")
128 };
129 let range_value = format!("bytes={}-{}", offset, offset + length.saturating_sub(1));
130 let request = format!(
131 "GET {path_query} HTTP/1.1\r\n\
132 Host: {host_header}\r\n\
133 Range: {range_value}\r\n\
134 Connection: close\r\n\
135 User-Agent: aria2-rust/1.0\r\n\
136 Accept: */*\r\n\
137 \r\n",
138 );
139 stream.write_all(request.as_bytes()).await?;
140 debug!(range = %range_value, "splice_download: request sent");
141
142 let mut header_buf = vec![0u8; MAX_HEADER_SIZE];
144 let mut header_len = 0usize;
145 let header_end_pos = loop {
146 if header_len >= MAX_HEADER_SIZE {
147 return Err(io::Error::new(
148 io::ErrorKind::InvalidData,
149 "response headers exceed 8 KB",
150 ));
151 }
152 let n = stream.read(&mut header_buf[header_len..]).await?;
153 if n == 0 {
154 return Err(io::Error::new(
155 io::ErrorKind::UnexpectedEof,
156 "connection closed before headers complete",
157 ));
158 }
159 header_len += n;
160 if let Some(pos) = find_header_end(&header_buf[..header_len]) {
161 break pos;
162 }
163 };
164
165 let header_bytes = &header_buf[..header_end_pos];
168 let header_str = std::str::from_utf8(header_bytes).map_err(|e| {
169 io::Error::new(io::ErrorKind::InvalidData, format!("non-UTF8 headers: {e}"))
170 })?;
171 let status = parse_status_code(header_str)?;
172 if status != 206 {
173 return Err(io::Error::other(format!(
174 "expected 206 Partial Content, got {status}"
175 )));
176 }
177
178 if is_chunked(header_str) {
180 return Err(io::Error::new(
181 io::ErrorKind::Unsupported,
182 "chunked transfer encoding not supported by splice",
183 ));
184 }
185 let content_length = parse_content_length(header_str)?
186 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing Content-Length"))?;
187
188 if content_length > length {
189 return Err(io::Error::new(
190 io::ErrorKind::InvalidData,
191 format!("Content-Length ({content_length}) exceeds requested length ({length})"),
192 ));
193 }
194
195 let body_start = header_end_pos + 4; let pre_read = &header_buf[body_start..header_len];
199 let pre_read_len = pre_read.len() as u64;
200
201 if pre_read_len > content_length {
202 return Err(io::Error::new(
203 io::ErrorKind::InvalidData,
204 format!(
205 "pre-read body bytes ({pre_read_len}) exceed Content-Length ({content_length})"
206 ),
207 ));
208 }
209
210 if pre_read_len > 0 {
211 write_all_at_offset(file, pre_read, file_offset)?;
212 }
213
214 let mut written = pre_read_len;
215 let remaining = content_length - pre_read_len;
216
217 if remaining > 0 {
223 let socket_fd = stream.as_raw_fd();
224 let file_fd = file.as_raw_fd();
225 let splice_file_offset = (file_offset + pre_read_len) as i64;
226 let splice_len = remaining as usize;
227
228 let splice_result = tokio::task::spawn_blocking(move || {
229 set_blocking(socket_fd);
234 splice_transfer(
235 socket_fd,
236 None,
237 file_fd,
238 Some(splice_file_offset),
239 splice_len,
240 )
241 })
242 .await
243 .map_err(|e| io::Error::other(format!("blocking task failed: {e}")))?
244 .map_err(|e| io::Error::other(format!("splice failed: {e}")))?;
245
246 let spliced = splice_result as u64;
247 if spliced < remaining {
248 return Err(io::Error::new(
249 io::ErrorKind::UnexpectedEof,
250 format!(
251 "splice EOF: transferred {} of {} body bytes (pre_read={}, spliced={})",
252 written + spliced,
253 content_length,
254 pre_read_len,
255 spliced
256 ),
257 ));
258 }
259 written += spliced;
260 }
261
262 debug!(
263 url,
264 bytes = written,
265 pre_read = pre_read_len,
266 spliced = written - pre_read_len,
267 "splice_download: complete"
268 );
269
270 Ok(written)
272}
273
274#[cfg(not(target_os = "linux"))]
276pub async fn try_splice_download(
277 _url: &str,
278 _offset: u64,
279 _length: u64,
280 _file: &std::fs::File,
281 _file_offset: u64,
282) -> io::Result<u64> {
283 Err(io::Error::new(
284 io::ErrorKind::Unsupported,
285 "splice not available on this platform",
286 ))
287}
288
289#[cfg(target_os = "linux")]
298fn find_header_end(buf: &[u8]) -> Option<usize> {
299 if buf.len() < 4 {
301 return None;
302 }
303 buf.windows(4).position(|w| w == b"\r\n\r\n")
306}
307
308#[cfg(target_os = "linux")]
312fn parse_status_code(header_str: &str) -> io::Result<u16> {
313 let first_line = header_str
314 .lines()
315 .next()
316 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "empty response"))?;
317
318 let mut parts = first_line.split_whitespace();
319 let _version = parts
320 .next()
321 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing HTTP version"))?;
322 let code_str = parts
323 .next()
324 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "missing status code"))?;
325
326 code_str.parse::<u16>().map_err(|_| {
327 io::Error::new(
328 io::ErrorKind::InvalidData,
329 format!("invalid status code: {code_str}"),
330 )
331 })
332}
333
334#[cfg(target_os = "linux")]
338fn is_chunked(header_str: &str) -> bool {
339 for line in header_str.lines().skip(1) {
340 if let Some((name, value)) = line.split_once(':')
341 && name.trim().eq_ignore_ascii_case("transfer-encoding")
342 && value.trim().eq_ignore_ascii_case("chunked")
343 {
344 return true;
345 }
346 }
347 false
348}
349
350#[cfg(target_os = "linux")]
354fn parse_content_length(header_str: &str) -> io::Result<Option<u64>> {
355 for line in header_str.lines().skip(1) {
356 if let Some((name, value)) = line.split_once(':')
357 && name.trim().eq_ignore_ascii_case("content-length")
358 {
359 let value = value.trim();
360 return value.parse::<u64>().map(Some).map_err(|_| {
361 io::Error::new(
362 io::ErrorKind::InvalidData,
363 format!("invalid Content-Length: {value}"),
364 )
365 });
366 }
367 }
368 Ok(None)
369}
370
371#[cfg(target_os = "linux")]
377fn write_all_at_offset(file: &std::fs::File, mut buf: &[u8], mut offset: u64) -> io::Result<()> {
378 use std::os::unix::fs::FileExt;
379 while !buf.is_empty() {
380 let n = file.write_at(buf, offset)?;
381 if n == 0 {
382 return Err(io::Error::new(
383 io::ErrorKind::WriteZero,
384 "pwrite returned 0 — failed to write all bytes",
385 ));
386 }
387 offset += n as u64;
388 buf = &buf[n..];
389 }
390 Ok(())
391}
392
393#[cfg(target_os = "linux")]
401fn set_blocking(fd: std::os::unix::io::RawFd) {
402 let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
406 if flags >= 0 {
407 unsafe {
408 libc::fcntl(fd, libc::F_SETFL, flags & !libc::O_NONBLOCK);
409 }
410 }
411}
412
413#[cfg(all(test, target_os = "linux"))]
418mod tests {
419 use super::*;
420 use tokio::io::{AsyncReadExt, AsyncWriteExt};
421
422 async fn spawn_mock_206_server(body: Vec<u8>) -> std::net::SocketAddr {
428 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
429 let addr = listener.local_addr().unwrap();
430
431 tokio::spawn(async move {
432 let (mut sock, _) = listener.accept().await.unwrap();
433
434 let mut req_buf = [0u8; 4096];
436 let _n = sock.read(&mut req_buf).await.unwrap();
437
438 let req_str = std::str::from_utf8(&req_buf).unwrap_or("");
440 let (offset, length) = parse_request_range(req_str, body.len());
441
442 let end = offset + length.saturating_sub(1);
443 let chunk = &body[offset..offset + length.min(body.len() - offset)];
444
445 let response = format!(
446 "HTTP/1.1 206 Partial Content\r\n\
447 Content-Length: {}\r\n\
448 Content-Range: bytes {}-{}/{}\r\n\
449 Connection: close\r\n\
450 \r\n",
451 chunk.len(),
452 offset,
453 end,
454 body.len()
455 );
456
457 sock.write_all(response.as_bytes()).await.unwrap();
458 sock.write_all(chunk).await.unwrap();
459 });
460
461 addr
462 }
463
464 fn parse_request_range(req: &str, total: usize) -> (usize, usize) {
466 for line in req.lines() {
467 if let Some(rest) = line.strip_prefix("Range:") {
468 let rest = rest.trim();
469 if let Some(range) = rest.strip_prefix("bytes=")
470 && let Some((start_s, end_s)) = range.split_once('-')
471 {
472 let start: usize = start_s.parse().unwrap_or(0);
473 let end: usize = end_s.parse().unwrap_or(total - 1);
474 let length = end.saturating_sub(start) + 1;
475 return (start, length);
476 }
477 }
478 }
479 (0, total)
480 }
481
482 #[tokio::test]
483 async fn test_splice_download_basic() {
484 let payload: Vec<u8> = (0..100_000u32).map(|i| (i % 256) as u8).collect();
486 let addr = spawn_mock_206_server(payload.clone()).await;
487
488 let dir = tempfile::tempdir().unwrap();
490 let out_path = dir.path().join("out.bin");
491 let file = std::fs::OpenOptions::new()
492 .write(true)
493 .create(true)
494 .truncate(true)
495 .open(&out_path)
496 .unwrap();
497
498 let url = format!("http://{addr}/test.bin");
499 let n = try_splice_download(&url, 0, payload.len() as u64, &file, 0)
500 .await
501 .expect("splice download should succeed");
502
503 assert_eq!(n, payload.len() as u64);
504
505 drop(file);
507 let content = std::fs::read(&out_path).unwrap();
508 assert_eq!(content, payload);
509 }
510
511 #[tokio::test]
512 async fn test_splice_download_range_offset() {
513 let payload: Vec<u8> = (0..200_000u32).map(|i| (i % 256) as u8).collect();
515 let addr = spawn_mock_206_server(payload.clone()).await;
516
517 let dir = tempfile::tempdir().unwrap();
518 let out_path = dir.path().join("out_range.bin");
519 let file = std::fs::OpenOptions::new()
520 .write(true)
521 .create(true)
522 .truncate(true)
523 .open(&out_path)
524 .unwrap();
525
526 let offset = 50_000u64;
527 let length = 80_000u64;
528 let url = format!("http://{addr}/test.bin");
529 let n = try_splice_download(&url, offset, length, &file, 0)
530 .await
531 .expect("splice range download should succeed");
532
533 assert_eq!(n, length);
534
535 drop(file);
537 let content = std::fs::read(&out_path).unwrap();
538 assert_eq!(content.len(), length as usize);
539 assert_eq!(
540 content,
541 &payload[offset as usize..(offset + length) as usize]
542 );
543 }
544
545 #[tokio::test]
546 async fn test_splice_download_with_file_offset() {
547 let payload: Vec<u8> = (0..64_000u32).map(|i| (i % 256) as u8).collect();
549 let addr = spawn_mock_206_server(payload.clone()).await;
550
551 let dir = tempfile::tempdir().unwrap();
552 let out_path = dir.path().join("out_foff.bin");
553 let file = std::fs::OpenOptions::new()
555 .write(true)
556 .create(true)
557 .truncate(true)
558 .open(&out_path)
559 .unwrap();
560 file.set_len(100_000).unwrap();
561
562 let file_offset = 30_000u64;
563 let url = format!("http://{addr}/test.bin");
564 let n = try_splice_download(&url, 0, payload.len() as u64, &file, file_offset)
565 .await
566 .expect("splice with file offset should succeed");
567
568 assert_eq!(n, payload.len() as u64);
569
570 drop(file);
572 let content = std::fs::read(&out_path).unwrap();
573 assert_eq!(content.len(), 100_000);
574 assert!(
576 content[..file_offset as usize].iter().all(|&b| b == 0),
577 "region before file_offset should be zero"
578 );
579 assert_eq!(
581 &content[file_offset as usize..file_offset as usize + payload.len()],
582 &payload[..]
583 );
584 assert!(
586 content[file_offset as usize + payload.len()..]
587 .iter()
588 .all(|&b| b == 0),
589 "region after spliced data should be zero"
590 );
591 }
592
593 #[tokio::test]
594 async fn test_splice_download_https_rejected() {
595 let file = tempfile::tempfile().unwrap();
596 let result = try_splice_download("https://example.com/file", 0, 100, &file, 0).await;
597 assert!(result.is_err());
598 let err = result.unwrap_err();
599 assert_eq!(err.kind(), io::ErrorKind::Unsupported);
600 }
601
602 #[tokio::test]
603 async fn test_splice_download_zero_length() {
604 let file = tempfile::tempfile().unwrap();
605 let result = try_splice_download("http://example.com/file", 0, 0, &file, 0).await;
606 assert!(result.is_ok());
607 assert_eq!(result.unwrap(), 0);
608 }
609
610 #[tokio::test]
611 async fn test_splice_download_non_206_falls_back() {
612 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
615 let addr = listener.local_addr().unwrap();
616
617 tokio::spawn(async move {
618 let (mut sock, _) = listener.accept().await.unwrap();
619 let mut buf = [0u8; 4096];
620 let _ = sock.read(&mut buf).await.unwrap();
621 sock.write_all(
622 b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nConnection: close\r\n\r\nhello",
623 )
624 .await
625 .unwrap();
626 });
627
628 let dir = tempfile::tempdir().unwrap();
629 let out_path = dir.path().join("out_200.bin");
630 let file = std::fs::OpenOptions::new()
631 .write(true)
632 .create(true)
633 .truncate(true)
634 .open(&out_path)
635 .unwrap();
636
637 let url = format!("http://{addr}/test.bin");
638 let result = try_splice_download(&url, 0, 5, &file, 0).await;
639 assert!(result.is_err(), "non-206 should return Err for fallback");
640 }
641
642 #[tokio::test]
643 async fn test_splice_download_small_body_in_header_buffer() {
644 let payload = b"tiny payload!".to_vec();
647 let addr = spawn_mock_206_server(payload.clone()).await;
648
649 let dir = tempfile::tempdir().unwrap();
650 let out_path = dir.path().join("out_tiny.bin");
651 let file = std::fs::OpenOptions::new()
652 .write(true)
653 .create(true)
654 .truncate(true)
655 .open(&out_path)
656 .unwrap();
657
658 let url = format!("http://{addr}/test.bin");
659 let n = try_splice_download(&url, 0, payload.len() as u64, &file, 0)
660 .await
661 .expect("splice tiny download should succeed");
662
663 assert_eq!(n, payload.len() as u64);
664
665 drop(file);
666 let content = std::fs::read(&out_path).unwrap();
667 assert_eq!(content, payload);
668 }
669
670 #[test]
671 fn test_find_header_end() {
672 assert_eq!(find_header_end(b""), None);
673 assert_eq!(find_header_end(b"HTTP/1.1 200 OK\r\n"), None);
674 assert_eq!(
675 find_header_end(b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\n\r\n"),
676 Some(34)
677 );
678 assert_eq!(find_header_end(b"HTTP/1.1 200 OK\r\n\r\nbody"), Some(15));
679 }
680
681 #[test]
682 fn test_parse_status_code() {
683 assert_eq!(
684 parse_status_code("HTTP/1.1 206 Partial Content\r\n").unwrap(),
685 206
686 );
687 assert_eq!(parse_status_code("HTTP/1.1 200 OK\r\n").unwrap(), 200);
688 assert!(parse_status_code("garbage").is_err());
689 assert!(parse_status_code("").is_err());
690 }
691
692 #[test]
693 fn test_is_chunked() {
694 assert!(is_chunked(
695 "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n"
696 ));
697 assert!(is_chunked(
698 "HTTP/1.1 200 OK\r\ntransfer-encoding: CHUNKED\r\n\r\n"
699 ));
700 assert!(!is_chunked(
701 "HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n"
702 ));
703 }
704
705 #[test]
706 fn test_parse_content_length() {
707 assert_eq!(
708 parse_content_length("HTTP/1.1 206\r\nContent-Length: 12345\r\n\r\n").unwrap(),
709 Some(12345)
710 );
711 assert_eq!(
712 parse_content_length("HTTP/1.1 206\r\ncontent-length: 0\r\n\r\n").unwrap(),
713 Some(0)
714 );
715 assert_eq!(parse_content_length("HTTP/1.1 206\r\n\r\n").unwrap(), None);
716 }
717
718 #[test]
719 fn test_write_all_at_offset() {
720 let dir = tempfile::tempdir().unwrap();
721 let path = dir.path().join("pwrite_test.bin");
722 let file = std::fs::OpenOptions::new()
723 .write(true)
724 .create(true)
725 .truncate(true)
726 .open(&path)
727 .unwrap();
728 file.set_len(100).unwrap();
729
730 write_all_at_offset(&file, b"hello", 10).unwrap();
731 write_all_at_offset(&file, b"world", 50).unwrap();
732
733 drop(file);
734 let content = std::fs::read(&path).unwrap();
735 assert_eq!(&content[10..15], b"hello");
736 assert_eq!(&content[50..55], b"world");
737 assert_eq!(&content[0..10], &[0u8; 10]);
738 }
739}
740
741#[cfg(all(test, not(target_os = "linux")))]
743mod tests {
744 use super::*;
745
746 #[tokio::test]
747 async fn test_splice_unsupported_on_non_linux() {
748 let file = tempfile::tempfile().unwrap();
749 let result = try_splice_download("http://example.com/file", 0, 100, &file, 0).await;
750 assert!(result.is_err());
751 assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Unsupported);
752 }
753}