<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::cmp;
use std::io;
use std::io::Write;

use crate::bindgen::config::{Braces, Config, Language};
use crate::bindgen::Bindings;

/// A type of way to format a list.
pub enum ListType<'a> {
    /// Join each adjacent item with a str.
    Join(&'a str),
    /// End each item with a str.
    Cap(&'a str),
}

/// An empty file used for creating a null source writer and measuring line
/// metrics for various code layouts.
pub struct NullFile;
impl Write for NullFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A utility wrapper to write unbuffered data and correctly adjust positions.
struct InnerWriter<'a, 'b: 'a, F: 'a + Write>(&'a mut SourceWriter<'b, F>);

impl<'a, 'b, F: Write> Write for InnerWriter<'a, 'b, F> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let writer = &mut self.0;

        if !writer.line_started {
            for _ in 0..writer.spaces() {
                write!(writer.out, " ").unwrap();
            }
            writer.line_started = true;
            writer.line_length += writer.spaces();
        }

        let written = writer.out.write(buf)?;
        writer.line_length += written;
        writer.max_line_length = cmp::max(writer.max_line_length, writer.line_length);
        Ok(written)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.out.flush()
    }
}

/// A utility writer for generating code easier.
pub struct SourceWriter<'a, F: Write> {
    out: F,
    bindings: &'a Bindings,
    spaces: Vec<usize>,
    line_started: bool,
    line_length: usize,
    line_number: usize,
    max_line_length: usize,
}

pub type MeasureWriter<'a> = SourceWriter<'a, NullFile>;

impl<'a, F: Write> SourceWriter<'a, F> {
    pub fn new(out: F, bindings: &'a Bindings) -> Self {
        SourceWriter {
            out,
            bindings,
            spaces: vec![0],
            line_started: false,
            line_length: 0,
            line_number: 1,
            max_line_length: 0,
        }
    }

    pub fn bindings(&self) -> &Bindings {
        &self.bindings
    }

    /// Takes a function that writes source and returns the maximum line length
    /// written.
    pub fn measure<T>(&self, func: T) -> usize
    where
        T: Fn(&mut MeasureWriter),
    {
        let mut measurer = SourceWriter {
            out: NullFile,
            bindings: self.bindings,
            spaces: self.spaces.clone(),
            line_started: self.line_started,
            line_length: self.line_length,
            line_number: self.line_number,
            max_line_length: self.line_length,
        };

        func(&mut measurer);

        measurer.max_line_length
    }

    fn spaces(&self) -> usize {
        *self.spaces.last().unwrap()
    }

    pub fn push_set_spaces(&mut self, spaces: usize) {
        self.spaces.push(spaces);
    }

    pub fn pop_set_spaces(&mut self) {
        self.pop_tab()
    }

    pub fn line_length_for_align(&self) -> usize {
        if self.line_started {
            self.line_length
        } else {
            self.line_length + self.spaces()
        }
    }

    pub fn push_tab(&mut self) {
        let spaces = self.spaces() - (self.spaces() % self.bindings.config.tab_width)
            + self.bindings.config.tab_width;
        self.spaces.push(spaces);
    }

    pub fn pop_tab(&mut self) {
        assert!(!self.spaces.is_empty());
        self.spaces.pop();
    }

    pub fn new_line(&mut self) {
        self.out
            .write_all(self.bindings.config.line_endings.as_str().as_bytes())
            .unwrap();
        self.line_started = false;
        self.line_length = 0;
        self.line_number += 1;
    }

    pub fn new_line_if_not_start(&mut self) {
        if self.line_number != 1 {
            self.new_line();
        }
    }

    pub fn open_brace(&mut self) {
        match self.bindings.config.language {
            Language::Cxx | Language::C => match self.bindings.config.braces {
                Braces::SameLine => {
                    self.write(" {");
                    self.push_tab();
                    self.new_line();
                }
                Braces::NextLine => {
                    self.new_line();
                    self.write("{");
                    self.push_tab();
                    self.new_line();
                }
            },
            Language::Cython => {
                self.write(":");
                self.new_line();
                self.push_tab();
            }
        }
    }

    pub fn close_brace(&mut self, semicolon: bool) {
        match self.bindings.config.language {
            Language::Cxx | Language::C => {
                self.pop_tab();
                self.new_line();
                if semicolon {
                    self.write("};");
                } else {
                    self.write("}");
                }
            }
            Language::Cython => {
                self.pop_tab();
            }
        }
    }

    pub fn write(&mut self, text: &'static str) {
        write!(self, "{}", text);
    }

    pub fn write_raw_block(&mut self, block: &str) {
        self.line_started = true;
        write!(self, "{}", block);
    }

    pub fn write_fmt(&mut self, fmt: ::std::fmt::Arguments) {
        InnerWriter(self).write_fmt(fmt).unwrap();
    }

    pub fn write_horizontal_source_list<'b, S: Source>(
        &mut self,
        items: &[S],
        list_type: ListType<'b>,
    ) {
        for (i, ref item) in items.iter().enumerate() {
            item.write(&self.bindings.config, self);

            match list_type {
                ListType::Join(text) => {
                    if i != items.len() - 1 {
                        write!(self, "{}", text);
                    }
                }
                ListType::Cap(text) => {
                    write!(self, "{}", text);
                }
            }
        }
    }

    pub fn write_vertical_source_list<'b, S: Source>(
        &mut self,
        items: &[S],
        list_type: ListType<'b>,
    ) {
        let align_length = self.line_length_for_align();
        self.push_set_spaces(align_length);
        for (i, ref item) in items.iter().enumerate() {
            item.write(&self.bindings.config, self);

            match list_type {
                ListType::Join(text) => {
                    if i != items.len() - 1 {
                        write!(self, "{}", text);
                    }
                }
                ListType::Cap(text) => {
                    write!(self, "{}", text);
                }
            }

            if i != items.len() - 1 {
                self.new_line();
            }
        }
        self.pop_tab();
    }
}

pub trait Source {
    fn write<F: Write>(&self, config: &Config, _: &mut SourceWriter<F>);
}
<commit_msg>Minor cleanup to `fn close_brace`<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::cmp;
use std::io;
use std::io::Write;

use crate::bindgen::config::{Braces, Config, Language};
use crate::bindgen::Bindings;

/// A type of way to format a list.
pub enum ListType<'a> {
    /// Join each adjacent item with a str.
    Join(&'a str),
    /// End each item with a str.
    Cap(&'a str),
}

/// An empty file used for creating a null source writer and measuring line
/// metrics for various code layouts.
pub struct NullFile;
impl Write for NullFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        Ok(buf.len())
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// A utility wrapper to write unbuffered data and correctly adjust positions.
struct InnerWriter<'a, 'b: 'a, F: 'a + Write>(&'a mut SourceWriter<'b, F>);

impl<'a, 'b, F: Write> Write for InnerWriter<'a, 'b, F> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let writer = &mut self.0;

        if !writer.line_started {
            for _ in 0..writer.spaces() {
                write!(writer.out, " ").unwrap();
            }
            writer.line_started = true;
            writer.line_length += writer.spaces();
        }

        let written = writer.out.write(buf)?;
        writer.line_length += written;
        writer.max_line_length = cmp::max(writer.max_line_length, writer.line_length);
        Ok(written)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.out.flush()
    }
}

/// A utility writer for generating code easier.
pub struct SourceWriter<'a, F: Write> {
    out: F,
    bindings: &'a Bindings,
    spaces: Vec<usize>,
    line_started: bool,
    line_length: usize,
    line_number: usize,
    max_line_length: usize,
}

pub type MeasureWriter<'a> = SourceWriter<'a, NullFile>;

impl<'a, F: Write> SourceWriter<'a, F> {
    pub fn new(out: F, bindings: &'a Bindings) -> Self {
        SourceWriter {
            out,
            bindings,
            spaces: vec![0],
            line_started: false,
            line_length: 0,
            line_number: 1,
            max_line_length: 0,
        }
    }

    pub fn bindings(&self) -> &Bindings {
        &self.bindings
    }

    /// Takes a function that writes source and returns the maximum line length
    /// written.
    pub fn measure<T>(&self, func: T) -> usize
    where
        T: Fn(&mut MeasureWriter),
    {
        let mut measurer = SourceWriter {
            out: NullFile,
            bindings: self.bindings,
            spaces: self.spaces.clone(),
            line_started: self.line_started,
            line_length: self.line_length,
            line_number: self.line_number,
            max_line_length: self.line_length,
        };

        func(&mut measurer);

        measurer.max_line_length
    }

    fn spaces(&self) -> usize {
        *self.spaces.last().unwrap()
    }

    pub fn push_set_spaces(&mut self, spaces: usize) {
        self.spaces.push(spaces);
    }

    pub fn pop_set_spaces(&mut self) {
        self.pop_tab()
    }

    pub fn line_length_for_align(&self) -> usize {
        if self.line_started {
            self.line_length
        } else {
            self.line_length + self.spaces()
        }
    }

    pub fn push_tab(&mut self) {
        let spaces = self.spaces() - (self.spaces() % self.bindings.config.tab_width)
            + self.bindings.config.tab_width;
        self.spaces.push(spaces);
    }

    pub fn pop_tab(&mut self) {
        assert!(!self.spaces.is_empty());
        self.spaces.pop();
    }

    pub fn new_line(&mut self) {
        self.out
            .write_all(self.bindings.config.line_endings.as_str().as_bytes())
            .unwrap();
        self.line_started = false;
        self.line_length = 0;
        self.line_number += 1;
    }

    pub fn new_line_if_not_start(&mut self) {
        if self.line_number != 1 {
            self.new_line();
        }
    }

    pub fn open_brace(&mut self) {
        match self.bindings.config.language {
            Language::Cxx | Language::C => match self.bindings.config.braces {
                Braces::SameLine => {
                    self.write(" {");
                    self.push_tab();
                    self.new_line();
                }
                Braces::NextLine => {
                    self.new_line();
                    self.write("{");
                    self.push_tab();
                    self.new_line();
                }
            },
            Language::Cython => {
                self.write(":");
                self.new_line();
                self.push_tab();
            }
        }
    }

    pub fn close_brace(&mut self, semicolon: bool) {
        self.pop_tab();
        match self.bindings.config.language {
            Language::Cxx | Language::C => {
                self.new_line();
                if semicolon {
                    self.write("};");
                } else {
                    self.write("}");
                }
            }
            Language::Cython => {}
        }
    }

    pub fn write(&mut self, text: &'static str) {
        write!(self, "{}", text);
    }

    pub fn write_raw_block(&mut self, block: &str) {
        self.line_started = true;
        write!(self, "{}", block);
    }

    pub fn write_fmt(&mut self, fmt: ::std::fmt::Arguments) {
        InnerWriter(self).write_fmt(fmt).unwrap();
    }

    pub fn write_horizontal_source_list<'b, S: Source>(
        &mut self,
        items: &[S],
        list_type: ListType<'b>,
    ) {
        for (i, ref item) in items.iter().enumerate() {
            item.write(&self.bindings.config, self);

            match list_type {
                ListType::Join(text) => {
                    if i != items.len() - 1 {
                        write!(self, "{}", text);
                    }
                }
                ListType::Cap(text) => {
                    write!(self, "{}", text);
                }
            }
        }
    }

    pub fn write_vertical_source_list<'b, S: Source>(
        &mut self,
        items: &[S],
        list_type: ListType<'b>,
    ) {
        let align_length = self.line_length_for_align();
        self.push_set_spaces(align_length);
        for (i, ref item) in items.iter().enumerate() {
            item.write(&self.bindings.config, self);

            match list_type {
                ListType::Join(text) => {
                    if i != items.len() - 1 {
                        write!(self, "{}", text);
                    }
                }
                ListType::Cap(text) => {
                    write!(self, "{}", text);
                }
            }

            if i != items.len() - 1 {
                self.new_line();
            }
        }
        self.pop_tab();
    }
}

pub trait Source {
    fn write<F: Write>(&self, config: &Config, _: &mut SourceWriter<F>);
}
<|endoftext|>


<commit_before>use Printer;
use With;
use XY;
use align::*;
use direction::Direction;
use event::*;

use unicode_width::UnicodeWidthStr;

use utils::{LinesIterator, Row};
use vec::Vec2;
use view::{SizeCache, View, ScrollBase, ScrollStrategy};

/// A simple view showing a fixed text
pub struct TextView {
    content: String,
    rows: Vec<Row>,

    align: Align,

    // If `false`, disable scrolling.
    scrollable: bool,

    // ScrollBase make many scrolling-related things easier
    scrollbase: ScrollBase,
    scroll_strategy: ScrollStrategy,
    last_size: Option<XY<SizeCache>>,
    width: Option<usize>,
}

// If the last character is a newline, strip it.
fn strip_last_newline(content: &mut String) {
    if content.ends_with('\n') {
        content.pop().unwrap();
    }
}

impl TextView {
    /// Creates a new TextView with the given content.
    pub fn new<S: Into<String>>(content: S) -> Self {
        let mut content = content.into();
        strip_last_newline(&mut content);
        TextView {
            content: content,
            rows: Vec::new(),
            scrollable: true,
            scrollbase: ScrollBase::new(),
            scroll_strategy: ScrollStrategy::KeepRow,
            align: Align::top_left(),
            last_size: None,
            width: None,
        }
    }

    /// Enable or disable the view's scrolling capabilities.
    ///
    /// When disabled, the view will never attempt to scroll
    /// (and will always ask for the full height).
    pub fn set_scrollable(&mut self, scrollable: bool) {
        self.scrollable = scrollable;
    }

    /// Enable or disable the view's scrolling capabilities.
    ///
    /// When disabled, the view will never attempt to scroll
    /// (and will always ask for the full height).
    ///
    /// Chainable variant.
    pub fn scrollable(self, scrollable: bool) -> Self {
        self.with(|s| s.set_scrollable(scrollable))
    }

    /// Sets the horizontal alignment for this view.
    pub fn h_align(mut self, h: HAlign) -> Self {
        self.align.h = h;

        self
    }

    /// Sets the vertical alignment for this view.
    pub fn v_align(mut self, v: VAlign) -> Self {
        self.align.v = v;

        self
    }

    /// Sets the alignment for this view.
    pub fn align(mut self, a: Align) -> Self {
        self.align = a;

        self
    }

    /// Center the text horizontally and vertically inside the view.
    pub fn center(mut self) -> Self {
        self.align = Align::center();
        self
    }

    /// Replace the text in this view.
    pub fn set_content<S: Into<String>>(&mut self, content: S) {
        let mut content = content.into();
        strip_last_newline(&mut content);
        self.content = content;
        self.invalidate();
    }

    /// Returns the current text in this view.
    pub fn get_content(&self) -> &str {
        &self.content
    }

    fn is_cache_valid(&self, size: Vec2) -> bool {
        match self.last_size {
            None => false,
            Some(ref last) => last.x.accept(size.x) && last.y.accept(size.y),
        }
    }

    /// Defines the way scrolling is adjusted on content or size change.
    ///
    /// The scroll strategy defines how the scrolling position is adjusted
    /// when the size of the view or the content change.
    ///
    /// It is reset to `ScrollStrategy::KeepRow` whenever the user scrolls
    /// manually.
    pub fn set_scroll_strategy(&mut self, strategy: ScrollStrategy) {
        self.scroll_strategy = strategy;
    }

    /// Defines the way scrolling is adjusted on content or size change.
    ///
    /// Chainable variant.
    pub fn scroll_strategy(self, strategy: ScrollStrategy) -> Self {
        self.with(|s| s.set_scroll_strategy(strategy))
    }

    fn adjust_scroll(&mut self) {
        match self.scroll_strategy {
            ScrollStrategy::StickToTop => self.scrollbase.scroll_top(),
            ScrollStrategy::StickToBottom => self.scrollbase.scroll_bottom(),
            ScrollStrategy::KeepRow => (),
        };
    }

    fn compute_rows(&mut self, size: Vec2) {
        if !self.is_cache_valid(size) {
            self.last_size = None;
            // Recompute

            if size.x == 0 {
                // Nothing we can do at this poing.
                return;
            }

            self.rows = LinesIterator::new(&self.content, size.x).collect();
            let mut scrollbar = 0;
            if self.scrollable && self.rows.len() > size.y {
                scrollbar = 2;
                if size.x < scrollbar {
                    // Again, this is a lost cause.
                    return;
                }

                // If we're too high, include a scrollbar
                self.rows = LinesIterator::new(&self.content,
                                               size.x - scrollbar)
                    .collect();
                if self.rows.is_empty() && !self.content.is_empty() {
                    return;
                }
            }

            // Desired width, including the scrollbar.
            self.width = self.rows
                .iter()
                .map(|row| row.width)
                .max()
                .map(|w| w + scrollbar);

            // Our resulting size.
            // We can't go lower, width-wise.

            let mut my_size = Vec2::new(self.width.unwrap_or(0),
                                        self.rows.len());

            if self.scrollable && my_size.y > size.y {
                my_size.y = size.y;
            }

            // println_stderr!("my: {:?} | si: {:?}", my_size, size);
            self.last_size = Some(SizeCache::build(my_size, size));
            self.scrollbase.set_heights(size.y, self.rows.len());
            self.adjust_scroll();
        }
    }

    // Invalidates the cache, so next call will recompute everything.
    fn invalidate(&mut self) {
        self.last_size = None;
    }
}

impl View for TextView {
    fn draw(&self, printer: &Printer) {

        let h = self.rows.len();
        let offset = self.align.v.get_offset(h, printer.size.y);
        let printer =
            &printer.sub_printer(Vec2::new(0, offset), printer.size, true);

        self.scrollbase.draw(printer, |printer, i| {
            let row = &self.rows[i];
            let text = &self.content[row.start..row.end];
            let l = text.width();
            let x = self.align.h.get_offset(l, printer.size.x);
            printer.print((x, 0), text);
        });
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        if !self.scrollbase.scrollable() {
            return EventResult::Ignored;
        }

        match event {
            Event::Key(Key::Home) => self.scrollbase.scroll_top(),
            Event::Key(Key::End) => self.scrollbase.scroll_bottom(),
            Event::Key(Key::Up) if self.scrollbase.can_scroll_up() => {
                self.scrollbase.scroll_up(1)
            }
            Event::Key(Key::Down) if self.scrollbase
                .can_scroll_down() => self.scrollbase.scroll_down(1),
            Event::Key(Key::PageDown) => self.scrollbase.scroll_down(10),
            Event::Key(Key::PageUp) => self.scrollbase.scroll_up(10),
            _ => return EventResult::Ignored,
        }

        // We just scrolled manually, so reset the scroll strategy.
        self.scroll_strategy = ScrollStrategy::KeepRow;
        EventResult::Consumed(None)
    }

    fn needs_relayout(&self) -> bool {
        self.last_size.is_none()
    }

    fn get_min_size(&mut self, size: Vec2) -> Vec2 {
        self.compute_rows(size);

        // This is what we'd like
        let mut ideal = Vec2::new(self.width.unwrap_or(0), self.rows.len());

        if self.scrollable && ideal.y > size.y {
            ideal.y = size.y;
        }

        ideal
    }

    fn take_focus(&mut self, _: Direction) -> bool {
        self.scrollbase.scrollable()
    }

    fn layout(&mut self, size: Vec2) {
        // Compute the text rows.
        self.compute_rows(size);
        self.scrollbase.set_heights(size.y, self.rows.len());
    }
}
<commit_msg>Adjust the scroll after setting the scroll strategy<commit_after>use Printer;
use With;
use XY;
use align::*;
use direction::Direction;
use event::*;

use unicode_width::UnicodeWidthStr;

use utils::{LinesIterator, Row};
use vec::Vec2;
use view::{SizeCache, View, ScrollBase, ScrollStrategy};

/// A simple view showing a fixed text
pub struct TextView {
    content: String,
    rows: Vec<Row>,

    align: Align,

    // If `false`, disable scrolling.
    scrollable: bool,

    // ScrollBase make many scrolling-related things easier
    scrollbase: ScrollBase,
    scroll_strategy: ScrollStrategy,
    last_size: Option<XY<SizeCache>>,
    width: Option<usize>,
}

// If the last character is a newline, strip it.
fn strip_last_newline(content: &mut String) {
    if content.ends_with('\n') {
        content.pop().unwrap();
    }
}

impl TextView {
    /// Creates a new TextView with the given content.
    pub fn new<S: Into<String>>(content: S) -> Self {
        let mut content = content.into();
        strip_last_newline(&mut content);
        TextView {
            content: content,
            rows: Vec::new(),
            scrollable: true,
            scrollbase: ScrollBase::new(),
            scroll_strategy: ScrollStrategy::KeepRow,
            align: Align::top_left(),
            last_size: None,
            width: None,
        }
    }

    /// Enable or disable the view's scrolling capabilities.
    ///
    /// When disabled, the view will never attempt to scroll
    /// (and will always ask for the full height).
    pub fn set_scrollable(&mut self, scrollable: bool) {
        self.scrollable = scrollable;
    }

    /// Enable or disable the view's scrolling capabilities.
    ///
    /// When disabled, the view will never attempt to scroll
    /// (and will always ask for the full height).
    ///
    /// Chainable variant.
    pub fn scrollable(self, scrollable: bool) -> Self {
        self.with(|s| s.set_scrollable(scrollable))
    }

    /// Sets the horizontal alignment for this view.
    pub fn h_align(mut self, h: HAlign) -> Self {
        self.align.h = h;

        self
    }

    /// Sets the vertical alignment for this view.
    pub fn v_align(mut self, v: VAlign) -> Self {
        self.align.v = v;

        self
    }

    /// Sets the alignment for this view.
    pub fn align(mut self, a: Align) -> Self {
        self.align = a;

        self
    }

    /// Center the text horizontally and vertically inside the view.
    pub fn center(mut self) -> Self {
        self.align = Align::center();
        self
    }

    /// Replace the text in this view.
    pub fn set_content<S: Into<String>>(&mut self, content: S) {
        let mut content = content.into();
        strip_last_newline(&mut content);
        self.content = content;
        self.invalidate();
    }

    /// Returns the current text in this view.
    pub fn get_content(&self) -> &str {
        &self.content
    }

    fn is_cache_valid(&self, size: Vec2) -> bool {
        match self.last_size {
            None => false,
            Some(ref last) => last.x.accept(size.x) && last.y.accept(size.y),
        }
    }

    /// Defines the way scrolling is adjusted on content or size change.
    ///
    /// The scroll strategy defines how the scrolling position is adjusted
    /// when the size of the view or the content change.
    ///
    /// It is reset to `ScrollStrategy::KeepRow` whenever the user scrolls
    /// manually.
    pub fn set_scroll_strategy(&mut self, strategy: ScrollStrategy) {
        self.scroll_strategy = strategy;
        self.adjust_scroll();
    }

    /// Defines the way scrolling is adjusted on content or size change.
    ///
    /// Chainable variant.
    pub fn scroll_strategy(self, strategy: ScrollStrategy) -> Self {
        self.with(|s| s.set_scroll_strategy(strategy))
    }

    fn adjust_scroll(&mut self) {
        match self.scroll_strategy {
            ScrollStrategy::StickToTop => self.scrollbase.scroll_top(),
            ScrollStrategy::StickToBottom => self.scrollbase.scroll_bottom(),
            ScrollStrategy::KeepRow => (),
        };
    }

    fn compute_rows(&mut self, size: Vec2) {
        if !self.is_cache_valid(size) {
            self.last_size = None;
            // Recompute

            if size.x == 0 {
                // Nothing we can do at this poing.
                return;
            }

            self.rows = LinesIterator::new(&self.content, size.x).collect();
            let mut scrollbar = 0;
            if self.scrollable && self.rows.len() > size.y {
                scrollbar = 2;
                if size.x < scrollbar {
                    // Again, this is a lost cause.
                    return;
                }

                // If we're too high, include a scrollbar
                self.rows = LinesIterator::new(&self.content,
                                               size.x - scrollbar)
                    .collect();
                if self.rows.is_empty() && !self.content.is_empty() {
                    return;
                }
            }

            // Desired width, including the scrollbar.
            self.width = self.rows
                .iter()
                .map(|row| row.width)
                .max()
                .map(|w| w + scrollbar);

            // Our resulting size.
            // We can't go lower, width-wise.

            let mut my_size = Vec2::new(self.width.unwrap_or(0),
                                        self.rows.len());

            if self.scrollable && my_size.y > size.y {
                my_size.y = size.y;
            }

            // println_stderr!("my: {:?} | si: {:?}", my_size, size);
            self.last_size = Some(SizeCache::build(my_size, size));
            self.scrollbase.set_heights(size.y, self.rows.len());
            self.adjust_scroll();
        }
    }

    // Invalidates the cache, so next call will recompute everything.
    fn invalidate(&mut self) {
        self.last_size = None;
    }
}

impl View for TextView {
    fn draw(&self, printer: &Printer) {

        let h = self.rows.len();
        let offset = self.align.v.get_offset(h, printer.size.y);
        let printer =
            &printer.sub_printer(Vec2::new(0, offset), printer.size, true);

        self.scrollbase.draw(printer, |printer, i| {
            let row = &self.rows[i];
            let text = &self.content[row.start..row.end];
            let l = text.width();
            let x = self.align.h.get_offset(l, printer.size.x);
            printer.print((x, 0), text);
        });
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        if !self.scrollbase.scrollable() {
            return EventResult::Ignored;
        }

        match event {
            Event::Key(Key::Home) => self.scrollbase.scroll_top(),
            Event::Key(Key::End) => self.scrollbase.scroll_bottom(),
            Event::Key(Key::Up) if self.scrollbase.can_scroll_up() => {
                self.scrollbase.scroll_up(1)
            }
            Event::Key(Key::Down) if self.scrollbase
                .can_scroll_down() => self.scrollbase.scroll_down(1),
            Event::Key(Key::PageDown) => self.scrollbase.scroll_down(10),
            Event::Key(Key::PageUp) => self.scrollbase.scroll_up(10),
            _ => return EventResult::Ignored,
        }

        // We just scrolled manually, so reset the scroll strategy.
        self.scroll_strategy = ScrollStrategy::KeepRow;
        EventResult::Consumed(None)
    }

    fn needs_relayout(&self) -> bool {
        self.last_size.is_none()
    }

    fn get_min_size(&mut self, size: Vec2) -> Vec2 {
        self.compute_rows(size);

        // This is what we'd like
        let mut ideal = Vec2::new(self.width.unwrap_or(0), self.rows.len());

        if self.scrollable && ideal.y > size.y {
            ideal.y = size.y;
        }

        ideal
    }

    fn take_focus(&mut self, _: Direction) -> bool {
        self.scrollbase.scrollable()
    }

    fn layout(&mut self, size: Vec2) {
        // Compute the text rows.
        self.compute_rows(size);
        self.scrollbase.set_heights(size.y, self.rows.len());
    }
}
<|endoftext|>


<commit_before>//! Locating various executables part of a C toolchain.

use std::path::PathBuf;

use rustc_middle::bug;
use rustc_session::Session;
use rustc_target::spec::LinkerFlavor;

/// Tries to infer the path of a binary for the target toolchain from the linker name.
pub(crate) fn get_toolchain_binary(sess: &Session, tool: &str) -> PathBuf {
    let (mut linker, _linker_flavor) = linker_and_flavor(sess);
    let linker_file_name = linker
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or_else(|| sess.fatal("couldn't extract file name from specified linker"));

    if linker_file_name == "ld.lld" {
        if tool != "ld" {
            linker.set_file_name(tool)
        }
    } else {
        let tool_file_name = linker_file_name
            .replace("ld", tool)
            .replace("gcc", tool)
            .replace("clang", tool)
            .replace("cc", tool);

        linker.set_file_name(tool_file_name)
    }

    linker
}

// Adapted from https://github.com/rust-lang/rust/blob/5db778affee7c6600c8e7a177c48282dab3f6292/src/librustc_codegen_ssa/back/link.rs#L848-L931
fn linker_and_flavor(sess: &Session) -> (PathBuf, LinkerFlavor) {
    fn infer_from(
        sess: &Session,
        linker: Option<PathBuf>,
        flavor: Option<LinkerFlavor>,
    ) -> Option<(PathBuf, LinkerFlavor)> {
        match (linker, flavor) {
            (Some(linker), Some(flavor)) => Some((linker, flavor)),
            // only the linker flavor is known; use the default linker for the selected flavor
            (None, Some(flavor)) => Some((
                PathBuf::from(match flavor {
                    LinkerFlavor::Em => {
                        if cfg!(windows) {
                            "emcc.bat"
                        } else {
                            "emcc"
                        }
                    }
                    LinkerFlavor::Gcc => {
                        if cfg!(any(target_os = "solaris", target_os = "illumos")) {
                            // On historical Solaris systems, "cc" may have
                            // been Sun Studio, which is not flag-compatible
                            // with "gcc".  This history casts a long shadow,
                            // and many modern illumos distributions today
                            // ship GCC as "gcc" without also making it
                            // available as "cc".
                            "gcc"
                        } else {
                            "cc"
                        }
                    }
                    LinkerFlavor::Ld => "ld",
                    LinkerFlavor::Msvc => "link.exe",
                    LinkerFlavor::Lld(_) => "lld",
                    LinkerFlavor::PtxLinker => "rust-ptx-linker",
                }),
                flavor,
            )),
            (Some(linker), None) => {
                let stem = linker.file_stem().and_then(|stem| stem.to_str()).unwrap_or_else(|| {
                    sess.fatal("couldn't extract file stem from specified linker")
                });

                let flavor = if stem == "emcc" {
                    LinkerFlavor::Em
                } else if stem == "gcc"
                    || stem.ends_with("-gcc")
                    || stem == "clang"
                    || stem.ends_with("-clang")
                {
                    LinkerFlavor::Gcc
                } else if stem == "ld" || stem == "ld.lld" || stem.ends_with("-ld") {
                    LinkerFlavor::Ld
                } else if stem == "link" || stem == "lld-link" {
                    LinkerFlavor::Msvc
                } else if stem == "lld" || stem == "rust-lld" {
                    LinkerFlavor::Lld(sess.target.lld_flavor)
                } else {
                    // fall back to the value in the target spec
                    sess.target.linker_flavor
                };

                Some((linker, flavor))
            }
            (None, None) => None,
        }
    }

    // linker and linker flavor specified via command line have precedence over what the target
    // specification specifies
    if let Some(ret) = infer_from(sess, sess.opts.cg.linker.clone(), sess.opts.cg.linker_flavor) {
        return ret;
    }

    if let Some(ret) = infer_from(
        sess,
        sess.target.linker.clone().map(PathBuf::from),
        Some(sess.target.linker_flavor),
    ) {
        return ret;
    }

    bug!("Not enough information provided to determine how to invoke the linker");
}
<commit_msg>Auto merge of #79608 - alessandrod:bpf, r=nagisa<commit_after>//! Locating various executables part of a C toolchain.

use std::path::PathBuf;

use rustc_middle::bug;
use rustc_session::Session;
use rustc_target::spec::LinkerFlavor;

/// Tries to infer the path of a binary for the target toolchain from the linker name.
pub(crate) fn get_toolchain_binary(sess: &Session, tool: &str) -> PathBuf {
    let (mut linker, _linker_flavor) = linker_and_flavor(sess);
    let linker_file_name = linker
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or_else(|| sess.fatal("couldn't extract file name from specified linker"));

    if linker_file_name == "ld.lld" {
        if tool != "ld" {
            linker.set_file_name(tool)
        }
    } else {
        let tool_file_name = linker_file_name
            .replace("ld", tool)
            .replace("gcc", tool)
            .replace("clang", tool)
            .replace("cc", tool);

        linker.set_file_name(tool_file_name)
    }

    linker
}

// Adapted from https://github.com/rust-lang/rust/blob/5db778affee7c6600c8e7a177c48282dab3f6292/src/librustc_codegen_ssa/back/link.rs#L848-L931
fn linker_and_flavor(sess: &Session) -> (PathBuf, LinkerFlavor) {
    fn infer_from(
        sess: &Session,
        linker: Option<PathBuf>,
        flavor: Option<LinkerFlavor>,
    ) -> Option<(PathBuf, LinkerFlavor)> {
        match (linker, flavor) {
            (Some(linker), Some(flavor)) => Some((linker, flavor)),
            // only the linker flavor is known; use the default linker for the selected flavor
            (None, Some(flavor)) => Some((
                PathBuf::from(match flavor {
                    LinkerFlavor::Em => {
                        if cfg!(windows) {
                            "emcc.bat"
                        } else {
                            "emcc"
                        }
                    }
                    LinkerFlavor::Gcc => {
                        if cfg!(any(target_os = "solaris", target_os = "illumos")) {
                            // On historical Solaris systems, "cc" may have
                            // been Sun Studio, which is not flag-compatible
                            // with "gcc".  This history casts a long shadow,
                            // and many modern illumos distributions today
                            // ship GCC as "gcc" without also making it
                            // available as "cc".
                            "gcc"
                        } else {
                            "cc"
                        }
                    }
                    LinkerFlavor::Ld => "ld",
                    LinkerFlavor::Msvc => "link.exe",
                    LinkerFlavor::Lld(_) => "lld",
                    LinkerFlavor::PtxLinker => "rust-ptx-linker",
                    LinkerFlavor::BpfLinker => "bpf-linker",
                }),
                flavor,
            )),
            (Some(linker), None) => {
                let stem = linker.file_stem().and_then(|stem| stem.to_str()).unwrap_or_else(|| {
                    sess.fatal("couldn't extract file stem from specified linker")
                });

                let flavor = if stem == "emcc" {
                    LinkerFlavor::Em
                } else if stem == "gcc"
                    || stem.ends_with("-gcc")
                    || stem == "clang"
                    || stem.ends_with("-clang")
                {
                    LinkerFlavor::Gcc
                } else if stem == "ld" || stem == "ld.lld" || stem.ends_with("-ld") {
                    LinkerFlavor::Ld
                } else if stem == "link" || stem == "lld-link" {
                    LinkerFlavor::Msvc
                } else if stem == "lld" || stem == "rust-lld" {
                    LinkerFlavor::Lld(sess.target.lld_flavor)
                } else {
                    // fall back to the value in the target spec
                    sess.target.linker_flavor
                };

                Some((linker, flavor))
            }
            (None, None) => None,
        }
    }

    // linker and linker flavor specified via command line have precedence over what the target
    // specification specifies
    if let Some(ret) = infer_from(sess, sess.opts.cg.linker.clone(), sess.opts.cg.linker_flavor) {
        return ret;
    }

    if let Some(ret) = infer_from(
        sess,
        sess.target.linker.clone().map(PathBuf::from),
        Some(sess.target.linker_flavor),
    ) {
        return ret;
    }

    bug!("Not enough information provided to determine how to invoke the linker");
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::collections::BTreeSet;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::ErrorKind;
use std::fs::{OpenOptions, read_dir};
use std::os::unix::prelude::AsRawFd;
use std::path::{Path, PathBuf};
use std::str::FromStr;

use time::Timespec;
use devicemapper::Device;
use uuid::Uuid;

use types::{Bytes, Sectors};
use engine::{EngineResult, EngineError, ErrorEnum};

use consts::*;
use super::metadata::{StaticHeader, BDA, validate_mda_size, BDA_STATIC_HDR_SECTORS};
use super::engine::DevOwnership;
pub use super::BlockDevSave;

type DevUuid = Uuid;
type PoolUuid = Uuid;

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi as u64);

ioctl!(read blkgetsize64 with 0x12, 114; u64);

pub fn blkdev_size(file: &File) -> EngineResult<Bytes> {
    let mut val: u64 = 0;

    match unsafe { blkgetsize64(file.as_raw_fd(), &mut val) } {
        Err(x) => Err(EngineError::Nix(x)),
        Ok(_) => Ok(Bytes(val)),
    }
}

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
pub fn resolve_devices(paths: &[&Path]) -> io::Result<BTreeSet<Device>> {
    let mut devices = BTreeSet::new();
    for path in paths {
        let dev = try!(Device::from_str(&path.to_string_lossy()));
        devices.insert(dev);
    }
    Ok(devices)
}

/// Find all Stratis Blockdevs.
///
/// Returns a map of pool uuids to maps of blockdev uuids to blockdevs.
pub fn find_all() -> EngineResult<BTreeMap<PoolUuid, BTreeMap<DevUuid, BlockDev>>> {

    /// If a Path refers to a valid Stratis blockdev, return a BlockDev
    /// struct. Otherwise, return None. Return an error if there was
    /// a problem inspecting the device.
    fn setup(devnode: &Path) -> EngineResult<Option<BlockDev>> {
        let dev = try!(Device::from_str(&devnode.to_string_lossy()));

        let mut f = try!(OpenOptions::new()
            .read(true)
            .open(devnode));

        let bda = try!(BDA::load(&mut f));

        Ok(Some(BlockDev {
            dev: dev,
            devnode: devnode.to_owned(),
            bda: bda,
        }))
    }

    let mut pool_map = BTreeMap::new();
    for dir_e in try!(read_dir("/dev")) {
        let devnode = match dir_e {
            Ok(d) => d.path(),
            Err(_) => continue,
        };

        match setup(&devnode) {
            Ok(Some(blockdev)) => {
                pool_map.entry(blockdev.bda.header.pool_uuid)
                    .or_insert_with(BTreeMap::new)
                    .insert(blockdev.bda.header.dev_uuid, blockdev);
            }
            _ => continue,
        };
    }

    Ok(pool_map)
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: BTreeSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<BTreeMap<PathBuf, BlockDev>> {

    /// Gets device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the blockdev's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.path().ok_or_else(|| {
            io::Error::new(ErrorKind::InvalidInput,
                           format!("could not get device node from dev {}", dev.dstr()))
        }));
        let mut f = try!(OpenOptions::new()
            .read(true)
            .write(true)
            .open(&devnode)
            .map_err(|_| {
                io::Error::new(ErrorKind::PermissionDenied,
                               format!("Could not open {}", devnode.display()))
            }));

        let dev_size = try!(blkdev_size(&f));

        let ownership = match StaticHeader::determine_ownership(&mut f) {
            Ok(ownership) => ownership,
            Err(err) => {
                let error_message = format!("{} for device {}", err, devnode.display());
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            }
        };

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            if dev_result.is_err() {
                return Err(dev_result.unwrap_err());
            }
            let (devnode, dev_size, ownership, f) = dev_result.unwrap();
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let error_str = format!("First 8K of {} not zeroed", devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    let mut bds = BTreeMap::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = try!(BDA::initialize(&mut f,
                                       pool_uuid,
                                       &Uuid::new_v4(),
                                       mda_size,
                                       dev_size.sectors()));

        let bd = BlockDev {
            dev: dev,
            devnode: devnode.clone(),
            bda: bda,
        };
        bds.insert(devnode, bd);
    }
    Ok(bds)
}

#[derive(Debug)]
pub struct BlockDev {
    pub dev: Device,
    pub devnode: PathBuf,
    pub bda: BDA,
}

impl BlockDev {
    pub fn to_save(&self) -> BlockDevSave {
        BlockDevSave {
            devnode: self.devnode.clone(),
            total_size: self.bda.header.blkdev_size,
        }
    }

    pub fn wipe_metadata(&mut self) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        BDA::wipe(&mut f)
    }

    /// Get the "x:y" device string for this blockdev
    pub fn dstr(&self) -> String {
        self.dev.dstr()
    }

    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        try!(self.bda.save_state(time, metadata, &mut f));

        Ok(())
    }

    /// List the available-for-upper-layer-use range in this blockdev.
    pub fn avail_range(&self) -> (Sectors, Sectors) {
        let start = BDA_STATIC_HDR_SECTORS + self.bda.header.mda_size +
                    self.bda.header.reserved_size;
        // Blockdev size is at least MIN_DEV_SIZE, so this can fail only if
        // size of metadata area exceeds 1 GiB. Initial metadata area size
        // is 4 MiB.
        assert!(start <= self.bda.header.blkdev_size);
        let length = self.bda.header.blkdev_size - start;
        (start, length)
    }
}
<commit_msg>squash with wait<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::collections::BTreeSet;
use std::collections::BTreeMap;
use std::fs::File;
use std::io::ErrorKind;
use std::fs::{OpenOptions, read_dir};
use std::os::unix::prelude::AsRawFd;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use time::Timespec;
use devicemapper::Device;
use uuid::Uuid;

use types::{Bytes, Sectors};
use engine::{EngineResult, EngineError, ErrorEnum};

use consts::*;
use super::metadata::{StaticHeader, BDA, validate_mda_size, BDA_STATIC_HDR_SECTORS};
use super::engine::DevOwnership;
pub use super::BlockDevSave;

type DevUuid = Uuid;
type PoolUuid = Uuid;

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi as u64);

ioctl!(read blkgetsize64 with 0x12, 114; u64);

pub fn blkdev_size(file: &File) -> EngineResult<Bytes> {
    let mut val: u64 = 0;

    match unsafe { blkgetsize64(file.as_raw_fd(), &mut val) } {
        Err(x) => Err(EngineError::Nix(x)),
        Ok(_) => Ok(Bytes(val)),
    }
}

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
pub fn resolve_devices(paths: &[&Path]) -> io::Result<BTreeSet<Device>> {
    let mut devices = BTreeSet::new();
    for path in paths {
        let dev = try!(Device::from_str(&path.to_string_lossy()));
        devices.insert(dev);
    }
    Ok(devices)
}

/// Find all Stratis Blockdevs.
///
/// Returns a map of pool uuids to maps of blockdev uuids to blockdevs.
pub fn find_all() -> EngineResult<BTreeMap<PoolUuid, BTreeMap<DevUuid, BlockDev>>> {

    /// If a Path refers to a valid Stratis blockdev, return a BlockDev
    /// struct. Otherwise, return None. Return an error if there was
    /// a problem inspecting the device.
    fn setup(devnode: &Path) -> EngineResult<Option<BlockDev>> {
        let dev = try!(Device::from_str(&devnode.to_string_lossy()));

        let mut f = try!(OpenOptions::new()
            .read(true)
            .open(devnode));

        let bda = try!(BDA::load(&mut f));

        Ok(Some(BlockDev {
            dev: dev,
            devnode: devnode.to_owned(),
            bda: bda,
        }))
    }

    let mut pool_map = BTreeMap::new();
    for dir_e in try!(read_dir("/dev")) {
        let devnode = match dir_e {
            Ok(d) => d.path(),
            Err(_) => continue,
        };

        match setup(&devnode) {
            Ok(Some(blockdev)) => {
                pool_map.entry(blockdev.bda.header.pool_uuid)
                    .or_insert_with(BTreeMap::new)
                    .insert(blockdev.bda.header.dev_uuid, blockdev);
            }
            _ => continue,
        };
    }

    Ok(pool_map)
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: BTreeSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<BTreeMap<PathBuf, BlockDev>> {

    /// Gets device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the blockdev's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.path().ok_or_else(|| {
            io::Error::new(ErrorKind::InvalidInput,
                           format!("could not get device node from dev {}", dev.dstr()))
        }));
        let mut f = try!(OpenOptions::new()
            .read(true)
            .write(true)
            .open(&devnode)
            .map_err(|_| {
                io::Error::new(ErrorKind::PermissionDenied,
                               format!("Could not open {}", devnode.display()))
            }));

        let dev_size = try!(blkdev_size(&f));

        let ownership = match StaticHeader::determine_ownership(&mut f) {
            Ok(ownership) => ownership,
            Err(err) => {
                let error_message = format!("{} for device {}", err, devnode.display());
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            }
        };

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            if dev_result.is_err() {
                return Err(dev_result.unwrap_err());
            }
            let (devnode, dev_size, ownership, f) = dev_result.unwrap();
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let error_str = format!("First 8K of {} not zeroed", devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    let mut bds = BTreeMap::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = try!(BDA::initialize(&mut f,
                                       pool_uuid,
                                       &Uuid::new_v4(),
                                       mda_size,
                                       dev_size.sectors()));

        let bd = BlockDev {
            dev: dev,
            devnode: devnode.clone(),
            bda: bda,
        };
        bds.insert(devnode, bd);
    }
    Ok(bds)
}

#[derive(Debug)]
pub struct BlockDev {
    pub dev: Device,
    pub devnode: PathBuf,
    pub bda: BDA,
}

impl BlockDev {
    pub fn to_save(&self) -> BlockDevSave {
        BlockDevSave {
            devnode: self.devnode.clone(),
            total_size: self.bda.header.blkdev_size,
        }
    }

    pub fn wipe_metadata(&mut self) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        BDA::wipe(&mut f)
    }

    /// Get the "x:y" device string for this blockdev
    pub fn dstr(&self) -> String {
        self.dev.dstr()
    }

    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        try!(self.bda.save_state(time, metadata, &mut f));

        Ok(())
    }

    /// List the available-for-upper-layer-use range in this blockdev.
    pub fn avail_range(&self) -> (Sectors, Sectors) {
        let start = BDA_STATIC_HDR_SECTORS + self.bda.header.mda_size +
                    self.bda.header.reserved_size;
        // Blockdev size is at least MIN_DEV_SIZE, so this can fail only if
        // size of metadata area exceeds 1 GiB. Initial metadata area size
        // is 4 MiB.
        assert!(start <= self.bda.header.blkdev_size);
        let length = self.bda.header.blkdev_size - start;
        (start, length)
    }

    /// The /dev/mapper/<name> device is not immediately available for use.
    /// TODO: Implement wait for event or poll.
    pub fn wait_for_dm() {
        thread::sleep(Duration::from_millis(500))
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix Windows regression<commit_after><|endoftext|>


<commit_before>#![allow(dead_code)]

// NOTE: Do not rustfmt this file. The expressions are specifically placed on a
// single line for testing purposes.

use std::fs::File;
use std::io::Read;

#[test]
fn simple_method_call() {
    let _ascii = "a".is_ascii();
    let _ascii = "a"
        .is_ascii();

    let _equal = "Ferrös".eq_ignore_ascii_case("FERRöS");
    let _equal = "Ferrös"
        .eq_ignore_ascii_case("FERRöS");

    let _is_empty = Vec::<i32>::new().is_empty();
    let _is_empty = Vec::<i32>::new()
        .is_empty();

    let _is_4_in_list = vec![1, 2, 3, 4].contains(&4);
    let _is_4_in_list = vec![1, 2, 3, 4]
        .contains(&4); // FIXME: This line is not covered.

    let _is_4_in_list = Vec::new().contains(&4);
    let _is_4_in_list = Vec::new()
        .contains(&4); // FIXME: This line is not covered.
}

#[test]
fn simple_method_call_with_try() -> Result<(), Box<dyn std::error::Error>> {
    let _num: i32 = "123".parse()?;

    let mut test = File::open("Cargo.toml")?;
    let mut string = String::new();
    test.read_to_string(&mut string)?;

    // Separated lines
    let _num: i32 = "123"
        .parse()?;

    let mut test = File::open("Cargo.toml")?;
    let mut string = String::new();
    test
        .read_to_string(&mut string)?; // FIXME: This line is not covered.

    Ok(())
}

#[test]
fn method_call_chain() {
    let num: &str = "123".clone().trim();
    let _truth = num.contains("23").to_string();
    let _is_4_in_list: bool = vec![1, 12, 4, 7, 5, 5, 0, 3].as_slice().iter().filter(|&x| x%2 != 0).map(|&x| x).collect::<Vec<_>>().contains(&&4);

    // Separated lines
    let num: &str = "123"
        .clone()
        .trim();

    let _truth = num
        .contains("23")
        .to_string();

    // Progressively building more complex chains for testing...
    let _is_4_in_list = Vec::new()
        .iter()
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .iter()
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string(); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string()
        .clone();

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string()
        .clone()
        .trim(); // FIXME: This line is not covered.
}

#[test]
fn method_call_chain_with_try() -> Result<(), Box<dyn std::error::Error>> {
    let _num: i32 = "123".clone().parse()?;
    let _num: i32 = "123"
        .clone()
        .parse()?;

    let _num: i32 = "123".parse::<i32>()?.clone();
    let _num: i32 = "123"
        .parse::<i32>()?
        .clone();

    let _num: i32 = "123".clone().parse::<i32>()?.clone();
    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .clone();

    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .reverse_bits()
        .clone()
        .reverse_bits()
        .clone()
        .to_string()
        .parse::<i32>()?; // FIXME: This line is not covered.

    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .reverse_bits()
        .clone()
        .reverse_bits()
        .clone()
        .to_string()
        .parse::<i32>()?
        .clone(); // FIXME: This line is not covered.

    Ok(())
}

// This test ensures that ".collect::<_>()?" should work okay under any
// circumstances that it is used.
#[test]
fn method_call_with_collect_try() -> Result<(), Box<dyn std::error::Error>> {
    let strings = vec!["93", "18"];

    let _collect_no_try = strings.iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>();
    let _collect_no_try = strings
        .iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>();

    let _collect_with_try = strings.iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>()?;
    let _collect_with_try = strings
        .iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>()?; // FIXME: This line is not covered.

    let _collect_no_try = vec!["93", "18"].into_iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>();
    let _collect_no_try = vec!["93", "18"]
        .into_iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>();

    let _collect_with_try = vec!["93", "18"].into_iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>()?;
    let _collect_with_try = vec!["93", "18"]
        .into_iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>()?;

    #[cfg_attr(tarpaulin, skip)]
    Ok(())
}
<commit_msg>Remove attribute on expression<commit_after>#![allow(dead_code)]

// NOTE: Do not rustfmt this file. The expressions are specifically placed on a
// single line for testing purposes.

use std::fs::File;
use std::io::Read;

#[test]
fn simple_method_call() {
    let _ascii = "a".is_ascii();
    let _ascii = "a"
        .is_ascii();

    let _equal = "Ferrös".eq_ignore_ascii_case("FERRöS");
    let _equal = "Ferrös"
        .eq_ignore_ascii_case("FERRöS");

    let _is_empty = Vec::<i32>::new().is_empty();
    let _is_empty = Vec::<i32>::new()
        .is_empty();

    let _is_4_in_list = vec![1, 2, 3, 4].contains(&4);
    let _is_4_in_list = vec![1, 2, 3, 4]
        .contains(&4); // FIXME: This line is not covered.

    let _is_4_in_list = Vec::new().contains(&4);
    let _is_4_in_list = Vec::new()
        .contains(&4); // FIXME: This line is not covered.
}

#[test]
fn simple_method_call_with_try() -> Result<(), Box<dyn std::error::Error>> {
    let _num: i32 = "123".parse()?;

    let mut test = File::open("Cargo.toml")?;
    let mut string = String::new();
    test.read_to_string(&mut string)?;

    // Separated lines
    let _num: i32 = "123"
        .parse()?;

    let mut test = File::open("Cargo.toml")?;
    let mut string = String::new();
    test
        .read_to_string(&mut string)?; // FIXME: This line is not covered.

    Ok(())
}

#[test]
fn method_call_chain() {
    let num: &str = "123".clone().trim();
    let _truth = num.contains("23").to_string();
    let _is_4_in_list: bool = vec![1, 12, 4, 7, 5, 5, 0, 3].as_slice().iter().filter(|&x| x%2 != 0).map(|&x| x).collect::<Vec<_>>().contains(&&4);

    // Separated lines
    let num: &str = "123"
        .clone()
        .trim();

    let _truth = num
        .contains("23")
        .to_string();

    // Progressively building more complex chains for testing...
    let _is_4_in_list = Vec::new()
        .iter()
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .iter()
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string(); // FIXME: This line is not covered.

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string()
        .clone();

    let _is_4_in_list = vec![1, 2, 3, 4]
        .as_slice()
        .iter()
        .filter(|&x| x%2 != 0)
        .map(|&x| x)
        .collect::<Vec<_>>()
        .contains(&&4)
        .to_string()
        .clone()
        .trim(); // FIXME: This line is not covered.
}

#[test]
fn method_call_chain_with_try() -> Result<(), Box<dyn std::error::Error>> {
    let _num: i32 = "123".clone().parse()?;
    let _num: i32 = "123"
        .clone()
        .parse()?;

    let _num: i32 = "123".parse::<i32>()?.clone();
    let _num: i32 = "123"
        .parse::<i32>()?
        .clone();

    let _num: i32 = "123".clone().parse::<i32>()?.clone();
    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .clone();

    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .reverse_bits()
        .clone()
        .reverse_bits()
        .clone()
        .to_string()
        .parse::<i32>()?; // FIXME: This line is not covered.

    let _num: i32 = "123"
        .clone()
        .parse::<i32>()?
        .reverse_bits()
        .clone()
        .reverse_bits()
        .clone()
        .to_string()
        .parse::<i32>()?
        .clone(); // FIXME: This line is not covered.

    Ok(())
}

// This test ensures that ".collect::<_>()?" should work okay under any
// circumstances that it is used.
#[test]
fn method_call_with_collect_try() -> Result<(), Box<dyn std::error::Error>> {
    let strings = vec!["93", "18"];

    let _collect_no_try = strings.iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>();
    let _collect_no_try = strings
        .iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>();

    let _collect_with_try = strings.iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>()?;
    let _collect_with_try = strings
        .iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>()?; // FIXME: This line is not covered.

    let _collect_no_try = vec!["93", "18"].into_iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>();
    let _collect_no_try = vec!["93", "18"]
        .into_iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>();

    let _collect_with_try = vec!["93", "18"].into_iter().map(|s| s.parse::<i32>()).collect::<Result<Vec<_>, _>>()?;
    let _collect_with_try = vec!["93", "18"]
        .into_iter()
        .map(|s| s.parse::<i32>())
        .collect::<Result<Vec<_>, _>>()?;
    Ok(())
}
<|endoftext|>


<commit_before>use main::{WINDOW_WIDTH, WINDOW_HEIGHT, RGB};
use opencl;
use opencl::mem::CLBuffer;

pub struct Scene {
  pub center: [f32; 3],
  pub radius: f32,
  pub camera: [f32; 3],
}

impl Scene {
  pub fn render(&self) -> Vec<RGB> {
    let (device, ctx, queue) = opencl::util::create_compute_context().unwrap();

    let len = WINDOW_WIDTH as usize * WINDOW_HEIGHT as usize;

    let output_buffer: CLBuffer<RGB> = ctx.create_buffer(len, opencl::cl::CL_MEM_WRITE_ONLY);

    let program = {
      let ker = format!("
        float toi(
          const float3 eye,
          const float3 look,

          const float3 sphere,
          const float radius)
        {{
          float a = dot(radius, radius);
          float b = 2 * dot(eye - sphere, look);
          float c = dot(sphere, sphere) + dot(eye, eye) - dot(eye, sphere) - radius * radius;

          float d = b*b - 4*a*c;

          if (d < 0) {{
            return -1;
          }}

          return (-sqrt(d) - b) / (2 * a);
        }}

        __kernel void render(
          const float3 obj_center,
          const float obj_radius,

          const float3 eye,

          __global float * output)
        {{
          int W = {};
          int H = {};

          float fov_x = 3.14 / 2;
          float fov_y = 3.14 / 2;

          int i = get_global_id(0);

          float x_pix = i % W;
          float y_pix = i / W;

          float t_x = fov_x * (x_pix / W - 0.5);
          float t_y = fov_y * (y_pix / H - 0.5);

          float c = -cos(t_y);
          float3 ray = {{c*sin(t_x), sin(t_y), c*cos(t_x)}};

          i = i * 3;
          if (toi(eye, ray, obj_center, obj_radius) >= 0) {{
            output[i+0] = 1;
            output[i+1] = 0;
            output[i+2] = 0;
          }} else {{
            output[i+0] = 0;
            output[i+1] = 0;
            output[i+2] = 0;
          }}
        }}",
        WINDOW_WIDTH,
        WINDOW_HEIGHT,
      );
      ctx.create_program_from_source(ker.as_slice())
    };
    program.build(&device).unwrap();

    let kernel = program.create_kernel("render");
    kernel.set_arg(0, &self.center);
    kernel.set_arg(1, &self.radius);
    kernel.set_arg(2, &self.camera);

    // This is sketchy; we "implicitly cast" output_buffer from a CLBuffer<RGB> to a CLBuffer<f32>.
    kernel.set_arg(3, &output_buffer);

    let event = queue.enqueue_async_kernel(&kernel, len, None, ());

    queue.get(&output_buffer, &event)
  }
}
<commit_msg>Fix ray orientation in X direction.<commit_after>use main::{WINDOW_WIDTH, WINDOW_HEIGHT, RGB};
use opencl;
use opencl::mem::CLBuffer;

pub struct Scene {
  pub center: [f32; 3],
  pub radius: f32,
  pub camera: [f32; 3],
}

impl Scene {
  pub fn render(&self) -> Vec<RGB> {
    let (device, ctx, queue) = opencl::util::create_compute_context().unwrap();

    let len = WINDOW_WIDTH as usize * WINDOW_HEIGHT as usize;

    let output_buffer: CLBuffer<RGB> = ctx.create_buffer(len, opencl::cl::CL_MEM_WRITE_ONLY);

    let program = {
      let ker = format!("
        float toi(
          const float3 eye,
          const float3 look,

          const float3 sphere,
          const float radius)
        {{
          float a = dot(radius, radius);
          float b = 2 * dot(eye - sphere, look);
          float c = dot(sphere, sphere) + dot(eye, eye) - dot(eye, sphere) - radius * radius;

          float d = b*b - 4*a*c;

          if (d < 0) {{
            return -1;
          }}

          return (-sqrt(d) - b) / (2 * a);
        }}

        __kernel void render(
          const float3 obj_center,
          const float obj_radius,

          const float3 eye,

          __global float * output)
        {{
          int W = {};
          int H = {};

          float fov_x = 3.14 / 2;
          float fov_y = 3.14 / 2;

          int i = get_global_id(0);

          float x_pix = i % W;
          float y_pix = i / W;

          float t_x = fov_x * (x_pix / W - 0.5);
          float t_y = fov_y * (y_pix / H - 0.5);

          float c = cos(t_y);
          float3 ray = {{c*sin(t_x), sin(t_y), -c*cos(t_x)}};

          i = i * 3;
          if (toi(eye, ray, obj_center, obj_radius) >= 0) {{
            output[i+0] = 1;
            output[i+1] = 0;
            output[i+2] = 0;
          }} else {{
            output[i+0] = 0;
            output[i+1] = 0;
            output[i+2] = 0;
          }}
        }}",
        WINDOW_WIDTH,
        WINDOW_HEIGHT,
      );
      ctx.create_program_from_source(ker.as_slice())
    };
    program.build(&device).unwrap();

    let kernel = program.create_kernel("render");
    kernel.set_arg(0, &self.center);
    kernel.set_arg(1, &self.radius);
    kernel.set_arg(2, &self.camera);

    // This is sketchy; we "implicitly cast" output_buffer from a CLBuffer<RGB> to a CLBuffer<f32>.
    kernel.set_arg(3, &output_buffer);

    let event = queue.enqueue_async_kernel(&kernel, len, None, ());

    queue.get(&output_buffer, &event)
  }
}
<|endoftext|>


<commit_before>use std::cmp::max;
use std::net::SocketAddr;

use mioco;
use mioco::sync::mpsc::channel;
use mioco::tcp::TcpListener;
use mioco::udp::UdpSocket;
use mioco::JoinHandle;
use trust_dns_proto::op::{Edns, Message, OpCode, ResponseCode};
use trust_dns_proto::rr::dnssec::rdata::DNSSECRecordType;
use trust_dns_proto::rr::{DNSClass, RecordType};
use trust_dns_proto::serialize::binary::BinEncodable;

use crate::query::EDNS_VER;
use crate::resolver::RcResolver;
use crate::transport::DnsTransport;
use crate::utils::{AsDisplay, Error, MessageExt, Result, WithTimeout};

/// This function should never fail, otherwise we have a panic
fn handle_request(resolver: &RcResolver, msg: &Message, should_truncate: bool) -> Result<Vec<u8>> {
    let mut ret: Message = msg.new_resp();
    ret.set_recursion_available(true);
    let mut have_dnssec = false;
    if let Some(req_edns) = msg.edns() {
        let mut resp_edns = Edns::new();
        resp_edns.set_version(EDNS_VER);
        have_dnssec = req_edns.dnssec_ok();
        resp_edns.set_dnssec_ok(have_dnssec);
        resp_edns.set_max_payload(max(512, req_edns.max_payload()));
        ret.set_edns(resp_edns);
        if req_edns.version() > EDNS_VER {
            warn!("Got EDNS version {}", req_edns.version());
            ret.set_response_code(ResponseCode::BADVERS);
            return ret.to_bytes().map_err(|e| e.into());
        }
    }
    if msg.queries().len() != 1 || msg.op_code() != OpCode::Query {
        // For simplicity, only support one question
        warn!("Refuse to serve: {:?}", msg);
        ret.set_response_code(ResponseCode::Refused);
        return ret.to_bytes().map_err(|e| e.into());
    }
    if msg.queries()[0].query_class() != DNSClass::IN {
        ret.set_response_code(ResponseCode::NotImp);
        return ret.to_bytes().map_err(|e| e.into());
    }
    match resolver.resolve(&mut ret) {
        Ok(_) => { /* Message is filled with answers */ }
        Err(e) => {
            warn!(
                "Resolver returned error for {}: {:?}",
                msg.queries()[0].as_disp(),
                e
            );
            ret.set_response_code(ResponseCode::ServFail);
        }
    };
    if !have_dnssec {
        ret = ret.strip_dnssec_records();
    } else {
        let is_authenticated = ret
            .answers()
            .iter()
            .chain(ret.name_servers())
            .any(|rec| rec.rr_type() == RecordType::DNSSEC(DNSSECRecordType::RRSIG));
        ret.set_authentic_data(is_authenticated);
    }
    let mut bytes = (ret.to_bytes())?;
    if should_truncate && bytes.len() > (msg.max_payload() as usize) {
        ret.set_truncated(true);
        ret.take_name_servers();
        ret.take_additionals();
        bytes = ret.to_bytes()?;
        while bytes.len() > (msg.max_payload() as usize) {
            let mut answers = ret.take_answers();
            let new_len = answers.len() / 2;
            answers.truncate(new_len);
            ret.insert_answers(answers);
            bytes = ret.to_bytes()?;
        }
    }
    Ok(bytes)
}

fn serve_transport_async<TRecv, TSend, F>(
    mut recv: TRecv,
    mut send: TSend,
    resolver: RcResolver,
    on_error: F,
) -> JoinHandle<()>
where
    TRecv: DnsTransport + Send + 'static,
    TSend: DnsTransport + Send + 'static,
    F: FnOnce(Error) + Send + 'static,
{
    let (sch, rch) = channel::<(Vec<u8>, Option<SocketAddr>)>();
    mioco::spawn(move || loop {
        match rch.recv() {
            Ok((buf, addr)) => {
                match send.send_msg_bytes(&buf, addr.as_ref()) {
                    Ok(_) => {}
                    Err(e) => {
                        warn!("Failed to respond to DNS request: {:?}", e);
                    }
                };
            }
            Err(_) => {
                trace!("Channel is broken");
                return;
            }
        };
    });
    let keepalive = resolver.to_keepalive();
    let resolver_weak = resolver.to_weak();
    mioco::spawn(move || {
        let _ = keepalive;
        let mut resolver_weak = resolver_weak;
        loop {
            let (msg, addr) = match recv.recv_msg(None) {
                Ok(x) => x,
                Err(e) => {
                    on_error(e);
                    return;
                }
            };
            let sch_req = sch.clone();
            let resolver = resolver_weak.upgrade().unwrap();
            resolver_weak = resolver.to_weak();
            mioco::spawn(move || {
                let resp = handle_request(&resolver, &msg, TSend::should_truncate())
                    .expect("handle_request should not return error");
                sch_req
                    .send((resp, addr))
                    .expect("Result pipe should not break here");
            });
        }
    })
}
pub fn serve_tcp(addr: &SocketAddr, resolver: RcResolver) -> Result<JoinHandle<()>> {
    let listener = (TcpListener::bind(addr))?;
    let keepalive = resolver.to_keepalive();
    let resolver_weak = resolver.to_weak();
    Ok(mioco::spawn(move || {
        let _ = keepalive;
        let mut resolver_weak = resolver_weak;
        loop {
            let sock = listener.accept().expect("Failed to accept TCP socket");
            sock.set_nodelay(true).unwrap_or(());
            let sock_clone = sock.try_clone().expect("Failed to clone TCP socket");
            let resolver = resolver_weak.upgrade().unwrap();
            resolver_weak = resolver.to_weak();
            serve_transport_async(
                sock.with_resetting_timeout(*resolver.config.serve.tcp_timeout),
                sock_clone.with_resetting_timeout(*resolver.config.serve.tcp_timeout),
                resolver,
                |e| trace!("Client TCP connection is broken: {:?}", e),
            );
        }
    }))
}
pub fn serve_udp(addr: &SocketAddr, resolver: RcResolver) -> Result<JoinHandle<()>> {
    if addr.ip().is_unspecified() {
        debug!("Binding UDP listener to all available IP addresses");
        let mut handles = Vec::<JoinHandle<()>>::new();
        for iface in pnet_datalink::interfaces() {
            for ipnetwork in iface.ips {
                let ip = ipnetwork.ip();
                if ip.is_unspecified() {
                    continue;
                }
                debug!("Binding to {}", ip);
                match serve_udp(&SocketAddr::new(ip, addr.port()), resolver.clone()) {
                    Err(e) => debug!("Failed to bind to {}: {}", ip, e),
                    Ok(handle) => handles.push(handle),
                }
            }
        }
        if handles.is_empty() {
            return Err(Error::Other("Failed to bind UDP listener".into(), None));
        }
        return Ok(mioco::spawn(move || {
            for h in handles {
                h.join().unwrap();
            }
        }));
    }
    let sock: UdpSocket = (UdpSocket::bound(addr))?;
    let sock_clone = (sock.try_clone())?;
    Ok(serve_transport_async(
        sock,
        sock_clone,
        resolver,
        |e| match e {
            Error::Proto(reason) => warn!("Invalid DNS request: {:?}", reason),
            e2 => panic!("UDP listener is broken: {:?}", e2),
        },
    ))
}
<commit_msg>Return SERVFAIL for broken CNAME chain<commit_after>use std::cmp::max;
use std::net::SocketAddr;

use mioco;
use mioco::sync::mpsc::channel;
use mioco::tcp::TcpListener;
use mioco::udp::UdpSocket;
use mioco::JoinHandle;
use trust_dns_proto::op::{Edns, Message, OpCode, ResponseCode};
use trust_dns_proto::rr::dnssec::rdata::DNSSECRecordType;
use trust_dns_proto::rr::{DNSClass, RecordType};
use trust_dns_proto::serialize::binary::BinEncodable;

use crate::query::EDNS_VER;
use crate::resolver::RcResolver;
use crate::transport::DnsTransport;
use crate::utils::{AsDisplay, Error, MessageExt, Result, WithTimeout};

/// This function should never fail, otherwise we have a panic
fn handle_request(resolver: &RcResolver, msg: &Message, should_truncate: bool) -> Result<Vec<u8>> {
    let mut ret: Message = msg.new_resp();
    ret.set_recursion_available(true);
    let mut have_dnssec = false;
    if let Some(req_edns) = msg.edns() {
        let mut resp_edns = Edns::new();
        resp_edns.set_version(EDNS_VER);
        have_dnssec = req_edns.dnssec_ok();
        resp_edns.set_dnssec_ok(have_dnssec);
        resp_edns.set_max_payload(max(512, req_edns.max_payload()));
        ret.set_edns(resp_edns);
        if req_edns.version() > EDNS_VER {
            warn!("Got EDNS version {}", req_edns.version());
            ret.set_response_code(ResponseCode::BADVERS);
            return ret.to_bytes().map_err(|e| e.into());
        }
    }
    if msg.queries().len() != 1 || msg.op_code() != OpCode::Query {
        // For simplicity, only support one question
        warn!("Refuse to serve: {:?}", msg);
        ret.set_response_code(ResponseCode::Refused);
        return ret.to_bytes().map_err(|e| e.into());
    }
    if msg.queries()[0].query_class() != DNSClass::IN {
        ret.set_response_code(ResponseCode::NotImp);
        return ret.to_bytes().map_err(|e| e.into());
    }
    match resolver.resolve(&mut ret) {
        Ok(_) => { /* Message is filled with answers */ }
        Err(e) => {
            warn!(
                "Resolver returned error for {}: {:?}",
                msg.queries()[0].as_disp(),
                e
            );
            ret.set_response_code(ResponseCode::ServFail);
        }
    };
    if !have_dnssec {
        ret = ret.strip_dnssec_records();
    } else {
        let is_authenticated = ret
            .answers()
            .iter()
            .chain(ret.name_servers())
            .any(|rec| rec.rr_type() == RecordType::DNSSEC(DNSSECRecordType::RRSIG));
        ret.set_authentic_data(is_authenticated);
    }
    let query = &msg.queries()[0];
    if ret.response_code() == ResponseCode::NoError
        && !ret
            .answers()
            .iter()
            .any(|x| x.name() == query.name() && x.rr_type() == query.query_type())
        && ret
            .answers()
            .iter()
            .any(|x| x.name() == query.name() && x.rr_type() == RecordType::CNAME)
    {
        warn!("Broken CNAME chain for {}", query.as_disp());
        ret.set_response_code(ResponseCode::ServFail);
    }
    let mut bytes = (ret.to_bytes())?;
    if should_truncate && bytes.len() > (msg.max_payload() as usize) {
        ret.set_truncated(true);
        ret.take_name_servers();
        ret.take_additionals();
        bytes = ret.to_bytes()?;
        while bytes.len() > (msg.max_payload() as usize) {
            let mut answers = ret.take_answers();
            let new_len = answers.len() / 2;
            answers.truncate(new_len);
            ret.insert_answers(answers);
            bytes = ret.to_bytes()?;
        }
    }
    Ok(bytes)
}

fn serve_transport_async<TRecv, TSend, F>(
    mut recv: TRecv,
    mut send: TSend,
    resolver: RcResolver,
    on_error: F,
) -> JoinHandle<()>
where
    TRecv: DnsTransport + Send + 'static,
    TSend: DnsTransport + Send + 'static,
    F: FnOnce(Error) + Send + 'static,
{
    let (sch, rch) = channel::<(Vec<u8>, Option<SocketAddr>)>();
    mioco::spawn(move || loop {
        match rch.recv() {
            Ok((buf, addr)) => {
                match send.send_msg_bytes(&buf, addr.as_ref()) {
                    Ok(_) => {}
                    Err(e) => {
                        warn!("Failed to respond to DNS request: {:?}", e);
                    }
                };
            }
            Err(_) => {
                trace!("Channel is broken");
                return;
            }
        };
    });
    let keepalive = resolver.to_keepalive();
    let resolver_weak = resolver.to_weak();
    mioco::spawn(move || {
        let _ = keepalive;
        let mut resolver_weak = resolver_weak;
        loop {
            let (msg, addr) = match recv.recv_msg(None) {
                Ok(x) => x,
                Err(e) => {
                    on_error(e);
                    return;
                }
            };
            let sch_req = sch.clone();
            let resolver = resolver_weak.upgrade().unwrap();
            resolver_weak = resolver.to_weak();
            mioco::spawn(move || {
                let resp = handle_request(&resolver, &msg, TSend::should_truncate())
                    .expect("handle_request should not return error");
                sch_req
                    .send((resp, addr))
                    .expect("Result pipe should not break here");
            });
        }
    })
}
pub fn serve_tcp(addr: &SocketAddr, resolver: RcResolver) -> Result<JoinHandle<()>> {
    let listener = (TcpListener::bind(addr))?;
    let keepalive = resolver.to_keepalive();
    let resolver_weak = resolver.to_weak();
    Ok(mioco::spawn(move || {
        let _ = keepalive;
        let mut resolver_weak = resolver_weak;
        loop {
            let sock = listener.accept().expect("Failed to accept TCP socket");
            sock.set_nodelay(true).unwrap_or(());
            let sock_clone = sock.try_clone().expect("Failed to clone TCP socket");
            let resolver = resolver_weak.upgrade().unwrap();
            resolver_weak = resolver.to_weak();
            serve_transport_async(
                sock.with_resetting_timeout(*resolver.config.serve.tcp_timeout),
                sock_clone.with_resetting_timeout(*resolver.config.serve.tcp_timeout),
                resolver,
                |e| trace!("Client TCP connection is broken: {:?}", e),
            );
        }
    }))
}
pub fn serve_udp(addr: &SocketAddr, resolver: RcResolver) -> Result<JoinHandle<()>> {
    if addr.ip().is_unspecified() {
        debug!("Binding UDP listener to all available IP addresses");
        let mut handles = Vec::<JoinHandle<()>>::new();
        for iface in pnet_datalink::interfaces() {
            for ipnetwork in iface.ips {
                let ip = ipnetwork.ip();
                if ip.is_unspecified() {
                    continue;
                }
                debug!("Binding to {}", ip);
                match serve_udp(&SocketAddr::new(ip, addr.port()), resolver.clone()) {
                    Err(e) => debug!("Failed to bind to {}: {}", ip, e),
                    Ok(handle) => handles.push(handle),
                }
            }
        }
        if handles.is_empty() {
            return Err(Error::Other("Failed to bind UDP listener".into(), None));
        }
        return Ok(mioco::spawn(move || {
            for h in handles {
                h.join().unwrap();
            }
        }));
    }
    let sock: UdpSocket = (UdpSocket::bound(addr))?;
    let sock_clone = (sock.try_clone())?;
    Ok(serve_transport_async(
        sock,
        sock_clone,
        resolver,
        |e| match e {
            Error::Proto(reason) => warn!("Invalid DNS request: {:?}", reason),
            e2 => panic!("UDP listener is broken: {:?}", e2),
        },
    ))
}
<|endoftext|>