Skip to main content

sheathe_package/
io.rs

1//! Output sinks and input sources for packaging (Phase 5 IO backends).
2//!
3//! - [`FileSink`] — write relative paths under an output directory (default).
4//! - [`HttpPushSink`] — HTTP/1.1 `PUT` each object to a base URL (pure `std`).
5//! - [`read_input`] — load bytes from a file path or `udp://host:port` (single
6//!   datagram / short capture window for live ingest demos).
7
8use anyhow::{Context, Result, bail};
9use std::fs;
10use std::io::{Read, Write};
11use std::net::{TcpStream, ToSocketAddrs, UdpSocket};
12use std::path::{Path, PathBuf};
13use std::time::Duration;
14
15/// Where packaged objects (segments, manifests) are written.
16pub trait ObjectSink: Send {
17    /// Write `data` at a path relative to the sink root.
18    fn put(&mut self, relative: &str, data: &[u8]) -> Result<()>;
19    /// Human-readable root (directory or base URL) for logging.
20    fn root(&self) -> String;
21}
22
23/// Local filesystem sink.
24#[derive(Debug, Clone)]
25pub struct FileSink {
26    /// Output directory (created on first write if missing).
27    pub dir: PathBuf,
28}
29
30impl FileSink {
31    /// Create a sink rooted at `dir`.
32    pub fn new(dir: impl Into<PathBuf>) -> Self {
33        Self { dir: dir.into() }
34    }
35}
36
37impl ObjectSink for FileSink {
38    fn put(&mut self, relative: &str, data: &[u8]) -> Result<()> {
39        let path = self.dir.join(relative);
40        if let Some(parent) = path.parent() {
41            fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?;
42        }
43        fs::write(&path, data).with_context(|| format!("writing {}", path.display()))
44    }
45
46    fn root(&self) -> String {
47        self.dir.display().to_string()
48    }
49}
50
51/// HTTP/1.1 PUT sink — pushes each object to `{base_url}/{relative}`.
52///
53/// Pure std: opens a TCP connection per object and writes a minimal request.
54/// Intended for origin CDN ingest / PUT-capable static hosts, not multipart.
55#[derive(Debug, Clone)]
56pub struct HttpPushSink {
57    /// Base URL without trailing slash, e.g. `http://127.0.0.1:8080/live`.
58    pub base_url: String,
59    /// Optional `Authorization` header value.
60    pub authorization: Option<String>,
61}
62
63impl HttpPushSink {
64    /// Create a push sink for `base_url`.
65    pub fn new(base_url: impl Into<String>) -> Self {
66        let mut base = base_url.into();
67        while base.ends_with('/') {
68            base.pop();
69        }
70        Self { base_url: base, authorization: None }
71    }
72}
73
74impl ObjectSink for HttpPushSink {
75    fn put(&mut self, relative: &str, data: &[u8]) -> Result<()> {
76        let url = format!("{}/{}", self.base_url, relative.trim_start_matches('/'));
77        http_put(&url, data, self.authorization.as_deref())
78            .with_context(|| format!("HTTP PUT {url}"))
79    }
80
81    fn root(&self) -> String {
82        self.base_url.clone()
83    }
84}
85
86/// Parse `http://host:port/path` and PUT `body`.
87fn http_put(url: &str, body: &[u8], authorization: Option<&str>) -> Result<()> {
88    let rest = url
89        .strip_prefix("http://")
90        .or_else(|| url.strip_prefix("https://"))
91        .context("only http:// and https:// push URLs are supported")?;
92    if url.starts_with("https://") {
93        bail!(
94            "HTTPS push requires TLS; use http:// for the pure-std push sink or terminate TLS externally"
95        );
96    }
97    let (hostport, path) =
98        rest.split_once('/').map(|(h, p)| (h, format!("/{p}"))).unwrap_or((rest, "/".into()));
99    let host = hostport.split(':').next().unwrap_or(hostport);
100    let addr = hostport
101        .to_socket_addrs()
102        .with_context(|| format!("resolving {hostport}"))?
103        .next()
104        .context("no addresses for push host")?;
105    let mut stream = TcpStream::connect_timeout(&addr, Duration::from_secs(10))
106        .with_context(|| format!("connecting to {hostport}"))?;
107    stream.set_write_timeout(Some(Duration::from_secs(30)))?;
108    stream.set_read_timeout(Some(Duration::from_secs(30)))?;
109
110    let mut req = format!(
111        "PUT {path} HTTP/1.1\r\nHost: {host}\r\nContent-Length: {}\r\nConnection: close\r\n",
112        body.len()
113    );
114    if let Some(auth) = authorization {
115        req.push_str(&format!("Authorization: {auth}\r\n"));
116    }
117    req.push_str("Content-Type: application/octet-stream\r\n\r\n");
118    stream.write_all(req.as_bytes())?;
119    stream.write_all(body)?;
120    stream.flush()?;
121
122    let mut resp = Vec::new();
123    stream.read_to_end(&mut resp).ok();
124    let text = String::from_utf8_lossy(&resp);
125    let status = text.lines().next().unwrap_or("");
126    if !(status.contains(" 200 ")
127        || status.contains(" 201 ")
128        || status.contains(" 204 ")
129        || status.contains(" 100 "))
130    {
131        // Some servers speak HTTP/1.0 without spaces the same way — accept 2xx.
132        let ok = text.contains(" 2") && status.starts_with("HTTP/");
133        if !ok {
134            bail!("push rejected: {status}");
135        }
136    }
137    Ok(())
138}
139
140/// Read an input path: plain filesystem, or `udp://bind_host:port` for a short
141/// live capture (collects datagrams for `duration` or until ~4 MiB).
142pub fn read_input(spec: &str) -> Result<(String, Vec<u8>)> {
143    if let Some(addr) = spec.strip_prefix("udp://") {
144        let data = udp_ingest(addr, Duration::from_secs(3), 4 * 1024 * 1024)?;
145        return Ok((spec.to_string(), data));
146    }
147    let path = Path::new(spec);
148    let data = fs::read(path).with_context(|| format!("reading {}", path.display()))?;
149    Ok((spec.to_string(), data))
150}
151
152/// Bind `addr` (e.g. `0.0.0.0:5000`) and collect UDP datagrams until `max_bytes`
153/// or `timeout` elapses with no further data after the first packet.
154fn udp_ingest(addr: &str, timeout: Duration, max_bytes: usize) -> Result<Vec<u8>> {
155    let sock = UdpSocket::bind(addr).with_context(|| format!("binding UDP {addr}"))?;
156    sock.set_read_timeout(Some(Duration::from_millis(500)))?;
157    let mut buf = vec![0u8; 65535];
158    let mut out = Vec::new();
159    let start = std::time::Instant::now();
160    let mut got_any = false;
161    while out.len() < max_bytes && start.elapsed() < timeout {
162        match sock.recv_from(&mut buf) {
163            Ok((n, _)) => {
164                out.extend_from_slice(&buf[..n]);
165                got_any = true;
166            }
167            Err(e)
168                if e.kind() == std::io::ErrorKind::WouldBlock
169                    || e.kind() == std::io::ErrorKind::TimedOut =>
170            {
171                if got_any {
172                    // Quiet period after first data — stop.
173                    break;
174                }
175            }
176            Err(e) => return Err(e).context("UDP recv"),
177        }
178    }
179    anyhow::ensure!(got_any, "UDP ingest on {addr}: no datagrams received within {timeout:?}");
180    Ok(out)
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use std::net::TcpListener;
187    use std::thread;
188
189    #[test]
190    fn file_sink_writes_relative() {
191        let dir = std::env::temp_dir().join(format!("sheathe-io-{}", std::process::id()));
192        let _ = fs::remove_dir_all(&dir);
193        let mut sink = FileSink::new(&dir);
194        sink.put("a/b.txt", b"hello").unwrap();
195        assert_eq!(fs::read(dir.join("a/b.txt")).unwrap(), b"hello");
196        let _ = fs::remove_dir_all(&dir);
197    }
198
199    #[test]
200    fn http_put_round_trip() {
201        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
202        let port = listener.local_addr().unwrap().port();
203        let handle = thread::spawn(move || {
204            let (mut sock, _) = listener.accept().unwrap();
205            let mut buf = Vec::new();
206            let mut chunk = [0u8; 1024];
207            // Read full request (headers + 9-byte body).
208            loop {
209                match sock.read(&mut chunk) {
210                    Ok(0) => break,
211                    Ok(n) => {
212                        buf.extend_from_slice(&chunk[..n]);
213                        if buf.windows(4).any(|w| w == b"\r\n\r\n")
214                            && buf.windows(9).any(|w| w == b"hello-seg")
215                        {
216                            break;
217                        }
218                    }
219                    Err(_) => break,
220                }
221            }
222            let req = String::from_utf8_lossy(&buf);
223            assert!(req.contains("PUT /live/seg.m4s "), "req={req}");
224            assert!(req.contains("hello-seg"), "req={req}");
225            sock.write_all(
226                b"HTTP/1.1 201 Created\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
227            )
228            .unwrap();
229        });
230        let mut sink = HttpPushSink::new(format!("http://127.0.0.1:{port}/live"));
231        sink.put("seg.m4s", b"hello-seg").unwrap();
232        handle.join().unwrap();
233    }
234}