<commit_before>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::{JoinHandle};
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory,
               Killable, Pollable, Poller, Terminal};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

#[derive(Debug)]
pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Irq(sys_util::Error),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

type Result<T> = std::result::Result<T, Error>;

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path())
                    .map_err(Error::Kernel)?
        },
        None => {
            return Err(Error::Kernel(
                    io::Error::new(io::ErrorKind::NotFound,
                                   "missing kernel path")))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => return Err(Error::Kernel(
                io::Error::new(io::ErrorKind::NotFound,
                               "missing kernel cmdline")))
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions)
            .map_err(Error::GuestMemory)?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr,
                               &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(vm.get_memory(),
                             kernel_start_addr,
                             cmdline_addr,
                             cmdline.to_bytes().len() + 1,
                             vcpu_count as u8)?;

    let mut io_bus = devices::Bus::new();
    let com_evt_1_3 = EventFd::new().map_err(Error::EventFd)?;
    let com_evt_2_4 = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial = Arc::new(Mutex::new(devices::Serial::new_out(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?,
            Box::new(stdout()))));

    io_bus.insert(stdio_serial.clone(), 0x3f8, 0x8).unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2f8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?))), 0x3e8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2e8, 0x8)
            .unwrap();

    vm.register_irqfd(&com_evt_1_3, 4)
            .map_err(Error::Irq)?;
    vm.register_irqfd(&com_evt_2_4, 3)
            .map_err(Error::Irq)?;

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(Error::Vcpu)?;
        x86_64::configure_vcpu(vm.get_memory(), kernel_start_addr, &kvm,
                               &vcpu, cpu_id as u64, vcpu_count as u64)?;
        vcpu_handles.push(thread::Builder::new()
                .name(format!("fc_vcpu{}", cpu_id))
                .spawn(move || {
            unsafe {
                extern "C" fn handle_signal() {}
                // Our signal handler does nothing and is trivially async signal safe.
                register_signal_handler(0, handle_signal)
                        .expect("failed to register vcpu signal handler");
            }

            vcpu_thread_barrier.wait();

            loop {
                match vcpu.run() {
                    Ok(run) => match run {
                        VcpuExit::IoIn(addr, data) => {
                            io_bus.read(addr as u64, data);
                        },
                        VcpuExit::IoOut(addr, data) => {
                            io_bus.write(addr as u64, data);
                        },
                        VcpuExit::MmioRead(_, _) => {},
                        VcpuExit::MmioWrite(_, _) => {},
                        VcpuExit::Hlt => {
                            println!("KVM_EXIT_HLT");
                            break;
                        },
                        VcpuExit::Shutdown => {
                            println!("KVM_EXIT_SHUTDOWN");
                            break;
                        },
                        r => {
                            println!("unexpected exit reason: {:?}", r);
                            break;
                        }
                    },
                    Err(e) => {
                        match e.errno() {
                            libc::EAGAIN | libc::EINTR => {},
                            _ => {
                                println!("vcpu hit unknown error: {:?}", e);
                                break;
                            }
                        }
                    }
                }

                if kill_signaled.load(Ordering::SeqCst) {
                    break;
                }
            }

            vcpu_exit_evt
                .write(1)
                .expect("failed to signal vcpu exit eventfd");

        }).map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(stdio_serial, exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(stdio_serial: Arc<Mutex<devices::Serial>>,
               exit_evt: EventFd,
               kill_signaled: Arc<AtomicBool>,
               vcpu_handles: Vec<JoinHandle<()>>) -> Result<()> {
    const EXIT: u32 = 0;
    const STDIN: u32 = 1;

    let stdin_handle = io::stdin();
    let stdin_lock = stdin_handle.lock();
    stdin_lock
            .set_raw_mode()
            .expect("failed to set terminal raw mode");

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));
    pollables.push((STDIN, &stdin_lock as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                },
                STDIN => {
                    let mut out = [0u8; 64];
                    match stdin_lock.read_raw(&mut out[..]) {
                        Ok(0) => {
                            // Zero-length read indicates EOF. Remove from pollables.
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Err(e) => {
                            println!("error while reading stdin: {:?}", e);
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Ok(count) => {
                            stdio_serial
                                    .lock()
                                    .unwrap()
                                    .queue_input_bytes(&out[..count])
                                    .expect("failed to queue bytes into serial port");
                        },
                    }
                }
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            },
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    stdin_lock
            .set_canon_mode()
            .expect("failed to restore canonical mode for terminal");

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
        0x00, 0xd8,       /* add %bl, %al */
        0x04, '0' as u8,  /* add $'0', %al */
        0xee,             /* out %al, (%dx) */
        0xb0, '\n' as u8, /* mov $'\n', %al */
        0xee,             /* out %al, (%dx) */
        0xf4,             /* hlt */
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            },
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break
            },
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

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

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<commit_msg>vmm: fix style manually as cargo fmt messes up the code<commit_after>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::{JoinHandle};
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory,
               Killable, Pollable, Poller, Terminal};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

#[derive(Debug)]
pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Irq(sys_util::Error),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

type Result<T> = std::result::Result<T, Error>;

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path())
                    .map_err(Error::Kernel)?
        },
        None => {
            return Err(Error::Kernel(
                    io::Error::new(io::ErrorKind::NotFound,
                                   "missing kernel path")))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => return Err(Error::Kernel(
                io::Error::new(io::ErrorKind::NotFound,
                               "missing kernel cmdline")))
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions)
            .map_err(Error::GuestMemory)?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr,
                               &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(vm.get_memory(),
                             kernel_start_addr,
                             cmdline_addr,
                             cmdline.to_bytes().len() + 1,
                             vcpu_count as u8)?;

    let mut io_bus = devices::Bus::new();
    let com_evt_1_3 = EventFd::new().map_err(Error::EventFd)?;
    let com_evt_2_4 = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial = Arc::new(Mutex::new(devices::Serial::new_out(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?,
            Box::new(stdout()))));

    io_bus.insert(stdio_serial.clone(), 0x3f8, 0x8).unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2f8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?))), 0x3e8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2e8, 0x8)
            .unwrap();

    vm.register_irqfd(&com_evt_1_3, 4)
            .map_err(Error::Irq)?;
    vm.register_irqfd(&com_evt_2_4, 3)
            .map_err(Error::Irq)?;

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(Error::Vcpu)?;
        x86_64::configure_vcpu(vm.get_memory(), kernel_start_addr, &kvm,
                               &vcpu, cpu_id as u64, vcpu_count as u64)?;
        vcpu_handles.push(thread::Builder::new()
                .name(format!("fc_vcpu{}", cpu_id))
                .spawn(move || {
            unsafe {
                extern "C" fn handle_signal() {}
                // Our signal handler does nothing and is trivially async signal safe.
                register_signal_handler(0, handle_signal)
                        .expect("failed to register vcpu signal handler");
            }

            vcpu_thread_barrier.wait();

            loop {
                match vcpu.run() {
                    Ok(run) => match run {
                        VcpuExit::IoIn(addr, data) => {
                            io_bus.read(addr as u64, data);
                        },
                        VcpuExit::IoOut(addr, data) => {
                            io_bus.write(addr as u64, data);
                        },
                        VcpuExit::MmioRead(_, _) => {},
                        VcpuExit::MmioWrite(_, _) => {},
                        VcpuExit::Hlt => {
                            println!("KVM_EXIT_HLT");
                            break;
                        },
                        VcpuExit::Shutdown => {
                            println!("KVM_EXIT_SHUTDOWN");
                            break;
                        },
                        r => {
                            println!("unexpected exit reason: {:?}", r);
                            break;
                        }
                    },
                    Err(e) => {
                        match e.errno() {
                            libc::EAGAIN | libc::EINTR => {},
                            _ => {
                                println!("vcpu hit unknown error: {:?}", e);
                                break;
                            }
                        }
                    }
                }

                if kill_signaled.load(Ordering::SeqCst) {
                    break;
                }
            }

            vcpu_exit_evt
                .write(1)
                .expect("failed to signal vcpu exit eventfd");

        }).map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(stdio_serial, exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(stdio_serial: Arc<Mutex<devices::Serial>>,
               exit_evt: EventFd,
               kill_signaled: Arc<AtomicBool>,
               vcpu_handles: Vec<JoinHandle<()>>) -> Result<()> {
    const EXIT: u32 = 0;
    const STDIN: u32 = 1;

    let stdin_handle = io::stdin();
    let stdin_lock = stdin_handle.lock();
    stdin_lock
            .set_raw_mode()
            .expect("failed to set terminal raw mode");

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));
    pollables.push((STDIN, &stdin_lock as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                },
                STDIN => {
                    let mut out = [0u8; 64];
                    match stdin_lock.read_raw(&mut out[..]) {
                        Ok(0) => {
                            // Zero-length read indicates EOF. Remove from pollables.
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Err(e) => {
                            println!("error while reading stdin: {:?}", e);
                            pollables.retain(|&pollable| pollable.0 != STDIN);
                        },
                        Ok(count) => {
                            stdio_serial
                                    .lock()
                                    .unwrap()
                                    .queue_input_bytes(&out[..count])
                                    .expect("failed to queue bytes into serial port");
                        },
                    }
                }
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            },
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    stdin_lock
            .set_canon_mode()
            .expect("failed to restore canonical mode for terminal");

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        /* mov $0x3f8, %dx */
        0xba,
        0xf8,
        0x03,
        /* add %bl, %al */
        0x00,
        0xd8,
        /* add $'0', %al */
        0x04,
        '0' as u8,
        /* out %al, (%dx) */
        0xee,
        /* mov $'\n', %al */
        0xb0,
        '\n' as u8,
        /* out %al, (%dx) */
        0xee,
        /* hlt */
        0xf4,
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            },
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break
            },
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

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

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<|endoftext|>


<commit_before>extern crate rspgod;
extern crate postgres;
extern crate rustc_serialize;

use postgres::{Connection, SslMode};
use rspgod::types::{Change};
use rustc_serialize::json;

#[derive(Debug)]
struct TestRecord {
    id:   i32,
    name: String,
}

#[test]
fn sanity_test() {
    with_clean_database(|c| {
        let records = vec![
            TestRecord { id: 1, name: "Michael Baker".to_string() },
            TestRecord { id: 2, name: "Josh Cheek".to_string()   },
        ];

        for record in records.iter() {
            create_record(c, &record);
        }

        fetch_records(c);
    });
}

#[test]
fn basic_insert() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string(), };
        create_record(c, &record);
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 1);
        let change:Change = json::decode(&updates[0][..]).unwrap();
        assert!(match change {
            Change::Insert {..} => { true },
            _                   => { false },
        });
    });
}

#[test]
fn basic_delete() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string() };
        create_record(c, &record);
        delete_record(c, 1);
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 2);
        let change:Change = json::decode(&updates[1][..]).unwrap();
        assert!(match change {
            Change::Delete {..} => { true },
            _                   => { false },
        });
    });
}

#[test]
fn basic_update() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string() };
        create_record(c, &record);
        update_record(c, TestRecord { id: 1, name: "Bichael Maker".to_string() });
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 2);
        let change:Change = json::decode(&updates[1][..]).unwrap();
        assert!(match change {
            Change::Update {..} => { true },
            _                   => { false },
        });
    });
}

//
// [TODO] I want to move a lot of these into a utility module when I can figure out how to do that
//

fn fetch_updates(c: &Connection) -> Vec<String> {
    let stmt = c.prepare("SELECT * FROM pg_logical_slot_peek_changes('slot', NULL, NULL)").unwrap();
    let mut result = vec![];
    for r in stmt.query(&[]).unwrap() {
        result.push(r.get(2));
    }
    result
}

fn create_slot(c: &Connection) {
    let stmt = c.prepare("select * from pg_create_logical_replication_slot('slot', 'thingy')").unwrap();
    stmt.execute(&[]).unwrap();
}

fn drop_slot(c: &Connection) {
    let stmt = c.prepare("select pg_drop_replication_slot('slot')").unwrap();
    match stmt.execute(&[]) {
        _ => {},
    }
}

fn with_slot<F>(f: F) where F:Fn(&Connection) -> () {
    with_clean_database(|c| {
        drop_slot(c);
        create_slot(c);
        f(c);
        drop_slot(c);
    });
}

fn with_clean_database<F>(f: F) where F:Fn(&Connection) -> () {
    let c = connection();
    reset_database(&c);
    f(&c);
    drop_database(&c);
}

fn reset_database(c: &Connection) {
    drop_database(&c);
    create_database(&c);
}

fn create_database(c: &Connection) {
    let stmt = c.prepare("create table test_table (id int primary key, name text)").unwrap();
    stmt.execute(&[]).unwrap();
}

fn fetch_records(c: &Connection) -> Vec<TestRecord> {
    let stmt = c.prepare("select id, name from test_table").unwrap();
    let results = stmt.query(&[]).unwrap();
    results.iter().map(|r| {
        TestRecord {
            id:   r.get(0),
            name: r.get(1),
        }
    }).collect()
}

fn create_record(c: &Connection, r: &TestRecord) {
    let stmt = c.prepare("insert into test_table (id, name) values ($1, $2)").unwrap();
    stmt.execute(&[
        &r.id,
        &r.name,
    ]).unwrap();
}

fn delete_record(c: &Connection, id: i32) {
    let stmt = c.prepare("delete from test_table where id = $1").unwrap();
    stmt.execute(&[&id]).unwrap();
}

fn update_record(c: &Connection, new_record: TestRecord) {
    let stmt = c.prepare("update test_table set name = $2 where id = $1").unwrap();
    stmt.execute(&[&new_record.id, &new_record.name]).unwrap();
}

fn drop_database(c: &Connection) {
    let stmt = c.prepare("drop table if exists test_table").unwrap();
    stmt.execute(&[]).unwrap();
}

fn connection_string() -> String {
    match std::env::var("POSTGRES_URL") {
        Ok(val) => val,
        Err(e)  => panic!(format!("You must set the POSTGRES_URL environment variable to point to a running Postgres test database when running automated tests: {}", e)),
    }
}

fn connection() -> Connection {
    match Connection::connect(&connection_string()[..], &SslMode::None) {
        Ok(c)  => c,
        Err(e) => panic!(format!("Could not connect to the test database: {}", e)),
    }
}
<commit_msg>Removes the dependency on rspgod to make the tests faster<commit_after>extern crate postgres;
extern crate rustc_serialize;

use postgres::{Connection, SslMode};
use rustc_serialize::json::Json;

#[derive(Debug)]
struct TestRecord {
    id:   i32,
    name: String,
}

#[test]
fn sanity_test() {
    with_clean_database(|c| {
        let records = vec![
            TestRecord { id: 1, name: "Michael Baker".to_string() },
            TestRecord { id: 2, name: "Josh Cheek".to_string()   },
        ];

        for record in records.iter() {
            create_record(c, &record);
        }

        fetch_records(c);
    });
}

#[test]
fn basic_insert() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string(), };
        create_record(c, &record);
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 1);
        let data = Json::from_str(&updates[0][..]).unwrap();
        let change = data.as_object().unwrap();
        let variant = change.get("variant").unwrap().as_string().unwrap();
        assert_eq!(variant, "Insert");
    });
}

#[test]
fn basic_delete() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string() };
        create_record(c, &record);
        delete_record(c, 1);
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 2);
        let data = Json::from_str(&updates[1][..]).unwrap();
        let change = data.as_object().unwrap();
        let variant = change.get("variant").unwrap().as_string().unwrap();
        assert_eq!(variant, "Delete");
    });
}

#[test]
fn basic_update() {
    with_slot(|c| {
        let record = TestRecord { id: 1, name: "Michael Baker".to_string() };
        create_record(c, &record);
        update_record(c, TestRecord { id: 1, name: "Bichael Maker".to_string() });
        let updates = fetch_updates(c);
        assert_eq!(updates.len(), 2);
        let data = Json::from_str(&updates[1][..]).unwrap();
        let change = data.as_object().unwrap();
        let variant = change.get("variant").unwrap().as_string().unwrap();
        assert_eq!(variant, "Update");
    });
}

//
// [TODO] I want to move a lot of these into a utility module when I can figure out how to do that
//

fn fetch_updates(c: &Connection) -> Vec<String> {
    let stmt = c.prepare("SELECT * FROM pg_logical_slot_peek_changes('slot', NULL, NULL)").unwrap();
    let mut result = vec![];
    for r in stmt.query(&[]).unwrap() {
        result.push(r.get(2));
    }
    result
}

fn create_slot(c: &Connection) {
    let stmt = c.prepare("select * from pg_create_logical_replication_slot('slot', 'thingy')").unwrap();
    stmt.execute(&[]).unwrap();
}

fn drop_slot(c: &Connection) {
    let stmt = c.prepare("select pg_drop_replication_slot('slot')").unwrap();
    match stmt.execute(&[]) {
        _ => {},
    }
}

fn with_slot<F>(f: F) where F:Fn(&Connection) -> () {
    with_clean_database(|c| {
        drop_slot(c);
        create_slot(c);
        f(c);
        drop_slot(c);
    });
}

fn with_clean_database<F>(f: F) where F:Fn(&Connection) -> () {
    let c = connection();
    reset_database(&c);
    f(&c);
    drop_database(&c);
}

fn reset_database(c: &Connection) {
    drop_database(&c);
    create_database(&c);
}

fn create_database(c: &Connection) {
    let stmt = c.prepare("create table test_table (id int primary key, name text)").unwrap();
    stmt.execute(&[]).unwrap();
}

fn fetch_records(c: &Connection) -> Vec<TestRecord> {
    let stmt = c.prepare("select id, name from test_table").unwrap();
    let results = stmt.query(&[]).unwrap();
    results.iter().map(|r| {
        TestRecord {
            id:   r.get(0),
            name: r.get(1),
        }
    }).collect()
}

fn create_record(c: &Connection, r: &TestRecord) {
    let stmt = c.prepare("insert into test_table (id, name) values ($1, $2)").unwrap();
    stmt.execute(&[
        &r.id,
        &r.name,
    ]).unwrap();
}

fn delete_record(c: &Connection, id: i32) {
    let stmt = c.prepare("delete from test_table where id = $1").unwrap();
    stmt.execute(&[&id]).unwrap();
}

fn update_record(c: &Connection, new_record: TestRecord) {
    let stmt = c.prepare("update test_table set name = $2 where id = $1").unwrap();
    stmt.execute(&[&new_record.id, &new_record.name]).unwrap();
}

fn drop_database(c: &Connection) {
    let stmt = c.prepare("drop table if exists test_table").unwrap();
    stmt.execute(&[]).unwrap();
}

fn connection_string() -> String {
    match std::env::var("POSTGRES_URL") {
        Ok(val) => val,
        Err(e)  => panic!(format!("You must set the POSTGRES_URL environment variable to point to a running Postgres test database when running automated tests: {}", e)),
    }
}

fn connection() -> Connection {
    match Connection::connect(&connection_string()[..], &SslMode::None) {
        Ok(c)  => c,
        Err(e) => panic!(format!("Could not connect to the test database: {}", e)),
    }
}
<|endoftext|>


<commit_before>use clippy_utils::diagnostics::span_lint_and_help;
use rustc_hir::{HirId, Item, ItemKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty::Const;
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::sym;

declare_clippy_lint! {
    /// ### What it does
    /// Displays a warning when a struct with a trailing zero-sized array is declared without a `repr` attribute.
    ///
    /// ### Why is this bad?
    /// Zero-sized arrays aren't very useful in Rust itself, so such a struct is likely being created to pass to C code or in some other situation where control over memory layout matters (for example, in conjuction with manual allocation to make it easy to compute the offset of the array). Either way, `#[repr(C)]` (or another `repr` attribute) is needed.
    ///
    /// ### Example
    /// ```rust
    /// struct RarelyUseful {
    ///     some_field: u32,
    ///     last: [u32; 0],
    /// }
    /// ```
    ///
    /// Use instead:
    /// ```rust
    /// #[repr(C)]
    /// struct MoreOftenUseful {
    ///     some_field: usize,
    ///     last: [u32; 0],
    /// }
    /// ```
    pub TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR,
    nursery,
    "struct with a trailing zero-sized array but without `#[repr(C)]` or another `repr` attribute"
}
declare_lint_pass!(TrailingZeroSizedArrayWithoutRepr => [TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR]);

impl<'tcx> LateLintPass<'tcx> for TrailingZeroSizedArrayWithoutRepr {
    fn check_item(&mut self, cx: &LateContext<'tcx>, item: &'tcx Item<'tcx>) {
        if is_struct_with_trailing_zero_sized_array(cx, item) && !has_repr_attr(cx, item.hir_id()) {
            span_lint_and_help(
                cx,
                TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR,
                item.span,
                "trailing zero-sized array in a struct which is not marked with a `repr` attribute",
                None,
                &format!(
                    "consider annotating `{}` with `#[repr(C)]` or another `repr` attribute",
                    cx.tcx.def_path_str(item.def_id.to_def_id())
                ),
            );
        }
    }
}

fn is_struct_with_trailing_zero_sized_array(cx: &LateContext<'tcx>, item: &'tcx Item<'tcx>) -> bool {
    if_chain! {
        // First check if last field is an array
        if let ItemKind::Struct(data, _) = &item.kind;
        if let Some(last_field) = data.fields().last();
        if let rustc_hir::TyKind::Array(_, length) = last_field.ty.kind;

        // Then check if that that array zero-sized
        let length_ldid = cx.tcx.hir().local_def_id(length.hir_id);
        let length = Const::from_anon_const(cx.tcx, length_ldid);
        let length = length.try_eval_usize(cx.tcx, cx.param_env);
        if let Some(length) = length;
        then {
            length == 0
        } else {
            false
        }
    }
}

fn has_repr_attr(cx: &LateContext<'_>, hir_id: HirId) -> bool {
    cx.tcx.hir().attrs(hir_id).iter().any(|attr| attr.has_name(sym::repr))
}
<commit_msg>formatting 🙃<commit_after>use clippy_utils::diagnostics::span_lint_and_help;
use rustc_hir::{HirId, Item, ItemKind};
use rustc_lint::{LateContext, LateLintPass};
use rustc_middle::ty::Const;
use rustc_session::{declare_lint_pass, declare_tool_lint};
use rustc_span::sym;

declare_clippy_lint! {
    /// ### What it does
    /// Displays a warning when a struct with a trailing zero-sized array is declared without a `repr` attribute.
    ///
    /// ### Why is this bad?
    /// Zero-sized arrays aren't very useful in Rust itself, so such a struct is likely being created to pass to C code or in some other situation where control over memory layout matters (for example, in conjuction with manual allocation to make it easy to compute the offset of the array). Either way, `#[repr(C)]` (or another `repr` attribute) is needed.
    ///
    /// ### Example
    /// ```rust
    /// struct RarelyUseful {
    ///     some_field: u32,
    ///     last: [u32; 0],
    /// }
    /// ```
    ///
    /// Use instead:
    /// ```rust
    /// #[repr(C)]
    /// struct MoreOftenUseful {
    ///     some_field: usize,
    ///     last: [u32; 0],
    /// }
    /// ```
    pub TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR,
    nursery,
    "struct with a trailing zero-sized array but without `#[repr(C)]` or another `repr` attribute"
}
declare_lint_pass!(TrailingZeroSizedArrayWithoutRepr => [TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR]);

impl<'tcx> LateLintPass<'tcx> for TrailingZeroSizedArrayWithoutRepr {
    fn check_item(&mut self, cx: &LateContext<'tcx>, item: &'tcx Item<'tcx>) {
        if is_struct_with_trailing_zero_sized_array(cx, item) && !has_repr_attr(cx, item.hir_id()) {
            span_lint_and_help(
                cx,
                TRAILING_ZERO_SIZED_ARRAY_WITHOUT_REPR,
                item.span,
                "trailing zero-sized array in a struct which is not marked with a `repr` attribute",
                None,
                &format!(
                    "consider annotating `{}` with `#[repr(C)]` or another `repr` attribute",
                    cx.tcx.def_path_str(item.def_id.to_def_id())
                ),
            );
        }
    }
}

fn is_struct_with_trailing_zero_sized_array(cx: &LateContext<'tcx>, item: &'tcx Item<'tcx>) -> bool {
    if_chain! {
        // First check if last field is an array
        if let ItemKind::Struct(data, _) = &item.kind;
        if let Some(last_field) = data.fields().last();
        if let rustc_hir::TyKind::Array(_, length) = last_field.ty.kind;

        // Then check if that that array zero-sized
        let length_ldid = cx.tcx.hir().local_def_id(length.hir_id);
        let length = Const::from_anon_const(cx.tcx, length_ldid);
        let length = length.try_eval_usize(cx.tcx, cx.param_env);
        if let Some(length) = length;
        then {
            length == 0
        } else {
            false
        }
    }
}

fn has_repr_attr(cx: &LateContext<'_>, hir_id: HirId) -> bool {
    cx.tcx.hir().attrs(hir_id).iter().any(|attr| attr.has_name(sym::repr))
}
<|endoftext|>


<commit_before>// Copyright 2015 Ted Mielczarek. See the COPYRIGHT
// file at the top-level directory of this distribution.

use chrono::{TimeZone, Utc};
use failure::Fail;

use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::ops::Deref;
use std::path::Path;

use minidump::{self, *};

use crate::process_state::{CallStack, CallStackInfo, LinuxStandardBase, ProcessState};
use crate::stackwalker;
use crate::symbols::*;
use crate::system_info::SystemInfo;

/// An error encountered during minidump processing.
#[derive(Debug, Fail)]
pub enum ProcessError {
    #[fail(display = "Failed to read minidump")]
    MinidumpReadError(minidump::Error),
    #[fail(display = "An unknown error occurred")]
    UnknownError,
    #[fail(display = "The system information stream was not found")]
    MissingSystemInfo,
    #[fail(display = "The thread list stream was not found")]
    MissingThreadList,
}

impl From<minidump::Error> for ProcessError {
    fn from(err: minidump::Error) -> ProcessError {
        ProcessError::MinidumpReadError(err)
    }
}

/// Unwind all threads in `dump` and return a `ProcessState`.
///
/// # Examples
///
/// ```
/// use minidump::Minidump;
/// use std::path::PathBuf;
/// use breakpad_symbols::{Symbolizer, SimpleSymbolSupplier};
///
/// # std::env::set_current_dir(env!("CARGO_MANIFEST_DIR"));
/// # fn foo() -> Result<(), minidump_processor::ProcessError> {
/// let mut dump = Minidump::read_path("../testdata/test.dmp")?;
/// let supplier = SimpleSymbolSupplier::new(vec!(PathBuf::from("../testdata/symbols")));
/// let symbolizer = Symbolizer::new(supplier);
/// let state = minidump_processor::process_minidump(&mut dump, &symbolizer)?;
/// assert_eq!(state.threads.len(), 2);
/// println!("Processed {} threads", state.threads.len());
/// # Ok(())
/// # }
/// # fn main() { foo().unwrap() }
/// ```
pub fn process_minidump<'a, T, P>(
    dump: &Minidump<'a, T>,
    symbol_provider: &P,
) -> Result<ProcessState, ProcessError>
where
    T: Deref<Target = [u8]> + 'a,
    P: SymbolProvider,
{
    // No Evil JSON Here!
    process_minidump_with_evil(dump, symbol_provider, None)
}

/// The same as `process_minidump` but with an extra evil little json file.
///
/// This is a hack to support mozilla's legacy workflow, just use `process_minidump`.
pub fn process_minidump_with_evil<'a, T, P>(
    dump: &Minidump<'a, T>,
    symbol_provider: &P,
    evil_json: Option<&Path>,
) -> Result<ProcessState, ProcessError>
where
    T: Deref<Target = [u8]> + 'a,
    P: SymbolProvider,
{
    // Thread list is required for processing.
    let thread_list = dump
        .get_stream::<MinidumpThreadList>()
        .or(Err(ProcessError::MissingThreadList))?;
    // Try to get thread names, but it's only a nice-to-have.
    let thread_names = dump
        .get_stream::<MinidumpThreadNames>()
        .unwrap_or_else(|_| MinidumpThreadNames::default());

    // System info is required for processing.
    let dump_system_info = dump
        .get_stream::<MinidumpSystemInfo>()
        .or(Err(ProcessError::MissingSystemInfo))?;

    let mut os_version = format!(
        "{}.{}.{}",
        dump_system_info.raw.major_version,
        dump_system_info.raw.minor_version,
        dump_system_info.raw.build_number
    );
    if let Some(csd_version) = dump_system_info.csd_version() {
        os_version.push(' ');
        os_version.push_str(&csd_version);
    }

    let linux_standard_base = dump.get_stream::<MinidumpLinuxLsbRelease>().ok();
    let linux_cpu_info = dump
        .get_stream::<MinidumpLinuxCpuInfo>()
        .unwrap_or_default();
    let _linux_environ = dump.get_stream::<MinidumpLinuxEnviron>().ok();
    let _linux_proc_status = dump.get_stream::<MinidumpLinuxProcStatus>().ok();

    // Extract everything we care about from linux streams here.
    // We don't eagerly process them in the minidump crate because there's just
    // tons of random information in there and it's not obvious what anyone
    // would care about. So just providing an iterator and letting minidump-processor
    // pull out the things it cares about is simple and effective.

    let mut cpu_microcode_version = None;
    for (key, val) in linux_cpu_info.iter() {
        if key.as_bytes() == b"microcode" {
            cpu_microcode_version = val
                .to_str()
                .ok()
                .and_then(|val| val.strip_prefix("0x"))
                .and_then(|val| u64::from_str_radix(val, 16).ok());
            break;
        }
    }

    let linux_standard_base = linux_standard_base.map(|linux_standard_base| {
        let mut lsb = LinuxStandardBase::default();
        for (key, val) in linux_standard_base.iter() {
            match key.as_bytes() {
                b"DISTRIB_ID" | b"ID" => lsb.id = val.to_string_lossy().into_owned(),
                b"DISTRIB_RELEASE" | b"VERSION_ID" => {
                    lsb.release = val.to_string_lossy().into_owned()
                }
                b"DISTRIB_CODENAME" | b"VERSION_CODENAME" => {
                    lsb.codename = val.to_string_lossy().into_owned()
                }
                b"DISTRIB_DESCRIPTION" | b"PRETTY_NAME" => {
                    lsb.description = val.to_string_lossy().into_owned()
                }
                _ => {}
            }
        }
        lsb
    });

    let cpu_info = dump_system_info
        .cpu_info()
        .map(|string| string.into_owned());

    let system_info = SystemInfo {
        os: dump_system_info.os,
        os_version: Some(os_version),
        cpu: dump_system_info.cpu,
        cpu_info,
        cpu_microcode_version,
        cpu_count: dump_system_info.raw.number_of_processors as usize,
    };

    let mac_crash_info = dump
        .get_stream::<MinidumpMacCrashInfo>()
        .ok()
        .map(|info| info.raw);
    // Process create time is optional.
    let (process_id, process_create_time) =
        if let Ok(misc_info) = dump.get_stream::<MinidumpMiscInfo>() {
            (
                misc_info.raw.process_id().cloned(),
                misc_info.process_create_time(),
            )
        } else {
            (None, None)
        };
    // If Breakpad info exists in dump, get dump and requesting thread ids.
    let breakpad_info = dump.get_stream::<MinidumpBreakpadInfo>();
    let (dump_thread_id, requesting_thread_id) = if let Ok(info) = breakpad_info {
        (info.dump_thread_id, info.requesting_thread_id)
    } else {
        (None, None)
    };
    // Get exception info if it exists.
    let exception_stream = dump.get_stream::<MinidumpException>().ok();
    let exception_ref = exception_stream.as_ref();
    let (crash_reason, crash_address, crashing_thread_id) = if let Some(exception) = exception_ref {
        (
            Some(exception.get_crash_reason(system_info.os, system_info.cpu)),
            Some(exception.get_crash_address(system_info.os)),
            Some(exception.get_crashing_thread_id()),
        )
    } else {
        (None, None, None)
    };
    let exception_context = exception_ref.and_then(|e| e.context.as_ref());
    // Get assertion
    let assertion = None;
    let modules = match dump.get_stream::<MinidumpModuleList>() {
        Ok(module_list) => module_list,
        // Just give an empty list, simplifies things.
        Err(_) => MinidumpModuleList::new(),
    };
    let unloaded_modules = match dump.get_stream::<MinidumpUnloadedModuleList>() {
        Ok(module_list) => module_list,
        // Just give an empty list, simplifies things.
        Err(_) => MinidumpUnloadedModuleList::new(),
    };
    let memory_list = dump.get_stream::<MinidumpMemoryList>().unwrap_or_default();
    let memory_info_list = dump.get_stream::<MinidumpMemoryInfoList>().ok();
    let linux_maps = dump.get_stream::<MinidumpLinuxMaps>().ok();
    let _memory_info = UnifiedMemoryInfoList::new(memory_info_list, linux_maps).unwrap_or_default();

    // Get memory list
    let mut threads = vec![];
    let mut requesting_thread = None;
    for (i, thread) in thread_list.threads.iter().enumerate() {
        // If this is the thread that wrote the dump, skip processing it.
        if dump_thread_id.is_some() && dump_thread_id.unwrap() == thread.raw.thread_id {
            threads.push(CallStack::with_info(CallStackInfo::DumpThreadSkipped));
            continue;
        }
        // If this thread requested the dump then try to use the exception
        // context if it exists. (prefer the exception stream's thread id over
        // the breakpad info stream's thread id.)
        let context = if crashing_thread_id
            .or(requesting_thread_id)
            .map(|id| id == thread.raw.thread_id)
            .unwrap_or(false)
        {
            requesting_thread = Some(i);
            exception_context.or_else(|| thread.context.as_ref())
        } else {
            thread.context.as_ref()
        };

        let stack = thread.stack.as_ref().or_else(|| {
            // Windows probably gave us null RVAs for our stack memory descriptors.
            // If this happens, then we need to look up the memory region by address.
            let stack_addr = thread.raw.stack.start_of_memory_range;
            memory_list.memory_at_address(stack_addr)
        });

        let mut stack = stackwalker::walk_stack(&context, stack, &modules, symbol_provider);

        let name = thread_names
            .get_name(thread.raw.thread_id)
            .map(|cow| cow.into_owned());
        stack.thread_name = name;

        stack.last_error_value = thread.last_error(system_info.cpu, &memory_list);

        threads.push(stack);
    }

    // Collect up info on unimplemented/unknown modules
    let unknown_streams = dump.unknown_streams().collect();
    let unimplemented_streams = dump.unimplemented_streams().collect();

    // Get symbol stats from the symbolizer
    let symbol_stats = symbol_provider.stats();

    // Finally, handle the evil JSON file (get module signing certs)
    let cert_info = evil_json.and_then(handle_evil).unwrap_or_else(HashMap::new);

    Ok(ProcessState {
        process_id,
        time: Utc.timestamp(dump.header.time_date_stamp as i64, 0),
        process_create_time,
        cert_info,
        crash_reason,
        crash_address,
        assertion,
        requesting_thread,
        system_info,
        linux_standard_base,
        mac_crash_info,
        threads,
        modules,
        unloaded_modules,
        unknown_streams,
        unimplemented_streams,
        symbol_stats,
    })
}

fn handle_evil(evil_path: &Path) -> Option<HashMap<String, String>> {
    use log::{error, warn};
    use serde_json::Value::{self, *};

    // Get the evil json
    let evil_json = File::open(evil_path)
        .map_err(|e| {
            error!("Could not load Extra JSON at {:?}", evil_path);
            e
        })
        .ok()?;

    let buf = BufReader::new(evil_json);
    let json: Value = serde_json::from_reader(buf)
        .map_err(|e| {
            error!("Could not parse Extra JSON (was not valid JSON)");
            e
        })
        .ok()?;

    // Get module signing info
    let temp_obj;
    let certs = match json.get("ModuleSignatureInfo") {
        Some(Object(obj)) => obj,
        Some(String(string)) => {
            // Possible the signature info was wrapped in a string by mistake,
            // So try to parse that string as an object.
            temp_obj = serde_json::from_str(string)
                .map_err(|e| {
                    error!("Could not parse Extra JSON's ModuleSignatureInfo (not an object)");
                    error!("ModuleSignatureInfo: {}", string);
                    e
                })
                .ok()?;
            &temp_obj
        }
        _ => {
            error!("Could not parse Extra JSON's ModuleSignatureInfo (not an object)");
            return None;
        }
    };

    // Each certificate lists the modules it applies to, but we want the
    // reverse mapping -- module names to certificates. Invert the map.
    let mut cert_map = HashMap::new();
    for (cert, modules) in certs {
        if let Array(modules) = modules {
            for module in modules {
                if let String(module) = module {
                    cert_map.insert(module.clone(), cert.clone());
                }
            }
        } else {
            warn!(
                "Extra JSON had corrupt entry -- \"{}\": {:?}",
                cert, modules
            );
        }
    }

    Some(cert_map)
}
<commit_msg>Use the thread names in the evil json<commit_after>// Copyright 2015 Ted Mielczarek. See the COPYRIGHT
// file at the top-level directory of this distribution.

use chrono::{TimeZone, Utc};
use failure::Fail;

use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::ops::Deref;
use std::path::Path;

use minidump::{self, *};

use crate::process_state::{CallStack, CallStackInfo, LinuxStandardBase, ProcessState};
use crate::stackwalker;
use crate::symbols::*;
use crate::system_info::SystemInfo;

/// An error encountered during minidump processing.
#[derive(Debug, Fail)]
pub enum ProcessError {
    #[fail(display = "Failed to read minidump")]
    MinidumpReadError(minidump::Error),
    #[fail(display = "An unknown error occurred")]
    UnknownError,
    #[fail(display = "The system information stream was not found")]
    MissingSystemInfo,
    #[fail(display = "The thread list stream was not found")]
    MissingThreadList,
}

impl From<minidump::Error> for ProcessError {
    fn from(err: minidump::Error) -> ProcessError {
        ProcessError::MinidumpReadError(err)
    }
}

/// Unwind all threads in `dump` and return a `ProcessState`.
///
/// # Examples
///
/// ```
/// use minidump::Minidump;
/// use std::path::PathBuf;
/// use breakpad_symbols::{Symbolizer, SimpleSymbolSupplier};
///
/// # std::env::set_current_dir(env!("CARGO_MANIFEST_DIR"));
/// # fn foo() -> Result<(), minidump_processor::ProcessError> {
/// let mut dump = Minidump::read_path("../testdata/test.dmp")?;
/// let supplier = SimpleSymbolSupplier::new(vec!(PathBuf::from("../testdata/symbols")));
/// let symbolizer = Symbolizer::new(supplier);
/// let state = minidump_processor::process_minidump(&mut dump, &symbolizer)?;
/// assert_eq!(state.threads.len(), 2);
/// println!("Processed {} threads", state.threads.len());
/// # Ok(())
/// # }
/// # fn main() { foo().unwrap() }
/// ```
pub fn process_minidump<'a, T, P>(
    dump: &Minidump<'a, T>,
    symbol_provider: &P,
) -> Result<ProcessState, ProcessError>
where
    T: Deref<Target = [u8]> + 'a,
    P: SymbolProvider,
{
    // No Evil JSON Here!
    process_minidump_with_evil(dump, symbol_provider, None)
}

/// The same as `process_minidump` but with an extra evil little json file.
///
/// This is a hack to support mozilla's legacy workflow, just use `process_minidump`.
pub fn process_minidump_with_evil<'a, T, P>(
    dump: &Minidump<'a, T>,
    symbol_provider: &P,
    evil_json: Option<&Path>,
) -> Result<ProcessState, ProcessError>
where
    T: Deref<Target = [u8]> + 'a,
    P: SymbolProvider,
{
    // Thread list is required for processing.
    let thread_list = dump
        .get_stream::<MinidumpThreadList>()
        .or(Err(ProcessError::MissingThreadList))?;
    // Try to get thread names, but it's only a nice-to-have.
    let thread_names = dump
        .get_stream::<MinidumpThreadNames>()
        .unwrap_or_else(|_| MinidumpThreadNames::default());

    // System info is required for processing.
    let dump_system_info = dump
        .get_stream::<MinidumpSystemInfo>()
        .or(Err(ProcessError::MissingSystemInfo))?;

    let mut os_version = format!(
        "{}.{}.{}",
        dump_system_info.raw.major_version,
        dump_system_info.raw.minor_version,
        dump_system_info.raw.build_number
    );
    if let Some(csd_version) = dump_system_info.csd_version() {
        os_version.push(' ');
        os_version.push_str(&csd_version);
    }

    let linux_standard_base = dump.get_stream::<MinidumpLinuxLsbRelease>().ok();
    let linux_cpu_info = dump
        .get_stream::<MinidumpLinuxCpuInfo>()
        .unwrap_or_default();
    let _linux_environ = dump.get_stream::<MinidumpLinuxEnviron>().ok();
    let _linux_proc_status = dump.get_stream::<MinidumpLinuxProcStatus>().ok();

    // Extract everything we care about from linux streams here.
    // We don't eagerly process them in the minidump crate because there's just
    // tons of random information in there and it's not obvious what anyone
    // would care about. So just providing an iterator and letting minidump-processor
    // pull out the things it cares about is simple and effective.

    let mut cpu_microcode_version = None;
    for (key, val) in linux_cpu_info.iter() {
        if key.as_bytes() == b"microcode" {
            cpu_microcode_version = val
                .to_str()
                .ok()
                .and_then(|val| val.strip_prefix("0x"))
                .and_then(|val| u64::from_str_radix(val, 16).ok());
            break;
        }
    }

    let linux_standard_base = linux_standard_base.map(|linux_standard_base| {
        let mut lsb = LinuxStandardBase::default();
        for (key, val) in linux_standard_base.iter() {
            match key.as_bytes() {
                b"DISTRIB_ID" | b"ID" => lsb.id = val.to_string_lossy().into_owned(),
                b"DISTRIB_RELEASE" | b"VERSION_ID" => {
                    lsb.release = val.to_string_lossy().into_owned()
                }
                b"DISTRIB_CODENAME" | b"VERSION_CODENAME" => {
                    lsb.codename = val.to_string_lossy().into_owned()
                }
                b"DISTRIB_DESCRIPTION" | b"PRETTY_NAME" => {
                    lsb.description = val.to_string_lossy().into_owned()
                }
                _ => {}
            }
        }
        lsb
    });

    let cpu_info = dump_system_info
        .cpu_info()
        .map(|string| string.into_owned());

    let system_info = SystemInfo {
        os: dump_system_info.os,
        os_version: Some(os_version),
        cpu: dump_system_info.cpu,
        cpu_info,
        cpu_microcode_version,
        cpu_count: dump_system_info.raw.number_of_processors as usize,
    };

    let mac_crash_info = dump
        .get_stream::<MinidumpMacCrashInfo>()
        .ok()
        .map(|info| info.raw);
    // Process create time is optional.
    let (process_id, process_create_time) =
        if let Ok(misc_info) = dump.get_stream::<MinidumpMiscInfo>() {
            (
                misc_info.raw.process_id().cloned(),
                misc_info.process_create_time(),
            )
        } else {
            (None, None)
        };
    // If Breakpad info exists in dump, get dump and requesting thread ids.
    let breakpad_info = dump.get_stream::<MinidumpBreakpadInfo>();
    let (dump_thread_id, requesting_thread_id) = if let Ok(info) = breakpad_info {
        (info.dump_thread_id, info.requesting_thread_id)
    } else {
        (None, None)
    };
    // Get exception info if it exists.
    let exception_stream = dump.get_stream::<MinidumpException>().ok();
    let exception_ref = exception_stream.as_ref();
    let (crash_reason, crash_address, crashing_thread_id) = if let Some(exception) = exception_ref {
        (
            Some(exception.get_crash_reason(system_info.os, system_info.cpu)),
            Some(exception.get_crash_address(system_info.os)),
            Some(exception.get_crashing_thread_id()),
        )
    } else {
        (None, None, None)
    };
    let exception_context = exception_ref.and_then(|e| e.context.as_ref());
    // Get assertion
    let assertion = None;
    let modules = match dump.get_stream::<MinidumpModuleList>() {
        Ok(module_list) => module_list,
        // Just give an empty list, simplifies things.
        Err(_) => MinidumpModuleList::new(),
    };
    let unloaded_modules = match dump.get_stream::<MinidumpUnloadedModuleList>() {
        Ok(module_list) => module_list,
        // Just give an empty list, simplifies things.
        Err(_) => MinidumpUnloadedModuleList::new(),
    };
    let memory_list = dump.get_stream::<MinidumpMemoryList>().unwrap_or_default();
    let memory_info_list = dump.get_stream::<MinidumpMemoryInfoList>().ok();
    let linux_maps = dump.get_stream::<MinidumpLinuxMaps>().ok();
    let _memory_info = UnifiedMemoryInfoList::new(memory_info_list, linux_maps).unwrap_or_default();

    // Get the evil JSON file (thread names and module certificates)
    let evil = evil_json.and_then(handle_evil).unwrap_or_default();

    let mut threads = vec![];
    let mut requesting_thread = None;
    for (i, thread) in thread_list.threads.iter().enumerate() {
        // If this is the thread that wrote the dump, skip processing it.
        if dump_thread_id.is_some() && dump_thread_id.unwrap() == thread.raw.thread_id {
            threads.push(CallStack::with_info(CallStackInfo::DumpThreadSkipped));
            continue;
        }
        // If this thread requested the dump then try to use the exception
        // context if it exists. (prefer the exception stream's thread id over
        // the breakpad info stream's thread id.)
        let context = if crashing_thread_id
            .or(requesting_thread_id)
            .map(|id| id == thread.raw.thread_id)
            .unwrap_or(false)
        {
            requesting_thread = Some(i);
            exception_context.or_else(|| thread.context.as_ref())
        } else {
            thread.context.as_ref()
        };

        let stack = thread.stack.as_ref().or_else(|| {
            // Windows probably gave us null RVAs for our stack memory descriptors.
            // If this happens, then we need to look up the memory region by address.
            let stack_addr = thread.raw.stack.start_of_memory_range;
            memory_list.memory_at_address(stack_addr)
        });

        let mut stack = stackwalker::walk_stack(&context, stack, &modules, symbol_provider);

        let name = thread_names
            .get_name(thread.raw.thread_id)
            .map(|cow| cow.into_owned())
            .or_else(|| evil.thread_names.get(&thread.raw.thread_id).cloned());
        stack.thread_name = name;

        stack.last_error_value = thread.last_error(system_info.cpu, &memory_list);

        threads.push(stack);
    }

    // Collect up info on unimplemented/unknown modules
    let unknown_streams = dump.unknown_streams().collect();
    let unimplemented_streams = dump.unimplemented_streams().collect();

    // Get symbol stats from the symbolizer
    let symbol_stats = symbol_provider.stats();

    Ok(ProcessState {
        process_id,
        time: Utc.timestamp(dump.header.time_date_stamp as i64, 0),
        process_create_time,
        cert_info: evil.certs,
        crash_reason,
        crash_address,
        assertion,
        requesting_thread,
        system_info,
        linux_standard_base,
        mac_crash_info,
        threads,
        modules,
        unloaded_modules,
        unknown_streams,
        unimplemented_streams,
        symbol_stats,
    })
}

/// Things extracted from the Evil JSON File
#[derive(Debug, Default)]
struct Evil {
    /// module name => cert
    certs: HashMap<String, String>,
    /// thread id => thread name
    thread_names: HashMap<u32, String>,
}

fn handle_evil(evil_path: &Path) -> Option<Evil> {
    use log::error;
    use serde_json::map::Map;
    use serde_json::Value;
    use std::str::FromStr;

    // Get the evil json
    let evil_json = File::open(evil_path)
        .map_err(|e| {
            error!("Could not load Extra JSON at {:?}", evil_path);
            e
        })
        .ok()?;

    let buf = BufReader::new(evil_json);
    let mut json: Map<String, Value> = serde_json::from_reader(buf)
        .map_err(|e| {
            error!("Could not parse Extra JSON (was not valid JSON)");
            e
        })
        .ok()?;

    // Of course evil json contains a string-that-can-be-parsed-as-a-json-object
    // instead of having a normal json object!
    fn evil_obj<K, V>(json: &mut Map<String, Value>, field_name: &str) -> Option<HashMap<K, V>>
    where
        K: for<'de> serde::de::Deserialize<'de> + Eq + std::hash::Hash,
        V: for<'de> serde::de::Deserialize<'de>,
    {
        json.remove(field_name).and_then(|val| {
            match val {
                Value::Object(_) => serde_json::from_value(val).ok(),
                Value::String(string) => serde_json::from_str(&string).ok(),
                _ => None,
            }
            .or_else(|| {
                error!("Could not parse Evil JSON's {} (not an object)", field_name);
                None
            })
        })
    }

    // Convert certs from
    // "cert_name1": ["module1", "module2", ...], "cert_name2": ...
    // to
    // "module1": "cert_name1", "module2": "cert_name1", ...
    let certs = evil_obj(&mut json, "ModuleSignatureInfo")
        .map(|certs: HashMap<String, Vec<String>>| {
            let mut cert_map = HashMap::new();
            for (cert, modules) in certs {
                for module in modules {
                    cert_map.insert(module, cert.clone());
                }
            }
            cert_map
        })
        .unwrap_or_default();

    // Get thread name mappings

    // In typical evil json fashion, this list doesn't conform to even the evil_obj format!
    // It's just a set of comma-separated int:string pairs, with a trailing comma.
    // This cannot be parsed as JSON at all, since the keys are not strings. So we just
    // do a sloppy `split` based parse and hope we don't encounter thread names with commas
    // in them because I hate this JSON file with a passion.
    //
    // ex: 123: "name1", 456: "name",
    let thread_names = json
        .remove("ThreadIdNameMapping")
        .unwrap_or_default()
        .as_str()
        .unwrap_or_default()
        .split(',')
        .filter_map(|entry| {
            entry.split_once(":").and_then(|(key, val)| {
                let key = u32::from_str(key).ok();
                let val = val
                    .strip_prefix('"')
                    .and_then(|val| val.strip_suffix('"'))
                    .map(String::from);
                key.zip(val)
            })
        })
        .collect();

    Some(Evil {
        certs,
        thread_names,
    })
}
<|endoftext|>


<commit_before>use std::{cmp::min, ffi::OsStr, fs::File, io, mem, num::Wrapping, path::PathBuf, time::Duration};

use clap::Parser;
use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 100;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    endpoint: String,
    filename: PathBuf,
    client: reqwest::blocking::Client,
    rng: SmallRng,
    spinner_idx: Wrapping<usize>,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(endpoint: &str, filename: PathBuf) -> Importer {
        Importer {
            endpoint: endpoint.to_owned(),
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            spinner_idx: Wrapping(0),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put(format!("{}/import/lichess", self.endpoint))
            .json(&self.batch)
            .send()
            .expect("send batch");

        self.spinner_idx += Wrapping(1);
        let spinner = &['⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'];

        println!(
            "{} {:?}: {}: {} - {}",
            spinner[self.spinner_idx.0 % spinner.len()],
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(true, |name| name == "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.00,

                _ if rating >= 2500 => 1.00,

                Speed::Rapid if rating >= 2200 => 1.00,
                Speed::Rapid if rating >= 2000 => 0.83,
                Speed::Rapid if rating >= 1800 => 0.46,
                Speed::Rapid if rating >= 1600 => 0.39,

                Speed::Blitz if rating >= 2200 => 0.38,
                Speed::Blitz if rating >= 2000 => 0.18,
                Speed::Blitz if rating >= 1600 => 0.13,

                Speed::Bullet if rating >= 2200 => 0.48,
                Speed::Bullet if rating >= 2000 => 0.27,
                Speed::Bullet if rating >= 1800 => 0.19,
                Speed::Bullet if rating >= 1600 => 0.18,

                Speed::UltraBullet => 1.00,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.00
            } else {
                0.50
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

#[derive(Parser)]
struct Args {
    #[clap(long, default_value = "http://localhost:9002")]
    endpoint: String,
    pgns: Vec<PathBuf>,
}

fn main() -> Result<(), io::Error> {
    let args = Args::parse();

    for arg in args.pgns {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.extension() == Some(OsStr::new("bz2")) {
            println!("Reading compressed {:?} ...", arg);
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            println!("Reading {:?} ...", arg);
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(&args.endpoint, arg);
        reader.read_all(&mut importer)?;
        importer.send();
    }

    Ok(())
}
<commit_msg>use new Outcome::from_ascii()<commit_after>use std::{cmp::min, ffi::OsStr, fs::File, io, mem, num::Wrapping, path::PathBuf, time::Duration};

use clap::Parser;
use pgn_reader::{BufferedReader, Color, Outcome, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, rngs::SmallRng, Rng, SeedableRng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 100;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    endpoint: String,
    filename: PathBuf,
    client: reqwest::blocking::Client,
    rng: SmallRng,
    spinner_idx: Wrapping<usize>,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(endpoint: &str, filename: PathBuf) -> Importer {
        Importer {
            endpoint: endpoint.to_owned(),
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            rng: SmallRng::from_seed([
                0x19, 0x29, 0xab, 0x17, 0xc6, 0xfa, 0xb0, 0xe9, 0x4b, 0x44, 0xd8, 0x07, 0x09, 0xbf,
                0x1d, 0x87, 0xbd, 0xd8, 0xb3, 0x2f, 0xe1, 0xe2, 0xa0, 0x1a, 0x9e, 0x30, 0x98, 0xd7,
                0xef, 0xd5, 0x7a, 0x1d,
            ]),
            spinner_idx: Wrapping(0),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put(format!("{}/import/lichess", self.endpoint))
            .json(&self.batch)
            .send()
            .expect("send batch");

        self.spinner_idx += Wrapping(1);
        let spinner = &['⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷'];

        println!(
            "{} {:?}: {}: {} - {}",
            spinner[self.spinner_idx.0 % spinner.len()],
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match Outcome::from_ascii(value.as_bytes()) {
                Ok(outcome) => self.current.winner = outcome.winner(),
                Err(_) => self.skip = true,
            }
        } else if key == b"FEN" {
            if value.as_bytes() == b"rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" {
                // https://github.com/ornicar/lichess-db/issues/40
                self.current.fen = None;
            } else {
                self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
            }
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(true, |name| name == "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.00,

                _ if rating >= 2500 => 1.00,

                Speed::Rapid if rating >= 2200 => 1.00,
                Speed::Rapid if rating >= 2000 => 0.83,
                Speed::Rapid if rating >= 1800 => 0.46,
                Speed::Rapid if rating >= 1600 => 0.39,

                Speed::Blitz if rating >= 2200 => 0.38,
                Speed::Blitz if rating >= 2000 => 0.18,
                Speed::Blitz if rating >= 1600 => 0.13,

                Speed::Bullet if rating >= 2200 => 0.48,
                Speed::Bullet if rating >= 2000 => 0.27,
                Speed::Bullet if rating >= 1800 => 0.19,
                Speed::Bullet if rating >= 1600 => 0.18,

                Speed::UltraBullet => 1.00,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.00
            } else {
                0.50
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= self.rng.sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

#[derive(Parser)]
struct Args {
    #[clap(long, default_value = "http://localhost:9002")]
    endpoint: String,
    pgns: Vec<PathBuf>,
}

fn main() -> Result<(), io::Error> {
    let args = Args::parse();

    for arg in args.pgns {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.extension() == Some(OsStr::new("bz2")) {
            println!("Reading compressed {:?} ...", arg);
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            println!("Reading {:?} ...", arg);
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(&args.endpoint, arg);
        reader.read_all(&mut importer)?;
        importer.send();
    }

    Ok(())
}
<|endoftext|>