GoldSrc demo repair tool.

main.rs

extern crate byteorder;

use byteorder::{ByteOrder, LittleEndian, ReadBytesExt, WriteBytesExt};
use std::env;
use std::ffi::OsString;
use std::fs::File;
use std::io;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::PathBuf;

mod types;
use types::*;

macro_rules! try_read {
    ( $x:expr ) => (
        match $x {
            Ok(x) => x,
            Err(err) => {
                if err.kind() == io::ErrorKind::UnexpectedEof {
                    break;
                } else {
                    return Err(err);
                }
            }
        }
    )
}

fn main() {
    if let Some(filename) = env::args().nth(1) {
        if let Err(err) = repair_demo(&filename, env::args().nth(2)) {
            println!("Error: {}", err);
        }
    } else {
        println!("Usage: ./demo-repair <path to demo> [path to output demo]");
    }
}

fn repair_demo(filename: &str, output: Option<String>) -> Result<(), io::Error> {
    let mut demo = try!(File::open(filename));

    let header_result = demo.read_header();

    if let Err(err) = header_result {
        if err.kind() == io::ErrorKind::UnexpectedEof {
            println!("This is not a GoldSource demo file.");
            return Ok(());
        }

        return Err(err);
    }

    let header = header_result.unwrap();

    if header.magic.chunks(6).next().unwrap() != b"HLDEMO" {
        println!("This is not a GoldSource demo file.");
        return Ok(());
    }

    if header.demo_protocol != 5 {
        println!("Unsupported demo protocol: {}. Only demo protocol 5 is supported.",
                 header.demo_protocol);
        return Ok(());
    }

    let mut new_demo = try!(if let Some(name) = output {
        File::create(name)
    } else {
        File::create(get_new_demo_path(filename))
    });
    try!(new_demo.write_header(header));

    let mut directory_entries = vec![DemoDirectoryEntry {
                                         entry_type: DemoDirectoryEntryType::Start,
                                         playback_time: 0f32,
                                         frame_count: 0,
                                         offset: 544,
                                     }];
    let mut current_entry = 0usize;
    let mut wrote_last_next_section = false;

    loop {
        let frame_type = try_read!(demo.read_u8());
        let frame_time = try_read!(demo.read_f32::<LittleEndian>());

        let remaining_bytes = match DemoFrameType::from(frame_type) {
            DemoFrameType::DemoStart => {
                let mut buf = Vec::with_capacity(4);
                buf.resize(4, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::ConsoleCommand => {
                let mut buf = Vec::with_capacity(68);
                buf.resize(68, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::ClientData => {
                let mut buf = Vec::with_capacity(36);
                buf.resize(36, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::NextSection => {
                let mut buf = Vec::with_capacity(4);
                buf.resize(4, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::Event => {
                let mut buf = Vec::with_capacity(88);
                buf.resize(88, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::WeaponAnim => {
                let mut buf = Vec::with_capacity(12);
                buf.resize(12, 0);

                try_read!(demo.read_exact(&mut buf));

                buf
            }

            DemoFrameType::Sound => {
                let mut buf = Vec::with_capacity(28);
                buf.resize(12, 0);

                try_read!(demo.read_exact(&mut buf));

                let length = LittleEndian::read_u32(buf.split_at(8).1);

                if length > 255 {
                    break;
                }

                buf.resize(28 + length as usize, 0);
                try_read!(demo.read_exact(buf.split_at_mut(12).1));

                buf
            }

            DemoFrameType::DemoBuffer => {
                let mut buf = Vec::with_capacity(8);
                buf.resize(8, 0);

                try_read!(demo.read_exact(&mut buf));

                let length = LittleEndian::read_u32(buf.split_at_mut(4).1);

                if length > 32768 {
                    break;
                }

                buf.resize(8 + length as usize, 0);
                try_read!(demo.read_exact(buf.split_at_mut(8).1));

                buf
            }

            DemoFrameType::NetMsg => {
                let mut buf = Vec::with_capacity(472);
                buf.resize(472, 0);

                try_read!(demo.read_exact(&mut buf));

                let length = LittleEndian::read_u32(buf.split_at(468).1);

                if length > 65536 {
                    break;
                }

                buf.resize(472 + length as usize, 0);
                try_read!(demo.read_exact(buf.split_at_mut(472).1));

                buf
            }
        };

        if wrote_last_next_section {
            directory_entries.push(DemoDirectoryEntry {
                entry_type: DemoDirectoryEntryType::Normal,
                playback_time: 0f32,
                frame_count: 0,
                offset: try!(new_demo.seek(SeekFrom::Current(0))) as i32,
            });

            current_entry += 1;
        }

        directory_entries[current_entry].playback_time =
            directory_entries[current_entry].playback_time.max(frame_time);
        directory_entries[current_entry].frame_count += 1;

        try!(new_demo.write_u8(frame_type));
        try!(new_demo.write_f32::<LittleEndian>(frame_time));
        try!(new_demo.write_all(&remaining_bytes));

        if DemoFrameType::from(frame_type) == DemoFrameType::NextSection {
            wrote_last_next_section = true;

            if current_entry == 1 {
                // Demos output by the engine contain 2 demo entries.
                break;
            }
        } else {
            wrote_last_next_section = false;
        }
    }

    if !wrote_last_next_section {
        try!(new_demo.write_u8(DemoFrameType::NextSection as u8));
        try!(new_demo.write_f32::<LittleEndian>(0f32));
        try!(new_demo.write_i32::<LittleEndian>(0));

        directory_entries[current_entry].frame_count += 1;
    }

    let directory_offset = try!(new_demo.seek(SeekFrom::Current(0))) as i32;

    try!(new_demo.write_directory(&directory_entries));

    try!(new_demo.seek(SeekFrom::Start(540)));
    try!(new_demo.write_i32::<LittleEndian>(directory_offset));

    println!("Done.");

    Ok(())
}

fn get_new_demo_path(filename: &str) -> PathBuf {
    let mut new_demo_path = PathBuf::from(filename);

    let extension = new_demo_path.extension().map(|x| x.to_os_string()).unwrap_or(OsString::new());
    new_demo_path.set_extension("");

    let mut new_demo_filename = new_demo_path.file_name().unwrap().to_os_string();
    new_demo_filename.push("_repaired");

    new_demo_path.set_file_name(new_demo_filename);
    new_demo_path.set_extension(extension);

    new_demo_path
}

types.rs

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::io;
use std::mem;

pub struct DemoHeader {
    pub magic: [u8; 8],
    pub demo_protocol: i32,
    remaining_bytes: [u8; 532],
}

pub enum DemoDirectoryEntryType {
    Start = 0,
    Normal = 1,
}

impl<'a> From<&'a DemoDirectoryEntryType> for i32 {
    fn from(n: &'a DemoDirectoryEntryType) -> Self {
        match *n {
            DemoDirectoryEntryType::Start => 0,
            DemoDirectoryEntryType::Normal => 1,
        }
    }
}

pub struct DemoDirectoryEntry {
    pub entry_type: DemoDirectoryEntryType,
    pub playback_time: f32,
    pub frame_count: i32,
    pub offset: i32,
}

#[derive(PartialEq)]
pub enum DemoFrameType {
    NetMsg = 1,
    DemoStart = 2,
    ConsoleCommand = 3,
    ClientData = 4,
    NextSection = 5,
    Event = 6,
    WeaponAnim = 7,
    Sound = 8,
    DemoBuffer = 9,
}

impl From<u8> for DemoFrameType {
    fn from(n: u8) -> Self {
        match n {
            2 => DemoFrameType::DemoStart,
            3 => DemoFrameType::ConsoleCommand,
            4 => DemoFrameType::ClientData,
            5 => DemoFrameType::NextSection,
            6 => DemoFrameType::Event,
            7 => DemoFrameType::WeaponAnim,
            8 => DemoFrameType::Sound,
            9 => DemoFrameType::DemoBuffer,
            _ => DemoFrameType::NetMsg,
        }
    }
}

pub trait CustomRead: ReadBytesExt {
    fn read_header(&mut self) -> io::Result<DemoHeader> {
        let mut rv = unsafe { mem::uninitialized::<DemoHeader>() };

        try!(self.read_exact(&mut rv.magic));
        rv.demo_protocol = try!(self.read_i32::<LittleEndian>());
        try!(self.read_exact(&mut rv.remaining_bytes));

        Ok(rv)
    }
}

impl<R: ReadBytesExt + ?Sized> CustomRead for R {}

pub trait CustomWrite: WriteBytesExt {
    fn write_header(&mut self, header: DemoHeader) -> io::Result<()> {
        try!(self.write_all(&header.magic));
        try!(self.write_i32::<LittleEndian>(header.demo_protocol));
        try!(self.write_all(&header.remaining_bytes));

        Ok(())
    }

    fn write_directory(&mut self, directory: &[DemoDirectoryEntry]) -> io::Result<()> {
        try!(self.write_i32::<LittleEndian>(directory.len() as i32));

        for entry in directory {
            try!(self.write_i32::<LittleEndian>(i32::from(&entry.entry_type)));

            let mut description = [0u8; 64];

            let description_str = match entry.entry_type {
                DemoDirectoryEntryType::Start => b"LOADING".as_ref(),
                DemoDirectoryEntryType::Normal => b"Playback".as_ref(),
            };

            description.chunks_mut(description_str.len())
                .next()
                .unwrap()
                .copy_from_slice(description_str);

            try!(self.write_all(&description));
            try!(self.write_i32::<LittleEndian>(0)); // flags
            try!(self.write_i32::<LittleEndian>(-1)); // CD track
            try!(self.write_f32::<LittleEndian>(entry.playback_time));
            try!(self.write_i32::<LittleEndian>(entry.frame_count));
            try!(self.write_i32::<LittleEndian>(entry.offset));
            try!(self.write_i32::<LittleEndian>(0)); // file length
        }

        Ok(())
    }
}

impl<W: WriteBytesExt + ?Sized> CustomWrite for W {}