<commit_before>use bounding_volume::AABB;
use math::Isometry;
use na::{Point2, Real, Vector3};
use partitioning::{BestFirstBVVisitStatus, BestFirstDataVisitStatus, BestFirstVisitor};
use query::{closest_points_internal, Ray, RayCast, RayIntersection};
use shape::{CompositeShape, HeightField, FeatureId};

#[cfg(feature = "dim2")]
impl<N: Real> RayCast<N> for HeightField<N> {
    #[inline]
    fn toi_and_normal_with_ray(
        &self,
        m: &Isometry<N>,
        ray: &Ray<N>,
        _: bool,
    ) -> Option<RayIntersection<N>>
    {
        let aabb = self.aabb();
        let ls_ray = ray.inverse_transform_by(m);
        let (min_t, max_t) = aabb.clip_ray_parameters(&ls_ray)?;
        let (clip_ray_a, clip_ray_b) = (ls_ray.point_at(min_t), ls_ray.point_at(max_t));
        let mut curr = match self.cell_at_point(&clip_ray_a) {
            Some(cell) => cell,
            // None may happen due to slight numerical errors.
            None => if ls_ray.origin.x > N::zero() { self.num_cells() - 1 } else {
                0
            }
        };

        /*
         * Test the segment under the ray.
         */
        if let Some(seg) = self.segment_at(curr) {
            let (s, t) = closest_points_internal::line_against_line_parameters(&ray.origin, &ray.dir, seg.a(), &seg.scaled_direction());
            if s >= N::zero() && t >= N::zero() && t <= N::one() {
                // Cast succeeded on the first element!
                // FIXME: return the correct feature and normal.
                return Some(RayIntersection::new(s, m * seg.normal().unwrap().unwrap(), FeatureId::Face(curr)))
            }
        }

        /*
         * Test other segments in the path of the ray.
         */
        if ls_ray.dir.x == N::zero() {
            return None;
        }

        let right = ls_ray.dir.x > N::zero();
        let cell_width = self.cell_width();
        let start_x = self.start_x();

        while (right && curr < self.num_cells()) || (!right && curr > 0) {
            let curr_param;

            if right {
                curr += 1;
                curr_param = (cell_width * na::convert(curr as f64) + start_x - ls_ray.origin.x) / ls_ray.dir.x;
            } else {
                curr_param = (ls_ray.origin.x - cell_width * na::convert(curr as f64) - start_x) / ls_ray.dir.x;
                curr -= 1;
            }

            if curr_param >= max_t {
                // The part of the ray after max_t is outside of the heightfield AABB.
                return None;
            }

            if let Some(seg) = self.segment_at(curr) {
                // TODO: test the y-coordinates (equivalent to an AABB test) before actually computing the intersection.
                let (s, t) = closest_points_internal::line_against_line_parameters(&ray.origin, &ray.dir, seg.a(), &seg.scaled_direction());
                if t >= N::zero() && t <= N::one() {
                    // FIXME: return the correct feature and normal.
                    return Some(RayIntersection::new(s, m * seg.normal().unwrap().unwrap(), FeatureId::Face(curr)))
                }
            }
        }

        None
    }
}

#[cfg(feature = "dim3")]
impl<N: Real> RayCast<N> for HeightField<N> {
    #[inline]
    fn toi_and_normal_with_ray(
        &self,
        m: &Isometry<N>,
        ray: &Ray<N>,
        solid: bool,
    ) -> Option<RayIntersection<N>>
    {
        let aabb = self.aabb();
        let ls_ray = ray.inverse_transform_by(m);
        let (min_t, max_t) = aabb.clip_ray_parameters(&ls_ray)?;
        let (clip_ray_a, clip_ray_b) = (ls_ray.point_at(min_t), ls_ray.point_at(max_t));
        let mut cell = match self.cell_at_point(&clip_ray_a) {
            Some(cell) => cell,
            // None may happen due to slight numerical errors.
            None => {
                let i = if ls_ray.origin.z > N::zero() {
                    self.nrows() - 1
                } else {
                    0
                };

                let j = if ls_ray.origin.x > N::zero() {
                    self.ncols() - 1
                } else {
                    0
                };

                (i, j)
            }
        };

        loop {
            let tris = self.triangles_at(cell.0, cell.1);
            let inter1 = tris.0.and_then(|tri| tri.toi_and_normal_with_ray(m, ray, solid));
            let inter2 = tris.1.and_then(|tri| tri.toi_and_normal_with_ray(m, ray, solid));

            match (inter1, inter2) {
                (Some(mut inter1), Some(mut inter2)) => {
                    if inter1.toi < inter2.toi {
                        inter1.feature = self.convert_triangle_feature_id(cell.0, cell.1, true, inter1.feature);
                        return Some(inter1);
                    } else {
                        inter2.feature = self.convert_triangle_feature_id(cell.0, cell.1, false, inter2.feature);
                        return Some(inter2);
                    }
                }
                (Some(mut inter), None) => {
                    inter.feature = self.convert_triangle_feature_id(cell.0, cell.1, true, inter.feature);
                    return Some(inter);
                }
                (None, Some(mut inter)) => {
                    inter.feature = self.convert_triangle_feature_id(cell.0, cell.1, false, inter.feature);
                    return Some(inter);
                }
                (None, None) => {}
            }

            /*
             * Find the next cell to cast the ray on.
             */
            let (toi_x, right) = if ls_ray.dir.x > N::zero() {
                let x = self.x_at(cell.1 + 1);
                ((x - ls_ray.origin.x) / ls_ray.dir.x, true)
            } else if ls_ray.dir.x < N::zero() {
                let x = self.x_at(cell.1 + 0);
                ((x - ls_ray.origin.x) / ls_ray.dir.x, false)
            } else {
                (N::max_value(), false)
            };

            let (toi_z, down) = if ls_ray.dir.z > N::zero() {
                let z = self.z_at(cell.0 + 1);
                ((z - ls_ray.origin.z) / ls_ray.dir.z, true)
            } else if ls_ray.dir.z < N::zero() {
                let z = self.z_at(cell.0 + 0);
                ((z - ls_ray.origin.z) / ls_ray.dir.z, false)
            } else {
                (N::max_value(), false)
            };

            if toi_x == N::max_value() && toi_z == N::max_value() {
                break;
            }

            if toi_x >= N::zero() && toi_x < toi_z {
                if right {
                    cell.1 += 1
                } else if cell.1 > 0 {
                    cell.1 -= 1
                } else {
                    break;
                }
            } else if toi_z >= N::zero() {
                if down {
                    cell.0 += 1
                } else if cell.0 > 0 {
                    cell.0 -= 1
                } else {
                    break
                }
            } else {
                break;
            }

            if cell.0 >= self.nrows() || cell.1 >= self.ncols() {
                break;
            }
        }

        None
    }
}
<commit_msg>Fix feature ID for 2D heightfield raycast.<commit_after>use bounding_volume::AABB;
use math::Isometry;
use na::{Point2, Real, Vector3};
use partitioning::{BestFirstBVVisitStatus, BestFirstDataVisitStatus, BestFirstVisitor};
use query::{closest_points_internal, Ray, RayCast, RayIntersection};
use shape::{CompositeShape, HeightField, FeatureId};

#[cfg(feature = "dim2")]
impl<N: Real> RayCast<N> for HeightField<N> {
    #[inline]
    fn toi_and_normal_with_ray(
        &self,
        m: &Isometry<N>,
        ray: &Ray<N>,
        _: bool,
    ) -> Option<RayIntersection<N>>
    {
        let aabb = self.aabb();
        let ls_ray = ray.inverse_transform_by(m);
        let (min_t, max_t) = aabb.clip_ray_parameters(&ls_ray)?;
        let (clip_ray_a, clip_ray_b) = (ls_ray.point_at(min_t), ls_ray.point_at(max_t));
        let mut curr = match self.cell_at_point(&clip_ray_a) {
            Some(cell) => cell,
            // None may happen due to slight numerical errors.
            None => if ls_ray.origin.x > N::zero() { self.num_cells() - 1 } else {
                0
            }
        };

        /*
         * Test the segment under the ray.
         */
        if let Some(seg) = self.segment_at(curr) {
            let (s, t) = closest_points_internal::line_against_line_parameters(&ray.origin, &ray.dir, seg.a(), &seg.scaled_direction());
            if s >= N::zero() && t >= N::zero() && t <= N::one() {
                // Cast succeeded on the first element!
                let n = seg.normal().unwrap().unwrap();
                let fid = if n.dot(&ls_ray.dir) > N::zero() {
                    // The ray hit the back face.
                    curr + self.num_cells()
                } else {
                    // The ray hit the front face.
                    curr
                };

                return Some(RayIntersection::new(s, m * n, FeatureId::Face(fid)))
            }
        }

        /*
         * Test other segments in the path of the ray.
         */
        if ls_ray.dir.x == N::zero() {
            return None;
        }

        let right = ls_ray.dir.x > N::zero();
        let cell_width = self.cell_width();
        let start_x = self.start_x();

        while (right && curr < self.num_cells()) || (!right && curr > 0) {
            let curr_param;

            if right {
                curr += 1;
                curr_param = (cell_width * na::convert(curr as f64) + start_x - ls_ray.origin.x) / ls_ray.dir.x;
            } else {
                curr_param = (ls_ray.origin.x - cell_width * na::convert(curr as f64) - start_x) / ls_ray.dir.x;
                curr -= 1;
            }

            if curr_param >= max_t {
                // The part of the ray after max_t is outside of the heightfield AABB.
                return None;
            }

            if let Some(seg) = self.segment_at(curr) {
                // TODO: test the y-coordinates (equivalent to an AABB test) before actually computing the intersection.
                let (s, t) = closest_points_internal::line_against_line_parameters(&ray.origin, &ray.dir, seg.a(), &seg.scaled_direction());
                if t >= N::zero() && t <= N::one() {
                    let n = seg.normal().unwrap().unwrap();
                    let fid = if n.dot(&ls_ray.dir) > N::zero() {
                        // The ray hit the back face.
                        curr + self.num_cells()
                    } else {
                        // The ray hit the front face.
                        curr
                    };
                    return Some(RayIntersection::new(s, m * n, FeatureId::Face(fid)))
                }
            }
        }

        None
    }
}

#[cfg(feature = "dim3")]
impl<N: Real> RayCast<N> for HeightField<N> {
    #[inline]
    fn toi_and_normal_with_ray(
        &self,
        m: &Isometry<N>,
        ray: &Ray<N>,
        solid: bool,
    ) -> Option<RayIntersection<N>>
    {
        let aabb = self.aabb();
        let ls_ray = ray.inverse_transform_by(m);
        let (min_t, max_t) = aabb.clip_ray_parameters(&ls_ray)?;
        let (clip_ray_a, clip_ray_b) = (ls_ray.point_at(min_t), ls_ray.point_at(max_t));
        let mut cell = match self.cell_at_point(&clip_ray_a) {
            Some(cell) => cell,
            // None may happen due to slight numerical errors.
            None => {
                let i = if ls_ray.origin.z > N::zero() {
                    self.nrows() - 1
                } else {
                    0
                };

                let j = if ls_ray.origin.x > N::zero() {
                    self.ncols() - 1
                } else {
                    0
                };

                (i, j)
            }
        };

        loop {
            let tris = self.triangles_at(cell.0, cell.1);
            let inter1 = tris.0.and_then(|tri| tri.toi_and_normal_with_ray(m, ray, solid));
            let inter2 = tris.1.and_then(|tri| tri.toi_and_normal_with_ray(m, ray, solid));

            match (inter1, inter2) {
                (Some(mut inter1), Some(mut inter2)) => {
                    if inter1.toi < inter2.toi {
                        inter1.feature = self.convert_triangle_feature_id(cell.0, cell.1, true, inter1.feature);
                        return Some(inter1);
                    } else {
                        inter2.feature = self.convert_triangle_feature_id(cell.0, cell.1, false, inter2.feature);
                        return Some(inter2);
                    }
                }
                (Some(mut inter), None) => {
                    inter.feature = self.convert_triangle_feature_id(cell.0, cell.1, true, inter.feature);
                    return Some(inter);
                }
                (None, Some(mut inter)) => {
                    inter.feature = self.convert_triangle_feature_id(cell.0, cell.1, false, inter.feature);
                    return Some(inter);
                }
                (None, None) => {}
            }

            /*
             * Find the next cell to cast the ray on.
             */
            let (toi_x, right) = if ls_ray.dir.x > N::zero() {
                let x = self.x_at(cell.1 + 1);
                ((x - ls_ray.origin.x) / ls_ray.dir.x, true)
            } else if ls_ray.dir.x < N::zero() {
                let x = self.x_at(cell.1 + 0);
                ((x - ls_ray.origin.x) / ls_ray.dir.x, false)
            } else {
                (N::max_value(), false)
            };

            let (toi_z, down) = if ls_ray.dir.z > N::zero() {
                let z = self.z_at(cell.0 + 1);
                ((z - ls_ray.origin.z) / ls_ray.dir.z, true)
            } else if ls_ray.dir.z < N::zero() {
                let z = self.z_at(cell.0 + 0);
                ((z - ls_ray.origin.z) / ls_ray.dir.z, false)
            } else {
                (N::max_value(), false)
            };

            if toi_x == N::max_value() && toi_z == N::max_value() {
                break;
            }

            if toi_x >= N::zero() && toi_x < toi_z {
                if right {
                    cell.1 += 1
                } else if cell.1 > 0 {
                    cell.1 -= 1
                } else {
                    break;
                }
            } else if toi_z >= N::zero() {
                if down {
                    cell.0 += 1
                } else if cell.0 > 0 {
                    cell.0 -= 1
                } else {
                    break
                }
            } else {
                break;
            }

            if cell.0 >= self.nrows() || cell.1 >= self.ncols() {
                break;
            }
        }

        None
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::error::Error;
use std::fs::File;
use std::io::Read;
use std::sync::Arc;

use crate::collections::HashSet;
use grpcio::{
    CertificateRequestType, Channel, ChannelBuilder, ChannelCredentialsBuilder, RpcContext,
    ServerBuilder, ServerCredentialsBuilder, ServerCredentialsFetcher,
};

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct SecurityConfig {
    pub ca_path: String,
    pub cert_path: String,
    pub key_path: String,
    // Test purpose only.
    #[serde(skip)]
    pub override_ssl_target: String,
    pub cipher_file: String,
    pub cert_allowed_cn: HashSet<String>,
}

impl Default for SecurityConfig {
    fn default() -> SecurityConfig {
        SecurityConfig {
            ca_path: String::new(),
            cert_path: String::new(),
            key_path: String::new(),
            override_ssl_target: String::new(),
            cipher_file: String::new(),
            cert_allowed_cn: HashSet::default(),
        }
    }
}

/// Checks and opens key file. Returns `Ok(None)` if the path is empty.
///
///  # Arguments
///
///  - `tag`: only used in the error message, like "ca key", "cert key", "private key", etc.
fn check_key_file(tag: &str, path: &str) -> Result<Option<File>, Box<dyn Error>> {
    if path.is_empty() {
        return Ok(None);
    }
    match File::open(path) {
        Err(e) => Err(format!("failed to open {} to load {}: {:?}", path, tag, e).into()),
        Ok(f) => Ok(Some(f)),
    }
}

/// Loads key file content. Returns `Ok(vec![])` if the path is empty.
fn load_key(tag: &str, path: &str) -> Result<Vec<u8>, Box<dyn Error>> {
    let mut key = vec![];
    let f = check_key_file(tag, path)?;
    match f {
        None => return Ok(vec![]),
        Some(mut f) => {
            if let Err(e) = f.read_to_end(&mut key) {
                return Err(format!("failed to load {} from path {}: {:?}", tag, path, e).into());
            }
        }
    }
    Ok(key)
}

type CertResult = Result<(Vec<u8>, Vec<u8>, Vec<u8>), Box<dyn Error>>;

impl SecurityConfig {
    /// Validates ca, cert and private key.
    pub fn validate(&self) -> Result<(), Box<dyn Error>> {
        check_key_file("ca key", &self.ca_path)?;
        check_key_file("cert key", &self.cert_path)?;
        check_key_file("private key", &self.key_path)?;
        // TODO: validate whether ca, cert and private key match.
        if (!self.ca_path.is_empty() || !self.cert_path.is_empty() || !self.key_path.is_empty())
            && (self.ca_path.is_empty() || self.cert_path.is_empty() || self.key_path.is_empty())
        {
            return Err("ca, cert and private key should be all configured.".into());
        }

        Ok(())
    }

    /// Load certificates from the given file path.
    /// Return ca, cert, key after successful read.
    fn load_certs(&self) -> CertResult {
        let ca = load_key("CA", &self.ca_path)?;
        let cert = load_key("certificate", &self.cert_path)?;
        let key = load_key("private key", &self.key_path)?;
        if ca.is_empty() || cert.is_empty() || key.is_empty() {
            return Err("ca, cert and private key should be all configured.".into());
        }
        Ok((ca, cert, key))
    }
}

#[derive(Default)]
pub struct SecurityManager {
    cfg: Arc<SecurityConfig>,
}

impl SecurityManager {
    pub fn new(cfg: &SecurityConfig) -> Result<SecurityManager, Box<dyn Error>> {
        Ok(SecurityManager {
            cfg: Arc::new(cfg.clone()),
        })
    }

    pub fn connect(&self, mut cb: ChannelBuilder, addr: &str) -> Channel {
        if self.cfg.ca_path.is_empty() {
            cb.connect(addr)
        } else {
            if !self.cfg.override_ssl_target.is_empty() {
                cb = cb.override_ssl_target(self.cfg.override_ssl_target.clone());
            }
            // Fill in empty certificate information if read fails.
            // Returning empty certificates delays error processing until
            // actual connection in grpc.
            let (ca, cert, key) = self.cfg.load_certs().unwrap_or_default();

            let cred = ChannelCredentialsBuilder::new()
                .root_cert(ca)
                .cert(cert, key)
                .build();
            cb.secure_connect(addr, cred)
        }
    }

    pub fn bind(&self, sb: ServerBuilder, addr: &str, port: u16) -> ServerBuilder {
        if self.cfg.ca_path.is_empty() {
            sb.bind(addr, port)
        } else {
            let fetcher = Box::new(Fetcher {
                cfg: self.cfg.clone(),
            });
            sb.bind_with_fetcher(
                addr,
                port,
                fetcher,
                CertificateRequestType::RequestAndRequireClientCertificateAndVerify,
            )
        }
    }

    pub fn cipher_file(&self) -> &str {
        &self.cfg.cipher_file
    }

    pub fn cert_allowed_cn(&self) -> &HashSet<String> {
        &self.cfg.cert_allowed_cn
    }
}

struct Fetcher {
    cfg: Arc<SecurityConfig>,
}

impl ServerCredentialsFetcher for Fetcher {
    fn fetch(&self) -> Result<Option<ServerCredentialsBuilder>, Box<dyn Error>> {
        let (ca, cert, key) = self.cfg.load_certs()?;
        let new_cred = ServerCredentialsBuilder::new()
            .add_cert(cert, key)
            .root_cert(
                ca,
                CertificateRequestType::RequestAndRequireClientCertificateAndVerify,
            );
        Ok(Some(new_cred))
    }
}

/// Check peer CN with cert-allowed-cn field.
/// Return true when the match is successful (support wildcard pattern).
/// Skip the check when cert-allowed-cn is not set or the secure channel is not used.
pub fn check_common_name(cert_allowed_cn: &HashSet<String>, ctx: &RpcContext) -> bool {
    if cert_allowed_cn.is_empty() {
        return true;
    }
    if let Some(auth_ctx) = ctx.auth_context() {
        if let Some(auth_property) = auth_ctx
            .into_iter()
            .find(|x| x.name() == "x509_common_name")
        {
            let peer_cn = auth_property.value_str().unwrap();
            match_peer_names(cert_allowed_cn, peer_cn)
        } else {
            false
        }
    } else {
        true
    }
}

/// Check peer CN with a set of allowed CN.
pub fn match_peer_names(allowed_cn: &HashSet<String>, name: &str) -> bool {
    for cn in allowed_cn {
        if cn == name {
            return true;
        }
    }
    false
}

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

    use std::fs;

    use tempfile::Builder;

    #[test]
    fn test_security() {
        let cfg = SecurityConfig::default();
        // default is disable secure connection.
        cfg.validate().unwrap();
        let mgr = SecurityManager::new(&cfg).unwrap();
        assert!(mgr.cfg.ca_path.is_empty());
        assert!(mgr.cfg.cert_path.is_empty());
        assert!(mgr.cfg.key_path.is_empty());

        let assert_cfg = |c: fn(&mut SecurityConfig), valid: bool| {
            let mut invalid_cfg = cfg.clone();
            c(&mut invalid_cfg);
            assert_eq!(invalid_cfg.validate().is_ok(), valid);
        };

        // invalid path should be rejected.
        assert_cfg(
            |c| {
                c.ca_path = "invalid ca path".to_owned();
                c.cert_path = "invalid cert path".to_owned();
                c.key_path = "invalid key path".to_owned();
            },
            false,
        );

        let temp = Builder::new().prefix("test_cred").tempdir().unwrap();
        let example_ca = temp.path().join("ca");
        let example_cert = temp.path().join("cert");
        let example_key = temp.path().join("key");
        for (id, f) in (&[&example_ca, &example_cert, &example_key])
            .iter()
            .enumerate()
        {
            fs::write(f, &[id as u8]).unwrap();
        }

        let mut c = cfg.clone();
        c.cert_path = format!("{}", example_cert.display());
        c.key_path = format!("{}", example_key.display());
        // incomplete configuration.
        c.validate().unwrap_err();

        // data should be loaded from file after validating.
        c.ca_path = format!("{}", example_ca.display());
        c.validate().unwrap();

        let (ca, cert, key) = c.load_certs().unwrap_or_default();
        assert_eq!(ca, vec![0]);
        assert_eq!(cert, vec![1]);
        assert_eq!(key, vec![2]);
    }
}
<commit_msg>Security: work around race condition between threads updating handshakers (#7251)<commit_after>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::error::Error;
use std::fs::{self, File};
use std::io::Read;
use std::sync::{Arc, Mutex};
use std::time::SystemTime;

use crate::collections::HashSet;
use grpcio::{
    CertificateRequestType, Channel, ChannelBuilder, ChannelCredentialsBuilder, RpcContext,
    ServerBuilder, ServerCredentialsBuilder, ServerCredentialsFetcher,
};

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct SecurityConfig {
    pub ca_path: String,
    pub cert_path: String,
    pub key_path: String,
    // Test purpose only.
    #[serde(skip)]
    pub override_ssl_target: String,
    pub cipher_file: String,
    pub cert_allowed_cn: HashSet<String>,
}

impl Default for SecurityConfig {
    fn default() -> SecurityConfig {
        SecurityConfig {
            ca_path: String::new(),
            cert_path: String::new(),
            key_path: String::new(),
            override_ssl_target: String::new(),
            cipher_file: String::new(),
            cert_allowed_cn: HashSet::default(),
        }
    }
}

/// Checks and opens key file. Returns `Ok(None)` if the path is empty.
///
///  # Arguments
///
///  - `tag`: only used in the error message, like "ca key", "cert key", "private key", etc.
fn check_key_file(tag: &str, path: &str) -> Result<Option<File>, Box<dyn Error>> {
    if path.is_empty() {
        return Ok(None);
    }
    match File::open(path) {
        Err(e) => Err(format!("failed to open {} to load {}: {:?}", path, tag, e).into()),
        Ok(f) => Ok(Some(f)),
    }
}

/// Loads key file content. Returns `Ok(vec![])` if the path is empty.
fn load_key(tag: &str, path: &str) -> Result<Vec<u8>, Box<dyn Error>> {
    let mut key = vec![];
    let f = check_key_file(tag, path)?;
    match f {
        None => return Ok(vec![]),
        Some(mut f) => {
            if let Err(e) = f.read_to_end(&mut key) {
                return Err(format!("failed to load {} from path {}: {:?}", tag, path, e).into());
            }
        }
    }
    Ok(key)
}

type CertResult = Result<(Vec<u8>, Vec<u8>, Vec<u8>), Box<dyn Error>>;

impl SecurityConfig {
    /// Validates ca, cert and private key.
    pub fn validate(&self) -> Result<(), Box<dyn Error>> {
        check_key_file("ca key", &self.ca_path)?;
        check_key_file("cert key", &self.cert_path)?;
        check_key_file("private key", &self.key_path)?;
        // TODO: validate whether ca, cert and private key match.
        if (!self.ca_path.is_empty() || !self.cert_path.is_empty() || !self.key_path.is_empty())
            && (self.ca_path.is_empty() || self.cert_path.is_empty() || self.key_path.is_empty())
        {
            return Err("ca, cert and private key should be all configured.".into());
        }

        Ok(())
    }

    /// Load certificates from the given file path.
    /// Return ca, cert, key after successful read.
    fn load_certs(&self) -> CertResult {
        let ca = load_key("CA", &self.ca_path)?;
        let cert = load_key("certificate", &self.cert_path)?;
        let key = load_key("private key", &self.key_path)?;
        if ca.is_empty() || cert.is_empty() || key.is_empty() {
            return Err("ca, cert and private key should be all configured.".into());
        }
        Ok((ca, cert, key))
    }

    /// Determine if the cert file has been modified.
    /// If modified, update the timestamp of this modification.
    fn is_modified(&self, last: &mut Option<SystemTime>) -> Result<bool, Box<dyn Error>> {
        let this = fs::metadata(&self.cert_path)?.modified()?;
        if let Some(last) = last {
            if *last == this {
                return Ok(false);
            }
        }
        *last = Some(this);
        Ok(true)
    }
}

#[derive(Default)]
pub struct SecurityManager {
    cfg: Arc<SecurityConfig>,
}

impl SecurityManager {
    pub fn new(cfg: &SecurityConfig) -> Result<SecurityManager, Box<dyn Error>> {
        Ok(SecurityManager {
            cfg: Arc::new(cfg.clone()),
        })
    }

    pub fn connect(&self, mut cb: ChannelBuilder, addr: &str) -> Channel {
        if self.cfg.ca_path.is_empty() {
            cb.connect(addr)
        } else {
            if !self.cfg.override_ssl_target.is_empty() {
                cb = cb.override_ssl_target(self.cfg.override_ssl_target.clone());
            }
            // Fill in empty certificate information if read fails.
            // Returning empty certificates delays error processing until
            // actual connection in grpc.
            let (ca, cert, key) = self.cfg.load_certs().unwrap_or_default();

            let cred = ChannelCredentialsBuilder::new()
                .root_cert(ca)
                .cert(cert, key)
                .build();
            cb.secure_connect(addr, cred)
        }
    }

    pub fn bind(&self, sb: ServerBuilder, addr: &str, port: u16) -> ServerBuilder {
        if self.cfg.ca_path.is_empty() {
            sb.bind(addr, port)
        } else {
            let fetcher = Box::new(Fetcher {
                cfg: self.cfg.clone(),
                last_modified_time: Arc::new(Mutex::new(None)),
            });
            sb.bind_with_fetcher(
                addr,
                port,
                fetcher,
                CertificateRequestType::RequestAndRequireClientCertificateAndVerify,
            )
        }
    }

    pub fn cipher_file(&self) -> &str {
        &self.cfg.cipher_file
    }

    pub fn cert_allowed_cn(&self) -> &HashSet<String> {
        &self.cfg.cert_allowed_cn
    }
}

struct Fetcher {
    cfg: Arc<SecurityConfig>,
    last_modified_time: Arc<Mutex<Option<SystemTime>>>,
}

impl ServerCredentialsFetcher for Fetcher {
    // Retrieves updated credentials. When returning `None` or
    // error, gRPC will continue to use the previous certificates
    // returned by the method.
    fn fetch(&self) -> Result<Option<ServerCredentialsBuilder>, Box<dyn Error>> {
        if let Ok(mut last) = self.last_modified_time.try_lock() {
            // Reload only when cert is modified.
            if self.cfg.is_modified(&mut last)? {
                let (ca, cert, key) = self.cfg.load_certs()?;
                let new_cred = ServerCredentialsBuilder::new()
                    .add_cert(cert, key)
                    .root_cert(
                        ca,
                        CertificateRequestType::RequestAndRequireClientCertificateAndVerify,
                    );
                Ok(Some(new_cred))
            } else {
                Ok(None)
            }
        } else {
            Ok(None)
        }
    }
}

/// Check peer CN with cert-allowed-cn field.
/// Return true when the match is successful (support wildcard pattern).
/// Skip the check when cert-allowed-cn is not set or the secure channel is not used.
pub fn check_common_name(cert_allowed_cn: &HashSet<String>, ctx: &RpcContext) -> bool {
    if cert_allowed_cn.is_empty() {
        return true;
    }
    if let Some(auth_ctx) = ctx.auth_context() {
        if let Some(auth_property) = auth_ctx
            .into_iter()
            .find(|x| x.name() == "x509_common_name")
        {
            let peer_cn = auth_property.value_str().unwrap();
            match_peer_names(cert_allowed_cn, peer_cn)
        } else {
            false
        }
    } else {
        true
    }
}

/// Check peer CN with a set of allowed CN.
pub fn match_peer_names(allowed_cn: &HashSet<String>, name: &str) -> bool {
    for cn in allowed_cn {
        if cn == name {
            return true;
        }
    }
    false
}

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

    use std::fs;
    use std::io::Write;
    use tempfile::Builder;

    #[test]
    fn test_security() {
        let cfg = SecurityConfig::default();
        // default is disable secure connection.
        cfg.validate().unwrap();
        let mgr = SecurityManager::new(&cfg).unwrap();
        assert!(mgr.cfg.ca_path.is_empty());
        assert!(mgr.cfg.cert_path.is_empty());
        assert!(mgr.cfg.key_path.is_empty());

        let assert_cfg = |c: fn(&mut SecurityConfig), valid: bool| {
            let mut invalid_cfg = cfg.clone();
            c(&mut invalid_cfg);
            assert_eq!(invalid_cfg.validate().is_ok(), valid);
        };

        // invalid path should be rejected.
        assert_cfg(
            |c| {
                c.ca_path = "invalid ca path".to_owned();
                c.cert_path = "invalid cert path".to_owned();
                c.key_path = "invalid key path".to_owned();
            },
            false,
        );

        let temp = Builder::new().prefix("test_cred").tempdir().unwrap();
        let example_ca = temp.path().join("ca");
        let example_cert = temp.path().join("cert");
        let example_key = temp.path().join("key");
        for (id, f) in (&[&example_ca, &example_cert, &example_key])
            .iter()
            .enumerate()
        {
            fs::write(f, &[id as u8]).unwrap();
        }

        let mut c = cfg.clone();
        c.cert_path = format!("{}", example_cert.display());
        c.key_path = format!("{}", example_key.display());
        // incomplete configuration.
        c.validate().unwrap_err();

        // data should be loaded from file after validating.
        c.ca_path = format!("{}", example_ca.display());
        c.validate().unwrap();

        let (ca, cert, key) = c.load_certs().unwrap_or_default();
        assert_eq!(ca, vec![0]);
        assert_eq!(cert, vec![1]);
        assert_eq!(key, vec![2]);
    }

    #[test]
    fn test_modify_file() {
        let mut file = Builder::new().prefix("test_modify").tempfile().unwrap();
        let mut cfg = SecurityConfig::default();
        cfg.cert_path = file.path().to_str().unwrap().to_owned();

        let mut last = None;
        assert!(cfg.is_modified(&mut last).unwrap());
        assert!(!cfg.is_modified(&mut last).unwrap());

        std::thread::sleep(std::time::Duration::from_millis(10));
        writeln!(file, "something").unwrap();
        assert!(cfg.is_modified(&mut last).unwrap());
        assert!(!cfg.is_modified(&mut last).unwrap());
    }
}
<|endoftext|>


<commit_before>use if_chain::if_chain;
use rustc::hir::{ImplItem, ImplItemKind};
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_lint_pass, declare_tool_lint};

use crate::utils::{
    get_trait_def_id, implements_trait, match_type, paths, return_ty, span_help_and_lint, trait_ref_of_method,
    walk_ptrs_ty,
};

declare_clippy_lint! {
    /// **What id does:** Checks for the definition of inherent methods with a signature of `to_string(&self) -> String`.
    ///
    /// **Why is this bad?** This method is also implicitly defined if a type implements the `Display` trait. As the functionality of `Display` is much more versatile, it should be preferred.
    ///
    /// **Known problems:** None
    ///
    /// ** Example:**
    ///
    /// ```rust
    /// // Bad
    /// pub struct A;
    ///
    /// impl A {
    ///     pub fn to_string(&self) -> String {
    ///         "I am A".to_string()
    ///     }
    /// }
    /// ```
    ///
    /// ```rust
    /// // Good
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A")
    ///     }
    /// }
    /// ```
    pub INHERENT_TO_STRING,
    style,
    "type implements inherent method `to_string()`, but should instead implement the `Display` trait"
}

declare_clippy_lint! {
    /// **What id does:** Checks for the definition of inherent methods with a signature of `to_string(&self) -> String` and if the type implementing this method also implements the `Display` trait.
    ///
    /// **Why is this bad?** This method is also implicitly defined if a type implements the `Display` trait. The less versatile inherent method will then shadow the implementation introduced by `Display`.
    ///
    /// **Known problems:** None
    ///
    /// ** Example:**
    ///
    /// ```rust
    /// // Bad
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl A {
    ///     pub fn to_string(&self) -> String {
    ///         "I am A".to_string()
    ///     }
    /// }
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A, too")
    ///     }
    /// }
    /// ```
    ///
    /// ```rust
    /// // Good
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A")
    ///     }
    /// }
    /// ```
    pub INHERENT_TO_STRING_SHADOW_DISPLAY,
    correctness,
    "type implements inherent method `to_string()`, which gets shadowed by the implementation of the `Display` trait "
}

declare_lint_pass!(InherentToString => [INHERENT_TO_STRING, INHERENT_TO_STRING_SHADOW_DISPLAY]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for InherentToString {
    fn check_impl_item(&mut self, cx: &LateContext<'a, 'tcx>, impl_item: &'tcx ImplItem) {
        if impl_item.span.from_expansion() {
            return;
        }

        if_chain! {
            // Check if item is a method, called to_string and has a parameter 'self'
            if let ImplItemKind::Method(ref signature, _) = impl_item.kind;
            if impl_item.ident.name.as_str() == "to_string";
            let decl = &signature.decl;
            if decl.implicit_self.has_implicit_self();
            if decl.inputs.len() == 1;

            // Check if return type is String
            if match_type(cx, return_ty(cx, impl_item.hir_id), &paths::STRING);

            // Filters instances of to_string which are required by a trait
            if trait_ref_of_method(cx, impl_item.hir_id).is_none();

            then {
                show_lint(cx, impl_item);
            }
        }
    }
}

fn show_lint(cx: &LateContext<'_, '_>, item: &ImplItem) {
    let display_trait_id =
        get_trait_def_id(cx, &["core", "fmt", "Display"]).expect("Failed to get trait ID of `Display`!");

    // Get the real type of 'self'
    let fn_def_id = cx.tcx.hir().local_def_id(item.hir_id);
    let self_type = cx.tcx.fn_sig(fn_def_id).input(0);
    let self_type = walk_ptrs_ty(self_type.skip_binder());

    // Emit either a warning or an error
    if implements_trait(cx, self_type, display_trait_id, &[]) {
        span_help_and_lint(
            cx,
            INHERENT_TO_STRING_SHADOW_DISPLAY,
            item.span,
            &format!(
                "type `{}` implements inherent method `to_string(&self) -> String` which shadows the implementation of `Display`",
                self_type.to_string()
            ),
            &format!("remove the inherent method from type `{}`", self_type.to_string())
        );
    } else {
        span_help_and_lint(
            cx,
            INHERENT_TO_STRING,
            item.span,
            &format!(
                "implementation of inherent method `to_string(&self) -> String` for type `{}`",
                self_type.to_string()
            ),
            &format!("implement trait `Display` for type `{}` instead", self_type.to_string()),
        );
    }
}
<commit_msg>Fix typo in inherent_to_string documentation<commit_after>use if_chain::if_chain;
use rustc::hir::{ImplItem, ImplItemKind};
use rustc::lint::{LateContext, LateLintPass, LintArray, LintPass};
use rustc::{declare_lint_pass, declare_tool_lint};

use crate::utils::{
    get_trait_def_id, implements_trait, match_type, paths, return_ty, span_help_and_lint, trait_ref_of_method,
    walk_ptrs_ty,
};

declare_clippy_lint! {
    /// **What it does:** Checks for the definition of inherent methods with a signature of `to_string(&self) -> String`.
    ///
    /// **Why is this bad?** This method is also implicitly defined if a type implements the `Display` trait. As the functionality of `Display` is much more versatile, it should be preferred.
    ///
    /// **Known problems:** None
    ///
    /// ** Example:**
    ///
    /// ```rust
    /// // Bad
    /// pub struct A;
    ///
    /// impl A {
    ///     pub fn to_string(&self) -> String {
    ///         "I am A".to_string()
    ///     }
    /// }
    /// ```
    ///
    /// ```rust
    /// // Good
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A")
    ///     }
    /// }
    /// ```
    pub INHERENT_TO_STRING,
    style,
    "type implements inherent method `to_string()`, but should instead implement the `Display` trait"
}

declare_clippy_lint! {
    /// **What it does:** Checks for the definition of inherent methods with a signature of `to_string(&self) -> String` and if the type implementing this method also implements the `Display` trait.
    ///
    /// **Why is this bad?** This method is also implicitly defined if a type implements the `Display` trait. The less versatile inherent method will then shadow the implementation introduced by `Display`.
    ///
    /// **Known problems:** None
    ///
    /// ** Example:**
    ///
    /// ```rust
    /// // Bad
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl A {
    ///     pub fn to_string(&self) -> String {
    ///         "I am A".to_string()
    ///     }
    /// }
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A, too")
    ///     }
    /// }
    /// ```
    ///
    /// ```rust
    /// // Good
    /// use std::fmt;
    ///
    /// pub struct A;
    ///
    /// impl fmt::Display for A {
    ///     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    ///         write!(f, "I am A")
    ///     }
    /// }
    /// ```
    pub INHERENT_TO_STRING_SHADOW_DISPLAY,
    correctness,
    "type implements inherent method `to_string()`, which gets shadowed by the implementation of the `Display` trait "
}

declare_lint_pass!(InherentToString => [INHERENT_TO_STRING, INHERENT_TO_STRING_SHADOW_DISPLAY]);

impl<'a, 'tcx> LateLintPass<'a, 'tcx> for InherentToString {
    fn check_impl_item(&mut self, cx: &LateContext<'a, 'tcx>, impl_item: &'tcx ImplItem) {
        if impl_item.span.from_expansion() {
            return;
        }

        if_chain! {
            // Check if item is a method, called to_string and has a parameter 'self'
            if let ImplItemKind::Method(ref signature, _) = impl_item.kind;
            if impl_item.ident.name.as_str() == "to_string";
            let decl = &signature.decl;
            if decl.implicit_self.has_implicit_self();
            if decl.inputs.len() == 1;

            // Check if return type is String
            if match_type(cx, return_ty(cx, impl_item.hir_id), &paths::STRING);

            // Filters instances of to_string which are required by a trait
            if trait_ref_of_method(cx, impl_item.hir_id).is_none();

            then {
                show_lint(cx, impl_item);
            }
        }
    }
}

fn show_lint(cx: &LateContext<'_, '_>, item: &ImplItem) {
    let display_trait_id =
        get_trait_def_id(cx, &["core", "fmt", "Display"]).expect("Failed to get trait ID of `Display`!");

    // Get the real type of 'self'
    let fn_def_id = cx.tcx.hir().local_def_id(item.hir_id);
    let self_type = cx.tcx.fn_sig(fn_def_id).input(0);
    let self_type = walk_ptrs_ty(self_type.skip_binder());

    // Emit either a warning or an error
    if implements_trait(cx, self_type, display_trait_id, &[]) {
        span_help_and_lint(
            cx,
            INHERENT_TO_STRING_SHADOW_DISPLAY,
            item.span,
            &format!(
                "type `{}` implements inherent method `to_string(&self) -> String` which shadows the implementation of `Display`",
                self_type.to_string()
            ),
            &format!("remove the inherent method from type `{}`", self_type.to_string())
        );
    } else {
        span_help_and_lint(
            cx,
            INHERENT_TO_STRING,
            item.span,
            &format!(
                "implementation of inherent method `to_string(&self) -> String` for type `{}`",
                self_type.to_string()
            ),
            &format!("implement trait `Display` for type `{}` instead", self_type.to_string()),
        );
    }
}
<|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 core::cmp::Ordering;
use std::fmt;
use std::path::PathBuf;

use intern::string_key::Intern;
use intern::string_key::StringKey;

use crate::span::Span;

/// The location of a source. Could be a standalone file (e.g. test.graphql),
/// an embedded source (GraphQL tag in a JS file) or generated code without a
/// location.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum SourceLocationKey {
    /// A source embedded within a file. The 0-based index is an index into the
    /// embedded sources. E.g. the second graphql tag has index 1.
    Embedded {
        path: StringKey,
        index: u16,
    },
    Standalone {
        path: StringKey,
    },
    Generated,
}

impl SourceLocationKey {
    /// Returns a `SourceLocationKey` that's not backed by a real file. In most
    /// cases it's preferred to use a related real file.
    pub fn generated() -> Self {
        SourceLocationKey::Generated
    }

    pub fn standalone(path: &str) -> Self {
        SourceLocationKey::Standalone {
            path: path.intern(),
        }
    }
    pub fn embedded(path: &str, index: usize) -> Self {
        SourceLocationKey::Embedded {
            path: path.intern(),
            index: index.try_into().unwrap(),
        }
    }

    pub fn path(self) -> &'static str {
        match self {
            SourceLocationKey::Embedded { path, .. } => path.lookup(),
            SourceLocationKey::Standalone { path } => path.lookup(),
            SourceLocationKey::Generated => "<generated>",
        }
    }

    pub fn get_dir(self) -> PathBuf {
        let mut path = PathBuf::from(self.path());
        path.pop();
        path
    }
}

/// An absolute source location describing both the file and position (span)
/// with that file.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Location {
    /// The source containing this location (e.g. embedded or standalone file).
    source_location: SourceLocationKey,

    /// Relative position with the file
    span: Span,
}

impl fmt::Debug for Location {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{:?}", self.source_location.path(), self.span)
    }
}

impl Location {
    /// Returns a location that's not backed by a real file. In most cases it's
    /// preferred to use a related real location.
    pub fn generated() -> Self {
        Location::new(SourceLocationKey::generated(), Span::new(0, 0))
    }

    pub fn new(source_location: SourceLocationKey, span: Span) -> Self {
        Self {
            source_location,
            span,
        }
    }

    pub fn source_location(&self) -> SourceLocationKey {
        self.source_location
    }

    pub fn span(&self) -> &Span {
        &self.span
    }

    pub fn with_span(&self, span: Span) -> Self {
        Self {
            source_location: self.source_location,
            span,
        }
    }

    pub fn contains(&self, subspan: Span) -> bool {
        self.span.contains(subspan)
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct WithLocation<T> {
    pub location: Location,
    pub item: T,
}

impl<T: Ord> Ord for WithLocation<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.item.cmp(&other.item)
    }
}

impl<T: PartialOrd> PartialOrd for WithLocation<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.item.partial_cmp(&other.item)
    }
}

impl<T> WithLocation<T> {
    pub fn from_span(source_location: SourceLocationKey, span: Span, item: T) -> Self {
        Self {
            location: Location::new(source_location, span),
            item,
        }
    }

    pub fn new(location: Location, item: T) -> Self {
        Self { location, item }
    }

    /// Wraps the given item in a `WithLocation`, without associating it with
    /// any real location. In most cases it's preferred to use a related real
    /// location the item was derived from.
    pub fn generated(item: T) -> Self {
        Self {
            location: Location::generated(),
            item,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn with_location_memory() {
        assert_eq!(20, std::mem::size_of::<WithLocation<StringKey>>());
    }

    #[test]
    fn location_debug_printing() {
        assert_eq!(
            format!(
                "{:?}",
                Location::new(
                    SourceLocationKey::embedded("example/file.js", 2),
                    Span::new(10, 30)
                )
            ),
            "example/file.js:10:30".to_string()
        );
    }
}
<commit_msg>Adding .map to WithLocation<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 core::cmp::Ordering;
use std::fmt;
use std::path::PathBuf;

use intern::string_key::Intern;
use intern::string_key::StringKey;

use crate::span::Span;

/// The location of a source. Could be a standalone file (e.g. test.graphql),
/// an embedded source (GraphQL tag in a JS file) or generated code without a
/// location.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum SourceLocationKey {
    /// A source embedded within a file. The 0-based index is an index into the
    /// embedded sources. E.g. the second graphql tag has index 1.
    Embedded {
        path: StringKey,
        index: u16,
    },
    Standalone {
        path: StringKey,
    },
    Generated,
}

impl SourceLocationKey {
    /// Returns a `SourceLocationKey` that's not backed by a real file. In most
    /// cases it's preferred to use a related real file.
    pub fn generated() -> Self {
        SourceLocationKey::Generated
    }

    pub fn standalone(path: &str) -> Self {
        SourceLocationKey::Standalone {
            path: path.intern(),
        }
    }
    pub fn embedded(path: &str, index: usize) -> Self {
        SourceLocationKey::Embedded {
            path: path.intern(),
            index: index.try_into().unwrap(),
        }
    }

    pub fn path(self) -> &'static str {
        match self {
            SourceLocationKey::Embedded { path, .. } => path.lookup(),
            SourceLocationKey::Standalone { path } => path.lookup(),
            SourceLocationKey::Generated => "<generated>",
        }
    }

    pub fn get_dir(self) -> PathBuf {
        let mut path = PathBuf::from(self.path());
        path.pop();
        path
    }
}

/// An absolute source location describing both the file and position (span)
/// with that file.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct Location {
    /// The source containing this location (e.g. embedded or standalone file).
    source_location: SourceLocationKey,

    /// Relative position with the file
    span: Span,
}

impl fmt::Debug for Location {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}:{:?}", self.source_location.path(), self.span)
    }
}

impl Location {
    /// Returns a location that's not backed by a real file. In most cases it's
    /// preferred to use a related real location.
    pub fn generated() -> Self {
        Location::new(SourceLocationKey::generated(), Span::new(0, 0))
    }

    pub fn new(source_location: SourceLocationKey, span: Span) -> Self {
        Self {
            source_location,
            span,
        }
    }

    pub fn source_location(&self) -> SourceLocationKey {
        self.source_location
    }

    pub fn span(&self) -> &Span {
        &self.span
    }

    pub fn with_span(&self, span: Span) -> Self {
        Self {
            source_location: self.source_location,
            span,
        }
    }

    pub fn contains(&self, subspan: Span) -> bool {
        self.span.contains(subspan)
    }
}

#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct WithLocation<T> {
    pub location: Location,
    pub item: T,
}

impl<T: Ord> Ord for WithLocation<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.item.cmp(&other.item)
    }
}

impl<T: PartialOrd> PartialOrd for WithLocation<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.item.partial_cmp(&other.item)
    }
}

impl<T> WithLocation<T> {
    pub fn from_span(source_location: SourceLocationKey, span: Span, item: T) -> Self {
        Self {
            location: Location::new(source_location, span),
            item,
        }
    }

    pub fn new(location: Location, item: T) -> Self {
        Self { location, item }
    }

    /// Wraps the given item in a `WithLocation`, without associating it with
    /// any real location. In most cases it's preferred to use a related real
    /// location the item was derived from.
    pub fn generated(item: T) -> Self {
        Self {
            location: Location::generated(),
            item,
        }
    }

    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> WithLocation<U> {
        WithLocation {
            location: self.location,
            item: f(self.item),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn with_location_memory() {
        assert_eq!(20, std::mem::size_of::<WithLocation<StringKey>>());
    }

    #[test]
    fn location_debug_printing() {
        assert_eq!(
            format!(
                "{:?}",
                Location::new(
                    SourceLocationKey::embedded("example/file.js", 2),
                    Span::new(10, 30)
                )
            ),
            "example/file.js:10:30".to_string()
        );
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

extern crate sysinfo;

use sysinfo::ProcessExt;
use sysinfo::SystemExt;

#[test]
fn test_process() {
    let mut s = sysinfo::System::new();
    assert_eq!(s.get_processes().len(), 0);
    s.refresh_processes();
    assert!(!s.get_processes().is_empty());
    #[cfg(not(windows))]
    assert!(s
        .get_processes()
        .values()
        .any(|p| !p.exe().to_str().unwrap_or("").is_empty()));
}

#[test]
fn test_process_refresh() {
    let mut s = sysinfo::System::new();
    assert_eq!(s.get_processes().len(), 0);
    s.refresh_process(sysinfo::get_current_pid().expect("failed to get current pid"));
    assert_eq!(
        s.get_process(sysinfo::get_current_pid().expect("failed to get current pid"))
            .is_some(),
        true
    );
}

#[test]
#[cfg(windows)]
fn test_get_cmd_line() {
    let p = std::process::Command::new("timeout")
        .arg("/t")
        .arg("3")
        .spawn()
        .unwrap();
    let mut s = sysinfo::System::new();
    assert!(s.get_processes().len() == 0);
    s.refresh_processes();
    assert!(s.get_processes().len() > 0);
    if let Some(process) = s.get_process(p.id() as sysinfo::Pid) {
        assert_eq!(process.cmd(), &["timeout", "/t", "3"]);
    } else {
        // We're very likely on a "linux-like" shell so let's try some unix command...
        unix_like_cmd();
    }
}

#[test]
#[cfg(not(windows))]
fn test_get_cmd_line() {
    unix_like_cmd();
}

fn unix_like_cmd() {
    let p = std::process::Command::new("sleep")
        .arg("3")
        .spawn()
        .unwrap();
    let mut s = sysinfo::System::new();
    assert!(s.get_processes().is_empty());
    s.refresh_processes();
    assert!(!s.get_processes().is_empty());
    let process = s.get_process(p.id() as sysinfo::Pid).unwrap();
    assert_eq!(process.cmd(), &["sleep", "3"]);
}

#[test]
fn test_process_disk_usage() {
    use std::fs;
    use std::fs::File;
    use std::io::prelude::*;
    use sysinfo::{get_current_pid, ProcessExt, SystemExt};
    {
        let mut file = File::create("test.txt").unwrap();
        file.write_all(b"This is a test file\nwith test data.\n")
            .unwrap();
    }
    fs::remove_file("test.txt").ok();
    let mut system = sysinfo::System::new();
    assert!(system.get_processes().is_empty());
    system.refresh_processes();
    assert!(!system.get_processes().is_empty());
    let p = system
        .get_process(get_current_pid().expect("Failed retrieving current pid."))
        .expect("failed to get process");

    assert!(
        p.disk_usage().total_written_bytes > 0,
        "found {} total written bytes...",
        p.disk_usage().total_written_bytes
    );
    assert!(
        p.disk_usage().written_bytes > 0,
        "found {} written bytes...",
        p.disk_usage().written_bytes
    );
}

#[test]
fn cpu_usage_is_not_nan() {
    let mut system = sysinfo::System::new();
    system.refresh_processes();

    let first_pids = system
        .get_processes()
        .iter()
        .take(10)
        .map(|(&pid, _)| pid)
        .collect::<Vec<_>>();
    let mut checked = 0;

    first_pids.into_iter().for_each(|pid| {
        system.refresh_process(pid);
        if let Some(p) = system.get_process(pid) {
            assert!(!p.cpu_usage().is_nan());
            checked += 1;
        }
    });
    assert!(checked > 0);
}
<commit_msg>Add sleep in linux test<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

extern crate sysinfo;

use sysinfo::ProcessExt;
use sysinfo::SystemExt;

#[test]
fn test_process() {
    let mut s = sysinfo::System::new();
    assert_eq!(s.get_processes().len(), 0);
    s.refresh_processes();
    assert!(!s.get_processes().is_empty());
    #[cfg(not(windows))]
    assert!(s
        .get_processes()
        .values()
        .any(|p| !p.exe().to_str().unwrap_or("").is_empty()));
}

#[test]
fn test_process_refresh() {
    let mut s = sysinfo::System::new();
    assert_eq!(s.get_processes().len(), 0);
    s.refresh_process(sysinfo::get_current_pid().expect("failed to get current pid"));
    assert_eq!(
        s.get_process(sysinfo::get_current_pid().expect("failed to get current pid"))
            .is_some(),
        true
    );
}

#[test]
#[cfg(windows)]
fn test_get_cmd_line() {
    let p = std::process::Command::new("timeout")
        .arg("/t")
        .arg("3")
        .spawn()
        .unwrap();
    let mut s = sysinfo::System::new();
    assert!(s.get_processes().len() == 0);
    s.refresh_processes();
    assert!(s.get_processes().len() > 0);
    if let Some(process) = s.get_process(p.id() as sysinfo::Pid) {
        assert_eq!(process.cmd(), &["timeout", "/t", "3"]);
    } else {
        // We're very likely on a "linux-like" shell so let's try some unix command...
        unix_like_cmd();
    }
}

#[test]
#[cfg(not(windows))]
fn test_get_cmd_line() {
    unix_like_cmd();
}

fn unix_like_cmd() {
    use std::{thread, time};

    let p = std::process::Command::new("sleep")
        .arg("3")
        .spawn()
        .unwrap();
    // To ensure that the system data are filled correctly...
    thread::sleep(time::Duration::from_millis(250));
    let mut s = sysinfo::System::new();
    assert!(s.get_processes().is_empty());
    s.refresh_processes();
    assert!(!s.get_processes().is_empty());
    let process = s.get_process(p.id() as sysinfo::Pid).unwrap();
    assert_eq!(process.cmd(), &["sleep", "3"]);
}

#[test]
fn test_process_disk_usage() {
    use std::fs;
    use std::fs::File;
    use std::io::prelude::*;
    use sysinfo::{get_current_pid, ProcessExt, SystemExt};
    {
        let mut file = File::create("test.txt").unwrap();
        file.write_all(b"This is a test file\nwith test data.\n")
            .unwrap();
    }
    fs::remove_file("test.txt").ok();
    let mut system = sysinfo::System::new();
    assert!(system.get_processes().is_empty());
    system.refresh_processes();
    assert!(!system.get_processes().is_empty());
    let p = system
        .get_process(get_current_pid().expect("Failed retrieving current pid."))
        .expect("failed to get process");

    assert!(
        p.disk_usage().total_written_bytes > 0,
        "found {} total written bytes...",
        p.disk_usage().total_written_bytes
    );
    assert!(
        p.disk_usage().written_bytes > 0,
        "found {} written bytes...",
        p.disk_usage().written_bytes
    );
}

#[test]
fn cpu_usage_is_not_nan() {
    let mut system = sysinfo::System::new();
    system.refresh_processes();

    let first_pids = system
        .get_processes()
        .iter()
        .take(10)
        .map(|(&pid, _)| pid)
        .collect::<Vec<_>>();
    let mut checked = 0;

    first_pids.into_iter().for_each(|pid| {
        system.refresh_process(pid);
        if let Some(p) = system.get_process(pid) {
            assert!(!p.cpu_usage().is_nan());
            checked += 1;
        }
    });
    assert!(checked > 0);
}
<|endoftext|>


<commit_before>use std::{mem, slice, fmt};
use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::ops::{Index, IndexMut};
use std::marker::PhantomData;
use std::collections::LinkedList;
use std::ops::{Deref, DerefMut};

/// Arc is the only valid way to access an item in
/// the pool. It is returned by alloc, and will automatically
/// release/retain when dropped/cloned. It implements Deref/DerefMut,
/// so all accesses can go through it.
pub struct Arc<T> {
    pool: *mut Pool<T>,
    index: usize,
}

/// Public functions
impl <T> Arc<T> {
    pub unsafe fn retain(&mut self) {
        self.get_pool().retain(self.index);
    }

    pub unsafe fn release(&mut self) {
        self.get_pool().release(self.index);
    }
}

/// Internal functions
impl <T> Arc<T> {
    fn new(index: usize, p: &Pool<T>) -> Arc<T> {
        Arc {
            pool: unsafe { mem::transmute(p) },
            index: index,
        }
    }

    fn get_pool(&self) -> &mut Pool<T> {
        unsafe {
            &mut *self.pool
        }
    }

    fn ref_count(&self) -> usize {
        self.get_pool().header_for(self.index).ref_count.load(Ordering::Relaxed)
    }
}

impl <T> Drop for Arc<T> {
    fn drop(&mut self) {
        self.get_pool().release(self.index);
    }
}

impl <T> Clone for Arc<T> {
    fn clone(&self) -> Self {
        self.get_pool().retain(self.index);
        Arc {
            pool: self.pool,
            index: self.index,
        }
    }
}

impl<T> Deref for Arc<T> {
    type Target = T;

    fn deref<'b>(&'b self) -> &'b T {
        &self.get_pool()[self.index]
    }
}

impl<T> DerefMut for Arc<T> {
    fn deref_mut<'b>(&'b mut self) -> &'b mut T {
        &mut self.get_pool()[self.index]
    }
}

impl <T> fmt::Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Arc{{ offset: {:?}, ref_count: {:?} }}", self.index, self.ref_count())
    }
}

impl <T> PartialEq for Arc<T> {
    fn eq(&self, other: &Arc<T>) -> bool {
        if self.index != other.index {
            false
        } else {
            unsafe {
                self.pool as *const _ == other.pool as *const _
            }
        }
    }
}

/// A pool represents a fixed number of ref-counted objects.
/// The pool treats all given space as an unallocated
/// pool of objects. Each object is prefixed with a header.
/// The header is formatted as follows:
/// * V1
///   - [0..2] ref_count: u16
///
pub struct Pool<T> {
    item_type: PhantomData<T>,

    buffer: *mut u8,
    buffer_size: usize,
    capacity: usize,

    tail: AtomicUsize, // One past the end index

    // Cached values
    slot_size: usize,
    header_size: usize,

    free_list: LinkedList<usize>,
}

struct SlotHeader {
    ref_count: AtomicUsize,
}

/// Public interface
impl <T> Pool<T> {
    pub fn new(mem: &mut [u8]) -> Pool<T> {
        let ptr: *mut u8 = mem.as_mut_ptr();
        let header_size = mem::size_of::<SlotHeader>();
        let slot_size = mem::size_of::<T>() + header_size;
        Pool {
            item_type: PhantomData,
            buffer: ptr,
            buffer_size: mem.len(),
            tail: AtomicUsize::new(0),
            slot_size: slot_size,
            capacity: mem.len() / slot_size,
            header_size: header_size,
            free_list: LinkedList::new(),
        }
    }

    /// Remove all objects from the pool
    /// and zero the memory
    pub unsafe fn clear(&mut self) {
        let mut i = self.buffer.clone();
        let end = self.buffer.clone().offset(self.buffer_size as isize);
        while i != end {
            *i = 0u8;
            i = i.offset(1);
        }
    }

    /// Try to allocate a new item from the pool.
    /// A mutable reference to the item is returned on success
    pub fn alloc(&mut self) -> Result<Arc<T>, &'static str> {
        let index = try!(self.claim_free_index());
        Ok(Arc::new(index, self))
    }

    // Increase the ref count for the cell at the given index
    pub fn retain(&mut self, index: usize) {
        let h = self.header_for(index);
        loop {
            let old = h.ref_count.load(Ordering::Relaxed);
            let swap = h.ref_count
            .compare_and_swap(old, old+1, Ordering::Relaxed);
            if swap == old {
                break
            }
        }
    }

    // Decrease the ref count for the cell at the given index
    pub fn release(&mut self, index: usize) {
        let mut is_free = false;
        { // Make the borrow checker happy
            let h = self.header_for(index);
            loop {
                let old = h.ref_count.load(Ordering::Relaxed);
                assert!(old > 0, "Release called on [{}] which has no refs!", index);

                let swap = h.ref_count
                .compare_and_swap(old, old-1, Ordering::Relaxed);
                if swap == old {
                    if old == 1 { // this was the last reference
                        is_free = true;
                    }
                    break
                }
            }
        }
        if is_free {
            self.free_list.push_back(index);
        }
    }

    /// Returns the number of live items. O(1) running time.
    pub fn live_count(&self) -> usize {
        self.tail.load(Ordering::Relaxed) - self.free_list.len()
    }
}

/// Internal Functions
impl <T> Pool<T> {
    // Returns an item from the free list, or
    // tries to allocate a new one from the buffer
    fn claim_free_index(&mut self) -> Result<usize, &'static str> {
        let index = match self.free_list.pop_front() {
            Some(i) => i,
            None => try!(self.push_back_alloc()),
        };
        self.retain(index);
        Ok(index)
    }

    // Pushes the end of the used space in the buffer back
    // returns the previous index
    fn push_back_alloc(&mut self) -> Result<usize, &'static str> {
        loop {
            let old_tail = self.tail.load(Ordering::Relaxed);
            let swap = self.tail.compare_and_swap(old_tail, old_tail+1, Ordering::Relaxed);
            // If we were the ones to claim this slot, or
            // we've overrun the buffer, return
            if old_tail >= self.capacity {
                return Err("OOM")
            } else if swap == old_tail {
                return Ok(old_tail)
            }
        }
    }

    fn header_for<'a>(&'a mut self, i: usize) -> &'a mut SlotHeader {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize);
            mem::transmute(ptr)
        }
    }
}

impl <T> Index<usize> for Pool<T> {
    type Output = T;

    fn index<'a>(&'a self, i: usize) -> &'a T {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize)
            .offset(self.header_size as isize);
            mem::transmute(ptr)
        }
    }
}

impl <T> IndexMut<usize> for Pool<T> {
    fn index_mut<'a>(&'a mut self, i: usize) -> &'a mut T {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize)
            .offset(self.header_size as isize);
            mem::transmute(ptr)
        }
    }
}

#[test]
fn release_frees() {
       let mut buf: [u8; 100] = [0; 100];
       let mut p = Pool::<u32>::new(&mut buf[..]);
       p.alloc().unwrap();
       p.alloc().unwrap();

       assert_eq!(2, p.live_count());

       p.release(0);
       assert_eq!(1, p.live_count());
       assert_eq!(1, p.free_list.len());
       assert_eq!(0, *p.free_list.front().unwrap());

       p.release(1);
       assert_eq!(0, p.live_count());
       assert_eq!(2, p.free_list.len());
}

#[test]
fn alloc_after_free_recycles() {
       let mut buf: [u8; 100] = [0; 100];
       let mut p = Pool::<u32>::new(&mut buf[..]);
       p.alloc().unwrap();
       assert_eq!(1, p.live_count());
       assert_eq!(1, p.tail.load(Ordering::Relaxed));

       p.release(0);
       assert_eq!(0, p.live_count());
       assert_eq!(1, p.free_list.len());

       p.alloc().unwrap();
       assert_eq!(1, p.tail.load(Ordering::Relaxed)); // Tail shouldn't move
       assert_eq!(1, p.live_count());
       assert_eq!(0, p.free_list.len());
}

fn arc_clone() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        assert_eq!(1, int1.ref_count());
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        {
            let mut int1_c:Arc<u32> = int1.clone(); // Should bump the ref count
            assert_eq!(2, int1.ref_count());
            assert_eq!(2, int1_c.ref_count());
            assert_eq!(2, p.header_for(0).ref_count.load(Ordering::Relaxed));
        }
        // Now, the clone should have been dropped, but no memory reclaimed
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        assert_eq!(0, p.free_list.len());
    }
    // Now, int1 should have been dropped, and all memory reclaimed
    assert_eq!(0, p.header_for(0).ref_count.load(Ordering::Relaxed));
    assert_eq!(1, p.free_list.len());
}

#[test]
fn arc_drop() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        assert_eq!(1, int1.ref_count());
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        {
            let mut int2:Arc<u32> = p.alloc().unwrap();
            assert_eq!(1, int2.ref_count());
            assert_eq!(1, p.header_for(1).ref_count.load(Ordering::Relaxed));
        }
        // Now, int2 should have been dropped
        assert_eq!(0, p.header_for(1).ref_count.load(Ordering::Relaxed));
        assert_eq!(1, p.free_list.len());
    }
    // Now, int1 should have been dropped
    assert_eq!(0, p.header_for(0).ref_count.load(Ordering::Relaxed));
    assert_eq!(2, p.free_list.len());
}

#[test]
fn construction() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);

    assert_eq!(100, p.buffer_size);
    assert_eq!(mem::size_of::<usize>(), p.header_size);

    let expected_size = mem::size_of::<usize>() + mem::size_of::<u32>();
    assert_eq!(expected_size, p.slot_size);
    assert_eq!(100/expected_size, p.capacity); // expected_size should be 8+4=12
    assert_eq!(8, p.capacity);
}

#[test]
fn free_list_alloc_works() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        *int1 = 42;
        // Check payload
        assert_eq!([42u8, 0u8, 0u8, 0u8][..], buf[8..12]);
        // Check ref_count
        assert_eq!([1u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8][..], buf[0..8]);
        assert_eq!(1, p.live_count());
    }
    // int1 is now out of scope, let's ensure the drop worked
    assert_eq!([0u8; 8][..], buf[0..8]);
}

#[test]
fn check_oom_error() {
    let mut buf: [u8; 1] = [0; 1];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    assert_eq!(Err("OOM"), p.alloc());
}

#[test]
fn multiple_allocations_work() {
    let mut buf: [u8; 120] = [0; 120];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    for i in 0..10 {
        let mut int1 = p.alloc().unwrap();
        *int1 = i;
        unsafe { int1.retain() }; // Make sure this stays around long enough to read later
   }
   assert_eq!(10, p.live_count());
   let expected_ref_count = [1u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
   for i in 0..10 {
       let start = 12*i;
       // Check ref_count
       assert_eq!(expected_ref_count[..], buf[start..start+8]);
       // Check payload
       assert_eq!([i as u8, 0u8, 0u8, 0u8][..], buf[start+8..start+12]);
    }
}
<commit_msg>Added internal alloc to test retain/release directly without Arc::Drop<commit_after>use std::{mem, slice, fmt};
use std::cell::UnsafeCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::ops::{Index, IndexMut};
use std::marker::PhantomData;
use std::collections::LinkedList;
use std::ops::{Deref, DerefMut};

/// Arc is the only valid way to access an item in
/// the pool. It is returned by alloc, and will automatically
/// release/retain when dropped/cloned. It implements Deref/DerefMut,
/// so all accesses can go through it.
pub struct Arc<T> {
    pool: *mut Pool<T>,
    index: usize,
}

/// Public functions
impl <T> Arc<T> {
    pub unsafe fn retain(&mut self) {
        self.get_pool().retain(self.index);
    }

    pub unsafe fn release(&mut self) {
        self.get_pool().release(self.index);
    }
}

/// Internal functions
impl <T> Arc<T> {
    fn new(index: usize, p: &Pool<T>) -> Arc<T> {
        Arc {
            pool: unsafe { mem::transmute(p) },
            index: index,
        }
    }

    fn get_pool(&self) -> &mut Pool<T> {
        unsafe {
            &mut *self.pool
        }
    }

    fn ref_count(&self) -> usize {
        self.get_pool().header_for(self.index).ref_count.load(Ordering::Relaxed)
    }
}

impl <T> Drop for Arc<T> {
    fn drop(&mut self) {
        self.get_pool().release(self.index);
    }
}

impl <T> Clone for Arc<T> {
    fn clone(&self) -> Self {
        self.get_pool().retain(self.index);
        Arc {
            pool: self.pool,
            index: self.index,
        }
    }
}

impl<T> Deref for Arc<T> {
    type Target = T;

    fn deref<'b>(&'b self) -> &'b T {
        &self.get_pool()[self.index]
    }
}

impl<T> DerefMut for Arc<T> {
    fn deref_mut<'b>(&'b mut self) -> &'b mut T {
        &mut self.get_pool()[self.index]
    }
}

impl <T> fmt::Debug for Arc<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Arc{{ offset: {:?}, ref_count: {:?} }}", self.index, self.ref_count())
    }
}

impl <T> PartialEq for Arc<T> {
    fn eq(&self, other: &Arc<T>) -> bool {
        if self.index != other.index {
            false
        } else {
            unsafe {
                self.pool as *const _ == other.pool as *const _
            }
        }
    }
}

/// A pool represents a fixed number of ref-counted objects.
/// The pool treats all given space as an unallocated
/// pool of objects. Each object is prefixed with a header.
/// The header is formatted as follows:
/// * V1
///   - [0..2] ref_count: u16
///
pub struct Pool<T> {
    item_type: PhantomData<T>,

    buffer: *mut u8,
    buffer_size: usize,
    capacity: usize,

    tail: AtomicUsize, // One past the end index

    // Cached values
    slot_size: usize,
    header_size: usize,

    free_list: LinkedList<usize>,
}

struct SlotHeader {
    ref_count: AtomicUsize,
}

/// Public interface
impl <T> Pool<T> {
    pub fn new(mem: &mut [u8]) -> Pool<T> {
        let ptr: *mut u8 = mem.as_mut_ptr();
        let header_size = mem::size_of::<SlotHeader>();
        let slot_size = mem::size_of::<T>() + header_size;
        Pool {
            item_type: PhantomData,
            buffer: ptr,
            buffer_size: mem.len(),
            tail: AtomicUsize::new(0),
            slot_size: slot_size,
            capacity: mem.len() / slot_size,
            header_size: header_size,
            free_list: LinkedList::new(),
        }
    }

    /// Remove all objects from the pool
    /// and zero the memory
    pub unsafe fn clear(&mut self) {
        let mut i = self.buffer.clone();
        let end = self.buffer.clone().offset(self.buffer_size as isize);
        while i != end {
            *i = 0u8;
            i = i.offset(1);
        }
    }

    /// Try to allocate a new item from the pool.
    /// A mutable reference to the item is returned on success
    pub fn alloc(&mut self) -> Result<Arc<T>, &'static str> {
        let index = try!(self.internal_alloc());
        Ok(Arc::new(index, self))
    }

    // Increase the ref count for the cell at the given index
    pub fn retain(&mut self, index: usize) {
        let h = self.header_for(index);
        loop {
            let old = h.ref_count.load(Ordering::Relaxed);
            let swap = h.ref_count
            .compare_and_swap(old, old+1, Ordering::Relaxed);
            if swap == old {
                break
            }
        }
    }

    // Decrease the ref count for the cell at the given index
    pub fn release(&mut self, index: usize) {
        let mut is_free = false;
        { // Make the borrow checker happy
            let h = self.header_for(index);
            loop {
                let old = h.ref_count.load(Ordering::Relaxed);
                assert!(old > 0, "Release called on [{}] which has no refs!", index);

                let swap = h.ref_count
                .compare_and_swap(old, old-1, Ordering::Relaxed);
                if swap == old {
                    if old == 1 { // this was the last reference
                        is_free = true;
                    }
                    break
                }
            }
        }
        if is_free {
            self.free_list.push_back(index);
        }
    }

    /// Returns the number of live items. O(1) running time.
    pub fn live_count(&self) -> usize {
        self.tail.load(Ordering::Relaxed) - self.free_list.len()
    }
}

/// Internal Functions
impl <T> Pool<T> {
    // Returns an item from the free list, or
    // tries to allocate a new one from the buffer
    fn claim_free_index(&mut self) -> Result<usize, &'static str> {
        let index = match self.free_list.pop_front() {
            Some(i) => i,
            None => try!(self.push_back_alloc()),
        };
        self.retain(index);
        Ok(index)
    }

    // Internal alloc that does not create an Arc but still claims a slot
    fn internal_alloc(&mut self) -> Result<usize, &'static str> {
        let index = try!(self.claim_free_index());
        Ok(index)
    }

    // Pushes the end of the used space in the buffer back
    // returns the previous index
    fn push_back_alloc(&mut self) -> Result<usize, &'static str> {
        loop {
            let old_tail = self.tail.load(Ordering::Relaxed);
            let swap = self.tail.compare_and_swap(old_tail, old_tail+1, Ordering::Relaxed);
            // If we were the ones to claim this slot, or
            // we've overrun the buffer, return
            if old_tail >= self.capacity {
                return Err("OOM")
            } else if swap == old_tail {
                return Ok(old_tail)
            }
        }
    }

    fn header_for<'a>(&'a mut self, i: usize) -> &'a mut SlotHeader {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize);
            mem::transmute(ptr)
        }
    }
}

impl <T> Index<usize> for Pool<T> {
    type Output = T;

    fn index<'a>(&'a self, i: usize) -> &'a T {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize)
            .offset(self.header_size as isize);
            mem::transmute(ptr)
        }
    }
}

impl <T> IndexMut<usize> for Pool<T> {
    fn index_mut<'a>(&'a mut self, i: usize) -> &'a mut T {
        unsafe {
            let ptr = self.buffer.clone()
            .offset((i * self.slot_size) as isize)
            .offset(self.header_size as isize);
            mem::transmute(ptr)
        }
    }
}

#[test]
fn release_frees() {
       let mut buf: [u8; 100] = [0; 100];
       let mut p = Pool::<u32>::new(&mut buf[..]);

       // Use internal_alloc so that the Arc doesn't drop
       // the reference immediately
       assert!(p.internal_alloc().is_ok());
       assert!(p.internal_alloc().is_ok());

       assert_eq!(2, p.live_count());

       p.release(0);
       assert_eq!(1, p.live_count());
       assert_eq!(1, p.free_list.len());
       assert_eq!(0, *p.free_list.front().unwrap());

       p.release(1);
       assert_eq!(0, p.live_count());
       assert_eq!(2, p.free_list.len());
}

#[test]
fn alloc_after_free_recycles() {
       let mut buf: [u8; 100] = [0; 100];
       let mut p = Pool::<u32>::new(&mut buf[..]);
       assert!(p.internal_alloc().is_ok());
       assert_eq!(1, p.live_count());
       assert_eq!(1, p.tail.load(Ordering::Relaxed));

       p.release(0);
       assert_eq!(0, p.live_count());
       assert_eq!(1, p.free_list.len());

       assert!(p.internal_alloc().is_ok());
       assert_eq!(1, p.tail.load(Ordering::Relaxed)); // Tail shouldn't move
       assert_eq!(1, p.live_count());
       assert_eq!(0, p.free_list.len());
}

fn arc_clone() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        assert_eq!(1, int1.ref_count());
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        {
            let mut int1_c:Arc<u32> = int1.clone(); // Should bump the ref count
            assert_eq!(2, int1.ref_count());
            assert_eq!(2, int1_c.ref_count());
            assert_eq!(2, p.header_for(0).ref_count.load(Ordering::Relaxed));
        }
        // Now, the clone should have been dropped, but no memory reclaimed
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        assert_eq!(0, p.free_list.len());
    }
    // Now, int1 should have been dropped, and all memory reclaimed
    assert_eq!(0, p.header_for(0).ref_count.load(Ordering::Relaxed));
    assert_eq!(1, p.free_list.len());
}

#[test]
fn arc_drop() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        assert_eq!(1, int1.ref_count());
        assert_eq!(1, p.header_for(0).ref_count.load(Ordering::Relaxed));
        {
            let mut int2:Arc<u32> = p.alloc().unwrap();
            assert_eq!(1, int2.ref_count());
            assert_eq!(1, p.header_for(1).ref_count.load(Ordering::Relaxed));
        }
        // Now, int2 should have been dropped
        assert_eq!(0, p.header_for(1).ref_count.load(Ordering::Relaxed));
        assert_eq!(1, p.free_list.len());
    }
    // Now, int1 should have been dropped
    assert_eq!(0, p.header_for(0).ref_count.load(Ordering::Relaxed));
    assert_eq!(2, p.free_list.len());
}

#[test]
fn construction() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);

    assert_eq!(100, p.buffer_size);
    assert_eq!(mem::size_of::<usize>(), p.header_size);

    let expected_size = mem::size_of::<usize>() + mem::size_of::<u32>();
    assert_eq!(expected_size, p.slot_size);
    assert_eq!(100/expected_size, p.capacity); // expected_size should be 8+4=12
    assert_eq!(8, p.capacity);
}

#[test]
fn free_list_alloc_works() {
    let mut buf: [u8; 100] = [0; 100];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    {
        let mut int1:Arc<u32> = p.alloc().unwrap();
        *int1 = 42;
        // Check payload
        assert_eq!([42u8, 0u8, 0u8, 0u8][..], buf[8..12]);
        // Check ref_count
        assert_eq!([1u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8][..], buf[0..8]);
        assert_eq!(1, p.live_count());
    }
    // int1 is now out of scope, let's ensure the drop worked
    assert_eq!([0u8; 8][..], buf[0..8]);
}

#[test]
fn check_oom_error() {
    let mut buf: [u8; 1] = [0; 1];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    assert_eq!(Err("OOM"), p.alloc());
}

#[test]
fn multiple_allocations_work() {
    let mut buf: [u8; 120] = [0; 120];
    let mut p = Pool::<u32>::new(&mut buf[..]);
    for i in 0..10 {
        let mut int1 = p.alloc().unwrap();
        *int1 = i;
        unsafe { int1.retain() }; // Make sure this stays around long enough to read later
   }
   assert_eq!(10, p.live_count());
   let expected_ref_count = [1u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
   for i in 0..10 {
       let start = 12*i;
       // Check ref_count
       assert_eq!(expected_ref_count[..], buf[start..start+8]);
       // Check payload
       assert_eq!([i as u8, 0u8, 0u8, 0u8][..], buf[start+8..start+12]);
    }
}
<|endoftext|>


<commit_before>//! Utility methods and constructors for Lua classes

use std::default::Default;
use std::rc::Rc;
use rlua::{self, Lua, ToLua, Table, UserData, AnyUserData, Value, Function};
use super::object::Object;
use super::property::Property;

pub type Allocator = Rc<Fn(&Lua) -> rlua::Result<Object>>;
pub type Collector = Rc<Fn(Object)>;
pub type Checker = Rc<Fn(Object) -> bool>;

#[derive(Debug)]
pub struct Class<'lua> {
    table: Table<'lua>
}

#[derive(Clone)]
pub struct ClassState {
    // NOTE That this is missing fields from the C version.
    // They stored in the meta table instead, to not have unsafety.
    // They are fetchable using getters.
    name: String,
    // TODO Move stringly types to metatable only, not have fields and just have setters / getters?
    // FIXME This __needs__ to happen to ensure gc doesn't collect!
    // Have test that tests that hypo ^
    signals_global: Option<String>,
    parent: Option<String>,
    allocator: Option<Allocator>,
    collector: Option<Collector>,
    checker: Option<Checker>,
    instances: u32
}

pub struct ClassBuilder<'lua>{
    lua: &'lua Lua,
    class: Class<'lua>
}

impl <'lua> ClassBuilder<'lua> {
    pub fn method(self, name: String, meth: rlua::Function)
                  -> rlua::Result<Self> {
        let meta = self.class.table.get_metatable()
            .expect("Class had no meta table!");
        meta.set(name, meth)?;
        Ok(self)
    }

    pub fn property(self, prop: Property<'lua>) -> rlua::Result<Self> {
        let properties = self.class.table.get::<_, Table>("properties")?;
        let length = properties.len().unwrap_or(0) + 1;
        properties.set(length, prop)?;
        Ok(self)
    }

    pub fn save_class(mut self, name: &str)
                      -> rlua::Result<Self> {
        self.lua.globals().set(name, self.class.table)?;
        self.class.table = self.lua.globals().get(name)?;
        Ok(self)
    }

    pub fn build(self) -> rlua::Result<Class<'lua>> {
        Ok(self.class)
    }
}

impl <'lua> ToLua<'lua> for Class<'lua> {
    fn to_lua(self, lua: &'lua Lua) -> rlua::Result<Value<'lua>> {
        self.table.to_lua(lua)
    }
}

impl Default for ClassState {
    fn default() -> Self {
        ClassState {
            name: String::default(),
            signals_global: Option::default(),
            parent: Option::default(),
            allocator: Option::default(),
            collector: Option::default(),
            checker: Option::default(),
            instances: 0
        }
    }
}

impl UserData for ClassState {}

impl <'lua> Class<'lua> {
    pub fn builder(lua: &'lua Lua,
                   allocator: Option<Allocator>,
                   collector: Option<Collector>,
                   checker: Option<Checker>)
                   -> rlua::Result<ClassBuilder<'lua>> {
        let mut class = ClassState::default();
        class.allocator = allocator;
        class.collector = collector;
        class.checker = checker;
        let table = lua.create_table();
        // Store in not meta table so we can't index it
        table.set("data", class)?;
        table.set("properties", Vec::<Property>::new().to_lua(lua)?)?;
        let meta = lua.create_table();
        meta.set("signals", lua.create_table())?;
        meta.set("set_index_miss_handler",
                 lua.create_function(set_index_miss_handler).bind(table.clone())?)?;
        meta.set("set_newindex_miss_handler",
                 lua.create_function(set_newindex_miss_handler).bind(table.clone())?)?;
        meta.set("__index", meta.clone())?;
        table.set_metatable(Some(meta));
        Ok(ClassBuilder{
            lua: lua,
            class: Class { table }
        })
    }

    #[allow(dead_code)]
    pub fn properties(&self) -> rlua::Result<Table<'lua>> {
        self.table.get("properties")
    }

    pub fn checker(&self) -> rlua::Result<Option<Checker>> {
        self.table.get::<_, ClassState>("data")
            .map(|state| state.checker)
    }
}

impl <'lua> From<Table<'lua>> for Class<'lua> {
    fn from(table: Table<'lua>) -> Self {
        Class { table }
    }
}

fn set_index_miss_handler<'lua>(_: &'lua Lua, (class, func): (Table, Function))
                                -> rlua::Result<()> {
    let meta = class.get_metatable()
        .expect("Object had no metatable");
    meta.set("__index_miss_handler", func)?;
    Ok(())
}
fn set_newindex_miss_handler<'lua>(_: &'lua Lua, (class, func): (Table, Function))
                                   -> rlua::Result<()> {
    let meta = class.get_metatable()
        .expect("Object had no metatable");
    meta.set("__newindex_miss_handler", func)?;
    Ok(())
}

pub fn button_class(lua: &Lua) -> rlua::Result<Class> {
    let table = lua.globals().get::<_, Table>("button")
        .expect("Button class was not set! Did you call button::init?");
    assert!(table.get::<_, AnyUserData>("data")?.is::<ClassState>(),
            "This table was not a class!");
    Ok(Class { table })
}
<commit_msg>Moved all lua class attributes to properties/table<commit_after>//! Utility methods and constructors for Lua classes

use std::default::Default;
use std::rc::Rc;
use rlua::{self, Lua, ToLua, Table, UserData, AnyUserData, Value, Function};
use super::object::Object;
use super::property::Property;

pub type Allocator = Rc<Fn(&Lua) -> rlua::Result<Object>>;
pub type Collector = Rc<Fn(Object)>;
pub type Checker = Rc<Fn(Object) -> bool>;

#[derive(Debug)]
pub struct Class<'lua> {
    table: Table<'lua>
}

#[derive(Clone)]
pub struct ClassState {
    // NOTE That this is missing fields from the C version.
    // They stored in the meta table instead, to not have unsafety.
    // They are fetchable using getters.
    name: String,
    allocator: Option<Allocator>,
    collector: Option<Collector>,
    checker: Option<Checker>,
    instances: u32
}

pub struct ClassBuilder<'lua>{
    lua: &'lua Lua,
    class: Class<'lua>
}

impl <'lua> ClassBuilder<'lua> {
    pub fn method(self, name: String, meth: rlua::Function)
                  -> rlua::Result<Self> {
        let meta = self.class.table.get_metatable()
            .expect("Class had no meta table!");
        meta.set(name, meth)?;
        Ok(self)
    }

    pub fn property(self, prop: Property<'lua>) -> rlua::Result<Self> {
        let properties = self.class.table.get::<_, Table>("properties")?;
        let length = properties.len().unwrap_or(0) + 1;
        properties.set(length, prop)?;
        Ok(self)
    }

    pub fn save_class(mut self, name: &str)
                      -> rlua::Result<Self> {
        self.lua.globals().set(name, self.class.table)?;
        self.class.table = self.lua.globals().get(name)?;
        Ok(self)
    }

    pub fn build(self) -> rlua::Result<Class<'lua>> {
        Ok(self.class)
    }
}

impl <'lua> ToLua<'lua> for Class<'lua> {
    fn to_lua(self, lua: &'lua Lua) -> rlua::Result<Value<'lua>> {
        self.table.to_lua(lua)
    }
}

impl Default for ClassState {
    fn default() -> Self {
        ClassState {
            name: String::default(),
            allocator: Option::default(),
            collector: Option::default(),
            checker: Option::default(),
            instances: 0
        }
    }
}

impl UserData for ClassState {}

impl <'lua> Class<'lua> {
    pub fn builder(lua: &'lua Lua,
                   allocator: Option<Allocator>,
                   collector: Option<Collector>,
                   checker: Option<Checker>)
                   -> rlua::Result<ClassBuilder<'lua>> {
        let mut class = ClassState::default();
        class.allocator = allocator;
        class.collector = collector;
        class.checker = checker;
        let table = lua.create_table();
        // Store in not meta table so we can't index it
        table.set("data", class)?;
        table.set("properties", Vec::<Property>::new().to_lua(lua)?)?;
        let meta = lua.create_table();
        meta.set("signals", lua.create_table())?;
        meta.set("set_index_miss_handler",
                 lua.create_function(set_index_miss_handler).bind(table.clone())?)?;
        meta.set("set_newindex_miss_handler",
                 lua.create_function(set_newindex_miss_handler).bind(table.clone())?)?;
        meta.set("__index", meta.clone())?;
        table.set_metatable(Some(meta));
        Ok(ClassBuilder{
            lua: lua,
            class: Class { table }
        })
    }

    #[allow(dead_code)]
    pub fn properties(&self) -> rlua::Result<Table<'lua>> {
        self.table.get("properties")
    }

    #[allow(dead_code)]
    pub fn parent(&self) -> rlua::Result<Option<Class>> {
        use rlua::Value;
        match self.table.get::<_, Value>("parent")? {
            Value::Table(table) => {
                let data = table.get::<_, AnyUserData>("data")?;
                if !data.is::<ClassState>() {
                    Ok(None)
                } else {
                    Ok(Some(table.into()))
                }
            },
            _ => Ok(None)
        }
    }

    pub fn checker(&self) -> rlua::Result<Option<Checker>> {
        self.table.get::<_, ClassState>("data")
            .map(|state| state.checker)
    }
}

impl <'lua> From<Table<'lua>> for Class<'lua> {
    fn from(table: Table<'lua>) -> Self {
        Class { table }
    }
}

fn set_index_miss_handler<'lua>(_: &'lua Lua, (class, func): (Table, Function))
                                -> rlua::Result<()> {
    let meta = class.get_metatable()
        .expect("Object had no metatable");
    meta.set("__index_miss_handler", func)?;
    Ok(())
}
fn set_newindex_miss_handler<'lua>(_: &'lua Lua, (class, func): (Table, Function))
                                   -> rlua::Result<()> {
    let meta = class.get_metatable()
        .expect("Object had no metatable");
    meta.set("__newindex_miss_handler", func)?;
    Ok(())
}

pub fn button_class(lua: &Lua) -> rlua::Result<Class> {
    let table = lua.globals().get::<_, Table>("button")
        .expect("Button class was not set! Did you call button::init?");
    assert!(table.get::<_, AnyUserData>("data")?.is::<ClassState>(),
            "This table was not a class!");
    Ok(Class { table })
}
<|endoftext|>