Genarrative/server-rs/crates/platform-audio/src/mp3.rs

use std::io::Cursor;

use bytes::Bytes;
use symphonia::{
    core::{
        codecs::CODEC_TYPE_MP3,
        errors::Error as SymphoniaError,
        formats::FormatOptions,
        io::{MediaSourceStream, MediaSourceStreamOptions},
        meta::MetadataOptions,
        probe::Hint,
    },
    default::get_probe,
};

pub const MAX_GENERATED_AUDIO_DURATION_SECONDS: f64 = 600.0;

pub(crate) fn probe_mp3_duration_seconds(audio_bytes: Bytes) -> Result<f64, String> {
    if audio_bytes.is_empty() {
        return Err("MP3 内容为空".to_string());
    }

    let source = Box::new(Cursor::new(audio_bytes));
    let media_source_stream = MediaSourceStream::new(
        source,
        MediaSourceStreamOptions {
            buffer_len: 64 * 1024,
        },
    );
    let mut hint = Hint::new();
    hint.with_extension("mp3");
    let probed = get_probe()
        .format(
            &hint,
            media_source_stream,
            &FormatOptions::default(),
            &MetadataOptions::default(),
        )
        .map_err(|error| format!("无法识别 MP3：{error}"))?;
    let mut format = probed.format;
    let track = format
        .default_track()
        .filter(|track| track.codec_params.codec == CODEC_TYPE_MP3)
        .ok_or_else(|| "响应不是 MP3 音频".to_string())?;
    let track_id = track.id;
    let time_base = track
        .codec_params
        .time_base
        .ok_or_else(|| "MP3 缺少时间基准".to_string())?;
    let mut end_timestamp = 0_u64;

    loop {
        match format.next_packet() {
            Ok(packet) => {
                if packet.track_id() == track_id {
                    end_timestamp = end_timestamp.max(packet.ts().saturating_add(packet.dur()));
                }
            }
            Err(SymphoniaError::IoError(error))
                if error.kind() == std::io::ErrorKind::UnexpectedEof =>
            {
                break;
            }
            Err(error) => return Err(format!("读取 MP3 帧失败：{error}")),
        }
    }

    if end_timestamp == 0 {
        return Err("MP3 不包含可用音频帧".to_string());
    }
    let duration = time_base.calc_time(end_timestamp);
    validate_generated_audio_duration_seconds(duration.seconds as f64 + duration.frac)
}

pub(crate) fn validate_generated_audio_duration_seconds(
    duration_seconds: f64,
) -> Result<f64, String> {
    if duration_seconds.is_finite()
        && duration_seconds > 0.0
        && duration_seconds <= MAX_GENERATED_AUDIO_DURATION_SECONDS
    {
        return Ok(duration_seconds);
    }
    Err(format!(
        "MP3 实际时长必须是有限正数且不超过 {MAX_GENERATED_AUDIO_DURATION_SECONDS} 秒"
    ))
}

#[cfg(test)]
mod tests {
    use base64::Engine;

    use super::*;

    const REAL_MP3_FRAME_OFFSET: usize = 1_374;
    const REAL_MP3_FRAME_LENGTH: usize = 156;

    fn repeated_real_mp3_frame(frame_count: usize) -> Bytes {
        let fixture = base64::engine::general_purpose::STANDARD
            .decode(include_str!("../tests/fixtures/vbr-id3.mp3.base64").trim())
            .expect("VBR MP3 fixture should decode");
        let frame = &fixture[REAL_MP3_FRAME_OFFSET..REAL_MP3_FRAME_OFFSET + REAL_MP3_FRAME_LENGTH];
        assert_eq!(&frame[..4], &[0xff, 0xfb, 0x30, 0xc4]);

        Bytes::from(frame.repeat(frame_count))
    }

    #[test]
    fn duration_guard_uses_only_the_independent_six_hundred_second_limit() {
        for duration in [0.001, 30.5, 60.0, 600.0] {
            assert_eq!(
                validate_generated_audio_duration_seconds(duration)
                    .expect("finite positive durations through 600 seconds should pass"),
                duration
            );
        }
        for duration in [0.0, -1.0, 600.000_001, f64::NAN, f64::INFINITY] {
            assert!(validate_generated_audio_duration_seconds(duration).is_err());
        }
    }

    #[test]
    fn probe_reads_the_tracked_id3_mp3_fixture() {
        let fixture = Bytes::from_static(include_bytes!(
            "../../../../public/wooden-fish/default-hit-sound.mp3"
        ));
        let duration = probe_mp3_duration_seconds(fixture)
            .expect("tracked ID3 MP3 fixture should have a duration");

        assert!((0.65..=0.75).contains(&duration), "duration={duration}");
    }

    #[test]
    fn probe_reads_vbr_mp3_with_id3_and_encoder_padding() {
        let fixture = base64::engine::general_purpose::STANDARD
            .decode(include_str!("../tests/fixtures/vbr-id3.mp3.base64").trim())
            .expect("VBR MP3 fixture should decode");
        let duration = probe_mp3_duration_seconds(Bytes::from(fixture))
            .expect("VBR MP3 fixture should have a duration");

        assert!((0.39..=0.46).contains(&duration), "duration={duration}");
    }

    #[test]
    fn probe_enforces_the_six_hundred_second_boundary_on_real_mp3_frames() {
        let accepted_duration = probe_mp3_duration_seconds(repeated_real_mp3_frame(22_968))
            .expect("real MP3 just below 600 seconds should pass");
        assert!(
            (599.9..=600.0).contains(&accepted_duration),
            "duration={accepted_duration}"
        );

        let error = probe_mp3_duration_seconds(repeated_real_mp3_frame(22_969))
            .expect_err("real MP3 just above 600 seconds should fail");
        assert!(error.contains("不超过 600 秒"), "error={error}");
    }

    #[test]
    fn probe_rejects_non_mp3_and_empty_payloads() {
        for payload in [Bytes::new(), Bytes::from_static(b"<html>error</html>")] {
            assert!(probe_mp3_duration_seconds(payload).is_err());
        }
    }
}