ractf-medea/src/main.rs

663 lines
16 KiB
Rust

use byteorder::{ReadBytesExt, LE};
use std::io::Cursor;
use std::convert::{TryInto, TryFrom};
use std::io::Read;
struct Machine {
RX: u16,
RY: u16,
RZ: u16,
RTRGT: u16,
RSTAT: Flags,
RCALL: u16,
RSK: u16,
RSR: u16,
SSK: Vec<u16>,
SIN: Vec<u16>,
SMAIN: Vec<u16>,
SCODE: Vec<u16>,
IP: u16,
}
impl Machine {
fn push_stack(&mut self, val: u16) {
self.SSK[self.RSK as usize] = val;
self.RSK -= 1;
}
fn pop_stack(&mut self) -> u16 {
let v = self.SSK[self.RSK as usize];
self.RSK += 1;
v
}
}
impl Machine {
fn write_reg(&mut self, reg: Register, val: u16) {
match reg {
Register::RX => self.RX = val,
Register::RY => self.RY = val,
Register::RZ => self.RZ = val,
Register::RTRGT => self.RTRGT = val,
Register::RSTAT => {}, // writes have no effect
Register::RCALL => {},
Register::NULL => {},
}
}
fn read_reg(&mut self, reg: Register) -> u16 {
match reg {
Register::RX => self.RX,
Register::RY => self.RY,
Register::RZ => self.RZ,
Register::RTRGT => 0,
Register::RSTAT => self.RSTAT.bits(),
Register::RCALL => self.RCALL,
Register::NULL => 0,
}
}
fn read(&mut self, a: Arg) -> u16 {
match a {
Arg::Register(r) => self.read_reg(r),
Arg::SIN(v) => self.SIN[v as usize],
Arg::SMAIN(v) => self.SMAIN[v as usize],
}
}
fn write(&mut self, a: Arg, val: u16) {
match a {
Arg::Register(r) => self.write_reg(r, val),
Arg::SIN(v) => self.SIN[v as usize] = val,
Arg::SMAIN(v) => self.SMAIN[v as usize] = val,
}
}
}
#[derive(Debug, Eq, PartialEq, Copy, Clone)]
enum Register {
RX,
RY,
RZ,
RTRGT,
RSTAT,
RCALL,
NULL,
}
impl Register {
fn into_register(x: u16) -> Self {
match x {
1 => Self::RX,
2 => Self::RY,
3 => Self::RZ,
4 => Self::RTRGT,
5 => Self::RSTAT,
6 => Self::RCALL,
_ => Self::NULL,
}
}
}
bitflags::bitflags! {
struct Flags: u16 {
const FZERO = 0b0000000000000001;
const FEQUL = 0b0000000000000010;
const FLT = 0b0000000000000100;
const FGT = 0b0000000000001000;
const FCRRY = 0b0000000000010000;
const FINF = 0b0000000000100000;
const FSE = 0b0000000001000000;
const FSF = 0b0000000010000000;
}
}
#[derive(Debug, Eq, PartialEq)]
enum OP {
HALT = 0x000 ,
NOOP = 0x001 ,
INC = 0x002 ,
DEC = 0x003 ,
ADD = 0x004 ,
SUB = 0x005 ,
MUL = 0x006 ,
DIV = 0x007 ,
ADDC = 0x008 ,
SUBC = 0x009 ,
READ = 0x00A ,
WRIT = 0x00B ,
CPY = 0x00C ,
MCPY = 0x00D ,
ICPY = 0x00E ,
CMP = 0x00F ,
AND = 0x010 ,
OR = 0x011 ,
CMPL = 0x012 ,
LSHF = 0x013 ,
RSHF = 0x014 ,
PUSH = 0x015 ,
POP = 0x016 ,
CFLG = 0x017 ,
CALL = 0x018 ,
RTRN = 0x019 ,
RTRV = 0x01A ,
RTL = 0x01B ,
RTR = 0x01C ,
CIP = 0x01D ,
BSWP = 0x01E ,
JUMP = 0x01F ,
JZRO = 0x020 ,
JEQU = 0x021 ,
JLT = 0x022 ,
JGT = 0x023 ,
JCRY = 0x024 ,
JINF = 0x025 ,
JSE = 0x026 ,
JSF = 0x027 ,
CZRO = 0x030 ,
CCRY = 0x034 ,
XOR = 0x040 ,
SWAP = 0x041 ,
RCPT = 0x042 ,
RCPF = 0x043 ,
}
impl std::convert::Into<OP> for u16 {
fn into(self) -> OP {
match self {
0x000 => OP::HALT,
0x001 => OP::NOOP,
0x002 => OP::INC,
0x003 => OP::DEC,
0x004 => OP::ADD,
0x005 => OP::SUB,
0x006 => OP::MUL,
0x007 => OP::DIV,
0x008 => OP::ADDC,
0x009 => OP::SUBC,
0x00A => OP::READ,
0x00B => OP::WRIT,
0x00C => OP::CPY,
0x00D => OP::MCPY,
0x00E => OP::ICPY,
0x00F => OP::CMP,
0x010 => OP::AND,
0x011 => OP::OR,
0x012 => OP::CMPL,
0x013 => OP::LSHF,
0x014 => OP::RSHF,
0x015 => OP::PUSH,
0x016 => OP::POP,
0x017 => OP::CFLG,
0x018 => OP::CALL,
0x019 => OP::RTRN,
0x01A => OP::RTRV,
0x01B => OP::RTL,
0x01C => OP::RTR,
0x01D => OP::CIP,
0x01E => OP::BSWP,
0x01F => OP::JUMP,
0x020 => OP::JZRO,
0x021 => OP::JEQU,
0x022 => OP::JLT,
0x023 => OP::JGT,
0x024 => OP::JCRY,
0x025 => OP::JINF,
0x026 => OP::JSE,
0x027 => OP::JSF,
0x030 => OP::CZRO,
0x034 => OP::CCRY,
0x040 => OP::XOR,
0x041 => OP::SWAP,
0x042 => OP::RCPT,
0x043 => OP::RCPF,
_ => panic!(),
}
}
}
#[derive(Debug, PartialEq, Eq)]
enum AFLG {
Register,
SIN,
SMAIN,
}
impl From<u16> for AFLG {
fn from(x: u16) -> AFLG {
match x {
0b00 => AFLG::Register,
0b01 => AFLG::Register,
0b10 => AFLG::SIN,
0b11 => AFLG::SMAIN,
_ => panic!(),
}
}
}
bitfield::bitfield! {
struct DefinitionWord(u16);
impl Debug;
into OP, opcode, set_opcode: 8, 0;
sign, set_sign: 9, 9;
into AFLG, aflg2, set_aflg2: 11, 10;
into AFLG, aflg1, set_aflg1: 13, 12;
into AFLG, aflg0, set_aflg0: 15, 14;
}
#[derive(Debug)]
struct Instruction {
def: DefinitionWord,
args: Vec<Arg>,
}
impl Default for Machine {
fn default() -> Self {
Self {
SSK: vec![0_u16; 65536_usize],
SIN: vec![0_u16; 65536_usize],
SMAIN: vec![0_u16; 65536_usize],
SCODE: vec![0_u16; 65536_usize],
RSTAT: Flags::FSE,
IP: 1,
RCALL: 0,
RSK: 0xffff,
RSR: 0,
RTRGT: 0,
RX: 0,
RY: 0,
RZ: 0,
}
}
}
#[derive(Debug)]
enum Chunk {
SIN(Vec<u16>),
SCODE(Vec<u16>),
SMAIN(Vec<u16>),
}
impl Chunk {
fn inner(&self) -> &[u16] {
match self {
Chunk::SIN(x) => x,
Chunk::SCODE(x) => x,
Chunk::SMAIN(x) => x,
}
}
}
fn parse_chunk(mut input: impl std::io::Read) -> Chunk {
let space_index = input.read_u8().unwrap();
let length = input.read_u16::<LE>().unwrap();
let mut data = Vec::new();
for _ in 0..length {
data.push(input.read_u16::<LE>().unwrap());
}
match space_index {
1 => Chunk::SIN(data),
2 => Chunk::SCODE(data),
3 => Chunk::SMAIN(data),
_ => panic!("unknown space index {}", space_index),
}
}
fn arg_count(opcode: OP) -> usize {
match opcode {
OP::HALT => 0,
OP::NOOP => 0,
OP::INC => 1,
OP::DEC => 1,
OP::ADD => 2,
OP::SUB => 2,
OP::MUL => 2,
OP::DIV => 2,
OP::ADDC => 2,
OP::SUBC => 2,
OP::READ => 1,
OP::WRIT => 1,
OP::CPY => 2,
OP::MCPY => 3,
OP::ICPY => 2,
OP::CMP => 2,
OP::AND => 2,
OP::OR => 2,
OP::CMPL => 1,
OP::LSHF => 2,
OP::RSHF => 2,
OP::PUSH => 1,
OP::POP => 1,
OP::CFLG => 0,
OP::CALL => 3,
OP::RTRN => 0,
OP::RTRV => 0,
OP::RTL => 2,
OP::RTR => 2,
OP::CIP => 1,
OP::BSWP => 1,
OP::JUMP => 0,
OP::JZRO => 0,
OP::JEQU => 0,
OP::JLT => 0,
OP::JGT => 0,
OP::JCRY => 0,
OP::JINF => 0,
OP::JSE => 0,
OP::JSF => 0,
OP::CZRO => 0,
OP::CCRY => 0,
OP::XOR => 2,
OP::SWAP => 2,
OP::RCPT => 2,
OP::RCPF => 2,
}
}
fn parse_instruction(mut bytes: &[u16]) -> (Instruction, usize) {
//println!("{:#02x?}", &bytes[0..8]);
let mut consumed = 0;
let def = DefinitionWord(bytes[0]);
bytes = &bytes[1..];
consumed += 1;
let ac = arg_count(def.opcode());
let mut args = Vec::new();
let mut register_count = 0;
if ac > 0 && def.aflg0() == AFLG::Register {
register_count += 1;
}
if ac > 1 && def.aflg1() == AFLG::Register {
register_count += 1;
}
if ac > 2 && def.aflg2() == AFLG::Register {
register_count += 1;
}
let mut registers = Vec::new();
if register_count > 0 {
let mut aw = bytes[0];
//dbg!(aw);
bytes = &bytes[1..];
consumed += 1;
//println!("{:016b}", aw);
aw >>= 4;
for i in 0..register_count {
registers.insert(0_usize, Register::into_register(aw & 0b1111));
aw >>= 4;
}
}
for idx in 0..ac {
let aflg = match idx {
0 => def.aflg0(),
1 => def.aflg1(),
2 => def.aflg2(),
_ => unreachable!(),
};
match aflg {
AFLG::Register => {
args.push(Arg::Register(registers.remove(0)))
},
AFLG::SIN => {
args.push(Arg::SIN(bytes[0]));
bytes = &bytes[1..];
consumed += 1;
}
AFLG::SMAIN => {
args.push(Arg::SMAIN(bytes[0]));
bytes = &bytes[1..];
consumed += 1;
}
}
}
(Instruction { def, args }, consumed)
}
#[derive(Debug, Copy, Clone)]
enum Arg {
Register(Register),
SIN(u16),
SMAIN(u16),
}
impl Arg {
fn as_immediate(self) -> u16 {
match self {
Arg::Register(_) => panic!("a register is not an immediate"),
Arg::SIN(x) => x,
Arg::SMAIN(x) => x,
}
}
fn offset(self, by: i32) -> Self {
match self {
Arg::Register(_) => panic!("offset must be memory space address"),
Arg::SIN(x) => Arg::SIN((x as i32 + by) as u16),
Arg::SMAIN(x) => Arg::SMAIN((x as i32 + by) as u16),
}
}
}
impl Machine {
fn step(&mut self) -> bool {
let (ci, inc_by) = parse_instruction(&self.SCODE[(self.IP as usize)..]);
//dbg!(&ci);
let oc = ci.def.opcode();
let mut should_inc_ip = true;
match oc {
OP::HALT => return false,
OP::NOOP => {},
OP::XOR => {
let arg0 = self.read(ci.args[0]);
let arg1 = self.read(ci.args[1]);
let xored = arg0 ^ arg1;
self.write(ci.args[0], xored);
self.RSTAT.set(Flags::FZERO, xored == 0);
}
OP::CALL => {
if self.RTRGT == 0 {
} else {
self.push_stack(self.RCALL);
self.push_stack(self.RSTAT.bits());
self.push_stack(self.RX);
self.push_stack(self.RY);
self.push_stack(self.RZ);
self.push_stack(u16::try_from(inc_by).unwrap());
self.RSR = self.RSK;
self.RX = self.read(ci.args[0]);
self.RY = self.read(ci.args[1]);
self.RZ = self.read(ci.args[2]);
self.RCALL = self.IP;
self.IP = self.RTRGT;
should_inc_ip = false;
self.RTRGT = 0;
self.RSTAT = Flags::empty();
}
}
OP::POP => {
let val = self.pop_stack();
self.write(ci.args[0], val);
self.RSTAT.set(Flags::FZERO, val == 0);
self.RSTAT.set(Flags::FSE, self.RSK == 0xffff);
}
OP::PUSH => {
let val = self.read(ci.args[0]);
self.push_stack(val);
self.RSTAT.set(Flags::FZERO, val == 0);
self.RSTAT.set(Flags::FSF, self.RSK == 0);
}
OP::ICPY => {
let arg0 = ci.args[0].as_immediate();
self.write(ci.args[1], arg0);
self.RSTAT.set(Flags::FZERO, arg0 == 0);
//eprintln!("copying {} to {:?}", arg0, ci.args[1]);
}
OP::CMP => {
let a = self.read(ci.args[0]);
let b = self.read(ci.args[1]);
self.RSTAT.set(Flags::FZERO, a == 0 || b == 0);
self.RSTAT.set(Flags::FEQUL, a == b);
self.RSTAT.set(Flags::FLT, a < b);
self.RSTAT.set(Flags::FGT, a > b);
}
OP::JEQU => {
if self.RSTAT.contains(Flags::FEQUL) {
if self.RTRGT == 0 {
return false;
}
self.IP = self.RTRGT;
should_inc_ip = false;
}
}
OP::JUMP => {
if self.RTRGT == 0 {
return false;
}
self.IP = self.RTRGT;
should_inc_ip = false;
}
OP::INC => {
let mut val = self.read(ci.args[0]);
val = val.wrapping_add(1);
self.write(ci.args[0], val);
self.RSTAT.set(Flags::FZERO, val == 0);
}
OP::READ => {
let input = 0;
std::io::stdin().read(&mut [input]).unwrap();
self.write(ci.args[0], input as u16);
self.RSTAT.set(Flags::FZERO, input == 0);
}
OP::CPY => {
let val = self.read(ci.args[0]);
self.write(ci.args[1], val);
self.RSTAT.set(Flags::FZERO, val == 0);
}
OP::RTRN => {
self.RSK = self.RSR;
self.RX = self.pop_stack();
self.IP = self.RCALL + self.RX;
should_inc_ip = false;
self.RZ = self.pop_stack();
self.RY = self.pop_stack();
self.RX = self.pop_stack();
self.RSTAT = Flags::from_bits(self.pop_stack()).unwrap();
self.RCALL = self.pop_stack();
self.push_stack(self.RTRGT);
self.RTRGT = 0;
}
OP::RCPT => {
let by: i32 = if ci.def.sign() == 1 {
self.RTRGT as i32
} else {
-(self.RTRGT as i32)
};
let from = self.read(ci.args[0]);
let to = ci.args[1].offset(by);
self.write(to, from);
self.RSTAT.set(Flags::FZERO, from == 0);
}
OP::RCPF => {
let by: i32 = if ci.def.sign() == 1 {
self.RTRGT as i32
} else {
-(self.RTRGT as i32)
};
let from = ci.args[0].offset(by);
let from_val = self.read(from);
self.RSTAT.set(Flags::FZERO, from_val == 0);
self.write(ci.args[1], from_val);
}
OP::CMPL => {
let mut val = self.read(ci.args[0]);
val = !val;
self.write(ci.args[0], val);
self.RSTAT.set(Flags::FZERO, val == 0);
}
OP::AND => {
let arg0 = self.read(ci.args[0]);
let arg1 = self.read(ci.args[1]);
let anded = arg0 & arg1;
self.write(ci.args[0], anded);
self.RSTAT.set(Flags::FZERO, anded == 0);
}
OP::OR => {
let arg0 = self.read(ci.args[0]);
let arg1 = self.read(ci.args[1]);
let ored = arg0 | arg1;
self.write(ci.args[0], ored);
self.RSTAT.set(Flags::FZERO, ored == 0);
}
OP::WRIT => {
let arg = self.read(ci.args[0]);
println!("ch: `{}`", arg);
}
_ => { eprintln!("unsupported opcode {:?}, continuing", oc); return false; }
}
if should_inc_ip {
self.IP += inc_by as u16;
}
true
}
}
fn dump_instructions(mut x: &[u16]) {
while !x.is_empty() {
let (i, inc) = parse_instruction(x);
x = &x[inc..];
//println!("{:?}", i);
}
}
fn main() {
let bin = include_bytes!("../challenge.mctf");
let (header, body) = bin.split_at(4);
assert_eq!(header, b"mCTZ");
let decompressed = zstd::decode_all(body).unwrap();
let mut csr = Cursor::new(decompressed);
let code = parse_chunk(&mut csr);
let sin = parse_chunk(&mut csr);
let mut codeinner = code.inner();
let mut m = Machine::default();
m.SIN[1..=sin.inner().len()].copy_from_slice(sin.inner());
m.SCODE[1..=code.inner().len()].copy_from_slice(code.inner());
// dump_instructions(&m.SCODE[1..=code.inner().len()]);
while m.step() {}
}