<commit_before>//! Error types and conversion functions.

use std;
use std::error::Error;
use std::fmt;

use gfx;
use gfx_window_sdl;

use app_dirs2::AppDirsError;
use image;
use rodio::decoder::DecoderError;
use sdl2;
use toml;
use zip;

/// An enum containing all kinds of game framework errors.
#[derive(Debug)]
pub enum GameError {
    /// An error in the filesystem layout
    FilesystemError(String),
    /// An error in the config file
    ConfigError(String),
    /// An error in some part of the underlying SDL library.
    SdlError(String),
    /// An error saying that a an integer overflow/underflow occured
    /// in an underlying library.
    IntegerError(String),
    /// An error trying to parse a resource
    ResourceLoadError(String),
    /// Unable to find a resource; the Vec is the paths it searched for and associated errors
    ResourceNotFound(String, Vec<(std::path::PathBuf, GameError)>),
    /// Something went wrong in the renderer
    RenderError(String),
    /// Something went wrong in the audio playback
    AudioError(String),
    /// Something went wrong trying to create a window
    WindowError(gfx_window_sdl::InitError),
    /// Something went wrong trying to read from a file
    IOError(std::io::Error),
    /// Something went wrong trying to load/render a font
    FontError(String),
    /// Something went wrong applying video settings.
    VideoError(String),
    /// Something went compiling shaders
    ShaderProgramError(gfx::shade::ProgramError),
    /// Something else happened; this is generally a bug.
    UnknownError(String),
}

impl fmt::Display for GameError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            GameError::ResourceNotFound(ref s, ref paths) => write!(
                f,
                "Resource not found: {}, searched in paths {:?}",
                s, paths
            ),
            GameError::ConfigError(ref s) => write!(f, "Config error: {}", s),
            GameError::ResourceLoadError(ref s) => write!(f, "Error loading resource: {}", s),
            _ => write!(f, "GameError {:?}", self),
        }
    }
}

impl Error for GameError {
    fn description(&self) -> &str {
        match *self {
            GameError::FilesystemError(_) => "Filesystem error",
            GameError::ConfigError(_) => "Config file error",
            GameError::SdlError(_) => "SDL error",
            GameError::IntegerError(_) => "Integer error",
            GameError::ResourceLoadError(_) => "Resource load error",
            GameError::ResourceNotFound(_, _) => "Resource not found",
            GameError::RenderError(_) => "Render error",
            GameError::AudioError(_) => "Audio error",
            GameError::WindowError(_) => "Window error",
            GameError::IOError(_) => "IO error",
            GameError::FontError(_) => "Font error",
            GameError::VideoError(_) => "Video error",
            GameError::ShaderProgramError(_) => "Shader program error",
            GameError::UnknownError(_) => "Unknown error",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            GameError::ShaderProgramError(ref e) => Some(e),
            GameError::IOError(ref e) => Some(e),
            _ => None,
        }
    }
}

/// A convenient result type consisting of a return type and a `GameError`
pub type GameResult<T> = Result<T, GameError>;

impl From<String> for GameError {
    fn from(s: String) -> GameError {
        GameError::UnknownError(s)
    }
}

impl From<gfx_window_sdl::InitError> for GameError {
    fn from(s: gfx_window_sdl::InitError) -> GameError {
        GameError::WindowError(s)
    }
}

impl From<sdl2::IntegerOrSdlError> for GameError {
    fn from(e: sdl2::IntegerOrSdlError) -> GameError {
        match e {
            sdl2::IntegerOrSdlError::IntegerOverflows(s, i) => {
                let message = format!("Integer overflow: {}, str {}", i, s);
                GameError::IntegerError(message)
            }
            sdl2::IntegerOrSdlError::SdlError(s) => GameError::SdlError(s),
        }
    }
}

impl From<sdl2::filesystem::PrefPathError> for GameError {
    fn from(e: sdl2::filesystem::PrefPathError) -> GameError {
        let msg = match e {
            sdl2::filesystem::PrefPathError::InvalidOrganizationName(e) => {
                format!("Invalid organization name, {}", e)
            }
            sdl2::filesystem::PrefPathError::InvalidApplicationName(e) => {
                format!("Invalid application name, {}", e)
            }
            sdl2::filesystem::PrefPathError::SdlError(e) => e,
        };
        GameError::ConfigError(msg)
    }
}

impl From<sdl2::render::TextureValueError> for GameError {
    fn from(e: sdl2::render::TextureValueError) -> GameError {
        let msg = e.description();
        GameError::ResourceLoadError(msg.to_owned())
    }
}

impl From<AppDirsError> for GameError {
    fn from(e: AppDirsError) -> GameError {
        let errmessage = format!("{}", e);
        GameError::FilesystemError(errmessage)
    }
}
impl From<std::io::Error> for GameError {
    fn from(e: std::io::Error) -> GameError {
        GameError::IOError(e)
    }
}

impl From<toml::de::Error> for GameError {
    fn from(e: toml::de::Error) -> GameError {
        let errstr = format!("TOML decode error: {}", e.description());

        GameError::ConfigError(errstr)
    }
}

impl From<toml::ser::Error> for GameError {
    fn from(e: toml::ser::Error) -> GameError {
        let errstr = format!("TOML error (possibly encoding?): {}", e.description());
        GameError::ConfigError(errstr)
    }
}

impl From<zip::result::ZipError> for GameError {
    fn from(e: zip::result::ZipError) -> GameError {
        let errstr = format!("Zip error: {}", e.description());
        GameError::ResourceLoadError(errstr)
    }
}

impl From<DecoderError> for GameError {
    fn from(e: DecoderError) -> GameError {
        let errstr = format!("Audio decoder error: {:?}", e);
        GameError::AudioError(errstr)
    }
}

impl From<image::ImageError> for GameError {
    fn from(e: image::ImageError) -> GameError {
        let errstr = format!("Image load error: {}", e.description());
        GameError::ResourceLoadError(errstr)
    }
}

impl From<gfx::PipelineStateError<std::string::String>> for GameError {
    fn from(e: gfx::PipelineStateError<std::string::String>) -> GameError {
        let errstr = format!(
            "Error constructing pipeline!\nThis should probably not be \
             happening; it probably means an error in a shader or \
             something.\nError was: {:?}",
            e
        );
        GameError::VideoError(errstr)
    }
}

impl From<gfx::mapping::Error> for GameError {
    fn from(e: gfx::mapping::Error) -> GameError {
        let errstr = format!("Buffer mapping error: {:?}", e);
        GameError::VideoError(errstr)
    }
}

impl<S, D> From<gfx::CopyError<S, D>> for GameError
where
    S: fmt::Debug,
    D: fmt::Debug,
{
    fn from(e: gfx::CopyError<S, D>) -> GameError {
        let errstr = format!("Memory copy error: {:?}", e);
        GameError::VideoError(errstr)
    }
}

impl From<gfx::CombinedError> for GameError {
    fn from(e: gfx::CombinedError) -> GameError {
        let errstr = format!("Texture+view load error: {}", e.description());
        GameError::VideoError(errstr)
    }
}

impl From<gfx::texture::CreationError> for GameError {
    fn from(e: gfx::texture::CreationError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl From<gfx::ResourceViewError> for GameError {
    fn from(e: gfx::ResourceViewError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl From<gfx::TargetViewError> for GameError {
    fn from(e: gfx::TargetViewError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl<T> From<gfx::UpdateError<T>> for GameError
where
    T: fmt::Debug + fmt::Display + 'static,
{
    fn from(e: gfx::UpdateError<T>) -> GameError {
        let errstr = format!("Buffer update error: {}", e);
        GameError::VideoError(errstr)
    }
}

impl From<gfx::shade::ProgramError> for GameError {
    fn from(e: gfx::shade::ProgramError) -> GameError {
        GameError::ShaderProgramError(e)
    }
}
<commit_msg>Work on `GameError`: removed `sdl2`, handled `winit` and `glutin`.<commit_after>//! Error types and conversion functions.

use std;
use std::error::Error;
use std::fmt;

use gfx;
use glutin;
use winit;

use app_dirs2::AppDirsError;
use image;
use rodio::decoder::DecoderError;
use toml;
use zip;

/// An enum containing all kinds of game framework errors.
#[derive(Debug)]
pub enum GameError {
    /// An error in the filesystem layout
    FilesystemError(String),
    /// An error in the config file
    ConfigError(String),
    /// Happens when an `EventsLoopProxy` attempts to
    /// wake up an `EventsLoop` that no longer exists.
    EventLoopError(String),
    /// An error trying to parse a resource
    ResourceLoadError(String),
    /// Unable to find a resource; the Vec is the paths it searched for and associated errors
    ResourceNotFound(String, Vec<(std::path::PathBuf, GameError)>),
    /// Something went wrong in the renderer
    RenderError(String),
    /// Something went wrong in the audio playback
    AudioError(String),
    /// Something went wrong trying to create a window
    WindowError(glutin::CreationError),
    /// Something went wrong trying to read from a file
    IOError(std::io::Error),
    /// Something went wrong trying to load/render a font
    FontError(String),
    /// Something went wrong applying video settings.
    VideoError(String),
    /// Something went compiling shaders
    ShaderProgramError(gfx::shade::ProgramError),
    /// Something else happened; this is generally a bug.
    UnknownError(String),
}

impl fmt::Display for GameError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            GameError::ConfigError(ref s) => write!(f, "Config error: {}", s),
            GameError::ResourceLoadError(ref s) => write!(f, "Error loading resource: {}", s),
            GameError::ResourceNotFound(ref s, ref paths) => write!(
                f,
                "Resource not found: {}, searched in paths {:?}",
                s, paths
            ),
            GameError::WindowError(ref e) => write!(f, "Window creation error: {}", e),
            _ => write!(f, "GameError {:?}", self),
        }
    }
}

impl Error for GameError {
    fn description(&self) -> &str {
        match *self {
            GameError::FilesystemError(_) => "Filesystem error",
            GameError::ConfigError(_) => "Config file error",
            GameError::EventLoopError(_) => "Event loop error",
            GameError::ResourceLoadError(_) => "Resource load error",
            GameError::ResourceNotFound(_, _) => "Resource not found",
            GameError::RenderError(_) => "Render error",
            GameError::AudioError(_) => "Audio error",
            GameError::WindowError(_) => "Window error",
            GameError::IOError(_) => "IO error",
            GameError::FontError(_) => "Font error",
            GameError::VideoError(_) => "Video error",
            GameError::ShaderProgramError(_) => "Shader program error",
            GameError::UnknownError(_) => "Unknown error",
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            GameError::WindowError(ref e) => Some(e),
            GameError::IOError(ref e) => Some(e),
            GameError::ShaderProgramError(ref e) => Some(e),
            _ => None,
        }
    }
}

/// A convenient result type consisting of a return type and a `GameError`
pub type GameResult<T> = Result<T, GameError>;

impl From<String> for GameError {
    fn from(s: String) -> GameError {
        GameError::UnknownError(s)
    }
}

impl From<AppDirsError> for GameError {
    fn from(e: AppDirsError) -> GameError {
        let errmessage = format!("{}", e);
        GameError::FilesystemError(errmessage)
    }
}
impl From<std::io::Error> for GameError {
    fn from(e: std::io::Error) -> GameError {
        GameError::IOError(e)
    }
}

impl From<toml::de::Error> for GameError {
    fn from(e: toml::de::Error) -> GameError {
        let errstr = format!("TOML decode error: {}", e.description());

        GameError::ConfigError(errstr)
    }
}

impl From<toml::ser::Error> for GameError {
    fn from(e: toml::ser::Error) -> GameError {
        let errstr = format!("TOML error (possibly encoding?): {}", e.description());
        GameError::ConfigError(errstr)
    }
}

impl From<zip::result::ZipError> for GameError {
    fn from(e: zip::result::ZipError) -> GameError {
        let errstr = format!("Zip error: {}", e.description());
        GameError::ResourceLoadError(errstr)
    }
}

impl From<DecoderError> for GameError {
    fn from(e: DecoderError) -> GameError {
        let errstr = format!("Audio decoder error: {:?}", e);
        GameError::AudioError(errstr)
    }
}

impl From<image::ImageError> for GameError {
    fn from(e: image::ImageError) -> GameError {
        let errstr = format!("Image load error: {}", e.description());
        GameError::ResourceLoadError(errstr)
    }
}

impl From<gfx::PipelineStateError<std::string::String>> for GameError {
    fn from(e: gfx::PipelineStateError<std::string::String>) -> GameError {
        let errstr = format!(
            "Error constructing pipeline!\nThis should probably not be \
             happening; it probably means an error in a shader or \
             something.\nError was: {:?}",
            e
        );
        GameError::VideoError(errstr)
    }
}

impl From<gfx::mapping::Error> for GameError {
    fn from(e: gfx::mapping::Error) -> GameError {
        let errstr = format!("Buffer mapping error: {:?}", e);
        GameError::VideoError(errstr)
    }
}

impl<S, D> From<gfx::CopyError<S, D>> for GameError
where
    S: fmt::Debug,
    D: fmt::Debug,
{
    fn from(e: gfx::CopyError<S, D>) -> GameError {
        let errstr = format!("Memory copy error: {:?}", e);
        GameError::VideoError(errstr)
    }
}

impl From<gfx::CombinedError> for GameError {
    fn from(e: gfx::CombinedError) -> GameError {
        let errstr = format!("Texture+view load error: {}", e.description());
        GameError::VideoError(errstr)
    }
}

impl From<gfx::texture::CreationError> for GameError {
    fn from(e: gfx::texture::CreationError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl From<gfx::ResourceViewError> for GameError {
    fn from(e: gfx::ResourceViewError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl From<gfx::TargetViewError> for GameError {
    fn from(e: gfx::TargetViewError) -> GameError {
        gfx::CombinedError::from(e).into()
    }
}

impl<T> From<gfx::UpdateError<T>> for GameError
where
    T: fmt::Debug + fmt::Display + 'static,
{
    fn from(e: gfx::UpdateError<T>) -> GameError {
        let errstr = format!("Buffer update error: {}", e);
        GameError::VideoError(errstr)
    }
}

impl From<gfx::shade::ProgramError> for GameError {
    fn from(e: gfx::shade::ProgramError) -> GameError {
        GameError::ShaderProgramError(e)
    }
}

impl From<winit::EventsLoopClosed> for GameError {
    fn from(_: glutin::EventsLoopClosed) -> GameError {
        let e = "An event loop proxy attempted to wake up an event loop that no longer exists."
            .to_owned();
        GameError::EventLoopError(e)
    }
}

impl From<glutin::CreationError> for GameError {
    fn from(s: glutin::CreationError) -> GameError {
        GameError::WindowError(s)
    }
}
<|endoftext|>


<commit_before>/// Compute the real-valued digamma function.
///
/// The formula is as follows:
///
/// ```math
///        d ln(Γ(x))
/// ψ(x) = ----------
///            dx
/// ```
///
/// where Γ is the gamma function. The computation is based on an approximation
/// as described in the reference below.
///
/// ## Examples
///
/// ```
/// use special::digamma;
///
/// const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
/// assert!((digamma(1.0) + EULER_MASCHERONI).abs() < 1e-15);
/// ```
///
/// ## References
///
/// 1. M. J. Beal, Variational algorithms for approximate Bayesian inference.
///    University of London, 2003, pp. 265–266.
pub fn digamma(x: f64)-> f64 {
    macro_rules! evaluate_polynomial(
        ($x:expr, $coefficients:expr) => (
            $coefficients.iter().rev().fold(0.0, |sum, &c| $x * sum + c)
        );
    );

    if x <= 8.0 {
        return digamma(x + 1.0) - x.recip();
    }

    let inv_x = x.recip();
    let inv_x_2 = inv_x * inv_x;
    x.ln() - 0.5 * inv_x - inv_x_2 * evaluate_polynomial!(inv_x_2, [
        1.0 / 12.0, -1.0 / 120.0, 1.0 / 252.0, -1.0 / 240.0,
        5.0 / 660.0, -691.0 / 32760.0, 1.0 / 12.0, -3617.0 / 8160.0,
    ])
}

/// Compute the gamma function.
#[inline]
pub fn gamma(x: f64) -> f64 {
    use m::tgamma;
    unsafe { tgamma(x) }
}

/// Compute the regularized lower incomplete gamma function.
///
/// The formula is as follows:
///
/// ```math
///           γ(x, p)    1   x
/// P(x, p) = ------- = ---- ∫ t^(p-1) e^(-t) dt
///            Γ(p)     Γ(p) 0
/// ```
///
/// where γ is the incomplete lower gamma function, and Γ is the complete gamma
/// function.
///
/// The code is based on a [C implementation][1] by John Burkardt. The original
/// algorithm was published in Applied Statistics and is known as
/// [Algorithm AS 239][2].
///
/// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa239/asa239.html
/// [2]: http://www.jstor.org/stable/2347328
pub fn inc_gamma(x: f64, p: f64) -> f64 {
    debug_assert!(x >= 0.0 && p > 0.0);

    const ELIMIT: f64 = -88.0;
    const OFLO: f64 = 1.0e+37;
    const TOL: f64 = 1.0e-14;
    const XBIG: f64 = 1.0e+08;

    if x == 0.0 {
        return 0.0;
    }

    // For `p ≥ 1000`, the original algorithm uses an approximation shown below.
    // However, it introduces a substantial accuracy loss.
    //
    // ```
    // use m::erf;
    // use std::f64::consts::FRAC_1_SQRT_2;
    //
    // const PLIMIT: f64 = 1000.0;
    //
    // if PLIMIT < p {
    //     let pn1 = 3.0 * p.sqrt() * ((x / p).powf(1.0 / 3.0) + 1.0 / (9.0 * p) - 1.0);
    //     return 0.5 * (1.0 + unsafe { erf(FRAC_1_SQRT_2 * pn1) });
    // }
    // ```

    if XBIG < x {
        return 1.0;
    }

    if x <= 1.0 || x < p {
        let mut arg = p * x.ln() - x - ln_gamma(p + 1.0).0;
        let mut c = 1.0;
        let mut value = 1.0;
        let mut a = p;

        loop {
            a += 1.0;
            c *= x / a;
            value += c;

            if c <= TOL {
                break;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return arg.exp();
        } else {
            return 0.0;
        };
    } else {
        let mut arg = p * x.ln() - x - ln_gamma(p).0;
        let mut a = 1.0 - p;
        let mut b = a + x + 1.0;
        let mut c = 0.0;
        let mut pn1 = 1.0;
        let mut pn2 = x;
        let mut pn3 = x + 1.0;
        let mut pn4 = x * b;
        let mut value = pn3 / pn4;

        loop {
            a += 1.0;
            b += 2.0;
            c += 1.0;
            let an = a * c;
            let pn5 = b * pn3 - an * pn1;
            let pn6 = b * pn4 - an * pn2;

            if pn6 != 0.0 {
                let rn = pn5 / pn6;
                if (value - rn).abs() <= TOL.min(TOL * rn) {
                    break;
                }
                value = rn;
            }

            pn1 = pn3;
            pn2 = pn4;
            pn3 = pn5;
            pn4 = pn6;

            if OFLO <= pn5.abs() {
                pn1 /= OFLO;
                pn2 /= OFLO;
                pn3 /= OFLO;
                pn4 /= OFLO;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return 1.0 - arg.exp();
        } else {
            return 1.0;
        }
    }
}

/// Compute the natural logarithm of the gamma function.
#[inline]
pub fn ln_gamma(x: f64) -> (f64, i32) {
    use m::lgamma_r;
    let mut sign: i32 = 0;
    let value = unsafe { lgamma_r(x, &mut sign) };
    (value, sign)
}

#[cfg(test)]
mod tests {
    use assert;

    #[test]
    fn digamma() {
        use std::f64::consts::{FRAC_PI_2, LN_2};
        const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
        assert_eq!(-FRAC_PI_2 - 3.0 * LN_2 - EULER_MASCHERONI, super::digamma(0.25));
    }

    #[test]
    fn gamma() {
        use std::f64::INFINITY;
        let x = vec![0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0];
        let y = vec![
                         INFINITY, 1.772453850905516e+00, 1.000000000000000e+00,
            8.862269254527581e-01, 1.000000000000000e+00, 1.329340388179137e+00,
            2.000000000000000e+00, 3.323350970447843e+00, 6.000000000000000e+00,
            1.163172839656745e+01, 2.400000000000000e+01, 5.234277778455353e+01,
            1.200000000000000e+02, 2.878852778150444e+02, 7.200000000000000e+02,
        ];

        let z = x.iter().map(|&x| super::gamma(x)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-13);
    }

    #[test]
    fn inc_gamma_small_p() {
        let p = 4.2;
        let x = vec![
            0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5,
            5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5,
        ];
        let y = vec![
            0.000000000000000e+00, 1.118818097571853e-03, 1.386936454406691e-02,
            5.173931341260211e-02, 1.186345016135447e-01, 2.092219100370761e-01,
            3.137074857845146e-01, 4.219736975484903e-01, 5.258832878449179e-01,
            6.200417684429613e-01, 7.016346974530381e-01, 7.698559854309833e-01,
            8.252531208548146e-01, 8.691540627528154e-01, 9.032348314835232e-01,
            9.292289044193606e-01, 9.487539837941871e-01, 9.632249236076943e-01,
            9.738240752336722e-01, 9.815062931491536e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-14);
    }

    #[test]
    fn inc_gamma_large_p() {
        let p = 1500.0;
        let x = vec![
            1400.0, 1410.0, 1420.0, 1430.0, 1440.0, 1450.0, 1460.0, 1470.0, 1480.0, 1490.0,
            1500.0, 1510.0, 1520.0, 1530.0, 1540.0, 1550.0, 1560.0, 1570.0, 1580.0, 1590.0,
        ];
        let y = vec![
            4.231080348517120e-03, 9.056183893278287e-03, 1.809200095269094e-02,
            3.380047876608895e-02, 5.917825204288418e-02, 9.731671903009434e-02,
            1.506860564044825e-01, 2.202940362119810e-01, 3.049926190158564e-01,
            4.012305141433336e-01, 5.034335611603484e-01, 6.049687028759408e-01,
            6.994146092246781e-01, 7.817404371226420e-01, 8.490443067241588e-01,
            9.006922020203059e-01, 9.379249194843459e-01, 9.631597573132408e-01,
            9.792520392189441e-01, 9.889149370075309e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-12);
    }
}
<commit_msg>Decrease the precision when testing gamma<commit_after>/// Compute the real-valued digamma function.
///
/// The formula is as follows:
///
/// ```math
///        d ln(Γ(x))
/// ψ(x) = ----------
///            dx
/// ```
///
/// where Γ is the gamma function. The computation is based on an approximation
/// as described in the reference below.
///
/// ## Examples
///
/// ```
/// use special::digamma;
///
/// const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
/// assert!((digamma(1.0) + EULER_MASCHERONI).abs() < 1e-15);
/// ```
///
/// ## References
///
/// 1. M. J. Beal, Variational algorithms for approximate Bayesian inference.
///    University of London, 2003, pp. 265–266.
pub fn digamma(x: f64)-> f64 {
    macro_rules! evaluate_polynomial(
        ($x:expr, $coefficients:expr) => (
            $coefficients.iter().rev().fold(0.0, |sum, &c| $x * sum + c)
        );
    );

    if x <= 8.0 {
        return digamma(x + 1.0) - x.recip();
    }

    let inv_x = x.recip();
    let inv_x_2 = inv_x * inv_x;
    x.ln() - 0.5 * inv_x - inv_x_2 * evaluate_polynomial!(inv_x_2, [
        1.0 / 12.0, -1.0 / 120.0, 1.0 / 252.0, -1.0 / 240.0,
        5.0 / 660.0, -691.0 / 32760.0, 1.0 / 12.0, -3617.0 / 8160.0,
    ])
}

/// Compute the gamma function.
#[inline]
pub fn gamma(x: f64) -> f64 {
    use m::tgamma;
    unsafe { tgamma(x) }
}

/// Compute the regularized lower incomplete gamma function.
///
/// The formula is as follows:
///
/// ```math
///           γ(x, p)    1   x
/// P(x, p) = ------- = ---- ∫ t^(p-1) e^(-t) dt
///            Γ(p)     Γ(p) 0
/// ```
///
/// where γ is the incomplete lower gamma function, and Γ is the complete gamma
/// function.
///
/// The code is based on a [C implementation][1] by John Burkardt. The original
/// algorithm was published in Applied Statistics and is known as
/// [Algorithm AS 239][2].
///
/// [1]: http://people.sc.fsu.edu/~jburkardt/c_src/asa239/asa239.html
/// [2]: http://www.jstor.org/stable/2347328
pub fn inc_gamma(x: f64, p: f64) -> f64 {
    debug_assert!(x >= 0.0 && p > 0.0);

    const ELIMIT: f64 = -88.0;
    const OFLO: f64 = 1.0e+37;
    const TOL: f64 = 1.0e-14;
    const XBIG: f64 = 1.0e+08;

    if x == 0.0 {
        return 0.0;
    }

    // For `p ≥ 1000`, the original algorithm uses an approximation shown below.
    // However, it introduces a substantial accuracy loss.
    //
    // ```
    // use m::erf;
    // use std::f64::consts::FRAC_1_SQRT_2;
    //
    // const PLIMIT: f64 = 1000.0;
    //
    // if PLIMIT < p {
    //     let pn1 = 3.0 * p.sqrt() * ((x / p).powf(1.0 / 3.0) + 1.0 / (9.0 * p) - 1.0);
    //     return 0.5 * (1.0 + unsafe { erf(FRAC_1_SQRT_2 * pn1) });
    // }
    // ```

    if XBIG < x {
        return 1.0;
    }

    if x <= 1.0 || x < p {
        let mut arg = p * x.ln() - x - ln_gamma(p + 1.0).0;
        let mut c = 1.0;
        let mut value = 1.0;
        let mut a = p;

        loop {
            a += 1.0;
            c *= x / a;
            value += c;

            if c <= TOL {
                break;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return arg.exp();
        } else {
            return 0.0;
        };
    } else {
        let mut arg = p * x.ln() - x - ln_gamma(p).0;
        let mut a = 1.0 - p;
        let mut b = a + x + 1.0;
        let mut c = 0.0;
        let mut pn1 = 1.0;
        let mut pn2 = x;
        let mut pn3 = x + 1.0;
        let mut pn4 = x * b;
        let mut value = pn3 / pn4;

        loop {
            a += 1.0;
            b += 2.0;
            c += 1.0;
            let an = a * c;
            let pn5 = b * pn3 - an * pn1;
            let pn6 = b * pn4 - an * pn2;

            if pn6 != 0.0 {
                let rn = pn5 / pn6;
                if (value - rn).abs() <= TOL.min(TOL * rn) {
                    break;
                }
                value = rn;
            }

            pn1 = pn3;
            pn2 = pn4;
            pn3 = pn5;
            pn4 = pn6;

            if OFLO <= pn5.abs() {
                pn1 /= OFLO;
                pn2 /= OFLO;
                pn3 /= OFLO;
                pn4 /= OFLO;
            }
        }

        arg += value.ln();

        if ELIMIT <= arg {
            return 1.0 - arg.exp();
        } else {
            return 1.0;
        }
    }
}

/// Compute the natural logarithm of the gamma function.
#[inline]
pub fn ln_gamma(x: f64) -> (f64, i32) {
    use m::lgamma_r;
    let mut sign: i32 = 0;
    let value = unsafe { lgamma_r(x, &mut sign) };
    (value, sign)
}

#[cfg(test)]
mod tests {
    use assert;

    #[test]
    fn digamma() {
        use std::f64::consts::{FRAC_PI_2, LN_2};
        const EULER_MASCHERONI: f64 = 0.57721566490153286060651209008240243104215933593992;
        assert_eq!(-FRAC_PI_2 - 3.0 * LN_2 - EULER_MASCHERONI, super::digamma(0.25));
    }

    #[test]
    fn gamma() {
        use std::f64::INFINITY;
        let x = vec![0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0];
        let y = vec![
                         INFINITY, 1.772453850905516e+00, 1.000000000000000e+00,
            8.862269254527581e-01, 1.000000000000000e+00, 1.329340388179137e+00,
            2.000000000000000e+00, 3.323350970447843e+00, 6.000000000000000e+00,
            1.163172839656745e+01, 2.400000000000000e+01, 5.234277778455353e+01,
            1.200000000000000e+02, 2.878852778150444e+02, 7.200000000000000e+02,
        ];

        let z = x.iter().map(|&x| super::gamma(x)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-12);
    }

    #[test]
    fn inc_gamma_small_p() {
        let p = 4.2;
        let x = vec![
            0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5,
            5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5,
        ];
        let y = vec![
            0.000000000000000e+00, 1.118818097571853e-03, 1.386936454406691e-02,
            5.173931341260211e-02, 1.186345016135447e-01, 2.092219100370761e-01,
            3.137074857845146e-01, 4.219736975484903e-01, 5.258832878449179e-01,
            6.200417684429613e-01, 7.016346974530381e-01, 7.698559854309833e-01,
            8.252531208548146e-01, 8.691540627528154e-01, 9.032348314835232e-01,
            9.292289044193606e-01, 9.487539837941871e-01, 9.632249236076943e-01,
            9.738240752336722e-01, 9.815062931491536e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-14);
    }

    #[test]
    fn inc_gamma_large_p() {
        let p = 1500.0;
        let x = vec![
            1400.0, 1410.0, 1420.0, 1430.0, 1440.0, 1450.0, 1460.0, 1470.0, 1480.0, 1490.0,
            1500.0, 1510.0, 1520.0, 1530.0, 1540.0, 1550.0, 1560.0, 1570.0, 1580.0, 1590.0,
        ];
        let y = vec![
            4.231080348517120e-03, 9.056183893278287e-03, 1.809200095269094e-02,
            3.380047876608895e-02, 5.917825204288418e-02, 9.731671903009434e-02,
            1.506860564044825e-01, 2.202940362119810e-01, 3.049926190158564e-01,
            4.012305141433336e-01, 5.034335611603484e-01, 6.049687028759408e-01,
            6.994146092246781e-01, 7.817404371226420e-01, 8.490443067241588e-01,
            9.006922020203059e-01, 9.379249194843459e-01, 9.631597573132408e-01,
            9.792520392189441e-01, 9.889149370075309e-01,
        ];

        let z = x.iter().map(|&x| super::inc_gamma(x, p)).collect::<Vec<_>>();
        assert::close(&z, &y, 1e-12);
    }
}
<|endoftext|>


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

use dom::bindings::codegen::Bindings::DissimilarOriginWindowBinding;
use dom::bindings::codegen::Bindings::DissimilarOriginWindowBinding::DissimilarOriginWindowMethods;
use dom::bindings::error::{Error, ErrorResult};
use dom::bindings::inheritance::Castable;
use dom::bindings::js::{JS, MutNullableJS, Root};
use dom::bindings::str::DOMString;
use dom::bindings::structuredclone::StructuredCloneData;
use dom::dissimilaroriginlocation::DissimilarOriginLocation;
use dom::globalscope::GlobalScope;
use dom::windowproxy::WindowProxy;
use dom_struct::dom_struct;
use ipc_channel::ipc;
use js::jsapi::{JSContext, HandleValue};
use js::jsval::{JSVal, UndefinedValue};
use msg::constellation_msg::PipelineId;
use script_traits::ScriptMsg;
use servo_url::ImmutableOrigin;
use servo_url::MutableOrigin;
use servo_url::ServoUrl;

/// Represents a dissimilar-origin `Window` that exists in another script thread.
///
/// Since the `Window` is in a different script thread, we cannot access it
/// directly, but some of its accessors (for example `window.parent`)
/// still need to function.
///
/// In `windowproxy.rs`, we create a custom window proxy for these windows,
/// that throws security exceptions for most accessors. This is not a replacement
/// for XOWs, but provides belt-and-braces security.
#[dom_struct]
pub struct DissimilarOriginWindow {
    /// The global for this window.
    globalscope: GlobalScope,

    /// The window proxy for this window.
    window_proxy: JS<WindowProxy>,

    /// The location of this window, initialized lazily.
    location: MutNullableJS<DissimilarOriginLocation>,
}

impl DissimilarOriginWindow {
    #[allow(unsafe_code)]
    pub fn new(global_to_clone_from: &GlobalScope, window_proxy: &WindowProxy) -> Root<DissimilarOriginWindow> {
        let cx = global_to_clone_from.get_cx();
        // Any timer events fired on this window are ignored.
        let (timer_event_chan, _) = ipc::channel().unwrap();
        let win = box DissimilarOriginWindow {
            globalscope: GlobalScope::new_inherited(PipelineId::new(),
                                                    global_to_clone_from.devtools_chan().cloned(),
                                                    global_to_clone_from.mem_profiler_chan().clone(),
                                                    global_to_clone_from.time_profiler_chan().clone(),
                                                    global_to_clone_from.script_to_constellation_chan().clone(),
                                                    global_to_clone_from.scheduler_chan().clone(),
                                                    global_to_clone_from.resource_threads().clone(),
                                                    timer_event_chan,
                                                    global_to_clone_from.origin().clone()),
            window_proxy: JS::from_ref(window_proxy),
            location: MutNullableJS::new(None),
        };
        unsafe { DissimilarOriginWindowBinding::Wrap(cx, win) }
    }

    #[allow(dead_code)]
    pub fn origin(&self) -> &MutableOrigin {
        self.globalscope.origin()
    }
}

impl DissimilarOriginWindowMethods for DissimilarOriginWindow {
    // https://html.spec.whatwg.org/multipage/#dom-window
    fn Window(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-self
    fn Self_(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-frames
    fn Frames(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-parent
    fn GetParent(&self) -> Option<Root<WindowProxy>> {
        // Steps 1-3.
        if self.window_proxy.is_browsing_context_discarded() {
            return None;
        }
        // Step 4.
        if let Some(parent) = self.window_proxy.parent() {
            return Some(Root::from_ref(parent));
        }
        // Step 5.
        Some(Root::from_ref(&*self.window_proxy))
    }

    // https://html.spec.whatwg.org/multipage/#dom-top
    fn GetTop(&self) -> Option<Root<WindowProxy>> {
        // Steps 1-3.
        if self.window_proxy.is_browsing_context_discarded() {
            return None;
        }
        // Steps 4-5.
        Some(Root::from_ref(self.window_proxy.top()))
    }

    // https://html.spec.whatwg.org/multipage/#dom-length
    fn Length(&self) -> u32 {
        // TODO: Implement x-origin length
        0
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-close
    fn Close(&self) {
        // TODO: Implement x-origin close
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-closed
    fn Closed(&self) -> bool {
        // TODO: Implement x-origin close
        false
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-window-postmessage
    unsafe fn PostMessage(&self, cx: *mut JSContext, message: HandleValue, origin: DOMString) -> ErrorResult {
        // Step 3-5.
        let origin = match &origin[..] {
            "*" => None,
            "/" => {
                // TODO: Should be the origin of the incumbent settings object.
                None
            },
            url => match ServoUrl::parse(&url) {
                Ok(url) => Some(url.origin()),
                Err(_) => return Err(Error::Syntax),
            }
        };

        // Step 1-2, 6-8.
        // TODO(#12717): Should implement the `transfer` argument.
        let data = StructuredCloneData::write(cx, message)?;

        // Step 9.
        self.post_message(origin, data);
        Ok(())
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-opener
    unsafe fn Opener(&self, _: *mut JSContext) -> JSVal {
        // TODO: Implement x-origin opener
        UndefinedValue()
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-opener
    unsafe fn SetOpener(&self, _: *mut JSContext, _: HandleValue) {
        // TODO: Implement x-origin opener
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-blur
    fn Blur(&self) {
        // TODO: Implement x-origin blur
    }

    // https://html.spec.whatwg.org/multipage/#dom-focus
    fn Focus(&self) {
        // TODO: Implement x-origin focus
    }

    // https://html.spec.whatwg.org/multipage/#dom-location
    fn Location(&self) -> Root<DissimilarOriginLocation> {
        self.location.or_init(|| DissimilarOriginLocation::new(self))
    }
}

impl DissimilarOriginWindow {
    pub fn post_message(&self, origin: Option<ImmutableOrigin>, data: StructuredCloneData) {
        let msg = ScriptMsg::PostMessage(self.window_proxy.browsing_context_id(),
                                                origin,
                                                data.move_to_arraybuffer());
        let _ = self.upcast::<GlobalScope>().script_to_constellation_chan().send(msg);
    }
}
<commit_msg>Reformat DissimilarOriginWindow::new<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use dom::bindings::codegen::Bindings::DissimilarOriginWindowBinding;
use dom::bindings::codegen::Bindings::DissimilarOriginWindowBinding::DissimilarOriginWindowMethods;
use dom::bindings::error::{Error, ErrorResult};
use dom::bindings::inheritance::Castable;
use dom::bindings::js::{JS, MutNullableJS, Root};
use dom::bindings::str::DOMString;
use dom::bindings::structuredclone::StructuredCloneData;
use dom::dissimilaroriginlocation::DissimilarOriginLocation;
use dom::globalscope::GlobalScope;
use dom::windowproxy::WindowProxy;
use dom_struct::dom_struct;
use ipc_channel::ipc;
use js::jsapi::{JSContext, HandleValue};
use js::jsval::{JSVal, UndefinedValue};
use msg::constellation_msg::PipelineId;
use script_traits::ScriptMsg;
use servo_url::ImmutableOrigin;
use servo_url::MutableOrigin;
use servo_url::ServoUrl;

/// Represents a dissimilar-origin `Window` that exists in another script thread.
///
/// Since the `Window` is in a different script thread, we cannot access it
/// directly, but some of its accessors (for example `window.parent`)
/// still need to function.
///
/// In `windowproxy.rs`, we create a custom window proxy for these windows,
/// that throws security exceptions for most accessors. This is not a replacement
/// for XOWs, but provides belt-and-braces security.
#[dom_struct]
pub struct DissimilarOriginWindow {
    /// The global for this window.
    globalscope: GlobalScope,

    /// The window proxy for this window.
    window_proxy: JS<WindowProxy>,

    /// The location of this window, initialized lazily.
    location: MutNullableJS<DissimilarOriginLocation>,
}

impl DissimilarOriginWindow {
    #[allow(unsafe_code)]
    pub fn new(
        global_to_clone_from: &GlobalScope,
        window_proxy: &WindowProxy,
    ) -> Root<Self> {
        let cx = global_to_clone_from.get_cx();
        // Any timer events fired on this window are ignored.
        let (timer_event_chan, _) = ipc::channel().unwrap();
        let win = box Self {
            globalscope: GlobalScope::new_inherited(
                PipelineId::new(),
                global_to_clone_from.devtools_chan().cloned(),
                global_to_clone_from.mem_profiler_chan().clone(),
                global_to_clone_from.time_profiler_chan().clone(),
                global_to_clone_from.script_to_constellation_chan().clone(),
                global_to_clone_from.scheduler_chan().clone(),
                global_to_clone_from.resource_threads().clone(),
                timer_event_chan,
                global_to_clone_from.origin().clone(),
            ),
            window_proxy: JS::from_ref(window_proxy),
            location: Default::default(),
        };
        unsafe { DissimilarOriginWindowBinding::Wrap(cx, win) }
    }

    #[allow(dead_code)]
    pub fn origin(&self) -> &MutableOrigin {
        self.globalscope.origin()
    }
}

impl DissimilarOriginWindowMethods for DissimilarOriginWindow {
    // https://html.spec.whatwg.org/multipage/#dom-window
    fn Window(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-self
    fn Self_(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-frames
    fn Frames(&self) -> Root<WindowProxy> {
        Root::from_ref(&*self.window_proxy)
    }

    // https://html.spec.whatwg.org/multipage/#dom-parent
    fn GetParent(&self) -> Option<Root<WindowProxy>> {
        // Steps 1-3.
        if self.window_proxy.is_browsing_context_discarded() {
            return None;
        }
        // Step 4.
        if let Some(parent) = self.window_proxy.parent() {
            return Some(Root::from_ref(parent));
        }
        // Step 5.
        Some(Root::from_ref(&*self.window_proxy))
    }

    // https://html.spec.whatwg.org/multipage/#dom-top
    fn GetTop(&self) -> Option<Root<WindowProxy>> {
        // Steps 1-3.
        if self.window_proxy.is_browsing_context_discarded() {
            return None;
        }
        // Steps 4-5.
        Some(Root::from_ref(self.window_proxy.top()))
    }

    // https://html.spec.whatwg.org/multipage/#dom-length
    fn Length(&self) -> u32 {
        // TODO: Implement x-origin length
        0
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-close
    fn Close(&self) {
        // TODO: Implement x-origin close
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-closed
    fn Closed(&self) -> bool {
        // TODO: Implement x-origin close
        false
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-window-postmessage
    unsafe fn PostMessage(&self, cx: *mut JSContext, message: HandleValue, origin: DOMString) -> ErrorResult {
        // Step 3-5.
        let origin = match &origin[..] {
            "*" => None,
            "/" => {
                // TODO: Should be the origin of the incumbent settings object.
                None
            },
            url => match ServoUrl::parse(&url) {
                Ok(url) => Some(url.origin()),
                Err(_) => return Err(Error::Syntax),
            }
        };

        // Step 1-2, 6-8.
        // TODO(#12717): Should implement the `transfer` argument.
        let data = StructuredCloneData::write(cx, message)?;

        // Step 9.
        self.post_message(origin, data);
        Ok(())
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-opener
    unsafe fn Opener(&self, _: *mut JSContext) -> JSVal {
        // TODO: Implement x-origin opener
        UndefinedValue()
    }

    #[allow(unsafe_code)]
    // https://html.spec.whatwg.org/multipage/#dom-opener
    unsafe fn SetOpener(&self, _: *mut JSContext, _: HandleValue) {
        // TODO: Implement x-origin opener
    }

    // https://html.spec.whatwg.org/multipage/#dom-window-blur
    fn Blur(&self) {
        // TODO: Implement x-origin blur
    }

    // https://html.spec.whatwg.org/multipage/#dom-focus
    fn Focus(&self) {
        // TODO: Implement x-origin focus
    }

    // https://html.spec.whatwg.org/multipage/#dom-location
    fn Location(&self) -> Root<DissimilarOriginLocation> {
        self.location.or_init(|| DissimilarOriginLocation::new(self))
    }
}

impl DissimilarOriginWindow {
    pub fn post_message(&self, origin: Option<ImmutableOrigin>, data: StructuredCloneData) {
        let msg = ScriptMsg::PostMessage(self.window_proxy.browsing_context_id(),
                                                origin,
                                                data.move_to_arraybuffer());
        let _ = self.upcast::<GlobalScope>().script_to_constellation_chan().send(msg);
    }
}
<|endoftext|>


<commit_before>//! Abstraction over blocking and unblocking the current thread.
//!
//! Provides an abstraction over blocking the current thread. This is similar to
//! the park / unpark constructs provided by [`std`] but made generic. This
//! allows embedding custom functionality to perform when the thread is blocked.
//!
//! A blocked [`Park`][p] instance is unblocked by calling [`unpark`] on its
//! [`Unpark`][up] handle.
//!
//! The [`ParkThread`] struct implements [`Park`][p] using
//! [`thread::park`][`std`] to put the thread to sleep. The Tokio reactor also
//! implements park, but uses [`mio::Poll`][mio] to block the thread instead.
//!
//! The [`Park`][p] trait is composable. A timer implementation might decorate a
//! [`Park`][p] implementation by checking if any timeouts have elapsed after
//! the inner [`Park`][p] implementation unblocks.
//!
//! # Model
//!
//! Conceptually, each [`Park`][p] instance has an associated token, which is
//! initially not present:
//!
//! * The [`park`] method blocks the current thread unless or until the token
//!   is available, at which point it atomically consumes the token.
//! * The [`unpark`] method atomically makes the token available if it wasn't
//!   already.
//!
//! Some things to note:
//!
//! * If [`unpark`] is called before [`park`], the next call to [`park`] will
//! **not** block the thread.
//! * **Spurious** wakeups are permitted, i.e., the [`park`] method may unblock
//!   even if [`unpark`] was not called.
//! * [`park_timeout`] does the same as [`park`] but allows specifying a maximum
//!   time to block the thread for.
//!
//! [`std`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`thread::park`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`ParkThread`]: struct.ParkThread.html
//! [p]: trait.Park.html
//! [`park`]: trait.Park.html#tymethod.park
//! [`park_timeout`]: trait.Park.html#tymethod.park_timeout
//! [`unpark`]: trait.Unpark.html#tymethod.unpark
//! [up]: trait.Unpark.html
//! [mio]: https://docs.rs/mio/0.6/mio/struct.Poll.html

use std::marker::PhantomData;
use std::rc::Rc;
use std::sync::{Arc, Mutex, Condvar};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

/// Block the current thread.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
pub trait Park {
    /// Unpark handle type for the `Park` implementation.
    type Unpark: Unpark;

    /// Error returned by `park`
    type Error;

    /// Get a new `Unpark` handle associated with this `Park` instance.
    fn unpark(&self) -> Self::Unpark;

    /// Block the current thread unless or until the token is available.
    ///
    /// A call to `park` does not guarantee that the thread will remain blocked
    /// forever, and callers should be prepared for this possibility. This
    /// function may wakeup spuriously for any reason.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Park` implementation
    ///
    /// [mod]: ../index.html
    fn park(&mut self) -> Result<(), Self::Error>;

    /// Park the current thread for at most `duration`.
    ///
    /// This function is the same as `park` but allows specifying a maximum time
    /// to block the thread for.
    ///
    /// Same as `park`, there is no guarantee that the thread will remain
    /// blocked for any amount of time. Spurious wakeups are permitted for any
    /// reason.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Park` implementation
    ///
    /// [mod]: ../index.html
    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error>;
}

/// Unblock a thread blocked by the associated [`Park`] instance.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
/// [`Park`]: trait.Park.html
pub trait Unpark: Sync + Send + 'static {
    /// Unblock a thread that is blocked by the associated `Park` handle.
    ///
    /// Calling `unpark` atomically makes available the unpark token, if it is
    /// not already available.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Unpark` implementation
    ///
    /// [mod]: ../index.html
    fn unpark(&self);
}

impl Unpark for Box<Unpark> {
    fn unpark(&self) {
        (**self).unpark()
    }
}

/// Blocks the current thread using a condition variable.
///
/// Implements the [`Park`] functionality by using a condition variable. An
/// atomic variable is also used to avoid using the condition variable if
/// possible.
///
/// The condition variable is cached in a thread-local variable and is shared
/// across all `ParkThread` instances created on the same thread. This also
/// means that an instance of `ParkThread` might be unblocked by a handle
/// associated with a different `ParkThread` instance.
#[derive(Debug)]
pub struct ParkThread {
    _anchor: PhantomData<Rc<()>>,
}

/// Error returned by [`ParkThread`]
///
/// This currently is never returned, but might at some point in the future.
///
/// [`ParkThread`]: struct.ParkThread.html
#[derive(Debug)]
pub struct ParkError {
    _p: (),
}

/// Unblocks a thread that was blocked by `ParkThread`.
#[derive(Clone, Debug)]
pub struct UnparkThread {
    inner: Arc<Inner>,
}

#[derive(Debug)]
struct Inner {
    state: AtomicUsize,
    mutex: Mutex<()>,
    condvar: Condvar,
}

const IDLE: usize = 0;
const NOTIFY: usize = 1;
const SLEEP: usize = 2;

thread_local! {
    static CURRENT_PARK_THREAD: Arc<Inner> = Arc::new(Inner {
        state: AtomicUsize::new(IDLE),
        mutex: Mutex::new(()),
        condvar: Condvar::new(),
    });
}

// ===== impl ParkThread =====

impl ParkThread {
    /// Create a new `ParkThread` handle for the current thread.
    ///
    /// This type cannot be moved to other threads, so it should be created on
    /// the thread that the caller intends to park.
    pub fn new() -> ParkThread {
        ParkThread {
            _anchor: PhantomData,
        }
    }

    /// Get a reference to the `ParkThread` handle for this thread.
    fn with_current<F, R>(&self, f: F) -> R
        where F: FnOnce(&Arc<Inner>) -> R,
    {
        CURRENT_PARK_THREAD.with(|inner| f(inner))
    }
}

impl Park for ParkThread {
    type Unpark = UnparkThread;
    type Error = ParkError;

    fn unpark(&self) -> Self::Unpark {
        let inner = self.with_current(|inner| inner.clone());
        UnparkThread { inner }
    }

    fn park(&mut self) -> Result<(), Self::Error> {
        self.with_current(|inner| inner.park(None))
    }

    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
        self.with_current(|inner| inner.park(Some(duration)))
    }
}

// ===== impl UnparkThread =====

impl Unpark for UnparkThread {
    fn unpark(&self) {
        self.inner.unpark();
    }
}

// ===== impl Inner =====

impl Inner {
    /// Park the current thread for at most `dur`.
    fn park(&self, timeout: Option<Duration>) -> Result<(), ParkError> {
        // If currently notified, then we skip sleeping. This is checked outside
        // of the lock to avoid acquiring a mutex if not necessary.
        match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) {
            NOTIFY => return Ok(()),
            IDLE => {},
            _ => unreachable!(),
        }

        // The state is currently idle, so obtain the lock and then try to
        // transition to a sleeping state.
        let mut m = self.mutex.lock().unwrap();

        // Transition to sleeping
        match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) {
            NOTIFY => {
                // Notified before we could sleep, consume the notification and
                // exit
                self.state.store(IDLE, Ordering::SeqCst);
                return Ok(());
            }
            IDLE => {},
            _ => unreachable!(),
        }

        m = match timeout {
            Some(timeout) => self.condvar.wait_timeout(m, timeout).unwrap().0,
            None => self.condvar.wait(m).unwrap(),
        };

        // Transition back to idle. If the state has transitioned to `NOTIFY`,
        // this will consume that notification
        self.state.store(IDLE, Ordering::SeqCst);

        // Explicitly drop the mutex guard. There is no real point in doing it
        // except that I find it helpful to make it explicit where we want the
        // mutex to unlock.
        drop(m);

        Ok(())
    }

    fn unpark(&self) {
        // First, try transitioning from IDLE -> NOTIFY, this does not require a
        // lock.
        match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) {
            IDLE | NOTIFY => return,
            SLEEP => {}
            _ => unreachable!(),
        }

        // The other half is sleeping, this requires a lock
        let _m = self.mutex.lock().unwrap();

        // Transition to NOTIFY
        match self.state.swap(NOTIFY, Ordering::SeqCst) {
            SLEEP => {}
            NOTIFY => return,
            IDLE => return,
            _ => unreachable!(),
        }

        // Wakeup the sleeper
        self.condvar.notify_one();
    }
}
<commit_msg>executor: impl Unpark for Arc<Unpark> (#802)<commit_after>//! Abstraction over blocking and unblocking the current thread.
//!
//! Provides an abstraction over blocking the current thread. This is similar to
//! the park / unpark constructs provided by [`std`] but made generic. This
//! allows embedding custom functionality to perform when the thread is blocked.
//!
//! A blocked [`Park`][p] instance is unblocked by calling [`unpark`] on its
//! [`Unpark`][up] handle.
//!
//! The [`ParkThread`] struct implements [`Park`][p] using
//! [`thread::park`][`std`] to put the thread to sleep. The Tokio reactor also
//! implements park, but uses [`mio::Poll`][mio] to block the thread instead.
//!
//! The [`Park`][p] trait is composable. A timer implementation might decorate a
//! [`Park`][p] implementation by checking if any timeouts have elapsed after
//! the inner [`Park`][p] implementation unblocks.
//!
//! # Model
//!
//! Conceptually, each [`Park`][p] instance has an associated token, which is
//! initially not present:
//!
//! * The [`park`] method blocks the current thread unless or until the token
//!   is available, at which point it atomically consumes the token.
//! * The [`unpark`] method atomically makes the token available if it wasn't
//!   already.
//!
//! Some things to note:
//!
//! * If [`unpark`] is called before [`park`], the next call to [`park`] will
//! **not** block the thread.
//! * **Spurious** wakeups are permitted, i.e., the [`park`] method may unblock
//!   even if [`unpark`] was not called.
//! * [`park_timeout`] does the same as [`park`] but allows specifying a maximum
//!   time to block the thread for.
//!
//! [`std`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`thread::park`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`ParkThread`]: struct.ParkThread.html
//! [p]: trait.Park.html
//! [`park`]: trait.Park.html#tymethod.park
//! [`park_timeout`]: trait.Park.html#tymethod.park_timeout
//! [`unpark`]: trait.Unpark.html#tymethod.unpark
//! [up]: trait.Unpark.html
//! [mio]: https://docs.rs/mio/0.6/mio/struct.Poll.html

use std::marker::PhantomData;
use std::rc::Rc;
use std::sync::{Arc, Mutex, Condvar};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

/// Block the current thread.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
pub trait Park {
    /// Unpark handle type for the `Park` implementation.
    type Unpark: Unpark;

    /// Error returned by `park`
    type Error;

    /// Get a new `Unpark` handle associated with this `Park` instance.
    fn unpark(&self) -> Self::Unpark;

    /// Block the current thread unless or until the token is available.
    ///
    /// A call to `park` does not guarantee that the thread will remain blocked
    /// forever, and callers should be prepared for this possibility. This
    /// function may wakeup spuriously for any reason.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Park` implementation
    ///
    /// [mod]: ../index.html
    fn park(&mut self) -> Result<(), Self::Error>;

    /// Park the current thread for at most `duration`.
    ///
    /// This function is the same as `park` but allows specifying a maximum time
    /// to block the thread for.
    ///
    /// Same as `park`, there is no guarantee that the thread will remain
    /// blocked for any amount of time. Spurious wakeups are permitted for any
    /// reason.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Park` implementation
    ///
    /// [mod]: ../index.html
    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error>;
}

/// Unblock a thread blocked by the associated [`Park`] instance.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
/// [`Park`]: trait.Park.html
pub trait Unpark: Sync + Send + 'static {
    /// Unblock a thread that is blocked by the associated `Park` handle.
    ///
    /// Calling `unpark` atomically makes available the unpark token, if it is
    /// not already available.
    ///
    /// See [module documentation][mod] for more details.
    ///
    /// # Panics
    ///
    /// This function **should** not panic, but ultimately, panics are left as
    /// an implementation detail. Refer to the documentation for the specific
    /// `Unpark` implementation
    ///
    /// [mod]: ../index.html
    fn unpark(&self);
}

impl Unpark for Box<Unpark> {
    fn unpark(&self) {
        (**self).unpark()
    }
}

impl Unpark for Arc<Unpark> {
    fn unpark(&self) {
        (**self).unpark()
    }
}

/// Blocks the current thread using a condition variable.
///
/// Implements the [`Park`] functionality by using a condition variable. An
/// atomic variable is also used to avoid using the condition variable if
/// possible.
///
/// The condition variable is cached in a thread-local variable and is shared
/// across all `ParkThread` instances created on the same thread. This also
/// means that an instance of `ParkThread` might be unblocked by a handle
/// associated with a different `ParkThread` instance.
#[derive(Debug)]
pub struct ParkThread {
    _anchor: PhantomData<Rc<()>>,
}

/// Error returned by [`ParkThread`]
///
/// This currently is never returned, but might at some point in the future.
///
/// [`ParkThread`]: struct.ParkThread.html
#[derive(Debug)]
pub struct ParkError {
    _p: (),
}

/// Unblocks a thread that was blocked by `ParkThread`.
#[derive(Clone, Debug)]
pub struct UnparkThread {
    inner: Arc<Inner>,
}

#[derive(Debug)]
struct Inner {
    state: AtomicUsize,
    mutex: Mutex<()>,
    condvar: Condvar,
}

const IDLE: usize = 0;
const NOTIFY: usize = 1;
const SLEEP: usize = 2;

thread_local! {
    static CURRENT_PARK_THREAD: Arc<Inner> = Arc::new(Inner {
        state: AtomicUsize::new(IDLE),
        mutex: Mutex::new(()),
        condvar: Condvar::new(),
    });
}

// ===== impl ParkThread =====

impl ParkThread {
    /// Create a new `ParkThread` handle for the current thread.
    ///
    /// This type cannot be moved to other threads, so it should be created on
    /// the thread that the caller intends to park.
    pub fn new() -> ParkThread {
        ParkThread {
            _anchor: PhantomData,
        }
    }

    /// Get a reference to the `ParkThread` handle for this thread.
    fn with_current<F, R>(&self, f: F) -> R
        where F: FnOnce(&Arc<Inner>) -> R,
    {
        CURRENT_PARK_THREAD.with(|inner| f(inner))
    }
}

impl Park for ParkThread {
    type Unpark = UnparkThread;
    type Error = ParkError;

    fn unpark(&self) -> Self::Unpark {
        let inner = self.with_current(|inner| inner.clone());
        UnparkThread { inner }
    }

    fn park(&mut self) -> Result<(), Self::Error> {
        self.with_current(|inner| inner.park(None))
    }

    fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
        self.with_current(|inner| inner.park(Some(duration)))
    }
}

// ===== impl UnparkThread =====

impl Unpark for UnparkThread {
    fn unpark(&self) {
        self.inner.unpark();
    }
}

// ===== impl Inner =====

impl Inner {
    /// Park the current thread for at most `dur`.
    fn park(&self, timeout: Option<Duration>) -> Result<(), ParkError> {
        // If currently notified, then we skip sleeping. This is checked outside
        // of the lock to avoid acquiring a mutex if not necessary.
        match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) {
            NOTIFY => return Ok(()),
            IDLE => {},
            _ => unreachable!(),
        }

        // The state is currently idle, so obtain the lock and then try to
        // transition to a sleeping state.
        let mut m = self.mutex.lock().unwrap();

        // Transition to sleeping
        match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) {
            NOTIFY => {
                // Notified before we could sleep, consume the notification and
                // exit
                self.state.store(IDLE, Ordering::SeqCst);
                return Ok(());
            }
            IDLE => {},
            _ => unreachable!(),
        }

        m = match timeout {
            Some(timeout) => self.condvar.wait_timeout(m, timeout).unwrap().0,
            None => self.condvar.wait(m).unwrap(),
        };

        // Transition back to idle. If the state has transitioned to `NOTIFY`,
        // this will consume that notification
        self.state.store(IDLE, Ordering::SeqCst);

        // Explicitly drop the mutex guard. There is no real point in doing it
        // except that I find it helpful to make it explicit where we want the
        // mutex to unlock.
        drop(m);

        Ok(())
    }

    fn unpark(&self) {
        // First, try transitioning from IDLE -> NOTIFY, this does not require a
        // lock.
        match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) {
            IDLE | NOTIFY => return,
            SLEEP => {}
            _ => unreachable!(),
        }

        // The other half is sleeping, this requires a lock
        let _m = self.mutex.lock().unwrap();

        // Transition to NOTIFY
        match self.state.swap(NOTIFY, Ordering::SeqCst) {
            SLEEP => {}
            NOTIFY => return,
            IDLE => return,
            _ => unreachable!(),
        }

        // Wakeup the sleeper
        self.condvar.notify_one();
    }
}
<|endoftext|>


<commit_before>//! Shadowsocks Local server serving on a Tun interface

#[cfg(unix)]
use std::os::unix::io::RawFd;
use std::{
    io::{self, Cursor, ErrorKind},
    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use byte_string::ByteStr;
use bytes::BytesMut;
use etherparse::{IpHeader, PacketHeaders, ReadError, TransportHeader};
use futures::future;
use ipnet::{IpNet, Ipv4Net};
use log::{error, info, trace, warn};
use shadowsocks::config::Mode;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tun::{AsyncDevice, Configuration as TunConfiguration, Device, Error as TunError, Layer};

use crate::local::{context::ServiceContext, loadbalancing::PingBalancer};

use self::{
    sys::{set_route_configuration, write_packet_with_pi, IFF_PI_PREFIX_LEN},
    tcp::TcpTun,
    udp::UdpTun,
};

mod sys;
mod tcp;
mod udp;

pub struct TunBuilder {
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    tun_config: TunConfiguration,
    udp_expiry_duration: Option<Duration>,
    udp_capacity: Option<usize>,
    mode: Mode,
}

impl TunBuilder {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer) -> TunBuilder {
        TunBuilder {
            context,
            balancer,
            tun_config: TunConfiguration::default(),
            udp_expiry_duration: None,
            udp_capacity: None,
            mode: Mode::TcpOnly,
        }
    }

    pub fn address(mut self, addr: IpNet) -> TunBuilder {
        self.tun_config.address(addr.addr()).netmask(addr.netmask());
        self
    }

    pub fn name(mut self, name: &str) -> TunBuilder {
        self.tun_config.name(name);
        self
    }

    #[cfg(unix)]
    pub fn file_descriptor(mut self, fd: RawFd) -> TunBuilder {
        self.tun_config.raw_fd(fd);
        self
    }

    pub fn udp_expiry_duration(mut self, udp_expiry_duration: Duration) -> TunBuilder {
        self.udp_expiry_duration = Some(udp_expiry_duration);
        self
    }

    pub fn udp_capacity(mut self, udp_capacity: usize) -> TunBuilder {
        self.udp_capacity = Some(udp_capacity);
        self
    }

    pub fn mode(mut self, mode: Mode) -> TunBuilder {
        self.mode = mode;
        self
    }

    pub async fn build(mut self) -> io::Result<Tun> {
        self.tun_config.layer(Layer::L3).up();

        #[cfg(any(target_os = "linux"))]
        self.tun_config.platform(|tun_config| {
            // IFF_NO_PI preventing excessive buffer reallocating
            tun_config.packet_information(false);
        });

        let device = match tun::create_as_async(&self.tun_config) {
            Ok(d) => d,
            Err(TunError::Io(err)) => return Err(err),
            Err(err) => return Err(io::Error::new(ErrorKind::Other, err)),
        };

        let tun_address = match device.get_ref().address() {
            Ok(t) => t,
            Err(err) => {
                error!(
                    "tun device doesn't have address, error: {}, set it by tun_interface_address",
                    err
                );
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let tun_netmask = match device.get_ref().netmask() {
            Ok(m) => m,
            Err(err) => {
                error!(
                    "tun device doesn't have netmask, error: {}, set it by tun_interface_address",
                    err
                );
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let tun_name = device.get_ref().name();

        trace!("tun {}, address: {}, netmask: {}", tun_name, tun_address, tun_netmask);

        if let Err(err) = set_route_configuration(device.get_ref()).await {
            warn!(
                "failed to set system route for {}, consider set it manually, error: {}",
                tun_name, err
            );
        }

        let tun_netmask_u32: u32 = tun_netmask.into();

        let tun_network = Ipv4Net::new(tun_address, tun_netmask_u32.leading_ones() as u8).expect("Ipv4Net::new");

        Ok(Tun {
            device,
            tcp: if self.mode.enable_tcp() {
                Some(TcpTun::new(self.context.clone(), tun_network.into(), self.balancer.clone()).await?)
            } else {
                None
            },
            udp: if self.mode.enable_udp() {
                Some(UdpTun::new(
                    self.context,
                    self.balancer,
                    self.udp_expiry_duration,
                    self.udp_capacity,
                ))
            } else {
                None
            },
            mode: self.mode,
        })
    }
}

pub struct Tun {
    device: AsyncDevice,
    tcp: Option<TcpTun>,
    udp: Option<UdpTun>,
    mode: Mode,
}

impl Tun {
    pub async fn run(mut self) -> io::Result<()> {
        let mtu = self.device.get_ref().mtu().expect("mtu");
        assert!(mtu > 0 && mtu as usize > IFF_PI_PREFIX_LEN);

        info!(
            "shadowsocks tun device {}, address {}, netmask {}, mtu {}, mode {:?}",
            self.device.get_ref().name(),
            self.device.get_ref().address().expect("address"),
            self.device.get_ref().netmask().expect("netmask"),
            mtu,
            self.mode,
        );

        let mut packet_buffer = vec![0u8; mtu as usize + IFF_PI_PREFIX_LEN].into_boxed_slice();

        loop {
            #[inline(always)]
            async fn udp_recv_packet(udp: &mut Option<UdpTun>) -> BytesMut {
                match *udp {
                    Some(ref mut udp) => udp.recv_packet().await,
                    None => future::pending().await,
                }
            }

            tokio::select! {
                // tun device
                n = self.device.read(&mut packet_buffer) => {
                    let n = n?;

                    if n <= IFF_PI_PREFIX_LEN {
                        error!(
                            "[TUN] packet too short, packet: {:?}",
                            ByteStr::new(&packet_buffer[..n])
                        );
                        continue;
                    }

                    let packet = &mut packet_buffer[IFF_PI_PREFIX_LEN..n];
                    trace!("[TUN] received IP packet {:?}", ByteStr::new(packet));

                    if self.handle_packet(packet).await? {
                        self.device.write_all(&packet_buffer[..n]).await?;
                    }
                }

                // channel sent back
                packet = udp_recv_packet(&mut self.udp) => {
                    if let Err(err) = write_packet_with_pi(&mut self.device, &packet).await {
                        error!("failed to set packet information, error: {}, {:?}", err, ByteStr::new(&packet));
                    }
                }
            }
        }
    }

    async fn handle_packet(&mut self, packet: &mut [u8]) -> io::Result<bool> {
        let mut ph = match PacketHeaders::from_ip_slice(packet) {
            Ok(ph) => ph,
            Err(ReadError::IoError(err)) => return Err(err),
            Err(err) => {
                error!("invalid IP packet, error: {:?}, {:?}", err, ByteStr::new(packet));
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let payload_len = ph.payload.len();

        let mut ip_header = match ph.ip {
            Some(ref mut i) => i,
            None => {
                error!("unrecognized ethernet packet {:?}", ph);
                return Err(io::Error::new(ErrorKind::Other, "unrecognized ethernet packet"));
            }
        };

        let (src_ip, dst_ip) = match *ip_header {
            IpHeader::Version4(ref v4) => (Ipv4Addr::from(v4.source).into(), Ipv4Addr::from(v4.destination).into()),
            IpHeader::Version6(ref v6) => (Ipv6Addr::from(v6.source).into(), Ipv6Addr::from(v6.destination).into()),
        };

        match ph.transport {
            Some(TransportHeader::Tcp(ref mut tcp_header)) => {
                let tcp = match self.tcp {
                    Some(ref mut tcp) => tcp,
                    None => return Ok(false),
                };

                let src_addr = SocketAddr::new(src_ip, tcp_header.source_port);
                let dst_addr = SocketAddr::new(dst_ip, tcp_header.destination_port);

                let (mod_src_addr, mod_dst_addr) = match tcp.handle_packet(src_addr, dst_addr, tcp_header).await {
                    Ok(Some(a)) => a,
                    Ok(None) => return Ok(false),
                    Err(err) => {
                        error!("handle TCP/IP packet failed, error: {}", err);
                        return Ok(false);
                    }
                };

                // Replaces IP_HEADER, TRANSPORT_HEADER directly into packet
                match (mod_src_addr, &mut ip_header) {
                    (SocketAddr::V4(v4addr), IpHeader::Version4(v4ip)) => v4ip.source = v4addr.ip().octets(),
                    (SocketAddr::V6(v6addr), IpHeader::Version6(v6ip)) => v6ip.source = v6addr.ip().octets(),
                    _ => unreachable!("modified saddr not match"),
                }
                tcp_header.source_port = mod_src_addr.port();
                match (mod_dst_addr, &mut ip_header) {
                    (SocketAddr::V4(v4addr), IpHeader::Version4(v4ip)) => v4ip.destination = v4addr.ip().octets(),
                    (SocketAddr::V6(v6addr), IpHeader::Version6(v6ip)) => v6ip.destination = v6addr.ip().octets(),
                    _ => unreachable!("modified daddr not match"),
                }
                tcp_header.destination_port = mod_dst_addr.port();
                match ip_header {
                    IpHeader::Version4(v4) => {
                        tcp_header.checksum = tcp_header
                            .calc_checksum_ipv4(v4, ph.payload)
                            .expect("calc_checksum_ipv4")
                    }
                    IpHeader::Version6(v6) => {
                        tcp_header.checksum = tcp_header
                            .calc_checksum_ipv6(v6, ph.payload)
                            .expect("calc_checksum_ipv6")
                    }
                }

                let (headers, _) = packet.split_at_mut(packet.len() - payload_len);
                let mut cursor = Cursor::new(headers);

                ip_header.write(&mut cursor).expect("ip_header.write");
                tcp_header.write(&mut cursor).expect("tcp_header.write");

                Ok(true)
            }
            Some(TransportHeader::Udp(ref udp_header)) => {
                // UDP proxies directly
                let udp = match self.udp {
                    Some(ref mut udp) => udp,
                    None => return Ok(false),
                };

                let src_addr = SocketAddr::new(src_ip, udp_header.source_port);
                let dst_addr = SocketAddr::new(dst_ip, udp_header.destination_port);

                if let Err(err) = udp.handle_packet(src_addr, dst_addr, ph.payload).await {
                    error!("handle UDP/IP packet failed, error: {}", err);
                }

                Ok(false)
            }
            None => {
                trace!("no transport layer in ethernet packet {:?}", ph);
                Ok(false)
            }
        }
    }
}
<commit_msg>show mode in configuration format<commit_after>//! Shadowsocks Local server serving on a Tun interface

#[cfg(unix)]
use std::os::unix::io::RawFd;
use std::{
    io::{self, Cursor, ErrorKind},
    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use byte_string::ByteStr;
use bytes::BytesMut;
use etherparse::{IpHeader, PacketHeaders, ReadError, TransportHeader};
use futures::future;
use ipnet::{IpNet, Ipv4Net};
use log::{error, info, trace, warn};
use shadowsocks::config::Mode;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tun::{AsyncDevice, Configuration as TunConfiguration, Device, Error as TunError, Layer};

use crate::local::{context::ServiceContext, loadbalancing::PingBalancer};

use self::{
    sys::{set_route_configuration, write_packet_with_pi, IFF_PI_PREFIX_LEN},
    tcp::TcpTun,
    udp::UdpTun,
};

mod sys;
mod tcp;
mod udp;

pub struct TunBuilder {
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    tun_config: TunConfiguration,
    udp_expiry_duration: Option<Duration>,
    udp_capacity: Option<usize>,
    mode: Mode,
}

impl TunBuilder {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer) -> TunBuilder {
        TunBuilder {
            context,
            balancer,
            tun_config: TunConfiguration::default(),
            udp_expiry_duration: None,
            udp_capacity: None,
            mode: Mode::TcpOnly,
        }
    }

    pub fn address(mut self, addr: IpNet) -> TunBuilder {
        self.tun_config.address(addr.addr()).netmask(addr.netmask());
        self
    }

    pub fn name(mut self, name: &str) -> TunBuilder {
        self.tun_config.name(name);
        self
    }

    #[cfg(unix)]
    pub fn file_descriptor(mut self, fd: RawFd) -> TunBuilder {
        self.tun_config.raw_fd(fd);
        self
    }

    pub fn udp_expiry_duration(mut self, udp_expiry_duration: Duration) -> TunBuilder {
        self.udp_expiry_duration = Some(udp_expiry_duration);
        self
    }

    pub fn udp_capacity(mut self, udp_capacity: usize) -> TunBuilder {
        self.udp_capacity = Some(udp_capacity);
        self
    }

    pub fn mode(mut self, mode: Mode) -> TunBuilder {
        self.mode = mode;
        self
    }

    pub async fn build(mut self) -> io::Result<Tun> {
        self.tun_config.layer(Layer::L3).up();

        #[cfg(any(target_os = "linux"))]
        self.tun_config.platform(|tun_config| {
            // IFF_NO_PI preventing excessive buffer reallocating
            tun_config.packet_information(false);
        });

        let device = match tun::create_as_async(&self.tun_config) {
            Ok(d) => d,
            Err(TunError::Io(err)) => return Err(err),
            Err(err) => return Err(io::Error::new(ErrorKind::Other, err)),
        };

        let tun_address = match device.get_ref().address() {
            Ok(t) => t,
            Err(err) => {
                error!(
                    "tun device doesn't have address, error: {}, set it by tun_interface_address",
                    err
                );
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let tun_netmask = match device.get_ref().netmask() {
            Ok(m) => m,
            Err(err) => {
                error!(
                    "tun device doesn't have netmask, error: {}, set it by tun_interface_address",
                    err
                );
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let tun_name = device.get_ref().name();

        trace!("tun {}, address: {}, netmask: {}", tun_name, tun_address, tun_netmask);

        if let Err(err) = set_route_configuration(device.get_ref()).await {
            warn!(
                "failed to set system route for {}, consider set it manually, error: {}",
                tun_name, err
            );
        }

        let tun_netmask_u32: u32 = tun_netmask.into();

        let tun_network = Ipv4Net::new(tun_address, tun_netmask_u32.leading_ones() as u8).expect("Ipv4Net::new");

        Ok(Tun {
            device,
            tcp: if self.mode.enable_tcp() {
                Some(TcpTun::new(self.context.clone(), tun_network.into(), self.balancer.clone()).await?)
            } else {
                None
            },
            udp: if self.mode.enable_udp() {
                Some(UdpTun::new(
                    self.context,
                    self.balancer,
                    self.udp_expiry_duration,
                    self.udp_capacity,
                ))
            } else {
                None
            },
            mode: self.mode,
        })
    }
}

pub struct Tun {
    device: AsyncDevice,
    tcp: Option<TcpTun>,
    udp: Option<UdpTun>,
    mode: Mode,
}

impl Tun {
    pub async fn run(mut self) -> io::Result<()> {
        let mtu = self.device.get_ref().mtu().expect("mtu");
        assert!(mtu > 0 && mtu as usize > IFF_PI_PREFIX_LEN);

        info!(
            "shadowsocks tun device {}, address {}, netmask {}, mtu {}, mode {}",
            self.device.get_ref().name(),
            self.device.get_ref().address().expect("address"),
            self.device.get_ref().netmask().expect("netmask"),
            mtu,
            self.mode,
        );

        let mut packet_buffer = vec![0u8; mtu as usize + IFF_PI_PREFIX_LEN].into_boxed_slice();

        loop {
            #[inline(always)]
            async fn udp_recv_packet(udp: &mut Option<UdpTun>) -> BytesMut {
                match *udp {
                    Some(ref mut udp) => udp.recv_packet().await,
                    None => future::pending().await,
                }
            }

            tokio::select! {
                // tun device
                n = self.device.read(&mut packet_buffer) => {
                    let n = n?;

                    if n <= IFF_PI_PREFIX_LEN {
                        error!(
                            "[TUN] packet too short, packet: {:?}",
                            ByteStr::new(&packet_buffer[..n])
                        );
                        continue;
                    }

                    let packet = &mut packet_buffer[IFF_PI_PREFIX_LEN..n];
                    trace!("[TUN] received IP packet {:?}", ByteStr::new(packet));

                    if self.handle_packet(packet).await? {
                        self.device.write_all(&packet_buffer[..n]).await?;
                    }
                }

                // channel sent back
                packet = udp_recv_packet(&mut self.udp) => {
                    if let Err(err) = write_packet_with_pi(&mut self.device, &packet).await {
                        error!("failed to set packet information, error: {}, {:?}", err, ByteStr::new(&packet));
                    }
                }
            }
        }
    }

    async fn handle_packet(&mut self, packet: &mut [u8]) -> io::Result<bool> {
        let mut ph = match PacketHeaders::from_ip_slice(packet) {
            Ok(ph) => ph,
            Err(ReadError::IoError(err)) => return Err(err),
            Err(err) => {
                error!("invalid IP packet, error: {:?}, {:?}", err, ByteStr::new(packet));
                return Err(io::Error::new(ErrorKind::Other, err));
            }
        };

        let payload_len = ph.payload.len();

        let mut ip_header = match ph.ip {
            Some(ref mut i) => i,
            None => {
                error!("unrecognized ethernet packet {:?}", ph);
                return Err(io::Error::new(ErrorKind::Other, "unrecognized ethernet packet"));
            }
        };

        let (src_ip, dst_ip) = match *ip_header {
            IpHeader::Version4(ref v4) => (Ipv4Addr::from(v4.source).into(), Ipv4Addr::from(v4.destination).into()),
            IpHeader::Version6(ref v6) => (Ipv6Addr::from(v6.source).into(), Ipv6Addr::from(v6.destination).into()),
        };

        match ph.transport {
            Some(TransportHeader::Tcp(ref mut tcp_header)) => {
                let tcp = match self.tcp {
                    Some(ref mut tcp) => tcp,
                    None => return Ok(false),
                };

                let src_addr = SocketAddr::new(src_ip, tcp_header.source_port);
                let dst_addr = SocketAddr::new(dst_ip, tcp_header.destination_port);

                let (mod_src_addr, mod_dst_addr) = match tcp.handle_packet(src_addr, dst_addr, tcp_header).await {
                    Ok(Some(a)) => a,
                    Ok(None) => return Ok(false),
                    Err(err) => {
                        error!("handle TCP/IP packet failed, error: {}", err);
                        return Ok(false);
                    }
                };

                // Replaces IP_HEADER, TRANSPORT_HEADER directly into packet
                match (mod_src_addr, &mut ip_header) {
                    (SocketAddr::V4(v4addr), IpHeader::Version4(v4ip)) => v4ip.source = v4addr.ip().octets(),
                    (SocketAddr::V6(v6addr), IpHeader::Version6(v6ip)) => v6ip.source = v6addr.ip().octets(),
                    _ => unreachable!("modified saddr not match"),
                }
                tcp_header.source_port = mod_src_addr.port();
                match (mod_dst_addr, &mut ip_header) {
                    (SocketAddr::V4(v4addr), IpHeader::Version4(v4ip)) => v4ip.destination = v4addr.ip().octets(),
                    (SocketAddr::V6(v6addr), IpHeader::Version6(v6ip)) => v6ip.destination = v6addr.ip().octets(),
                    _ => unreachable!("modified daddr not match"),
                }
                tcp_header.destination_port = mod_dst_addr.port();
                match ip_header {
                    IpHeader::Version4(v4) => {
                        tcp_header.checksum = tcp_header
                            .calc_checksum_ipv4(v4, ph.payload)
                            .expect("calc_checksum_ipv4")
                    }
                    IpHeader::Version6(v6) => {
                        tcp_header.checksum = tcp_header
                            .calc_checksum_ipv6(v6, ph.payload)
                            .expect("calc_checksum_ipv6")
                    }
                }

                let (headers, _) = packet.split_at_mut(packet.len() - payload_len);
                let mut cursor = Cursor::new(headers);

                ip_header.write(&mut cursor).expect("ip_header.write");
                tcp_header.write(&mut cursor).expect("tcp_header.write");

                Ok(true)
            }
            Some(TransportHeader::Udp(ref udp_header)) => {
                // UDP proxies directly
                let udp = match self.udp {
                    Some(ref mut udp) => udp,
                    None => return Ok(false),
                };

                let src_addr = SocketAddr::new(src_ip, udp_header.source_port);
                let dst_addr = SocketAddr::new(dst_ip, udp_header.destination_port);

                if let Err(err) = udp.handle_packet(src_addr, dst_addr, ph.payload).await {
                    error!("handle UDP/IP packet failed, error: {}", err);
                }

                Ok(false)
            }
            None => {
                trace!("no transport layer in ethernet packet {:?}", ph);
                Ok(false)
            }
        }
    }
}
<|endoftext|>