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        crate::tls::http_put(&url, data, &crate::tls::HttpOptions::default())
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` (plain HTTP, no TLS).
87#[allow(dead_code)]
88fn http_put_plain(url: &str, body: &[u8], authorization: Option<&str>) -> Result<()> {
89    let rest = url
90        .strip_prefix("http://")
91        .or_else(|| url.strip_prefix("https://"))
92        .context("only http:// and https:// push URLs are supported")?;
93    if url.starts_with("https://") {
94        return crate::tls::http_put(url, body, &crate::tls::HttpOptions::default());
95    }
96    let (hostport, path) =
97        rest.split_once('/').map(|(h, p)| (h, format!("/{p}"))).unwrap_or((rest, "/".into()));
98    let host = hostport.split(':').next().unwrap_or(hostport);
99    let addr = hostport
100        .to_socket_addrs()
101        .with_context(|| format!("resolving {hostport}"))?
102        .next()
103        .context("no addresses for push host")?;
104    let mut stream = TcpStream::connect_timeout(&addr, Duration::from_secs(10))
105        .with_context(|| format!("connecting to {hostport}"))?;
106    stream.set_write_timeout(Some(Duration::from_secs(30)))?;
107    stream.set_read_timeout(Some(Duration::from_secs(30)))?;
108
109    let mut req = format!(
110        "PUT {path} HTTP/1.1\r\nHost: {host}\r\nContent-Length: {}\r\nConnection: close\r\n",
111        body.len()
112    );
113    if let Some(auth) = authorization {
114        req.push_str(&format!("Authorization: {auth}\r\n"));
115    }
116    req.push_str("Content-Type: application/octet-stream\r\n\r\n");
117    stream.write_all(req.as_bytes())?;
118    stream.write_all(body)?;
119    stream.flush()?;
120
121    let mut resp = Vec::new();
122    stream.read_to_end(&mut resp).ok();
123    let text = String::from_utf8_lossy(&resp);
124    let status = text.lines().next().unwrap_or("");
125    if !(status.contains(" 200 ")
126        || status.contains(" 201 ")
127        || status.contains(" 204 ")
128        || status.contains(" 100 "))
129    {
130        // Some servers speak HTTP/1.0 without spaces the same way — accept 2xx.
131        let ok = text.contains(" 2") && status.starts_with("HTTP/");
132        if !ok {
133            bail!("push rejected: {status}");
134        }
135    }
136    Ok(())
137}
138
139/// Read an input path: plain filesystem, or `udp://bind_host:port` for a short
140/// live capture (collects datagrams for `duration` or until ~4 MiB).
141pub fn read_input(spec: &str) -> Result<(String, Vec<u8>)> {
142    if let Some(addr) = spec.strip_prefix("udp://") {
143        let data = udp_ingest(addr, Duration::from_secs(3), 4 * 1024 * 1024)?;
144        return Ok((spec.to_string(), data));
145    }
146    let path = Path::new(spec);
147    let data = fs::read(path).with_context(|| format!("reading {}", path.display()))?;
148    Ok((spec.to_string(), data))
149}
150
151/// Bind `addr` (e.g. `0.0.0.0:5000`) and collect UDP datagrams until `max_bytes`
152/// or `timeout` elapses with no further data after the first packet.
153fn udp_ingest(addr: &str, timeout: Duration, max_bytes: usize) -> Result<Vec<u8>> {
154    let sock = UdpSocket::bind(addr).with_context(|| format!("binding UDP {addr}"))?;
155    sock.set_read_timeout(Some(Duration::from_millis(500)))?;
156    let mut buf = vec![0u8; 65535];
157    let mut out = Vec::new();
158    let start = std::time::Instant::now();
159    let mut got_any = false;
160    while out.len() < max_bytes && start.elapsed() < timeout {
161        match sock.recv_from(&mut buf) {
162            Ok((n, _)) => {
163                out.extend_from_slice(&buf[..n]);
164                got_any = true;
165            }
166            Err(e)
167                if e.kind() == std::io::ErrorKind::WouldBlock
168                    || e.kind() == std::io::ErrorKind::TimedOut =>
169            {
170                if got_any {
171                    // Quiet period after first data — stop.
172                    break;
173                }
174            }
175            Err(e) => return Err(e).context("UDP recv"),
176        }
177    }
178    anyhow::ensure!(got_any, "UDP ingest on {addr}: no datagrams received within {timeout:?}");
179    Ok(out)
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use std::net::TcpListener;
186    use std::thread;
187
188    #[test]
189    fn file_sink_writes_relative() {
190        let dir = std::env::temp_dir().join(format!("sheathe-io-{}", std::process::id()));
191        let _ = fs::remove_dir_all(&dir);
192        let mut sink = FileSink::new(&dir);
193        sink.put("a/b.txt", b"hello").unwrap();
194        assert_eq!(fs::read(dir.join("a/b.txt")).unwrap(), b"hello");
195        let _ = fs::remove_dir_all(&dir);
196    }
197
198    #[test]
199    fn http_put_round_trip() {
200        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
201        let port = listener.local_addr().unwrap().port();
202        let handle = thread::spawn(move || {
203            let (mut sock, _) = listener.accept().unwrap();
204            let mut buf = Vec::new();
205            let mut chunk = [0u8; 1024];
206            // Read full request (headers + 9-byte body).
207            loop {
208                match sock.read(&mut chunk) {
209                    Ok(0) => break,
210                    Ok(n) => {
211                        buf.extend_from_slice(&chunk[..n]);
212                        if buf.windows(4).any(|w| w == b"\r\n\r\n")
213                            && buf.windows(9).any(|w| w == b"hello-seg")
214                        {
215                            break;
216                        }
217                    }
218                    Err(_) => break,
219                }
220            }
221            let req = String::from_utf8_lossy(&buf);
222            assert!(req.contains("PUT /live/seg.m4s "), "req={req}");
223            assert!(req.contains("hello-seg"), "req={req}");
224            sock.write_all(
225                b"HTTP/1.1 201 Created\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
226            )
227            .unwrap();
228        });
229        let mut sink = HttpPushSink::new(format!("http://127.0.0.1:{port}/live"));
230        sink.put("seg.m4s", b"hello-seg").unwrap();
231        handle.join().unwrap();
232    }
233}