<commit_before>extern crate rs6502;

use std::io::{self, BufRead, Write};
use std::sync::mpsc::{channel, Receiver};
use std::thread;

use rs6502::{Cpu, Disassembler};

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    code_offset: u16,
    clock_rate: Option<u32>,
    receiver: Receiver<String>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
}

impl VirtualMachine {
    pub fn new<CR>(cpu: Cpu, code_offset: u16, clock_rate: CR) -> VirtualMachine
        where CR: Into<Option<u32>>
    {
        let (tx, rx) = channel();

        thread::spawn(move || {
            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            loop {
                let mut line = String::new();
                let stdin = io::stdin();
                stdin.lock().read_line(&mut line).expect("Could not read line");
                tx.send(line).unwrap();
            }
        });

        VirtualMachine {
            cpu: cpu,
            code_offset: code_offset,
            clock_rate: clock_rate.into(),
            receiver: rx,
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
        }
    }

    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while n < clock_rate && !self.broken {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    println!("");
                    println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
                }
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                println!("");
                println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
            }
        }
    }

    pub fn try_execute_command(&mut self) {
        if let Ok(input) = self.receiver.try_recv() {
            let mut input: String = input.trim().into();

            if input.len() == 0 {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    input = self.last_command.clone();
                }
            }

            if input == "source" {
                self.dump_disassembly();
            }

            if input == "list" {
                self.dump_local_disassembly();
            }

            if input.starts_with("monitor") {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    self.enable_memory_monitor(&input);
                }
            }

            if input.starts_with("memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else if parts.len() == 3 {
                    if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                        if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                            self.cpu.memory[dst as usize] = src;
                        } else {
                            println!("ERR: Unable to parse source byte value");
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                } else {
                    if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..],
                                                               16) {
                        for p in &parts[2..] {
                            if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                                self.cpu.memory[dst] = byte;
                                dst += 0x01;
                            } else {
                                println!("ERR: Unable to parse source byte value");
                            }
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                }
            }

            if input.starts_with("memdmp") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    let page_number: usize =
                        parts[1].parse().expect("Unable to parse memory page number");

                    self.dump_memory_page(page_number);
                }
            }

            if input.starts_with("break") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                        if self.breakpoints[addr] > 0 {
                            self.breakpoints[addr] = 0;
                        } else {
                            self.breakpoints[addr] = 1;
                        }
                    } else {
                        println!("Err: Unable to parse breakpoint address");
                    }
                } else {
                    println!("ERR: Requires 1 argument");
                }
            }

            if input == "continue" {
                self.broken = false;
                println!("Execution resumed");
            }

            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            // Don't assign a blank command as a last command
            if input.len() > 0 {
                self.last_command = input.clone();
            }
        }

    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            println!("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;

            println!("{:?}", self.monitor);
        }
    }

    fn dump_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pairs = disassembler.disassemble_with_addresses(self.cpu.get_code());
        let result = self.highlight_lines(self.cpu.registers.PC as usize, pairs, false).join("");
        print!("{}", result);
    }

    fn dump_local_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pc = self.cpu.registers.PC as usize;
        let code = self.cpu.get_code();
        let pairs = disassembler.disassemble_with_addresses(code);
        let result = self.highlight_lines(pc, pairs, true).join("");
        print!("{}", result);
    }

    pub fn dump_memory_page(&self, page: usize) {
        let mut addr = page * 0x100;
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            print!("{:04X}: ", addr);
            for b in chunk {
                print!("{:02X} ", *b);
            }
            addr += 0x08;
            println!("");
        }
    }

    pub fn dump_memory(&self) {
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                print!("{:02X} ", *b);
            }
            println!("");
        }
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc - self.code_offset as usize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as usize) < start || (pair.1 as usize) > base + 0x0A {
                    continue;
                }
            }
            let breakpoint = self.breakpoints[self.code_offset as usize + pair.1 as usize] > 0x00;
            let current_line = pc as u16 == self.code_offset + pair.1;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}<commit_msg>Add stepping line by line<commit_after>extern crate rs6502;

use std::io::{self, BufRead, Write};
use std::sync::mpsc::{channel, Receiver};
use std::thread;

use rs6502::{Cpu, Disassembler};

#[derive(Debug)]
pub struct MemoryMonitor {
    pub enabled: bool,
    start_addr: usize,
    end_addr: usize,
}

pub struct VirtualMachine {
    pub cpu: Cpu,
    pub monitor: MemoryMonitor,
    code_offset: u16,
    clock_rate: Option<u32>,
    receiver: Receiver<String>,
    last_command: String,
    breakpoints: [u8; 64 * 1024],
    broken: bool,
    step: bool,
}

impl VirtualMachine {
    pub fn new<CR>(cpu: Cpu, code_offset: u16, clock_rate: CR) -> VirtualMachine
        where CR: Into<Option<u32>>
    {
        let (tx, rx) = channel();

        thread::spawn(move || {
            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            loop {
                let mut line = String::new();
                let stdin = io::stdin();
                stdin.lock().read_line(&mut line).expect("Could not read line");
                tx.send(line).unwrap();
            }
        });

        VirtualMachine {
            cpu: cpu,
            code_offset: code_offset,
            clock_rate: clock_rate.into(),
            receiver: rx,
            monitor: MemoryMonitor {
                enabled: false,
                start_addr: 0,
                end_addr: 0,
            },
            last_command: "".into(),
            breakpoints: [0; 64 * 1024],
            broken: false,
            step: false,
        }
    }

    pub fn cycle(&mut self) {
        if let Some(clock_rate) = self.clock_rate {
            let mut n = 0;
            while (n < clock_rate && !self.broken) || self.step {
                n += self.cpu.step().expect("SEGFAULT") as u32;
                if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                    self.broken = true;
                    println!("");
                    println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
                }
                self.step = false;
            }
        } else {
            self.cpu.step().expect("SEGFAULT");
            self.step = false;
            if self.breakpoints[self.cpu.registers.PC as usize] > 0 {
                self.broken = true;
                println!("");
                println!("BREAKPOINT hit at {:04X}", self.cpu.registers.PC);
            }
        }
    }

    pub fn try_execute_command(&mut self) {
        if let Ok(input) = self.receiver.try_recv() {
            let mut input: String = input.trim().into();

            if input.len() == 0 {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    input = self.last_command.clone();
                }
            }

            if input == "source" {
                self.dump_disassembly();
            }

            if input == "list" {
                self.dump_local_disassembly();
            }

            if input.starts_with("monitor") {
                if self.monitor.enabled {
                    self.monitor.enabled = false;
                } else {
                    self.enable_memory_monitor(&input);
                }
            }

            if input.starts_with("memset") {
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() < 3 {
                    println!("ERR: Requires 2 arguments. Example: memset 0x00 0x01 to store 0x01 \
                              in 0x00.");
                } else if parts.len() == 3 {
                    if let Ok(dst) = u16::from_str_radix(&parts[1].replace("0x", "")[..], 16) {
                        if let Ok(src) = u8::from_str_radix(&parts[2].replace("0x", "")[..], 16) {
                            self.cpu.memory[dst as usize] = src;
                        } else {
                            println!("ERR: Unable to parse source byte value");
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                } else {
                    if let Ok(mut dst) = usize::from_str_radix(&parts[1].replace("0x", "")[..],
                                                               16) {
                        for p in &parts[2..] {
                            if let Ok(byte) = u8::from_str_radix(&p.replace("0x", "")[..], 16) {
                                self.cpu.memory[dst] = byte;
                                dst += 0x01;
                            } else {
                                println!("ERR: Unable to parse source byte value");
                            }
                        }
                    } else {
                        println!("ERR: Unable to parse destination byte value");
                    }
                }
            }

            if input.starts_with("memdmp") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    let page_number: usize =
                        parts[1].parse().expect("Unable to parse memory page number");

                    self.dump_memory_page(page_number);
                }
            }

            if input.starts_with("break") {
                // 1 argument assumes 1 memory "page"
                let parts: Vec<&str> = input.split(" ").collect();
                if parts.len() == 2 {
                    if let Ok(addr) = usize::from_str_radix(&parts[1][..], 16) {
                        if self.breakpoints[addr] > 0 {
                            self.breakpoints[addr] = 0;
                        } else {
                            self.breakpoints[addr] = 1;
                        }
                    } else {
                        println!("Err: Unable to parse breakpoint address");
                    }
                } else {
                    println!("ERR: Requires 1 argument");
                }
            }

            if input == "continue" {
                self.broken = false;
                println!("Execution resumed");
            }

            if input == "step" {
                self.step = true;
                self.dump_local_disassembly();
            }

            std::io::stdout().write(b"hakka> ").unwrap();
            std::io::stdout().flush().unwrap();

            // Don't assign a blank command as a last command
            if input.len() > 0 {
                self.last_command = input.clone();
            }
        }

    }

    fn enable_memory_monitor(&mut self, input: &str) {
        let parts: Vec<&str> = input.split(" ").collect();
        if parts.len() < 3 {
            println!("ERR: Requires 2 arguments. Example: monitor 0x00 0xFF");
        } else {
            let start = usize::from_str_radix(&parts[1].replace("0x", "")[..], 16).unwrap();
            let end = usize::from_str_radix(&parts[2].replace("0x", "")[..], 16).unwrap();

            self.monitor.start_addr = start;
            self.monitor.end_addr = end;
            self.monitor.enabled = true;

            println!("{:?}", self.monitor);
        }
    }

    fn dump_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pairs = disassembler.disassemble_with_addresses(self.cpu.get_code());
        let result = self.highlight_lines(self.cpu.registers.PC as usize, pairs, false).join("");
        print!("{}", result);
    }

    fn dump_local_disassembly(&self) {
        println!("");
        println!("-- Disassembly --");

        let disassembler = Disassembler::with_offset(self.code_offset);
        let pc = self.cpu.registers.PC as usize;
        let code = self.cpu.get_code();
        let pairs = disassembler.disassemble_with_addresses(code);
        let result = self.highlight_lines(pc, pairs, true).join("");
        print!("{}", result);
    }

    pub fn dump_memory_page(&self, page: usize) {
        let mut addr = page * 0x100;
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[page * 0x100..(page * 0x100) + 0x100].chunks(8) {
            print!("{:04X}: ", addr);
            for b in chunk {
                print!("{:02X} ", *b);
            }
            addr += 0x08;
            println!("");
        }
    }

    pub fn dump_memory(&self) {
        println!("-- Memory dump --");
        for chunk in self.cpu.memory[self.monitor.start_addr..self.monitor.end_addr + 0x01]
            .chunks(8) {
            for b in chunk {
                print!("{:02X} ", *b);
            }
            println!("");
        }
    }

    fn highlight_lines(&self,
                       pc: usize,
                       pairs: Vec<(String, u16)>,
                       limit_results: bool)
                       -> Vec<String> {
        let mut result = Vec::new();

        let base = pc - self.code_offset as usize;

        for pair in pairs {
            if limit_results {
                let start = if base > 0x0A { base - 0x0A } else { 0 };
                if (pair.1 as usize) < start || (pair.1 as usize) > base + 0x0A {
                    continue;
                }
            }
            let current_line = pc as u16 == self.code_offset + pair.1;
            let breakpoint = self.breakpoints[self.code_offset as usize + pair.1 as usize] > 0x00;

            if breakpoint && current_line {
                result.push(format!("> * {}", pair.0));
            } else if breakpoint && !current_line {
                result.push(format!("  * {}", pair.0));
            } else if !breakpoint && current_line {
                result.push(format!(">   {}", pair.0));
            } else {
                result.push(format!("    {}", pair.0));
            }
        }

        result
    }
}<|endoftext|>


<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, ensure, Result};
use sha2::{Digest, Sha256};
use std::time::{Duration, Instant};
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::impl_serializable_error;
use crate::io::uart::Uart;
use crate::transport::Capability;

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 1024 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const MAGIC_HEADER: [u8; 4] = [0xfd, 0xff, 0xff, 0xff];

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::digest(&frame[Frame::HASH_LEN..]);
        sha.into()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let mut digest = Sha256::digest(self.as_bytes());
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest.into()
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Result<Vec<Frame>> {
        // The given payload will contain up to four sections concatenated together:
        // RO_A, RW_A optionally follwed by RO_B, RW_B
        // Each section starts with a magic number on at least a 256 byte boundary.

        // This rescue protocol uses the RW_A section only, which will start at the second
        // occurrance of the magic value, and end at the third occurrence or at the end of the
        // file.

        ensure!(
            payload.starts_with(&Self::MAGIC_HEADER),
            RescueError::ImageFormatError
        );

        // Find second occurrence of magic value.
        let min_addr = match payload[256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER)
        {
            Some(n) => (n + 1) * 256,
            None => bail!(RescueError::ImageFormatError),
        };

        // Find third occurrence of magic value.
        let max_addr = match payload[min_addr + 256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER)
        {
            Some(n) => (n + 1) * 256 + min_addr,
            None => payload.len(),
        };

        // Trim trailing 0xff bytes.
        let max_addr = (payload[..max_addr]
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frames = Vec::new();
        let mut frame_num = 0;
        let mut addr = min_addr;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: addr as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        Ok(frames)
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum RescueError {
    #[error("Unrecognized image file format")]
    ImageFormatError,
    #[error("Synchronization error communicating with boot rom")]
    SyncError,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}
impl_serializable_error!(RescueError);

/// Implements the UART rescue protocol of Google Ti50 firmware.
pub struct Rescue {}

impl Rescue {
    /// Abort if a block has not been accepted after this number of retries.
    const MAX_CONSECUTIVE_ERRORS: u32 = 50;
    /// Take some measure to regain protocol synchronization, in case of this number of retries
    /// of the same block.
    const RESYNC_AFTER_CONSECUTIVE_ERRORS: u32 = 3;

    /// Creates a new `Rescue` protocol updater from `options`.
    pub fn new(_options: &BootstrapOptions) -> Self {
        Self {}
    }

    /// Waits for some time for a character, returns None on timeout.
    fn read_char(&self, uart: &dyn Uart) -> Option<char> {
        let mut buf = [0u8; 1];
        match uart.read_timeout(&mut buf, Duration::from_millis(100)) {
            Ok(1) => Some(buf[0] as char),
            Ok(_) => None,
            _ => None,
        }
    }

    /// Waits some time for data, returning true if the given string was seen in full, or false
    /// as soon as a non-matching character is received or on timeout.
    fn expect_string(&self, uart: &dyn Uart, s: &str) -> bool {
        for expected_ch in s.chars() {
            match self.read_char(uart) {
                Some(ch) if ch == expected_ch => (),
                _ => return false,
            }
        }
        true
    }

    /// Reads and discards any characters in the receive buffer, waiting a little while for any
    /// more which will also be discarded.
    fn flush_rx(&self, uart: &dyn Uart) {
        let mut response = [0u8; Frame::HASH_LEN];
        loop {
            match uart.read_timeout(&mut response, Duration::from_millis(500)) {
                Ok(0) | Err(_) => break,
                Ok(_) => continue,
            }
        }
    }

    /// As the 1024 byte blocks sent to the chip have no discernible header, the sender and
    /// receiver could be "out of sync".  This is resolved by sending one byte at a time, and
    /// observing when the chip sends a response (which will be a rejection due to checksum).
    fn synchronize(&self, uart: &dyn Uart) -> Result<()> {
        // Most likely, only a few "extra" bytes have been sent during initial negotiation.
        // Send almost a complete block in one go, and then send each of the last 16 bytes one
        // at a time, slowly enough to detect a response before sending the next byte.
        uart.write(&[0u8; 1008])?;
        let mut response = [0u8; 1];
        let limit = match uart.read_timeout(&mut response, Duration::from_millis(50)) {
            Ok(0) | Err(_) => 16,
            Ok(_) => {
                // A response at this point must mean that more than 16 bytes had already been
                // sent before entering this method.  This will be resolved by doing another
                // slower round of 1024 bytes with delay in between every one.
                self.flush_rx(uart);
                1024
            }
        };
        for _ in 0..limit {
            uart.write(&[0u8; 1])?;
            match uart.read_timeout(&mut response, Duration::from_millis(50)) {
                Ok(0) | Err(_) => (),
                Ok(_) => {
                    self.flush_rx(uart);
                    return Ok(());
                }
            }
        }
        Err(RescueError::SyncError.into())
    }

    /// Reset the chip and send the magic 'r' character at the opportune moment during boot in
    /// order to enter rescue more, repeat if necessary.
    fn enter_rescue_mode(&self, container: &Bootstrap, uart: &dyn Uart) -> Result<()> {
        // Attempt getting the attention of the bootloader.
        let timeout = Duration::from_millis(2000);
        for _ in 0..Self::MAX_CONSECUTIVE_ERRORS {
            eprint!("Resetting...");
            container.reset_pin.write(false)?; // Low active
            std::thread::sleep(container.reset_delay);
            container.reset_pin.write(true)?; // Release reset
            let stopwatch = Instant::now();
            while stopwatch.elapsed() < timeout {
                if !self.expect_string(uart, "Bldr |") {
                    continue;
                }
                uart.write(&['r' as u8])?;
                eprint!("a.");
                while stopwatch.elapsed() < timeout {
                    if !self.expect_string(uart, "oops?|") {
                        continue;
                    }
                    uart.write(&['r' as u8])?;
                    eprint!("b.");
                    if self.expect_string(uart, "escue") {
                        eprintln!("c: Entered rescue mode!");
                        self.synchronize(uart)?;
                        return Ok(());
                    }
                }
            }
            eprintln!(" Failed to enter rescue mode.");
        }
        Err(RescueError::RepeatedErrors.into())
    }
}

impl UpdateProtocol for Rescue {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::UART)
            .ok()?;
        Ok(())
    }

    /// Returns false, in order to as the containing Bootstrap struct to not perform standard
    /// BOOTSTRAP/RESET sequence.
    fn uses_common_bootstrap_reset(&self) -> bool {
        false
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
    ) -> Result<()> {
        let frames = Frame::from_payload(payload)?;
        let uart = container.uart_params.create(transport)?;

        self.enter_rescue_mode(container, &*uart)?;

        // Send frames one at a time.
        'next_block: for (idx, frame) in frames.iter().enumerate() {
            for consecutive_errors in 0..Self::MAX_CONSECUTIVE_ERRORS {
                eprint!("{}.", idx);
                uart.write(frame.as_bytes())?;
                let mut response = [0u8; Frame::HASH_LEN];
                let mut index = 0;
                while index < Frame::HASH_LEN {
                    let timeout = if index == 0 {
                        Duration::from_millis(1000)
                    } else {
                        Duration::from_millis(10)
                    };
                    match uart.read_timeout(&mut response[index..], timeout) {
                        Ok(0) | Err(_) => break,
                        Ok(n) => index += n,
                    }
                }
                if index < Frame::HASH_LEN {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                    continue;
                }
                if response[4..].chunks(4).all(|x| x == &response[..4]) {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                } else if response == frame.frame_hash() {
                    continue 'next_block;
                } else {
                    self.flush_rx(&*uart);
                    if consecutive_errors >= Self::RESYNC_AFTER_CONSECUTIVE_ERRORS {
                        eprint!("sync.");
                        self.synchronize(&*uart)?;
                    }
                }
            }
            bail!(RescueError::RepeatedErrors);
        }

        // Reset, in order to leave rescue mode.
        container.reset_pin.write(false)?; // Low active
        std::thread::sleep(container.reset_delay);
        container.reset_pin.write(true)?; // Release reset
        eprintln!("Success!");
        Ok(())
    }
}
<commit_msg>[opentitantool] Tweak to Ti50 rescue bootstrapping<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, ensure, Result};
use sha2::{Digest, Sha256};
use std::time::{Duration, Instant};
use thiserror::Error;
use zerocopy::AsBytes;

use crate::app::TransportWrapper;
use crate::bootstrap::{Bootstrap, BootstrapOptions, UpdateProtocol};
use crate::impl_serializable_error;
use crate::io::uart::Uart;
use crate::transport::Capability;

#[derive(AsBytes, Debug, Default)]
#[repr(C)]
struct FrameHeader {
    hash: [u8; Frame::HASH_LEN],
    frame_num: u32,
    flash_offset: u32,
}

#[derive(AsBytes, Debug)]
#[repr(C)]
struct Frame {
    header: FrameHeader,
    data: [u8; Frame::DATA_LEN],
}

impl Default for Frame {
    fn default() -> Self {
        Frame {
            header: Default::default(),
            data: [0xff; Frame::DATA_LEN],
        }
    }
}

impl Frame {
    const EOF: u32 = 0x8000_0000;
    const FLASH_SECTOR_SIZE: usize = 2048;
    const FLASH_SECTOR_MASK: usize = Self::FLASH_SECTOR_SIZE - 1;
    const FLASH_BUFFER_SIZE: usize = 128;
    const FLASH_BUFFER_MASK: usize = Self::FLASH_BUFFER_SIZE - 1;
    const DATA_LEN: usize = 1024 - std::mem::size_of::<FrameHeader>();
    const HASH_LEN: usize = 32;
    const MAGIC_HEADER: [u8; 4] = [0xfd, 0xff, 0xff, 0xff];
    const CRYPTOLIB_TELL: [u8; 4] = [0x53, 0x53, 0x53, 0x53];

    /// Computes the hash in the header.
    fn header_hash(&self) -> [u8; Frame::HASH_LEN] {
        let frame = self.as_bytes();
        let sha = Sha256::digest(&frame[Frame::HASH_LEN..]);
        sha.into()
    }

    /// Computes the hash over the entire frame.
    fn frame_hash(&self) -> [u8; Frame::HASH_LEN] {
        let mut digest = Sha256::digest(self.as_bytes());
        // Touch up zeroes into ones, as that is what the old chips are doing.
        for b in &mut digest {
            if *b == 0 {
                *b = 1;
            }
        }
        digest.into()
    }

    /// Creates a sequence of frames based on a `payload` binary.
    fn from_payload(payload: &[u8]) -> Result<Vec<Frame>> {
        // The given payload will contain up to four sections concatenated together:
        // RO_A, RW_A optionally follwed by RO_B, RW_B
        // Each section starts with a magic number on at least a 256 byte boundary.

        // This rescue protocol uses the RW_A section only, which will start at the second
        // occurrance of the magic value, and end at the third occurrence or at the end of the
        // file.

        ensure!(
            payload.starts_with(&Self::MAGIC_HEADER),
            RescueError::ImageFormatError
        );

        // Find second occurrence of magic value, not followed by signature of encrypted
        // cryptolib.
        let min_addr = match payload[256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER && &c[4..8] != &Self::CRYPTOLIB_TELL)
        {
            Some(n) => (n + 1) * 256,
            None => bail!(RescueError::ImageFormatError),
        };

        // Find third occurrence of magic value.
        let max_addr = match payload[min_addr + 256..]
            .chunks(256)
            .position(|c| &c[0..4] == &Self::MAGIC_HEADER)
        {
            Some(n) => (n + 1) * 256 + min_addr,
            None => payload.len(),
        };

        // Trim trailing 0xff bytes.
        let max_addr = (payload[..max_addr]
            .chunks(4)
            .rposition(|c| c != &[0xff; 4])
            .unwrap_or(0)
            + 1)
            * 4;

        let mut frames = Vec::new();
        let mut frame_num = 0;
        let mut addr = min_addr;
        while addr < max_addr {
            // Try skipping over 0xffffffff words.
            let nonempty_addr = addr
                + payload[addr..]
                    .chunks(4)
                    .position(|c| c != &[0xff; 4])
                    .unwrap()
                    * 4;
            let skip_addr = nonempty_addr & !Self::FLASH_SECTOR_MASK;
            if skip_addr > addr && (addr == 0 || addr & Self::FLASH_BUFFER_MASK != 0) {
                // Can only skip from the start or if the last addr wasn't an exact multiple of
                // 128 (per H1D boot rom).
                addr = skip_addr;
            }

            let mut frame = Frame {
                header: FrameHeader {
                    frame_num,
                    flash_offset: addr as u32,
                    ..Default::default()
                },
                ..Default::default()
            };
            let slice_size = Self::DATA_LEN.min(payload.len() - addr);
            frame.data[..slice_size].copy_from_slice(&payload[addr..addr + slice_size]);
            frames.push(frame);

            addr += Self::DATA_LEN;
            frame_num += 1;
        }
        frames.last_mut().map(|f| f.header.frame_num |= Self::EOF);
        frames
            .iter_mut()
            .for_each(|f| f.header.hash = f.header_hash());
        Ok(frames)
    }
}

#[derive(Debug, Error, serde::Serialize, serde::Deserialize)]
pub enum RescueError {
    #[error("Unrecognized image file format")]
    ImageFormatError,
    #[error("Synchronization error communicating with boot rom")]
    SyncError,
    #[error("Repeated errors communicating with boot rom")]
    RepeatedErrors,
}
impl_serializable_error!(RescueError);

/// Implements the UART rescue protocol of Google Ti50 firmware.
pub struct Rescue {}

impl Rescue {
    /// Abort if a block has not been accepted after this number of retries.
    const MAX_CONSECUTIVE_ERRORS: u32 = 50;
    /// Take some measure to regain protocol synchronization, in case of this number of retries
    /// of the same block.
    const RESYNC_AFTER_CONSECUTIVE_ERRORS: u32 = 3;

    /// Creates a new `Rescue` protocol updater from `options`.
    pub fn new(_options: &BootstrapOptions) -> Self {
        Self {}
    }

    /// Waits for some time for a character, returns None on timeout.
    fn read_char(&self, uart: &dyn Uart) -> Option<char> {
        let mut buf = [0u8; 1];
        match uart.read_timeout(&mut buf, Duration::from_millis(100)) {
            Ok(1) => Some(buf[0] as char),
            Ok(_) => None,
            _ => None,
        }
    }

    /// Waits some time for data, returning true if the given string was seen in full, or false
    /// as soon as a non-matching character is received or on timeout.
    fn expect_string(&self, uart: &dyn Uart, s: &str) -> bool {
        for expected_ch in s.chars() {
            match self.read_char(uart) {
                Some(ch) if ch == expected_ch => (),
                _ => return false,
            }
        }
        true
    }

    /// Reads and discards any characters in the receive buffer, waiting a little while for any
    /// more which will also be discarded.
    fn flush_rx(&self, uart: &dyn Uart) {
        let mut response = [0u8; Frame::HASH_LEN];
        loop {
            match uart.read_timeout(&mut response, Duration::from_millis(500)) {
                Ok(0) | Err(_) => break,
                Ok(_) => continue,
            }
        }
    }

    /// As the 1024 byte blocks sent to the chip have no discernible header, the sender and
    /// receiver could be "out of sync".  This is resolved by sending one byte at a time, and
    /// observing when the chip sends a response (which will be a rejection due to checksum).
    fn synchronize(&self, uart: &dyn Uart) -> Result<()> {
        // Most likely, only a few "extra" bytes have been sent during initial negotiation.
        // Send almost a complete block in one go, and then send each of the last 16 bytes one
        // at a time, slowly enough to detect a response before sending the next byte.
        uart.write(&[0u8; 1008])?;
        let mut response = [0u8; 1];
        let limit = match uart.read_timeout(&mut response, Duration::from_millis(50)) {
            Ok(0) | Err(_) => 16,
            Ok(_) => {
                // A response at this point must mean that more than 16 bytes had already been
                // sent before entering this method.  This will be resolved by doing another
                // slower round of 1024 bytes with delay in between every one.
                self.flush_rx(uart);
                1024
            }
        };
        for _ in 0..limit {
            uart.write(&[0u8; 1])?;
            match uart.read_timeout(&mut response, Duration::from_millis(50)) {
                Ok(0) | Err(_) => (),
                Ok(_) => {
                    self.flush_rx(uart);
                    return Ok(());
                }
            }
        }
        Err(RescueError::SyncError.into())
    }

    /// Reset the chip and send the magic 'r' character at the opportune moment during boot in
    /// order to enter rescue more, repeat if necessary.
    fn enter_rescue_mode(&self, container: &Bootstrap, uart: &dyn Uart) -> Result<()> {
        // Attempt getting the attention of the bootloader.
        let timeout = Duration::from_millis(2000);
        for _ in 0..Self::MAX_CONSECUTIVE_ERRORS {
            eprint!("Resetting...");
            container.reset_pin.write(false)?; // Low active
            std::thread::sleep(container.reset_delay);
            container.reset_pin.write(true)?; // Release reset
            let stopwatch = Instant::now();
            while stopwatch.elapsed() < timeout {
                if !self.expect_string(uart, "Bldr |") {
                    continue;
                }
                uart.write(&['r' as u8])?;
                eprint!("a.");
                while stopwatch.elapsed() < timeout {
                    if !self.expect_string(uart, "oops?|") {
                        continue;
                    }
                    uart.write(&['r' as u8])?;
                    eprint!("b.");
                    if self.expect_string(uart, "escue") {
                        eprintln!("c: Entered rescue mode!");
                        self.synchronize(uart)?;
                        return Ok(());
                    }
                }
            }
            eprintln!(" Failed to enter rescue mode.");
        }
        Err(RescueError::RepeatedErrors.into())
    }
}

impl UpdateProtocol for Rescue {
    fn verify_capabilities(
        &self,
        _container: &Bootstrap,
        transport: &TransportWrapper,
    ) -> Result<()> {
        transport
            .capabilities()?
            .request(Capability::GPIO | Capability::UART)
            .ok()?;
        Ok(())
    }

    /// Returns false, in order to as the containing Bootstrap struct to not perform standard
    /// BOOTSTRAP/RESET sequence.
    fn uses_common_bootstrap_reset(&self) -> bool {
        false
    }

    /// Performs the update protocol using the `transport` with the firmware `payload`.
    fn update(
        &self,
        container: &Bootstrap,
        transport: &TransportWrapper,
        payload: &[u8],
    ) -> Result<()> {
        let frames = Frame::from_payload(payload)?;
        let uart = container.uart_params.create(transport)?;

        self.enter_rescue_mode(container, &*uart)?;

        // Send frames one at a time.
        'next_block: for (idx, frame) in frames.iter().enumerate() {
            for consecutive_errors in 0..Self::MAX_CONSECUTIVE_ERRORS {
                eprint!("{}.", idx);
                uart.write(frame.as_bytes())?;
                let mut response = [0u8; Frame::HASH_LEN];
                let mut index = 0;
                while index < Frame::HASH_LEN {
                    let timeout = if index == 0 {
                        Duration::from_millis(1000)
                    } else {
                        Duration::from_millis(10)
                    };
                    match uart.read_timeout(&mut response[index..], timeout) {
                        Ok(0) | Err(_) => break,
                        Ok(n) => index += n,
                    }
                }
                if index < Frame::HASH_LEN {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                    continue;
                }
                if response[4..].chunks(4).all(|x| x == &response[..4]) {
                    eprint!("sync.");
                    self.synchronize(&*uart)?;
                } else if response == frame.frame_hash() {
                    continue 'next_block;
                } else {
                    self.flush_rx(&*uart);
                    if consecutive_errors >= Self::RESYNC_AFTER_CONSECUTIVE_ERRORS {
                        eprint!("sync.");
                        self.synchronize(&*uart)?;
                    }
                }
            }
            bail!(RescueError::RepeatedErrors);
        }

        // Reset, in order to leave rescue mode.
        container.reset_pin.write(false)?; // Low active
        std::thread::sleep(container.reset_delay);
        container.reset_pin.write(true)?; // Release reset
        eprintln!("Success!");
        Ok(())
    }
}
<|endoftext|>


<commit_before>pub trait Summarizable {
    fn author_summary(&self) -> String;

    fn summary(&self) -> String {
        String::from("Read more from {}...", self.author_summary())
    }
}

pub struct NewsArticle {
    pub headline: String,
    pub location: String,
    pub author: String,
    pub content: String,
}

impl Summarizable for NewsArticle {
    fn summary(&self) -> String {
        format!("{}, by {} ({})", self.headline, self.author, self.location)
    }
}

pub struct Tweet {
    pub username: String,
    pub content: String,
    pub reply: bool,
    pub retweet: bool,
}

impl Summarizable for Tweet {
    fn summary(&self) -> String {
        format!("{}: {}", self.username, self.content)
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<commit_msg>Fix summary default and implement author summary on structs<commit_after>pub trait Summarizable {
    fn author_summary(&self) -> String;

    fn summary(&self) -> String {
        format!("Read more from {}...", self.author_summary())
    }
}

pub struct NewsArticle {
    pub headline: String,
    pub location: String,
    pub author: String,
    pub content: String,
}

impl Summarizable for NewsArticle {
    fn author_summary(&self) -> String {
        format!("{}", self.author)
    }

    fn summary(&self) -> String {
        format!("{}, by {} ({})", self.headline, self.author, self.location)
    }
}

pub struct Tweet {
    pub username: String,
    pub content: String,
    pub reply: bool,
    pub retweet: bool,
}

impl Summarizable for Tweet {
    fn author_summary(&self) -> String {
        format!("{}", self.username)
    }

    fn summary(&self) -> String {
        format!("{}: {}", self.username, self.content)
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Issue #7526: lowercase static constants in patterns look like bindings

// This is similar to compile-fail/match-static-const-lc, except it
// shows the expected usual workaround (choosing a different name for
// the static definition) and also demonstrates that one can work
// around this problem locally by renaming the constant in the `use`
// form to an uppercase identifier that placates the lint.

#[deny(non_uppercase_pattern_statics)];

pub static A : int = 97;

fn f() {
    let r = match (0,0) {
        (0, A) => 0,
        (x, y) => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,97) {
        (0, A) => 0,
        (x, y) => 1 + x + y,
    };
    assert!(r == 0);
}

mod m {
    pub static aha : int = 7;
}

fn g() {
    use AHA = self::m::aha;
    let r = match (0,0) {
        (0, AHA) => 0,
        (x, y)   => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,7) {
        (0, AHA) => 0,
        (x, y)   => 1 + x + y,
    };
    assert!(r == 0);
}

fn h() {
    let r = match (0,0) {
        (0, self::m::aha) => 0,
        (x, y)      => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,7) {
        (0, self::m::aha) => 0,
        (x, y)      => 1 + x + y,
    };
    assert!(r == 0);
}

fn main () {
    f();
    g();
    h();
}
<commit_msg>fix tests for check-fast.<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Issue #7526: lowercase static constants in patterns look like bindings

// This is similar to compile-fail/match-static-const-lc, except it
// shows the expected usual workaround (choosing a different name for
// the static definition) and also demonstrates that one can work
// around this problem locally by renaming the constant in the `use`
// form to an uppercase identifier that placates the lint.

#[deny(non_uppercase_pattern_statics)];

pub static A : int = 97;

fn f() {
    let r = match (0,0) {
        (0, A) => 0,
        (x, y) => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,97) {
        (0, A) => 0,
        (x, y) => 1 + x + y,
    };
    assert!(r == 0);
}

mod m {
    pub static aha : int = 7;
}

fn g() {
    use AHA = self::m::aha;
    let r = match (0,0) {
        (0, AHA) => 0,
        (x, y)   => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,7) {
        (0, AHA) => 0,
        (x, y)   => 1 + x + y,
    };
    assert!(r == 0);
}

fn h() {
    let r = match (0,0) {
        (0, self::m::aha) => 0,
        (x, y)      => 1 + x + y,
    };
    assert!(r == 1);
    let r = match (0,7) {
        (0, self::m::aha) => 0,
        (x, y)      => 1 + x + y,
    };
    assert!(r == 0);
}

pub fn main () {
    f();
    g();
    h();
}
<|endoftext|>


<commit_before>use std::cast;
use sync::{Arc, Mutex};
use conn::{MyInnerConn, MyOpts, MyResult, Stmt, QueryResult};

struct MyInnerPool {
    opts: MyOpts,
    pool: Vec<*()>,
    cap: uint,
    count: uint
}

impl Drop for MyInnerPool {
    fn drop(&mut self) {
        loop {
            match self.pool.pop() {
                Some(conn) => unsafe {
                    drop(cast::transmute::<*(), Box<MyInnerConn>>(conn));
                },
                None => break
            }
        }
    }
}

impl MyInnerPool {
    fn new_conn(&mut self) -> MyResult<()> {
        match MyInnerConn::new(self.opts.clone()) {
            Ok(conn) => {
                unsafe { self.pool.push(cast::transmute(box conn)) };
                Ok(())
            },
            Err(err) => Err(err)
        }
    }
}

#[deriving(Clone)]
pub struct MyPool {
    pool: Arc<Mutex<MyInnerPool>>
}

impl MyPool {
    pub fn new(cap: uint, opts: MyOpts) -> MyPool {
        let pool = MyInnerPool {
            opts: opts,
            pool: Vec::new(),
            cap: cap,
            count: 0
        };
        MyPool{ pool: Arc::new(Mutex::new(pool)) }
    }

    pub fn get_conn(&self) -> MyResult<MyPooledConn> {
        let mut pool = self.pool.lock();

        while pool.pool.is_empty() {
            if pool.cap == 0 || pool.count < pool.cap {
                match pool.new_conn() {
                    Ok(()) => {
                        pool.count += 1;
                        break;
                    },
                    Err(err) => return Err(err)
                }
            } else {
                pool.cond.wait();
            }
        }

        Ok(MyPooledConn {
            pool: self.clone(),
            conn: unsafe { cast::transmute(pool.pool.pop().unwrap()) }
        })
    }
}

pub struct MyPooledConn {
    pool: MyPool,
    conn: Option<Box<MyInnerConn>>
}

#[unsafe_destructor]
impl Drop for MyPooledConn {
    fn drop(&mut self) {
        let conn = unsafe { cast::transmute(self.conn.take_unwrap()) };
        let mut pool = self.pool.pool.lock();
        pool.pool.push(conn);
        pool.cond.signal();
    }
}

impl MyPooledConn {
    fn query<'a>(&'a mut self, query: &str) -> MyResult<QueryResult<'a>> {
        self.conn.get_mut_ref().query(query)
    }
    fn prepare<'a>(&'a mut self, query: &str) -> MyResult<Stmt<'a>> {
        self.conn.get_mut_ref().prepare(query)
    }
}

#[cfg(test)]
mod test {
    use conn::{MyOpts};
    use std::default::{Default};
    use super::{MyPool};

    #[test]
    fn test_query() {
        let pool = MyPool::new(3, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                assert!(conn.query("SELECT 1").is_ok());
            });
        }
    }

    #[test]
    fn test_prepared_query() {
        let pool = MyPool::new(3, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                let stmt = conn.prepare("SELECT 1");
                assert!(stmt.is_ok());
                let mut stmt = stmt.unwrap();
                assert!(stmt.execute([]).is_ok());
            });
        }
    }

    #[test]
    fn test_zero_cap() {
        let pool = MyPool::new(0, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                assert!(conn.query("SELECT 1").is_ok());
            });
        }
    }
}
<commit_msg>Make `MyPooledConn::query` and `MyPooledConn::prepare` public<commit_after>use std::cast;
use sync::{Arc, Mutex};
use conn::{MyInnerConn, MyOpts, MyResult, Stmt, QueryResult};

struct MyInnerPool {
    opts: MyOpts,
    pool: Vec<*()>,
    cap: uint,
    count: uint
}

impl Drop for MyInnerPool {
    fn drop(&mut self) {
        loop {
            match self.pool.pop() {
                Some(conn) => unsafe {
                    drop(cast::transmute::<*(), Box<MyInnerConn>>(conn));
                },
                None => break
            }
        }
    }
}

impl MyInnerPool {
    fn new_conn(&mut self) -> MyResult<()> {
        match MyInnerConn::new(self.opts.clone()) {
            Ok(conn) => {
                unsafe { self.pool.push(cast::transmute(box conn)) };
                Ok(())
            },
            Err(err) => Err(err)
        }
    }
}

#[deriving(Clone)]
pub struct MyPool {
    pool: Arc<Mutex<MyInnerPool>>
}

impl MyPool {
    pub fn new(cap: uint, opts: MyOpts) -> MyPool {
        let pool = MyInnerPool {
            opts: opts,
            pool: Vec::new(),
            cap: cap,
            count: 0
        };
        MyPool{ pool: Arc::new(Mutex::new(pool)) }
    }

    pub fn get_conn(&self) -> MyResult<MyPooledConn> {
        let mut pool = self.pool.lock();

        while pool.pool.is_empty() {
            if pool.cap == 0 || pool.count < pool.cap {
                match pool.new_conn() {
                    Ok(()) => {
                        pool.count += 1;
                        break;
                    },
                    Err(err) => return Err(err)
                }
            } else {
                pool.cond.wait();
            }
        }

        Ok(MyPooledConn {
            pool: self.clone(),
            conn: unsafe { cast::transmute(pool.pool.pop().unwrap()) }
        })
    }
}

pub struct MyPooledConn {
    pool: MyPool,
    conn: Option<Box<MyInnerConn>>
}

#[unsafe_destructor]
impl Drop for MyPooledConn {
    fn drop(&mut self) {
        let conn = unsafe { cast::transmute(self.conn.take_unwrap()) };
        let mut pool = self.pool.pool.lock();
        pool.pool.push(conn);
        pool.cond.signal();
    }
}

impl MyPooledConn {
    pub fn query<'a>(&'a mut self, query: &str) -> MyResult<QueryResult<'a>> {
        self.conn.get_mut_ref().query(query)
    }
    pub fn prepare<'a>(&'a mut self, query: &str) -> MyResult<Stmt<'a>> {
        self.conn.get_mut_ref().prepare(query)
    }
}

#[cfg(test)]
mod test {
    use conn::{MyOpts};
    use std::default::{Default};
    use super::{MyPool};

    #[test]
    fn test_query() {
        let pool = MyPool::new(3, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                assert!(conn.query("SELECT 1").is_ok());
            });
        }
    }

    #[test]
    fn test_prepared_query() {
        let pool = MyPool::new(3, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                let stmt = conn.prepare("SELECT 1");
                assert!(stmt.is_ok());
                let mut stmt = stmt.unwrap();
                assert!(stmt.execute([]).is_ok());
            });
        }
    }

    #[test]
    fn test_zero_cap() {
        let pool = MyPool::new(0, MyOpts{user: Some("root".to_owned()),
                                         ..Default::default()});
        for _ in range(0, 10) {
            let pool = pool.clone();
            spawn(proc() {
                let conn = pool.get_conn();
                assert!(conn.is_ok());
                let mut conn = conn.unwrap();
                assert!(conn.query("SELECT 1").is_ok());
            });
        }
    }
}
<|endoftext|>


<commit_before>extern crate syntax;
extern crate rustc;

use rustc::session;
use rustc::session::config;

use syntax::diagnostics::registry;
use syntax::parse::lexer;
use syntax::parse::token;
use syntax::parse::token::{Lit, Token};
use syntax::parse;

use std::io;
use std::iter;

/* This is a simple standalone lexer based on rustc's lexer. The main
 * difference is a custom to_str function for tokens that prints in
 * the same format as the flex lexer in this project, so the results
 * between the two are easily comparable.
 */

fn token_to_string(tok: token::Token) -> String {
    match tok {
        Token::Literal(l, _) => match l {
            Lit::Byte(c) => {
                format!("Byte(b'{}')", c.as_str())
            },
            Lit::Char(c) => {
                format!("Char('{}')", c.as_str())
            },
            Lit::Integer(c) => {
                format!("Integer({})", c.as_str().to_string())
            },
            Lit::Float(c) => {
                format!("Float({})", c.as_str().to_string())
            },
            Lit::Str_(s) => {
                format!("Str(\"{}\")", token::get_name(s).get().escape_default())
            },
            Lit::StrRaw(s, n) => {
                format!("StrRaw(r{delim}\"{string}\"{delim})",
                        delim=iter::repeat("#").take(n).collect::<String>(),
                        string=token::get_name(s))
            },
            Lit::Binary(s) => {
                format!("Binary(b\"{}\")", s.as_str())
            },
            Lit::BinaryRaw(s, n) => {
                format!("BinaryRaw(br{delim}\"{string}\"{delim})",
                        delim=iter::repeat("#").take(n).collect::<String>(),
                        string=token::get_name(s))
            },
        },
        Token::Ident(s, _) => {
            format!("Ident({})", token::get_ident(s).get().to_string())
        },
        Token::Lifetime(s) => {
            format!("Lifetime({})", token::get_ident(s).get().to_string())
        },
        Token::DocComment(s) => {
            format!("DocComment({})", token::get_name(s).get().to_string())
        },
        Token::Whitespace | Token::Comment => "".to_string(),
        Token::Shebang(s) => {
            format!("Shebang({})", token::get_name(s).get().to_string())
        },
        t => {
            format!("{}", t)
        },
    }
}

fn main() {
    let line = io::stdin().read_to_string().unwrap();

    let options = config::basic_options();
    let session = session::build_session(options, None,
                                         registry::Registry::new(&rustc::DIAGNOSTICS));
    let filemap = parse::string_to_filemap(&session.parse_sess,
                                           line,
                                           String::from_str("<n/a>"));
    let mut lexer = lexer::StringReader::new(session.diagnostic(), filemap);

    {
        use syntax::parse::lexer::Reader;
        while !lexer.is_eof() {
            let s = token_to_string(lexer.next_token().tok);
            if !s.is_empty() {
                println!("{}", s);
            }
        }
    }
}
<commit_msg>Fix rlex for language/library changes<commit_after>extern crate syntax;
extern crate rustc;

use rustc::session;
use rustc::session::config;

use syntax::diagnostics::registry;
use syntax::parse::lexer;
use syntax::parse::token;
use syntax::parse::token::{Lit, Token};
use syntax::parse;

use std::io;
use std::iter;

/* This is a simple standalone lexer based on rustc's lexer. The main
 * difference is a custom to_str function for tokens that prints in
 * the same format as the flex lexer in this project, so the results
 * between the two are easily comparable.
 */

fn token_to_string(tok: token::Token) -> String {
    match tok {
        Token::Literal(l, _) => match l {
            Lit::Byte(c) => {
                format!("Byte(b'{}')", c.as_str())
            },
            Lit::Char(c) => {
                format!("Char('{}')", c.as_str())
            },
            Lit::Integer(c) => {
                format!("Integer({})", c.as_str().to_string())
            },
            Lit::Float(c) => {
                format!("Float({})", c.as_str().to_string())
            },
            Lit::Str_(s) => {
                format!("Str(\"{}\")", token::get_name(s).get().escape_default())
            },
            Lit::StrRaw(s, n) => {
                format!("StrRaw(r{delim}\"{string}\"{delim})",
                        delim=iter::repeat("#").take(n).collect::<String>(),
                        string=token::get_name(s))
            },
            Lit::Binary(s) => {
                format!("Binary(b\"{}\")", s.as_str())
            },
            Lit::BinaryRaw(s, n) => {
                format!("BinaryRaw(br{delim}\"{string}\"{delim})",
                        delim=iter::repeat("#").take(n).collect::<String>(),
                        string=token::get_name(s))
            },
        },
        Token::Ident(s, _) => {
            format!("Ident({})", token::get_ident(s).get().to_string())
        },
        Token::Lifetime(s) => {
            format!("Lifetime({})", token::get_ident(s).get().to_string())
        },
        Token::DocComment(s) => {
            format!("DocComment({})", token::get_name(s).get().to_string())
        },
        Token::Whitespace | Token::Comment => "".to_string(),
        Token::Shebang(s) => {
            format!("Shebang({})", token::get_name(s).get().to_string())
        },
        t => {
            format!("{:?}", t)
        },
    }
}

fn main() {
    let line = io::stdin().read_to_string().unwrap();

    let options = config::basic_options();
    let session = session::build_session(options, None,
                                         registry::Registry::new(&rustc::DIAGNOSTICS));
    let filemap = parse::string_to_filemap(&session.parse_sess,
                                           line,
                                           String::from_str("<n/a>"));
    let mut lexer = lexer::StringReader::new(session.diagnostic(), filemap);

    {
        use syntax::parse::lexer::Reader;
        while !lexer.is_eof() {
            let s = token_to_string(lexer.next_token().tok);
            if !s.is_empty() {
                println!("{}", s);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Data types for addresses to nodes in DICOM networks.
//!
//! This module provides the definitions for [`FullAeAddr`] and [`AeAddr`],
//! which enable consumers to couple a socket address with an expected
//! application entity (AE) title.
//!
//! The syntax is `«ae_title»@«network_address»:«port»`,
//! which works not only with IPv4 and IPv6 addresses,
//! but also with domain names.
use std::{
    convert::TryFrom,
    net::{SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
    str::FromStr,
};

use snafu::{ensure, AsErrorSource, ResultExt, Snafu};

/// A specification for a full address to the target SCP:
/// an application entity title, plus a generic  address,
/// typically a socket address.
///
/// These addresses can be serialized and parsed
/// with the syntax `{ae_title}@{address}`,
/// where the socket address is parsed according to
/// the expectations of the parameter type `T`.
///
/// For the version of the struct without a mandatory AE title,
/// see [`AeAddr`].
///
/// # Example
///
/// ```
/// # use dicom_ul::FullAeAddr;
/// # use std::net::SocketAddr;
/// #
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # // socket address can be a string
/// let addr: FullAeAddr<String> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), "SCP-STORAGE");
/// assert_eq!(addr.socket_addr(), "127.0.0.1:104");
/// # // or anything else which can be parsed into a socket address
/// let addr: FullAeAddr<SocketAddr> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), "SCP-STORAGE");
/// assert_eq!(addr.socket_addr(), &SocketAddr::from(([127, 0, 0, 1], 104)));
/// assert_eq!(&addr.to_string(), "SCP-STORAGE@127.0.0.1:104");
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct FullAeAddr<T> {
    ae_title: String,
    socket_addr: T,
}

impl<T> FullAeAddr<T> {
    /// Create an AE address from its bare constituent parts.
    pub fn new(ae_title: impl Into<String>, socket_addr: T) -> Self {
        FullAeAddr {
            ae_title: ae_title.into(),
            socket_addr,
        }
    }

    /// Retrieve the application entity title portion.
    pub fn ae_title(&self) -> &str {
        &self.ae_title
    }

    /// Retrieve the network address portion.
    pub fn socket_addr(&self) -> &T {
        &self.socket_addr
    }

    /// Convert the full address into its constituent parts.
    pub fn into_parts(self) -> (String, T) {
        (self.ae_title, self.socket_addr)
    }
}

impl<T> From<(String, T)> for FullAeAddr<T> {
    fn from((ae_title, socket_addr): (String, T)) -> Self {
        Self::new(ae_title, socket_addr)
    }
}

/// A error which occurred when parsing an AE address.
#[derive(Debug, Clone, Eq, PartialEq, Snafu)]
pub enum ParseAeAddressError<E>
where
    E: std::fmt::Debug + AsErrorSource,
{
    /// Missing `@` in full AE address
    MissingPart,

    /// Could not parse network socket address
    ParseSocketAddress { source: E },
}

impl<T> FromStr for FullAeAddr<T>
where
    T: FromStr,
    T::Err: std::fmt::Debug + AsErrorSource,
{
    type Err = ParseAeAddressError<<T as FromStr>::Err>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // !!! there should be a way to escape the `@`
        if let Some((ae_title, addr)) = s.split_once('@') {
            ensure!(!ae_title.is_empty(), MissingPartSnafu);
            Ok(FullAeAddr {
                ae_title: ae_title.to_string(),
                socket_addr: addr.parse().context(ParseSocketAddressSnafu)?,
            })
        } else {
            Err(ParseAeAddressError::MissingPart)
        }
    }
}

impl<T> ToSocketAddrs for FullAeAddr<T>
where
    T: ToSocketAddrs,
{
    type Iter = T::Iter;

    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
        self.socket_addr.to_socket_addrs()
    }
}

impl<T> std::fmt::Display for FullAeAddr<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.ae_title.replace("@", "\\@"))?;
        f.write_str("@")?;
        std::fmt::Display::fmt(&self.socket_addr, f)
    }
}

/// A specification for an address to the target SCP:
/// a generic network socket address
/// which may also include an application entity title.
///
/// These addresses can be serialized and parsed
/// with the syntax `{ae_title}@{address}`,
/// where the socket address is parsed according to
/// the expectations of the parameter type `T`.
///
/// For the version of the struct in which the AE title part is mandatory,
/// see [`FullAeAddr`].
///
/// # Example
///
/// ```
/// # use dicom_ul::{AeAddr, FullAeAddr};
/// # use std::net::{SocketAddr, SocketAddrV4};
/// #
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let addr: AeAddr<SocketAddrV4> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), Some("SCP-STORAGE"));
/// assert_eq!(addr.socket_addr(), &SocketAddrV4::new([127, 0, 0, 1].into(), 104));
/// assert_eq!(&addr.to_string(), "SCP-STORAGE@127.0.0.1:104");
///
/// // AE title can be missing
/// let addr: AeAddr<String> = "192.168.1.99:1045".parse()?;
/// assert_eq!(addr.ae_title(), None);
/// // but can be provided later
/// let full_addr: FullAeAddr<_> = addr.with_ae_title("SCP-QUERY");
/// assert_eq!(full_addr.ae_title(), "SCP-QUERY");
/// assert_eq!(&full_addr.to_string(), "SCP-QUERY@192.168.1.99:1045");
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct AeAddr<T> {
    ae_title: Option<String>,
    socket_addr: T,
}

impl<T> AeAddr<T> {
    /// Create an AE address from its bare constituent parts.
    pub fn new(ae_title: impl Into<String>, socket_addr: T) -> Self {
        AeAddr {
            ae_title: Some(ae_title.into()),
            socket_addr,
        }
    }

    /// Create an address with a missing AE title.
    pub fn new_socket_addr(socket_addr: T) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }

    /// Retrieve the application entity title portion, if present.
    pub fn ae_title(&self) -> Option<&str> {
        self.ae_title.as_deref()
    }

    /// Retrieve the socket address portion.
    pub fn socket_addr(&self) -> &T {
        &self.socket_addr
    }

    /// Create a new address with the full application entity target,
    /// discarding any potentially existing AE title.
    pub fn with_ae_title(self, ae_title: impl Into<String>) -> FullAeAddr<T> {
        FullAeAddr {
            ae_title: ae_title.into(),
            socket_addr: self.socket_addr,
        }
    }

    /// Create a new address with the full application entity target,
    /// using the given AE title if it is missing.
    pub fn with_default_ae_title(self, ae_title: impl Into<String>) -> FullAeAddr<T> {
        FullAeAddr {
            ae_title: self.ae_title.unwrap_or_else(|| ae_title.into()),
            socket_addr: self.socket_addr,
        }
    }

    /// Convert the address into its constituent parts.
    pub fn into_parts(self) -> (Option<String>, T) {
        (self.ae_title, self.socket_addr)
    }
}

/// This conversion provides a socket address without an AE title.
impl From<SocketAddr> for AeAddr<SocketAddr> {
    fn from(socket_addr: SocketAddr) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

/// This conversion provides an IPv4 socket address without an AE title.
impl From<SocketAddrV4> for AeAddr<SocketAddrV4> {
    fn from(socket_addr: SocketAddrV4) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

/// This conversion provides an IPv6 socket address without an AE title.
impl From<SocketAddrV6> for AeAddr<SocketAddrV6> {
    fn from(socket_addr: SocketAddrV6) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

impl<T> From<FullAeAddr<T>> for AeAddr<T> {
    fn from(full: FullAeAddr<T>) -> Self {
        AeAddr {
            ae_title: Some(full.ae_title),
            socket_addr: full.socket_addr,
        }
    }
}

impl<T> FromStr for AeAddr<T>
where
    T: FromStr,
{
    type Err = <T as FromStr>::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // !!! there should be a way to escape the `@`
        if let Some((ae_title, address)) = s.split_once('@') {
            Ok(AeAddr {
                ae_title: Some(ae_title)
                    .filter(|s| !s.is_empty())
                    .map(|s| s.to_string()),
                socket_addr: address.parse()?,
            })
        } else {
            Ok(AeAddr {
                ae_title: None,
                socket_addr: s.parse()?,
            })
        }
    }
}

impl<'a> TryFrom<&'a str> for AeAddr<String> {
    type Error = <AeAddr<String> as FromStr>::Err;

    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        s.parse()
    }
}

impl<T> ToSocketAddrs for AeAddr<T>
where
    T: ToSocketAddrs,
{
    type Iter = T::Iter;

    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
        self.socket_addr.to_socket_addrs()
    }
}

impl<T> std::fmt::Display for AeAddr<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let socket_addr = self.socket_addr.to_string();
        if let Some(ae_title) = &self.ae_title {
            f.write_str(&ae_title.replace("@", "\\@"))?;
            f.write_str("@")?;
        } else if socket_addr.contains("@") {
            // if formatted socket address contains a `@`,
            // we need to start the output with `@`
            // so that the start of the socket address
            // is not interpreted as an AE title
            f.write_str("@")?;
        }

        std::fmt::Display::fmt(&socket_addr, f)
    }
}

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

    #[test]
    fn ae_addr_parse() {
        // socket address can be a string
        let addr: FullAeAddr<String> = "SCP-STORAGE@127.0.0.1:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "SCP-STORAGE");
        assert_eq!(addr.socket_addr(), "127.0.0.1:104");

        // or anything else which can be parsed into a socket address
        let addr: FullAeAddr<SocketAddr> = "SCP_STORAGE@127.0.0.1:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "SCP_STORAGE");
        assert_eq!(addr.socket_addr(), &SocketAddr::from(([127, 0, 0, 1], 104)));
        assert_eq!(&addr.to_string(), "SCP_STORAGE@127.0.0.1:104");

        // IPv4 socket address
        let addr: FullAeAddr<SocketAddrV4> = "MAMMOSTORE@10.0.0.11:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "MAMMOSTORE");
        assert_eq!(
            addr.socket_addr(),
            &SocketAddrV4::new([10, 0, 0, 11].into(), 104)
        );
        assert_eq!(&addr.to_string(), "MAMMOSTORE@10.0.0.11:104");
    }

    /// test addresses without an AE title
    #[test]
    fn ae_addr_parse_no_ae() {
        // should fail
        let res = FullAeAddr::<String>::from_str("pacs.hospital.example.com:104");
        assert!(matches!(res, Err(ParseAeAddressError::MissingPart)));
        // should also fail (AE title can't be empty)
        let res = FullAeAddr::<String>::from_str("@pacs.hospital.example.com:104");
        assert!(matches!(res, Err(ParseAeAddressError::MissingPart)));

        // should return an ae addr with no AE title
        let addr: AeAddr<String> = "pacs.hospital.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), None);
        assert_eq!(addr.socket_addr(), "pacs.hospital.example.com:104");
        // should also return an ae addr with no AE title
        let addr: AeAddr<String> = "@pacs.hospital.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), None);
        assert_eq!(addr.socket_addr(), "pacs.hospital.example.com:104");
    }

    #[test]
    fn ae_addr_parse_weird_scenarios() {
        // can parse addresses with multiple @'s
        let addr: FullAeAddr<String> = "ABC@DICOM@pacs.archive.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "ABC");
        assert_eq!(addr.socket_addr(), "DICOM@pacs.archive.example.com:104");
        assert_eq!(&addr.to_string(), "ABC@DICOM@pacs.archive.example.com:104");
    }
}
<commit_msg>[ul] minor improvements to *AeAddress types<commit_after>//! Data types for addresses to nodes in DICOM networks.
//!
//! This module provides the definitions for [`FullAeAddr`] and [`AeAddr`],
//! which enable consumers to couple a socket address with an expected
//! application entity (AE) title.
//!
//! The syntax is `«ae_title»@«network_address»:«port»`,
//! which works not only with IPv4 and IPv6 addresses,
//! but also with domain names.
use std::{
    convert::TryFrom,
    net::{SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
    str::FromStr,
};

use snafu::{ensure, AsErrorSource, ResultExt, Snafu};

/// A specification for a full address to the target SCP:
/// an application entity title, plus a generic  address,
/// typically a socket address.
///
/// These addresses can be serialized and parsed
/// with the syntax `{ae_title}@{address}`,
/// where the socket address is parsed according to
/// the expectations of the parameter type `T`.
///
/// For the version of the struct without a mandatory AE title,
/// see [`AeAddr`].
///
/// # Example
///
/// ```
/// # use dicom_ul::FullAeAddr;
/// # use std::net::SocketAddr;
/// #
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # // socket address can be a string
/// let addr: FullAeAddr<String> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), "SCP-STORAGE");
/// assert_eq!(addr.socket_addr(), "127.0.0.1:104");
/// # // or anything else which can be parsed into a socket address
/// let addr: FullAeAddr<SocketAddr> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), "SCP-STORAGE");
/// assert_eq!(addr.socket_addr(), &SocketAddr::from(([127, 0, 0, 1], 104)));
/// assert_eq!(&addr.to_string(), "SCP-STORAGE@127.0.0.1:104");
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone, Eq, Hash, PartialEq)]
pub struct FullAeAddr<T> {
    ae_title: String,
    socket_addr: T,
}

impl<T> FullAeAddr<T> {
    /// Create an AE address from its bare constituent parts.
    pub fn new(ae_title: impl Into<String>, socket_addr: T) -> Self {
        FullAeAddr {
            ae_title: ae_title.into(),
            socket_addr,
        }
    }

    /// Retrieve the application entity title portion.
    pub fn ae_title(&self) -> &str {
        &self.ae_title
    }

    /// Retrieve the network address portion.
    pub fn socket_addr(&self) -> &T {
        &self.socket_addr
    }

    /// Convert the full address into its constituent parts.
    pub fn into_parts(self) -> (String, T) {
        (self.ae_title, self.socket_addr)
    }
}

impl<T> From<(String, T)> for FullAeAddr<T> {
    fn from((ae_title, socket_addr): (String, T)) -> Self {
        Self::new(ae_title, socket_addr)
    }
}

/// A error which occurred when parsing an AE address.
#[derive(Debug, Clone, Eq, PartialEq, Snafu)]
pub enum ParseAeAddressError<E>
where
    E: std::fmt::Debug + AsErrorSource,
{
    /// Missing `@` in full AE address
    MissingPart,

    /// Could not parse network socket address
    ParseSocketAddress { source: E },
}

impl<T> FromStr for FullAeAddr<T>
where
    T: FromStr,
    T::Err: std::fmt::Debug + AsErrorSource,
{
    type Err = ParseAeAddressError<<T as FromStr>::Err>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // !!! there should be a way to escape the `@`
        if let Some((ae_title, addr)) = s.split_once('@') {
            ensure!(!ae_title.is_empty(), MissingPartSnafu);
            Ok(FullAeAddr {
                ae_title: ae_title.to_string(),
                socket_addr: addr.parse().context(ParseSocketAddressSnafu)?,
            })
        } else {
            Err(ParseAeAddressError::MissingPart)
        }
    }
}

impl<T> ToSocketAddrs for FullAeAddr<T>
where
    T: ToSocketAddrs,
{
    type Iter = T::Iter;

    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
        self.socket_addr.to_socket_addrs()
    }
}

impl<T> std::fmt::Display for FullAeAddr<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.ae_title.replace('@', "\\@"))?;
        f.write_str("@")?;
        std::fmt::Display::fmt(&self.socket_addr, f)
    }
}

/// A specification for an address to the target SCP:
/// a generic network socket address
/// which may also include an application entity title.
///
/// These addresses can be serialized and parsed
/// with the syntax `{ae_title}@{address}`,
/// where the socket address is parsed according to
/// the expectations of the parameter type `T`.
///
/// For the version of the struct in which the AE title part is mandatory,
/// see [`FullAeAddr`].
///
/// # Example
///
/// ```
/// # use dicom_ul::{AeAddr, FullAeAddr};
/// # use std::net::{SocketAddr, SocketAddrV4};
/// #
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let addr: AeAddr<SocketAddrV4> = "SCP-STORAGE@127.0.0.1:104".parse()?;
/// assert_eq!(addr.ae_title(), Some("SCP-STORAGE"));
/// assert_eq!(addr.socket_addr(), &SocketAddrV4::new([127, 0, 0, 1].into(), 104));
/// assert_eq!(&addr.to_string(), "SCP-STORAGE@127.0.0.1:104");
///
/// // AE title can be missing
/// let addr: AeAddr<String> = "192.168.1.99:1045".parse()?;
/// assert_eq!(addr.ae_title(), None);
/// // but can be provided later
/// let full_addr: FullAeAddr<_> = addr.with_ae_title("SCP-QUERY");
/// assert_eq!(full_addr.ae_title(), "SCP-QUERY");
/// assert_eq!(&full_addr.to_string(), "SCP-QUERY@192.168.1.99:1045");
/// # Ok(())
/// # }
/// ```
#[derive(Debug, Clone, Eq, Hash, PartialEq)]
pub struct AeAddr<T> {
    ae_title: Option<String>,
    socket_addr: T,
}

impl<T> AeAddr<T> {
    /// Create an AE address from its bare constituent parts.
    pub fn new(ae_title: impl Into<String>, socket_addr: T) -> Self {
        AeAddr {
            ae_title: Some(ae_title.into()),
            socket_addr,
        }
    }

    /// Create an address with a missing AE title.
    pub fn new_socket_addr(socket_addr: T) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }

    /// Retrieve the application entity title portion, if present.
    pub fn ae_title(&self) -> Option<&str> {
        self.ae_title.as_deref()
    }

    /// Retrieve the socket address portion.
    pub fn socket_addr(&self) -> &T {
        &self.socket_addr
    }

    /// Create a new address with the full application entity target,
    /// discarding any potentially existing AE title.
    pub fn with_ae_title(self, ae_title: impl Into<String>) -> FullAeAddr<T> {
        FullAeAddr {
            ae_title: ae_title.into(),
            socket_addr: self.socket_addr,
        }
    }

    /// Create a new address with the full application entity target,
    /// using the given AE title if it is missing.
    pub fn with_default_ae_title(self, ae_title: impl Into<String>) -> FullAeAddr<T> {
        FullAeAddr {
            ae_title: self.ae_title.unwrap_or_else(|| ae_title.into()),
            socket_addr: self.socket_addr,
        }
    }

    /// Convert the address into its constituent parts.
    pub fn into_parts(self) -> (Option<String>, T) {
        (self.ae_title, self.socket_addr)
    }
}

/// This conversion provides a socket address without an AE title.
impl From<SocketAddr> for AeAddr<SocketAddr> {
    fn from(socket_addr: SocketAddr) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

/// This conversion provides an IPv4 socket address without an AE title.
impl From<SocketAddrV4> for AeAddr<SocketAddrV4> {
    fn from(socket_addr: SocketAddrV4) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

/// This conversion provides an IPv6 socket address without an AE title.
impl From<SocketAddrV6> for AeAddr<SocketAddrV6> {
    fn from(socket_addr: SocketAddrV6) -> Self {
        AeAddr {
            ae_title: None,
            socket_addr,
        }
    }
}

impl<T> From<FullAeAddr<T>> for AeAddr<T> {
    fn from(full: FullAeAddr<T>) -> Self {
        AeAddr {
            ae_title: Some(full.ae_title),
            socket_addr: full.socket_addr,
        }
    }
}

impl<T> FromStr for AeAddr<T>
where
    T: FromStr,
{
    type Err = <T as FromStr>::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // !!! there should be a way to escape the `@`
        if let Some((ae_title, address)) = s.split_once('@') {
            Ok(AeAddr {
                ae_title: Some(ae_title)
                    .filter(|s| !s.is_empty())
                    .map(|s| s.to_string()),
                socket_addr: address.parse()?,
            })
        } else {
            Ok(AeAddr {
                ae_title: None,
                socket_addr: s.parse()?,
            })
        }
    }
}

impl<'a> TryFrom<&'a str> for AeAddr<String> {
    type Error = <AeAddr<String> as FromStr>::Err;

    fn try_from(s: &'a str) -> Result<Self, Self::Error> {
        s.parse()
    }
}

impl<T> ToSocketAddrs for AeAddr<T>
where
    T: ToSocketAddrs,
{
    type Iter = T::Iter;

    fn to_socket_addrs(&self) -> std::io::Result<Self::Iter> {
        self.socket_addr.to_socket_addrs()
    }
}

impl<T> std::fmt::Display for AeAddr<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let socket_addr = self.socket_addr.to_string();
        if let Some(ae_title) = &self.ae_title {
            f.write_str(&ae_title.replace('@', "\\@"))?;
            f.write_str("@")?;
        } else if socket_addr.contains('@') {
            // if formatted socket address contains a `@`,
            // we need to start the output with `@`
            // so that the start of the socket address
            // is not interpreted as an AE title
            f.write_str("@")?;
        }

        std::fmt::Display::fmt(&socket_addr, f)
    }
}

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

    #[test]
    fn ae_addr_parse() {
        // socket address can be a string
        let addr: FullAeAddr<String> = "SCP-STORAGE@127.0.0.1:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "SCP-STORAGE");
        assert_eq!(addr.socket_addr(), "127.0.0.1:104");

        // or anything else which can be parsed into a socket address
        let addr: FullAeAddr<SocketAddr> = "SCP_STORAGE@127.0.0.1:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "SCP_STORAGE");
        assert_eq!(addr.socket_addr(), &SocketAddr::from(([127, 0, 0, 1], 104)));
        assert_eq!(&addr.to_string(), "SCP_STORAGE@127.0.0.1:104");

        // IPv4 socket address
        let addr: FullAeAddr<SocketAddrV4> = "MAMMOSTORE@10.0.0.11:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "MAMMOSTORE");
        assert_eq!(
            addr.socket_addr(),
            &SocketAddrV4::new([10, 0, 0, 11].into(), 104)
        );
        assert_eq!(&addr.to_string(), "MAMMOSTORE@10.0.0.11:104");
    }

    /// test addresses without an AE title
    #[test]
    fn ae_addr_parse_no_ae() {
        // should fail
        let res = FullAeAddr::<String>::from_str("pacs.hospital.example.com:104");
        assert!(matches!(res, Err(ParseAeAddressError::MissingPart)));
        // should also fail (AE title can't be empty)
        let res = FullAeAddr::<String>::from_str("@pacs.hospital.example.com:104");
        assert!(matches!(res, Err(ParseAeAddressError::MissingPart)));

        // should return an ae addr with no AE title
        let addr: AeAddr<String> = "pacs.hospital.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), None);
        assert_eq!(addr.socket_addr(), "pacs.hospital.example.com:104");
        // should also return an ae addr with no AE title
        let addr: AeAddr<String> = "@pacs.hospital.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), None);
        assert_eq!(addr.socket_addr(), "pacs.hospital.example.com:104");
    }

    #[test]
    fn ae_addr_parse_weird_scenarios() {
        // can parse addresses with multiple @'s
        let addr: FullAeAddr<String> = "ABC@DICOM@pacs.archive.example.com:104".parse().unwrap();
        assert_eq!(addr.ae_title(), "ABC");
        assert_eq!(addr.socket_addr(), "DICOM@pacs.archive.example.com:104");
        assert_eq!(&addr.to_string(), "ABC@DICOM@pacs.archive.example.com:104");
    }
}
<|endoftext|>