PoC
This commit is contained in:
commit
59a265fe24
|
|
@ -0,0 +1 @@
|
|||
/target
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,17 @@
|
|||
[package]
|
||||
name = "rust-mjpeg-proxy"
|
||||
version = "0.1.0"
|
||||
authors = ["Astro <astro@spaceboyz.net>"]
|
||||
edition = "2018"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
bytes = "1"
|
||||
http = "0.2"
|
||||
hyper = "0.14"
|
||||
mpart-async = "0.5"
|
||||
futures = "0.3"
|
||||
tokio = { version = "1", features = ["macros", "rt"] }
|
||||
reqwest = { version = "0.11", features = ["stream"] }
|
||||
warp = "0.3"
|
||||
|
|
@ -0,0 +1,79 @@
|
|||
use std::sync::{Arc, RwLock};
|
||||
use http::header::HeaderMap;
|
||||
use tokio::sync::broadcast::{self, Sender, Receiver};
|
||||
use mpart_async::server::ParseOutput;
|
||||
|
||||
pub type Payload = Arc<ParseOutput>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct App {
|
||||
state: Arc<RwLock<Option<State>>>,
|
||||
}
|
||||
|
||||
impl App {
|
||||
pub fn new() -> App {
|
||||
App {
|
||||
state: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn source(&self, headers: HeaderMap, boundary: Vec<u8>) -> AppSource {
|
||||
{
|
||||
let mut state = self.state.write().unwrap();
|
||||
*state = Some(State {
|
||||
headers: Arc::new(headers),
|
||||
boundary,
|
||||
next_broadcast: broadcast::channel(128),
|
||||
});
|
||||
}
|
||||
|
||||
AppSource {
|
||||
app: self.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn subscribe(&self) -> Option<Receiver<Payload>> {
|
||||
let state = self.state.read().unwrap();
|
||||
state.as_ref().map(|state| {
|
||||
let tx = &state.next_broadcast.0;
|
||||
tx.subscribe()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn boundary(&self) -> Vec<u8> {
|
||||
let state = self.state.read().unwrap();
|
||||
state.as_ref().unwrap().boundary.clone()
|
||||
}
|
||||
|
||||
pub fn headers(&self) -> Option<Arc<HeaderMap>> {
|
||||
let state = self.state.read().unwrap();
|
||||
state.as_ref().map(|state| {
|
||||
state.headers.clone()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct AppSource {
|
||||
app: App,
|
||||
}
|
||||
|
||||
impl Drop for AppSource {
|
||||
fn drop(&mut self) {
|
||||
let mut state = self.app.state.write().unwrap();
|
||||
*state = None;
|
||||
}
|
||||
}
|
||||
|
||||
impl AppSource {
|
||||
pub fn begin_part(&mut self) -> Sender<Payload> {
|
||||
let mut state = self.app.state.write().unwrap();
|
||||
let (tx, _rx) = std::mem::replace(&mut state.as_mut().unwrap().next_broadcast, broadcast::channel(128));
|
||||
tx
|
||||
}
|
||||
}
|
||||
|
||||
struct State {
|
||||
headers: Arc<HeaderMap>,
|
||||
boundary: Vec<u8>,
|
||||
next_broadcast: (Sender<Payload>, Receiver<Payload>),
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
use std::sync::Arc;
|
||||
use futures::stream::TryStreamExt;
|
||||
use mpart_async::server::{
|
||||
MultipartParser,
|
||||
ParseOutput::{Headers, Bytes},
|
||||
};
|
||||
use reqwest::StatusCode;
|
||||
use crate::app::App;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Error {
|
||||
InvalidStatus(StatusCode),
|
||||
InvalidMultipart,
|
||||
}
|
||||
|
||||
impl std::error::Error for Error {}
|
||||
|
||||
impl std::fmt::Display for Error {
|
||||
fn fmt(&self, fmt: &mut std::fmt::Formatter) -> Result<(), std::fmt::Error> {
|
||||
write!(fmt, "Invalid multipart")
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_once(app: &App) -> Result<(), Box<dyn std::error::Error>> {
|
||||
// let res = reqwest::get("http://172.20.78.164:81/stream")
|
||||
let res = reqwest::get("http://localhost:8080/?action=stream")
|
||||
.await?;
|
||||
if res.status() != StatusCode::OK {
|
||||
Err(Error::InvalidStatus(res.status()))?;
|
||||
}
|
||||
|
||||
let content_type = res.headers().get("content-type")
|
||||
.ok_or(Error::InvalidMultipart)?;
|
||||
let mut boundary = None;
|
||||
for field in content_type.as_ref().split(|b| *b == ';' as u8) {
|
||||
if field.starts_with(b"boundary=") {
|
||||
boundary = Some(field[9..].to_owned());
|
||||
}
|
||||
}
|
||||
let boundary = boundary.ok_or(Error::InvalidMultipart)?;
|
||||
let headers = res.headers().clone();
|
||||
let mut stream = MultipartParser::new(
|
||||
boundary.clone(),
|
||||
res.bytes_stream()
|
||||
.map_ok(|buf| buf),
|
||||
);
|
||||
let mut source = app.source(headers, boundary);
|
||||
let mut tx = source.begin_part();
|
||||
while let Ok(Some(part)) = stream.try_next().await {
|
||||
match &part {
|
||||
Headers(_headers) =>
|
||||
tx = source.begin_part(),
|
||||
Bytes(_bytes) => {}
|
||||
}
|
||||
let _ = tx.send(Arc::new(part));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn run(app: App) {
|
||||
loop {
|
||||
run_once(&app).await
|
||||
.unwrap_or_else(|e| println!("Client: {:?}", e));
|
||||
|
||||
// TODO: delay with backoff
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
use futures::prelude::future::FutureExt;
|
||||
|
||||
mod app;
|
||||
mod client;
|
||||
mod server;
|
||||
|
||||
#[tokio::main(flavor = "current_thread")]
|
||||
async fn main() {
|
||||
let app = app::App::new();
|
||||
let c = tokio::spawn(client::run(app.clone()));
|
||||
let s = tokio::spawn(server::run(app));
|
||||
futures::select! {
|
||||
_ = c.fuse() => panic!("client exit"),
|
||||
_ = s.fuse() => panic!("server exit"),
|
||||
};
|
||||
}
|
||||
|
|
@ -0,0 +1,162 @@
|
|||
use std::{convert::Infallible};
|
||||
use bytes::{BufMut, Bytes, BytesMut};
|
||||
use futures::stream;
|
||||
use tokio::sync::broadcast::{Receiver};
|
||||
use http::{
|
||||
header::HeaderMap,
|
||||
status::StatusCode,
|
||||
};
|
||||
use hyper::Body;
|
||||
use mpart_async::server::ParseOutput;
|
||||
use warp::{
|
||||
Filter,
|
||||
reply::{Reply, Response},
|
||||
};
|
||||
use crate::app::{App, Payload};
|
||||
|
||||
struct PartStream {
|
||||
app: App,
|
||||
}
|
||||
|
||||
impl PartStream {
|
||||
async fn new(app: App) -> Result<impl warp::Reply, Infallible> {
|
||||
Ok(PartStream { app })
|
||||
}
|
||||
}
|
||||
|
||||
impl Reply for PartStream {
|
||||
fn into_response(self) -> Response {
|
||||
let (rx, headers) = match (self.app.subscribe(), self.app.headers()) {
|
||||
(Some(rx), Some(headers)) => (rx, headers),
|
||||
_ => {
|
||||
let mut res = Response::new(Body::empty());
|
||||
*res.status_mut() = StatusCode::BAD_GATEWAY;
|
||||
return res;
|
||||
}
|
||||
};
|
||||
|
||||
let body_stream = stream::unfold((self, rx), |(this, mut rx)| async move {
|
||||
let result = match rx.recv().await {
|
||||
Ok(payload) => {
|
||||
let bytes = match payload.as_ref() {
|
||||
ParseOutput::Headers(headers) => {
|
||||
let mut buf = BytesMut::new();
|
||||
buf.put_slice(b"--");
|
||||
buf.put_slice(&this.app.boundary());
|
||||
buf.put_slice(b"\r\n");
|
||||
for (name, value) in headers {
|
||||
buf.put_slice(name.as_ref());
|
||||
buf.put_slice(b": ");
|
||||
buf.put_slice(value.as_ref());
|
||||
buf.put_slice(b"\r\n");
|
||||
}
|
||||
buf.put_slice(b"\r\n");
|
||||
buf.freeze()
|
||||
}
|
||||
ParseOutput::Bytes(bytes) => {
|
||||
bytes.clone()
|
||||
}
|
||||
};
|
||||
Some(Ok(bytes))
|
||||
}
|
||||
Err(e) => {
|
||||
match this.app.subscribe() {
|
||||
Some(next_rx) => {
|
||||
rx = next_rx;
|
||||
Some(Ok(Bytes::from("")))
|
||||
}
|
||||
None =>
|
||||
Some(Err(e)),
|
||||
}
|
||||
}
|
||||
};
|
||||
result.map(|result| (result, (this, rx)))
|
||||
});
|
||||
let mut res = Response::new(Body::wrap_stream(body_stream));
|
||||
for (name, value) in headers.as_ref() {
|
||||
res.headers_mut()
|
||||
.insert(name, value.clone());
|
||||
}
|
||||
res
|
||||
}
|
||||
}
|
||||
|
||||
struct PartCapture {
|
||||
rx: Option<Receiver<Payload>>,
|
||||
headers: Option<HeaderMap>,
|
||||
}
|
||||
|
||||
impl PartCapture {
|
||||
// async fn start(app: App) -> PartCapture {
|
||||
async fn start(app: App) -> Result<impl warp::Reply, Infallible> {
|
||||
let mut rx = app.subscribe();
|
||||
let headers = match (&mut rx, app.headers()) {
|
||||
(Some(ref mut rx), Some(ref stream_headers)) => {
|
||||
match rx.recv().await {
|
||||
Ok(output) => match output.as_ref() {
|
||||
ParseOutput::Headers(part_headers) => {
|
||||
let mut headers = stream_headers.as_ref().clone();
|
||||
for (name, value) in part_headers {
|
||||
headers.insert(name, value.clone());
|
||||
}
|
||||
Some(headers)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
Ok(PartCapture { rx, headers })
|
||||
}
|
||||
}
|
||||
|
||||
impl Reply for PartCapture {
|
||||
fn into_response(self) -> Response {
|
||||
let (rx, headers) = match (self.rx, self.headers) {
|
||||
(Some(rx), Some(headers)) => (rx, headers),
|
||||
_ => {
|
||||
let mut res = Response::new(Body::empty());
|
||||
*res.status_mut() = StatusCode::BAD_GATEWAY;
|
||||
return res;
|
||||
}
|
||||
};
|
||||
|
||||
let body_stream = stream::unfold(rx, |mut rx| async move {
|
||||
match rx.recv().await {
|
||||
Ok(payload) => {
|
||||
match payload.as_ref() {
|
||||
ParseOutput::Headers(_headers) => {
|
||||
None
|
||||
}
|
||||
ParseOutput::Bytes(bytes) => {
|
||||
Some((Ok(bytes.clone()), rx))
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
println!("e: {:?}", e);
|
||||
Some((Err(e), rx))
|
||||
},
|
||||
}
|
||||
});
|
||||
|
||||
let mut res = Response::new(Body::wrap_stream(body_stream));
|
||||
*res.headers_mut() = headers;
|
||||
res
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run(app: App) {
|
||||
let app_ = app.clone();
|
||||
let stream = warp::path!("stream.mjpeg")
|
||||
.and_then(move || PartStream::new(app_.clone()));
|
||||
let capture = warp::path!("capture.jpg")
|
||||
.map(move || app.clone())
|
||||
.and_then(PartCapture::start);
|
||||
let routes = warp::get()
|
||||
.and(stream)
|
||||
.or(capture);
|
||||
warp::serve(routes).run(([127, 0, 0, 1], 3030)).await;
|
||||
}
|
||||
Loading…
Reference in New Issue