buenosairesam
e9e1d14e6b
- AudioParams.framing field: client declares "raw" or "adts" - Client strips ADTS from audio before sending (strip_adts) - Client does H.264 NAL inspection for keyframe detection (h264_is_keyframe) - Server uses declared sample_rate/channels for ADTS synthesis instead of hardcoded 48kHz/stereo - Server gates ADTS wrapping on framing field instead of per-packet sniffing New backends only need to pipe output to demux_and_send() — server and Python unchanged. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
356 lines
14 KiB
Rust
356 lines
14 KiB
Rust
//! Session: manages the ffmpeg recording subprocess for one client connection.
|
|
//!
|
|
//! Receives raw H.264 NAL units and AAC audio from the transport:
|
|
//! - Video: piped into ffmpeg → fragmented MP4 + UDP relay for live display
|
|
//! - Audio: written to raw AAC file for Python post-processing
|
|
//!
|
|
//! Also provides a Unix domain socket at `stream/scene.sock` carrying a copy
|
|
//! of the raw H.264 stream for Python's GPU scene detection. The socket is
|
|
//! fire-and-forget: if nobody connects, data is silently dropped; if the
|
|
//! reader is slow, old frames are dropped rather than stalling recording.
|
|
//!
|
|
//! Creates the session directory and writes its path to `data/active-session`
|
|
//! so the Python app can pick it up for SessionProcessor (audio extraction, etc).
|
|
|
|
use std::fs::{self, File};
|
|
use std::io::Write;
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::{Child, ChildStdin, Command, Stdio};
|
|
use std::thread;
|
|
|
|
use anyhow::{Context, Result};
|
|
use cht_common::protocol::AudioParams;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tracing::{debug, info, warn};
|
|
|
|
// Written next to the sessions/ directory so everything stays under data/.
|
|
// Python reads this to discover the session dir created by cht-server.
|
|
const ACTIVE_SESSION_FILENAME: &str = "active-session";
|
|
const RELAY_URL: &str = "udp://127.0.0.1:4445";
|
|
const SCENE_SOCKET_NAME: &str = "scene.sock";
|
|
|
|
struct ScenePacket {
|
|
data: Vec<u8>,
|
|
keyframe: bool,
|
|
}
|
|
|
|
/// ADTS configuration derived from AudioParams at session start.
|
|
struct AdtsConfig {
|
|
/// Whether to wrap audio with ADTS headers (false if client sends ADTS).
|
|
wrap: bool,
|
|
sr_idx: u8,
|
|
ch_cfg: u8,
|
|
}
|
|
|
|
impl AdtsConfig {
|
|
fn from_params(params: &AudioParams) -> Self {
|
|
let wrap = params.framing == "raw";
|
|
let sr_idx = match params.sample_rate {
|
|
96000 => 0, 88200 => 1, 64000 => 2, 48000 => 3,
|
|
44100 => 4, 32000 => 5, 24000 => 6, 22050 => 7,
|
|
16000 => 8, 12000 => 9, 11025 => 10, 8000 => 11,
|
|
_ => 3, // default 48kHz
|
|
};
|
|
let ch_cfg = params.channels.min(7) as u8;
|
|
Self { wrap, sr_idx, ch_cfg }
|
|
}
|
|
}
|
|
|
|
pub struct Session {
|
|
#[allow(dead_code)]
|
|
session_dir: PathBuf,
|
|
active_session_file: PathBuf,
|
|
ffmpeg: Child,
|
|
video_stdin: Option<ChildStdin>,
|
|
audio_file: Option<File>,
|
|
audio_config: AdtsConfig,
|
|
scene_tx: Option<tokio::sync::mpsc::Sender<ScenePacket>>,
|
|
#[allow(dead_code)]
|
|
fps: u32,
|
|
}
|
|
|
|
impl Session {
|
|
pub fn start(session_id: &str, sessions_dir: &Path, fps: u32, audio_params: &AudioParams) -> Result<Self> {
|
|
let active_session_file = sessions_dir
|
|
.parent()
|
|
.unwrap_or(sessions_dir)
|
|
.join(ACTIVE_SESSION_FILENAME);
|
|
let session_dir = sessions_dir.join(session_id);
|
|
let stream_dir = session_dir.join("stream");
|
|
fs::create_dir_all(&stream_dir)
|
|
.with_context(|| format!("create session dir: {}", stream_dir.display()))?;
|
|
|
|
let recording_path = stream_dir.join("recording_000.mp4");
|
|
let audio_path = stream_dir.join("audio.aac");
|
|
|
|
info!("Session {session_id}: recording → {}", recording_path.display());
|
|
|
|
let mut child = Command::new("ffmpeg")
|
|
.args([
|
|
"-fflags", "nobuffer",
|
|
"-flags", "low_delay",
|
|
"-f", "h264",
|
|
"-framerate", &fps.to_string(),
|
|
"-i", "pipe:0",
|
|
// fMP4 — same flags as Python StreamRecorder
|
|
"-c:v", "copy",
|
|
"-f", "mp4",
|
|
"-movflags", "frag_keyframe+empty_moov+default_base_moof",
|
|
"-flush_packets", "1",
|
|
recording_path.to_str().unwrap(),
|
|
// UDP relay for live display
|
|
"-c:v", "copy",
|
|
"-f", "mpegts",
|
|
"-flush_packets", "1",
|
|
RELAY_URL,
|
|
"-hide_banner", "-loglevel", "warning",
|
|
])
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::null())
|
|
.stderr(Stdio::piped())
|
|
.spawn()
|
|
.context("spawn ffmpeg recorder")?;
|
|
|
|
let video_stdin = child.stdin.take().expect("stdin piped");
|
|
|
|
// Drain stderr so ffmpeg never blocks on a full pipe.
|
|
let stderr = child.stderr.take().expect("stderr piped");
|
|
let sid = session_id.to_string();
|
|
thread::Builder::new()
|
|
.name("ffmpeg-recorder-stderr".into())
|
|
.spawn(move || {
|
|
use std::io::{BufRead, BufReader};
|
|
for line in BufReader::new(stderr).lines().map_while(Result::ok) {
|
|
if !line.is_empty() {
|
|
debug!("[recorder/{sid}] {line}");
|
|
}
|
|
}
|
|
})
|
|
.expect("spawn stderr thread");
|
|
|
|
// Open audio file for raw AAC frames from client
|
|
let audio_file = File::create(&audio_path)
|
|
.map(Some)
|
|
.unwrap_or_else(|e| {
|
|
warn!("Could not create audio file: {e}");
|
|
None
|
|
});
|
|
|
|
// Scene relay: Unix socket at data/scene.sock (fixed path).
|
|
// Python always connects here — no need to discover per-session paths.
|
|
let data_dir = sessions_dir.parent().unwrap_or(sessions_dir);
|
|
let socket_path = data_dir.join(SCENE_SOCKET_NAME);
|
|
let (scene_tx, scene_rx) = tokio::sync::mpsc::channel(32);
|
|
info!("Scene relay: spawning for {}", socket_path.display());
|
|
tokio::spawn(scene_relay_task(socket_path, scene_rx));
|
|
|
|
// Tell Python which session dir to watch.
|
|
if let Err(e) = fs::write(&active_session_file, session_dir.to_str().unwrap_or("")) {
|
|
warn!("Could not write {}: {e}", active_session_file.display());
|
|
}
|
|
|
|
info!("Session {session_id}: ffmpeg pid={}, audio → {}",
|
|
child.id(), audio_path.display());
|
|
|
|
Ok(Self {
|
|
session_dir,
|
|
active_session_file,
|
|
ffmpeg: child,
|
|
video_stdin: Some(video_stdin),
|
|
audio_file,
|
|
audio_config: AdtsConfig::from_params(audio_params),
|
|
scene_tx: Some(scene_tx),
|
|
fps,
|
|
})
|
|
}
|
|
|
|
pub fn write_video(&mut self, data: &[u8], keyframe: bool) -> Result<()> {
|
|
if let Some(stdin) = &mut self.video_stdin {
|
|
stdin.write_all(data).context("write H.264 to ffmpeg")?;
|
|
}
|
|
// Best-effort relay to scene detector — drop if channel full.
|
|
if let Some(tx) = &self.scene_tx {
|
|
let _ = tx.try_send(ScenePacket { data: data.to_vec(), keyframe });
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
pub fn write_audio(&mut self, data: &[u8]) -> Result<()> {
|
|
if let Some(f) = &mut self.audio_file {
|
|
if self.audio_config.wrap {
|
|
// Client sends raw AAC — wrap with ADTS using declared params.
|
|
write_adts_frame(f, data, &self.audio_config)?;
|
|
} else {
|
|
// Client sends ADTS-framed audio — write as-is.
|
|
f.write_all(data).context("write ADTS audio")?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub fn session_dir(&self) -> &Path {
|
|
&self.session_dir
|
|
}
|
|
|
|
pub fn close(mut self) {
|
|
// Drop stdin → ffmpeg gets EOF → flushes and exits cleanly.
|
|
drop(self.video_stdin.take());
|
|
drop(self.audio_file.take());
|
|
// Drop scene_tx → relay task sees channel closed → exits.
|
|
drop(self.scene_tx.take());
|
|
match self.ffmpeg.wait() {
|
|
Ok(s) => info!("ffmpeg recorder exited: {s}"),
|
|
Err(e) => warn!("ffmpeg recorder wait error: {e}"),
|
|
}
|
|
// Clear the active session marker only if it still points to our session.
|
|
// Another session may have overwritten it if the server restarted.
|
|
if let Ok(content) = fs::read_to_string(&self.active_session_file) {
|
|
if content.trim() == self.session_dir.to_str().unwrap_or("") {
|
|
let _ = fs::remove_file(&self.active_session_file);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for Session {
|
|
fn drop(&mut self) {
|
|
if self.video_stdin.is_some() {
|
|
drop(self.video_stdin.take());
|
|
drop(self.audio_file.take());
|
|
drop(self.scene_tx.take());
|
|
let _ = self.ffmpeg.kill();
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Scene relay: serves raw H.264 over a Unix domain socket
|
|
// ---------------------------------------------------------------------------
|
|
|
|
async fn scene_relay_task(
|
|
socket_path: PathBuf,
|
|
mut rx: tokio::sync::mpsc::Receiver<ScenePacket>,
|
|
) {
|
|
// Remove stale socket from a previous session.
|
|
let _ = fs::remove_file(&socket_path);
|
|
|
|
let listener = match tokio::net::UnixListener::bind(&socket_path) {
|
|
Ok(l) => l,
|
|
Err(e) => {
|
|
warn!("Scene relay: bind failed on {}: {e}", socket_path.display());
|
|
return;
|
|
}
|
|
};
|
|
info!("Scene relay: listening on {}", socket_path.display());
|
|
|
|
let mut client: Option<tokio::net::UnixStream> = None;
|
|
// Buffer the latest keyframe so new clients start with a valid decoder state.
|
|
let mut last_keyframe: Option<Vec<u8>> = None;
|
|
|
|
loop {
|
|
if client.is_some() {
|
|
// We have a connected reader — forward data.
|
|
match rx.recv().await {
|
|
Some(pkt) => {
|
|
if pkt.keyframe {
|
|
last_keyframe = Some(pkt.data.clone());
|
|
}
|
|
let stream = client.as_mut().unwrap();
|
|
// Use a short timeout so a slow reader doesn't stall us.
|
|
// A stalled relay would queue old frames — better to drop.
|
|
let write_result = tokio::time::timeout(
|
|
std::time::Duration::from_millis(100),
|
|
stream.write_all(&pkt.data),
|
|
).await;
|
|
match write_result {
|
|
Ok(Ok(())) => {}
|
|
Ok(Err(_)) => {
|
|
info!("Scene relay: client disconnected");
|
|
client = None;
|
|
}
|
|
Err(_) => {
|
|
// Timeout — reader too slow, drop this packet.
|
|
debug!("Scene relay: slow reader, dropping packet");
|
|
}
|
|
}
|
|
}
|
|
None => break, // Channel closed, session ending.
|
|
}
|
|
} else {
|
|
// No reader — accept connections while draining the channel.
|
|
tokio::select! {
|
|
biased;
|
|
result = listener.accept() => {
|
|
match result {
|
|
Ok((mut stream, _)) => {
|
|
info!("Scene relay: client connected");
|
|
// Send the last keyframe so the decoder can initialize.
|
|
if let Some(ref kf) = last_keyframe {
|
|
if stream.write_all(kf).await.is_err() {
|
|
warn!("Scene relay: failed to send keyframe");
|
|
continue;
|
|
}
|
|
info!("Scene relay: sent keyframe ({} bytes)", kf.len());
|
|
}
|
|
client = Some(stream);
|
|
}
|
|
Err(e) => warn!("Scene relay: accept error: {e}"),
|
|
}
|
|
}
|
|
pkt = rx.recv() => {
|
|
match pkt {
|
|
Some(pkt) => {
|
|
if pkt.keyframe {
|
|
last_keyframe = Some(pkt.data);
|
|
}
|
|
// Discard — no reader connected.
|
|
}
|
|
None => break, // Channel closed.
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
drop(client);
|
|
let _ = fs::remove_file(&socket_path);
|
|
info!("Scene relay: stopped");
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ADTS header for raw AAC framing
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Write a raw AAC frame wrapped in a 7-byte ADTS header.
|
|
///
|
|
/// Uses sample rate and channel count from the session's AudioParams
|
|
/// rather than hardcoded values, so any backend can declare its format.
|
|
fn write_adts_frame(w: &mut impl Write, aac_data: &[u8], cfg: &AdtsConfig) -> Result<()> {
|
|
const PROFILE_MINUS1: u8 = 1; // AAC-LC (object_type 2, stored as 2-1=1)
|
|
|
|
let sr_idx = cfg.sr_idx;
|
|
let ch_cfg = cfg.ch_cfg;
|
|
let frame_len = (aac_data.len() + 7) as u16;
|
|
|
|
let header: [u8; 7] = [
|
|
// byte 0-1: syncword(12) | ID(1)=0(MPEG4) | layer(2)=0 | protection(1)=1(no CRC)
|
|
0xFF,
|
|
0xF1,
|
|
// byte 2: profile(2) | sr_idx(4) | private(1)=0 | ch_cfg[2](1)
|
|
(PROFILE_MINUS1 << 6) | (sr_idx << 2) | ((ch_cfg >> 2) & 1),
|
|
// byte 3: ch_cfg[1:0](2) | orig(1)=0 | home(1)=0 | copyright_id(1)=0 | copyright_start(1)=0 | frame_len[12:11](2)
|
|
((ch_cfg & 3) << 6) | ((frame_len >> 11) as u8 & 0x03),
|
|
// byte 4: frame_len[10:3](8)
|
|
((frame_len >> 3) & 0xFF) as u8,
|
|
// byte 5: frame_len[2:0](3) | buffer_fullness[10:6](5)
|
|
((frame_len & 0x07) << 5) as u8 | 0x1F,
|
|
// byte 6: buffer_fullness[5:0](6) | num_aac_frames_minus1(2)=0
|
|
0xFC,
|
|
];
|
|
|
|
w.write_all(&header).context("ADTS header")?;
|
|
w.write_all(aac_data).context("AAC frame")?;
|
|
Ok(())
|
|
}
|