1use 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
15pub trait ObjectSink: Send {
17 fn put(&mut self, relative: &str, data: &[u8]) -> Result<()>;
19 fn root(&self) -> String;
21}
22
23#[derive(Debug, Clone)]
25pub struct FileSink {
26 pub dir: PathBuf,
28}
29
30impl FileSink {
31 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#[derive(Debug, Clone)]
56pub struct HttpPushSink {
57 pub base_url: String,
59 pub authorization: Option<String>,
61}
62
63impl HttpPushSink {
64 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#[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 let ok = text.contains(" 2") && status.starts_with("HTTP/");
132 if !ok {
133 bail!("push rejected: {status}");
134 }
135 }
136 Ok(())
137}
138
139pub 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
151fn 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 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 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}