<commit_before>use std::cmp;
use std::cmp::Ord;
use num::traits::Zero;
use num::traits::One;
use num::traits::ToPrimitive;

mod iterator;
mod test;

/// Marker trait for types we allow (namely, u8-u64)
pub trait HistogramCount : Ord + Zero + One + ToPrimitive + Copy {}

impl HistogramCount for u8 {}
impl HistogramCount for u16 {}
impl HistogramCount for u32 {}
impl HistogramCount for u64 {}

///
/// This struct essentially encapsulates the "instance variables" of the histogram
///
#[derive(Debug)]
pub struct SimpleHdrHistogram<T:HistogramCount> {
    leading_zeros_count_base: usize,
    sub_bucket_mask: u64,
    unit_magnitude: u32,
    sub_bucket_count: usize,
    sub_bucket_half_count: usize,
    sub_bucket_half_count_magnitude: u32,
    counts: Vec<T>,
    // is at most counts.len(), so i32 is plenty since counts scales exponentially
    normalizing_index_offset: i32,
    max_value: u64,
    min_non_zero_value: u64,
    unit_magnitude_mask: u64,
    total_count: u64,
}

pub trait HistogramBase<T: HistogramCount> {
    //TODO this block should be default impl of this trait
    fn record_single_value(&mut self, value: u64) -> Result<(), String>;
    fn counts_array_index(&self, value: u64) -> usize;
    // in the Java impl, the functions above/below have same name but are overloaded, which
    //  Rust does not allow, thus the name change
    fn counts_array_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> usize;
    fn get_bucket_index(&self, value: u64) -> usize;
    fn get_sub_bucket_index(&self, value: u64, bucket_index: usize) -> usize;
    fn get_count(&self) -> u64;
    fn get_max(&self) -> u64;
    fn get_min_non_zero(&self) -> u64;
    fn update_min_and_max(&mut self, value: u64);
    fn update_max_value(&mut self, value: u64);
    fn update_min_non_zero_value(&mut self, value: u64);
    fn get_count_at_value(&self, value: u64) -> Result<T, String>;
    fn get_value_at_percentile(&self, percentile: f64) -> u64;
    fn value_from_index(&self, index: usize) -> u64;
    // in the Java impl, the functions above/below have same name but are overloaded, which
    //  Rust does not allow, thus the name change
    fn value_from_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> u64;
    fn lowest_equivalent_value(&self, value: u64) -> u64;
    fn highest_equivalent_value(&self, value: u64) -> u64;
    fn next_non_equivalent_value(&self, value: u64) -> u64;
    fn size_of_equivalent_value_range(&self, value: u64) -> u64;
    fn get_mean(&self) -> f64;
    // end TODO

    fn increment_count_at_index(&mut self, index: usize) -> Result<(), String>;
    fn normalize_index(&self, index: usize, normalizing_index_offset: i32, array_length: usize) ->
        Result<usize, String>;
    fn increment_total_count(&mut self);
    fn get_count_at_index(&self, index: usize) -> Result<T, String>;

}

impl<T: HistogramCount> HistogramBase<T> for SimpleHdrHistogram<T> {

    fn get_mean(&self) -> f64 {

        if self.get_count() == 0 { 0.0 } else {
            //TODO stuff
            0.0
        }
    }

    fn next_non_equivalent_value(&self, value: u64) -> u64 {
        self.lowest_equivalent_value(value) + self.size_of_equivalent_value_range(value)
    }

    fn size_of_equivalent_value_range(&self, value: u64) -> u64 {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        let distance_to_next_value =
            1 << (self.unit_magnitude
                    + bucket_index as u32
                    + if sub_bucket_index >= self.sub_bucket_count {1} else {0});
        distance_to_next_value as u64
    }

    fn highest_equivalent_value(&self, value: u64) -> u64 {
        self.next_non_equivalent_value(value) - 1
    }

    fn lowest_equivalent_value(&self, value: u64) -> u64 {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        let this_value_base_level = self.value_from_index_sub(bucket_index, sub_bucket_index);
        this_value_base_level
    }

    fn value_from_index(&self, index: usize) -> u64 {
        let mut bucket_index = (index as u32 >> self.sub_bucket_half_count_magnitude) - 1;
        let mut sub_bucket_index = (index as u32 & (self.sub_bucket_half_count as u32 - 1)) + self.sub_bucket_half_count as u32;
        if bucket_index < 0 {
            sub_bucket_index -= self.sub_bucket_half_count as u32;
            bucket_index = 0;
        }
        self.value_from_index_sub(bucket_index as usize, sub_bucket_index as usize)
    }

    fn value_from_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> u64 {
        (sub_bucket_index as u64) << (bucket_index as u32 + self.unit_magnitude)
    }

    fn get_value_at_percentile(&self, percentile: f64) -> u64 {
        let requested_percentile = percentile.min(100.0);
        let mut count_at_percentile = (((requested_percentile / 100.0) * self.get_count() as f64) + 0.5) as u64;
        count_at_percentile = cmp::max(count_at_percentile, 1);
        let mut total_to_current_index: u64 = 0;
        for i in 0..self.counts.len() {
            let count_at_index = self.get_count_at_index(i as usize);
            match count_at_index {
                Ok(the_index) => {
                    // we only use u8 - u64 types, so this must always work
                    total_to_current_index += the_index.to_u64().unwrap();
                    if total_to_current_index >= count_at_percentile {
                        let value_at_index = self.value_from_index(i as usize);
                        return if percentile == 0.0 {
                            self.lowest_equivalent_value(value_at_index)
                        } else {
                            self.highest_equivalent_value(value_at_index)
                        }
                    }
                }
                Err(_) => { return 0 }
            }
        }
        0
    }

    fn get_count_at_index(&self, index: usize) -> Result<T, String> {
        let normalized_index =
            self.normalize_index(index, self.normalizing_index_offset, self.counts.len());
        match normalized_index {
            Ok(the_index) =>
                Ok(self.counts[the_index]),
            Err(err) =>
                Err(err)
        }
    }

    fn get_count_at_value(&self, value: u64) -> Result<T, String> {
        // TODO is it ok to just clamp to max value rathe than saying it's inexpressible?
        let index = cmp::min(cmp::max(0, self.counts_array_index(value)), self.counts.len() - 1);
        self.get_count_at_index(index)
    }

    fn increment_total_count(&mut self) {
        self.total_count += 1;
    }

    fn get_max(&self) -> u64 {
        self.max_value
    }

    fn get_min_non_zero(&self) -> u64 {
        self.min_non_zero_value
    }

    fn get_count(&self) -> u64 {
        self.total_count
    }

    fn update_max_value(&mut self, value: u64) {
        let internal_value = value | self.unit_magnitude_mask;
        self.max_value = internal_value;
    }

    fn update_min_non_zero_value(&mut self, value: u64) {
        if value <= self.unit_magnitude_mask {
            return
        }
        let internal_value = value & !self.unit_magnitude_mask;
        self.min_non_zero_value = internal_value;
    }

    fn update_min_and_max(&mut self, value: u64) {
        if value > self.max_value {
            self.update_max_value(value);
        }
        if value < self.min_non_zero_value && value != 0 {
            self.update_min_non_zero_value(value);
        }
    }

    fn record_single_value(&mut self, value: u64) -> Result<(), String> {
        let counts_index = self.counts_array_index(value);
            match self.increment_count_at_index(counts_index) {
                Ok(_) => {
                    self.update_min_and_max(value);
                    self.increment_total_count();
                    Ok(())
                }
                Err(err) => {
                    Err(String::from(format!("Could not increment count at index due to: {}", err)))
                }
            }
    }

    fn increment_count_at_index(&mut self, index: usize) -> Result<(), String> {
        let normalized_index =
        self.normalize_index(index, self.normalizing_index_offset, self.counts.len());
        match normalized_index {
            Ok(the_index) => {
                // TODO express exceeding the counts size as an error here?
                self.counts[the_index] = self.counts[the_index] + T::one();
                Ok(())
            }
            Err(err) =>
            Err(err)
        }
    }

    fn normalize_index(&self, index: usize, normalizing_index_offset: i32, array_length: usize) ->
    Result<usize, String> {
        match normalizing_index_offset {
            0 => Ok(index),
            _ =>
            if index > array_length {
                Err(String::from("index out of covered range"))
            } else {
                // indices are always pretty small since they scale to values exponentially
                let array_length_i32 : i32 = array_length as i32;
                let mut normalized_index: i32 = index as i32 - normalizing_index_offset;

                // Calculate unsigned remainder.
                // When offset is positive (right shifted), it's checked at the time of the
                // right shift op to not shift anything into the bottom half of the first
                // bucket or below. Therefore, we know that the offset is less than array
                // length. So, we can simply check for sum < 0 and add one array length.
                // When offset is negative (left shifted), it's also checked at left shift time
                // to ensure it won't cause anything to be shifted past the end of the array.
                // Similarly, we know that |offset| < array length, so if sum > array length,
                // we can just subtract 1 array length.
                if normalized_index < 0 {
                    normalized_index += array_length_i32
                } else if normalized_index >= array_length_i32 {
                    normalized_index -= array_length_i32
                }

                Ok(normalized_index as usize)
            }
        }
    }

    fn counts_array_index(&self, value: u64) -> usize {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        return self.counts_array_index_sub(bucket_index, sub_bucket_index);
    }

    fn get_sub_bucket_index(&self, value: u64, bucket_index: usize) -> usize {
        // safe cast: sub bucket indexes are at most 2 * 10^precision, so can fit in usize.
        // bucket_indexes are even smaller, so can certainly fit in u32.
        (value >> (bucket_index as u32 + self.unit_magnitude)) as usize
    }

    fn get_bucket_index(&self, value: u64) -> usize {
        let value_orred = value | self.sub_bucket_mask;
        self.leading_zeros_count_base - (value_orred.leading_zeros() as usize)
    }

    fn counts_array_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> usize {
        assert!(sub_bucket_index < self.sub_bucket_count);
        assert!(bucket_index == 0 || (sub_bucket_index >= self.sub_bucket_half_count));

        let bucket_base_index = (bucket_index + 1) << self.sub_bucket_half_count_magnitude;

        // offset_in_bucket can be negative for bucket 0.
        // these casts are safe: sub_bucket_index is at most sub_bucket_count, and sub_bucket_count
        // is at most 2 * 10^precision.
        let offset_in_bucket: i32 = sub_bucket_index as i32 - self.sub_bucket_half_count as i32;

        // buckets scale with 2^x * (sub bucket count), so bucket index could be at most the bit
        // length of the value datatype (e.g. 64 bits), and since sub bucket count is > 1 in
        // practice, it's even smaller. Thus, this case to signed is safe.

        let bucket_base_signed: i32 = bucket_base_index as i32;

        // this always works out to be non-negative: when offset_in_bucket is negative for bucket
        // 0, bucket_base_index is still at sub bucket half count, so the sum is positive.
        (bucket_base_signed + offset_in_bucket) as usize
    }
}

<commit_msg>Reorg functions so only public ones are in trait<commit_after>use std::cmp;
use std::cmp::Ord;
use num::traits::Zero;
use num::traits::One;
use num::traits::ToPrimitive;

mod iterator;
mod test;

/// Marker trait for types we allow (namely, u8-u64)
pub trait HistogramCount : Ord + Zero + One + ToPrimitive + Copy {}

impl HistogramCount for u8 {}
impl HistogramCount for u16 {}
impl HistogramCount for u32 {}
impl HistogramCount for u64 {}

///
/// This struct essentially encapsulates the "instance variables" of the histogram
///
#[derive(Debug)]
pub struct SimpleHdrHistogram<T:HistogramCount> {
    leading_zeros_count_base: usize,
    sub_bucket_mask: u64,
    unit_magnitude: u32,
    sub_bucket_count: usize,
    sub_bucket_half_count: usize,
    sub_bucket_half_count_magnitude: u32,
    counts: Vec<T>,
    // is at most counts.len(), so i32 is plenty since counts scales exponentially
    normalizing_index_offset: i32,
    max_value: u64,
    min_non_zero_value: u64,
    unit_magnitude_mask: u64,
    total_count: u64,
}

pub trait HistogramBase<T: HistogramCount> {
    fn record_single_value(&mut self, value: u64) -> Result<(), String>;

    fn get_count(&self) -> u64;
    fn get_count_at_value(&self, value: u64) -> Result<T, String>;

    fn get_max(&self) -> u64;
    fn get_min_non_zero(&self) -> u64;
    fn get_mean(&self) -> f64;

    fn get_value_at_percentile(&self, percentile: f64) -> u64;

    fn lowest_equivalent_value(&self, value: u64) -> u64;
    fn highest_equivalent_value(&self, value: u64) -> u64;
    fn next_non_equivalent_value(&self, value: u64) -> u64;
    fn size_of_equivalent_value_range(&self, value: u64) -> u64;
}

impl<T: HistogramCount> HistogramBase<T> for SimpleHdrHistogram<T> {

    fn get_mean(&self) -> f64 {

        if self.get_count() == 0 { 0.0 } else {
            //TODO stuff
            0.0
        }
    }

    fn next_non_equivalent_value(&self, value: u64) -> u64 {
        self.lowest_equivalent_value(value) + self.size_of_equivalent_value_range(value)
    }

    fn size_of_equivalent_value_range(&self, value: u64) -> u64 {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        let distance_to_next_value =
            1 << (self.unit_magnitude
                    + bucket_index as u32
                    + if sub_bucket_index >= self.sub_bucket_count {1} else {0});
        distance_to_next_value as u64
    }

    fn highest_equivalent_value(&self, value: u64) -> u64 {
        self.next_non_equivalent_value(value) - 1
    }

    fn lowest_equivalent_value(&self, value: u64) -> u64 {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        let this_value_base_level = self.value_from_index_sub(bucket_index, sub_bucket_index);
        this_value_base_level
    }

    fn get_value_at_percentile(&self, percentile: f64) -> u64 {
        let requested_percentile = percentile.min(100.0);
        let mut count_at_percentile = (((requested_percentile / 100.0) * self.get_count() as f64) + 0.5) as u64;
        count_at_percentile = cmp::max(count_at_percentile, 1);
        let mut total_to_current_index: u64 = 0;
        for i in 0..self.counts.len() {
            let count_at_index = self.get_count_at_index(i as usize);
            match count_at_index {
                Ok(the_index) => {
                    // we only use u8 - u64 types, so this must always work
                    total_to_current_index += the_index.to_u64().unwrap();
                    if total_to_current_index >= count_at_percentile {
                        let value_at_index = self.value_from_index(i as usize);
                        return if percentile == 0.0 {
                            self.lowest_equivalent_value(value_at_index)
                        } else {
                            self.highest_equivalent_value(value_at_index)
                        }
                    }
                }
                Err(_) => { return 0 }
            }
        }
        0
    }

    fn get_count_at_value(&self, value: u64) -> Result<T, String> {
        // TODO is it ok to just clamp to max value rathe than saying it's inexpressible?
        let index = cmp::min(cmp::max(0, self.counts_array_index(value)), self.counts.len() - 1);
        self.get_count_at_index(index)
    }

    fn get_max(&self) -> u64 {
        self.max_value
    }

    fn get_min_non_zero(&self) -> u64 {
        self.min_non_zero_value
    }

    fn get_count(&self) -> u64 {
        self.total_count
    }

    fn record_single_value(&mut self, value: u64) -> Result<(), String> {
        let counts_index = self.counts_array_index(value);
            match self.increment_count_at_index(counts_index) {
                Ok(_) => {
                    self.update_min_and_max(value);
                    self.increment_total_count();
                    Ok(())
                }
                Err(err) => {
                    Err(String::from(format!("Could not increment count at index due to: {}", err)))
                }
            }
    }

}

impl<T: HistogramCount> SimpleHdrHistogram<T> {

    fn value_from_index(&self, index: usize) -> u64 {
        let mut bucket_index = (index as u32 >> self.sub_bucket_half_count_magnitude) - 1;
        let mut sub_bucket_index = (index as u32 & (self.sub_bucket_half_count as u32 - 1)) + self.sub_bucket_half_count as u32;
        if bucket_index < 0 {
            sub_bucket_index -= self.sub_bucket_half_count as u32;
            bucket_index = 0;
        }
        self.value_from_index_sub(bucket_index as usize, sub_bucket_index as usize)
    }

    fn get_count_at_index(&self, index: usize) -> Result<T, String> {
        let normalized_index =
        self.normalize_index(index, self.normalizing_index_offset, self.counts.len());
        match normalized_index {
            Ok(the_index) =>
            Ok(self.counts[the_index]),
            Err(err) =>
            Err(err)
        }
    }

    fn value_from_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> u64 {
        (sub_bucket_index as u64) << (bucket_index as u32 + self.unit_magnitude)
    }

    fn increment_total_count(&mut self) {
        self.total_count += 1;
    }

    fn update_max_value(&mut self, value: u64) {
        let internal_value = value | self.unit_magnitude_mask;
        self.max_value = internal_value;
    }

    fn update_min_non_zero_value(&mut self, value: u64) {
        if value <= self.unit_magnitude_mask {
            return
        }
        let internal_value = value & !self.unit_magnitude_mask;
        self.min_non_zero_value = internal_value;
    }

    fn update_min_and_max(&mut self, value: u64) {
        if value > self.max_value {
            self.update_max_value(value);
        }
        if value < self.min_non_zero_value && value != 0 {
            self.update_min_non_zero_value(value);
        }
    }

    fn increment_count_at_index(&mut self, index: usize) -> Result<(), String> {
        let normalized_index =
        self.normalize_index(index, self.normalizing_index_offset, self.counts.len());
        match normalized_index {
            Ok(the_index) => {
                // TODO express exceeding the counts size as an error here?
                self.counts[the_index] = self.counts[the_index] + T::one();
                Ok(())
            }
            Err(err) =>
            Err(err)
        }
    }

    fn normalize_index(&self, index: usize, normalizing_index_offset: i32, array_length: usize) ->
    Result<usize, String> {
        match normalizing_index_offset {
            0 => Ok(index),
            _ =>
            if index > array_length {
                Err(String::from("index out of covered range"))
            } else {
                // indices are always pretty small since they scale to values exponentially
                let array_length_i32 : i32 = array_length as i32;
                let mut normalized_index: i32 = index as i32 - normalizing_index_offset;

                // Calculate unsigned remainder.
                // When offset is positive (right shifted), it's checked at the time of the
                // right shift op to not shift anything into the bottom half of the first
                // bucket or below. Therefore, we know that the offset is less than array
                // length. So, we can simply check for sum < 0 and add one array length.
                // When offset is negative (left shifted), it's also checked at left shift time
                // to ensure it won't cause anything to be shifted past the end of the array.
                // Similarly, we know that |offset| < array length, so if sum > array length,
                // we can just subtract 1 array length.
                if normalized_index < 0 {
                    normalized_index += array_length_i32
                } else if normalized_index >= array_length_i32 {
                    normalized_index -= array_length_i32
                }

                Ok(normalized_index as usize)
            }
        }
    }

    fn counts_array_index(&self, value: u64) -> usize {
        let bucket_index = self.get_bucket_index(value);
        let sub_bucket_index = self.get_sub_bucket_index(value, bucket_index);
        return self.counts_array_index_sub(bucket_index, sub_bucket_index);
    }

    fn get_sub_bucket_index(&self, value: u64, bucket_index: usize) -> usize {
        // safe cast: sub bucket indexes are at most 2 * 10^precision, so can fit in usize.
        // bucket_indexes are even smaller, so can certainly fit in u32.
        (value >> (bucket_index as u32 + self.unit_magnitude)) as usize
    }

    fn get_bucket_index(&self, value: u64) -> usize {
        let value_orred = value | self.sub_bucket_mask;
        self.leading_zeros_count_base - (value_orred.leading_zeros() as usize)
    }

    fn counts_array_index_sub(&self, bucket_index: usize, sub_bucket_index: usize) -> usize {
        assert!(sub_bucket_index < self.sub_bucket_count);
        assert!(bucket_index == 0 || (sub_bucket_index >= self.sub_bucket_half_count));

        let bucket_base_index = (bucket_index + 1) << self.sub_bucket_half_count_magnitude;

        // offset_in_bucket can be negative for bucket 0.
        // these casts are safe: sub_bucket_index is at most sub_bucket_count, and sub_bucket_count
        // is at most 2 * 10^precision.
        let offset_in_bucket: i32 = sub_bucket_index as i32 - self.sub_bucket_half_count as i32;

        // buckets scale with 2^x * (sub bucket count), so bucket index could be at most the bit
        // length of the value datatype (e.g. 64 bits), and since sub bucket count is > 1 in
        // practice, it's even smaller. Thus, this case to signed is safe.

        let bucket_base_signed: i32 = bucket_base_index as i32;

        // this always works out to be non-negative: when offset_in_bucket is negative for bucket
        // 0, bucket_base_index is still at sub bucket half count, so the sum is positive.
        (bucket_base_signed + offset_in_bucket) as usize
    }
}<|endoftext|>


<commit_before>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt::Result as FmtResult;
use std::fmt::Write;

use common::DirectiveName;
use intern::string_key::StringKey;
use intern::Lookup;

use crate::definitions::Directive;
use crate::definitions::*;

pub trait Schema {
    fn query_type(&self) -> Option<Type>;

    fn mutation_type(&self) -> Option<Type>;

    fn subscription_type(&self) -> Option<Type>;

    fn clientid_field(&self) -> FieldID;

    fn strongid_field(&self) -> FieldID;

    fn typename_field(&self) -> FieldID;

    fn fetch_token_field(&self) -> FieldID;

    fn is_fulfilled_field(&self) -> FieldID;

    fn get_type(&self, type_name: StringKey) -> Option<Type>;

    fn get_directive(&self, name: DirectiveName) -> Option<&Directive>;

    fn input_object(&self, id: InputObjectID) -> &InputObject;
    fn input_objects<'a>(&'a self) -> Box<dyn Iterator<Item = &'a InputObject> + 'a>;

    fn enum_(&self, id: EnumID) -> &Enum;
    fn enums<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Enum> + 'a>;

    fn scalar(&self, id: ScalarID) -> &Scalar;
    fn scalars<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Scalar> + 'a>;

    fn field(&self, id: FieldID) -> &Field;
    fn fields<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Field> + 'a>;

    fn object(&self, id: ObjectID) -> &Object;
    fn objects<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Object> + 'a>;

    fn union(&self, id: UnionID) -> ⋃
    fn unions<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Union> + 'a>;

    fn interface(&self, id: InterfaceID) -> &Interface;
    fn interfaces<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Interface> + 'a>;

    fn get_type_name(&self, type_: Type) -> StringKey;

    fn is_extension_type(&self, type_: Type) -> bool;

    fn is_string(&self, type_: Type) -> bool;

    fn is_id(&self, type_: Type) -> bool;

    fn named_field(&self, parent_type: Type, name: StringKey) -> Option<FieldID>;

    fn unchecked_argument_type_sentinel(&self) -> &TypeReference<Type>;

    fn snapshot_print(&self) -> String;

    // Is `maybe_subtype` usable as — but not equal to — `super_type`.
    // Examples:
    // is_type_strict_subtype_of(Int!, Int) => true, Int! is usable where an Int is expected
    // is_type_strict_subtype_of(Int, Int) => false, same type
    // is_type_strict_subtype_of(InterfaceThatExtendsNode, Node) => true
    // is_type_strict_subtype_of(Node, Node) => false, same type
    fn is_type_strict_subtype_of(
        &self,
        maybe_subtype: &TypeReference<Type>,
        super_type: &TypeReference<Type>,
    ) -> bool {
        match (maybe_subtype, super_type) {
            (TypeReference::NonNull(of_sub), TypeReference::NonNull(of_super)) => {
                self.is_type_strict_subtype_of(of_sub, of_super)
            }
            // If supertype is non-null, maybeSubType must be non-nullable too
            (_, TypeReference::NonNull(_)) => false,
            // If supertype is nullable, maybeSubType may be non-nullable or nullable
            (TypeReference::NonNull(of_sub), _) => {
                // `T!` is a strict subtype of `U` so long as T is a (non-strict) subtype of U
                self.is_type_subtype_of(of_sub, super_type)
            }
            (TypeReference::List(of_sub), TypeReference::List(of_super)) => {
                self.is_type_strict_subtype_of(of_sub, of_super)
            }
            // If supertype is a list, maybeSubType must be a list too
            (_, TypeReference::List(_)) => false,
            // If supertype is not a list, maybeSubType must also not be a list
            (TypeReference::List(_), _) => false,
            (TypeReference::Named(named_subtype), TypeReference::Named(named_supertype)) => {
                self.is_named_type_strict_subtype_of(*named_subtype, *named_supertype)
            }
        }
    }

    fn is_type_subtype_of(
        &self,
        maybe_subtype: &TypeReference<Type>,
        super_type: &TypeReference<Type>,
    ) -> bool {
        match (maybe_subtype, super_type) {
            (TypeReference::NonNull(of_sub), TypeReference::NonNull(of_super)) => {
                self.is_type_subtype_of(of_sub, of_super)
            }
            // If supertype is non-null, maybeSubType must be non-nullable too
            (_, TypeReference::NonNull(_)) => false,
            // If supertype is nullable, maybeSubType may be non-nullable or nullable
            (TypeReference::NonNull(of_sub), _) => self.is_type_subtype_of(of_sub, super_type),
            (TypeReference::List(of_sub), TypeReference::List(of_super)) => {
                self.is_type_subtype_of(of_sub, of_super)
            }
            // If supertype is a list, maybeSubType must be a list too
            (_, TypeReference::List(_)) => false,
            // If supertype is not a list, maybeSubType must also not be a list
            (TypeReference::List(_), _) => false,
            (TypeReference::Named(named_subtype), TypeReference::Named(named_supertype)) => {
                self.is_named_type_subtype_of(*named_subtype, *named_supertype)
            }
        }
    }

    // Is `maybe_subtype` usable as — but not equal to — `super_type`.
    // Examples:
    // is_type_strict_subtype_of(InterfaceThatExtendsNode, Node) => true
    // is_type_strict_subtype_of(Node, Node) => false, same type
    fn is_named_type_strict_subtype_of(&self, maybe_subtype: Type, super_type: Type) -> bool {
        match (maybe_subtype, super_type) {
            (Type::Object(sub_id), Type::Interface(super_id)) => {
                // does object implement the interface
                let object = self.object(sub_id);
                object.interfaces.contains(&super_id)
            }
            (Type::Object(sub_id), Type::Union(super_id)) => {
                // is object a member of the union
                let union = self.union(super_id);
                union.members.contains(&sub_id)
            }
            (Type::Interface(sub_id), Type::Interface(super_id)) => {
                // does sub interface implement the super interface
                let interface = self.interface(sub_id);
                interface.interfaces.contains(&super_id)
            }
            _ => false,
        }
    }

    fn is_named_type_subtype_of(&self, maybe_subtype: Type, super_type: Type) -> bool {
        match (maybe_subtype, super_type) {
            (Type::Object(sub_id), Type::Interface(super_id)) => {
                // does object implement the interface
                let object = self.object(sub_id);
                object.interfaces.contains(&super_id)
            }
            (Type::Object(sub_id), Type::Union(super_id)) => {
                // is object a member of the union
                let union = self.union(super_id);
                union.members.contains(&sub_id)
            }
            (Type::Interface(sub_id), Type::Interface(super_id)) => {
                // does interface implement the interface
                let interface = self.interface(sub_id);
                sub_id == super_id || interface.interfaces.contains(&super_id)
            }
            _ => maybe_subtype == super_type,
        }
    }

    fn are_overlapping_types(&self, a: Type, b: Type) -> bool {
        fn overlapping_objects(a: &[ObjectID], b: &[ObjectID]) -> bool {
            a.iter().any(|item| b.contains(item))
        }

        fn overlapping_interfaces(a: &[InterfaceID], b: &[InterfaceID]) -> bool {
            a.iter().any(|item| b.contains(item))
        }

        if a == b {
            return true;
        };
        match (a, b) {
            (Type::Interface(a), Type::Interface(b)) => {
                overlapping_objects(
                    &self.interface(a).implementing_objects,
                    &self.interface(b).implementing_objects,
                ) || overlapping_interfaces(
                    &self.interface(a).implementing_interfaces,
                    &self.interface(b).implementing_interfaces,
                )
            }

            (Type::Union(a), Type::Union(b)) => {
                overlapping_objects(&self.union(a).members, &self.union(b).members)
            }

            (Type::Union(union_id), Type::Interface(interface_id))
            | (Type::Interface(interface_id), Type::Union(union_id)) => overlapping_objects(
                &self.union(union_id).members,
                &self.interface(interface_id).implementing_objects,
            ),

            (Type::Interface(interface_id), Type::Object(object_id))
            | (Type::Object(object_id), Type::Interface(interface_id)) => self
                .interface(interface_id)
                .implementing_objects
                .contains(&object_id),

            (Type::Union(union_id), Type::Object(object_id))
            | (Type::Object(object_id), Type::Union(union_id)) => {
                self.union(union_id).members.contains(&object_id)
            }

            _ => false,
        }
    }

    fn write_type_string(&self, writer: &mut String, type_: &TypeReference<Type>) -> FmtResult {
        match type_ {
            TypeReference::Named(inner) => {
                write!(writer, "{}", self.get_type_name(*inner).lookup())
            }
            TypeReference::NonNull(of) => {
                self.write_type_string(writer, of)?;
                write!(writer, "!")
            }
            TypeReference::List(of) => {
                write!(writer, "[")?;
                self.write_type_string(writer, of)?;
                write!(writer, "]")
            }
        }
    }

    fn get_type_string(&self, type_: &TypeReference<Type>) -> String {
        let mut result = String::new();
        self.write_type_string(&mut result, type_).unwrap();
        result
    }

    fn is_extension_directive(&self, name: DirectiveName) -> bool {
        if let Some(directive) = self.get_directive(name) {
            directive.is_extension
        } else {
            panic!("Unknown directive {}.", name.0.lookup())
        }
    }
}
<commit_msg>Add utilities for accessing schema directives on a given type<commit_after>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use std::fmt::Result as FmtResult;
use std::fmt::Write;

use common::DirectiveName;
use intern::string_key::StringKey;
use intern::Lookup;

use crate::definitions::Directive;
use crate::definitions::*;

pub trait Schema {
    fn query_type(&self) -> Option<Type>;

    fn mutation_type(&self) -> Option<Type>;

    fn subscription_type(&self) -> Option<Type>;

    fn clientid_field(&self) -> FieldID;

    fn strongid_field(&self) -> FieldID;

    fn typename_field(&self) -> FieldID;

    fn fetch_token_field(&self) -> FieldID;

    fn is_fulfilled_field(&self) -> FieldID;

    fn get_type(&self, type_name: StringKey) -> Option<Type>;

    fn get_directive(&self, name: DirectiveName) -> Option<&Directive>;

    fn input_object(&self, id: InputObjectID) -> &InputObject;
    fn input_objects<'a>(&'a self) -> Box<dyn Iterator<Item = &'a InputObject> + 'a>;

    fn enum_(&self, id: EnumID) -> &Enum;
    fn enums<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Enum> + 'a>;

    fn scalar(&self, id: ScalarID) -> &Scalar;
    fn scalars<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Scalar> + 'a>;

    fn field(&self, id: FieldID) -> &Field;
    fn fields<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Field> + 'a>;

    fn object(&self, id: ObjectID) -> &Object;
    fn objects<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Object> + 'a>;

    fn union(&self, id: UnionID) -> ⋃
    fn unions<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Union> + 'a>;

    fn interface(&self, id: InterfaceID) -> &Interface;
    fn interfaces<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Interface> + 'a>;

    fn get_type_name(&self, type_: Type) -> StringKey;

    fn is_extension_type(&self, type_: Type) -> bool;

    fn is_string(&self, type_: Type) -> bool;

    fn is_id(&self, type_: Type) -> bool;

    fn named_field(&self, parent_type: Type, name: StringKey) -> Option<FieldID>;

    fn unchecked_argument_type_sentinel(&self) -> &TypeReference<Type>;

    fn snapshot_print(&self) -> String;

    // Is `maybe_subtype` usable as — but not equal to — `super_type`.
    // Examples:
    // is_type_strict_subtype_of(Int!, Int) => true, Int! is usable where an Int is expected
    // is_type_strict_subtype_of(Int, Int) => false, same type
    // is_type_strict_subtype_of(InterfaceThatExtendsNode, Node) => true
    // is_type_strict_subtype_of(Node, Node) => false, same type
    fn is_type_strict_subtype_of(
        &self,
        maybe_subtype: &TypeReference<Type>,
        super_type: &TypeReference<Type>,
    ) -> bool {
        match (maybe_subtype, super_type) {
            (TypeReference::NonNull(of_sub), TypeReference::NonNull(of_super)) => {
                self.is_type_strict_subtype_of(of_sub, of_super)
            }
            // If supertype is non-null, maybeSubType must be non-nullable too
            (_, TypeReference::NonNull(_)) => false,
            // If supertype is nullable, maybeSubType may be non-nullable or nullable
            (TypeReference::NonNull(of_sub), _) => {
                // `T!` is a strict subtype of `U` so long as T is a (non-strict) subtype of U
                self.is_type_subtype_of(of_sub, super_type)
            }
            (TypeReference::List(of_sub), TypeReference::List(of_super)) => {
                self.is_type_strict_subtype_of(of_sub, of_super)
            }
            // If supertype is a list, maybeSubType must be a list too
            (_, TypeReference::List(_)) => false,
            // If supertype is not a list, maybeSubType must also not be a list
            (TypeReference::List(_), _) => false,
            (TypeReference::Named(named_subtype), TypeReference::Named(named_supertype)) => {
                self.is_named_type_strict_subtype_of(*named_subtype, *named_supertype)
            }
        }
    }

    fn is_type_subtype_of(
        &self,
        maybe_subtype: &TypeReference<Type>,
        super_type: &TypeReference<Type>,
    ) -> bool {
        match (maybe_subtype, super_type) {
            (TypeReference::NonNull(of_sub), TypeReference::NonNull(of_super)) => {
                self.is_type_subtype_of(of_sub, of_super)
            }
            // If supertype is non-null, maybeSubType must be non-nullable too
            (_, TypeReference::NonNull(_)) => false,
            // If supertype is nullable, maybeSubType may be non-nullable or nullable
            (TypeReference::NonNull(of_sub), _) => self.is_type_subtype_of(of_sub, super_type),
            (TypeReference::List(of_sub), TypeReference::List(of_super)) => {
                self.is_type_subtype_of(of_sub, of_super)
            }
            // If supertype is a list, maybeSubType must be a list too
            (_, TypeReference::List(_)) => false,
            // If supertype is not a list, maybeSubType must also not be a list
            (TypeReference::List(_), _) => false,
            (TypeReference::Named(named_subtype), TypeReference::Named(named_supertype)) => {
                self.is_named_type_subtype_of(*named_subtype, *named_supertype)
            }
        }
    }

    // Is `maybe_subtype` usable as — but not equal to — `super_type`.
    // Examples:
    // is_type_strict_subtype_of(InterfaceThatExtendsNode, Node) => true
    // is_type_strict_subtype_of(Node, Node) => false, same type
    fn is_named_type_strict_subtype_of(&self, maybe_subtype: Type, super_type: Type) -> bool {
        match (maybe_subtype, super_type) {
            (Type::Object(sub_id), Type::Interface(super_id)) => {
                // does object implement the interface
                let object = self.object(sub_id);
                object.interfaces.contains(&super_id)
            }
            (Type::Object(sub_id), Type::Union(super_id)) => {
                // is object a member of the union
                let union = self.union(super_id);
                union.members.contains(&sub_id)
            }
            (Type::Interface(sub_id), Type::Interface(super_id)) => {
                // does sub interface implement the super interface
                let interface = self.interface(sub_id);
                interface.interfaces.contains(&super_id)
            }
            _ => false,
        }
    }

    fn is_named_type_subtype_of(&self, maybe_subtype: Type, super_type: Type) -> bool {
        match (maybe_subtype, super_type) {
            (Type::Object(sub_id), Type::Interface(super_id)) => {
                // does object implement the interface
                let object = self.object(sub_id);
                object.interfaces.contains(&super_id)
            }
            (Type::Object(sub_id), Type::Union(super_id)) => {
                // is object a member of the union
                let union = self.union(super_id);
                union.members.contains(&sub_id)
            }
            (Type::Interface(sub_id), Type::Interface(super_id)) => {
                // does interface implement the interface
                let interface = self.interface(sub_id);
                sub_id == super_id || interface.interfaces.contains(&super_id)
            }
            _ => maybe_subtype == super_type,
        }
    }

    fn are_overlapping_types(&self, a: Type, b: Type) -> bool {
        fn overlapping_objects(a: &[ObjectID], b: &[ObjectID]) -> bool {
            a.iter().any(|item| b.contains(item))
        }

        fn overlapping_interfaces(a: &[InterfaceID], b: &[InterfaceID]) -> bool {
            a.iter().any(|item| b.contains(item))
        }

        if a == b {
            return true;
        };
        match (a, b) {
            (Type::Interface(a), Type::Interface(b)) => {
                overlapping_objects(
                    &self.interface(a).implementing_objects,
                    &self.interface(b).implementing_objects,
                ) || overlapping_interfaces(
                    &self.interface(a).implementing_interfaces,
                    &self.interface(b).implementing_interfaces,
                )
            }

            (Type::Union(a), Type::Union(b)) => {
                overlapping_objects(&self.union(a).members, &self.union(b).members)
            }

            (Type::Union(union_id), Type::Interface(interface_id))
            | (Type::Interface(interface_id), Type::Union(union_id)) => overlapping_objects(
                &self.union(union_id).members,
                &self.interface(interface_id).implementing_objects,
            ),

            (Type::Interface(interface_id), Type::Object(object_id))
            | (Type::Object(object_id), Type::Interface(interface_id)) => self
                .interface(interface_id)
                .implementing_objects
                .contains(&object_id),

            (Type::Union(union_id), Type::Object(object_id))
            | (Type::Object(object_id), Type::Union(union_id)) => {
                self.union(union_id).members.contains(&object_id)
            }

            _ => false,
        }
    }

    fn write_type_string(&self, writer: &mut String, type_: &TypeReference<Type>) -> FmtResult {
        match type_ {
            TypeReference::Named(inner) => {
                write!(writer, "{}", self.get_type_name(*inner).lookup())
            }
            TypeReference::NonNull(of) => {
                self.write_type_string(writer, of)?;
                write!(writer, "!")
            }
            TypeReference::List(of) => {
                write!(writer, "[")?;
                self.write_type_string(writer, of)?;
                write!(writer, "]")
            }
        }
    }

    fn get_type_string(&self, type_: &TypeReference<Type>) -> String {
        let mut result = String::new();
        self.write_type_string(&mut result, type_).unwrap();
        result
    }

    fn is_extension_directive(&self, name: DirectiveName) -> bool {
        if let Some(directive) = self.get_directive(name) {
            directive.is_extension
        } else {
            panic!("Unknown directive {}.", name.0.lookup())
        }
    }

    fn directives_for_type(&self, type_: Type) -> &Vec<DirectiveValue> {
        match type_ {
            Type::Enum(id) => &self.enum_(id).directives,
            Type::InputObject(id) => &self.input_object(id).directives,
            Type::Interface(id) => &self.interface(id).directives,
            Type::Object(id) => &self.object(id).directives,
            Type::Scalar(id) => &self.scalar(id).directives,
            Type::Union(id) => &self.union(id).directives,
        }
    }
}
<|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/.

#[macro_use]
mod macros;

mod activate;
mod consts;
mod handle;
mod initialize;
mod metadata_handle;
mod shared;

pub use self::{
    activate::CryptActivationHandle,
    consts::CLEVIS_TANG_TRUST_URL,
    handle::CryptHandle,
    initialize::CryptInitializer,
    metadata_handle::CryptMetadataHandle,
    shared::{
        back_up_luks_header, crypt_metadata_size, interpret_clevis_config, restore_luks_header,
        set_up_crypt_logging,
    },
};

#[cfg(test)]
mod tests {
    use std::{
        env,
        error::Error,
        ffi::CString,
        fs::{File, OpenOptions},
        io::{self, Read, Write},
        mem::MaybeUninit,
        path::Path,
        ptr, slice,
    };

    use devicemapper::{Bytes, Sectors, IEC};
    use libcryptsetup_rs::{
        consts::vals::{CryptStatusInfo, EncryptionFormat},
        CryptInit, Either,
    };

    use crate::{
        engine::{
            strat_engine::{
                backstore::crypt::{
                    consts::{
                        CLEVIS_LUKS_TOKEN_ID, DEFAULT_CRYPT_KEYSLOTS_SIZE,
                        DEFAULT_CRYPT_METADATA_SIZE, LUKS2_TOKEN_ID, STRATIS_MEK_SIZE,
                    },
                    shared::acquire_crypt_device,
                },
                ns::{unshare_namespace, MemoryFilesystem},
                tests::{crypt, loopbacked, real},
            },
            types::{DevUuid, DevicePath, KeyDescription, PoolUuid, UnlockMethod},
        },
        stratis::StratisError,
    };

    use super::*;

    /// If this method is called without a key with the specified key description
    /// in the kernel ring, it should always fail and allow us to test the rollback
    /// of failed initializations.
    fn test_failed_init(paths: &[&Path]) {
        assert_eq!(paths.len(), 1);

        let path = paths.get(0).expect("There must be exactly one path");
        let key_description =
            KeyDescription::try_from("I am not a key".to_string()).expect("no semi-colons");

        let pool_uuid = PoolUuid::new_v4();
        let dev_uuid = DevUuid::new_v4();

        let result = CryptInitializer::new(DevicePath::new(path).unwrap(), pool_uuid, dev_uuid)
            .initialize(Some(&key_description), None);

        // Initialization cannot occur with a non-existent key
        assert!(result.is_err());

        assert!(CryptHandle::setup(path).unwrap().is_none());

        // TODO: Check actual superblock with libblkid
    }

    #[test]
    fn loop_test_failed_init() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_failed_init,
        );
    }

    #[test]
    fn real_test_failed_init() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_failed_init,
        );
    }

    /// Test the method `can_unlock` works on an initialized device in both
    /// active and inactive states.
    fn test_can_unlock(paths: &[&Path]) {
        fn crypt_test(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> std::result::Result<(), Box<dyn Error>> {
            let mut handles = vec![];

            let pool_uuid = PoolUuid::new_v4();
            for path in paths {
                let dev_uuid = DevUuid::new_v4();

                let handle = CryptInitializer::new(DevicePath::new(path)?, pool_uuid, dev_uuid)
                    .initialize(Some(key_desc), None)?;
                handles.push(handle);
            }

            for path in paths {
                if !CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should be able to be unlocked".to_string(),
                    )));
                }
            }

            for handle in handles.iter_mut() {
                handle.deactivate()?;
            }

            for path in paths {
                if !CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should be able to be unlocked".to_string(),
                    )));
                }
            }

            for handle in handles.iter_mut() {
                handle.wipe()?;
            }

            for path in paths {
                if CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should no longer be able to be unlocked".to_string(),
                    )));
                }
            }

            Ok(())
        }

        crypt::insert_and_cleanup_key(paths, crypt_test)
    }

    #[test]
    fn loop_test_can_unlock() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_can_unlock,
        );
    }

    #[test]
    fn real_test_can_unlock() {
        real::test_with_spec(
            &real::DeviceLimits::Range(1, 3, None, None),
            test_can_unlock,
        );
    }

    /// Test initializing and activating an encrypted device using
    /// the utilities provided here.
    ///
    /// The overall format of the test involves generating a random byte buffer
    /// of size 1 MiB, encrypting it on disk, and then ensuring that the plaintext
    /// cannot be found on the encrypted disk by doing a scan of the disk using
    /// a sliding window.
    ///
    /// The sliding window size of 1 MiB was chosen to lower the number of
    /// searches that need to be done compared to a smaller sliding window
    /// and also to decrease the probability of the random sequence being found
    /// on the disk due to leftover data from other tests.
    // TODO: Rewrite libc calls using nix crate.
    fn test_crypt_device_ops(paths: &[&Path]) {
        fn crypt_test(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> std::result::Result<(), Box<dyn Error>> {
            let path = paths.get(0).ok_or_else(|| {
                Box::new(StratisError::Msg(
                    "This test only accepts a single device".to_string(),
                ))
            })?;

            let pool_uuid = PoolUuid::new_v4();
            let dev_uuid = DevUuid::new_v4();

            let handle = CryptInitializer::new(DevicePath::new(path)?, pool_uuid, dev_uuid)
                .initialize(Some(key_desc), None)?;
            let logical_path = handle.activated_device_path();

            const WINDOW_SIZE: usize = 1024 * 1024;
            let mut devicenode = OpenOptions::new().write(true).open(logical_path)?;
            let mut random_buffer = Box::new(vec![0; WINDOW_SIZE].into_boxed_slice());
            File::open("/dev/urandom")?.read_exact(&mut random_buffer)?;
            devicenode.write_all(&random_buffer)?;
            std::mem::drop(devicenode);

            let dev_path_cstring = CString::new(path.to_str().ok_or_else(|| {
                Box::new(io::Error::new(
                    io::ErrorKind::Other,
                    "Failed to convert path to string",
                ))
            })?)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
            let fd = unsafe { libc::open(dev_path_cstring.as_ptr(), libc::O_RDONLY) };
            if fd < 0 {
                return Err(Box::new(io::Error::last_os_error()));
            }

            let mut stat: MaybeUninit<libc::stat> = MaybeUninit::zeroed();
            let fstat_result = unsafe { libc::fstat(fd, stat.as_mut_ptr()) };
            if fstat_result < 0 {
                return Err(Box::new(io::Error::last_os_error()));
            }
            let device_size =
                convert_int!(unsafe { stat.assume_init() }.st_size, libc::off_t, usize)?;
            let mapped_ptr = unsafe {
                libc::mmap(
                    ptr::null_mut(),
                    device_size,
                    libc::PROT_READ,
                    libc::MAP_SHARED,
                    fd,
                    0,
                )
            };
            if mapped_ptr.is_null() {
                return Err(Box::new(io::Error::new(
                    io::ErrorKind::Other,
                    "mmap failed",
                )));
            }

            {
                let disk_buffer =
                    unsafe { slice::from_raw_parts(mapped_ptr as *const u8, device_size) };
                for window in disk_buffer.windows(WINDOW_SIZE) {
                    if window == &*random_buffer as &[u8] {
                        unsafe {
                            libc::munmap(mapped_ptr, device_size);
                            libc::close(fd);
                        };
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            "Disk was not encrypted!",
                        )));
                    }
                }
            }

            unsafe {
                libc::munmap(mapped_ptr, device_size);
                libc::close(fd);
            };

            let device_name = handle.name().to_owned();
            loop {
                match libcryptsetup_rs::status(
                    Some(&mut handle.acquire_crypt_device().unwrap()),
                    &device_name,
                ) {
                    Ok(CryptStatusInfo::Busy) => (),
                    Ok(CryptStatusInfo::Active) => break,
                    Ok(s) => {
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            format!("Crypt device is in invalid state {:?}", s),
                        )))
                    }
                    Err(e) => {
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            format!("Checking device status returned error: {}", e),
                        )))
                    }
                }
            }

            handle.deactivate()?;

            let handle =
                CryptActivationHandle::setup(path, UnlockMethod::Keyring)?.ok_or_else(|| {
                    Box::new(io::Error::new(
                        io::ErrorKind::Other,
                        format!(
                            "Device {} no longer appears to be a LUKS2 device",
                            path.display(),
                        ),
                    ))
                })?;
            handle.wipe()?;

            Ok(())
        }

        assert_eq!(paths.len(), 1);

        crypt::insert_and_cleanup_key(paths, crypt_test);
    }

    #[test]
    fn real_test_crypt_device_ops() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(2 * IEC::Mi))),
            test_crypt_device_ops,
        );
    }

    #[test]
    fn loop_test_crypt_metadata_defaults() {
        fn test_defaults(paths: &[&Path]) {
            let mut context = CryptInit::init(paths[0]).unwrap();
            context
                .context_handle()
                .format::<()>(
                    EncryptionFormat::Luks2,
                    ("aes", "xts-plain64"),
                    None,
                    Either::Right(STRATIS_MEK_SIZE),
                    None,
                )
                .unwrap();
            let (metadata, keyslot) = context.settings_handle().get_metadata_size().unwrap();
            assert_eq!(DEFAULT_CRYPT_METADATA_SIZE, Bytes::from(*metadata));
            assert_eq!(DEFAULT_CRYPT_KEYSLOTS_SIZE, Bytes::from(*keyslot));
        }

        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Exactly(1, None), test_defaults);
    }

    fn test_both_initialize(paths: &[&Path]) {
        fn both_initialize(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> Result<(), Box<dyn Error>> {
            unshare_namespace()?;
            let _memfs = MemoryFilesystem::new()?;
            let path = paths
                .get(0)
                .copied()
                .ok_or_else(|| StratisError::Msg("Expected exactly one path".to_string()))?;
            let handle = CryptInitializer::new(
                DevicePath::new(path)?,
                PoolUuid::new_v4(),
                DevUuid::new_v4(),
            )
            .initialize(
                Some(key_desc),
                Some(&(
                    "tang".to_string(),
                    json!({"url": env::var("TANG_URL")?, "stratis:tang:trust_url": true}),
                )),
            )?;

            let mut device = acquire_crypt_device(handle.luks2_device_path())?;
            device.token_handle().json_get(LUKS2_TOKEN_ID)?;
            device.token_handle().json_get(CLEVIS_LUKS_TOKEN_ID)?;
            Ok(())
        }

        crypt::insert_and_cleanup_key(paths, both_initialize);
    }

    #[test]
    fn clevis_real_test_both_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_both_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_real_should_fail_test_both_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_both_initialize,
        );
    }

    #[test]
    fn clevis_loop_test_both_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_both_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_loop_should_fail_test_both_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_both_initialize,
        );
    }

    fn test_clevis_initialize(paths: &[&Path]) {
        unshare_namespace().unwrap();
        let _memfs = MemoryFilesystem::new().unwrap();
        let path = paths[0];
        let handle = CryptInitializer::new(
            DevicePath::new(path).unwrap(),
            PoolUuid::new_v4(),
            DevUuid::new_v4(),
        )
        .initialize(
            None,
            Some(&(
                "tang".to_string(),
                json!({"url": env::var("TANG_URL").unwrap(), "stratis:tang:trust_url": true}),
            )),
        )
        .unwrap();

        let mut device = acquire_crypt_device(handle.luks2_device_path()).unwrap();
        assert!(device.token_handle().json_get(CLEVIS_LUKS_TOKEN_ID).is_ok());
        assert!(device.token_handle().json_get(LUKS2_TOKEN_ID).is_err());
    }

    #[test]
    fn clevis_real_test_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_clevis_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_real_should_fail_test_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_clevis_initialize,
        );
    }

    #[test]
    fn clevis_loop_test_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_clevis_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_loop_should_fail_test_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_clevis_initialize,
        );
    }
}
<commit_msg>Fix double-box error<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/.

#[macro_use]
mod macros;

mod activate;
mod consts;
mod handle;
mod initialize;
mod metadata_handle;
mod shared;

pub use self::{
    activate::CryptActivationHandle,
    consts::CLEVIS_TANG_TRUST_URL,
    handle::CryptHandle,
    initialize::CryptInitializer,
    metadata_handle::CryptMetadataHandle,
    shared::{
        back_up_luks_header, crypt_metadata_size, interpret_clevis_config, restore_luks_header,
        set_up_crypt_logging,
    },
};

#[cfg(test)]
mod tests {
    use std::{
        env,
        error::Error,
        ffi::CString,
        fs::{File, OpenOptions},
        io::{self, Read, Write},
        mem::MaybeUninit,
        path::Path,
        ptr, slice,
    };

    use devicemapper::{Bytes, Sectors, IEC};
    use libcryptsetup_rs::{
        consts::vals::{CryptStatusInfo, EncryptionFormat},
        CryptInit, Either,
    };

    use crate::{
        engine::{
            strat_engine::{
                backstore::crypt::{
                    consts::{
                        CLEVIS_LUKS_TOKEN_ID, DEFAULT_CRYPT_KEYSLOTS_SIZE,
                        DEFAULT_CRYPT_METADATA_SIZE, LUKS2_TOKEN_ID, STRATIS_MEK_SIZE,
                    },
                    shared::acquire_crypt_device,
                },
                ns::{unshare_namespace, MemoryFilesystem},
                tests::{crypt, loopbacked, real},
            },
            types::{DevUuid, DevicePath, KeyDescription, PoolUuid, UnlockMethod},
        },
        stratis::StratisError,
    };

    use super::*;

    /// If this method is called without a key with the specified key description
    /// in the kernel ring, it should always fail and allow us to test the rollback
    /// of failed initializations.
    fn test_failed_init(paths: &[&Path]) {
        assert_eq!(paths.len(), 1);

        let path = paths.get(0).expect("There must be exactly one path");
        let key_description =
            KeyDescription::try_from("I am not a key".to_string()).expect("no semi-colons");

        let pool_uuid = PoolUuid::new_v4();
        let dev_uuid = DevUuid::new_v4();

        let result = CryptInitializer::new(DevicePath::new(path).unwrap(), pool_uuid, dev_uuid)
            .initialize(Some(&key_description), None);

        // Initialization cannot occur with a non-existent key
        assert!(result.is_err());

        assert!(CryptHandle::setup(path).unwrap().is_none());

        // TODO: Check actual superblock with libblkid
    }

    #[test]
    fn loop_test_failed_init() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_failed_init,
        );
    }

    #[test]
    fn real_test_failed_init() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_failed_init,
        );
    }

    /// Test the method `can_unlock` works on an initialized device in both
    /// active and inactive states.
    fn test_can_unlock(paths: &[&Path]) {
        fn crypt_test(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> std::result::Result<(), Box<dyn Error>> {
            let mut handles = vec![];

            let pool_uuid = PoolUuid::new_v4();
            for path in paths {
                let dev_uuid = DevUuid::new_v4();

                let handle = CryptInitializer::new(DevicePath::new(path)?, pool_uuid, dev_uuid)
                    .initialize(Some(key_desc), None)?;
                handles.push(handle);
            }

            for path in paths {
                if !CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should be able to be unlocked".to_string(),
                    )));
                }
            }

            for handle in handles.iter_mut() {
                handle.deactivate()?;
            }

            for path in paths {
                if !CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should be able to be unlocked".to_string(),
                    )));
                }
            }

            for handle in handles.iter_mut() {
                handle.wipe()?;
            }

            for path in paths {
                if CryptActivationHandle::can_unlock(path, true, false) {
                    return Err(Box::new(StratisError::Msg(
                        "All devices should no longer be able to be unlocked".to_string(),
                    )));
                }
            }

            Ok(())
        }

        crypt::insert_and_cleanup_key(paths, crypt_test)
    }

    #[test]
    fn loop_test_can_unlock() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Range(1, 3, None),
            test_can_unlock,
        );
    }

    #[test]
    fn real_test_can_unlock() {
        real::test_with_spec(
            &real::DeviceLimits::Range(1, 3, None, None),
            test_can_unlock,
        );
    }

    /// Test initializing and activating an encrypted device using
    /// the utilities provided here.
    ///
    /// The overall format of the test involves generating a random byte buffer
    /// of size 1 MiB, encrypting it on disk, and then ensuring that the plaintext
    /// cannot be found on the encrypted disk by doing a scan of the disk using
    /// a sliding window.
    ///
    /// The sliding window size of 1 MiB was chosen to lower the number of
    /// searches that need to be done compared to a smaller sliding window
    /// and also to decrease the probability of the random sequence being found
    /// on the disk due to leftover data from other tests.
    // TODO: Rewrite libc calls using nix crate.
    fn test_crypt_device_ops(paths: &[&Path]) {
        fn crypt_test(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> std::result::Result<(), Box<dyn Error>> {
            let path = paths.get(0).ok_or_else(|| {
                Box::new(StratisError::Msg(
                    "This test only accepts a single device".to_string(),
                ))
            })?;

            let pool_uuid = PoolUuid::new_v4();
            let dev_uuid = DevUuid::new_v4();

            let handle = CryptInitializer::new(DevicePath::new(path)?, pool_uuid, dev_uuid)
                .initialize(Some(key_desc), None)?;
            let logical_path = handle.activated_device_path();

            const WINDOW_SIZE: usize = 1024 * 1024;
            let mut devicenode = OpenOptions::new().write(true).open(logical_path)?;
            let mut random_buffer = vec![0; WINDOW_SIZE].into_boxed_slice();
            File::open("/dev/urandom")?.read_exact(&mut random_buffer)?;
            devicenode.write_all(&random_buffer)?;
            std::mem::drop(devicenode);

            let dev_path_cstring = CString::new(path.to_str().ok_or_else(|| {
                Box::new(io::Error::new(
                    io::ErrorKind::Other,
                    "Failed to convert path to string",
                ))
            })?)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
            let fd = unsafe { libc::open(dev_path_cstring.as_ptr(), libc::O_RDONLY) };
            if fd < 0 {
                return Err(Box::new(io::Error::last_os_error()));
            }

            let mut stat: MaybeUninit<libc::stat> = MaybeUninit::zeroed();
            let fstat_result = unsafe { libc::fstat(fd, stat.as_mut_ptr()) };
            if fstat_result < 0 {
                return Err(Box::new(io::Error::last_os_error()));
            }
            let device_size =
                convert_int!(unsafe { stat.assume_init() }.st_size, libc::off_t, usize)?;
            let mapped_ptr = unsafe {
                libc::mmap(
                    ptr::null_mut(),
                    device_size,
                    libc::PROT_READ,
                    libc::MAP_SHARED,
                    fd,
                    0,
                )
            };
            if mapped_ptr.is_null() {
                return Err(Box::new(io::Error::new(
                    io::ErrorKind::Other,
                    "mmap failed",
                )));
            }

            {
                let disk_buffer =
                    unsafe { slice::from_raw_parts(mapped_ptr as *const u8, device_size) };
                for window in disk_buffer.windows(WINDOW_SIZE) {
                    if window == &*random_buffer as &[u8] {
                        unsafe {
                            libc::munmap(mapped_ptr, device_size);
                            libc::close(fd);
                        };
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            "Disk was not encrypted!",
                        )));
                    }
                }
            }

            unsafe {
                libc::munmap(mapped_ptr, device_size);
                libc::close(fd);
            };

            let device_name = handle.name().to_owned();
            loop {
                match libcryptsetup_rs::status(
                    Some(&mut handle.acquire_crypt_device().unwrap()),
                    &device_name,
                ) {
                    Ok(CryptStatusInfo::Busy) => (),
                    Ok(CryptStatusInfo::Active) => break,
                    Ok(s) => {
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            format!("Crypt device is in invalid state {:?}", s),
                        )))
                    }
                    Err(e) => {
                        return Err(Box::new(io::Error::new(
                            io::ErrorKind::Other,
                            format!("Checking device status returned error: {}", e),
                        )))
                    }
                }
            }

            handle.deactivate()?;

            let handle =
                CryptActivationHandle::setup(path, UnlockMethod::Keyring)?.ok_or_else(|| {
                    Box::new(io::Error::new(
                        io::ErrorKind::Other,
                        format!(
                            "Device {} no longer appears to be a LUKS2 device",
                            path.display(),
                        ),
                    ))
                })?;
            handle.wipe()?;

            Ok(())
        }

        assert_eq!(paths.len(), 1);

        crypt::insert_and_cleanup_key(paths, crypt_test);
    }

    #[test]
    fn real_test_crypt_device_ops() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(2 * IEC::Mi))),
            test_crypt_device_ops,
        );
    }

    #[test]
    fn loop_test_crypt_metadata_defaults() {
        fn test_defaults(paths: &[&Path]) {
            let mut context = CryptInit::init(paths[0]).unwrap();
            context
                .context_handle()
                .format::<()>(
                    EncryptionFormat::Luks2,
                    ("aes", "xts-plain64"),
                    None,
                    Either::Right(STRATIS_MEK_SIZE),
                    None,
                )
                .unwrap();
            let (metadata, keyslot) = context.settings_handle().get_metadata_size().unwrap();
            assert_eq!(DEFAULT_CRYPT_METADATA_SIZE, Bytes::from(*metadata));
            assert_eq!(DEFAULT_CRYPT_KEYSLOTS_SIZE, Bytes::from(*keyslot));
        }

        loopbacked::test_with_spec(&loopbacked::DeviceLimits::Exactly(1, None), test_defaults);
    }

    fn test_both_initialize(paths: &[&Path]) {
        fn both_initialize(
            paths: &[&Path],
            key_desc: &KeyDescription,
        ) -> Result<(), Box<dyn Error>> {
            unshare_namespace()?;
            let _memfs = MemoryFilesystem::new()?;
            let path = paths
                .get(0)
                .copied()
                .ok_or_else(|| StratisError::Msg("Expected exactly one path".to_string()))?;
            let handle = CryptInitializer::new(
                DevicePath::new(path)?,
                PoolUuid::new_v4(),
                DevUuid::new_v4(),
            )
            .initialize(
                Some(key_desc),
                Some(&(
                    "tang".to_string(),
                    json!({"url": env::var("TANG_URL")?, "stratis:tang:trust_url": true}),
                )),
            )?;

            let mut device = acquire_crypt_device(handle.luks2_device_path())?;
            device.token_handle().json_get(LUKS2_TOKEN_ID)?;
            device.token_handle().json_get(CLEVIS_LUKS_TOKEN_ID)?;
            Ok(())
        }

        crypt::insert_and_cleanup_key(paths, both_initialize);
    }

    #[test]
    fn clevis_real_test_both_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_both_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_real_should_fail_test_both_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_both_initialize,
        );
    }

    #[test]
    fn clevis_loop_test_both_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_both_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_loop_should_fail_test_both_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_both_initialize,
        );
    }

    fn test_clevis_initialize(paths: &[&Path]) {
        unshare_namespace().unwrap();
        let _memfs = MemoryFilesystem::new().unwrap();
        let path = paths[0];
        let handle = CryptInitializer::new(
            DevicePath::new(path).unwrap(),
            PoolUuid::new_v4(),
            DevUuid::new_v4(),
        )
        .initialize(
            None,
            Some(&(
                "tang".to_string(),
                json!({"url": env::var("TANG_URL").unwrap(), "stratis:tang:trust_url": true}),
            )),
        )
        .unwrap();

        let mut device = acquire_crypt_device(handle.luks2_device_path()).unwrap();
        assert!(device.token_handle().json_get(CLEVIS_LUKS_TOKEN_ID).is_ok());
        assert!(device.token_handle().json_get(LUKS2_TOKEN_ID).is_err());
    }

    #[test]
    fn clevis_real_test_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_clevis_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_real_should_fail_test_initialize() {
        real::test_with_spec(
            &real::DeviceLimits::Exactly(1, None, Some(Sectors(1024 * 1024 * 1024 / 512))),
            test_clevis_initialize,
        );
    }

    #[test]
    fn clevis_loop_test_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_clevis_initialize,
        );
    }

    #[test]
    #[should_panic]
    fn clevis_loop_should_fail_test_initialize() {
        loopbacked::test_with_spec(
            &loopbacked::DeviceLimits::Exactly(1, None),
            test_clevis_initialize,
        );
    }
}
<|endoftext|>


<commit_before>use colored::*;
use regex::Regex;
use std::env;
use std::path::PathBuf;
use ui_test::{Config, Mode, OutputConflictHandling};

fn miri_path() -> PathBuf {
    PathBuf::from(option_env!("MIRI").unwrap_or(env!("CARGO_BIN_EXE_miri")))
}

fn run_tests(mode: Mode, path: &str, target: Option<String>) {
    let in_rustc_test_suite = option_env!("RUSTC_STAGE").is_some();

    // Add some flags we always want.
    let mut flags = Vec::new();
    flags.push("--edition".to_owned());
    flags.push("2018".to_owned());
    if in_rustc_test_suite {
        // Less aggressive warnings to make the rustc toolstate management less painful.
        // (We often get warnings when e.g. a feature gets stabilized or some lint gets added/improved.)
        flags.push("-Astable-features".to_owned());
    } else {
        flags.push("-Dwarnings".to_owned());
        flags.push("-Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    }
    if let Ok(sysroot) = env::var("MIRI_SYSROOT") {
        flags.push("--sysroot".to_string());
        flags.push(sysroot);
    }
    if let Ok(extra_flags) = env::var("MIRIFLAGS") {
        for flag in extra_flags.split_whitespace() {
            flags.push(flag.to_string());
        }
    }
    flags.push("-Zui-testing".to_string());
    if let Some(target) = &target {
        flags.push("--target".to_string());
        flags.push(target.clone());
    }

    let skip_ui_checks = in_rustc_test_suite || env::var_os("MIRI_SKIP_UI_CHECKS").is_some();

    let output_conflict_handling = match (env::var_os("MIRI_BLESS").is_some(), skip_ui_checks) {
        (false, false) => OutputConflictHandling::Error,
        (true, false) => OutputConflictHandling::Bless,
        (false, true) => OutputConflictHandling::Ignore,
        (true, true) => panic!("cannot use MIRI_BLESS and MIRI_SKIP_UI_CHECKS at the same time"),
    };

    // Pass on all arguments as filters.
    let path_filter = std::env::args().skip(1);

    let config = Config {
        args: flags,
        target,
        stderr_filters: STDERR.clone(),
        stdout_filters: STDOUT.clone(),
        root_dir: PathBuf::from(path),
        mode,
        path_filter: path_filter.collect(),
        program: miri_path(),
        output_conflict_handling,
    };
    ui_test::run_tests(config)
}

macro_rules! regexes {
    ($name:ident: $($regex:expr => $replacement:expr,)*) => {lazy_static::lazy_static! {
        static ref $name: Vec<(Regex, &'static str)> = vec![
            $((Regex::new($regex).unwrap(), $replacement),)*
        ];
    }};
}

regexes! {
    STDOUT:
    // Windows file paths
    r"\\"                           => "/",
}

regexes! {
    STDERR:
    // erase line and column info
    r"\.rs:[0-9]+:[0-9]+(: [0-9]+:[0-9]+)?" => ".rs:LL:CC",
    // erase alloc ids
    "alloc[0-9]+"                    => "ALLOC",
    // erase Stacked Borrows tags
    "<[0-9]+>"                       => "<TAG>",
    // erase whitespace that differs between platforms
    r" +at (.*\.rs)"                 => " at $1",
    // erase generics in backtraces
    "([0-9]+: .*)::<.*>"             => "$1",
    // erase addresses in backtraces
    "([0-9]+: ) +0x[0-9a-f]+ - (.*)" => "$1$2",
    // erase long hexadecimals
    r"0x[0-9a-fA-F]+[0-9a-fA-F]{2,2}" => "$$HEX",
    // erase clocks
    r"VClock\(\[[^\]]+\]\)"      => "VClock",
    // erase specific alignments
    "alignment [0-9]+"               => "alignment ALIGN",
    // erase thread caller ids
    r"\(call [0-9]+\)"              => "(call ID)",
    // erase platform module paths
    "sys::[a-z]+::"                  => "sys::PLATFORM::",
    // Windows file paths
    r"\\"                           => "/",
    // erase Rust stdlib path
    "[^ `]*/(rust[^/]*|checkout)/library/" => "RUSTLIB/",
    // erase platform file paths
    "sys/[a-z]+/"                    => "sys/PLATFORM/",
}

fn ui(mode: Mode, path: &str) {
    let target = get_target();

    let msg = format!(
        "## Running ui tests in {path} against miri for {}",
        target.as_deref().unwrap_or("host")
    );
    eprintln!("{}", msg.green().bold());

    run_tests(mode, path, target);
}

fn get_target() -> Option<String> {
    env::var("MIRI_TEST_TARGET").ok()
}

fn main() {
    // Add a test env var to do environment communication tests.
    env::set_var("MIRI_ENV_VAR_TEST", "0");
    // Let the tests know where to store temp files (they might run for a different target, which can make this hard to find).
    env::set_var("MIRI_TEMP", env::temp_dir());

    ui(Mode::Pass, "tests/pass");
    ui(Mode::Panic, "tests/panic");
    ui(Mode::Fail, "tests/fail");
}
<commit_msg>Auto merge of #2179 - RalfJung:rustc, r=oli-obk<commit_after>use colored::*;
use regex::Regex;
use std::env;
use std::path::PathBuf;
use ui_test::{Config, Mode, OutputConflictHandling};

fn miri_path() -> PathBuf {
    PathBuf::from(option_env!("MIRI").unwrap_or(env!("CARGO_BIN_EXE_miri")))
}

fn run_tests(mode: Mode, path: &str, target: Option<String>) {
    let in_rustc_test_suite = option_env!("RUSTC_STAGE").is_some();

    // Add some flags we always want.
    let mut flags = Vec::new();
    flags.push("--edition".to_owned());
    flags.push("2018".to_owned());
    if in_rustc_test_suite {
        // Less aggressive warnings to make the rustc toolstate management less painful.
        // (We often get warnings when e.g. a feature gets stabilized or some lint gets added/improved.)
        flags.push("-Astable-features".to_owned());
    } else {
        flags.push("-Dwarnings".to_owned());
        flags.push("-Dunused".to_owned()); // overwrite the -Aunused in compiletest-rs
    }
    if let Ok(sysroot) = env::var("MIRI_SYSROOT") {
        flags.push("--sysroot".to_string());
        flags.push(sysroot);
    }
    if let Ok(extra_flags) = env::var("MIRIFLAGS") {
        for flag in extra_flags.split_whitespace() {
            flags.push(flag.to_string());
        }
    }
    flags.push("-Zui-testing".to_string());
    if let Some(target) = &target {
        flags.push("--target".to_string());
        flags.push(target.clone());
    }

    let skip_ui_checks = env::var_os("MIRI_SKIP_UI_CHECKS").is_some();

    let output_conflict_handling = match (env::var_os("MIRI_BLESS").is_some(), skip_ui_checks) {
        (false, false) => OutputConflictHandling::Error,
        (true, false) => OutputConflictHandling::Bless,
        (false, true) => OutputConflictHandling::Ignore,
        (true, true) => panic!("cannot use MIRI_BLESS and MIRI_SKIP_UI_CHECKS at the same time"),
    };

    // Pass on all arguments as filters.
    let path_filter = std::env::args().skip(1);

    let config = Config {
        args: flags,
        target,
        stderr_filters: STDERR.clone(),
        stdout_filters: STDOUT.clone(),
        root_dir: PathBuf::from(path),
        mode,
        path_filter: path_filter.collect(),
        program: miri_path(),
        output_conflict_handling,
    };
    ui_test::run_tests(config)
}

macro_rules! regexes {
    ($name:ident: $($regex:expr => $replacement:expr,)*) => {lazy_static::lazy_static! {
        static ref $name: Vec<(Regex, &'static str)> = vec![
            $((Regex::new($regex).unwrap(), $replacement),)*
        ];
    }};
}

regexes! {
    STDOUT:
    // Windows file paths
    r"\\"                           => "/",
}

regexes! {
    STDERR:
    // erase line and column info
    r"\.rs:[0-9]+:[0-9]+(: [0-9]+:[0-9]+)?" => ".rs:LL:CC",
    // erase alloc ids
    "alloc[0-9]+"                    => "ALLOC",
    // erase Stacked Borrows tags
    "<[0-9]+>"                       => "<TAG>",
    // erase whitespace that differs between platforms
    r" +at (.*\.rs)"                 => " at $1",
    // erase generics in backtraces
    "([0-9]+: .*)::<.*>"             => "$1",
    // erase addresses in backtraces
    "([0-9]+: ) +0x[0-9a-f]+ - (.*)" => "$1$2",
    // erase long hexadecimals
    r"0x[0-9a-fA-F]+[0-9a-fA-F]{2,2}" => "$$HEX",
    // erase clocks
    r"VClock\(\[[^\]]+\]\)"      => "VClock",
    // erase specific alignments
    "alignment [0-9]+"               => "alignment ALIGN",
    // erase thread caller ids
    r"\(call [0-9]+\)"              => "(call ID)",
    // erase platform module paths
    "sys::[a-z]+::"                  => "sys::PLATFORM::",
    // Windows file paths
    r"\\"                           => "/",
    // erase Rust stdlib path
    "[^ `]*/(rust[^/]*|checkout)/library/" => "RUSTLIB/",
    // erase platform file paths
    "sys/[a-z]+/"                    => "sys/PLATFORM/",
}

fn ui(mode: Mode, path: &str) {
    let target = get_target();

    let msg = format!(
        "## Running ui tests in {path} against miri for {}",
        target.as_deref().unwrap_or("host")
    );
    eprintln!("{}", msg.green().bold());

    run_tests(mode, path, target);
}

fn get_target() -> Option<String> {
    env::var("MIRI_TEST_TARGET").ok()
}

fn main() {
    // Add a test env var to do environment communication tests.
    env::set_var("MIRI_ENV_VAR_TEST", "0");
    // Let the tests know where to store temp files (they might run for a different target, which can make this hard to find).
    env::set_var("MIRI_TEMP", env::temp_dir());

    ui(Mode::Pass, "tests/pass");
    ui(Mode::Panic, "tests/panic");
    ui(Mode::Fail, "tests/fail");
}
<|endoftext|>