<commit_before>//! Writes audio samples to files. (Reads too, this module needs a rename)
//!
//! ### Audacity PCM import settings
//!
//! * File > Import > Raw Data...
//! * Signed 16-bit PCM
//! * Little-endian
//! * 1 Channel (Mono)
//! * Sample rate: 44_100Hz (or whatever your samples generated have)

use std::fs::OpenOptions;
use std::io::{BufReader, Error, Read, Result, Write};
use std::path::Path;

use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};

/// Creates a file at `filename` and writes a bunch of `&[i16]` samples to it as a PCM file.
/// See module documentation for PCM settings.
///
/// ```
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_pcm_file;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// write_pcm_file(
///     "out/sine.wav",
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_pcm_file(filename: &str, samples: &[i16]) -> Result<()> {
    let path = Path::new(filename);
    let mut f = OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(&path)?;
    write_pcm(&mut f, samples)
}

/// Writes a bunch of `&[i16]` samples to a `Write` as raw PCM.
/// See module documentation for PCM settings.
///
/// Also see `synthrs::writer::write_pcm_file`.
///
/// ```
/// use std::io::Cursor;
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_pcm;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// let mut output_buffer: Vec<u8> = Vec::new();
/// let mut output_writer = Cursor::new(output_buffer);
///
/// // You can use whatever implements `Write`, such as a `File`.
/// write_pcm(
///     &mut output_writer,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_pcm<W>(writer: &mut W, samples: &[i16]) -> Result<()>
where
    W: Write,
{
    for &sample in samples.iter() {
        writer.write_i16::<LittleEndian>(sample)?;
    }

    Ok(())
}

/// Creates a file at `filename` and writes a bunch of `&[i16]` samples to it as a WAVE file
/// ```
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_wav_file;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// write_wav_file(
///     "out/sine.wav",
///     44_100,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_wav_file(filename: &str, sample_rate: usize, samples: &[i16]) -> Result<()> {
    let path = Path::new(filename);
    let mut f = OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(&path)?;
    write_wav(&mut f, sample_rate, samples)
}

/// Writes a bunch of `&[i16]` samples to a `Write`. Also see `synthrs::writer::write_wav_file`.
/// ```
/// use std::io::Cursor;
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_wav;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// let mut output_buffer: Vec<u8> = Vec::new();
/// let mut output_writer = Cursor::new(output_buffer);
///
/// // You can use whatever implements `Write`, such as a `File`.
/// write_wav(
///     &mut output_writer,
///     44_100,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_wav<W>(writer: &mut W, sample_rate: usize, samples: &[i16]) -> Result<()>
where
    W: Write,
{
    // See: http://www-mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html
    // Some WAV header fields
    let channels = 1;
    let bit_depth = 16;
    let subchunk_2_size = samples.len() * channels * bit_depth / 8;
    let chunk_size = 36 + subchunk_2_size as i32;
    let byte_rate = (sample_rate * channels * bit_depth / 8) as i32;
    let block_align = (channels * bit_depth / 8) as i16;

    writer.write_i32::<BigEndian>(0x5249_4646)?; // ChunkID, RIFF
    writer.write_i32::<LittleEndian>(chunk_size)?; // ChunkSize
    writer.write_i32::<BigEndian>(0x5741_5645)?; // Format, WAVE

    writer.write_i32::<BigEndian>(0x666d_7420)?; // Subchunk1ID, fmt
    writer.write_i32::<LittleEndian>(16)?; // Subchunk1Size, 16 for PCM
    writer.write_i16::<LittleEndian>(1)?; // AudioFormat, PCM = 1 (linear quantization)
    writer.write_i16::<LittleEndian>(channels as i16)?; // NumChannels
    writer.write_i32::<LittleEndian>(sample_rate as i32)?; // SampleRate
    writer.write_i32::<LittleEndian>(byte_rate)?; // ByteRate
    writer.write_i16::<LittleEndian>(block_align)?; // BlockAlign
    writer.write_i16::<LittleEndian>(bit_depth as i16)?; // BitsPerSample

    writer.write_i32::<BigEndian>(0x6461_7461)?; // Subchunk2ID, data
    writer.write_i32::<LittleEndian>(subchunk_2_size as i32)?; // Subchunk2Size, number of bytes in the data

    for sample in samples {
        writer.write_i16::<LittleEndian>(*sample)?
    }

    Ok(())
}

// Borrowing of packed &wave.pcm is unsafe
// #[repr(C, packed)]
/// Representation of a WAV file. Does not contain fields for extended WAV formats.
#[repr(C)]
#[derive(Debug, Clone)]
pub struct Wave {
    pub chunk_id: i32,
    pub chunk_size: i32,
    pub format: i32,
    pub subchunk_1_id: i32,
    pub subchunk_1_size: i32,
    /// 1 = PCM
    pub audio_format: i16,
    pub num_channels: i16,
    pub sample_rate: i32,
    pub byte_rate: i32,
    pub block_align: i16,
    pub bits_per_sample: i16,
    pub subchunk_2_id: i32,
    pub subchunk_2_size: i32,
    pub pcm: Vec<i16>,
}

/// Reads a wave file given a file path. Convenience wrapper around `crate::writer::read_wav_file`.
/// ```
/// use synthrs::writer;
///
/// let wave = writer::read_wav_file("./tests/assets/sine.wav").unwrap();
/// assert_eq!(wave.num_channels, 1);
/// ```
pub fn read_wav_file(filename: &str) -> Result<Wave> {
    let path = Path::new(filename);
    let file = OpenOptions::new().read(true).open(&path)?;
    let mut reader = BufReader::new(file);
    read_wav(&mut reader)
}

/// Reads a wave file. Only supports mono 16-bit, little-endian, signed PCM WAV files
///
/// ### Useful commands:
///
/// * Use `ffmpeg -i .\example.wav` to inspect a wav
/// * Use `ffmpeg -i "dirty.wav" -f wav -flags +bitexact -acodec pcm_s16le -ar 44100 -ac 1 "clean.wav"` to clean a WAV
///
/// ```
/// use std::path::Path;
/// use std::fs::OpenOptions;
/// use std::io::BufReader;
/// use synthrs::writer;
///
/// let path = Path::new("./tests/assets/sine.wav");
/// let file = OpenOptions::new().read(true).open(&path).unwrap();
/// let mut reader = BufReader::new(file);
///
/// let wave = writer::read_wav(&mut reader).unwrap();
/// assert_eq!(wave.num_channels, 1);
/// ```
// For -acodec values, see https://trac.ffmpeg.org/wiki/audio%20types
pub fn read_wav<R>(reader: &mut R) -> Result<Wave>
where
    R: Read,
{
    let chunk_id = reader.read_i32::<BigEndian>()?; // ChunkID, RIFF
    let chunk_size = reader.read_i32::<LittleEndian>()?; // ChunkSize
    let format = reader.read_i32::<BigEndian>()?; // Format, WAVE
    if format != 0x5741_5645 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            "file is not a WAV".to_string(),
        ));
    }

    let subchunk_1_id = reader.read_i32::<BigEndian>()?; // Subchunk1ID, fmt
    let subchunk_1_size = reader.read_i32::<LittleEndian>()?; // Subchunk1Size, Chunk size: 16, 18 or 40

    let audio_format = reader.read_i16::<LittleEndian>()?; // AudioFormat, PCM = 1 (linear quantization)
    if audio_format != 1 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            format!(
                "only 16-bit little-endian signed PCM WAV supported, audio_format: {}",
                audio_format
            ),
        ));
    }

    let num_channels = reader.read_i16::<LittleEndian>()?; // NumChannels
    let sample_rate = reader.read_i32::<LittleEndian>()?; // SampleRate
    let byte_rate = reader.read_i32::<LittleEndian>()?; // ByteRate
    let block_align = reader.read_i16::<LittleEndian>()?; // BlockAlign
    let bits_per_sample = reader.read_i16::<LittleEndian>()?; // BitsPerSample

    let extra_bytes = subchunk_1_size - 16;

    if extra_bytes > 0 {
        let extension_size = reader.read_i16::<LittleEndian>()?; // Size of the extension (0 or 22)
        if extension_size == 22 {
            // TODO: Add this to `Wave` as optionals
            let _valid_bits_per_sample = reader.read_i16::<LittleEndian>()?;
            let _speaker_mask = reader.read_i16::<LittleEndian>()?;
            let _subformat = reader.read_i16::<BigEndian>()?;
        } else if extension_size != 0 {
            return Err(Error::new(
                std::io::ErrorKind::InvalidInput,
                format!(
                    "unexpected fmt chunk extension size: should be 0 or 22 but got: {}",
                    extension_size
                ),
            ));
        }
    }

    let subchunk_2_id = reader.read_i32::<BigEndian>()?; // Subchunk2ID, data
    if subchunk_2_id != 0x6461_7461 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            format!(
                "unexpected subchunk name: expecting data, got: {}",
                subchunk_2_id
            ),
        ));
    }
    let subchunk_2_size = reader.read_i32::<LittleEndian>()?; // Subchunk2Size, number of bytes in the data

    let mut pcm: Vec<i16> = Vec::with_capacity(2 * subchunk_2_size as usize);

    // `reader.read_into_i16(&pcm)` doesn't seem to work here due to bad input?
    // It just does nothing and &pcm is left empty after that.
    for _i in 0..subchunk_2_size {
        let sample = reader.read_i16::<LittleEndian>().unwrap_or(0);
        pcm.push(sample);
    }

    let wave = Wave {
        audio_format,
        bits_per_sample,
        block_align,
        byte_rate,
        chunk_id,
        chunk_size,
        format,
        num_channels,
        pcm,
        sample_rate,
        subchunk_1_id,
        subchunk_1_size,
        subchunk_2_id,
        subchunk_2_size,
    };

    Ok(wave)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_read_wav() {
        // 1 second at 44,100Hz, PCM 16 (2-byte)
        let wave = read_wav_file("./tests/assets/sine.wav").unwrap();

        assert_eq!(wave.format, 0x5741_5645); // WAVE
        assert_eq!(wave.audio_format, 1);
        assert_eq!(wave.num_channels, 1);
        assert_eq!(wave.sample_rate, 44_100);
        assert_eq!(wave.byte_rate, 88_200);
        assert_eq!(wave.block_align, 2);
        assert_eq!(wave.bits_per_sample, 16);
        assert_eq!(wave.subchunk_2_size, 88_200);
        assert_eq!(wave.pcm.len(), 88_200);
    }

    #[test]
    fn test_write_read_wav() {
        use crate::synthesizer::{make_samples, quantize_samples};
        use crate::wave::sine_wave;
        use crate::writer::write_wav;
        use std::io::{Cursor, Seek, SeekFrom};

        let output_buffer: Vec<u8> = Vec::new();
        let mut output_writer = Cursor::new(output_buffer);

        write_wav(
            &mut output_writer,
            44_100,
            &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
        )
        .unwrap();

        let _ = output_writer.seek(SeekFrom::Start(0));
        let wave = read_wav(&mut output_writer).unwrap();
        assert_eq!(wave.format, 0x5741_5645); // WAVE
        assert_eq!(wave.audio_format, 1);
        assert_eq!(wave.num_channels, 1);
        assert_eq!(wave.sample_rate, 44_100);
        assert_eq!(wave.byte_rate, 88_200);
        assert_eq!(wave.block_align, 2);
        assert_eq!(wave.bits_per_sample, 16);
        assert_eq!(wave.subchunk_2_size, 8820);
        assert_eq!(wave.pcm.len(), 8820);
    }
}
<commit_msg>Fix new cargo clippy warning<commit_after>//! Writes audio samples to files. (Reads too, this module needs a rename)
//!
//! ### Audacity PCM import settings
//!
//! * File > Import > Raw Data...
//! * Signed 16-bit PCM
//! * Little-endian
//! * 1 Channel (Mono)
//! * Sample rate: 44_100Hz (or whatever your samples generated have)

use std::fs::OpenOptions;
use std::io::{BufReader, Error, Read, Result, Write};
use std::path::Path;

use byteorder::{BigEndian, LittleEndian, ReadBytesExt, WriteBytesExt};

/// Creates a file at `filename` and writes a bunch of `&[i16]` samples to it as a PCM file.
/// See module documentation for PCM settings.
///
/// ```
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_pcm_file;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// write_pcm_file(
///     "out/sine.wav",
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_pcm_file(filename: &str, samples: &[i16]) -> Result<()> {
    let path = Path::new(filename);
    let mut f = OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(&path)?;
    write_pcm(&mut f, samples)
}

/// Writes a bunch of `&[i16]` samples to a `Write` as raw PCM.
/// See module documentation for PCM settings.
///
/// Also see `synthrs::writer::write_pcm_file`.
///
/// ```
/// use std::io::Cursor;
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_pcm;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// let mut output_buffer: Vec<u8> = Vec::new();
/// let mut output_writer = Cursor::new(output_buffer);
///
/// // You can use whatever implements `Write`, such as a `File`.
/// write_pcm(
///     &mut output_writer,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_pcm<W>(writer: &mut W, samples: &[i16]) -> Result<()>
where
    W: Write,
{
    for &sample in samples.iter() {
        writer.write_i16::<LittleEndian>(sample)?;
    }

    Ok(())
}

/// Creates a file at `filename` and writes a bunch of `&[i16]` samples to it as a WAVE file
/// ```
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_wav_file;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// write_wav_file(
///     "out/sine.wav",
///     44_100,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_wav_file(filename: &str, sample_rate: usize, samples: &[i16]) -> Result<()> {
    let path = Path::new(filename);
    let mut f = OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(&path)?;
    write_wav(&mut f, sample_rate, samples)
}

/// Writes a bunch of `&[i16]` samples to a `Write`. Also see `synthrs::writer::write_wav_file`.
/// ```
/// use std::io::Cursor;
/// use synthrs::wave::sine_wave;
/// use synthrs::writer::write_wav;
/// use synthrs::synthesizer::{quantize_samples, make_samples};
///
/// let mut output_buffer: Vec<u8> = Vec::new();
/// let mut output_writer = Cursor::new(output_buffer);
///
/// // You can use whatever implements `Write`, such as a `File`.
/// write_wav(
///     &mut output_writer,
///     44_100,
///     &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
/// ).expect("failed to write wav");
/// ```
pub fn write_wav<W>(writer: &mut W, sample_rate: usize, samples: &[i16]) -> Result<()>
where
    W: Write,
{
    // See: http://www-mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html
    // Some WAV header fields
    let channels = 1;
    let bit_depth = 16;
    let subchunk_2_size = samples.len() * channels * bit_depth / 8;
    let chunk_size = 36 + subchunk_2_size as i32;
    let byte_rate = (sample_rate * channels * bit_depth / 8) as i32;
    let block_align = (channels * bit_depth / 8) as i16;

    writer.write_i32::<BigEndian>(0x5249_4646)?; // ChunkID, RIFF
    writer.write_i32::<LittleEndian>(chunk_size)?; // ChunkSize
    writer.write_i32::<BigEndian>(0x5741_5645)?; // Format, WAVE

    writer.write_i32::<BigEndian>(0x666d_7420)?; // Subchunk1ID, fmt
    writer.write_i32::<LittleEndian>(16)?; // Subchunk1Size, 16 for PCM
    writer.write_i16::<LittleEndian>(1)?; // AudioFormat, PCM = 1 (linear quantization)
    writer.write_i16::<LittleEndian>(channels as i16)?; // NumChannels
    writer.write_i32::<LittleEndian>(sample_rate as i32)?; // SampleRate
    writer.write_i32::<LittleEndian>(byte_rate)?; // ByteRate
    writer.write_i16::<LittleEndian>(block_align)?; // BlockAlign
    writer.write_i16::<LittleEndian>(bit_depth as i16)?; // BitsPerSample

    writer.write_i32::<BigEndian>(0x6461_7461)?; // Subchunk2ID, data
    writer.write_i32::<LittleEndian>(subchunk_2_size as i32)?; // Subchunk2Size, number of bytes in the data

    for sample in samples {
        writer.write_i16::<LittleEndian>(*sample)?
    }

    Ok(())
}

// Borrowing of packed &wave.pcm is unsafe
// #[repr(C, packed)]
/// Representation of a WAV file. Does not contain fields for extended WAV formats.
#[repr(C)]
#[derive(Debug, Clone)]
pub struct Wave {
    pub chunk_id: i32,
    pub chunk_size: i32,
    pub format: i32,
    pub subchunk_1_id: i32,
    pub subchunk_1_size: i32,
    /// 1 = PCM
    pub audio_format: i16,
    pub num_channels: i16,
    pub sample_rate: i32,
    pub byte_rate: i32,
    pub block_align: i16,
    pub bits_per_sample: i16,
    pub subchunk_2_id: i32,
    pub subchunk_2_size: i32,
    pub pcm: Vec<i16>,
}

/// Reads a wave file given a file path. Convenience wrapper around `crate::writer::read_wav_file`.
/// ```
/// use synthrs::writer;
///
/// let wave = writer::read_wav_file("./tests/assets/sine.wav").unwrap();
/// assert_eq!(wave.num_channels, 1);
/// ```
pub fn read_wav_file(filename: &str) -> Result<Wave> {
    let path = Path::new(filename);
    let file = OpenOptions::new().read(true).open(&path)?;
    let mut reader = BufReader::new(file);
    read_wav(&mut reader)
}

/// Reads a wave file. Only supports mono 16-bit, little-endian, signed PCM WAV files
///
/// ### Useful commands:
///
/// * Use `ffmpeg -i .\example.wav` to inspect a wav
/// * Use `ffmpeg -i "dirty.wav" -f wav -flags +bitexact -acodec pcm_s16le -ar 44100 -ac 1 "clean.wav"` to clean a WAV
///
/// ```
/// use std::path::Path;
/// use std::fs::OpenOptions;
/// use std::io::BufReader;
/// use synthrs::writer;
///
/// let path = Path::new("./tests/assets/sine.wav");
/// let file = OpenOptions::new().read(true).open(&path).unwrap();
/// let mut reader = BufReader::new(file);
///
/// let wave = writer::read_wav(&mut reader).unwrap();
/// assert_eq!(wave.num_channels, 1);
/// ```
// For -acodec values, see https://trac.ffmpeg.org/wiki/audio%20types
pub fn read_wav<R>(reader: &mut R) -> Result<Wave>
where
    R: Read,
{
    let chunk_id = reader.read_i32::<BigEndian>()?; // ChunkID, RIFF
    let chunk_size = reader.read_i32::<LittleEndian>()?; // ChunkSize
    let format = reader.read_i32::<BigEndian>()?; // Format, WAVE
    if format != 0x5741_5645 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            "file is not a WAV".to_string(),
        ));
    }

    let subchunk_1_id = reader.read_i32::<BigEndian>()?; // Subchunk1ID, fmt
    let subchunk_1_size = reader.read_i32::<LittleEndian>()?; // Subchunk1Size, Chunk size: 16, 18 or 40

    let audio_format = reader.read_i16::<LittleEndian>()?; // AudioFormat, PCM = 1 (linear quantization)
    if audio_format != 1 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            format!(
                "only 16-bit little-endian signed PCM WAV supported, audio_format: {}",
                audio_format
            ),
        ));
    }

    let num_channels = reader.read_i16::<LittleEndian>()?; // NumChannels
    let sample_rate = reader.read_i32::<LittleEndian>()?; // SampleRate
    let byte_rate = reader.read_i32::<LittleEndian>()?; // ByteRate
    let block_align = reader.read_i16::<LittleEndian>()?; // BlockAlign
    let bits_per_sample = reader.read_i16::<LittleEndian>()?; // BitsPerSample

    let extra_bytes = subchunk_1_size - 16;

    if extra_bytes > 0 {
        let extension_size = reader.read_i16::<LittleEndian>()?; // Size of the extension (0 or 22)
        if extension_size == 22 {
            // TODO: Add this to `Wave` as optionals
            let _valid_bits_per_sample = reader.read_i16::<LittleEndian>()?;
            let _speaker_mask = reader.read_i16::<LittleEndian>()?;
            let _subformat = reader.read_i16::<BigEndian>()?;
        } else if extension_size != 0 {
            return Err(Error::new(
                std::io::ErrorKind::InvalidInput,
                format!(
                    "unexpected fmt chunk extension size: should be 0 or 22 but got: {}",
                    extension_size
                ),
            ));
        }
    }

    let subchunk_2_id = reader.read_i32::<BigEndian>()?; // Subchunk2ID, data
    if subchunk_2_id != 0x6461_7461 {
        return Err(Error::new(
            std::io::ErrorKind::InvalidInput,
            format!(
                "unexpected subchunk name: expecting data, got: {}",
                subchunk_2_id
            ),
        ));
    }
    let subchunk_2_size = reader.read_i32::<LittleEndian>()?; // Subchunk2Size, number of bytes in the data

    let mut pcm: Vec<i16> = Vec::with_capacity(2 * subchunk_2_size as usize);

    // `reader.read_into_i16(&pcm)` doesn't seem to work here due to bad input?
    // It just does nothing and &pcm is left empty after that.
    for _i in 0..subchunk_2_size {
        let sample = reader.read_i16::<LittleEndian>().unwrap_or(0);
        pcm.push(sample);
    }

    let wave = Wave {
        chunk_id,
        chunk_size,
        format,
        subchunk_1_id,
        subchunk_1_size,
        audio_format,
        num_channels,
        sample_rate,
        byte_rate,
        block_align,
        bits_per_sample,
        subchunk_2_id,
        subchunk_2_size,
        pcm,
    };

    Ok(wave)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_read_wav() {
        // 1 second at 44,100Hz, PCM 16 (2-byte)
        let wave = read_wav_file("./tests/assets/sine.wav").unwrap();

        assert_eq!(wave.format, 0x5741_5645); // WAVE
        assert_eq!(wave.audio_format, 1);
        assert_eq!(wave.num_channels, 1);
        assert_eq!(wave.sample_rate, 44_100);
        assert_eq!(wave.byte_rate, 88_200);
        assert_eq!(wave.block_align, 2);
        assert_eq!(wave.bits_per_sample, 16);
        assert_eq!(wave.subchunk_2_size, 88_200);
        assert_eq!(wave.pcm.len(), 88_200);
    }

    #[test]
    fn test_write_read_wav() {
        use crate::synthesizer::{make_samples, quantize_samples};
        use crate::wave::sine_wave;
        use crate::writer::write_wav;
        use std::io::{Cursor, Seek, SeekFrom};

        let output_buffer: Vec<u8> = Vec::new();
        let mut output_writer = Cursor::new(output_buffer);

        write_wav(
            &mut output_writer,
            44_100,
            &quantize_samples::<i16>(&make_samples(0.1, 44_100, sine_wave(440.0))),
        )
        .unwrap();

        let _ = output_writer.seek(SeekFrom::Start(0));
        let wave = read_wav(&mut output_writer).unwrap();
        assert_eq!(wave.format, 0x5741_5645); // WAVE
        assert_eq!(wave.audio_format, 1);
        assert_eq!(wave.num_channels, 1);
        assert_eq!(wave.sample_rate, 44_100);
        assert_eq!(wave.byte_rate, 88_200);
        assert_eq!(wave.block_align, 2);
        assert_eq!(wave.bits_per_sample, 16);
        assert_eq!(wave.subchunk_2_size, 8820);
        assert_eq!(wave.pcm.len(), 8820);
    }
}
<|endoftext|>


<commit_before>#![feature(macro_rules)]
use std::any::{Any, TypeId};
use std::ops::{Add, Sub};

struct DistanceMatrix<T> {
    length: usize,
    inner_max: Vec<Vec<Option<T>>>,
}

type Matrix<T> = Vec<Vec<Option<T>>>;

struct nil;

impl<T> DistanceMatrix<T>
where
    T: Add + Sub,
{
    fn new(input: Matrix<T>) -> Self {
        DistanceMatrix {
            length: input.len(),
            inner_max: input,
        }
    }
}

fn is_nil<T: ?Sized + Any>(o: &T) -> bool {
    TypeId::of::<i32>() == TypeId::of::<T>()
}

// new_distance_matrix!(3,[[1,2,nil],[2,3,nil],[4,nil,6]])
macro_rules! new_distance_matrix {
    ($([$($x:expr),*]),*) => {{
        vec![$(vec![$(
            if is_nil($x) {
                None
            }else{
                Some($x)
            }
        ),*]),*]
    }};
    ($($x:expr),*) => {{

    }};
}

fn main() {
    let test0 = new_distance_matrix![[&1, &2, &3], [&4, &5, &6], [&7, &8, &_]];
    println!("{:?}", test0);
    //println!("{:?}", new_distance_matrix![[1, 2, 3], [4, 5, 6], [7, 8]]);
    println!("{:?}", is_nil(&nil {}));
}
<commit_msg>a bit<commit_after>//#![feature(macro_rules)]
//use std::any::{Any, TypeId};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::ops::{Add, Sub};

/*
{start: }
*/
type Matrix<T> = HashMap<&'static str, HashMap<&'static str, Option<T>>>;

struct DistanceMatrix<T> {
    length: usize,
    inner_max: Matrix<T>,
}

impl<T> DistanceMatrix<T>
where
    T: Add + Sub + Ord,
{
    fn new(input: Matrix<T>) -> Self {
        DistanceMatrix {
            length: input.len(),
            inner_max: input,
        }
    }

    fn new_dijkstra_on(&mut self, end: &'static str) {}
}

//struct nil;
// type check function template
/*
fn is_nil<T: ?Sized + Any>(o: &T) -> bool {
    TypeId::of::<i32>() == TypeId::of::<T>()
}*/

// new_distance_matrix!()
// cannot check type in macro body. so do not need this macro
macro_rules! new_distance_matrix {
    ($([$($x:expr),*]),*) => {{
        vec![$(vec![$(
            if is_nil($x) {
                None
            }else{
                Some($x)
            }
        ),*]),*]
    }};
    ($($x:expr),*) => {{

    }};
}

fn main() {
    let test0: Matrix<u32> = [(
        "a",
        [("b", Some(6)), ("c", Some(3))] as HashMap<&'static str, Option<u32>>,
    )];
    println!("{:?}", test0);
    //println!("{:?}", new_distance_matrix![[1, 2, 3], [4, 5, 6], [7, 8]]);
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

extern crate cc;
extern crate glsl_to_cxx;
extern crate webrender_build;

use std::collections::HashSet;
use std::fmt::Write;
use webrender_build::shader::{ShaderFeatureFlags, get_shader_features};

// Shader key is in "name feature,feature" format.
// File name needs to be formatted as "name_feature_feature".
fn shader_file(shader_key: &str) -> String {
    shader_key.replace(' ', "_").replace(',', "_")
}

fn write_load_shader(shader_keys: &[String]) {
    let mut load_shader = String::new();
    for s in shader_keys {
        let _ = write!(load_shader, "#include \"{}.h\"\n", shader_file(s));
    }
    load_shader.push_str("ProgramLoader load_shader(const char* name) {\n");
    for s in shader_keys {
        let _ = write!(load_shader, "  if (!strcmp(name, \"{}\")) {{ return {}_program::loader; }}\n",
                       s, shader_file(s));
    }
    load_shader.push_str("  return nullptr;\n}\n");
    std::fs::write(std::env::var("OUT_DIR").unwrap() + "/load_shader.h", load_shader).unwrap();
}

fn process_imports(shader_dir: &str, shader: &str, included: &mut HashSet<String>, output: &mut String) {
    if !included.insert(shader.into()) {
        return;
    }
    println!("cargo:rerun-if-changed={}/{}.glsl", shader_dir, shader);
    let source = std::fs::read_to_string(format!("{}/{}.glsl", shader_dir, shader)).unwrap();
    for line in source.lines() {
        if line.starts_with("#include ") {
            let imports = line["#include ".len() ..].split(',');
            for import in imports {
                process_imports(shader_dir, import, included, output);
            }
        } else if line.starts_with("#version ") || line.starts_with("#extension ") {
            // ignore
        } else {
            output.push_str(line);
            output.push('\n');
        }
    }
}

fn translate_shader(shader_key: &str, shader_dir: &str) {
    let mut imported = String::from("#define SWGL 1\n#define __VERSION__ 150\n");
    let _ = write!(imported, "#define WR_MAX_VERTEX_TEXTURE_WIDTH {}U\n",
                   webrender_build::MAX_VERTEX_TEXTURE_WIDTH);

    let (basename, features) =
        shader_key.split_at(shader_key.find(' ').unwrap_or(shader_key.len()));
    if !features.is_empty() {
        for feature in features.trim().split(',') {
            let _ = write!(imported, "#define WR_FEATURE_{}\n", feature);
        }
    }

    process_imports(shader_dir, basename, &mut HashSet::new(), &mut imported);

    let shader = shader_file(shader_key);

    let out_dir = std::env::var("OUT_DIR").unwrap();
    let imp_name = format!("{}/{}.c", out_dir, shader);
    std::fs::write(&imp_name, imported).unwrap();

    let mut build = cc::Build::new();
    build.no_default_flags(true);
    if build.get_compiler().is_like_msvc() {
        build.flag("/EP").flag("/u");
    } else {
        build.flag("-xc").flag("-P").flag("-undef");
    }
    build.file(&imp_name);
    let vs = build.clone()
        .define("WR_VERTEX_SHADER", Some("1"))
        .expand();
    let fs = build.clone()
        .define("WR_FRAGMENT_SHADER", Some("1"))
        .expand();
    let vs_name = format!("{}/{}.vert", out_dir, shader);
    let fs_name = format!("{}/{}.frag", out_dir, shader);
    std::fs::write(&vs_name, vs).unwrap();
    std::fs::write(&fs_name, fs).unwrap();

    let args = vec![
        "glsl_to_cxx".to_string(),
        vs_name,
        fs_name,
    ];
    let result = glsl_to_cxx::translate(&mut args.into_iter());
    std::fs::write(format!("{}/{}.h", out_dir, shader), result).unwrap();
}

fn main() {
    let shader_dir = match std::env::var("MOZ_SRC") {
        Ok(dir) => dir + "/gfx/wr/webrender/res",
        Err(_) => std::env::var("CARGO_MANIFEST_DIR").unwrap() + "/../webrender/res",
    };

    let shader_flags =
        ShaderFeatureFlags::GL |
        ShaderFeatureFlags::DUAL_SOURCE_BLENDING |
        ShaderFeatureFlags::ADVANCED_BLEND_EQUATION |
        ShaderFeatureFlags::DEBUG;
    let mut shaders: Vec<String> = Vec::new();
    for (name, features) in get_shader_features(shader_flags) {
        shaders.extend(features.iter().map(|f| {
            if f.is_empty() { name.to_owned() } else { format!("{} {}", name, f) }
        }));
    }

    shaders.sort();

    for shader in &shaders {
        translate_shader(shader, &shader_dir);
    }

    write_load_shader(&shaders);

    println!("cargo:rerun-if-changed=src/blend.h");
    println!("cargo:rerun-if-changed=src/composite.h");
    println!("cargo:rerun-if-changed=src/gl_defs.h");
    println!("cargo:rerun-if-changed=src/glsl.h");
    println!("cargo:rerun-if-changed=src/program.h");
    println!("cargo:rerun-if-changed=src/rasterize.h");
    println!("cargo:rerun-if-changed=src/swgl_ext.h");
    println!("cargo:rerun-if-changed=src/texture.h");
    println!("cargo:rerun-if-changed=src/vector_type.h");
    println!("cargo:rerun-if-changed=src/gl.cc");
    let mut build = cc::Build::new();
    build.cpp(true);

    // SWGL relies heavily on inlining for performance so override -Oz with -O2
    if build.get_compiler().args().contains(&std::ffi::OsString::from("-Oz")) {
        build.flag("-O2");
    }

    build.file("src/gl.cc")
        .flag("-std=c++17")
        .flag("-UMOZILLA_CONFIG_H")
        .flag("-fno-exceptions")
        .flag("-fno-rtti")
        .flag("-fno-math-errno")
        .define("_GLIBCXX_USE_CXX11_ABI", Some("0"))
        .include(shader_dir)
        .include("src")
        .include(std::env::var("OUT_DIR").unwrap())
        .compile("gl_cc");
}
<commit_msg>Bug 1700339 - Use /clang:-undef to undefine predefined macros when clang-cl is used. r=lsalzman<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

extern crate cc;
extern crate glsl_to_cxx;
extern crate webrender_build;

use std::collections::HashSet;
use std::fmt::Write;
use webrender_build::shader::{ShaderFeatureFlags, get_shader_features};

// Shader key is in "name feature,feature" format.
// File name needs to be formatted as "name_feature_feature".
fn shader_file(shader_key: &str) -> String {
    shader_key.replace(' ', "_").replace(',', "_")
}

fn write_load_shader(shader_keys: &[String]) {
    let mut load_shader = String::new();
    for s in shader_keys {
        let _ = write!(load_shader, "#include \"{}.h\"\n", shader_file(s));
    }
    load_shader.push_str("ProgramLoader load_shader(const char* name) {\n");
    for s in shader_keys {
        let _ = write!(load_shader, "  if (!strcmp(name, \"{}\")) {{ return {}_program::loader; }}\n",
                       s, shader_file(s));
    }
    load_shader.push_str("  return nullptr;\n}\n");
    std::fs::write(std::env::var("OUT_DIR").unwrap() + "/load_shader.h", load_shader).unwrap();
}

fn process_imports(shader_dir: &str, shader: &str, included: &mut HashSet<String>, output: &mut String) {
    if !included.insert(shader.into()) {
        return;
    }
    println!("cargo:rerun-if-changed={}/{}.glsl", shader_dir, shader);
    let source = std::fs::read_to_string(format!("{}/{}.glsl", shader_dir, shader)).unwrap();
    for line in source.lines() {
        if line.starts_with("#include ") {
            let imports = line["#include ".len() ..].split(',');
            for import in imports {
                process_imports(shader_dir, import, included, output);
            }
        } else if line.starts_with("#version ") || line.starts_with("#extension ") {
            // ignore
        } else {
            output.push_str(line);
            output.push('\n');
        }
    }
}

fn translate_shader(shader_key: &str, shader_dir: &str) {
    let mut imported = String::from("#define SWGL 1\n#define __VERSION__ 150\n");
    let _ = write!(imported, "#define WR_MAX_VERTEX_TEXTURE_WIDTH {}U\n",
                   webrender_build::MAX_VERTEX_TEXTURE_WIDTH);

    let (basename, features) =
        shader_key.split_at(shader_key.find(' ').unwrap_or(shader_key.len()));
    if !features.is_empty() {
        for feature in features.trim().split(',') {
            let _ = write!(imported, "#define WR_FEATURE_{}\n", feature);
        }
    }

    process_imports(shader_dir, basename, &mut HashSet::new(), &mut imported);

    let shader = shader_file(shader_key);

    let out_dir = std::env::var("OUT_DIR").unwrap();
    let imp_name = format!("{}/{}.c", out_dir, shader);
    std::fs::write(&imp_name, imported).unwrap();

    let mut build = cc::Build::new();
    build.no_default_flags(true);
    if build.get_compiler().is_like_msvc() {
        build.flag("/EP").flag("/clang:-undef");
    } else {
        build.flag("-xc").flag("-P").flag("-undef");
    }
    build.file(&imp_name);
    let vs = build.clone()
        .define("WR_VERTEX_SHADER", Some("1"))
        .expand();
    let fs = build.clone()
        .define("WR_FRAGMENT_SHADER", Some("1"))
        .expand();
    let vs_name = format!("{}/{}.vert", out_dir, shader);
    let fs_name = format!("{}/{}.frag", out_dir, shader);
    std::fs::write(&vs_name, vs).unwrap();
    std::fs::write(&fs_name, fs).unwrap();

    let args = vec![
        "glsl_to_cxx".to_string(),
        vs_name,
        fs_name,
    ];
    let result = glsl_to_cxx::translate(&mut args.into_iter());
    std::fs::write(format!("{}/{}.h", out_dir, shader), result).unwrap();
}

fn main() {
    let shader_dir = match std::env::var("MOZ_SRC") {
        Ok(dir) => dir + "/gfx/wr/webrender/res",
        Err(_) => std::env::var("CARGO_MANIFEST_DIR").unwrap() + "/../webrender/res",
    };

    let shader_flags =
        ShaderFeatureFlags::GL |
        ShaderFeatureFlags::DUAL_SOURCE_BLENDING |
        ShaderFeatureFlags::ADVANCED_BLEND_EQUATION |
        ShaderFeatureFlags::DEBUG;
    let mut shaders: Vec<String> = Vec::new();
    for (name, features) in get_shader_features(shader_flags) {
        shaders.extend(features.iter().map(|f| {
            if f.is_empty() { name.to_owned() } else { format!("{} {}", name, f) }
        }));
    }

    shaders.sort();

    for shader in &shaders {
        translate_shader(shader, &shader_dir);
    }

    write_load_shader(&shaders);

    println!("cargo:rerun-if-changed=src/blend.h");
    println!("cargo:rerun-if-changed=src/composite.h");
    println!("cargo:rerun-if-changed=src/gl_defs.h");
    println!("cargo:rerun-if-changed=src/glsl.h");
    println!("cargo:rerun-if-changed=src/program.h");
    println!("cargo:rerun-if-changed=src/rasterize.h");
    println!("cargo:rerun-if-changed=src/swgl_ext.h");
    println!("cargo:rerun-if-changed=src/texture.h");
    println!("cargo:rerun-if-changed=src/vector_type.h");
    println!("cargo:rerun-if-changed=src/gl.cc");
    let mut build = cc::Build::new();
    build.cpp(true);

    // SWGL relies heavily on inlining for performance so override -Oz with -O2
    if build.get_compiler().args().contains(&std::ffi::OsString::from("-Oz")) {
        build.flag("-O2");
    }

    build.file("src/gl.cc")
        .flag("-std=c++17")
        .flag("-UMOZILLA_CONFIG_H")
        .flag("-fno-exceptions")
        .flag("-fno-rtti")
        .flag("-fno-math-errno")
        .define("_GLIBCXX_USE_CXX11_ABI", Some("0"))
        .include(shader_dir)
        .include("src")
        .include(std::env::var("OUT_DIR").unwrap())
        .compile("gl_cc");
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::collections::btree_map::Entry;
use std::collections::BTreeMap;
use std::fs::File;

use libvda::encode::{EncodeCapabilities, VeaImplType, VeaInstance};

use base::{error, warn, IntoRawDescriptor};

use super::*;
use crate::virtio::video::format::{Format, FormatDesc, FormatRange, FrameFormat, Level, Profile};
use crate::virtio::video::{
    encoder::{encoder::*, EncoderDevice},
    error::{VideoError, VideoResult},
    Tube,
};

pub struct LibvdaEncoder {
    instance: VeaInstance,
    capabilities: EncoderCapabilities,
}

impl LibvdaEncoder {
    fn new() -> VideoResult<Self> {
        let instance = VeaInstance::new(VeaImplType::Gavea)?;

        let EncodeCapabilities {
            input_formats,
            output_formats,
        } = instance.get_capabilities();

        if input_formats.len() == 0 || output_formats.len() == 0 {
            error!("No input or output formats.");
            return Err(VideoError::InvalidFormat);
        }

        let input_format_descs: Vec<FormatDesc> = input_formats
            .iter()
            .map(|input_format| {
                let format = match input_format {
                    libvda::PixelFormat::NV12 => Format::NV12,
                    libvda::PixelFormat::YV12 => Format::YUV420,
                };

                // VEA's GetSupportedProfiles does not return resolution information.
                // The input formats are retrieved by querying minigbm.
                // TODO(alexlau): Populate this with real information.

                FormatDesc {
                    mask: !(u64::MAX << output_formats.len()),
                    format,
                    frame_formats: vec![FrameFormat {
                        width: FormatRange {
                            min: 2,
                            max: 4096,
                            step: 1,
                        },
                        height: FormatRange {
                            min: 2,
                            max: 4096,
                            step: 1,
                        },
                        bitrates: vec![FormatRange {
                            min: 0,
                            max: 8000,
                            step: 1,
                        }],
                    }],
                }
            })
            .collect();

        if input_format_descs
            .iter()
            .find(|fd| fd.format == Format::NV12)
            .is_none()
        {
            // NV12 is currently the only supported pixel format for libvda.
            error!("libvda encoder does not support NV12.");
            return Err(VideoError::InvalidFormat);
        }

        struct ParsedFormat {
            profiles: Vec<Profile>,
            max_width: u32,
            max_height: u32,
        }
        let mut parsed_formats: BTreeMap<Format, ParsedFormat> = BTreeMap::new();

        for output_format in output_formats.iter() {
            // TODO(alexlau): Consider using `max_framerate_numerator` and `max_framerate_denominator`.
            let libvda::encode::OutputProfile {
                profile: libvda_profile,
                max_width,
                max_height,
                ..
            } = output_format;

            let profile = match Profile::from_libvda_profile(*libvda_profile) {
                Some(p) => p,
                None => {
                    warn!("Skipping unsupported libvda profile: {:?}", libvda_profile);
                    continue;
                }
            };

            match parsed_formats.entry(profile.to_format()) {
                Entry::Occupied(mut occupied_entry) => {
                    let parsed_format = occupied_entry.get_mut();
                    parsed_format.profiles.push(profile);
                    // If we get different libvda profiles of the same VIRTIO_VIDEO_FORMAT
                    // (Format) that have different max resolutions or bitrates, take the
                    // minimum between all of the different profiles.
                    parsed_format.max_width = std::cmp::min(*max_width, parsed_format.max_width);
                    parsed_format.max_height = std::cmp::min(*max_height, parsed_format.max_height);
                }
                Entry::Vacant(vacant_entry) => {
                    vacant_entry.insert(ParsedFormat {
                        profiles: vec![profile],
                        max_width: *max_width,
                        max_height: *max_height,
                    });
                }
            }
        }

        let mut output_format_descs = vec![];
        let mut coded_format_profiles = BTreeMap::new();
        for (format, parsed_format) in parsed_formats.into_iter() {
            let ParsedFormat {
                mut profiles,
                max_width,
                max_height,
            } = parsed_format;

            output_format_descs.push(FormatDesc {
                mask: !(u64::MAX << output_formats.len()),
                format,
                frame_formats: vec![FrameFormat {
                    width: FormatRange {
                        min: 2,
                        max: max_width,
                        step: 1,
                    },
                    height: FormatRange {
                        min: 2,
                        max: max_height,
                        step: 1,
                    },
                    bitrates: vec![FormatRange {
                        min: 0,
                        max: 8000,
                        step: 1,
                    }],
                }],
            });

            profiles.sort_unstable();
            coded_format_profiles.insert(format, profiles);
        }

        Ok(LibvdaEncoder {
            instance,
            capabilities: EncoderCapabilities {
                input_format_descs,
                output_format_descs,
                coded_format_profiles,
            },
        })
    }
}

impl Encoder for LibvdaEncoder {
    type Session = LibvdaEncoderSession;

    fn query_capabilities(&self) -> VideoResult<EncoderCapabilities> {
        Ok(self.capabilities.clone())
    }

    fn start_session(&mut self, config: SessionConfig) -> VideoResult<LibvdaEncoderSession> {
        if config.dst_params.format.is_none() {
            return Err(VideoError::InvalidArgument);
        }

        let input_format = match config
            .src_params
            .format
            .ok_or(VideoError::InvalidArgument)?
        {
            Format::NV12 => libvda::PixelFormat::NV12,
            Format::YUV420 => libvda::PixelFormat::YV12,
            unsupported_format => {
                error!("Unsupported libvda format: {}", unsupported_format);
                return Err(VideoError::InvalidArgument);
            }
        };

        let output_profile = match config.dst_profile.to_libvda_profile() {
            Some(p) => p,
            None => {
                error!("Unsupported libvda profile");
                return Err(VideoError::InvalidArgument);
            }
        };

        let config = libvda::encode::Config {
            input_format,
            input_visible_width: config.src_params.frame_width,
            input_visible_height: config.src_params.frame_height,
            output_profile,
            initial_bitrate: config.dst_bitrate,
            initial_framerate: if config.frame_rate == 0 {
                None
            } else {
                Some(config.frame_rate)
            },
            h264_output_level: config.dst_h264_level.map(|level| {
                // This value is aligned to the H264 standard definition of SPS.level_idc.
                match level {
                    Level::H264_1_0 => 10,
                    Level::H264_1_1 => 11,
                    Level::H264_1_2 => 12,
                    Level::H264_1_3 => 13,
                    Level::H264_2_0 => 20,
                    Level::H264_2_1 => 21,
                    Level::H264_2_2 => 22,
                    Level::H264_3_0 => 30,
                    Level::H264_3_1 => 31,
                    Level::H264_3_2 => 32,
                    Level::H264_4_0 => 40,
                    Level::H264_4_1 => 41,
                    Level::H264_4_2 => 42,
                    Level::H264_5_0 => 50,
                    Level::H264_5_1 => 51,
                }
            }),
        };

        let session = self.instance.open_session(config)?;

        Ok(LibvdaEncoderSession {
            session,
            next_input_buffer_id: 1,
            next_output_buffer_id: 1,
        })
    }

    fn stop_session(&mut self, _session: LibvdaEncoderSession) -> VideoResult<()> {
        // Resources will be freed when `_session` is dropped.
        Ok(())
    }
}

pub struct LibvdaEncoderSession {
    session: libvda::encode::Session,
    next_input_buffer_id: InputBufferId,
    next_output_buffer_id: OutputBufferId,
}

impl EncoderSession for LibvdaEncoderSession {
    fn encode(
        &mut self,
        resource: File,
        planes: &[VideoFramePlane],
        timestamp: u64,
        force_keyframe: bool,
    ) -> VideoResult<InputBufferId> {
        let input_buffer_id = self.next_input_buffer_id;

        let libvda_planes = planes
            .iter()
            .map(|plane| libvda::FramePlane {
                offset: plane.offset as i32,
                stride: plane.stride as i32,
            })
            .collect::<Vec<_>>();

        self.session.encode(
            input_buffer_id as i32,
            resource.into_raw_descriptor(),
            &libvda_planes,
            timestamp as i64,
            force_keyframe,
        )?;

        self.next_input_buffer_id = self.next_input_buffer_id.wrapping_add(1);

        Ok(input_buffer_id)
    }

    fn use_output_buffer(
        &mut self,
        file: File,
        offset: u32,
        size: u32,
    ) -> VideoResult<OutputBufferId> {
        let output_buffer_id = self.next_output_buffer_id;
        self.next_output_buffer_id = self.next_output_buffer_id.wrapping_add(1);

        self.session.use_output_buffer(
            output_buffer_id as i32,
            file.into_raw_descriptor(),
            offset,
            size,
        )?;

        Ok(output_buffer_id)
    }

    fn flush(&mut self) -> VideoResult<()> {
        self.session.flush().map_err(VideoError::from)
    }

    fn request_encoding_params_change(&mut self, bitrate: u32, framerate: u32) -> VideoResult<()> {
        self.session
            .request_encoding_params_change(bitrate, framerate)
            .map_err(VideoError::from)
    }

    fn event_pipe(&self) -> &File {
        self.session.pipe()
    }

    fn read_event(&mut self) -> VideoResult<EncoderEvent> {
        let event = self.session.read_event()?;

        use libvda::encode::Event::*;
        let encoder_event = match event {
            RequireInputBuffers {
                input_count,
                input_frame_width,
                input_frame_height,
                output_buffer_size,
            } => EncoderEvent::RequireInputBuffers {
                input_count,
                input_frame_width,
                input_frame_height,
                output_buffer_size,
            },
            ProcessedInputBuffer(id) => EncoderEvent::ProcessedInputBuffer { id: id as u32 },
            ProcessedOutputBuffer {
                output_buffer_id,
                payload_size,
                key_frame,
                timestamp,
                ..
            } => EncoderEvent::ProcessedOutputBuffer {
                id: output_buffer_id as u32,
                bytesused: payload_size,
                keyframe: key_frame,
                timestamp: timestamp as u64,
            },
            FlushResponse { flush_done } => EncoderEvent::FlushResponse { flush_done },
            NotifyError(err) => EncoderEvent::NotifyError {
                error: VideoError::BackendFailure(Box::new(err)),
            },
        };

        Ok(encoder_event)
    }
}

/// Create a new encoder instance using a Libvda encoder instance to perform
/// the encoding.
impl EncoderDevice<LibvdaEncoder> {
    pub fn new(resource_bridge: Tube) -> VideoResult<Self> {
        let vea = LibvdaEncoder::new()?;
        EncoderDevice::from_backend(vea, resource_bridge)
    }
}
<commit_msg>virtio: video: encoder: vda: increase buffer ID on success<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::collections::btree_map::Entry;
use std::collections::BTreeMap;
use std::fs::File;

use libvda::encode::{EncodeCapabilities, VeaImplType, VeaInstance};

use base::{error, warn, IntoRawDescriptor};

use super::*;
use crate::virtio::video::format::{Format, FormatDesc, FormatRange, FrameFormat, Level, Profile};
use crate::virtio::video::{
    encoder::{encoder::*, EncoderDevice},
    error::{VideoError, VideoResult},
    Tube,
};

pub struct LibvdaEncoder {
    instance: VeaInstance,
    capabilities: EncoderCapabilities,
}

impl LibvdaEncoder {
    fn new() -> VideoResult<Self> {
        let instance = VeaInstance::new(VeaImplType::Gavea)?;

        let EncodeCapabilities {
            input_formats,
            output_formats,
        } = instance.get_capabilities();

        if input_formats.len() == 0 || output_formats.len() == 0 {
            error!("No input or output formats.");
            return Err(VideoError::InvalidFormat);
        }

        let input_format_descs: Vec<FormatDesc> = input_formats
            .iter()
            .map(|input_format| {
                let format = match input_format {
                    libvda::PixelFormat::NV12 => Format::NV12,
                    libvda::PixelFormat::YV12 => Format::YUV420,
                };

                // VEA's GetSupportedProfiles does not return resolution information.
                // The input formats are retrieved by querying minigbm.
                // TODO(alexlau): Populate this with real information.

                FormatDesc {
                    mask: !(u64::MAX << output_formats.len()),
                    format,
                    frame_formats: vec![FrameFormat {
                        width: FormatRange {
                            min: 2,
                            max: 4096,
                            step: 1,
                        },
                        height: FormatRange {
                            min: 2,
                            max: 4096,
                            step: 1,
                        },
                        bitrates: vec![FormatRange {
                            min: 0,
                            max: 8000,
                            step: 1,
                        }],
                    }],
                }
            })
            .collect();

        if input_format_descs
            .iter()
            .find(|fd| fd.format == Format::NV12)
            .is_none()
        {
            // NV12 is currently the only supported pixel format for libvda.
            error!("libvda encoder does not support NV12.");
            return Err(VideoError::InvalidFormat);
        }

        struct ParsedFormat {
            profiles: Vec<Profile>,
            max_width: u32,
            max_height: u32,
        }
        let mut parsed_formats: BTreeMap<Format, ParsedFormat> = BTreeMap::new();

        for output_format in output_formats.iter() {
            // TODO(alexlau): Consider using `max_framerate_numerator` and `max_framerate_denominator`.
            let libvda::encode::OutputProfile {
                profile: libvda_profile,
                max_width,
                max_height,
                ..
            } = output_format;

            let profile = match Profile::from_libvda_profile(*libvda_profile) {
                Some(p) => p,
                None => {
                    warn!("Skipping unsupported libvda profile: {:?}", libvda_profile);
                    continue;
                }
            };

            match parsed_formats.entry(profile.to_format()) {
                Entry::Occupied(mut occupied_entry) => {
                    let parsed_format = occupied_entry.get_mut();
                    parsed_format.profiles.push(profile);
                    // If we get different libvda profiles of the same VIRTIO_VIDEO_FORMAT
                    // (Format) that have different max resolutions or bitrates, take the
                    // minimum between all of the different profiles.
                    parsed_format.max_width = std::cmp::min(*max_width, parsed_format.max_width);
                    parsed_format.max_height = std::cmp::min(*max_height, parsed_format.max_height);
                }
                Entry::Vacant(vacant_entry) => {
                    vacant_entry.insert(ParsedFormat {
                        profiles: vec![profile],
                        max_width: *max_width,
                        max_height: *max_height,
                    });
                }
            }
        }

        let mut output_format_descs = vec![];
        let mut coded_format_profiles = BTreeMap::new();
        for (format, parsed_format) in parsed_formats.into_iter() {
            let ParsedFormat {
                mut profiles,
                max_width,
                max_height,
            } = parsed_format;

            output_format_descs.push(FormatDesc {
                mask: !(u64::MAX << output_formats.len()),
                format,
                frame_formats: vec![FrameFormat {
                    width: FormatRange {
                        min: 2,
                        max: max_width,
                        step: 1,
                    },
                    height: FormatRange {
                        min: 2,
                        max: max_height,
                        step: 1,
                    },
                    bitrates: vec![FormatRange {
                        min: 0,
                        max: 8000,
                        step: 1,
                    }],
                }],
            });

            profiles.sort_unstable();
            coded_format_profiles.insert(format, profiles);
        }

        Ok(LibvdaEncoder {
            instance,
            capabilities: EncoderCapabilities {
                input_format_descs,
                output_format_descs,
                coded_format_profiles,
            },
        })
    }
}

impl Encoder for LibvdaEncoder {
    type Session = LibvdaEncoderSession;

    fn query_capabilities(&self) -> VideoResult<EncoderCapabilities> {
        Ok(self.capabilities.clone())
    }

    fn start_session(&mut self, config: SessionConfig) -> VideoResult<LibvdaEncoderSession> {
        if config.dst_params.format.is_none() {
            return Err(VideoError::InvalidArgument);
        }

        let input_format = match config
            .src_params
            .format
            .ok_or(VideoError::InvalidArgument)?
        {
            Format::NV12 => libvda::PixelFormat::NV12,
            Format::YUV420 => libvda::PixelFormat::YV12,
            unsupported_format => {
                error!("Unsupported libvda format: {}", unsupported_format);
                return Err(VideoError::InvalidArgument);
            }
        };

        let output_profile = match config.dst_profile.to_libvda_profile() {
            Some(p) => p,
            None => {
                error!("Unsupported libvda profile");
                return Err(VideoError::InvalidArgument);
            }
        };

        let config = libvda::encode::Config {
            input_format,
            input_visible_width: config.src_params.frame_width,
            input_visible_height: config.src_params.frame_height,
            output_profile,
            initial_bitrate: config.dst_bitrate,
            initial_framerate: if config.frame_rate == 0 {
                None
            } else {
                Some(config.frame_rate)
            },
            h264_output_level: config.dst_h264_level.map(|level| {
                // This value is aligned to the H264 standard definition of SPS.level_idc.
                match level {
                    Level::H264_1_0 => 10,
                    Level::H264_1_1 => 11,
                    Level::H264_1_2 => 12,
                    Level::H264_1_3 => 13,
                    Level::H264_2_0 => 20,
                    Level::H264_2_1 => 21,
                    Level::H264_2_2 => 22,
                    Level::H264_3_0 => 30,
                    Level::H264_3_1 => 31,
                    Level::H264_3_2 => 32,
                    Level::H264_4_0 => 40,
                    Level::H264_4_1 => 41,
                    Level::H264_4_2 => 42,
                    Level::H264_5_0 => 50,
                    Level::H264_5_1 => 51,
                }
            }),
        };

        let session = self.instance.open_session(config)?;

        Ok(LibvdaEncoderSession {
            session,
            next_input_buffer_id: 1,
            next_output_buffer_id: 1,
        })
    }

    fn stop_session(&mut self, _session: LibvdaEncoderSession) -> VideoResult<()> {
        // Resources will be freed when `_session` is dropped.
        Ok(())
    }
}

pub struct LibvdaEncoderSession {
    session: libvda::encode::Session,
    next_input_buffer_id: InputBufferId,
    next_output_buffer_id: OutputBufferId,
}

impl EncoderSession for LibvdaEncoderSession {
    fn encode(
        &mut self,
        resource: File,
        planes: &[VideoFramePlane],
        timestamp: u64,
        force_keyframe: bool,
    ) -> VideoResult<InputBufferId> {
        let input_buffer_id = self.next_input_buffer_id;

        let libvda_planes = planes
            .iter()
            .map(|plane| libvda::FramePlane {
                offset: plane.offset as i32,
                stride: plane.stride as i32,
            })
            .collect::<Vec<_>>();

        self.session.encode(
            input_buffer_id as i32,
            resource.into_raw_descriptor(),
            &libvda_planes,
            timestamp as i64,
            force_keyframe,
        )?;

        self.next_input_buffer_id = self.next_input_buffer_id.wrapping_add(1);

        Ok(input_buffer_id)
    }

    fn use_output_buffer(
        &mut self,
        file: File,
        offset: u32,
        size: u32,
    ) -> VideoResult<OutputBufferId> {
        let output_buffer_id = self.next_output_buffer_id;

        self.session.use_output_buffer(
            output_buffer_id as i32,
            file.into_raw_descriptor(),
            offset,
            size,
        )?;

        self.next_output_buffer_id = self.next_output_buffer_id.wrapping_add(1);

        Ok(output_buffer_id)
    }

    fn flush(&mut self) -> VideoResult<()> {
        self.session.flush().map_err(VideoError::from)
    }

    fn request_encoding_params_change(&mut self, bitrate: u32, framerate: u32) -> VideoResult<()> {
        self.session
            .request_encoding_params_change(bitrate, framerate)
            .map_err(VideoError::from)
    }

    fn event_pipe(&self) -> &File {
        self.session.pipe()
    }

    fn read_event(&mut self) -> VideoResult<EncoderEvent> {
        let event = self.session.read_event()?;

        use libvda::encode::Event::*;
        let encoder_event = match event {
            RequireInputBuffers {
                input_count,
                input_frame_width,
                input_frame_height,
                output_buffer_size,
            } => EncoderEvent::RequireInputBuffers {
                input_count,
                input_frame_width,
                input_frame_height,
                output_buffer_size,
            },
            ProcessedInputBuffer(id) => EncoderEvent::ProcessedInputBuffer { id: id as u32 },
            ProcessedOutputBuffer {
                output_buffer_id,
                payload_size,
                key_frame,
                timestamp,
                ..
            } => EncoderEvent::ProcessedOutputBuffer {
                id: output_buffer_id as u32,
                bytesused: payload_size,
                keyframe: key_frame,
                timestamp: timestamp as u64,
            },
            FlushResponse { flush_done } => EncoderEvent::FlushResponse { flush_done },
            NotifyError(err) => EncoderEvent::NotifyError {
                error: VideoError::BackendFailure(Box::new(err)),
            },
        };

        Ok(encoder_event)
    }
}

/// Create a new encoder instance using a Libvda encoder instance to perform
/// the encoding.
impl EncoderDevice<LibvdaEncoder> {
    pub fn new(resource_bridge: Tube) -> VideoResult<Self> {
        let vea = LibvdaEncoder::new()?;
        EncoderDevice::from_backend(vea, resource_bridge)
    }
}
<|endoftext|>


<commit_before>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines. Any line that starts with a # is ignored.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
        }
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)) {
        Err(error) => {
            println!("Error creating a UDP socket at {}:{}. {}.",ip,BROADCAST_PORT_UDP,error);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(error) => {
            println!("Error listening to TCP at {}:{}. {}.",ip,SERVER_LIST_PORT_TCP,error);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = time::now().to_timespec().sec;
            servers.retain(|x| x.last_alive + DROP_TIME > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                blacklist_ref.clear();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        match reader.lines().collect() {
                            Err(_) => {},
                            Ok(t) => {
                                let lines : Vec<String> = t;
                                for line in lines.iter().filter(|x| !x.starts_with("#")) {
                                    blacklist_ref.push(line.clone());
                                }
                            }
                        }
                    },
                    Err(_) => {}
                };
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => ip.ip_string()
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\n');
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

            let mut servers = servers_mut_udp.lock().unwrap();

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = time::now().to_timespec().sec;
                    match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                        None => {
                            if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                                let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                                (*servers).push(serverness);
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if i.ip == client_ip && i.port == source.port() {
                    i.last_alive = time::now().to_timespec().sec;
                    break;
                }
            }
        }
    }
}
<commit_msg>Improve blacklist.txt parsing * Make code more compact/immutable * Improve efficiency due to collect()'ing at the very end * Ignore empty lines<commit_after>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines. Any line that starts with a # is ignored.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
        }
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)) {
        Err(error) => {
            println!("Error creating a UDP socket at {}:{}. {}.",ip,BROADCAST_PORT_UDP,error);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(error) => {
            println!("Error listening to TCP at {}:{}. {}.",ip,SERVER_LIST_PORT_TCP,error);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = time::now().to_timespec().sec;
            servers.retain(|x| x.last_alive + DROP_TIME > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        let valid_line = |x: &str| -> bool { x.trim().len() > 0 && !x.starts_with("#") };
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                *blacklist_ref =
                    File::open(BLACKLIST_FILE).
                    map(|file|
                        BufReader::new(&file).lines().
                        filter_map(|line| line.ok().and_then(|x| if valid_line(&x) { Some(x) } else { None })).
                        collect()
                    ).
                    unwrap_or_else(|_| Vec::new());

                println!("{:?}", *blacklist_ref);
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => ip.ip_string()
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\n');
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

            let mut servers = servers_mut_udp.lock().unwrap();

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = time::now().to_timespec().sec;
                    match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                        None => {
                            if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                                let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                                (*servers).push(serverness);
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if i.ip == client_ip && i.port == source.port() {
                    i.last_alive = time::now().to_timespec().sec;
                    break;
                }
            }
        }
    }
}
<|endoftext|>