<commit_before>// Copyright 2018 The xi-editor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Module for searching text.

use std::cmp::{max, min};
use std::iter;

use crate::annotations::{AnnotationSlice, AnnotationType, ToAnnotation};
use crate::selection::{InsertDrift, SelRegion, Selection};
use crate::view::View;
use crate::word_boundaries::WordCursor;
use regex::{Regex, RegexBuilder};
use xi_rope::delta::DeltaRegion;
use xi_rope::find::{find, is_multiline_regex, CaseMatching};
use xi_rope::{Cursor, Interval, LinesMetric, Metric, Rope, RopeDelta};

const REGEX_SIZE_LIMIT: usize = 1000000;

/// Information about search queries and number of matches for find
#[derive(Serialize, Deserialize, Debug)]
pub struct FindStatus {
    /// Identifier for the current search query.
    id: usize,

    /// The current search query.
    chars: Option<String>,

    /// Whether the active search is case matching.
    case_sensitive: Option<bool>,

    /// Whether the search query is considered as regular expression.
    is_regex: Option<bool>,

    /// Query only matches whole words.
    whole_words: Option<bool>,

    /// Total number of matches.
    matches: usize,

    /// Line numbers which have find results.
    lines: Vec<usize>,
}

/// Contains logic to search text
pub struct Find {
    /// Uniquely identifies this search query.
    id: usize,

    /// The occurrences, which determine the highlights, have been updated.
    hls_dirty: bool,

    /// The currently active search string.
    search_string: Option<String>,

    /// The case matching setting for the currently active search.
    case_matching: CaseMatching,

    /// The search query should be considered as regular expression.
    regex: Option<Regex>,

    /// Query matches only whole words.
    whole_words: bool,

    /// The set of all known find occurrences (highlights).
    occurrences: Selection,
}

impl Find {
    pub fn new(id: usize) -> Find {
        Find {
            id,
            hls_dirty: true,
            search_string: None,
            case_matching: CaseMatching::CaseInsensitive,
            regex: None,
            whole_words: false,
            occurrences: Selection::new(),
        }
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn occurrences(&self) -> &Selection {
        &self.occurrences
    }

    pub fn hls_dirty(&self) -> bool {
        self.hls_dirty
    }

    pub fn find_status(&self, view: &View, text: &Rope, matches_only: bool) -> FindStatus {
        if matches_only {
            FindStatus {
                id: self.id,
                chars: None,
                case_sensitive: None,
                is_regex: None,
                whole_words: None,
                matches: self.occurrences.len(),
                lines: self
                    .occurrences
                    .iter()
                    .map(|o| view.offset_to_line_col(text, o.min()).0 + 1)
                    .collect(),
            }
        } else {
            FindStatus {
                id: self.id,
                chars: self.search_string.clone(),
                case_sensitive: Some(self.case_matching == CaseMatching::Exact),
                is_regex: Some(self.regex.is_some()),
                whole_words: Some(self.whole_words),
                matches: self.occurrences.len(),
                lines: self
                    .occurrences
                    .iter()
                    .map(|o| view.offset_to_line_col(text, o.min()).0 + 1)
                    .collect(),
            }
        }
    }

    pub fn set_hls_dirty(&mut self, is_dirty: bool) {
        self.hls_dirty = is_dirty
    }

    pub fn update_highlights(&mut self, text: &Rope, delta: &RopeDelta) {
        // update search highlights for changed regions
        if self.search_string.is_some() {
            // invalidate occurrences around deletion positions
            for DeltaRegion { old_offset, len, .. } in delta.iter_deletions() {
                self.occurrences.delete_range(old_offset, old_offset + len, false);
            }

            self.occurrences = self.occurrences.apply_delta(delta, false, InsertDrift::Default);

            // invalidate occurrences around insert positions
            for DeltaRegion { new_offset, len, .. } in delta.iter_inserts() {
                // also invalidate previous occurrence since it might expand after insertion
                // eg. for regex .* every insertion after match will be part of match
                self.occurrences.delete_range(
                    new_offset.checked_sub(1).unwrap_or(0),
                    new_offset + len,
                    false,
                );
            }

            // update find for the whole delta and everything after
            let (iv, _) = delta.summary();

            // get last valid occurrence that was unaffected by the delta
            let start = match self.occurrences.regions_in_range(0, iv.start()).last() {
                Some(reg) => reg.end,
                None => 0,
            };

            // invalidate all search results from the point of the last valid search result until ...
            let is_multi_line =
                LinesMetric::next(self.search_string.as_ref().unwrap(), 0).is_some();
            let is_multi_line_regex =
                self.regex.is_some() && is_multiline_regex(self.search_string.as_ref().unwrap());

            if is_multi_line || is_multi_line_regex {
                // ... the end of the file
                self.occurrences.delete_range(iv.start(), text.len(), false);
                self.update_find(text, start, text.len(), false);
            } else {
                // ... the end of the line including line break
                let mut cursor = Cursor::new(&text, iv.end());

                let end_of_line = match cursor.next::<LinesMetric>() {
                    Some(end) => end,
                    None if cursor.pos() == text.len() => cursor.pos(),
                    _ => return,
                };

                self.occurrences.delete_range(iv.start(), end_of_line, false);
                self.update_find(text, start, end_of_line, false);
            }
        }
    }

    /// Set search parameters and executes the search.
    pub fn do_find(
        &mut self,
        text: &Rope,
        search_string: &str,
        case_sensitive: bool,
        is_regex: bool,
        whole_words: bool,
    ) {
        if search_string.is_empty() {
            self.unset();
        }

        self.set_find(search_string, case_sensitive, is_regex, whole_words);
        self.update_find(text, 0, text.len(), false);
    }

    /// Unsets the search and removes all highlights from the view.
    pub fn unset(&mut self) {
        self.search_string = None;
        self.occurrences = Selection::new();
        self.hls_dirty = true;
    }

    /// Sets find parameters and search query.
    fn set_find(
        &mut self,
        search_string: &str,
        case_sensitive: bool,
        is_regex: bool,
        whole_words: bool,
    ) {
        let case_matching =
            if case_sensitive { CaseMatching::Exact } else { CaseMatching::CaseInsensitive };

        if let Some(ref s) = self.search_string {
            if s == search_string
                && case_matching == self.case_matching
                && self.regex.is_some() == is_regex
                && self.whole_words == whole_words
            {
                // search parameters did not change
                return;
            }
        }

        self.unset();

        self.search_string = Some(search_string.to_string());
        self.case_matching = case_matching;
        self.whole_words = whole_words;

        // create regex from untrusted input
        self.regex = match is_regex {
            false => None,
            true => RegexBuilder::new(search_string)
                .size_limit(REGEX_SIZE_LIMIT)
                .case_insensitive(case_matching == CaseMatching::CaseInsensitive)
                .build()
                .ok(),
        };
    }

    /// Execute the search on the provided text in the range provided by `start` and `end`.
    pub fn update_find(&mut self, text: &Rope, start: usize, end: usize, include_slop: bool) {
        if self.search_string.is_none() {
            return;
        }

        // extend the search by twice the string length (twice, because case matching may increase
        // the length of an occurrence)
        let slop = if include_slop { self.search_string.as_ref().unwrap().len() * 2 } else { 0 };

        let search_string = self.search_string.as_ref().unwrap();

        // expand region to be able to find occurrences around the region's edges
        let from = max(start, slop) - slop;
        let to = min(end + slop, text.len());

        // TODO: this interval might cut a unicode codepoint, make sure it is
        // aligned to codepoint boundaries.
        let sub_text = text.subseq(Interval::new(0, to));
        let mut find_cursor = Cursor::new(&sub_text, from);
        let mut raw_lines = text.lines_raw(from..to);

        while let Some(start) =
            find(&mut find_cursor, &mut raw_lines, self.case_matching, &search_string, &self.regex)
        {
            let end = find_cursor.pos();

            if self.whole_words && !self.is_matching_whole_words(text, start, end) {
                raw_lines = text.lines_raw(find_cursor.pos()..to);
                continue;
            }

            let region = SelRegion::new(start, end);
            let (_, e) = self.occurrences.add_range_distinct(region);
            // in case of ambiguous search results (e.g. search "aba" in "ababa"),
            // the search result closer to the beginning of the file wins
            if e != end {
                // Skip the search result and keep the occurrence that is closer to
                // the beginning of the file. Re-align the cursor to the kept
                // occurrence
                find_cursor.set(e);
                raw_lines = text.lines_raw(find_cursor.pos()..to);
                continue;
            }

            // in case current cursor matches search result (for example query a* matches)
            // all cursor positions, then cursor needs to be increased so that search
            // continues at next position. Otherwise, search will result in overflow since
            // search will always repeat at current cursor position.
            if start == end {
                // determine whether end of text is reached and stop search or increase
                // cursor manually
                if end + 1 >= text.len() {
                    break;
                } else {
                    find_cursor.set(end + 1);
                }
            }

            // update line iterator so that line starts at current cursor position
            raw_lines = text.lines_raw(find_cursor.pos()..to);
        }

        self.hls_dirty = true;
    }

    /// Return the occurrence closest to the provided selection `sel`. If searched is reversed then
    /// the occurrence closest to the start of the selection is returned. `wrapped` indicates that
    /// if the end of the text is reached the search continues from the start.
    pub fn next_occurrence(
        &self,
        text: &Rope,
        reverse: bool,
        wrapped: bool,
        sel: &Selection,
    ) -> Option<SelRegion> {
        if self.occurrences.len() == 0 {
            return None;
        }

        let (sel_start, sel_end) = match sel.last() {
            Some(last) if last.is_caret() =>
            // if last selection is caret then allow the current position to be part of the occurrence
            {
                (last.min(), last.max())
            }
            Some(last) if !last.is_caret() =>
            // if the last selection is not a caret then continue searching after the caret
            {
                (last.min(), last.max() + 1)
            }
            _ => (0, 0),
        };

        if reverse {
            let next_occurrence = match sel_start.checked_sub(1) {
                Some(search_end) => self.occurrences.regions_in_range(0, search_end).last(),
                None => None,
            };

            if next_occurrence.is_none() && !wrapped {
                // get previous unselected occurrence
                return self
                    .occurrences
                    .regions_in_range(0, text.len())
                    .iter()
                    .cloned()
                    .filter(|o| sel.regions_in_range(o.min(), o.max()).is_empty())
                    .collect::<Vec<SelRegion>>()
                    .last()
                    .cloned();
            }

            next_occurrence.cloned()
        } else {
            let next_occurrence = self.occurrences.regions_in_range(sel_end, text.len()).first();

            if next_occurrence.is_none() && !wrapped {
                // get next unselected occurrence
                return self
                    .occurrences
                    .regions_in_range(0, text.len())
                    .iter()
                    .cloned()
                    .filter(|o| sel.regions_in_range(o.min(), o.max()).is_empty())
                    .collect::<Vec<SelRegion>>()
                    .first()
                    .cloned();
            }

            next_occurrence.cloned()
        }
    }

    /// Checks if the start and end of a match is matching whole words.
    fn is_matching_whole_words(&self, text: &Rope, start: usize, end: usize) -> bool {
        let mut word_end_cursor = WordCursor::new(text, end - 1);
        let mut word_start_cursor = WordCursor::new(text, start + 1);

        if let Some(start_boundary) = word_start_cursor.prev_boundary() {
            if start_boundary != start {
                return false;
            }
        }

        if let Some(end_boundary) = word_end_cursor.next_boundary() {
            if end_boundary != end {
                return false;
            }
        }

        true
    }
}

/// Implementing the `ToAnnotation` trait allows to convert finds to annotations.
impl ToAnnotation for Find {
    fn get_annotations(&self, interval: Interval, view: &View, text: &Rope) -> AnnotationSlice {
        let regions = self.occurrences.regions_in_range(interval.start(), interval.end());
        let ranges = regions
            .iter()
            .map(|region| {
                let (start_line, start_col) = view.offset_to_line_col(text, region.min());
                let (end_line, end_col) = view.offset_to_line_col(text, region.max());
                [start_line, start_col, end_line, end_col]
            })
            .collect::<Vec<[usize; 4]>>();

        let payload = iter::repeat(json!({"id": self.id})).take(ranges.len()).collect::<Vec<_>>();

        AnnotationSlice::new(AnnotationType::Find, ranges, Some(payload))
    }
}
<commit_msg>Bugfix: update find highlights for pasted text<commit_after>// Copyright 2018 The xi-editor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Module for searching text.

use std::cmp::{max, min};
use std::iter;

use crate::annotations::{AnnotationSlice, AnnotationType, ToAnnotation};
use crate::selection::{InsertDrift, SelRegion, Selection};
use crate::view::View;
use crate::word_boundaries::WordCursor;
use regex::{Regex, RegexBuilder};
use xi_rope::delta::DeltaRegion;
use xi_rope::find::{find, is_multiline_regex, CaseMatching};
use xi_rope::{Cursor, Interval, LinesMetric, Metric, Rope, RopeDelta};

const REGEX_SIZE_LIMIT: usize = 1000000;

/// Information about search queries and number of matches for find
#[derive(Serialize, Deserialize, Debug)]
pub struct FindStatus {
    /// Identifier for the current search query.
    id: usize,

    /// The current search query.
    chars: Option<String>,

    /// Whether the active search is case matching.
    case_sensitive: Option<bool>,

    /// Whether the search query is considered as regular expression.
    is_regex: Option<bool>,

    /// Query only matches whole words.
    whole_words: Option<bool>,

    /// Total number of matches.
    matches: usize,

    /// Line numbers which have find results.
    lines: Vec<usize>,
}

/// Contains logic to search text
pub struct Find {
    /// Uniquely identifies this search query.
    id: usize,

    /// The occurrences, which determine the highlights, have been updated.
    hls_dirty: bool,

    /// The currently active search string.
    search_string: Option<String>,

    /// The case matching setting for the currently active search.
    case_matching: CaseMatching,

    /// The search query should be considered as regular expression.
    regex: Option<Regex>,

    /// Query matches only whole words.
    whole_words: bool,

    /// The set of all known find occurrences (highlights).
    occurrences: Selection,
}

impl Find {
    pub fn new(id: usize) -> Find {
        Find {
            id,
            hls_dirty: true,
            search_string: None,
            case_matching: CaseMatching::CaseInsensitive,
            regex: None,
            whole_words: false,
            occurrences: Selection::new(),
        }
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn occurrences(&self) -> &Selection {
        &self.occurrences
    }

    pub fn hls_dirty(&self) -> bool {
        self.hls_dirty
    }

    pub fn find_status(&self, view: &View, text: &Rope, matches_only: bool) -> FindStatus {
        if matches_only {
            FindStatus {
                id: self.id,
                chars: None,
                case_sensitive: None,
                is_regex: None,
                whole_words: None,
                matches: self.occurrences.len(),
                lines: self
                    .occurrences
                    .iter()
                    .map(|o| view.offset_to_line_col(text, o.min()).0 + 1)
                    .collect(),
            }
        } else {
            FindStatus {
                id: self.id,
                chars: self.search_string.clone(),
                case_sensitive: Some(self.case_matching == CaseMatching::Exact),
                is_regex: Some(self.regex.is_some()),
                whole_words: Some(self.whole_words),
                matches: self.occurrences.len(),
                lines: self
                    .occurrences
                    .iter()
                    .map(|o| view.offset_to_line_col(text, o.min()).0 + 1)
                    .collect(),
            }
        }
    }

    pub fn set_hls_dirty(&mut self, is_dirty: bool) {
        self.hls_dirty = is_dirty
    }

    pub fn update_highlights(&mut self, text: &Rope, delta: &RopeDelta) {
        // update search highlights for changed regions
        if self.search_string.is_some() {
            // invalidate occurrences around deletion positions
            for DeltaRegion { old_offset, len, .. } in delta.iter_deletions() {
                self.occurrences.delete_range(old_offset, old_offset + len, false);
            }

            self.occurrences = self.occurrences.apply_delta(delta, false, InsertDrift::Default);

            // invalidate occurrences around insert positions
            for DeltaRegion { new_offset, len, .. } in delta.iter_inserts() {
                // also invalidate previous occurrence since it might expand after insertion
                // eg. for regex .* every insertion after match will be part of match
                self.occurrences.delete_range(
                    new_offset.checked_sub(1).unwrap_or(0),
                    new_offset + len,
                    false,
                );
            }

            // update find for the whole delta and everything after
            let (iv, new_len) = delta.summary();

            // get last valid occurrence that was unaffected by the delta
            let start = match self.occurrences.regions_in_range(0, iv.start()).last() {
                Some(reg) => reg.end,
                None => 0,
            };

            // invalidate all search results from the point of the last valid search result until ...
            let is_multi_line =
                LinesMetric::next(self.search_string.as_ref().unwrap(), 0).is_some();
            let is_multi_line_regex =
                self.regex.is_some() && is_multiline_regex(self.search_string.as_ref().unwrap());

            if is_multi_line || is_multi_line_regex {
                // ... the end of the file
                self.occurrences.delete_range(iv.start(), text.len(), false);
                self.update_find(text, start, text.len(), false);
            } else {
                // ... the end of the line including line break
                let mut cursor = Cursor::new(&text, iv.end() + new_len);

                let end_of_line = match cursor.next::<LinesMetric>() {
                    Some(end) => end,
                    None if cursor.pos() == text.len() => cursor.pos(),
                    _ => return,
                };

                self.occurrences.delete_range(iv.start(), end_of_line, false);
                self.update_find(text, start, end_of_line, false);
            }
        }
    }

    /// Set search parameters and executes the search.
    pub fn do_find(
        &mut self,
        text: &Rope,
        search_string: &str,
        case_sensitive: bool,
        is_regex: bool,
        whole_words: bool,
    ) {
        if search_string.is_empty() {
            self.unset();
        }

        self.set_find(search_string, case_sensitive, is_regex, whole_words);
        self.update_find(text, 0, text.len(), false);
    }

    /// Unsets the search and removes all highlights from the view.
    pub fn unset(&mut self) {
        self.search_string = None;
        self.occurrences = Selection::new();
        self.hls_dirty = true;
    }

    /// Sets find parameters and search query.
    fn set_find(
        &mut self,
        search_string: &str,
        case_sensitive: bool,
        is_regex: bool,
        whole_words: bool,
    ) {
        let case_matching =
            if case_sensitive { CaseMatching::Exact } else { CaseMatching::CaseInsensitive };

        if let Some(ref s) = self.search_string {
            if s == search_string
                && case_matching == self.case_matching
                && self.regex.is_some() == is_regex
                && self.whole_words == whole_words
            {
                // search parameters did not change
                return;
            }
        }

        self.unset();

        self.search_string = Some(search_string.to_string());
        self.case_matching = case_matching;
        self.whole_words = whole_words;

        // create regex from untrusted input
        self.regex = match is_regex {
            false => None,
            true => RegexBuilder::new(search_string)
                .size_limit(REGEX_SIZE_LIMIT)
                .case_insensitive(case_matching == CaseMatching::CaseInsensitive)
                .build()
                .ok(),
        };
    }

    /// Execute the search on the provided text in the range provided by `start` and `end`.
    pub fn update_find(&mut self, text: &Rope, start: usize, end: usize, include_slop: bool) {
        if self.search_string.is_none() {
            return;
        }

        // extend the search by twice the string length (twice, because case matching may increase
        // the length of an occurrence)
        let slop = if include_slop { self.search_string.as_ref().unwrap().len() * 2 } else { 0 };

        let search_string = self.search_string.as_ref().unwrap();

        // expand region to be able to find occurrences around the region's edges
        let from = max(start, slop) - slop;
        let to = min(end + slop, text.len());

        // TODO: this interval might cut a unicode codepoint, make sure it is
        // aligned to codepoint boundaries.
        let sub_text = text.subseq(Interval::new(0, to));
        let mut find_cursor = Cursor::new(&sub_text, from);
        let mut raw_lines = text.lines_raw(from..to);

        while let Some(start) =
            find(&mut find_cursor, &mut raw_lines, self.case_matching, &search_string, &self.regex)
        {
            let end = find_cursor.pos();

            if self.whole_words && !self.is_matching_whole_words(text, start, end) {
                raw_lines = text.lines_raw(find_cursor.pos()..to);
                continue;
            }

            let region = SelRegion::new(start, end);
            let (_, e) = self.occurrences.add_range_distinct(region);
            // in case of ambiguous search results (e.g. search "aba" in "ababa"),
            // the search result closer to the beginning of the file wins
            if e != end {
                // Skip the search result and keep the occurrence that is closer to
                // the beginning of the file. Re-align the cursor to the kept
                // occurrence
                find_cursor.set(e);
                raw_lines = text.lines_raw(find_cursor.pos()..to);
                continue;
            }

            // in case current cursor matches search result (for example query a* matches)
            // all cursor positions, then cursor needs to be increased so that search
            // continues at next position. Otherwise, search will result in overflow since
            // search will always repeat at current cursor position.
            if start == end {
                // determine whether end of text is reached and stop search or increase
                // cursor manually
                if end + 1 >= text.len() {
                    break;
                } else {
                    find_cursor.set(end + 1);
                }
            }

            // update line iterator so that line starts at current cursor position
            raw_lines = text.lines_raw(find_cursor.pos()..to);
        }

        self.hls_dirty = true;
    }

    /// Return the occurrence closest to the provided selection `sel`. If searched is reversed then
    /// the occurrence closest to the start of the selection is returned. `wrapped` indicates that
    /// if the end of the text is reached the search continues from the start.
    pub fn next_occurrence(
        &self,
        text: &Rope,
        reverse: bool,
        wrapped: bool,
        sel: &Selection,
    ) -> Option<SelRegion> {
        if self.occurrences.len() == 0 {
            return None;
        }

        let (sel_start, sel_end) = match sel.last() {
            Some(last) if last.is_caret() =>
            // if last selection is caret then allow the current position to be part of the occurrence
            {
                (last.min(), last.max())
            }
            Some(last) if !last.is_caret() =>
            // if the last selection is not a caret then continue searching after the caret
            {
                (last.min(), last.max() + 1)
            }
            _ => (0, 0),
        };

        if reverse {
            let next_occurrence = match sel_start.checked_sub(1) {
                Some(search_end) => self.occurrences.regions_in_range(0, search_end).last(),
                None => None,
            };

            if next_occurrence.is_none() && !wrapped {
                // get previous unselected occurrence
                return self
                    .occurrences
                    .regions_in_range(0, text.len())
                    .iter()
                    .cloned()
                    .filter(|o| sel.regions_in_range(o.min(), o.max()).is_empty())
                    .collect::<Vec<SelRegion>>()
                    .last()
                    .cloned();
            }

            next_occurrence.cloned()
        } else {
            let next_occurrence = self.occurrences.regions_in_range(sel_end, text.len()).first();

            if next_occurrence.is_none() && !wrapped {
                // get next unselected occurrence
                return self
                    .occurrences
                    .regions_in_range(0, text.len())
                    .iter()
                    .cloned()
                    .filter(|o| sel.regions_in_range(o.min(), o.max()).is_empty())
                    .collect::<Vec<SelRegion>>()
                    .first()
                    .cloned();
            }

            next_occurrence.cloned()
        }
    }

    /// Checks if the start and end of a match is matching whole words.
    fn is_matching_whole_words(&self, text: &Rope, start: usize, end: usize) -> bool {
        let mut word_end_cursor = WordCursor::new(text, end - 1);
        let mut word_start_cursor = WordCursor::new(text, start + 1);

        if let Some(start_boundary) = word_start_cursor.prev_boundary() {
            if start_boundary != start {
                return false;
            }
        }

        if let Some(end_boundary) = word_end_cursor.next_boundary() {
            if end_boundary != end {
                return false;
            }
        }

        true
    }
}

/// Implementing the `ToAnnotation` trait allows to convert finds to annotations.
impl ToAnnotation for Find {
    fn get_annotations(&self, interval: Interval, view: &View, text: &Rope) -> AnnotationSlice {
        let regions = self.occurrences.regions_in_range(interval.start(), interval.end());
        let ranges = regions
            .iter()
            .map(|region| {
                let (start_line, start_col) = view.offset_to_line_col(text, region.min());
                let (end_line, end_col) = view.offset_to_line_col(text, region.max());
                [start_line, start_col, end_line, end_col]
            })
            .collect::<Vec<[usize; 4]>>();

        let payload = iter::repeat(json!({"id": self.id})).take(ranges.len()).collect::<Vec<_>>();

        AnnotationSlice::new(AnnotationType::Find, ranges, Some(payload))
    }
}

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

    #[test]
    fn find() {
        let base_text = Rope::from("hello world");
        let mut find = Find::new(1);
        find.do_find(&base_text, "world", false, false, false);
        assert_eq!(find.occurrences().len(), 1);
        assert_eq!(find.occurrences().first(), Some(&SelRegion::new(6, 11)));

        // todo: more tests with all options
    }

    #[test]
    fn update_find_edit() {
        let base_text = Rope::from("a b a c");
        let mut find = Find::new(1);
        find.do_find(&base_text, "a", false, false, false);
        let mut builder = DeltaBuilder::new(base_text.len());
        builder.replace(0..0, "a ".into());

        assert_eq!(find.occurrences().len(), 2);
        assert_eq!(find.occurrences().first(), Some(&SelRegion::new(0, 1)));
        assert_eq!(find.occurrences().last(), Some(&SelRegion::new(4, 5)));
    }

    #[test]
    fn update_find_multi_line_edit() {
        let base_text = Rope::from("x\n a\n b\n a\n c");
        let mut find = Find::new(1);
        find.do_find(&base_text, "a", false, false, false);
        let mut builder = DeltaBuilder::new(base_text.len());
        builder.replace(2..2, " a\n b\n a\n".into());

        assert_eq!(find.occurrences().len(), 2);
        assert_eq!(find.occurrences().first(), Some(&SelRegion::new(3, 4)));
        assert_eq!(find.occurrences().last(), Some(&SelRegion::new(9, 10)));
    }
}
<|endoftext|>


<commit_before>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![no_std]
#![no_main]
#![feature(abi_x86_interrupt)]

use core::{arch::asm, ops::Deref, panic::PanicInfo};
use lazy_static::lazy_static;
use oak_linux_boot_params::{BootParams, CCBlobSevInfo, CCSetupData, SetupDataType};
use sev_guest::{
    cpuid::CpuidPage,
    msr::{get_sev_status, SevStatus},
    secrets::SecretsPage,
};
use x86_64::{
    instructions::{hlt, interrupts::int3},
    structures::idt::{InterruptDescriptorTable, InterruptStackFrame},
};

mod asm;
mod ghcb;
mod serial;

lazy_static! {
    static ref IDT: InterruptDescriptorTable = {
        let mut idt = InterruptDescriptorTable::new();
        idt.general_protection_fault.set_handler_fn(gp_handler);
        idt
    };
}

/// Crude implementation of a #GP handler that skips any `rdmsr` instructions that trigger #GP.
///
/// Note that this doesn't properly emulate reading the MSR! It doesn't change any register values;
/// the instruction is just skipped.
extern "x86-interrupt" fn gp_handler(mut frame: InterruptStackFrame, _: u64) {
    unsafe {
        // RDMSR can cause a #GP if you try to read a non-existing MSR
        if *frame.deref().instruction_pointer.as_ptr::<u16>() == 0x320Fu16 {
            // just skip the instruction
            frame.as_mut().update(|val| {
                val.instruction_pointer += 2u64;
            });
        } else {
            // not a RDMSR, trigger a double fault
            asm!("int 8", options(noreturn));
        }
    }
}

#[no_mangle]
pub extern "C" fn rust64_start(_: u64, boot_params: &BootParams) -> ! {
    IDT.load();
    serial::init_logging();
    log::info!("Hello World!");

    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ENABLED) {
        log::info!("SEV enabled");
    }
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        log::info!("SEV-ES enabled");
    }
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        log::info!("SNP active");

        // Walk the null-terminated setup_data list until we find the CC blob.
        let mut setup_data = boot_params.hdr.setup_data;

        // Safety: there's a lot of dereferences of raw pointers here; unfortunately necessary as we
        // need to access the C data structures set up by the bootloader. It is safe as long as the
        // bootloader populated the zero page correctly.
        while !setup_data.is_null() {
            let type_ = unsafe { (*setup_data).type_ };
            if type_ == SetupDataType::CCBlob {
                break;
            }
            setup_data = unsafe { (*setup_data).next };
        }

        if !setup_data.is_null() {
            let cc_setup_data = unsafe { &*(setup_data as *const CCSetupData) };
            let cc_blob = unsafe { &*(cc_setup_data.cc_blob_address as *const CCBlobSevInfo) };
            let magic = cc_blob.magic;
            if magic != oak_linux_boot_params::CC_BLOB_SEV_INFO_MAGIC {
                panic!("CCBlobSevInfo data structure has invalid magic: {}", magic);
            }
            let cpuid = unsafe { &*(cc_blob.cpuid_phys as *const CpuidPage) };
            let secrets = unsafe { &*(cc_blob.secrets_phys as *const SecretsPage) };

            log::info!("CPUID page: {:?}", cpuid);
            log::info!("Secrets page: {:?}", secrets);
        } else {
            log::warn!("CCSetupData data structure not found in setup_data");
        }
    }
    panic!();
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        let mut ghcb_protocol = crate::ghcb::init_ghcb(sev_status.contains(SevStatus::SNP_ACTIVE));
        // Use the i8042 device to shut down the machine. Assumes the VMM exposes the device.
        // This is safe as both qemu and crosvm expose the i8042 device by default.
        let _ = ghcb_protocol.io_write_u8(0x64, 0xFE);
    } else {
        // Trigger a breakpoint exception. As we don't have a #BP handler, this will triple fault
        // and terminate the program.
        int3();
    }

    // But if we're still here for some reason, just go in an infinite halt loop.
    loop {
        hlt();
    }
}
<commit_msg>Replace the #GP handler in hello world test with one that emulates `RDMSR` properly.<commit_after>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#![no_std]
#![no_main]
#![feature(abi_x86_interrupt)]
#![feature(naked_functions)]

use core::{arch::asm, panic::PanicInfo};
use lazy_static::lazy_static;
use oak_linux_boot_params::{BootParams, CCBlobSevInfo, CCSetupData, SetupDataType};
use sev_guest::{
    cpuid::CpuidPage,
    msr::{get_sev_status, SevStatus},
    secrets::SecretsPage,
};
use x86_64::{
    instructions::{hlt, interrupts::int3},
    structures::idt::{InterruptDescriptorTable, InterruptStackFrame},
};

mod asm;
mod ghcb;
mod serial;

lazy_static! {
    static ref IDT: InterruptDescriptorTable = {
        let mut idt = InterruptDescriptorTable::new();
        idt.general_protection_fault.set_handler_fn(gp_handler);
        idt
    };
}

/// Crude implementation of a #GP handler that emulates any `rdmsr` instructions that trigger #GP.
///
/// This function emulates the `rdmsr` instruction by returning 0 in RAX and RDX.
#[naked]
extern "x86-interrupt" fn gp_handler(_: InterruptStackFrame, _: u64) {
    unsafe {
        asm! {
            "mov 8(%rsp), %rax",    // rax = rsp + 16 (address of the return RIP)
            "cmpw $0x320F, (%rax)", // is RIP pointing to 0x320F (RDMSR)?
            "jne 2f",               // if not, jump to label 2
            "add $2, %rax",         // increment rax by 2 (size of RDMSR instruction)
            "add $16, %rsp",        // drop the error code and old RIP from stack
            "push %rax",            // put Return RIP back onto the stack for iretq
            "xor %rax, %rax",       // zero out RAX
            "xor %rdx, %rdx",       // zero out RDX
            "iretq",                // and go back claiming we've executed the RDMSR
            "2:",                   // it wasn't RDMSR
            "int $8",               // cause a double fault
            options(att_syntax, noreturn)
        }
    }
}

#[no_mangle]
pub extern "C" fn rust64_start(_: u64, boot_params: &BootParams) -> ! {
    IDT.load();
    serial::init_logging();
    log::info!("Hello World!");

    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ENABLED) {
        log::info!("SEV enabled");
    }
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        log::info!("SEV-ES enabled");
    }
    if sev_status.contains(SevStatus::SNP_ACTIVE) {
        log::info!("SNP active");

        // Walk the null-terminated setup_data list until we find the CC blob.
        let mut setup_data = boot_params.hdr.setup_data;

        // Safety: there's a lot of dereferences of raw pointers here; unfortunately necessary as we
        // need to access the C data structures set up by the bootloader. It is safe as long as the
        // bootloader populated the zero page correctly.
        while !setup_data.is_null() {
            let type_ = unsafe { (*setup_data).type_ };
            if type_ == SetupDataType::CCBlob {
                break;
            }
            setup_data = unsafe { (*setup_data).next };
        }

        if !setup_data.is_null() {
            let cc_setup_data = unsafe { &*(setup_data as *const CCSetupData) };
            let cc_blob = unsafe { &*(cc_setup_data.cc_blob_address as *const CCBlobSevInfo) };
            let magic = cc_blob.magic;
            if magic != oak_linux_boot_params::CC_BLOB_SEV_INFO_MAGIC {
                panic!("CCBlobSevInfo data structure has invalid magic: {}", magic);
            }
            let cpuid = unsafe { &*(cc_blob.cpuid_phys as *const CpuidPage) };
            let secrets = unsafe { &*(cc_blob.secrets_phys as *const SecretsPage) };

            log::info!("CPUID page: {:?}", cpuid);
            log::info!("Secrets page: {:?}", secrets);
        } else {
            log::warn!("CCSetupData data structure not found in setup_data");
        }
    }
    panic!();
}

#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    let sev_status = get_sev_status().unwrap_or(SevStatus::empty());
    if sev_status.contains(SevStatus::SEV_ES_ENABLED) {
        let mut ghcb_protocol = crate::ghcb::init_ghcb(sev_status.contains(SevStatus::SNP_ACTIVE));
        // Use the i8042 device to shut down the machine. Assumes the VMM exposes the device.
        // This is safe as both qemu and crosvm expose the i8042 device by default.
        let _ = ghcb_protocol.io_write_u8(0x64, 0xFE);
    } else {
        // Trigger a breakpoint exception. As we don't have a #BP handler, this will triple fault
        // and terminate the program.
        int3();
    }

    // But if we're still here for some reason, just go in an infinite halt loop.
    loop {
        hlt();
    }
}
<|endoftext|>


<commit_before>// crypto_auth_hmacsha512256.h

pub const crypto_auth_hmacsha512256_BYTES: usize = 32;
pub const crypto_auth_hmacsha512256_KEYBYTES: usize = 32;

extern {
    pub fn crypto_auth_hmacsha512256_bytes() -> size_t;
    pub fn crypto_auth_hmacsha512256_keybytes() -> size_t;
    pub fn crypto_auth_hmacsha512256(
        a: *mut [u8; crypto_auth_hmacsha512256_BYTES],
        m: *const u8,
        mlen: c_ulonglong,
        k: *const [u8; crypto_auth_hmacsha512256_KEYBYTES]) -> c_int;
    pub fn crypto_auth_hmacsha512256_verify(
        a: *const [u8; crypto_auth_hmacsha512256_BYTES],
        m: *const u8,
        mlen: c_ulonglong,
        k: *const [u8; crypto_auth_hmacsha512256_KEYBYTES]) -> c_int;
}

#[test]
fn test_crypto_auth_hmacsha512256_bytes() {
    assert!(unsafe { crypto_auth_hmacsha512256_bytes() as usize } ==
            crypto_auth_hmacsha512256_BYTES)
}
#[test]
fn test_crypto_auth_hmacsha512256_keybytes() {
    assert!(unsafe { crypto_auth_hmacsha512256_keybytes() as usize } ==
            crypto_auth_hmacsha512256_KEYBYTES)
}
<commit_msg>libsodium-sys: definitions from crypto_auth_hmacsha512256.h complete<commit_after>// crypto_auth_hmacsha512256.h

#[allow(non_camel_case_types)]
pub type crypto_auth_hmacsha512256_state = crypto_auth_hmacsha512_state;

pub const crypto_auth_hmacsha512256_BYTES: usize = 32;
pub const crypto_auth_hmacsha512256_KEYBYTES: usize = 32;

extern {
    pub fn crypto_auth_hmacsha512256_bytes() -> size_t;
    pub fn crypto_auth_hmacsha512256_keybytes() -> size_t;
    pub fn crypto_auth_hmacsha512256(
        a: *mut [u8; crypto_auth_hmacsha512256_BYTES],
        m: *const u8,
        mlen: c_ulonglong,
        k: *const [u8; crypto_auth_hmacsha512256_KEYBYTES]) -> c_int;
    pub fn crypto_auth_hmacsha512256_verify(
        a: *const [u8; crypto_auth_hmacsha512256_BYTES],
        m: *const u8,
        mlen: c_ulonglong,
        k: *const [u8; crypto_auth_hmacsha512256_KEYBYTES]) -> c_int;
    pub fn crypto_auth_hmacsha512256_init(
        state: *mut crypto_auth_hmacsha512256_state,
        key: *const u8,
        keylen: size_t) -> c_int;
    pub fn crypto_auth_hmacsha512256_update(
        state: *mut crypto_auth_hmacsha512256_state,
        m: *const u8,
        mlen: c_ulonglong) -> c_int;
    pub fn crypto_auth_hmacsha512256_final(
        state: *mut crypto_auth_hmacsha512256_state,
        a: *mut [u8; crypto_auth_hmacsha512256_BYTES]) -> c_int;
}

#[test]
fn test_crypto_auth_hmacsha512256_bytes() {
    assert!(unsafe { crypto_auth_hmacsha512256_bytes() as usize } ==
            crypto_auth_hmacsha512256_BYTES)
}
#[test]
fn test_crypto_auth_hmacsha512256_keybytes() {
    assert!(unsafe { crypto_auth_hmacsha512256_keybytes() as usize } ==
            crypto_auth_hmacsha512256_KEYBYTES)
}
<|endoftext|>


<commit_before>
use std::mem::transmute;
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};

use super::{Abomonation, decode};

/// A type wrapping owned decoded abomonated data.
///
/// This type ensures that decoding and pointer correction has already happened,
/// and implements `Deref<Target=T>` using a pointer cast (transmute).
///
/// #Safety
///
/// The safety of this type, and in particular its `transute` implementation of
/// the `Deref` trait, relies on the owned bytes not being externally mutated
/// once provided. You could imagine a new type implementing `DerefMut` as required,
/// but which also retains the ability (e.g. through `RefCell`) to mutate the bytes.
/// This would be very bad, but seems hard to prevent in the type system. Please
/// don't do this.
///
/// #Examples
///
/// ```
/// use std::ops::Deref;
/// use abomonation::{encode, decode};
/// use abomonation::abomonated::Abomonated;
///
/// // create some test data out of abomonation-approved types
/// let vector = (0..256u64).map(|i| (i, format!("{}", i)))
///                         .collect::<Vec<_>>();
///
/// // encode a Vec<(u64, String)> into a Vec<u8>
/// let mut bytes = Vec::new();
/// unsafe { encode(&vector, &mut bytes); }
///
/// // attempt to decode `bytes` into a `&Vec<(u64, String)>`.
/// let maybe_decoded = unsafe { Abomonated::<Vec<(u64, String)>,_>::new(bytes) };
///
/// if let Some(decoded) = maybe_decoded {
///     // each `deref()` call is now just a pointer cast.
///     assert!(decoded.deref() == &vector);
/// }
/// else {
///     panic!("failed to decode");
/// }
/// ```
pub struct Abomonated<T: Abomonation, S: DerefMut<Target=[u8]>> {
    phantom: PhantomData<T>,
    decoded: S,
}

impl<T: Abomonation, S: DerefMut<Target=[u8]>> Abomonated<T, S> {

    /// Attempts to create decoded data from owned mutable bytes.
    ///
    /// This method will return `None` if it is unable to decode the data with
    /// type `T`.
    ///
    /// #Examples
    ///
    /// ```
    /// use std::ops::Deref;
    /// use abomonation::{encode, decode};
    /// use abomonation::abomonated::Abomonated;
    ///
    /// // create some test data out of abomonation-approved types
    /// let vector = (0..256u64).map(|i| (i, format!("{}", i)))
    ///                         .collect::<Vec<_>>();
    ///
    /// // encode a Vec<(u64, String)> into a Vec<u8>
    /// let mut bytes = Vec::new();
    /// unsafe { encode(&vector, &mut bytes); }
    ///
    /// // attempt to decode `bytes` into a `&Vec<(u64, String)>`.
    /// let maybe_decoded = unsafe { Abomonated::<Vec<(u64, String)>,_>::new(bytes) };
    ///
    /// if let Some(decoded) = maybe_decoded {
    ///     // each `deref()` call is now just a pointer cast.
    ///     assert!(decoded.deref() == &vector);
    /// }
    /// else {
    ///     panic!("failed to decode");
    /// }
    /// ```
    pub unsafe fn new(mut bytes: S) -> Option<Self> {

        // performs the underlying pointer correction, indicates success.
        let decoded = decode::<T>(bytes.deref_mut()).is_some();

        if decoded {
            Some(Abomonated {
                phantom: PhantomData,
                decoded: bytes,
            })
        }
        else {
            None
        }
    }
}

impl<T: Abomonation, S: DerefMut<Target=[u8]>> Deref for Abomonated<T, S> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &T {
        let result: &T = unsafe { transmute(self.decoded.get_unchecked(0)) };
        result
    }
}
<commit_msg>removed constraints, added comments<commit_after>
use std::mem::transmute;
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};

use super::{Abomonation, decode};

/// A type wrapping owned decoded abomonated data.
///
/// This type ensures that decoding and pointer correction has already happened,
/// and implements `Deref<Target=T>` using a pointer cast (transmute).
///
/// #Safety
///
/// The safety of this type, and in particular its `transute` implementation of
/// the `Deref` trait, relies on the owned bytes not being externally mutated
/// once provided. You could imagine a new type implementing `DerefMut` as required,
/// but which also retains the ability (e.g. through `RefCell`) to mutate the bytes.
/// This would be very bad, but seems hard to prevent in the type system. Please
/// don't do this.
///
/// You must also use a type `S` whose bytes have a fixed location in memory.
/// Otherwise moving an instance of `Abomonated<T, S>` may invalidate decoded
/// pointers, and everything goes badly.
///
/// #Examples
///
/// ```
/// use std::ops::Deref;
/// use abomonation::{encode, decode};
/// use abomonation::abomonated::Abomonated;
///
/// // create some test data out of abomonation-approved types
/// let vector = (0..256u64).map(|i| (i, format!("{}", i)))
///                         .collect::<Vec<_>>();
///
/// // encode a Vec<(u64, String)> into a Vec<u8>
/// let mut bytes = Vec::new();
/// unsafe { encode(&vector, &mut bytes); }
///
/// // attempt to decode `bytes` into a `&Vec<(u64, String)>`.
/// let maybe_decoded = unsafe { Abomonated::<Vec<(u64, String)>,_>::new(bytes) };
///
/// if let Some(decoded) = maybe_decoded {
///     // each `deref()` call is now just a pointer cast.
///     assert!(decoded.deref() == &vector);
/// }
/// else {
///     panic!("failed to decode");
/// }
/// ```
pub struct Abomonated<T, S: DerefMut<Target=[u8]>> {
    phantom: PhantomData<T>,
    decoded: S,
}

impl<T: Abomonation, S: DerefMut<Target=[u8]>> Abomonated<T, S> {

    /// Attempts to create decoded data from owned mutable bytes.
    ///
    /// This method will return `None` if it is unable to decode the data with
    /// type `T`.
    ///
    /// #Examples
    ///
    /// ```
    /// use std::ops::Deref;
    /// use abomonation::{encode, decode};
    /// use abomonation::abomonated::Abomonated;
    ///
    /// // create some test data out of abomonation-approved types
    /// let vector = (0..256u64).map(|i| (i, format!("{}", i)))
    ///                         .collect::<Vec<_>>();
    ///
    /// // encode a Vec<(u64, String)> into a Vec<u8>
    /// let mut bytes = Vec::new();
    /// unsafe { encode(&vector, &mut bytes); }
    ///
    /// // attempt to decode `bytes` into a `&Vec<(u64, String)>`.
    /// let maybe_decoded = unsafe { Abomonated::<Vec<(u64, String)>,_>::new(bytes) };
    ///
    /// if let Some(decoded) = maybe_decoded {
    ///     // each `deref()` call is now just a pointer cast.
    ///     assert!(decoded.deref() == &vector);
    /// }
    /// else {
    ///     panic!("failed to decode");
    /// }
    /// ```
    ///
    /// #Safety
    ///
    /// The type `S` must have its bytes at a fixed location, which will
    /// not change if the `bytes: S` instance is moved. Good examples are
    /// `Vec<u8>` whereas bad examples are `[u8; 16]`.
    pub unsafe fn new(mut bytes: S) -> Option<Self> {

        // performs the underlying pointer correction, indicates success.
        let decoded = decode::<T>(bytes.deref_mut()).is_some();

        if decoded {
            Some(Abomonated {
                phantom: PhantomData,
                decoded: bytes,
            })
        }
        else {
            None
        }
    }
}

impl<T, S: DerefMut<Target=[u8]>> Abomonated<T, S> {
    pub fn as_bytes(&self) -> &[u8] {
        &self.decoded
    }
}

impl<T, S: DerefMut<Target=[u8]>> Deref for Abomonated<T, S> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &T {
        let result: &T = unsafe { transmute(self.decoded.get_unchecked(0)) };
        result
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str;
use std::task::{Context, Poll};

use futures::prelude::*;
use futures::Future;

use hyper::Client;

use tokio::net::UdpSocket;
use tokio::time::timeout;

use hyper::body::Bytes;

use crate::aio::Gateway;
use crate::common::{messages, parsing, SearchOptions};
use crate::errors::SearchError;

const MAX_RESPONSE_SIZE: usize = 1500;

/// Search for a gateway with the provided options
pub async fn search_gateway(options: SearchOptions) -> Result<Gateway, SearchError> {
    // Create socket for future calls
    let socket = UdpSocket::bind(&options.bind_addr).await?;

    let search_future = SearchFuture::search(socket, options.broadcast_address).await?;
    // Create future and issue request
    match options.timeout {
        Some(t) => timeout(t, search_future).await?,
        None => search_future.await,
    }
}

pub struct SearchFuture {
    socket: UdpSocket,
    pending: HashMap<SocketAddr, RequestState>,
    state: SearchState,
}

enum RequestState {
    Connecting(String, Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>),
    Done(String),
    Error,
}

enum SearchState {
    Start,
    AwaitingConnections,
    AwaitingSchemas(
        Gateway,
        Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>,
    ),
    End,
}

impl SearchFuture {
    // Create a new search
    async fn search(mut socket: UdpSocket, addr: SocketAddr) -> Result<SearchFuture, SearchError> {
        debug!(
            "sending broadcast request to: {} on interface: {:?}",
            addr,
            socket.local_addr()
        );
        socket.send_to(messages::SEARCH_REQUEST.as_bytes(), &addr).await?;
        Ok(SearchFuture {
            socket,
            pending: HashMap::new(),
            state: SearchState::Start,
        })
    }

    // Handle a UDP response message
    fn handle_broadcast_resp(from: SocketAddr, data: &[u8]) -> Result<(SocketAddr, String), SearchError> {
        debug!("handling broadcast response from: {}", from);

        // Convert response to text
        let text = str::from_utf8(&data).map_err(|e| SearchError::from(e))?;

        // Parse socket address and path
        let (addr, root_url) = parsing::parse_search_result(text)?;

        Ok((SocketAddr::V4(addr), root_url))
    }

    // Issue a control URL request over HTTP using the provided
    fn request_control_url(
        addr: SocketAddr,
        path: String,
    ) -> Result<Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>, SearchError> {
        let client = Client::new();

        let uri = match format!("http://{}{}", addr, path).parse() {
            Ok(uri) => uri,
            Err(err) => return Err(SearchError::from(err)),
        };

        debug!("requesting control url from: {}", uri);

        Ok(Box::pin(
            client
                .get(uri)
                .and_then(|resp| hyper::body::to_bytes(resp.into_body()))
                .map_err(|e| SearchError::from(e)),
        ))
    }

    // Process a control response to extract the control URL
    fn handle_control_resp(addr: SocketAddr, resp: Bytes) -> Result<(String, String), SearchError> {
        debug!("handling control response from: {}", addr);

        // Create a cursor over the response data
        let c = std::io::Cursor::new(&resp);

        // Parse control URL out of body
        let urls = parsing::parse_control_urls(c)?;

        Ok(urls)
    }

    fn request_schemas(
        addr: SocketAddr,
        control_schema_url: String,
    ) -> Result<Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>, SearchError> {
        let client = Client::new();
        let uri = match format!("http://{}{}", addr, control_schema_url).parse() {
            Ok(uri) => uri,
            Err(err) => return Err(SearchError::from(err)),
        };

        Ok(Box::pin(
            client
                .get(uri)
                .and_then(|resp| hyper::body::to_bytes(resp.into_body()))
                .map_err(|e| SearchError::from(e)),
        ))
    }

    fn handle_schema_resp(resp: Bytes) -> Result<HashMap<String, Vec<String>>, SearchError> {
        debug!("handling schema response");
        let c = std::io::Cursor::new(&resp);
        parsing::parse_schemas(c)
    }
}

impl Future for SearchFuture {
    type Output = Result<Gateway, SearchError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<Gateway, SearchError>> {
        'state_check: loop {
            match &mut self.state {
                SearchState::Start => self.state = SearchState::AwaitingConnections,

                SearchState::AwaitingConnections => {
                    // Poll for (and handle) incoming messages
                    let mut buff = [0u8; MAX_RESPONSE_SIZE];
                    let resp = self.socket.poll_recv_from(cx, &mut buff);
                    if let Poll::Ready(Ok((n, from))) = resp {
                        // Try handle response messages
                        if let Ok((addr, root_url)) = Self::handle_broadcast_resp(from, &buff[0..n]) {
                            if !self.pending.contains_key(&addr) {
                                debug!("received broadcast response from: {}", from);

                                // Issue control request
                                match Self::request_control_url(addr, root_url.clone()) {
                                    // Store pending requests
                                    Ok(f) => {
                                        self.pending.insert(addr, RequestState::Connecting(root_url, f));
                                    }
                                    Err(e) => return Poll::Ready(Err(e)),
                                }
                            } else {
                                debug!("received duplicate broadcast response from: {}, dropping", from);
                            }
                        }
                    }
                    if let Poll::Ready(Err(err)) = resp {
                        return Poll::Ready(Err(err.into()));
                    }

                    // Poll on any outstanding control requests
                    for (addr, state) in &mut self.pending {
                        // Poll if we're in the connecting state
                        let (root_url, resp) = {
                            let (root_url, c) = match state {
                                RequestState::Connecting(root_url, c) => (root_url, c),
                                _ => continue,
                            };

                            match c.as_mut().poll(cx)? {
                                Poll::Ready(resp) => (root_url.clone(), resp),
                                _ => continue,
                            }
                        };

                        // Handle any responses
                        if let Ok((control_schema_url, control_url)) = Self::handle_control_resp(*addr, resp) {
                            debug!("received control url from: {} (url: {})", addr, control_url);
                            let c = match Self::request_schemas(*addr, control_schema_url.clone()) {
                                Ok(c) => c,
                                Err(e) => return Poll::Ready(Err(e)),
                            };
                            *state = RequestState::Done(control_url.clone());

                            match addr {
                                SocketAddr::V4(a) => {
                                    let g = Gateway {
                                        addr: *a,
                                        root_url,
                                        control_url,
                                        control_schema_url,
                                        control_schema: HashMap::new(),
                                    };
                                    self.state = SearchState::AwaitingSchemas(g, c);
                                    continue 'state_check;
                                }
                                _ => warn!("unsupported IPv6 gateway response from addr: {}", addr),
                            }
                        } else {
                            *state = RequestState::Error;
                        }
                    }

                    return Poll::Pending;
                }

                SearchState::AwaitingSchemas(g, c) => {
                    let resp = {
                        match c.as_mut().poll(cx)? {
                            Poll::Ready(resp) => resp,
                            Poll::Pending => return Poll::Pending,
                        }
                    };

                    let control_schema = match Self::handle_schema_resp(resp) {
                        Ok(control_schema) => control_schema,
                        Err(e) => return Poll::Ready(Err(e)),
                    };

                    let g = Gateway {
                        control_schema,
                        ..g.clone()
                    };
                    self.state = SearchState::End;
                    return Poll::Ready(Ok(g));
                }

                SearchState::End => panic!("Calling poll on a finished Future"),
            }
        }
    }
}
<commit_msg>Rewrite aio search without impl Future<commit_after>use std::collections::HashMap;
use std::net::SocketAddr;

use futures::prelude::*;
use hyper::Client;
use tokio::net::UdpSocket;
use tokio::time::timeout;

use crate::aio::Gateway;
use crate::common::{messages, parsing, SearchOptions};
use crate::errors::SearchError;

const MAX_RESPONSE_SIZE: usize = 1500;

/// Search for a gateway with the provided options
pub async fn search_gateway(options: SearchOptions) -> Result<Gateway, SearchError> {
    // Create socket for future calls
    let mut socket = UdpSocket::bind(&options.bind_addr).await?;

    send_search_request(&mut socket, options.broadcast_address).await?;

    let search_response = receive_search_response(&mut socket);

    // Receive search response, optionally with a timeout
    let (response_body, from) = match options.timeout {
        Some(t) => timeout(t, search_response).await?,
        None => search_response.await,
    }?;

    let (addr, root_url) = handle_broadcast_resp(&from, &response_body)?;

    let (control_schema_url, control_url) = get_control_urls(&addr, &root_url).await?;
    let control_schema = get_control_schemas(&addr, &control_schema_url).await?;

    let addr = match addr {
        SocketAddr::V4(a) => Ok(a),
        _ => {
            warn!("unsupported IPv6 gateway response from addr: {}", addr);
            Err(SearchError::InvalidResponse)
        }
    }?;

    Ok(Gateway {
        addr,
        root_url,
        control_url,
        control_schema_url,
        control_schema,
    })
}

// Create a new search
async fn send_search_request(socket: &mut UdpSocket, addr: SocketAddr) -> Result<(), SearchError> {
    debug!(
        "sending broadcast request to: {} on interface: {:?}",
        addr,
        socket.local_addr()
    );
    socket
        .send_to(messages::SEARCH_REQUEST.as_bytes(), &addr)
        .map_ok(|_| ())
        .map_err(|e| SearchError::from(e))
        .await
}

async fn receive_search_response(socket: &mut UdpSocket) -> Result<(Vec<u8>, SocketAddr), SearchError> {
    let mut buff = [0u8; MAX_RESPONSE_SIZE];
    let (n, from) = socket.recv_from(&mut buff).map_err(|e| SearchError::from(e)).await?;
    debug!("received broadcast response from: {}", from);
    Ok((buff[..n].to_vec(), from))
}

// Handle a UDP response message
fn handle_broadcast_resp(from: &SocketAddr, data: &[u8]) -> Result<(SocketAddr, String), SearchError> {
    debug!("handling broadcast response from: {}", from);

    // Convert response to text
    let text = std::str::from_utf8(&data).map_err(|e| SearchError::from(e))?;

    // Parse socket address and path
    let (addr, root_url) = parsing::parse_search_result(text)?;

    Ok((SocketAddr::V4(addr), root_url))
}

async fn get_control_urls(addr: &SocketAddr, path: &String) -> Result<(String, String), SearchError> {
    let uri = match format!("http://{}{}", addr, path).parse() {
        Ok(uri) => uri,
        Err(err) => return Err(SearchError::from(err)),
    };

    debug!("requesting control url from: {}", uri);
    let client = Client::new();
    let resp = hyper::body::to_bytes(client.get(uri).await?.into_body())
        .map_err(|e| SearchError::from(e))
        .await?;

    debug!("handling control response from: {}", addr);
    let c = std::io::Cursor::new(&resp);
    parsing::parse_control_urls(c)
}

async fn get_control_schemas(
    addr: &SocketAddr,
    control_schema_url: &String,
) -> Result<HashMap<String, Vec<String>>, SearchError> {
    let uri = match format!("http://{}{}", addr, control_schema_url).parse() {
        Ok(uri) => uri,
        Err(err) => return Err(SearchError::from(err)),
    };

    debug!("requesting control schema from: {}", uri);
    let client = Client::new();
    let resp = hyper::body::to_bytes(client.get(uri).await?.into_body())
        .map_err(|e| SearchError::from(e))
        .await?;

    debug!("handling schema response from: {}", addr);
    let c = std::io::Cursor::new(&resp);
    parsing::parse_schemas(c)
}
<|endoftext|>


<commit_before>use expectest::prelude::*;
use super::*;
use super::match_request;
use pact_matching::models::{Interaction, Request, Response, OptionalBody};
use pact_matching::Mismatch;

#[test]
fn match_request_returns_a_match_for_identical_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone()];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_not_found_for_no_interactions() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestNotFound(request)));
}

#[test]
fn match_request_returns_a_match_for_multiple_identical_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(),
        Interaction { description: s!("test2"), provider_state: None,
            request: request.clone(), response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_match_for_multiple_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let request2 = Request { method: s!("POST"), path: s!("/post"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(),
        Interaction { description: s!("test2"), provider_state: None,
            request: request2.clone(), response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_mismatch_for_incorrect_request() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result.match_key()).to(be_equal_to(s!("Request-Mismatch")));
}

#[test]
fn match_request_returns_request_not_found_if_method_or_path_do_not_match() {
    let request = Request { method: s!("GET"), path: s!("/path"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("POST"), path: s!("/otherpath"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestNotFound(request)));
}

#[test]
fn match_request_returns_the_most_appropriate_mismatch_for_multiple_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Present(s!("This is a body")), matching_rules: None };
    let request2 = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Present(s!("This is a body")), matching_rules: None };
    let request3 = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interaction2 = Interaction { description: s!("test2"), provider_state: None,
            request: request2.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(), interaction2.clone()];
    let result = match_request(&request3, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMismatch(interaction2,
        vec![Mismatch::BodyMismatch { path: s!("/"), expected: Some(s!("This is a body")), actual: None,
        mismatch: s!("Expected body \'This is a body\' but was missing") }])));
}
<commit_msg>added some V2 matcher tests<commit_after>use expectest::prelude::*;
use super::*;
use super::match_request;
use pact_matching::models::{Interaction, Request, Response, OptionalBody};
use pact_matching::Mismatch;

#[test]
fn match_request_returns_a_match_for_identical_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone()];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_not_found_for_no_interactions() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestNotFound(request)));
}

#[test]
fn match_request_returns_a_match_for_multiple_identical_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(),
        Interaction { description: s!("test2"), provider_state: None,
            request: request.clone(), response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_match_for_multiple_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let request2 = Request { method: s!("POST"), path: s!("/post"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(),
        Interaction { description: s!("test2"), provider_state: None,
            request: request2.clone(), response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_returns_a_mismatch_for_incorrect_request() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result.match_key()).to(be_equal_to(s!("Request-Mismatch")));
}

#[test]
fn match_request_returns_request_not_found_if_method_or_path_do_not_match() {
    let request = Request { method: s!("GET"), path: s!("/path"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("POST"), path: s!("/otherpath"), query: None,
        headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interactions = vec![Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() }];
    let result = match_request(&request, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestNotFound(request)));
}

#[test]
fn match_request_returns_the_most_appropriate_mismatch_for_multiple_requests() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: None, body: OptionalBody::Present(s!("This is a body")), matching_rules: None };
    let request2 = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Present(s!("This is a body")), matching_rules: None };
    let request3 = Request { method: s!("GET"), path: s!("/"), query: Some(hashmap!{
        s!("QueryA") => vec![s!("Value A")]
        }), headers: None, body: OptionalBody::Missing, matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: request.clone(), response: response.clone() };
    let interaction2 = Interaction { description: s!("test2"), provider_state: None,
            request: request2.clone(), response: response.clone() };
    let interactions = vec![interaction.clone(), interaction2.clone()];
    let result = match_request(&request3, &interactions);
    expect!(result).to(be_equal_to(MatchResult::RequestMismatch(interaction2,
        vec![Mismatch::BodyMismatch { path: s!("/"), expected: Some(s!("This is a body")), actual: None,
        mismatch: s!("Expected body \'This is a body\' but was missing") }])));
}

#[test]
fn match_request_supports_v2_matchers() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: Some(hashmap!{ s!("Content-Type") => s!("application/json") }), body: OptionalBody::Present(
            s!(r#"
            {
                "a": 100,
                "b": "one hundred"
            }
            "#)
        ), matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: Some(hashmap!{ s!("Content-Type") => s!("application/json") }),
        body: OptionalBody::Present(
            s!(r#"
            {
                "a": 1000,
                "b": "One Thousand"
            }
            "#)
        ), matching_rules: Some(hashmap!{
            s!("$.body.*") => hashmap!{ s!("match") => s!("type") }
        }) };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() };
    let result = match_request(&request, &vec![interaction.clone()]);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}

#[test]
fn match_request_supports_v2_matchers_with_xml() {
    let request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: Some(hashmap!{ s!("Content-Type") => s!("application/xml") }), body: OptionalBody::Present(
            s!(r#"<?xml version="1.0" encoding="UTF-8"?>
            <foo>hello<bar/>world</foo>
            "#)
        ), matching_rules: None };
    let response = Response { status: 200, headers: None, body: OptionalBody::Missing, matching_rules: None };
    let expected_request = Request { method: s!("GET"), path: s!("/"), query: None,
        headers: Some(hashmap!{ s!("Content-Type") => s!("application/xml") }),
        body: OptionalBody::Present(
            s!(r#"<?xml version="1.0" encoding="UTF-8"?>
            <foo>hello<bar/>mars </foo>
            "#)
        ), matching_rules: Some(hashmap!{
            s!("$.body.foo['#text']") => hashmap!{ s!("match") => s!("regex"), s!("regex") => s!("[a-z]+") }
        }) };
    let interaction = Interaction { description: s!("test"), provider_state: None,
        request: expected_request, response: response.clone() };
    let result = match_request(&request, &vec![interaction.clone()]);
    expect!(result).to(be_equal_to(MatchResult::RequestMatch(interaction)));
}
<|endoftext|>


<commit_before>//! Supply the command line.

// Standard library
use std::path::PathBuf;

// Third party
use clap::Clap;

#[derive(Clap)]
#[clap(
    name = "Lightning (lx)",
    about = "A very fast, very opinionated static site generator",
    version = "1.0",
    author = "Chris Krycho <hello@@chriskrycho.com>",
    setting = clap::AppSettings::ArgRequiredElseHelp,
    setting = clap::AppSettings::ColoredHelp
)]
pub(crate) enum Command {
    /// Build the site.
    #[clap(name = "build")]
    Build {
        /// The root of the site (if different from the current directory).
        site_directory: Option<PathBuf>,
    },
}

impl Command {
    pub(crate) fn cli() -> Self {
        Self::parse()
    }
}
<commit_msg>Global clap color settings<commit_after>//! Supply the command line.

// Standard library
use std::path::PathBuf;

// Third party
use clap::Clap;

#[derive(Clap)]
#[clap(
    name = "Lightning (lx)",
    about = "A very fast, very opinionated static site generator",
    version = "1.0",
    author = "Chris Krycho <hello@@chriskrycho.com>",
    setting = clap::AppSettings::ArgRequiredElseHelp,
    global_setting = clap::AppSettings::ColoredHelp,
    global_setting = clap::AppSettings::ColorAuto
)]
pub(crate) enum Command {
    /// Build the site.
    #[clap(name = "build")]
    Build {
        /// The root of the site (if different from the current directory).
        site_directory: Option<PathBuf>,
    },
}

impl Command {
    pub(crate) fn cli() -> Self {
        Self::parse()
    }
}
<|endoftext|>


<commit_before>use std::u32;
use std::sync::Arc;
use parking_lot::RwLock;
use elements::MemoryType;
use interpreter::Error;
use interpreter::module::check_limits;

/// Linear memory page size.
pub const LINEAR_MEMORY_PAGE_SIZE: u32 = 65536;
/// Maximal number of pages.
const LINEAR_MEMORY_MAX_PAGES: u32 = 65536;

/// Linear memory instance.
pub struct MemoryInstance {
	/// Linear memory buffer.
	buffer: RwLock<Vec<u8>>,
	/// Maximum buffer size.
	maximum_size: u32,
}

struct CheckedRegion<'a, B: 'a> where B: ::std::ops::Deref<Target=Vec<u8>> {
	buffer: &'a B,
	offset: usize,
	size: usize,
}

impl<'a, B: 'a> CheckedRegion<'a, B> where B: ::std::ops::Deref<Target=Vec<u8>> {
	fn range(&self) -> ::std::ops::Range<usize> {
		self.offset..self.offset+self.size
	}

	fn slice(&self) -> &[u8] {
		&self.buffer[self.range()]
	}
}

impl MemoryInstance {
	/// Create new linear memory instance.
	pub fn new(memory_type: &MemoryType) -> Result<Arc<Self>, Error> {
		check_limits(memory_type.limits())?;

		if let Some(maximum_pages) = memory_type.limits().maximum() {
			if maximum_pages > LINEAR_MEMORY_MAX_PAGES {
				return Err(Error::Memory(format!("memory size must be at most 65536 pages")));
			}
		}

		let memory = MemoryInstance {
			buffer: RwLock::new(Vec::new()), // TODO: with_capacity
			maximum_size: memory_type.limits().maximum()
				.map(|s| s.saturating_mul(LINEAR_MEMORY_PAGE_SIZE))
				.unwrap_or(u32::MAX),
		};
		if memory.grow(memory_type.limits().initial())? == u32::MAX {
			return Err(Error::Memory(format!("error initializing {}-pages linear memory region", memory_type.limits().initial())));
		}
		Ok(Arc::new(memory))
	}

	/// Return linear memory size (in pages).
	pub fn size(&self) -> u32 {
		self.buffer.read().len() as u32 / LINEAR_MEMORY_PAGE_SIZE
	}

	/// Get data at given offset.
	pub fn get(&self, offset: u32, size: usize) -> Result<Vec<u8>, Error> {
		let buffer = self.buffer.read();
		let region = self.checked_region(&buffer, offset as usize, size)?;

		Ok(region.slice().to_vec())
	}

	/// Set data at given offset.
	pub fn set(&self, offset: u32, value: &[u8]) -> Result<(), Error> {
		let mut buffer = self.buffer.write();
		let range = self.checked_region(&buffer, offset as usize, value.len())?.range();

		buffer[range].copy_from_slice(value);

		Ok(())
	}

	/// Increases the size of the linear memory by given number of pages.
	/// Returns -1 if allocation fails or previous memory size, if succeeds.
	pub fn grow(&self, pages: u32) -> Result<u32, Error> {
		let mut buffer = self.buffer.write();
		let old_size = buffer.len() as u32;
		match pages.checked_mul(LINEAR_MEMORY_PAGE_SIZE).and_then(|bytes| old_size.checked_add(bytes)) {
			None => Ok(u32::MAX),
			Some(new_size) if new_size > self.maximum_size => Ok(u32::MAX),
			Some(new_size) => {
				buffer.extend(vec![0; (new_size - old_size) as usize]);
				Ok(old_size / LINEAR_MEMORY_PAGE_SIZE)
			},
		}
	}

	fn checked_region<'a, B>(&self, buffer: &'a B, offset: usize, size: usize) -> Result<CheckedRegion<'a, B>, Error>
		where B: ::std::ops::Deref<Target=Vec<u8>>
	{
		let end = offset.checked_add(size)
			.ok_or(Error::Memory(format!("trying to access memory block of size {} from offset {}", size, offset)))?;

		if end > buffer.len() {
			return Err(Error::Memory(format!("trying to access region [{}..{}] in memory [0..{}]", offset, end, buffer.len())));
		}

		Ok(CheckedRegion {
			buffer: buffer,
			offset: offset,
			size: size,
		})
	}

	/// Copy memory region
	pub fn copy(&self, src_offset: usize, dst_offset: usize, len: usize) -> Result<(), Error> {
		let buffer = self.buffer.write();

		let read_region = self.checked_region(&buffer, src_offset, len)?;
		let write_region = self.checked_region(&buffer, dst_offset, len)?;

		unsafe { ::std::ptr::copy(
			buffer[read_region.range()].as_ptr(),
			buffer[write_region.range()].as_ptr() as *mut _,
			len,
		)}

		Ok(())
	}
}
<commit_msg>also mem zero<commit_after>use std::u32;
use std::sync::Arc;
use parking_lot::RwLock;
use elements::MemoryType;
use interpreter::Error;
use interpreter::module::check_limits;

/// Linear memory page size.
pub const LINEAR_MEMORY_PAGE_SIZE: u32 = 65536;
/// Maximal number of pages.
const LINEAR_MEMORY_MAX_PAGES: u32 = 65536;

/// Linear memory instance.
pub struct MemoryInstance {
	/// Linear memory buffer.
	buffer: RwLock<Vec<u8>>,
	/// Maximum buffer size.
	maximum_size: u32,
}

struct CheckedRegion<'a, B: 'a> where B: ::std::ops::Deref<Target=Vec<u8>> {
	buffer: &'a B,
	offset: usize,
	size: usize,
}

impl<'a, B: 'a> CheckedRegion<'a, B> where B: ::std::ops::Deref<Target=Vec<u8>> {
	fn range(&self) -> ::std::ops::Range<usize> {
		self.offset..self.offset+self.size
	}

	fn slice(&self) -> &[u8] {
		&self.buffer[self.range()]
	}
}

impl MemoryInstance {
	/// Create new linear memory instance.
	pub fn new(memory_type: &MemoryType) -> Result<Arc<Self>, Error> {
		check_limits(memory_type.limits())?;

		if let Some(maximum_pages) = memory_type.limits().maximum() {
			if maximum_pages > LINEAR_MEMORY_MAX_PAGES {
				return Err(Error::Memory(format!("memory size must be at most 65536 pages")));
			}
		}

		let memory = MemoryInstance {
			buffer: RwLock::new(Vec::new()), // TODO: with_capacity
			maximum_size: memory_type.limits().maximum()
				.map(|s| s.saturating_mul(LINEAR_MEMORY_PAGE_SIZE))
				.unwrap_or(u32::MAX),
		};
		if memory.grow(memory_type.limits().initial())? == u32::MAX {
			return Err(Error::Memory(format!("error initializing {}-pages linear memory region", memory_type.limits().initial())));
		}
		Ok(Arc::new(memory))
	}

	/// Return linear memory size (in pages).
	pub fn size(&self) -> u32 {
		self.buffer.read().len() as u32 / LINEAR_MEMORY_PAGE_SIZE
	}

	/// Get data at given offset.
	pub fn get(&self, offset: u32, size: usize) -> Result<Vec<u8>, Error> {
		let buffer = self.buffer.read();
		let region = self.checked_region(&buffer, offset as usize, size)?;

		Ok(region.slice().to_vec())
	}

	/// Set data at given offset.
	pub fn set(&self, offset: u32, value: &[u8]) -> Result<(), Error> {
		let mut buffer = self.buffer.write();
		let range = self.checked_region(&buffer, offset as usize, value.len())?.range();

		buffer[range].copy_from_slice(value);

		Ok(())
	}

	/// Increases the size of the linear memory by given number of pages.
	/// Returns -1 if allocation fails or previous memory size, if succeeds.
	pub fn grow(&self, pages: u32) -> Result<u32, Error> {
		let mut buffer = self.buffer.write();
		let old_size = buffer.len() as u32;
		match pages.checked_mul(LINEAR_MEMORY_PAGE_SIZE).and_then(|bytes| old_size.checked_add(bytes)) {
			None => Ok(u32::MAX),
			Some(new_size) if new_size > self.maximum_size => Ok(u32::MAX),
			Some(new_size) => {
				buffer.extend(vec![0; (new_size - old_size) as usize]);
				Ok(old_size / LINEAR_MEMORY_PAGE_SIZE)
			},
		}
	}

	fn checked_region<'a, B>(&self, buffer: &'a B, offset: usize, size: usize) -> Result<CheckedRegion<'a, B>, Error>
		where B: ::std::ops::Deref<Target=Vec<u8>>
	{
		let end = offset.checked_add(size)
			.ok_or(Error::Memory(format!("trying to access memory block of size {} from offset {}", size, offset)))?;

		if end > buffer.len() {
			return Err(Error::Memory(format!("trying to access region [{}..{}] in memory [0..{}]", offset, end, buffer.len())));
		}

		Ok(CheckedRegion {
			buffer: buffer,
			offset: offset,
			size: size,
		})
	}

	/// Copy memory region
	pub fn copy(&self, src_offset: usize, dst_offset: usize, len: usize) -> Result<(), Error> {
		let buffer = self.buffer.write();

		let read_region = self.checked_region(&buffer, src_offset, len)?;
		let write_region = self.checked_region(&buffer, dst_offset, len)?;

		unsafe { ::std::ptr::copy(
			buffer[read_region.range()].as_ptr(),
			buffer[write_region.range()].as_ptr() as *mut _,
			len,
		)}

		Ok(())
	}

	/// Zero memory region
	pub fn zero(&self, offset: usize, len: usize) -> Result<(), Error> {
		let mut buffer = self.buffer.write();

		let range = self.checked_region(&buffer, offset, len)?.range();
		for val in &mut buffer[range] { *val = 0 }
		Ok(())
	}
}
<|endoftext|>