1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188

use crate::FromBytes;
use crate::ToBytes;
use crate::{
    errors::{ReadError, ReadResult},
    ksmfile::sections::CodeSection,
};
use flate2::write::GzEncoder;
use flate2::{read::GzDecoder, Compression};
use std::{
    io::{Read, Write},
    iter::Peekable,
    slice::Iter,
};

use self::sections::{ArgumentSection, DebugSection};

pub mod sections;

pub mod instructions;

pub use instructions::Instr;

const KSM_MAGIC_NUMBER: u32 = 0x4558036b;

pub struct KSMFile {
    header: KSMHeader,
    arg_section: ArgumentSection,
    code_sections: Vec<CodeSection>,
    debug_section: DebugSection,
}

impl KSMFile {
    pub fn new() -> Self {
        KSMFile {
            header: KSMHeader::new(),
            arg_section: ArgumentSection::new(),
            code_sections: Vec::new(),
            debug_section: DebugSection::new(1),
        }
    }

    pub fn arg_section(&self) -> &ArgumentSection {
        &self.arg_section
    }

    pub fn arg_section_mut(&mut self) -> &mut ArgumentSection {
        &mut self.arg_section
    }

    pub fn code_sections(&self) -> Iter<CodeSection> {
        self.code_sections.iter()
    }

    pub fn add_code_section(&mut self, code_section: CodeSection) {
        self.code_sections.push(code_section);
    }

    pub fn debug_section(&self) -> &DebugSection {
        &self.debug_section
    }

    pub fn debug_section_mut(&mut self) -> &mut DebugSection {
        &mut self.debug_section
    }
}

impl ToBytes for KSMFile {
    fn to_bytes(&self, buf: &mut Vec<u8>) {
        let mut uncompressed_buf = Vec::with_capacity(2048);
        let mut zipped_contents = Vec::with_capacity(2048);

        self.header.to_bytes(&mut uncompressed_buf);

        self.arg_section.to_bytes(&mut uncompressed_buf);

        let mut encoder = GzEncoder::new(&mut zipped_contents, Compression::best());

        for code_section in self.code_sections.iter() {
            code_section.to_bytes(&mut uncompressed_buf, self.arg_section.num_index_bytes());
        }

        self.debug_section.to_bytes(&mut uncompressed_buf);

        encoder
            .write_all(&uncompressed_buf)
            .expect("Error compressing KSM file");

        encoder
            .finish()
            .expect("Error finishing KSM file compression");

        buf.append(&mut zipped_contents);
    }
}

impl FromBytes for KSMFile {
    fn from_bytes(source: &mut Peekable<Iter<u8>>, debug: bool) -> ReadResult<Self>
    where
        Self: Sized,
    {
        let zipped_contents: Vec<u8> = source.map(|b| *b).collect();
        let mut decoder = GzDecoder::new(zipped_contents.as_slice());

        let mut decompressed: Vec<u8> = Vec::with_capacity(2048);

        decoder
            .read_to_end(&mut decompressed)
            .map_err(|e| ReadError::KSMDecompressionError(e))?;

        let mut decompressed_source = decompressed.iter().peekable();

        let header = KSMHeader::from_bytes(&mut decompressed_source, debug)?;

        let arg_section = ArgumentSection::from_bytes(&mut decompressed_source, debug)?;

        let mut code_sections = Vec::new();

        loop {
            let delimiter = u8::from_bytes(&mut decompressed_source, debug)
                .map_err(|_| ReadError::MissingCodeSectionError)?;

            if delimiter != b'%' {
                return Err(ReadError::ExpectedCodeSectionError(delimiter));
            }

            if let Some(next) = decompressed_source.peek() {
                if **next == b'D' {
                    break;
                }

                let code_section = CodeSection::from_bytes(
                    &mut decompressed_source,
                    debug,
                    arg_section.num_index_bytes(),
                )?;

                code_sections.push(code_section);
            } else {
                return Err(ReadError::MissingCodeSectionError);
            }
        }

        let debug_section = DebugSection::from_bytes(&mut decompressed_source, debug)?;

        let ksm = KSMFile {
            header,
            arg_section,
            code_sections,
            debug_section,
        };

        Ok(ksm)
    }
}

pub struct KSMHeader {
    magic: u32,
}

impl KSMHeader {
    pub fn new() -> Self {
        KSMHeader {
            magic: KSM_MAGIC_NUMBER,
        }
    }
}

impl ToBytes for KSMHeader {
    fn to_bytes(&self, buf: &mut Vec<u8>) {
        self.magic.to_bytes(buf);
    }
}

impl FromBytes for KSMHeader {
    fn from_bytes(source: &mut Peekable<Iter<u8>>, debug: bool) -> ReadResult<Self>
    where
        Self: Sized,
    {
        let magic = u32::from_bytes(source, debug)
            .map_err(|_| ReadError::KSMHeaderReadError("file magic"))?;

        if magic != KSM_MAGIC_NUMBER {
            return Err(ReadError::InvalidKSMFileMagicError);
        }

        Ok(KSMHeader::new())
    }
}