ort_openrouter_cli/net/
chunked.rs

1//! ort: Open Router CLI
2//! https://github.com/grahamking/ort
3//!
4//! MIT License
5//! Copyright (c) 2025 Graham King
6
7extern crate alloc;
8use alloc::ffi::CString;
9use alloc::string::{String, ToString};
10use alloc::vec::Vec;
11
12use crate::{ErrorKind, OrtResult, Read, common::buf_read, libc, ort_error};
13
14/// Read a transfer encoding chunked body, chunk by chunk.
15///
16/// This normally returns the chunks as provided by upstream, except if that
17/// would split a mutli-byte char in which case we return N chunks at once.
18pub fn read<R: Read, const MAX_CHUNK_SIZE: usize>(
19    r: buf_read::OrtBufReader<R>,
20) -> ChunkedIterator<R, MAX_CHUNK_SIZE> {
21    ChunkedIterator::new(r)
22}
23
24pub struct ChunkedIterator<R: Read, const MAX_CHUNK_SIZE: usize> {
25    r: buf_read::OrtBufReader<R>,
26    size_buf: String,
27    data_buf: Vec<u8>,
28}
29
30/// Lending Iterator. This doesn't implement Iterator because that doesn't allow the Item
31/// to borrow from the iterator (so Item couldn't be &str).
32///
33/// max_chunk_size: Estimated size of the biggest chunk we will receive.
34/// Ideally a power of 2. It's OK if this is wrong, we will realloc.
35impl<R: Read, const MAX_CHUNK_SIZE: usize> ChunkedIterator<R, MAX_CHUNK_SIZE> {
36    fn new(r: buf_read::OrtBufReader<R>) -> ChunkedIterator<R, MAX_CHUNK_SIZE> {
37        ChunkedIterator {
38            r,
39            size_buf: String::with_capacity(16),
40            data_buf: Vec::with_capacity(MAX_CHUNK_SIZE),
41        }
42    }
43
44    pub fn next_chunk(&mut self) -> Option<OrtResult<&str>> {
45        let mut bytes_read = 0;
46        // Usually we only go through the loop once per call.
47        // Exceptions are the initial blank line, and splitting a multi-byte char.
48        loop {
49            // Read size line
50            // The size is always valid UTF-8. It's an ASCII hex number.
51            self.size_buf.clear();
52            match self.r.read_line(&mut self.size_buf) {
53                Ok(0) => {
54                    return Some(Err(ort_error(ErrorKind::ChunkedEofInSize, "")));
55                }
56                Ok(_) => {}
57                Err(err) => {
58                    err.debug_print();
59                    return Some(Err(ort_error(ErrorKind::ChunkedSizeReadError, "")));
60                }
61            }
62            let size_str = self.size_buf.trim();
63            if size_str.is_empty() {
64                // Skip initial blank line
65                continue;
66            }
67            let size = match usize::from_str_radix(size_str, 16) {
68                Ok(n) => n,
69                Err(_err) => {
70                    let c_s = CString::new("ERROR invalid chunked size: ".to_string() + size_str)
71                        .unwrap();
72                    unsafe {
73                        libc::write(2, c_s.as_ptr().cast(), c_s.count_bytes());
74                    }
75                    return Some(Err(ort_error(ErrorKind::ChunkedInvalidSize, "")));
76                }
77            };
78            if size == 0 {
79                // How transfer-encoding chunked signals EOF
80                return None;
81            }
82
83            // Ensure buffer capacity (do not shrink)
84            if bytes_read == 0 {
85                self.data_buf.clear();
86            }
87            // no-op if already enough space, so we don't need to check
88            self.data_buf.reserve_exact(size);
89            unsafe { self.data_buf.set_len(size + bytes_read) };
90
91            if let Err(_err) = self.r.read_exact(&mut self.data_buf[bytes_read..]) {
92                // Original included err detail
93                return Some(Err(ort_error(ErrorKind::ChunkedDataReadError, "")));
94            };
95            bytes_read += size;
96
97            // If we split a UTF-8 multi-byte character on the end of the chunk,
98            // fetch the next chunk. This really happens.
99            let last_byte = self.data_buf[self.data_buf.len() - 1];
100            if (last_byte & 0b1000_0000) != 0 {
101                //let c_s = CString::new("SPLIT MULTI-BYTE CHAR\n").unwrap();
102                //unsafe { libc::write(2, c_s.as_ptr().cast(), c_s.count_bytes()) };
103                continue;
104            }
105            break;
106        }
107        Some(Ok(unsafe { str::from_utf8_unchecked(&self.data_buf) }))
108    }
109}
ort_openrouter_cli/net/chunked.rs

ort_openrouter_cli/net/
chunked.rs
