<commit_before>//
extern crate termion;

//
use self::termion::event::{Event, Key, MouseEvent};
use self::termion::input::{TermRead, MouseTerminal};
use self::termion::raw::IntoRawMode;
use self::termion::terminal_size;

use self::termion::cursor::DetectCursorPos;

use std::io::{self, Write, stdin, Stdout};

//
use core::screen::Screen;
use core::codepointinfo::CodepointInfo;

use core::editor::Editor;

fn fill_screen(screen: &mut Screen, data: &[u8]) {

    screen.clear();

    let mut offset = 0;
    for c in data {

        let mut displayed_cp: char = *c as char;
        if *c as char == '\n' as char {
            displayed_cp = ' ';
        }

        let cpi = CodepointInfo {
            cp: *c as char,
            displayed_cp: displayed_cp,
            offset: offset,
        };

        let (ok, _) = screen.push(cpi);
        offset += 1;
        if ok == false {
            break;
        }
    }
}

fn draw_screen(screen: &Screen, stdout: &mut Stdout) {

    write!(stdout, "{}", termion::cursor::Goto(1, 1)).unwrap();

    for l in 0..screen.height {
        let line = &screen.line[l];

        for c in 0..line.width {
            let cpi = &line.chars[c];
            write!(stdout, "{}", cpi.displayed_cp).unwrap();
        }

        if l < screen.height - 1 {
            write!(stdout, "\r\n").unwrap();
        }
    }

    stdout.flush().unwrap();
}

fn draw_buffer(buf: &Option<&Box<::core::buffer::Buffer>>,
               mut scr: &mut Screen,
               mut stdout: &mut Stdout) {
    match *buf {
        Some(ref b) => {
            match b.byte_buffer {

                Some(ref byte_buf) => {
                    fill_screen(&mut scr, &byte_buf.data);
                    draw_screen(&scr, &mut stdout);
                }
                _ => {}
            }
        }
        _ => {}
    }
}

fn terminal_clear_current_line(mut stdout: &mut Stdout, line_width: u16) {
    let (current_line, _) = match stdout.cursor_pos() {
        Err(_) => { (0, 0) },
        Ok((c, _)) => { (c as  u16, 0) },
    };

    for _ in 0..line_width {
        write!(stdout, " ").unwrap();
    }
    terminal_goto_line(&mut stdout, current_line);
}

fn terminal_goto_line(mut stdout: &mut Stdout, y: u16) {
    write!(stdout, "{}", termion::cursor::Goto(1, y)).unwrap();
}

/*
fn terminal_cursor_to(mut stdout: &mut Stdout, x: u16, y: u16) {
    write!(stdout, "{}", termion::cursor::Goto(x, y)).unwrap();
}
*/

fn terminal_clear_screen(stdout: &mut Stdout, clear_toggle_flag: &mut bool) {
    if *clear_toggle_flag == true {
        write!(stdout, "{}", termion::clear::All).unwrap();
        stdout.flush().unwrap();
    }
    *clear_toggle_flag = false;
}

pub fn main_loop(editor: &mut Editor) {

    let mut stdout = MouseTerminal::from(io::stdout().into_raw_mode().unwrap());

    write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
    stdout.flush().unwrap();

    let mut keys: Vec<Event> = Vec::new();

    let mut quit = false;
    let mut clear_toggle_flag = true;

    //
    let display_status = true;

    // select file
    let mut bid = 2;
    let mut buf = editor.buffer_map.get(&bid);

    while !quit {
        let (width, height) = terminal_size().unwrap();
        let mut scr = Screen::new(width as usize, height as usize);

        terminal_clear_screen(&mut stdout, &mut clear_toggle_flag);
        write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
        stdout.flush().unwrap();

        draw_buffer(&buf, &mut scr, &mut stdout);
        if display_status == true {
            display_status_line(&buf, "", height, width, &mut stdout);
        }

        for evt in stdin().events() {

            keys.push(evt.unwrap());
            let evt = keys[keys.len() - 1].clone();

            terminal_clear_screen(&mut stdout, &mut clear_toggle_flag);
            write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
            stdout.flush().unwrap();

            draw_buffer(&buf, &mut scr, &mut stdout);

            // Print recieved Events...
            let mut status = String::new();
            match evt {

                Event::Key(k) => {
                    match k {
                        // Exit.
                        Key::Ctrl('r') => {
                            keys.clear();
                        }

                        Key::Ctrl('c') => {
                            if keys.len() > 1 {
                                if let Event::Key(prev_event) = keys[keys.len() - 2] {
                                    if let Key::Ctrl(prev_char) = prev_event {
                                        if prev_char == 'x' {
                                            quit = true;
                                            break;
                                        }
                                    }
                                }
                            }
                        }

                        Key::Char(c) => {
                            if c == '\n' {
                                status = format!("'{}'", "<newline>");
                            } else {
                                status = format!("'{}'", c);
                            }
                        }
                        Key::Alt(c) => status = format!("Alt-{}", c),
                        Key::Ctrl(c) => status = format!("Ctrl-{}", c),

                        Key::F(1) => {
                            if bid > 0 {
                                bid -= 1;
                            }
                            buf = editor.buffer_map.get(&bid);
                            clear_toggle_flag = true;
                            break;
                        }

                        Key::F(2) => {
                            bid = ::std::cmp::min(bid + 1, (editor.buffer_map.len() - 1) as u64);
                            buf = editor.buffer_map.get(&bid);
                            clear_toggle_flag = true;
                            break;
                        }

                        Key::F(f) => status = format!("F{:?}", f),
                        Key::Left => status = format!("<left>"),
                        Key::Right => status = format!("<right>"),
                        Key::Up => status = format!("<up>"),
                        Key::Down => status = format!("<down>"),
                        Key::Backspace => status = format!("<backspace>"),
                        Key::Home => status = format!("<Home>"),
                        Key::End => status = format!("<End>"),
                        Key::PageUp => status = format!("<PageUp>"),
                        Key::PageDown => status = format!("<PageDown>"),
                        Key::Delete => status = format!("<Delete>"),
                        Key::Insert => status = format!("<Insert>"),
                        Key::Esc => status = format!("<Esc>"),
                        _ => status = format!("Other"),
                    }
                }

                Event::Mouse(m) => {
                    match m {
                        MouseEvent::Press(mb, x, y) => {
                            status = format!("MouseEvent::Press => MouseButton {:?} @ ({}, {})",
                                             mb,
                                             x,
                                             y);
                        }

                        MouseEvent::Release(x, y) => {
                            status = format!("MouseEvent::Release => @ ({}, {})", x, y);
                        }

                        MouseEvent::Hold(x, y) => {
                            status = format!("MouseEvent::Hold => @ ({}, {})", x, y);
                        }
                    };
                }

                Event::Unsupported(_) => {}
            }

            if display_status == true {
                display_status_line(&buf, &status, height, width, &mut stdout);
            }
        }
    }

    // quit
    // clear, restore cursor
    write!(stdout, "{}{}", termion::clear::All, termion::cursor::Show).unwrap();
    stdout.flush().unwrap();
}

fn display_status_line(buf: &Option<&Box<::core::buffer::Buffer>>,
                       status: &str,
                       line: u16,
                       width: u16,
                       mut stdout: &mut Stdout) {
    // select/clear last line
    let name = match *buf {
        Some(ref b) => b.name.as_str(),
        None => "",
    };

    let file_name = match *buf {
        Some(ref b) => {
            match b.byte_buffer {
                Some(ref bb) => bb.file_name.as_str(),
                None => "",
            }
        }
        None => "",
    };

    terminal_goto_line(&mut stdout, line);
    terminal_clear_current_line(&mut stdout, width);

    let status_str = format!("buffer_name '{}', file: '{}', event '{}'",
                             name,
                             file_name,
                             status);

    print!("{}", status_str);
    stdout.flush().unwrap();
}
<commit_msg>do not use stdout.cursor_pos(), it will corrput arrow escape sequences<commit_after>//
extern crate termion;

//
use self::termion::event::{Event, Key, MouseEvent};
use self::termion::input::{TermRead, MouseTerminal};
use self::termion::raw::IntoRawMode;
use self::termion::terminal_size;

use self::termion::cursor::DetectCursorPos;

use std::io::{self, Write, stdin, Stdout};

//
use core::screen::Screen;
use core::codepointinfo::CodepointInfo;

use core::editor::Editor;

fn fill_screen(screen: &mut Screen, data: &[u8]) {

    screen.clear();

    let mut offset = 0;
    for c in data {

        let mut displayed_cp: char = *c as char;
        if *c as char == '\n' as char {
            displayed_cp = ' ';
        }

        let cpi = CodepointInfo {
            cp: *c as char,
            displayed_cp: displayed_cp,
            offset: offset,
        };

        let (ok, _) = screen.push(cpi);
        offset += 1;
        if ok == false {
            break;
        }
    }
}

fn draw_screen(screen: &Screen, stdout: &mut Stdout) {

    write!(stdout, "{}", termion::cursor::Goto(1, 1)).unwrap();

    for l in 0..screen.height {
        let line = &screen.line[l];

        for c in 0..line.width {
            let cpi = &line.chars[c];
            write!(stdout, "{}", cpi.displayed_cp).unwrap();
        }

        if l < screen.height - 1 {
            write!(stdout, "\r\n").unwrap();
        }
    }

    stdout.flush().unwrap();
}

fn draw_buffer(buf: &Option<&Box<::core::buffer::Buffer>>,
               mut scr: &mut Screen,
               mut stdout: &mut Stdout) {
    match *buf {
        Some(ref b) => {
            match b.byte_buffer {

                Some(ref byte_buf) => {
                    fill_screen(&mut scr, &byte_buf.data);
                    draw_screen(&scr, &mut stdout);
                }
                _ => {}
            }
        }
        _ => {}
    }
}

fn terminal_clear_current_line(mut stdout: &mut Stdout, line_width: u16) {
    for _ in 0..line_width {
        write!(stdout, " ").unwrap();
    }
}

fn terminal_goto_line(mut stdout: &mut Stdout, y: u16) {
    write!(stdout, "{}", termion::cursor::Goto(1, y)).unwrap();
}

fn terminal_cursor_to(mut stdout: &mut Stdout, x: u16, y: u16) {
    write!(stdout, "{}", termion::cursor::Goto(x, y)).unwrap();
}

fn terminal_clear_screen(stdout: &mut Stdout, clear_toggle_flag: &mut bool) {
    if *clear_toggle_flag == true {
        write!(stdout, "{}", termion::clear::All).unwrap();
        stdout.flush().unwrap();
    }
    *clear_toggle_flag = false;
}

pub fn main_loop(editor: &mut Editor) {

    let mut stdout = MouseTerminal::from(io::stdout().into_raw_mode().unwrap());

    write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
    stdout.flush().unwrap();

    let mut keys: Vec<Event> = Vec::new();

    let mut quit = false;
    let mut clear_toggle_flag = true;

    //
    let display_status = true;

    // select file
    let mut bid = 2;
    let mut buf = editor.buffer_map.get(&bid);

    while !quit {
        let (width, height) = terminal_size().unwrap();
        let mut scr = Screen::new(width as usize, height as usize);

        terminal_clear_screen(&mut stdout, &mut clear_toggle_flag);
        write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
        stdout.flush().unwrap();

        draw_buffer(&buf, &mut scr, &mut stdout);
        if display_status == true {
            display_status_line(&buf, "", height, width, &mut stdout);
        }

        for evt in stdin().events() {

            keys.push(evt.unwrap());
            let evt = keys[keys.len() - 1].clone();

            terminal_clear_screen(&mut stdout, &mut clear_toggle_flag);
            write!(stdout, "{}{}", termion::cursor::Hide, termion::clear::All).unwrap();
            stdout.flush().unwrap();

            draw_buffer(&buf, &mut scr, &mut stdout);

            // Print recieved Events...
            let mut status = String::new();
            match evt {

                Event::Key(k) => {
                    match k {
                        // Exit.
                        Key::Ctrl('r') => {
                            keys.clear();
                        }

                        Key::Ctrl('c') => {
                            if keys.len() > 1 {
                                if let Event::Key(prev_event) = keys[keys.len() - 2] {
                                    if let Key::Ctrl(prev_char) = prev_event {
                                        if prev_char == 'x' {
                                            quit = true;
                                            break;
                                        }
                                    }
                                }
                            }
                        }

                        Key::Char(c) => {
                            if c == '\n' {
                                status = format!("'{}'", "<newline>");
                            } else {
                                status = format!("'{}'", c);
                            }
                        }
                        Key::Alt(c) => status = format!("Alt-{}", c),
                        Key::Ctrl(c) => status = format!("Ctrl-{}", c),

                        Key::F(1) => {
                            if bid > 0 {
                                bid -= 1;
                            }
                            buf = editor.buffer_map.get(&bid);
                            clear_toggle_flag = true;
                            break;
                        }

                        Key::F(2) => {
                            bid = ::std::cmp::min(bid + 1, (editor.buffer_map.len() - 1) as u64);
                            buf = editor.buffer_map.get(&bid);
                            clear_toggle_flag = true;
                            break;
                        }

                        Key::F(f) => status = format!("F{:?}", f),
                        Key::Left => status = format!("<left>"),
                        Key::Right => status = format!("<right>"),
                        Key::Up => status = format!("<up>"),
                        Key::Down => status = format!("<down>"),
                        Key::Backspace => status = format!("<backspace>"),
                        Key::Home => status = format!("<Home>"),
                        Key::End => status = format!("<End>"),
                        Key::PageUp => status = format!("<PageUp>"),
                        Key::PageDown => status = format!("<PageDown>"),
                        Key::Delete => status = format!("<Delete>"),
                        Key::Insert => status = format!("<Insert>"),
                        Key::Esc => status = format!("<Esc>"),
                        _ => status = format!("Other"),
                    }
                }

                Event::Mouse(m) => {
                    match m {
                        MouseEvent::Press(mb, x, y) => {
                            status = format!("MouseEvent::Press => MouseButton {:?} @ ({}, {})",
                                             mb,
                                             x,
                                             y);
                        }

                        MouseEvent::Release(x, y) => {
                            status = format!("MouseEvent::Release => @ ({}, {})", x, y);
                        }

                        MouseEvent::Hold(x, y) => {
                            status = format!("MouseEvent::Hold => @ ({}, {})", x, y);
                        }
                    };
                }

                Event::Unsupported(_) => {}
            }

            if display_status == true {
                display_status_line(&buf, &status, height, width, &mut stdout);
            }
        }
    }

    // quit
    // clear, restore cursor
    write!(stdout, "{}{}", termion::clear::All, termion::cursor::Show).unwrap();
    stdout.flush().unwrap();
}

fn display_status_line(buf: &Option<&Box<::core::buffer::Buffer>>,
                       status: &str,
                       line: u16,
                       width: u16,
                       mut stdout: &mut Stdout) {
    // select/clear last line
    let name = match *buf {
        Some(ref b) => b.name.as_str(),
        None => "",
    };

    let file_name = match *buf {
        Some(ref b) => {
            match b.byte_buffer {
                Some(ref bb) => bb.file_name.as_str(),
                None => "",
            }
        }
        None => "",
    };

    terminal_cursor_to(&mut stdout, 1, line);
    terminal_clear_current_line(&mut stdout, width);
    terminal_cursor_to(&mut stdout, 1, line);

    let status_str = format!("buffer_name '{}', file: '{}', event '{}'",
                             name,
                             file_name,
                             status);

    print!("{}", status_str);
    stdout.flush().unwrap();
}
<|endoftext|>


<commit_before>use std::io::buffered::BufferedStream;
use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;
use std::io::Stream;
use std::vec::bytes::push_bytes;
use std::vec;

pub enum Result {
  Nil,
  Int(int),
  Data(~[u8]),
  List(~[Result]),
  Error(~str),
  Status(~str),
  ProtocolError(&'static str)
}

fn read_char<T: Stream>(io: &mut BufferedStream<T>) -> Option<char> {
  match io.read_byte() {
    Some(ch) => Some(ch as char),
    None     => None
  }
}

fn parse_data<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let res =
    if (len > 0) {
      let bytes = io.read_bytes(len);
      if bytes.len() != len {
        return ProtocolError("Invalid number of bytes")
      } else {
        Data(bytes)
      }
    } else {
      Data(~[])
    };

  if (read_char(io) != Some('\r')) {
    return ProtocolError("Carriage return expected"); // TODO: ignore
  }

  if (read_char(io) != Some('\n')) {
    return ProtocolError("Newline expected");
  }

  return res;
}

fn parse_list<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let mut list: ~[Result] = vec::with_capacity(len);

  for _ in range(0, len) {
    match parse(io) {
      ProtocolError(err) => {
        return ProtocolError(err);
      }
      other => {
        list.push(other);
      }
    }
  }

  List(list)
}

fn parse_int_line<T: Stream>(io: &mut BufferedStream<T>) -> Option<int> {
  let mut i: int = 0;
  let mut digits: uint = 0;
  let mut negative: bool = false;

  loop {
    match read_char(io) {
      None => { return None }
      Some(ch) => {
        match ch {
          '0' .. '9' => {
            digits += 1;
            i = (i * 10) + (ch as int - '0' as int);
          }
          '-' => {
            if negative { return None }
            negative = true
          }
          '\r' => {
            if read_char(io) != Some('\n') { return None }
            break
          }
          '\n' => { break }
          _ => { return None }
        }
      }
    }
  }

  if digits == 0 { return None }

  if negative { Some(-i) }
  else { Some(i) }
}

fn parse_n<T: Stream>(io: &mut BufferedStream<T>, f: |uint, &mut BufferedStream<T>| -> Result) -> Result {
  match parse_int_line(io) {
    Some(-1) => Nil,
    Some(len) if len >= 0 => f(len as uint, io),
    _ => ProtocolError("Invalid number")
  }
}

fn parse_status<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Status(line),
    None       => ProtocolError("Invalid status line")
  }
}

fn parse_error<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Error(line),
    None       => ProtocolError("Invalid error line")
  }
}

pub fn parse<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match read_char(io) {
    Some('$') => parse_n(io, parse_data),
    Some('*') => parse_n(io, parse_list),
    Some('+') => parse_status(io),
    Some('-') => parse_error(io),
    Some(':') => {
      match parse_int_line(io) {
        Some(i) => Int(i),
        None => ProtocolError("Invalid number")
      }
    },
    Some(_)   => ProtocolError("Invalid character"),
    None      => ProtocolError("Invalid EOF")
  }
}

pub struct CommandWriter {
  buf: ~[u8]
}

impl CommandWriter {
  pub fn new() -> CommandWriter {
    CommandWriter { buf: ~[] }
  }

  pub fn args<'a>(&'a mut self, n: uint) -> &'a mut CommandWriter {
    self.write_char('*');
    self.write_uint(n);
    self.write_crnl();
    self
  }

  pub fn arg_bin<'a>(&'a mut self, arg: &[u8]) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write(arg);
    self.write_crnl();
    self
  }

  pub fn nil(&mut self) {
    self.write_str("$-1");
    self.write_crnl();
  }

  pub fn arg_str<'a>(&'a mut self, arg: &str) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write_str(arg);
    self.write_crnl();
    self
  }

  pub fn error(&mut self, err: &str) {
    self.write_char('-');
    self.write_str(err);
    self.write_crnl();
  }

  pub fn status(&mut self, status: &str) {
    self.write_char('+');
    self.write_str(status);
    self.write_crnl();
  }

  fn write_crnl(&mut self) {
    self.write_char('\r');
    self.write_char('\n');
  }

  fn write_uint(&mut self, n: uint) {
    if n < 10 {
      self.write_byte('0' as u8 + (n as u8));
    }
    else {
      push_bytes(&mut self.buf, n.to_str().into_bytes());
    }
  }

  fn write_str(&mut self, s: &str) {
    push_bytes(&mut self.buf, s.as_bytes());
  }

  fn write(&mut self, s: &[u8]) {
    push_bytes(&mut self.buf, s);
  }

  fn write_char(&mut self, s: char) {
    self.buf.push(s as u8);
  }

  fn write_byte(&mut self, b: u8) {
    self.buf.push(b);
  }

  pub fn with_buf<T>(&self, f: |&[u8]| -> T) -> T {
    f(self.buf.as_slice())
  }
}

fn execute<T: Stream>(cmd: &[u8], io: &mut BufferedStream<T>) -> Result {
  io.write(cmd);
  io.flush();
  parse(io)
}

pub struct Client<T> {
  priv io: BufferedStream<T>
}

impl Client<TcpStream> {
  pub fn new(sock_addr: &str) -> Client<TcpStream> {
    let addr = from_str::<SocketAddr>(sock_addr).unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    Client::new_from_stream(tcp_stream)
  }
}

impl<T: Stream> Client<T> {
  pub fn new_from_stream(io: T) -> Client<T> {
    Client { io: BufferedStream::new(io) }
  }

  pub fn get(&mut self, key: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(2).
        arg_str("GET").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn set(&mut self, key: &str, val: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(3).
        arg_str("SET").
        arg_str(key).
        arg_str(val).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }
}
<commit_msg>Add get_int/get_str and set_int commands<commit_after>use std::io::buffered::BufferedStream;
use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;
use std::io::Stream;
use std::vec::bytes::push_bytes;
use std::vec;
use std::str::from_utf8;

pub enum Result {
  Nil,
  Int(int),
  Data(~[u8]),
  List(~[Result]),
  Error(~str),
  Status(~str),
  ProtocolError(&'static str)
}

fn read_char<T: Stream>(io: &mut BufferedStream<T>) -> Option<char> {
  match io.read_byte() {
    Some(ch) => Some(ch as char),
    None     => None
  }
}

fn parse_data<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let res =
    if (len > 0) {
      let bytes = io.read_bytes(len);
      if bytes.len() != len {
        return ProtocolError("Invalid number of bytes")
      } else {
        Data(bytes)
      }
    } else {
      Data(~[])
    };

  if (read_char(io) != Some('\r')) {
    return ProtocolError("Carriage return expected"); // TODO: ignore
  }

  if (read_char(io) != Some('\n')) {
    return ProtocolError("Newline expected");
  }

  return res;
}

fn parse_list<T: Stream>(len: uint, io: &mut BufferedStream<T>) -> Result {
  let mut list: ~[Result] = vec::with_capacity(len);

  for _ in range(0, len) {
    match parse(io) {
      ProtocolError(err) => {
        return ProtocolError(err);
      }
      other => {
        list.push(other);
      }
    }
  }

  List(list)
}

fn parse_int_line<T: Stream>(io: &mut BufferedStream<T>) -> Option<int> {
  let mut i: int = 0;
  let mut digits: uint = 0;
  let mut negative: bool = false;

  loop {
    match read_char(io) {
      None => { return None }
      Some(ch) => {
        match ch {
          '0' .. '9' => {
            digits += 1;
            i = (i * 10) + (ch as int - '0' as int);
          }
          '-' => {
            if negative { return None }
            negative = true
          }
          '\r' => {
            if read_char(io) != Some('\n') { return None }
            break
          }
          '\n' => { break }
          _ => { return None }
        }
      }
    }
  }

  if digits == 0 { return None }

  if negative { Some(-i) }
  else { Some(i) }
}

fn parse_n<T: Stream>(io: &mut BufferedStream<T>, f: |uint, &mut BufferedStream<T>| -> Result) -> Result {
  match parse_int_line(io) {
    Some(-1) => Nil,
    Some(len) if len >= 0 => f(len as uint, io),
    _ => ProtocolError("Invalid number")
  }
}

fn parse_status<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Status(line),
    None       => ProtocolError("Invalid status line")
  }
}

fn parse_error<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match io.read_line() {
    Some(line) => Error(line),
    None       => ProtocolError("Invalid error line")
  }
}

pub fn parse<T: Stream>(io: &mut BufferedStream<T>) -> Result {
  match read_char(io) {
    Some('$') => parse_n(io, parse_data),
    Some('*') => parse_n(io, parse_list),
    Some('+') => parse_status(io),
    Some('-') => parse_error(io),
    Some(':') => {
      match parse_int_line(io) {
        Some(i) => Int(i),
        None => ProtocolError("Invalid number")
      }
    },
    Some(_)   => ProtocolError("Invalid character"),
    None      => ProtocolError("Invalid EOF")
  }
}

pub struct CommandWriter {
  buf: ~[u8]
}

impl CommandWriter {
  pub fn new() -> CommandWriter {
    CommandWriter { buf: ~[] }
  }

  pub fn args<'a>(&'a mut self, n: uint) -> &'a mut CommandWriter {
    self.write_char('*');
    self.write_uint(n);
    self.write_crnl();
    self
  }

  pub fn arg_bin<'a>(&'a mut self, arg: &[u8]) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write(arg);
    self.write_crnl();
    self
  }

  pub fn nil(&mut self) {
    self.write_str("$-1");
    self.write_crnl();
  }

  pub fn arg_str<'a>(&'a mut self, arg: &str) -> &'a mut CommandWriter {
    self.write_char('$');
    self.write_uint(arg.len());
    self.write_crnl();
    self.write_str(arg);
    self.write_crnl();
    self
  }

  pub fn error(&mut self, err: &str) {
    self.write_char('-');
    self.write_str(err);
    self.write_crnl();
  }

  pub fn status(&mut self, status: &str) {
    self.write_char('+');
    self.write_str(status);
    self.write_crnl();
  }

  fn write_crnl(&mut self) {
    self.write_char('\r');
    self.write_char('\n');
  }

  fn write_uint(&mut self, n: uint) {
    if n < 10 {
      self.write_byte('0' as u8 + (n as u8));
    }
    else {
      push_bytes(&mut self.buf, n.to_str().into_bytes());
    }
  }

  fn write_str(&mut self, s: &str) {
    push_bytes(&mut self.buf, s.as_bytes());
  }

  fn write(&mut self, s: &[u8]) {
    push_bytes(&mut self.buf, s);
  }

  fn write_char(&mut self, s: char) {
    self.buf.push(s as u8);
  }

  fn write_byte(&mut self, b: u8) {
    self.buf.push(b);
  }

  pub fn with_buf<T>(&self, f: |&[u8]| -> T) -> T {
    f(self.buf.as_slice())
  }
}

fn execute<T: Stream>(cmd: &[u8], io: &mut BufferedStream<T>) -> Result {
  io.write(cmd);
  io.flush();
  parse(io)
}

pub struct Client<T> {
  priv io: BufferedStream<T>
}

impl Client<TcpStream> {
  pub fn new(sock_addr: &str) -> Client<TcpStream> {
    let addr = from_str::<SocketAddr>(sock_addr).unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    Client::new_from_stream(tcp_stream)
  }
}

impl<T: Stream> Client<T> {
  pub fn new_from_stream(io: T) -> Client<T> {
    Client { io: BufferedStream::new(io) }
  }

  pub fn get(&mut self, key: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(2).
        arg_str("GET").
        arg_str(key).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn get_str(&mut self, key: &str) -> Option<~str> {
    match self.get(key) {
      Nil => None,
      Int(i) => Some(i.to_str()),
      Data(ref bytes) => Some(from_utf8(*bytes).to_owned()),
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn get_int(&mut self, key: &str) -> Option<int> {
    match self.get(key) {
      Nil => None,
      Int(i) => Some(i),
      _ => fail!("Invalid result type from Redis")
    }
  }

  pub fn set(&mut self, key: &str, val: &str) -> Result {
    let mut cwr = CommandWriter::new();
    cwr.args(3).
        arg_str("SET").
        arg_str(key).
        arg_str(val).
        with_buf(|cmd| execute(cmd, &mut self.io))
  }

  pub fn set_int(&mut self, key: &str, val: int) -> Result {
    self.set(key, val.to_str())
  }
}
<|endoftext|>


<commit_before>// Copyright 2019 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io;
use std::mem;
use std::sync::Arc;
use std::thread;

use base::error;
use base::warn;
use base::AsRawDescriptor;
use base::Error as SysError;
use base::Event;
use base::RawDescriptor;
use base::Tube;
use data_model::DataInit;
use data_model::Le32;
use remain::sorted;
use resources::Alloc;
use sync::Mutex;
use thiserror::Error;
use vm_control::FsMappingRequest;
use vm_control::VmResponse;
use vm_memory::GuestMemory;

use crate::pci::PciAddress;
use crate::pci::PciBarConfiguration;
use crate::pci::PciBarPrefetchable;
use crate::pci::PciBarRegionType;
use crate::pci::PciCapability;
use crate::virtio::copy_config;
use crate::virtio::DescriptorError;
use crate::virtio::DeviceType;
use crate::virtio::Interrupt;
use crate::virtio::PciCapabilityType;
use crate::virtio::Queue;
use crate::virtio::VirtioDevice;
use crate::virtio::VirtioPciShmCap;
use crate::Suspendable;

mod caps;
mod multikey;
pub mod passthrough;
mod read_dir;
mod worker;

use fuse::Server;
use passthrough::PassthroughFs;
pub use worker::process_fs_queue;
use worker::Worker;

// The fs device does not have a fixed number of queues.
pub const QUEUE_SIZE: u16 = 1024;

const FS_BAR_NUM: u8 = 4;
const FS_BAR_OFFSET: u64 = 0;
const FS_BAR_SIZE: u64 = 1 << 33;

/// Defined in kernel/include/uapi/linux/virtio_fs.h.
const VIRTIO_FS_SHMCAP_ID_CACHE: u8 = 0;

/// The maximum allowable length of the tag used to identify a specific virtio-fs device.
pub const FS_MAX_TAG_LEN: usize = 36;

/// kernel/include/uapi/linux/virtio_fs.h
#[repr(C, packed)]
#[derive(Clone, Copy)]
pub struct virtio_fs_config {
    /// Filesystem name (UTF-8, not NUL-terminated, padded with NULs)
    pub tag: [u8; FS_MAX_TAG_LEN],
    /// Number of request queues
    pub num_request_queues: Le32,
}

// Safe because all members are plain old data and any value is valid.
unsafe impl DataInit for virtio_fs_config {}

/// Errors that may occur during the creation or operation of an Fs device.
#[sorted]
#[derive(Error, Debug)]
pub enum Error {
    /// Failed to create the file system.
    #[error("failed to create file system: {0}")]
    CreateFs(io::Error),
    /// Creating WaitContext failed.
    #[error("failed to create WaitContext: {0}")]
    CreateWaitContext(SysError),
    /// Error happened in FUSE.
    #[error("fuse error: {0}")]
    FuseError(fuse::Error),
    /// Failed to get the securebits for the worker thread.
    #[error("failed to get securebits for the worker thread: {0}")]
    GetSecurebits(SysError),
    /// The `len` field of the header is too small.
    #[error("DescriptorChain is invalid: {0}")]
    InvalidDescriptorChain(DescriptorError),
    /// A request is missing readable descriptors.
    #[error("request does not have any readable descriptors")]
    NoReadableDescriptors,
    /// A request is missing writable descriptors.
    #[error("request does not have any writable descriptors")]
    NoWritableDescriptors,
    /// Error while reading from the virtio queue's Event.
    #[error("failed to read from virtio queue Event: {0}")]
    ReadQueueEvent(SysError),
    /// Failed to set the securebits for the worker thread.
    #[error("failed to set securebits for the worker thread: {0}")]
    SetSecurebits(SysError),
    /// Failed to signal the virio used queue.
    #[error("failed to signal used queue: {0}")]
    SignalUsedQueue(SysError),
    /// The tag for the Fs device was too long to fit in the config space.
    #[error("Fs device tag is too long: len = {0}, max = {}", FS_MAX_TAG_LEN)]
    TagTooLong(usize),
    /// Calling unshare to disassociate FS attributes from parent failed.
    #[error("failed to unshare fs from parent: {0}")]
    UnshareFromParent(SysError),
    /// Error while polling for events.
    #[error("failed to wait for events: {0}")]
    WaitError(SysError),
}

impl From<fuse::Error> for Error {
    fn from(err: fuse::Error) -> Error {
        Error::FuseError(err)
    }
}

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

pub struct Fs {
    cfg: virtio_fs_config,
    fs: Option<PassthroughFs>,
    queue_sizes: Box<[u16]>,
    avail_features: u64,
    acked_features: u64,
    pci_bar: Option<Alloc>,
    tube: Option<Tube>,
    workers: Vec<(Event, thread::JoinHandle<Result<()>>)>,
}

impl Fs {
    pub fn new(
        base_features: u64,
        tag: &str,
        num_workers: usize,
        fs_cfg: passthrough::Config,
        tube: Tube,
    ) -> Result<Fs> {
        if tag.len() > FS_MAX_TAG_LEN {
            return Err(Error::TagTooLong(tag.len()));
        }

        let mut cfg_tag = [0u8; FS_MAX_TAG_LEN];
        cfg_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let cfg = virtio_fs_config {
            tag: cfg_tag,
            num_request_queues: Le32::from(num_workers as u32),
        };

        let fs = PassthroughFs::new(fs_cfg).map_err(Error::CreateFs)?;

        // There is always a high priority queue in addition to the request queues.
        let num_queues = num_workers + 1;

        Ok(Fs {
            cfg,
            fs: Some(fs),
            queue_sizes: vec![QUEUE_SIZE; num_queues].into_boxed_slice(),
            avail_features: base_features,
            acked_features: 0,
            pci_bar: None,
            tube: Some(tube),
            workers: Vec::with_capacity(num_workers + 1),
        })
    }

    fn stop_workers(&mut self) {
        for (kill_evt, handle) in mem::take(&mut self.workers) {
            if let Err(e) = kill_evt.signal() {
                error!("failed to kill virtio-fs worker thread: {}", e);
                continue;
            }

            // Only wait on the child thread if we were able to send it a kill event.
            match handle.join() {
                Ok(r) => {
                    if let Err(e) = r {
                        error!("virtio-fs worker thread exited with error: {}", e)
                    }
                }
                Err(e) => error!("virtio-fs worker thread panicked: {:?}", e),
            }
        }
    }
}

impl VirtioDevice for Fs {
    fn keep_rds(&self) -> Vec<RawDescriptor> {
        let mut fds = self
            .fs
            .as_ref()
            .map(PassthroughFs::keep_rds)
            .unwrap_or_else(Vec::new);
        if let Some(rd) = self.tube.as_ref().map(|s| s.as_raw_descriptor()) {
            fds.push(rd);
        }

        fds
    }

    fn device_type(&self) -> DeviceType {
        DeviceType::Fs
    }

    fn queue_max_sizes(&self) -> &[u16] {
        &self.queue_sizes
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, mut v: u64) {
        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("virtio_fs got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn activate(
        &mut self,
        guest_mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() != self.queue_sizes.len() || queue_evts.len() != self.queue_sizes.len() {
            return;
        }

        let fs = self.fs.take().expect("missing file system implementation");
        let use_dax = fs.cfg().use_dax;

        let server = Arc::new(Server::new(fs));
        let socket = self.tube.take().expect("missing mapping socket");
        let mut slot = 0;

        // Set up shared memory for DAX.
        // TODO(b/176129399): Remove cfg! once DAX is supported on ARM.
        if cfg!(any(target_arch = "x86", target_arch = "x86_64")) && use_dax {
            // Create the shared memory region now before we start processing requests.
            let request = FsMappingRequest::AllocateSharedMemoryRegion(
                self.pci_bar.as_ref().cloned().expect("No pci_bar"),
            );
            socket
                .send(&request)
                .expect("failed to send allocation message");
            slot = match socket.recv() {
                Ok(VmResponse::RegisterMemory { pfn: _, slot }) => slot,
                Ok(VmResponse::Err(e)) => panic!("failed to allocate shared memory region: {}", e),
                r => panic!(
                    "unexpected response to allocate shared memory region: {:?}",
                    r
                ),
            };
        }

        let socket = Arc::new(Mutex::new(socket));
        let mut watch_resample_event = true;
        for (idx, (queue, evt)) in queues.into_iter().zip(queue_evts.into_iter()).enumerate() {
            let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e)))
            {
                Ok(v) => v,
                Err(e) => {
                    error!("fs: failed creating kill Event pair: {}", e);
                    self.stop_workers();
                    return;
                }
            };

            let mem = guest_mem.clone();
            let server = server.clone();
            let irq = interrupt.clone();
            let socket = Arc::clone(&socket);

            let worker_result =
                thread::Builder::new()
                    .name(format!("v_fs:{idx}"))
                    .spawn(move || {
                        let mut worker = Worker::new(mem, queue, server, irq, socket, slot);
                        worker.run(evt, kill_evt, watch_resample_event)
                    });

            if watch_resample_event {
                watch_resample_event = false;
            }

            match worker_result {
                Ok(worker) => self.workers.push((self_kill_evt, worker)),
                Err(e) => {
                    error!("fs: failed to spawn virtio_fs worker: {}", e);
                    self.stop_workers();
                    return;
                }
            }
        }
    }

    fn get_device_bars(&mut self, address: PciAddress) -> Vec<PciBarConfiguration> {
        if self.fs.as_ref().map_or(false, |fs| !fs.cfg().use_dax) {
            return vec![];
        }

        self.pci_bar = Some(Alloc::PciBar {
            bus: address.bus,
            dev: address.dev,
            func: address.func,
            bar: FS_BAR_NUM,
        });

        vec![PciBarConfiguration::new(
            FS_BAR_NUM as usize,
            FS_BAR_SIZE,
            PciBarRegionType::Memory64BitRegion,
            PciBarPrefetchable::Prefetchable,
        )]
    }

    fn get_device_caps(&self) -> Vec<Box<dyn PciCapability>> {
        if self.fs.as_ref().map_or(false, |fs| !fs.cfg().use_dax) {
            return vec![];
        }

        vec![Box::new(VirtioPciShmCap::new(
            PciCapabilityType::SharedMemoryConfig,
            FS_BAR_NUM,
            FS_BAR_OFFSET,
            FS_BAR_SIZE,
            VIRTIO_FS_SHMCAP_ID_CACHE,
        ))]
    }
}

impl Suspendable for Fs {}

impl Drop for Fs {
    fn drop(&mut self) {
        self.stop_workers()
    }
}
<commit_msg>devices: virtio: fs: Include tag name in worker thread name<commit_after>// Copyright 2019 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::io;
use std::mem;
use std::sync::Arc;
use std::thread;

use base::error;
use base::warn;
use base::AsRawDescriptor;
use base::Error as SysError;
use base::Event;
use base::RawDescriptor;
use base::Tube;
use data_model::DataInit;
use data_model::Le32;
use remain::sorted;
use resources::Alloc;
use sync::Mutex;
use thiserror::Error;
use vm_control::FsMappingRequest;
use vm_control::VmResponse;
use vm_memory::GuestMemory;

use crate::pci::PciAddress;
use crate::pci::PciBarConfiguration;
use crate::pci::PciBarPrefetchable;
use crate::pci::PciBarRegionType;
use crate::pci::PciCapability;
use crate::virtio::copy_config;
use crate::virtio::DescriptorError;
use crate::virtio::DeviceType;
use crate::virtio::Interrupt;
use crate::virtio::PciCapabilityType;
use crate::virtio::Queue;
use crate::virtio::VirtioDevice;
use crate::virtio::VirtioPciShmCap;
use crate::Suspendable;

mod caps;
mod multikey;
pub mod passthrough;
mod read_dir;
mod worker;

use fuse::Server;
use passthrough::PassthroughFs;
pub use worker::process_fs_queue;
use worker::Worker;

// The fs device does not have a fixed number of queues.
pub const QUEUE_SIZE: u16 = 1024;

const FS_BAR_NUM: u8 = 4;
const FS_BAR_OFFSET: u64 = 0;
const FS_BAR_SIZE: u64 = 1 << 33;

/// Defined in kernel/include/uapi/linux/virtio_fs.h.
const VIRTIO_FS_SHMCAP_ID_CACHE: u8 = 0;

/// The maximum allowable length of the tag used to identify a specific virtio-fs device.
pub const FS_MAX_TAG_LEN: usize = 36;

/// kernel/include/uapi/linux/virtio_fs.h
#[repr(C, packed)]
#[derive(Clone, Copy)]
pub struct virtio_fs_config {
    /// Filesystem name (UTF-8, not NUL-terminated, padded with NULs)
    pub tag: [u8; FS_MAX_TAG_LEN],
    /// Number of request queues
    pub num_request_queues: Le32,
}

// Safe because all members are plain old data and any value is valid.
unsafe impl DataInit for virtio_fs_config {}

/// Errors that may occur during the creation or operation of an Fs device.
#[sorted]
#[derive(Error, Debug)]
pub enum Error {
    /// Failed to create the file system.
    #[error("failed to create file system: {0}")]
    CreateFs(io::Error),
    /// Creating WaitContext failed.
    #[error("failed to create WaitContext: {0}")]
    CreateWaitContext(SysError),
    /// Error happened in FUSE.
    #[error("fuse error: {0}")]
    FuseError(fuse::Error),
    /// Failed to get the securebits for the worker thread.
    #[error("failed to get securebits for the worker thread: {0}")]
    GetSecurebits(SysError),
    /// The `len` field of the header is too small.
    #[error("DescriptorChain is invalid: {0}")]
    InvalidDescriptorChain(DescriptorError),
    /// A request is missing readable descriptors.
    #[error("request does not have any readable descriptors")]
    NoReadableDescriptors,
    /// A request is missing writable descriptors.
    #[error("request does not have any writable descriptors")]
    NoWritableDescriptors,
    /// Error while reading from the virtio queue's Event.
    #[error("failed to read from virtio queue Event: {0}")]
    ReadQueueEvent(SysError),
    /// Failed to set the securebits for the worker thread.
    #[error("failed to set securebits for the worker thread: {0}")]
    SetSecurebits(SysError),
    /// Failed to signal the virio used queue.
    #[error("failed to signal used queue: {0}")]
    SignalUsedQueue(SysError),
    /// The tag for the Fs device was too long to fit in the config space.
    #[error("Fs device tag is too long: len = {0}, max = {}", FS_MAX_TAG_LEN)]
    TagTooLong(usize),
    /// Calling unshare to disassociate FS attributes from parent failed.
    #[error("failed to unshare fs from parent: {0}")]
    UnshareFromParent(SysError),
    /// Error while polling for events.
    #[error("failed to wait for events: {0}")]
    WaitError(SysError),
}

impl From<fuse::Error> for Error {
    fn from(err: fuse::Error) -> Error {
        Error::FuseError(err)
    }
}

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

pub struct Fs {
    cfg: virtio_fs_config,
    tag: String,
    fs: Option<PassthroughFs>,
    queue_sizes: Box<[u16]>,
    avail_features: u64,
    acked_features: u64,
    pci_bar: Option<Alloc>,
    tube: Option<Tube>,
    workers: Vec<(Event, thread::JoinHandle<Result<()>>)>,
}

impl Fs {
    pub fn new(
        base_features: u64,
        tag: &str,
        num_workers: usize,
        fs_cfg: passthrough::Config,
        tube: Tube,
    ) -> Result<Fs> {
        if tag.len() > FS_MAX_TAG_LEN {
            return Err(Error::TagTooLong(tag.len()));
        }

        let mut cfg_tag = [0u8; FS_MAX_TAG_LEN];
        cfg_tag[..tag.len()].copy_from_slice(tag.as_bytes());

        let cfg = virtio_fs_config {
            tag: cfg_tag,
            num_request_queues: Le32::from(num_workers as u32),
        };

        let fs = PassthroughFs::new(fs_cfg).map_err(Error::CreateFs)?;

        // There is always a high priority queue in addition to the request queues.
        let num_queues = num_workers + 1;

        Ok(Fs {
            cfg,
            tag: tag.to_string(),
            fs: Some(fs),
            queue_sizes: vec![QUEUE_SIZE; num_queues].into_boxed_slice(),
            avail_features: base_features,
            acked_features: 0,
            pci_bar: None,
            tube: Some(tube),
            workers: Vec::with_capacity(num_workers + 1),
        })
    }

    fn stop_workers(&mut self) {
        for (kill_evt, handle) in mem::take(&mut self.workers) {
            if let Err(e) = kill_evt.signal() {
                error!("failed to kill virtio-fs worker thread: {}", e);
                continue;
            }

            // Only wait on the child thread if we were able to send it a kill event.
            match handle.join() {
                Ok(r) => {
                    if let Err(e) = r {
                        error!("virtio-fs worker thread exited with error: {}", e)
                    }
                }
                Err(e) => error!("virtio-fs worker thread panicked: {:?}", e),
            }
        }
    }
}

impl VirtioDevice for Fs {
    fn keep_rds(&self) -> Vec<RawDescriptor> {
        let mut fds = self
            .fs
            .as_ref()
            .map(PassthroughFs::keep_rds)
            .unwrap_or_else(Vec::new);
        if let Some(rd) = self.tube.as_ref().map(|s| s.as_raw_descriptor()) {
            fds.push(rd);
        }

        fds
    }

    fn device_type(&self) -> DeviceType {
        DeviceType::Fs
    }

    fn queue_max_sizes(&self) -> &[u16] {
        &self.queue_sizes
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, mut v: u64) {
        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let unrequested_features = v & !self.avail_features;
        if unrequested_features != 0 {
            warn!("virtio_fs got unknown feature ack: {:x}", v);

            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.acked_features |= v;
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn activate(
        &mut self,
        guest_mem: GuestMemory,
        interrupt: Interrupt,
        queues: Vec<Queue>,
        queue_evts: Vec<Event>,
    ) {
        if queues.len() != self.queue_sizes.len() || queue_evts.len() != self.queue_sizes.len() {
            return;
        }

        let fs = self.fs.take().expect("missing file system implementation");
        let use_dax = fs.cfg().use_dax;

        let server = Arc::new(Server::new(fs));
        let socket = self.tube.take().expect("missing mapping socket");
        let mut slot = 0;

        // Set up shared memory for DAX.
        // TODO(b/176129399): Remove cfg! once DAX is supported on ARM.
        if cfg!(any(target_arch = "x86", target_arch = "x86_64")) && use_dax {
            // Create the shared memory region now before we start processing requests.
            let request = FsMappingRequest::AllocateSharedMemoryRegion(
                self.pci_bar.as_ref().cloned().expect("No pci_bar"),
            );
            socket
                .send(&request)
                .expect("failed to send allocation message");
            slot = match socket.recv() {
                Ok(VmResponse::RegisterMemory { pfn: _, slot }) => slot,
                Ok(VmResponse::Err(e)) => panic!("failed to allocate shared memory region: {}", e),
                r => panic!(
                    "unexpected response to allocate shared memory region: {:?}",
                    r
                ),
            };
        }

        let socket = Arc::new(Mutex::new(socket));
        let mut watch_resample_event = true;
        for (idx, (queue, evt)) in queues.into_iter().zip(queue_evts.into_iter()).enumerate() {
            let (self_kill_evt, kill_evt) = match Event::new().and_then(|e| Ok((e.try_clone()?, e)))
            {
                Ok(v) => v,
                Err(e) => {
                    error!("fs: failed creating kill Event pair: {}", e);
                    self.stop_workers();
                    return;
                }
            };

            let mem = guest_mem.clone();
            let server = server.clone();
            let irq = interrupt.clone();
            let socket = Arc::clone(&socket);

            let worker_result = thread::Builder::new()
                .name(format!("v_fs:{}:{}", self.tag, idx))
                .spawn(move || {
                    let mut worker = Worker::new(mem, queue, server, irq, socket, slot);
                    worker.run(evt, kill_evt, watch_resample_event)
                });

            if watch_resample_event {
                watch_resample_event = false;
            }

            match worker_result {
                Ok(worker) => self.workers.push((self_kill_evt, worker)),
                Err(e) => {
                    error!("fs: failed to spawn virtio_fs worker: {}", e);
                    self.stop_workers();
                    return;
                }
            }
        }
    }

    fn get_device_bars(&mut self, address: PciAddress) -> Vec<PciBarConfiguration> {
        if self.fs.as_ref().map_or(false, |fs| !fs.cfg().use_dax) {
            return vec![];
        }

        self.pci_bar = Some(Alloc::PciBar {
            bus: address.bus,
            dev: address.dev,
            func: address.func,
            bar: FS_BAR_NUM,
        });

        vec![PciBarConfiguration::new(
            FS_BAR_NUM as usize,
            FS_BAR_SIZE,
            PciBarRegionType::Memory64BitRegion,
            PciBarPrefetchable::Prefetchable,
        )]
    }

    fn get_device_caps(&self) -> Vec<Box<dyn PciCapability>> {
        if self.fs.as_ref().map_or(false, |fs| !fs.cfg().use_dax) {
            return vec![];
        }

        vec![Box::new(VirtioPciShmCap::new(
            PciCapabilityType::SharedMemoryConfig,
            FS_BAR_NUM,
            FS_BAR_OFFSET,
            FS_BAR_SIZE,
            VIRTIO_FS_SHMCAP_ID_CACHE,
        ))]
    }
}

impl Suspendable for Fs {}

impl Drop for Fs {
    fn drop(&mut self) {
        self.stop_workers()
    }
}
<|endoftext|>


<commit_before>import std::io;
import std::io::writer_util;

export write_markdown;

type ctxt = {
    w: io::writer,
    mutable depth: uint
};

fn write_markdown(
    doc: doc::cratedoc,
    writer: io::writer
) {
    let ctxt = {
        w: writer,
        mutable depth: 1u
    };

    write_crate(ctxt, doc);
}

fn write_header(ctxt: ctxt, title: str) {
    let hashes = str::from_chars(vec::init_elt('#', ctxt.depth));
    ctxt.w.write_line(#fmt("%s %s", hashes, title));
}

fn write_crate(
    ctxt: ctxt,
    doc: doc::cratedoc
) {
    ctxt.depth = 1u;
    write_header(ctxt, #fmt("Crate %s", doc.topmod.name));
    write_top_module(ctxt, doc.topmod);
}

fn write_top_module(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod_contents(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    for fndoc in *moddoc.fns {
        write_fn(ctxt, fndoc);
    }

    for moddoc in *moddoc.mods {
        write_mod(ctxt, moddoc);
    }
}

fn write_fn(
    ctxt: ctxt,
    doc: doc::fndoc
) {
    ctxt.depth = 2u;
    write_header(ctxt, #fmt("Function `%s`", doc.name));
    alt doc.brief {
      some(brief) {
        ctxt.w.write_line(brief);
      }
      none. { }
    }
    alt doc.desc {
        some(_d) {
            ctxt.w.write_line("");
            ctxt.w.write_line(_d);
            ctxt.w.write_line("");
        }
        none. { }
    }
    for (arg, desc) in doc.args {
        ctxt.w.write_str("### Argument `" + arg + "`: ");
        ctxt.w.write_str(desc)
    }
    alt doc.return {
      some(doc) {
        alt doc.ty {
          some(ty) {
            ctxt.w.write_line("### Returns `" + ty + "`");
            alt doc.desc {
              some(d) {
                ctxt.w.write_line(d);
              }
              none. { }
            }
          }
          none. { fail "unimplemented"; }
        }
      }
      none. { }
    }
}

#[cfg(test)]
mod tests {
    fn write_markdown_str(
        doc: doc::cratedoc
    ) -> str {
        let buffer = io::mk_mem_buffer();
        let writer = io::mem_buffer_writer(buffer);
        write_markdown(doc, writer);
        ret io::mem_buffer_str(buffer);
    }

    #[test]
    fn write_markdown_should_write_crate_header() {
        let source = "";
        let ast = parse::from_str(source);
        let doc = extract::extract(ast, "belch");
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "# Crate belch\n");
    }

    #[test]
    fn write_markdown_should_write_function_header() {
        let source = "fn func() { }";
        let ast = parse::from_str(source);
        let doc = extract::extract(ast, "");
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "## Function `func`");
    }
}<commit_msg>rustdoc: Add gen::subsection for managing header levels<commit_after>import std::io;
import std::io::writer_util;

export write_markdown;

type ctxt = {
    w: io::writer,
    mutable depth: uint
};

fn write_markdown(
    doc: doc::cratedoc,
    writer: io::writer
) {
    let ctxt = {
        w: writer,
        mutable depth: 1u
    };

    write_crate(ctxt, doc);
}

fn write_header(ctxt: ctxt, title: str) {
    let hashes = str::from_chars(vec::init_elt('#', ctxt.depth));
    ctxt.w.write_line(#fmt("%s %s", hashes, title));
}

fn subsection(ctxt: ctxt, f: fn&()) {
    ctxt.depth += 1u;
    f();
    ctxt.depth -= 1u;
}

fn write_crate(
    ctxt: ctxt,
    doc: doc::cratedoc
) {
    write_header(ctxt, #fmt("Crate %s", doc.topmod.name));
    write_top_module(ctxt, doc.topmod);
}

fn write_top_module(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    write_mod_contents(ctxt, moddoc);
}

fn write_mod_contents(
    ctxt: ctxt,
    moddoc: doc::moddoc
) {
    for fndoc in *moddoc.fns {
        subsection(ctxt) {||
            write_fn(ctxt, fndoc);
        }
    }

    for moddoc in *moddoc.mods {
        subsection(ctxt) {||
            write_mod(ctxt, moddoc);
        }
    }
}

fn write_fn(
    ctxt: ctxt,
    doc: doc::fndoc
) {
    write_header(ctxt, #fmt("Function `%s`", doc.name));
    alt doc.brief {
      some(brief) {
        ctxt.w.write_line(brief);
      }
      none. { }
    }
    alt doc.desc {
        some(_d) {
            ctxt.w.write_line("");
            ctxt.w.write_line(_d);
            ctxt.w.write_line("");
        }
        none. { }
    }
    for (arg, desc) in doc.args {
        ctxt.w.write_str("### Argument `" + arg + "`: ");
        ctxt.w.write_str(desc)
    }
    alt doc.return {
      some(doc) {
        alt doc.ty {
          some(ty) {
            ctxt.w.write_line("### Returns `" + ty + "`");
            alt doc.desc {
              some(d) {
                ctxt.w.write_line(d);
              }
              none. { }
            }
          }
          none. { fail "unimplemented"; }
        }
      }
      none. { }
    }
}

#[cfg(test)]
mod tests {
    fn write_markdown_str(
        doc: doc::cratedoc
    ) -> str {
        let buffer = io::mk_mem_buffer();
        let writer = io::mem_buffer_writer(buffer);
        write_markdown(doc, writer);
        ret io::mem_buffer_str(buffer);
    }

    #[test]
    fn write_markdown_should_write_crate_header() {
        let source = "";
        let ast = parse::from_str(source);
        let doc = extract::extract(ast, "belch");
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "# Crate belch\n");
    }

    #[test]
    fn write_markdown_should_write_function_header() {
        let source = "fn func() { }";
        let ast = parse::from_str(source);
        let doc = extract::extract(ast, "");
        let markdown = write_markdown_str(doc);
        assert str::contains(markdown, "## Function `func`");
    }
}<|endoftext|>


<commit_before>#![cfg(not(windows))]

use screen::Key;
use screen::Key::*;
use std::io::{Read, Write};
use std::fs::{File, OpenOptions};
use std::os::unix::io::AsRawFd;
use std::path::*;
use std::process::Command;
use std::process::Stdio;
use std::iter::repeat;
use std::cmp::min;
use ansi;
use ::NEWLINE;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;

#[allow(non_camel_case_types)] type c_ushort = u16;
#[allow(non_camel_case_types)] type c_int = i32;
#[allow(non_camel_case_types)] type c_ulong = u64;

pub struct Screen {
    tty: Terminal,
    original_stty_state: Vec<u8>,
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut tty = Terminal::open_terminal();
        let current_stty_state = tty.stty(&["-g"]);
        tty.initialize();
        let (cols, rows) = tty.winsize().unwrap();
        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        for _ in 0..visible_choices {
            tty.write(NEWLINE.as_bytes());
        }
        tty.write(&ansi::save_cursor());

        Screen {
            tty: tty,
            original_stty_state: current_stty_state,
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
        }
    }

    fn restore_tty(&mut self) {
        self.tty.stty(&[&String::from_utf8(self.original_stty_state.clone()).unwrap()]);
    }

    fn reset_cursor(&mut self) {
        self.tty.write(&ansi::restore_cursor());
        self.tty.write(&ansi::cursor_up(self.visible_choices));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        self.reset_cursor();
        self.tty.write(&ansi::cursor_right(col));
    }

    pub fn blank_screen(&mut self) {
        self.reset_cursor();
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.visible_choices + 1 {
            self.tty.write(blank_line.as_bytes());
        }
        self.reset_cursor();
    }

    pub fn show_cursor(&mut self) {
        self.tty.write(&ansi::show_cursor());
    }

    pub fn hide_cursor(&mut self) {
        self.tty.write(&ansi::hide_cursor());
    }

    pub fn write(&mut self, s: &str) {
        self.tty.write(s.as_bytes());
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_red(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    // Return all buffered keystrokes, or the next key if buffer is empty.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(byte) = self.tty.input.try_recv() {
            ret.push(byte);
        }
        while ret.is_empty() || Screen::more_bytes_needed(&ret) {
            ret.push(self.tty.input.recv().unwrap());
        }
        Terminal::translate_bytes(ret.clone())
    }

    fn more_bytes_needed(bytes: &Vec<u8>) -> bool {
        if let Err(_) = String::from_utf8(bytes.clone()) {
            return true;
        }
        false
    }
}

impl Drop for Screen {
    fn drop(&mut self) {
        self.tty.write(&ansi::cursor_up(1));
        self.restore_tty();
    }
}

struct Terminal {
    input: Receiver<u8>,
    input_fd: i32,
    output: File,
}

impl Terminal {
    fn open_terminal() -> Terminal {
        let term_path = Path::new("/dev/tty");
        let mut input_file = File::open(&term_path).unwrap();
        let output_file = OpenOptions::new().write(true).open(&term_path).unwrap();
        let input_fd = input_file.as_raw_fd();
        let (tx, rx) = mpsc::channel();
        thread::spawn(move || {
            loop {
                let mut buf = [0];
                if input_file.read(&mut buf).unwrap() != 1 {
                    panic!("Failed to read a single byte from tty");
                }
                tx.send(buf[0]).unwrap();
            }
        });
        Terminal {
            input: rx,
            input_fd: input_fd,
            output: output_file,
        }
    }

    fn initialize(&mut self) {
        self.stty(&["raw", "-echo", "cbreak", "opost"]);
    }

    fn stty(&mut self, args: &[&str]) -> Vec<u8> {
        extern { fn dup2(src: c_int, dst: c_int) -> c_int;  }
        unsafe {
            // This is a hack until a replacement for InheritFd from old_io is available.
            dup2(self.input_fd, 0);
        }
        let mut process = match Command::new("stty")
            .args(args)
            .stdin(Stdio::inherit())
            .stdout(Stdio::piped())
            .spawn()
        {
            Ok(p) => p,
            Err(e) => panic!("Spawn failed: {}", e),
        };

        let exit = process.wait();

        let mut buf = Vec::new();
        if exit.unwrap().success() {
            process.stdout.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            let mut str = String::from_utf8(buf).unwrap();

            // The output from `stty -g` may include a newline, which we have to strip off. Otherwise,
            // when we go to restore the tty, stty (on some platforms) will fail with an "invalid
            // argument" error.
            ::trim(&mut str);

            str.into_bytes()
        } else {
            process.stderr.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            panic!(String::from_utf8(buf).unwrap());
        }
    }

    fn translate_bytes(bytes: Vec<u8>) -> Vec<Key> {
        let chars = String::from_utf8(bytes).unwrap().chars().collect::<Vec<char>>();
        chars.into_iter().map(|c| Terminal::translate_char(c)).collect()
    }

    fn translate_char(c: char) -> Key {
        let numeric_char = c as u32;
        if c == '\r' {
            Enter
        } else if numeric_char == 9 {
            Tab
        } else if numeric_char == 127 {
            Backspace
        } else if numeric_char & 96 == 0 && numeric_char <= 128u32 {
            let c = c as u8;
            Control((c + 96u8) as char)
        } else {
            Char(c)
        }
    }

    fn write(&mut self, s: &[u8]) {
        self.output.write_all(&s).unwrap();
    }

    fn writeln(&mut self, s: &str) {
        self.output.write(s.as_bytes()).unwrap();
        self.output.write("\n".as_bytes()).unwrap();
    }

    fn winsize(&self) -> Option<(u16, u16)> {
        extern { fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int; }
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(self.output.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            Some((size.cols as u16, size.rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    // Travis-CI builds run without a tty, making this test impossible.
    if option_env!("TRAVIS").is_some() {
        // TODO: It should be made obvious from the output that this test was skipped
        return;
    }
    let term = Terminal::open_terminal();
    let (cols, rows) = term.winsize().expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<commit_msg>Remove unnecessary cursor tweak on POSIX<commit_after>#![cfg(not(windows))]

use screen::Key;
use screen::Key::*;
use std::io::{Read, Write};
use std::fs::{File, OpenOptions};
use std::os::unix::io::AsRawFd;
use std::path::*;
use std::process::Command;
use std::process::Stdio;
use std::iter::repeat;
use std::cmp::min;
use ansi;
use ::NEWLINE;

use std::thread;
use std::sync::mpsc::Receiver;
use std::sync::mpsc;

#[allow(non_camel_case_types)] type c_ushort = u16;
#[allow(non_camel_case_types)] type c_int = i32;
#[allow(non_camel_case_types)] type c_ulong = u64;

pub struct Screen {
    tty: Terminal,
    original_stty_state: Vec<u8>,
    pub height: u16,
    pub width: u16,
    pub visible_choices: u16,
    start_line: u16,
}

impl Screen {
    pub fn open_screen() -> Screen {
        let mut tty = Terminal::open_terminal();
        let current_stty_state = tty.stty(&["-g"]);
        tty.initialize();
        let (cols, rows) = tty.winsize().unwrap();
        let visible_choices = min(20, rows - 1);
        let start_line = rows - visible_choices - 1;
        for _ in 0..visible_choices {
            tty.write(NEWLINE.as_bytes());
        }
        tty.write(&ansi::save_cursor());

        Screen {
            tty: tty,
            original_stty_state: current_stty_state,
            height: rows,
            width: cols,
            visible_choices: visible_choices,
            start_line: start_line,
        }
    }

    fn restore_tty(&mut self) {
        self.tty.stty(&[&String::from_utf8(self.original_stty_state.clone()).unwrap()]);
    }

    fn reset_cursor(&mut self) {
        self.tty.write(&ansi::restore_cursor());
        self.tty.write(&ansi::cursor_up(self.visible_choices));
    }

    pub fn move_cursor_to_prompt_line(&mut self, col: u16) {
        self.reset_cursor();
        self.tty.write(&ansi::cursor_right(col));
    }

    pub fn blank_screen(&mut self) {
        self.reset_cursor();
        let blank_line = repeat(' ').take(self.width as usize).collect::<String>();
        for _ in 0..self.visible_choices + 1 {
            self.tty.write(blank_line.as_bytes());
        }
        self.reset_cursor();
    }

    pub fn show_cursor(&mut self) {
        self.tty.write(&ansi::show_cursor());
    }

    pub fn hide_cursor(&mut self) {
        self.tty.write(&ansi::hide_cursor());
    }

    pub fn write(&mut self, s: &str) {
        self.tty.write(s.as_bytes());
    }

    pub fn write_red_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_red(&mut self, s: &str) {
        self.tty.write(&ansi::red());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    pub fn write_inverted(&mut self, s: &str) {
        self.tty.write(&ansi::inverse());
        self.tty.write(s.as_bytes());
        self.tty.write(&ansi::reset());
    }

    // Return all buffered keystrokes, or the next key if buffer is empty.
    pub fn get_buffered_keys(&mut self) -> Vec<Key> {
        let mut ret = Vec::new();
        while let Ok(byte) = self.tty.input.try_recv() {
            ret.push(byte);
        }
        while ret.is_empty() || Screen::more_bytes_needed(&ret) {
            ret.push(self.tty.input.recv().unwrap());
        }
        Terminal::translate_bytes(ret.clone())
    }

    fn more_bytes_needed(bytes: &Vec<u8>) -> bool {
        if let Err(_) = String::from_utf8(bytes.clone()) {
            return true;
        }
        false
    }
}

impl Drop for Screen {
    fn drop(&mut self) {
        self.restore_tty();
    }
}

struct Terminal {
    input: Receiver<u8>,
    input_fd: i32,
    output: File,
}

impl Terminal {
    fn open_terminal() -> Terminal {
        let term_path = Path::new("/dev/tty");
        let mut input_file = File::open(&term_path).unwrap();
        let output_file = OpenOptions::new().write(true).open(&term_path).unwrap();
        let input_fd = input_file.as_raw_fd();
        let (tx, rx) = mpsc::channel();
        thread::spawn(move || {
            loop {
                let mut buf = [0];
                if input_file.read(&mut buf).unwrap() != 1 {
                    panic!("Failed to read a single byte from tty");
                }
                tx.send(buf[0]).unwrap();
            }
        });
        Terminal {
            input: rx,
            input_fd: input_fd,
            output: output_file,
        }
    }

    fn initialize(&mut self) {
        self.stty(&["raw", "-echo", "cbreak", "opost"]);
    }

    fn stty(&mut self, args: &[&str]) -> Vec<u8> {
        extern { fn dup2(src: c_int, dst: c_int) -> c_int;  }
        unsafe {
            // This is a hack until a replacement for InheritFd from old_io is available.
            dup2(self.input_fd, 0);
        }
        let mut process = match Command::new("stty")
            .args(args)
            .stdin(Stdio::inherit())
            .stdout(Stdio::piped())
            .spawn()
        {
            Ok(p) => p,
            Err(e) => panic!("Spawn failed: {}", e),
        };

        let exit = process.wait();

        let mut buf = Vec::new();
        if exit.unwrap().success() {
            process.stdout.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            let mut str = String::from_utf8(buf).unwrap();

            // The output from `stty -g` may include a newline, which we have to strip off. Otherwise,
            // when we go to restore the tty, stty (on some platforms) will fail with an "invalid
            // argument" error.
            ::trim(&mut str);

            str.into_bytes()
        } else {
            process.stderr.as_mut().unwrap().read_to_end(&mut buf).unwrap();
            panic!(String::from_utf8(buf).unwrap());
        }
    }

    fn translate_bytes(bytes: Vec<u8>) -> Vec<Key> {
        let chars = String::from_utf8(bytes).unwrap().chars().collect::<Vec<char>>();
        chars.into_iter().map(|c| Terminal::translate_char(c)).collect()
    }

    fn translate_char(c: char) -> Key {
        let numeric_char = c as u32;
        if c == '\r' {
            Enter
        } else if numeric_char == 9 {
            Tab
        } else if numeric_char == 127 {
            Backspace
        } else if numeric_char & 96 == 0 && numeric_char <= 128u32 {
            let c = c as u8;
            Control((c + 96u8) as char)
        } else {
            Char(c)
        }
    }

    fn write(&mut self, s: &[u8]) {
        self.output.write_all(&s).unwrap();
    }

    fn writeln(&mut self, s: &str) {
        self.output.write(s.as_bytes()).unwrap();
        self.output.write("\n".as_bytes()).unwrap();
    }

    fn winsize(&self) -> Option<(u16, u16)> {
        extern { fn ioctl(fd: c_int, request: c_ulong, ...) -> c_int; }
        #[cfg(any(target_os = "macos", target_os = "freebsd"))]
        const TIOCGWINSZ: c_ulong = 0x40087468;

        #[cfg(any(target_os = "linux", target_os = "android"))]
        const TIOCGWINSZ: c_ulong = 0x00005413;

        #[repr(C)]
        struct TermSize {
            rows: c_ushort,
            cols: c_ushort,
            x: c_ushort,
            y: c_ushort,
        }

        let size = TermSize { rows: 0, cols: 0, x: 0, y: 0 };
        if unsafe { ioctl(self.output.as_raw_fd(), TIOCGWINSZ, &size) } == 0 {
            Some((size.cols as u16, size.rows as u16))
        } else {
            None
        }
    }
}

#[test]
fn winsize_test() {
    // Travis-CI builds run without a tty, making this test impossible.
    if option_env!("TRAVIS").is_some() {
        // TODO: It should be made obvious from the output that this test was skipped
        return;
    }
    let term = Terminal::open_terminal();
    let (cols, rows) = term.winsize().expect("Failed to get window size!");
    // We don't know the window size a priori, but we can at least
    // assert that it is within some kind of sensible range.
    assert!(cols > 40);
    assert!(rows > 40);
    assert!(cols < 1000);
    assert!(rows < 1000);
}
<|endoftext|>


<commit_before>pub use self::{
	browse::*,
	connection::*,
	enumerate_domains::*,
	query_record::*,
	records::Record,
	register::*,
	resolve::*,
	resolve_host::*,
};

mod browse;
mod connection;
mod enumerate_domains;
mod query_record;
mod records;
mod register;
mod resolve;
mod resolve_host;

use crate::dns_consts::{
	Class,
	Type,
};
use std::os::raw::c_char;

/// Purge record from cache
///
/// See [`DNSServiceReconfirmRecord`](https://developer.apple.com/documentation/dnssd/1804726-dnsservicereconfirmrecord).
#[doc(alias = "DNSServiceReconfirmRecord")]
pub fn reconfirm_record(
	interface: crate::interface::Interface,
	fullname: &str,
	rr_type: Type,
	rr_class: Class,
	rdata: &[u8],
) -> ::std::io::Result<()> {
	crate::init();

	let fullname = crate::cstr::CStr::from(&fullname)?;
	crate::inner::reconfirm_record(
		0, // no flags
		interface.into_raw(),
		&fullname,
		rr_type,
		rr_class,
		rdata,
	);

	Ok(())
}

/// Full name consiting of (up to) three parts
pub struct FullName<'a> {
	/// (unescaped) service name (becomes single label in full name)
	pub service: Option<&'a str>,
	/// registration type (valid names don't need escaping)
	pub reg_type: &'a str,
	/// (escaped) domain name (most names don't need escaping)
	pub domain: &'a str,
}

impl<'a> FullName<'a> {
	/// Escape and concatenate all three parts to a full name
	///
	/// See [`DNSServiceConstructFullName`](https://developer.apple.com/documentation/dnssd/1804753-dnsserviceconstructfullname)
	#[doc(alias = "DNSServiceConstructFullName")]
	pub fn construct(&self) -> ::std::io::Result<String> {
		use std::io;

		let service = crate::cstr::NullableCStr::from(&self.service)?;
		let reg_type = crate::cstr::CStr::from(&self.reg_type)?;
		let domain = crate::cstr::CStr::from(&self.domain)?;

		const SIZE: usize = crate::ffi::MAX_DOMAIN_NAME + 200;
		let mut buf: Vec<u8> = Vec::new();
		buf.reserve(SIZE);
		let len = unsafe {
			crate::ffi::DNSServiceConstructFullName(
				buf.as_mut_ptr() as *mut c_char,
				service.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
			)
		};

		if len < 0 {
			return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid input"));
		}

		unsafe {
			buf.set_len(len as usize);
		}

		String::from_utf8(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
	}
}
<commit_msg>FullName::construct: fix wrong handling of ffi return value; result was always empty (fixes #21)<commit_after>pub use self::{
	browse::*,
	connection::*,
	enumerate_domains::*,
	query_record::*,
	records::Record,
	register::*,
	resolve::*,
	resolve_host::*,
};

mod browse;
mod connection;
mod enumerate_domains;
mod query_record;
mod records;
mod register;
mod resolve;
mod resolve_host;

use crate::dns_consts::{
	Class,
	Type,
};
use std::os::raw::c_char;

/// Purge record from cache
///
/// See [`DNSServiceReconfirmRecord`](https://developer.apple.com/documentation/dnssd/1804726-dnsservicereconfirmrecord).
#[doc(alias = "DNSServiceReconfirmRecord")]
pub fn reconfirm_record(
	interface: crate::interface::Interface,
	fullname: &str,
	rr_type: Type,
	rr_class: Class,
	rdata: &[u8],
) -> ::std::io::Result<()> {
	crate::init();

	let fullname = crate::cstr::CStr::from(&fullname)?;
	crate::inner::reconfirm_record(
		0, // no flags
		interface.into_raw(),
		&fullname,
		rr_type,
		rr_class,
		rdata,
	);

	Ok(())
}

/// Full name consiting of (up to) three parts
pub struct FullName<'a> {
	/// (unescaped) service name (becomes single label in full name)
	pub service: Option<&'a str>,
	/// registration type (valid names don't need escaping)
	pub reg_type: &'a str,
	/// (escaped) domain name (most names don't need escaping)
	pub domain: &'a str,
}

impl<'a> FullName<'a> {
	/// Escape and concatenate all three parts to a full name
	///
	/// See [`DNSServiceConstructFullName`](https://developer.apple.com/documentation/dnssd/1804753-dnsserviceconstructfullname)
	#[doc(alias = "DNSServiceConstructFullName")]
	pub fn construct(&self) -> ::std::io::Result<String> {
		use std::io;

		let service = crate::cstr::NullableCStr::from(&self.service)?;
		let reg_type = crate::cstr::CStr::from(&self.reg_type)?;
		let domain = crate::cstr::CStr::from(&self.domain)?;

		const SIZE: usize = crate::ffi::MAX_DOMAIN_NAME;
		let mut buf: Vec<u8> = Vec::new();
		buf.reserve(SIZE);
		let result = unsafe {
			crate::ffi::DNSServiceConstructFullName(
				buf.as_mut_ptr() as *mut c_char,
				service.as_ptr(),
				reg_type.as_ptr(),
				domain.as_ptr(),
			)
		};

		if result != 0 {
			// manual only mentions a single possible error (kDNSServiceErr_BadParam), so we use a static io::Error here for now
			// TODO: convert to our normal `Error` type?
			return Err(io::Error::new(io::ErrorKind::InvalidInput, "BadParam"));
		}

		unsafe {
			// ensure NUL termination (MAX_DOMAIN_NAME includes space for trailing NUL, so content must fit)
			buf.as_mut_ptr().wrapping_add(SIZE).write(0);
			buf.set_len(libc::strlen(buf.as_ptr() as *const libc::c_char));
		};

		String::from_utf8(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
struct Fake {
    replies: Vec<Box<dyn std::any::Any>>,
    responses: std::collections::HashMap<Sink, Vec<Box<dyn std::any::Any>>>,
}

#[cfg(test)]
impl Fake {

    pub fn new() -> Fake {
        Fake {
            replies: Vec::new(),
            responses: std::collections::HashMap::new(),
        }
    }

    pub fn get<R>(&self, index: usize) -> &R
    where
        R: action::Response + 'static
    {
        let reply = &self.replies[index];
        reply.downcast_ref().unwrap()
    }

    pub fn get_sink<R>(&self, sink: Sink, index: usize) -> &R
    where
        R: action::Response + 'static
    {
        let responses = match self.responses.get(&sink) {
            Some(responses) => responses,
            None => panic!("no responses for sink '{:?}'", sink),
        };

        let response = &responses[index];
        response.downcast_ref().unwrap()
    }
}

#[cfg(test)]
impl Session for Fake {

    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response + 'static,
    {
        self.replies.push(Box::new(response));
        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response + 'static,
    {
        let responses = self.responses.entry(sink).or_insert_with(Vec::new);
        responses.push(Box::new(response));

        Ok(())
    }
}
<commit_msg>Improve panic messages of the fake session.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! Utilities for working with sessions.
//!
//! Sessions are created per action request and should be used as the only way
//! to communicate with the server. Sessions are responsible for handling action
//! errors (when something goes wrong) or notifying the server that the action
//! finished (by sending appropriate status signal).
//!
//! They also keep track of various statistics (such as number of transferred
//! bytes, action runtime, etc.) and stop the execution if they exceed limits
//! for a particular request.

mod demand;
mod error;
mod response;
mod sink;

use std::convert::TryInto;

use log::{error, info};

use crate::action;
use crate::message;
pub use self::demand::{Demand, Header, Payload};
pub use self::error::{Error, ParseError, MissingFieldError};
use self::response::{Response, Status};
pub use self::sink::{Sink};

/// A specialized `Result` type for sessions.
pub type Result<T> = std::result::Result<T, Error>;

/// Object associating a session with particular action request.
///
/// This is just a convenience type used to avoid threading large numbers of
/// parameters through different function calls.
///
/// Note that the payload can be quite large object, so it should not be stored
/// with the session itself. The payload should be consumed as soon as the
/// request is dispatched to a particular handler and parsed to a concrete type.
pub struct Task<'s, S: Session> {
    /// A session that this task should be run on.
    pub session: &'s mut S,
    /// Serialized data of the action request.
    pub payload: Payload,
}

impl<'s, S: Session> Task<'s, S> {

    /// Executes the task with a particular action handler.
    pub fn execute<R, H>(self, handler: H) -> Result<()>
    where
        R: action::Request,
        H: FnOnce(&mut S, R) -> Result<()>,
    {
        let request = self.payload.parse()?;
        handler(self.session, request)
    }
}

/// Processes given message, handling all errors.
///
/// This function takes a message from the GRR server and interprets it as an
/// action request. It creates a new session object and dispatches the request
/// to an appropriate action handler. Once the action finishes, the function
/// will send a final status message, informing the server about a success
/// or a failure.
///
/// Note that if action execution fails, this function deals with all the errors
/// by sending appropriate information to the server (if possible), logging them
/// and failing hard if a critical error (e.g. communication failure) occurred.
pub fn handle<M>(message: M)
where
    M: TryInto<Demand, Error=ParseError>,
{
    let demand = match message.try_into() {
        Ok(demand) => {
            info!("requested to execute the '{}' action", demand.action);
            demand
        }
        Err(error) => {
            error!("failed to parse the message: {}", error);
            return;
        }
    };

    let mut session = Action::from_demand(&demand);

    let result = action::dispatch(&demand.action, Task {
        session: &mut session,
        payload: demand.payload,
    });

    if let Err(ref error) = result {
        error!("failed to execute the '{}' action: {}", demand.action, error);
    } else {
        info!("finished executing the '{}' action", demand.action);
    }

    let message = match session.status(result).try_into() {
        Ok(message) => message,
        Err(error) => {
            // If we cannot encode the final status message, there is nothing
            // we can do to notify the server, as status is responsible for
            // reporting errors. We can only log the error and carry on.
            error!("failed to encode status message: {}", error);
            return;
        }
    };

    message::send(message);
}

/// Abstraction for various kinds of sessions.
pub trait Session {
    /// Sends a reply to the flow that call the action.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a message to a particular sink.
    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where R: action::Response + 'static;

    /// Sends a heartbeat signal to the Fleetspeak process.
    fn heartbeat(&mut self) {
        // TODO: Create a real implementation.
    }
}

/// A session type for unrequested action executions.
///
/// Certain kind of actions are executed not only when a server flow decides to
/// do so, but also upon particular kind of events (e.g. the agent's startup).
/// In such cases, when one needs to trigger action execution manually, ad-hoc
/// sessions should be used.
pub struct Adhoc;

impl Session for Adhoc {

    // TODO: Session trait should be probably split into two traits and then
    // make the actions that do not care about the `reply` method implement the
    // simpler one.
    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response,
    {
        error!("attempted to reply to an ad-hoc session, dropping response");
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// A session type for ordinary action requests.
///
/// This is a normal session type that that is associated with some flow on the
/// server. It keeps track of the responses it sends and collects statistics
/// about network and runtime utilization to kill the action if it is needed.
pub struct Action {
    header: Header,
    next_response_id: u64,
}

impl Action {

    /// Constructs a new session for the given `demand` object.
    pub fn from_demand(demand: &Demand) -> Action {
        // Response identifiers that GRR agents use start at 1. Unfortunately,
        // the server uses this assumption (to determine the number of expected
        // responses when status message is received), so we have to follow this
        // behaviour in RRG as well.
        Action {
            header: demand.header.clone(),
            next_response_id: 1,
        }
    }

    /// Wraps an action response to a session-specific response.
    fn wrap<R>(&self, response: R) -> Response<R>
    where
        R: action::Response
    {
        Response {
            session_id: self.header.session_id.clone(),
            request_id: Some(self.header.request_id),
            response_id: Some(self.next_response_id),
            data: response,
        }
    }

    /// Wraps an action result to a session-specific status response.
    ///
    /// Note that this method consumes the session. The reason for this is that
    /// status response should be obtained as the very last response and no
    /// further replies should be possible after that.
    fn status(self, result: Result<()>) -> Status {
        Status {
            session_id: self.header.session_id,
            request_id: self.header.request_id,
            response_id: self.next_response_id,
            result: result,
        }
    }
}

impl Session for Action {

    fn reply<R: action::Response>(&mut self, response: R) -> Result<()> {
        send(self.wrap(response))?;
        self.next_response_id += 1;

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response,
    {
        send(sink.wrap(response))?;

        Ok(())
    }
}

/// Sends a session response to the server.
///
/// Note that this function is not exposed on purpose. Actions should send
/// responses through session objects which introduce a layer of safety. `send`
/// is a low-level utility supposed to be used internally.
fn send<R>(response: Response<R>) -> Result<()>
where
    R: action::Response,
{
    let message = response.try_into()?;
    message::send(message);

    Ok(())
}

#[cfg(test)]
struct Fake {
    replies: Vec<Box<dyn std::any::Any>>,
    responses: std::collections::HashMap<Sink, Vec<Box<dyn std::any::Any>>>,
}

#[cfg(test)]
impl Fake {

    pub fn new() -> Fake {
        Fake {
            replies: Vec::new(),
            responses: std::collections::HashMap::new(),
        }
    }

    pub fn get<R>(&self, index: usize) -> &R
    where
        R: action::Response + 'static
    {
        let reply = match self.replies.get(index) {
            Some(reply) => reply,
            None => panic!("no reply #{}", index),
        };

        reply.downcast_ref().expect("unexpected reply type")
    }

    pub fn get_sink<R>(&self, sink: Sink, index: usize) -> &R
    where
        R: action::Response + 'static
    {
        let responses = match self.responses.get(&sink) {
            Some(responses) => responses,
            None => panic!("no responses for sink '{:?}'", sink),
        };

        let response = match responses.get(index) {
            Some(response) => response,
            None => panic!("no response #{} for sink '{:?}'", index, sink),
        };

        match response.downcast_ref() {
            Some(response) => response,
            None => panic!("unexpected response type for sink '{:?}'", sink),
        }
    }
}

#[cfg(test)]
impl Session for Fake {

    fn reply<R>(&mut self, response: R) -> Result<()>
    where
        R: action::Response + 'static,
    {
        self.replies.push(Box::new(response));
        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> Result<()>
    where
        R: action::Response + 'static,
    {
        let responses = self.responses.entry(sink).or_insert_with(Vec::new);
        responses.push(Box::new(response));

        Ok(())
    }
}
<|endoftext|>