<commit_before>//! Contains types related to 2D image data.

use std::ops;
use std::mem;
use std::slice;
use gl;

/// A trait for types that that contain 2D image data, which can
/// be uploaded to a texture using a [`Texture2dBuilder`]
/// (../context/texture_context/struct.Texture2dBuilder)
/// or using the [`gl.image_2d`]
/// (../context/texture_context/trait.ContextTextureExt.html#method.image_2d)
/// method.
///
/// See [`Pixels`](struct.Pixels.html) for a simple `Image2d`.
pub trait Image2d {
    /// Get the width of the image, in texels.
    fn width(&self) -> usize;

    /// Get the height of the image, in texels.
    fn height(&self) -> usize;

    /// Get the format of the image data that is returned by the
    /// [`texel_bytes`](trait.Image2d.html#tymethod.texel_bytes) method.
    fn format(&self) -> ImageFormat;

    /// Get the raw texel data of the image data, as a `u8` slice.
    fn texel_bytes(&self) -> &[u8];
}

/// A single OpenGL color value, with `u8` components laid out
/// as a C struct in RGBA order. This type is the simplest implementation
/// of [`Image2d`](trait.Image2d.html), which allows it to be uploaded
/// to a texture.
#[derive(Debug, Clone, Copy, Default)]
#[repr(C)]
pub struct Pixel {
    /// The red component.
    pub r: u8,

    /// The green component.
    pub g: u8,

    /// The blue component.
    pub b: u8,

    /// The alpha component.
    pub a: u8
}

impl Pixel {
    /// Create an RGBA color, provided each color component.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::r_g_b_a(0xAA, 0xBB, 0xCC, 0xDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn r_g_b_a(r: u8, g: u8, b: u8, a: u8) -> Self {
        Pixel { r: r, g: g, b: b, a: a }
    }

    /// Create an RGBA color, provided the RGB components and using
    /// `0xFF` as the A value.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::r_g_b(0xAA, 0xBB, 0xCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xFF);
    /// ```
    pub fn r_g_b(r: u8, g: u8, b: u8) -> Self {
        Pixel::r_g_b_a(r, g, b, 0xFF)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value. The `u32` value will be read as `0x00RRGGBB`.
    ///
    /// # Examples
    /// ```
    /// // The `0x55` component gets discared
    /// let pixel = glitter::Pixel::rgb(0x55AABBCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xFF);
    /// ```
    pub fn rgb(rgb: u32) -> Self {
        Pixel::rgb_a(rgb, 0xFF)
    }

    /// Create an RGBA color value, provided the components as a packed
    /// `u32` value. The `u32` value will be read as `0xAARRGGBB`.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::argb(0xDDAABBCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn argb(argb: u32) -> Self {
        let a = (argb & 0xFF000000) >> 24;
        Pixel::rgb_a(argb, a as u8)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value. The `u32` value will be read as `0xRRGGBBAA`.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::rgba(0xAABBCCDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn rgba(rgba: u32) -> Self {
        let r = (rgba & 0xFF000000) >> 24;
        let g = (rgba & 0x00FF0000) >> 16;
        let b = (rgba & 0x0000FF00) >> 8;
        let a =  rgba & 0x000000FF;
        Pixel::r_g_b_a(r as u8, g as u8, b as u8, a as u8)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value, and a separate A component. The `u32` value
    /// will be read as `0x00RRGGBB`.
    ///
    /// # Examples
    /// ```
    /// // The `0x55` component gets discarded
    /// let pixel = glitter::Pixel::rgb_a(0x55AABBCC, 0xDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn rgb_a(rgb: u32, a: u8) -> Self {
        let r = (rgb & 0xFF0000) >> 16;
        let g = (rgb & 0x00FF00) >> 8;
        let b =  rgb & 0x0000FF;
        Pixel::r_g_b_a(r as u8, g as u8, b as u8, a)
    }
}

/// A (heap-allocated) 2D image composed of a list of pixels.
///
/// # Example
///
/// ```
/// // Draw a red circle with a 50px radius in a 100px * 100px black image.
/// extern crate glitter;
///
/// let width = 100;
/// let height = 100;
/// let radius = 50;
///
// let (center_x, center_y) = (width as f32/2.0, height as f32/2.0);
//
// let mut pixels = glitter::Pixels::new(width, height);
// for x in 0..width {
//     for y in 0..height {
//         let dx = center_x - x as f32;
//         let dy = center_y - y as f32;
//         let distance = (dx*dx + dy*dy).sqrt();
//
//         let color = if distance < radius {
//             // The point is within the circle, so it should be red
//             glitter::Pixel::rgb(0xFF0000)
//         }
//         else {
//             // The point is outside the circle, so it should be black
//             glitter::Pixel::rgb(0x000000)
//         };
//         pixels[y][x] = color;
//     }
// }
// ```
#[derive(Clone)]
pub struct Pixels {
    width: usize,
    height: usize,
    pixels: Vec<Pixel>
}

impl Pixels {
    /// Create a new image with the desired width and height. Pixels
    /// are initialized with [`Pixel::default()`]
    /// (https://doc.rust-lang.org/nightly/core/default/trait.Default.html#tymethod.default).
    pub fn new(width: usize, height: usize) -> Self {
        Pixels {
            width: width,
            height: height,
            pixels: vec![Pixel::default(); width * height]
        }
    }
}

impl ops::Index<usize> for Pixels {
    type Output = [Pixel];

    fn index(&self, row: usize) -> &[Pixel] {
        debug_assert!(row < self.height);

        &self.pixels[(row*self.width)..((row+1)*self.width)]
    }
}

impl ops::IndexMut<usize> for Pixels {
    fn index_mut(&mut self, row: usize) -> &mut [Pixel] {
        debug_assert!(row < self.height);

        &mut self.pixels[(row*self.width)..((row+1)*self.width)]
    }
}

impl Image2d for Pixels {
    fn width(&self) -> usize {
        self.width
    }

    fn height(&self) -> usize {
        self.height
    }

    fn format(&self) -> ImageFormat {
        ImageFormat::rgba8()
    }

    fn texel_bytes(&self) -> &[u8] {
        let len = self.pixels.len() * mem::size_of::<Pixel>();
        unsafe {
            slice::from_raw_parts(mem::transmute(&self.pixels[0]), len)
        }
    }
}

gl_enum! {
    /// The data types for a texel.
    pub gl_enum TexelType {
        /// Each texel is 4 * 8 bits.
        pub const UnsignedByte as UNSIGNED_BYTE_TEXEL =
            gl::UNSIGNED_BYTE,

        /// Each texel is 16 bits, with 5-bit/6-bit/5-bit components.
        pub const UnsignedShort565 as UNSIGNED_SHORT_5_6_5 =
            gl::UNSIGNED_SHORT_5_6_5,

        /// Each texel is 16 bits, with 4, 4-bit components.
        pub const UnsignedShort4444 as UNSIGNED_SHORT_4_4_4_4 =
            gl::UNSIGNED_SHORT_4_4_4_4,

        /// Each texel is 16 bits, with 5-bit/5-bit/5-bit/1-bit components.
        pub const UnsignedShort5551 as UNSIGNED_SHORT_5_5_5_1 =
            gl::UNSIGNED_SHORT_5_5_5_1
    }
}

gl_enum! {
    /// The different texel formats.
    pub gl_enum TexelFormat {
        /// A texel contains only an alpha component.
        pub const Alpha as ALPHA = gl::ALPHA,

        /// A texel contains red, green, and blue components.
        pub const RGB as RGB = gl::RGB,

        /// A texel contains red, green, blue, and alpha components.
        pub const RGBA as RGBA = gl::RGBA
    }
}

gl_enum! {
    /// The various image formats of a renderbuffer.
    pub gl_enum RenderbufferFormat {
        /// The red, green, blue, and alpha channels are all stored with 4 bits.
        pub const RGBA4 as RGBA4 = gl::RGBA4,

        /// The red, green, and blue channels are stored with 5 bits, 6 bits,
        /// and 5 bits, respectively.
        pub const RGB565 as RGB565 = gl::RGB565,

        /// The red, green, and blue channels are stored with 5 bits, and the
        /// alpha channel is stored with 1 bit.
        pub const RGB5A1 as RGB5_A1 = gl::RGB5_A1,

        /// The renderbuffer stores a 16-bit depth component.
        pub const DepthComponent16 as DEPTH_COMPONENT16 = gl::DEPTH_COMPONENT16,

        /// The renderbuffer stores an 8-bit stencil component.
        pub const StencilIndex8 as STENCIL_INDEX8 = gl::STENCIL_INDEX8
    }
}

/// A type that contains the texel type and format that make up a 2D image.
#[derive(Debug, Clone, Copy)]
pub struct ImageFormat {
    /// The texel type of the image.
    pub texel_type: TexelType,

        /// The texel format of the image.
    pub texel_format: TexelFormat
}

impl ImageFormat {
    /// Returns the RGBA image format with 8 bits per component.
    pub fn rgba8() -> Self {
        ImageFormat {
            texel_type: TexelType::UnsignedByte,
            texel_format: TexelFormat::RGBA
        }
    }
}
<commit_msg>Docs: fix `Pixels` doc comment<commit_after>//! Contains types related to 2D image data.

use std::ops;
use std::mem;
use std::slice;
use gl;

/// A trait for types that that contain 2D image data, which can
/// be uploaded to a texture using a [`Texture2dBuilder`]
/// (../context/texture_context/struct.Texture2dBuilder)
/// or using the [`gl.image_2d`]
/// (../context/texture_context/trait.ContextTextureExt.html#method.image_2d)
/// method.
///
/// See [`Pixels`](struct.Pixels.html) for a simple `Image2d`.
pub trait Image2d {
    /// Get the width of the image, in texels.
    fn width(&self) -> usize;

    /// Get the height of the image, in texels.
    fn height(&self) -> usize;

    /// Get the format of the image data that is returned by the
    /// [`texel_bytes`](trait.Image2d.html#tymethod.texel_bytes) method.
    fn format(&self) -> ImageFormat;

    /// Get the raw texel data of the image data, as a `u8` slice.
    fn texel_bytes(&self) -> &[u8];
}

/// A single OpenGL color value, with `u8` components laid out
/// as a C struct in RGBA order. This type is the simplest implementation
/// of [`Image2d`](trait.Image2d.html), which allows it to be uploaded
/// to a texture.
#[derive(Debug, Clone, Copy, Default)]
#[repr(C)]
pub struct Pixel {
    /// The red component.
    pub r: u8,

    /// The green component.
    pub g: u8,

    /// The blue component.
    pub b: u8,

    /// The alpha component.
    pub a: u8
}

impl Pixel {
    /// Create an RGBA color, provided each color component.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::r_g_b_a(0xAA, 0xBB, 0xCC, 0xDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn r_g_b_a(r: u8, g: u8, b: u8, a: u8) -> Self {
        Pixel { r: r, g: g, b: b, a: a }
    }

    /// Create an RGBA color, provided the RGB components and using
    /// `0xFF` as the A value.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::r_g_b(0xAA, 0xBB, 0xCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xFF);
    /// ```
    pub fn r_g_b(r: u8, g: u8, b: u8) -> Self {
        Pixel::r_g_b_a(r, g, b, 0xFF)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value. The `u32` value will be read as `0x00RRGGBB`.
    ///
    /// # Examples
    /// ```
    /// // The `0x55` component gets discared
    /// let pixel = glitter::Pixel::rgb(0x55AABBCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xFF);
    /// ```
    pub fn rgb(rgb: u32) -> Self {
        Pixel::rgb_a(rgb, 0xFF)
    }

    /// Create an RGBA color value, provided the components as a packed
    /// `u32` value. The `u32` value will be read as `0xAARRGGBB`.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::argb(0xDDAABBCC);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn argb(argb: u32) -> Self {
        let a = (argb & 0xFF000000) >> 24;
        Pixel::rgb_a(argb, a as u8)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value. The `u32` value will be read as `0xRRGGBBAA`.
    ///
    /// # Examples
    /// ```
    /// let pixel = glitter::Pixel::rgba(0xAABBCCDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn rgba(rgba: u32) -> Self {
        let r = (rgba & 0xFF000000) >> 24;
        let g = (rgba & 0x00FF0000) >> 16;
        let b = (rgba & 0x0000FF00) >> 8;
        let a =  rgba & 0x000000FF;
        Pixel::r_g_b_a(r as u8, g as u8, b as u8, a as u8)
    }

    /// Create an RGBA color, provided the RGB components as a packed
    /// `u32` value, and a separate A component. The `u32` value
    /// will be read as `0x00RRGGBB`.
    ///
    /// # Examples
    /// ```
    /// // The `0x55` component gets discarded
    /// let pixel = glitter::Pixel::rgb_a(0x55AABBCC, 0xDD);
    /// assert_eq!(pixel.r, 0xAA);
    /// assert_eq!(pixel.g, 0xBB);
    /// assert_eq!(pixel.b, 0xCC);
    /// assert_eq!(pixel.a, 0xDD);
    /// ```
    pub fn rgb_a(rgb: u32, a: u8) -> Self {
        let r = (rgb & 0xFF0000) >> 16;
        let g = (rgb & 0x00FF00) >> 8;
        let b =  rgb & 0x0000FF;
        Pixel::r_g_b_a(r as u8, g as u8, b as u8, a)
    }
}

/// A (heap-allocated) 2D image composed of a list of pixels.
///
/// # Example
///
/// ```
/// // Draw a red circle with a 50px radius in a 100px * 100px black image.
/// extern crate glitter;
///
/// let width = 100;
/// let height = 100;
/// let radius = 50;
///
/// let (center_x, center_y) = (width as f32/2.0, height as f32/2.0);
///
/// let mut pixels = glitter::Pixels::new(width, height);
/// for x in 0..width {
///     for y in 0..height {
///         let dx = center_x - x as f32;
///         let dy = center_y - y as f32;
///         let distance = (dx*dx + dy*dy).sqrt();
///
///         let color = if distance < radius {
///             // The point is within the circle, so it should be red
///             glitter::Pixel::rgb(0xFF0000)
///         }
///         else {
///             // The point is outside the circle, so it should be black
///             glitter::Pixel::rgb(0x000000)
///         };
///         pixels[y][x] = color;
///     }
/// }
/// ```
#[derive(Clone)]
pub struct Pixels {
    width: usize,
    height: usize,
    pixels: Vec<Pixel>
}

impl Pixels {
    /// Create a new image with the desired width and height. Pixels
    /// are initialized with [`Pixel::default()`]
    /// (https://doc.rust-lang.org/nightly/core/default/trait.Default.html#tymethod.default).
    pub fn new(width: usize, height: usize) -> Self {
        Pixels {
            width: width,
            height: height,
            pixels: vec![Pixel::default(); width * height]
        }
    }
}

impl ops::Index<usize> for Pixels {
    type Output = [Pixel];

    fn index(&self, row: usize) -> &[Pixel] {
        debug_assert!(row < self.height);

        &self.pixels[(row*self.width)..((row+1)*self.width)]
    }
}

impl ops::IndexMut<usize> for Pixels {
    fn index_mut(&mut self, row: usize) -> &mut [Pixel] {
        debug_assert!(row < self.height);

        &mut self.pixels[(row*self.width)..((row+1)*self.width)]
    }
}

impl Image2d for Pixels {
    fn width(&self) -> usize {
        self.width
    }

    fn height(&self) -> usize {
        self.height
    }

    fn format(&self) -> ImageFormat {
        ImageFormat::rgba8()
    }

    fn texel_bytes(&self) -> &[u8] {
        let len = self.pixels.len() * mem::size_of::<Pixel>();
        unsafe {
            slice::from_raw_parts(mem::transmute(&self.pixels[0]), len)
        }
    }
}

gl_enum! {
    /// The data types for a texel.
    pub gl_enum TexelType {
        /// Each texel is 4 * 8 bits.
        pub const UnsignedByte as UNSIGNED_BYTE_TEXEL =
            gl::UNSIGNED_BYTE,

        /// Each texel is 16 bits, with 5-bit/6-bit/5-bit components.
        pub const UnsignedShort565 as UNSIGNED_SHORT_5_6_5 =
            gl::UNSIGNED_SHORT_5_6_5,

        /// Each texel is 16 bits, with 4, 4-bit components.
        pub const UnsignedShort4444 as UNSIGNED_SHORT_4_4_4_4 =
            gl::UNSIGNED_SHORT_4_4_4_4,

        /// Each texel is 16 bits, with 5-bit/5-bit/5-bit/1-bit components.
        pub const UnsignedShort5551 as UNSIGNED_SHORT_5_5_5_1 =
            gl::UNSIGNED_SHORT_5_5_5_1
    }
}

gl_enum! {
    /// The different texel formats.
    pub gl_enum TexelFormat {
        /// A texel contains only an alpha component.
        pub const Alpha as ALPHA = gl::ALPHA,

        /// A texel contains red, green, and blue components.
        pub const RGB as RGB = gl::RGB,

        /// A texel contains red, green, blue, and alpha components.
        pub const RGBA as RGBA = gl::RGBA
    }
}

gl_enum! {
    /// The various image formats of a renderbuffer.
    pub gl_enum RenderbufferFormat {
        /// The red, green, blue, and alpha channels are all stored with 4 bits.
        pub const RGBA4 as RGBA4 = gl::RGBA4,

        /// The red, green, and blue channels are stored with 5 bits, 6 bits,
        /// and 5 bits, respectively.
        pub const RGB565 as RGB565 = gl::RGB565,

        /// The red, green, and blue channels are stored with 5 bits, and the
        /// alpha channel is stored with 1 bit.
        pub const RGB5A1 as RGB5_A1 = gl::RGB5_A1,

        /// The renderbuffer stores a 16-bit depth component.
        pub const DepthComponent16 as DEPTH_COMPONENT16 = gl::DEPTH_COMPONENT16,

        /// The renderbuffer stores an 8-bit stencil component.
        pub const StencilIndex8 as STENCIL_INDEX8 = gl::STENCIL_INDEX8
    }
}

/// A type that contains the texel type and format that make up a 2D image.
#[derive(Debug, Clone, Copy)]
pub struct ImageFormat {
    /// The texel type of the image.
    pub texel_type: TexelType,

        /// The texel format of the image.
    pub texel_format: TexelFormat
}

impl ImageFormat {
    /// Returns the RGBA image format with 8 bits per component.
    pub fn rgba8() -> Self {
        ImageFormat {
            texel_type: TexelType::UnsignedByte,
            texel_format: TexelFormat::RGBA
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, SocketAddr},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::mpsc,
};

#[cfg(target_os = "android")]
use std::path::PathBuf;
#[cfg(target_os = "android")]
use tokio::net::UnixStream;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use byteorder::{BigEndian, ByteOrder};
use log::{debug, error, info};
use rand::Rng;
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{Name, RData, RecordType},
};

use crate::{
    config::{ConfigType, ServerConfig},
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        socks5::Address,
        sys::create_udp_socket,
        tcprelay::ProxyStream,
        utils::try_timeout,
    },
};

async fn stream_lookup<T>(qname: &Name, qtype: RecordType, stream: &mut T) -> io::Result<Message>
where
    T: AsyncReadExt + AsyncWriteExt + Unpin,
{
    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut send_buffer = vec![0; size + 2];

    BigEndian::write_u16(&mut send_buffer[0..2], size as u16);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    let mut res_buffer = vec![0; 2];
    stream.read_exact(&mut res_buffer[0..2]).await?;

    let size = BigEndian::read_u16(&res_buffer[0..2]) as usize;
    let mut res_buffer = vec![0; size];
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

#[cfg(target_os = "android")]
async fn local_lookup(qname: &Name, qtype: RecordType, path: &PathBuf) -> io::Result<Message> {
    let mut stream = UnixStream::connect(path).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

#[cfg(not(target_os = "android"))]
async fn local_lookup(qname: &Name, qtype: RecordType, server: &SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

async fn proxy_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    ns: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<Message> {
    let mut stream = ProxyStream::connect_proxied(context, svr_cfg, ns).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

async fn acl_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    #[cfg(target_os = "android")] local: &PathBuf,
    #[cfg(not(target_os = "android"))] local: &SocketAddr,
    remote: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<(Message, bool)> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        qtype, qname, local, remote
    );

    // remove the last dot from fqdn name
    let mut name = qname.to_ascii();
    name.pop();
    let addr = Address::DomainNameAddress(name, 0);
    let qname_in_proxy_list = context.check_qname_in_proxy_list(&addr);

    let remote_response_fut = async {
        match qname_in_proxy_list {
            Some(false) => None,
            _ => {
                let timeout = Some(Duration::new(3, 0));
                try_timeout(proxy_lookup(context.clone(), svr_cfg, remote, qname, qtype), timeout)
                    .await
                    .ok()
            }
        }
    };

    let local_response = match qname_in_proxy_list {
        Some(true) => None,
        _ => {
            let timeout = Some(Duration::new(3, 0));
            try_timeout(local_lookup(qname, qtype, local), timeout).await.ok()
        }
    }
    .unwrap_or_else(Message::new);

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    if local_response.answer_count() == 0 {
        let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
        return Ok((remote_response, true));
    }

    for rec in local_response.answers() {
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            _ => true,
        };
        if !forward {
            debug!("pick local response (ip): {:?}", local_response);
            return Ok((local_response, false));
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
    debug!("pick remote response (ip): {:?}", remote_response);
    Ok((remote_response, true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let mut qtx = qtx.clone();
        let server = balancer.pick_server();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let config = context.config();

                #[cfg(target_os = "android")]
                let local_addr = config.local_dns_path.as_ref().expect("local query DNS path");
                #[cfg(not(target_os = "android"))]
                let local_addr = config.local_dns_addr.as_ref().expect("local query DNS address");

                let remote_addr = config.remote_dns_addr.as_ref().expect("remote query DNS address");

                let question = &request.queries()[0];

                let qname = question.name();
                let qtype = question.query_type();
                let svr_cfg = server.server_config();

                let r = acl_lookup(context.clone(), svr_cfg, &local_addr, &remote_addr, qname, qtype).await;

                if let Ok((result, forward)) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Simplify interface to use Query<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::mpsc,
};

#[cfg(target_os = "android")]
use std::path::PathBuf;
#[cfg(target_os = "android")]
use tokio::net::UnixStream;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use byteorder::{BigEndian, ByteOrder};
use log::{debug, error, info};
use rand::Rng;
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::{ConfigType, ServerConfig},
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        socks5::Address,
        sys::create_udp_socket,
        tcprelay::ProxyStream,
        utils::try_timeout,
    },
};

async fn stream_lookup<T>(query: &Query, stream: &mut T) -> io::Result<Message>
where
    T: AsyncReadExt + AsyncWriteExt + Unpin,
{
    let mut message = Message::new();

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query.clone());

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut send_buffer = vec![0; size + 2];

    BigEndian::write_u16(&mut send_buffer[0..2], size as u16);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    let mut res_buffer = vec![0; 2];
    stream.read_exact(&mut res_buffer[0..2]).await?;

    let size = BigEndian::read_u16(&res_buffer[0..2]) as usize;
    let mut res_buffer = vec![0; size];
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

#[cfg(target_os = "android")]
async fn local_lookup(query: &Query, path: &PathBuf) -> io::Result<Message> {
    let mut stream = UnixStream::connect(path).await?;
    stream_lookup(query, &mut stream).await
}

#[cfg(not(target_os = "android"))]
async fn local_lookup(query: &Query, server: &SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query.clone());

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

async fn proxy_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    ns: &Address,
    query: &Query,
) -> io::Result<Message> {
    let mut stream = ProxyStream::connect_proxied(context, svr_cfg, ns).await?;
    stream_lookup(query, &mut stream).await
}

async fn acl_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    #[cfg(target_os = "android")] local: &PathBuf,
    #[cfg(not(target_os = "android"))] local: &SocketAddr,
    remote: &Address,
    query: &Query,
) -> io::Result<(Message, bool)> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    // remove the last dot from fqdn name
    let mut name = query.name().to_ascii();
    name.pop();
    let addr = Address::DomainNameAddress(name, 0);
    let qname_in_proxy_list = context.check_qname_in_proxy_list(&addr);

    let remote_response_fut = async {
        match qname_in_proxy_list {
            Some(false) => None,
            _ => {
                let timeout = Some(Duration::new(3, 0));
                try_timeout(proxy_lookup(context.clone(), svr_cfg, remote, query), timeout)
                    .await
                    .ok()
            }
        }
    };

    let local_response = match qname_in_proxy_list {
        Some(true) => None,
        _ => {
            let timeout = Some(Duration::new(3, 0));
            try_timeout(local_lookup(query, local), timeout).await.ok()
        }
    }
    .unwrap_or_else(Message::new);

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    if local_response.answer_count() == 0 {
        let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
        return Ok((remote_response, true));
    }

    for rec in local_response.answers() {
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            _ => true,
        };
        if !forward {
            debug!("pick local response (ip): {:?}", local_response);
            return Ok((local_response, false));
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
    debug!("pick remote response (ip): {:?}", remote_response);
    Ok((remote_response, true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let mut qtx = qtx.clone();
        let server = balancer.pick_server();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let config = context.config();

                #[cfg(target_os = "android")]
                let local_addr = config.local_dns_path.as_ref().expect("local query DNS path");
                #[cfg(not(target_os = "android"))]
                let local_addr = config.local_dns_addr.as_ref().expect("local query DNS address");

                let remote_addr = config.remote_dns_addr.as_ref().expect("remote query DNS address");

                let question = &request.queries()[0];

                let svr_cfg = server.server_config();

                let r = acl_lookup(context.clone(), svr_cfg, &local_addr, &remote_addr, question).await;

                if let Ok((result, forward)) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match context.check_query_in_proxy_list(query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (qname): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (qname): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;
    for rec in local_response.unwrap_or_else(|_| Message::new()).answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return (local_response, false);
        }
    }

    let remote_response = remote_response_fut.await;
    debug!("pick remote response (response): {:?}", remote_response);
    (remote_response, true)
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
    let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
    let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.clone(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Fix the build<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::RData,
};

use crate::{
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

async fn acl_lookup<Local, Remote>(
    context: SharedContext,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match context.check_query_in_proxy_list(query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (qname): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (qname): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await.unwrap_or_else(|_| Message::new());
    for rec in local_response.answers() {
        if rec.record_type() != query.query_type() {
            warn!("local DNS response has inconsistent answer type {} for query {}", rec.record_type(), query);
            break
        }
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::PTR(_) => panic!("PTR records should not reach here"),
            _ => context.is_default_in_proxy_list(),
        };
        if !forward {
            debug!("pick local response (response): {:?}", local_response);
            return (Ok(local_response), false);
        }
    }

    let remote_response = remote_response_fut.await;
    debug!("pick remote response (response): {:?}", remote_response);
    (remote_response, true)
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
    let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
    let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.clone(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>extern crate phf;

#[cfg(feature = "flate2")]
extern crate flate2;

use std::borrow::{Borrow, Cow};
use std::io::{self, BufReader, Cursor, Error, ErrorKind, Read};
use std::fs::File;
use std::sync::atomic::{AtomicBool, Ordering};

#[cfg(feature = "flate2")]
use flate2::bufread::GzDecoder;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Compression {
    None,
    Gzip
}

/// Runtime access to the included files
pub struct Files {
    // Do not access these fields, they are only public to allow for code generation!
    #[doc(hidden)]
    pub files: phf::Map<&'static str, (Compression, &'static [u8])>,
    #[doc(hidden)]
    pub passthrough: AtomicBool
}

#[cfg(windows)]
fn as_key(path: &str) -> Cow<str> {
    Cow::Owned(path.replace("\\", "/"))
}

#[cfg(not(windows))]
fn as_key(path: &str) -> Cow<str> {
    Cow::Borrowed(path)
}

impl Files {
    pub fn set_passthrough(&self, enabled: bool) {
        self.passthrough.store(enabled, Ordering::Relaxed);
    }

    pub fn is_available(&self, path: &str) -> bool {
        self.files.contains_key(path)
    }

    /// Returns an iterator over all available file names.  Does not
    /// decompress any compressed data.
    pub fn file_names(&'static self) -> FileNames {
        FileNames { iter: self.files.keys() }
    }

    pub fn get(&self, path: &str) -> io::Result<Cow<'static, [u8]>> {
        match self.get_raw(path) {
            Ok((Compression::None, data)) => Ok(data),
            #[cfg(feature = "flate2")]
            Ok((Compression::Gzip, compressed)) => {
                let mut r = GzDecoder::new(Cursor::new(compressed));
                let mut v = Vec::new();
                r.read_to_end(&mut v)?;
                Ok(Cow::Owned(v))
            },
            #[cfg(not(feature = "flate2"))]
            Ok((Compression::Gzip, _)) => panic!("Feature 'flate2' not enabled"),

            Err(e) => Err(e)
        }
    }

    pub fn get_raw(&self, path: &str) -> io::Result<(Compression, Cow<'static, [u8]>)> {
        if self.passthrough.load(Ordering::Relaxed) {
            let mut r = BufReader::new(File::open(path)?);
            let mut v = Vec::new();
            r.read_to_end(&mut v)?;
            return Ok((Compression::None, Cow::Owned(v)))
        }

        let key = as_key(path);

        self.files.get(&*key)
            .map(|&(c,d)| (c, Cow::Owned(d.to_owned())))
            .ok_or_else(|| Error::new(ErrorKind::NotFound, "Key not found"))
    }

    pub fn read(&self, path: &str) -> io::Result<Box<Read>> {
        if self.passthrough.load(Ordering::Relaxed) {
            return Ok(Box::new(BufReader::new(File::open(path)?)))
        }

        let key = as_key(path);
        match self.files.get(key.borrow() as &str) {
            Some(b) => {
                match b.0 {
                    Compression::None => Ok(Box::new(Cursor::new(b.1))),
                    #[cfg(feature = "flate2")]
                    Compression::Gzip => Ok(Box::new(GzDecoder::new(Cursor::new(b.1)))),
                    #[cfg(not(feature = "flate2"))]
                    Compression::Gzip => panic!("Feature 'flate2' not enabled"),
                }
            }
            None => Err(Error::new(ErrorKind::NotFound, "Key not found")),
        }
    }
}

/// Iterates over the file names available for `Files` object.
pub struct FileNames {
    // Our internal iterator.  We wrap this in a nice struct so our
    // caller doesn't need to know the details.
    iter: phf::map::Keys<'static, &'static str, (Compression, &'static [u8])>,
}

impl Iterator for FileNames {
    type Item = &'static str;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|s| *s)
    }
}
<commit_msg>Use .cloned() instead of manual cloning<commit_after>extern crate phf;

#[cfg(feature = "flate2")]
extern crate flate2;

use std::borrow::{Borrow, Cow};
use std::io::{self, BufReader, Cursor, Error, ErrorKind, Read};
use std::fs::File;
use std::sync::atomic::{AtomicBool, Ordering};

#[cfg(feature = "flate2")]
use flate2::bufread::GzDecoder;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Compression {
    None,
    Gzip
}

/// Runtime access to the included files
pub struct Files {
    // Do not access these fields, they are only public to allow for code generation!
    #[doc(hidden)]
    pub files: phf::Map<&'static str, (Compression, &'static [u8])>,
    #[doc(hidden)]
    pub passthrough: AtomicBool
}

#[cfg(windows)]
fn as_key(path: &str) -> Cow<str> {
    Cow::Owned(path.replace("\\", "/"))
}

#[cfg(not(windows))]
fn as_key(path: &str) -> Cow<str> {
    Cow::Borrowed(path)
}

impl Files {
    pub fn set_passthrough(&self, enabled: bool) {
        self.passthrough.store(enabled, Ordering::Relaxed);
    }

    pub fn is_available(&self, path: &str) -> bool {
        self.files.contains_key(path)
    }

    /// Returns an iterator over all available file names.  Does not
    /// decompress any compressed data.
    pub fn file_names(&'static self) -> FileNames {
        FileNames { iter: self.files.keys() }
    }

    pub fn get(&self, path: &str) -> io::Result<Cow<'static, [u8]>> {
        match self.get_raw(path) {
            Ok((Compression::None, data)) => Ok(data),
            #[cfg(feature = "flate2")]
            Ok((Compression::Gzip, compressed)) => {
                let mut r = GzDecoder::new(Cursor::new(compressed));
                let mut v = Vec::new();
                r.read_to_end(&mut v)?;
                Ok(Cow::Owned(v))
            },
            #[cfg(not(feature = "flate2"))]
            Ok((Compression::Gzip, _)) => panic!("Feature 'flate2' not enabled"),

            Err(e) => Err(e)
        }
    }

    pub fn get_raw(&self, path: &str) -> io::Result<(Compression, Cow<'static, [u8]>)> {
        if self.passthrough.load(Ordering::Relaxed) {
            let mut r = BufReader::new(File::open(path)?);
            let mut v = Vec::new();
            r.read_to_end(&mut v)?;
            return Ok((Compression::None, Cow::Owned(v)))
        }

        let key = as_key(path);

        self.files.get(&*key)
            .map(|&(c,d)| (c, Cow::Owned(d.to_owned())))
            .ok_or_else(|| Error::new(ErrorKind::NotFound, "Key not found"))
    }

    pub fn read(&self, path: &str) -> io::Result<Box<Read>> {
        if self.passthrough.load(Ordering::Relaxed) {
            return Ok(Box::new(BufReader::new(File::open(path)?)))
        }

        let key = as_key(path);
        match self.files.get(key.borrow() as &str) {
            Some(b) => {
                match b.0 {
                    Compression::None => Ok(Box::new(Cursor::new(b.1))),
                    #[cfg(feature = "flate2")]
                    Compression::Gzip => Ok(Box::new(GzDecoder::new(Cursor::new(b.1)))),
                    #[cfg(not(feature = "flate2"))]
                    Compression::Gzip => panic!("Feature 'flate2' not enabled"),
                }
            }
            None => Err(Error::new(ErrorKind::NotFound, "Key not found")),
        }
    }
}

/// Iterates over the file names available for `Files` object.
pub struct FileNames {
    // Our internal iterator.  We wrap this in a nice struct so our
    // caller doesn't need to know the details.
    iter: phf::map::Keys<'static, &'static str, (Compression, &'static [u8])>,
}

impl Iterator for FileNames {
    type Item = &'static str;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().cloned()
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of the install aspects of the compiler.
//!
//! This module is responsible for installing the standard library,
//! compiler, and documentation.

use std::env;
use std::fs;
use std::path::{Path, PathBuf, Component};
use std::process::Command;

use dist::{self, pkgname, sanitize_sh, tmpdir};

use Build;
use builder::{Builder, Step};

pub struct Installer<'a> {
    build: &'a Build,
    prefix: PathBuf,
    sysconfdir: PathBuf,
    docdir: PathBuf,
    bindir: PathBuf,
    libdir: PathBuf,
    mandir: PathBuf,
    empty_dir: PathBuf,
}

impl<'a> Drop for Installer<'a> {
    fn drop(&mut self) {
        t!(fs::remove_dir_all(&self.empty_dir));
    }
}

impl<'a> Installer<'a> {
    pub fn new(build: &'a Build) -> Installer<'a> {
        let prefix_default = PathBuf::from("/usr/local");
        let sysconfdir_default = PathBuf::from("/etc");
        let docdir_default = PathBuf::from("share/doc/rust");
        let bindir_default = PathBuf::from("bin");
        let libdir_default = PathBuf::from("lib");
        let mandir_default = PathBuf::from("share/man");
        let prefix = build.config.prefix.as_ref().unwrap_or(&prefix_default);
        let sysconfdir = build.config.sysconfdir.as_ref().unwrap_or(&sysconfdir_default);
        let docdir = build.config.docdir.as_ref().unwrap_or(&docdir_default);
        let bindir = build.config.bindir.as_ref().unwrap_or(&bindir_default);
        let libdir = build.config.libdir.as_ref().unwrap_or(&libdir_default);
        let mandir = build.config.mandir.as_ref().unwrap_or(&mandir_default);

        let sysconfdir = prefix.join(sysconfdir);
        let docdir = prefix.join(docdir);
        let bindir = prefix.join(bindir);
        let libdir = prefix.join(libdir);
        let mandir = prefix.join(mandir);

        let destdir = env::var_os("DESTDIR").map(PathBuf::from);

        let prefix = add_destdir(&prefix, &destdir);
        let sysconfdir = add_destdir(&sysconfdir, &destdir);
        let docdir = add_destdir(&docdir, &destdir);
        let bindir = add_destdir(&bindir, &destdir);
        let libdir = add_destdir(&libdir, &destdir);
        let mandir = add_destdir(&mandir, &destdir);

        let empty_dir = build.out.join("tmp/empty_dir");

        t!(fs::create_dir_all(&empty_dir));

        Installer {
            build,
            prefix,
            sysconfdir,
            docdir,
            bindir,
            libdir,
            mandir,
            empty_dir,
        }
    }

    pub fn install_docs(&self, stage: u32, host: &str) {
        self.install_sh("docs", "rust-docs", stage, Some(host));
    }

    pub fn install_std(&self, stage: u32) {
        for target in self.build.config.target.iter() {
            self.install_sh("std", "rust-std", stage, Some(target));
        }
    }

    pub fn install_cargo(&self, stage: u32, host: &str) {
        self.install_sh("cargo", "cargo", stage, Some(host));
    }

    pub fn install_rls(&self, stage: u32, host: &str) {
        self.install_sh("rls", "rls", stage, Some(host));
    }

    pub fn install_analysis(&self, stage: u32, host: &str) {
        self.install_sh("analysis", "rust-analysis", stage, Some(host));
    }

    pub fn install_src(&self, stage: u32) {
        self.install_sh("src", "rust-src", stage, None);
    }
    pub fn install_rustc(&self, stage: u32, host: &str) {
        self.install_sh("rustc", "rustc", stage, Some(host));
    }

    fn install_sh(&self, package: &str, name: &str, stage: u32, host: Option<&str>) {
        println!("Install {} stage{} ({:?})", package, stage, host);
        let package_name = if let Some(host) = host {
            format!("{}-{}", pkgname(self.build, name), host)
        } else {
            pkgname(self.build, name)
        };

        let mut cmd = Command::new("sh");
        cmd.current_dir(&self.empty_dir)
           .arg(sanitize_sh(&tmpdir(self.build).join(&package_name).join("install.sh")))
           .arg(format!("--prefix={}", sanitize_sh(&self.prefix)))
           .arg(format!("--sysconfdir={}", sanitize_sh(&self.sysconfdir)))
           .arg(format!("--docdir={}", sanitize_sh(&self.docdir)))
           .arg(format!("--bindir={}", sanitize_sh(&self.bindir)))
           .arg(format!("--libdir={}", sanitize_sh(&self.libdir)))
           .arg(format!("--mandir={}", sanitize_sh(&self.mandir)))
           .arg("--disable-ldconfig");
        self.build.run(&mut cmd);
    }
}

fn add_destdir(path: &Path, destdir: &Option<PathBuf>) -> PathBuf {
    let mut ret = match *destdir {
        Some(ref dest) => dest.clone(),
        None => return path.to_path_buf(),
    };
    for part in path.components() {
        match part {
            Component::Normal(s) => ret.push(s),
            _ => {}
        }
    }
    ret
}

macro_rules! install {
    ($($name:ident,
       $path:expr,
       $default_cond:expr,
       only_hosts: $only_hosts:expr,
       ($sel:ident, $builder:ident),
       $run_item:block $(, $c:ident)*;)+) => {
        $(#[derive(Serialize)]
        pub struct $name<'a> {
            pub stage: u32,
            pub target: &'a str,
            pub host: &'a str,
        }

        impl<'a> Step<'a> for $name<'a> {
            type Output = ();
            const DEFAULT: bool = true;
            const ONLY_BUILD_TARGETS: bool = true;
            const ONLY_HOSTS: bool = $only_hosts;
            $(const $c: bool = true;)*

            fn should_run(_builder: &Builder, path: &Path) -> bool {
                path.ends_with($path)
            }

            fn make_run($builder: &Builder, path: Option<&Path>, host: &str, target: &str) {
                if path.is_none() && !($default_cond) {
                    return;
                }
                $builder.ensure($name {
                    stage: $builder.top_stage,
                    target,
                    host,
                });
            }

            fn run($sel, $builder: &Builder) {
                $run_item
            }
        })+
    }
}

install!(
    // rules.install("install-docs", "src/doc")
    //      .default(build.config.docs)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-docs"))
    //      .run(move |s| install::Installer::new(build).install_docs(s.stage, s.target));
    Docs, "src/doc", builder.build.config.docs, only_hosts: false, (self, builder), {
        builder.ensure(dist::Docs { stage: self.stage, target: self.target });
        Installer::new(builder.build).install_docs(self.stage, self.target);
    };
    // rules.install("install-std", "src/libstd")
    //      .default(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-std"))
    //      .run(move |s| install::Installer::new(build).install_std(s.stage));
    Std, "src/libstd", true, only_hosts: true, (self, builder), {
        builder.ensure(dist::Std {
            compiler: builder.compiler(self.stage, self.host),
            target: self.target
        });
        Installer::new(builder.build).install_std(self.stage);
    };
    // rules.install("install-cargo", "cargo")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-cargo"))
    //      .run(move |s| install::Installer::new(build).install_cargo(s.stage, s.target));
    Cargo, "cargo", builder.build.config.extended, only_hosts: true, (self, builder), {
        builder.ensure(dist::Cargo { stage: self.stage, target: self.target });
        Installer::new(builder.build).install_cargo(self.stage, self.target);
    };
    // rules.install("install-rls", "rls")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-rls"))
    //      .run(move |s| install::Installer::new(build).install_rls(s.stage, s.target));
    Rls, "rls", builder.build.config.extended, only_hosts: true, (self, builder), {
        builder.ensure(dist::Rls { stage: self.stage, target: self.target });
        Installer::new(builder.build).install_rls(self.stage, self.target);
    };
    // rules.install("install-analysis", "analysis")
    //      .default(build.config.extended)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-analysis"))
    //      .run(move |s| install::Installer::new(build).install_analysis(s.stage, s.target));
    Analysis, "analysis", builder.build.config.extended, only_hosts: false, (self, builder), {
        builder.ensure(dist::Analysis {
            compiler: builder.compiler(self.stage, self.host),
            target: self.target
        });
        Installer::new(builder.build).install_analysis(self.stage, self.target);
    };
    // rules.install("install-src", "src")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_build(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-src"))
    //      .run(move |s| install::Installer::new(build).install_src(s.stage));
    Src, "src", builder.build.config.extended, only_hosts: true, (self, builder), {
        builder.ensure(dist::Src);
        Installer::new(builder.build).install_src(self.stage);
    }, ONLY_BUILD;
    // rules.install("install-rustc", "src/librustc")
    //      .default(true)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-rustc"))
    //      .run(move |s| install::Installer::new(build).install_rustc(s.stage, s.target));
    Rustc, "src/librustc", builder.build.config.extended, only_hosts: true, (self, builder), {
        builder.ensure(dist::Rustc { stage: self.stage, target: self.target });
        Installer::new(builder.build).install_rustc(self.stage, self.target);
    };
);
<commit_msg>Clean up install<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of the install aspects of the compiler.
//!
//! This module is responsible for installing the standard library,
//! compiler, and documentation.

use std::env;
use std::fs;
use std::path::{Path, PathBuf, Component};
use std::process::Command;

use dist::{self, pkgname, sanitize_sh, tmpdir};

use builder::{Builder, Step};

pub fn install_docs(builder: &Builder, stage: u32, host: &str) {
    install_sh(builder, "docs", "rust-docs", stage, Some(host));
}

pub fn install_std(builder: &Builder, stage: u32) {
    for target in builder.build.config.target.iter() {
        install_sh(builder, "std", "rust-std", stage, Some(target));
    }
}

pub fn install_cargo(builder: &Builder, stage: u32, host: &str) {
    install_sh(builder, "cargo", "cargo", stage, Some(host));
}

pub fn install_rls(builder: &Builder, stage: u32, host: &str) {
    install_sh(builder, "rls", "rls", stage, Some(host));
}

pub fn install_analysis(builder: &Builder, stage: u32, host: &str) {
    install_sh(builder, "analysis", "rust-analysis", stage, Some(host));
}

pub fn install_src(builder: &Builder, stage: u32) {
    install_sh(builder, "src", "rust-src", stage, None);
}
pub fn install_rustc(builder: &Builder, stage: u32, host: &str) {
    install_sh(builder, "rustc", "rustc", stage, Some(host));
}

fn install_sh(builder: &Builder, package: &str, name: &str, stage: u32, host: Option<&str>) {
    let build = builder.build;
    println!("Install {} stage{} ({:?})", package, stage, host);

    let prefix_default = PathBuf::from("/usr/local");
    let sysconfdir_default = PathBuf::from("/etc");
    let docdir_default = PathBuf::from("share/doc/rust");
    let bindir_default = PathBuf::from("bin");
    let libdir_default = PathBuf::from("lib");
    let mandir_default = PathBuf::from("share/man");
    let prefix = build.config.prefix.as_ref().unwrap_or(&prefix_default);
    let sysconfdir = build.config.sysconfdir.as_ref().unwrap_or(&sysconfdir_default);
    let docdir = build.config.docdir.as_ref().unwrap_or(&docdir_default);
    let bindir = build.config.bindir.as_ref().unwrap_or(&bindir_default);
    let libdir = build.config.libdir.as_ref().unwrap_or(&libdir_default);
    let mandir = build.config.mandir.as_ref().unwrap_or(&mandir_default);

    let sysconfdir = prefix.join(sysconfdir);
    let docdir = prefix.join(docdir);
    let bindir = prefix.join(bindir);
    let libdir = prefix.join(libdir);
    let mandir = prefix.join(mandir);

    let destdir = env::var_os("DESTDIR").map(PathBuf::from);

    let prefix = add_destdir(&prefix, &destdir);
    let sysconfdir = add_destdir(&sysconfdir, &destdir);
    let docdir = add_destdir(&docdir, &destdir);
    let bindir = add_destdir(&bindir, &destdir);
    let libdir = add_destdir(&libdir, &destdir);
    let mandir = add_destdir(&mandir, &destdir);

    let empty_dir = build.out.join("tmp/empty_dir");

    t!(fs::create_dir_all(&empty_dir));
    let package_name = if let Some(host) = host {
        format!("{}-{}", pkgname(build, name), host)
    } else {
        pkgname(build, name)
    };

    let mut cmd = Command::new("sh");
    cmd.current_dir(&empty_dir)
        .arg(sanitize_sh(&tmpdir(build).join(&package_name).join("install.sh")))
        .arg(format!("--prefix={}", sanitize_sh(&prefix)))
        .arg(format!("--sysconfdir={}", sanitize_sh(&sysconfdir)))
        .arg(format!("--docdir={}", sanitize_sh(&docdir)))
        .arg(format!("--bindir={}", sanitize_sh(&bindir)))
        .arg(format!("--libdir={}", sanitize_sh(&libdir)))
        .arg(format!("--mandir={}", sanitize_sh(&mandir)))
        .arg("--disable-ldconfig");
    build.run(&mut cmd);
    t!(fs::remove_dir_all(&empty_dir));
}

fn add_destdir(path: &Path, destdir: &Option<PathBuf>) -> PathBuf {
    let mut ret = match *destdir {
        Some(ref dest) => dest.clone(),
        None => return path.to_path_buf(),
    };
    for part in path.components() {
        match part {
            Component::Normal(s) => ret.push(s),
            _ => {}
        }
    }
    ret
}

macro_rules! install {
    (($sel:ident, $builder:ident),
       $($name:ident,
       $path:expr,
       $default_cond:expr,
       only_hosts: $only_hosts:expr,
       $run_item:block $(, $c:ident)*;)+) => {
        $(#[derive(Serialize)]
        pub struct $name<'a> {
            pub stage: u32,
            pub target: &'a str,
            pub host: &'a str,
        }

        impl<'a> Step<'a> for $name<'a> {
            type Output = ();
            const DEFAULT: bool = true;
            const ONLY_BUILD_TARGETS: bool = true;
            const ONLY_HOSTS: bool = $only_hosts;
            $(const $c: bool = true;)*

            fn should_run(_builder: &Builder, path: &Path) -> bool {
                path.ends_with($path)
            }

            fn make_run($builder: &Builder, path: Option<&Path>, host: &str, target: &str) {
                if path.is_none() && !($default_cond) {
                    return;
                }
                $builder.ensure($name {
                    stage: $builder.top_stage,
                    target,
                    host,
                });
            }

            fn run($sel, $builder: &Builder) {
                $run_item
            }
        })+
    }
}

install!((self, builder),
    // rules.install("install-docs", "src/doc")
    //      .default(build.config.docs)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-docs"))
    //      .run(move |s| install::Installer::new(build).install_docs(s.stage, s.target));
    Docs, "src/doc", builder.build.config.docs, only_hosts: false, {
        builder.ensure(dist::Docs { stage: self.stage, target: self.target });
        install_docs(builder, self.stage, self.target);
    };
    // rules.install("install-std", "src/libstd")
    //      .default(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-std"))
    //      .run(move |s| install::Installer::new(build).install_std(s.stage));
    Std, "src/libstd", true, only_hosts: true, {
        builder.ensure(dist::Std {
            compiler: builder.compiler(self.stage, self.host),
            target: self.target
        });
        install_std(builder, self.stage);
    };
    // rules.install("install-cargo", "cargo")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-cargo"))
    //      .run(move |s| install::Installer::new(build).install_cargo(s.stage, s.target));
    Cargo, "cargo", builder.build.config.extended, only_hosts: true, {
        builder.ensure(dist::Cargo { stage: self.stage, target: self.target });
        install_cargo(builder, self.stage, self.target);
    };
    // rules.install("install-rls", "rls")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-rls"))
    //      .run(move |s| install::Installer::new(build).install_rls(s.stage, s.target));
    Rls, "rls", builder.build.config.extended, only_hosts: true, {
        builder.ensure(dist::Rls { stage: self.stage, target: self.target });
        install_rls(builder, self.stage, self.target);
    };
    // rules.install("install-analysis", "analysis")
    //      .default(build.config.extended)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-analysis"))
    //      .run(move |s| install::Installer::new(build).install_analysis(s.stage, s.target));
    Analysis, "analysis", builder.build.config.extended, only_hosts: false, {
        builder.ensure(dist::Analysis {
            compiler: builder.compiler(self.stage, self.host),
            target: self.target
        });
        install_analysis(builder, self.stage, self.target);
    };
    // rules.install("install-src", "src")
    //      .default(build.config.extended)
    //      .host(true)
    //      .only_build(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-src"))
    //      .run(move |s| install::Installer::new(build).install_src(s.stage));
    Src, "src", builder.build.config.extended, only_hosts: true, {
        builder.ensure(dist::Src);
        install_src(builder, self.stage);
    }, ONLY_BUILD;
    // rules.install("install-rustc", "src/librustc")
    //      .default(true)
    //      .host(true)
    //      .only_host_build(true)
    //      .dep(|s| s.name("dist-rustc"))
    //      .run(move |s| install::Installer::new(build).install_rustc(s.stage, s.target));
    Rustc, "src/librustc", builder.build.config.extended, only_hosts: true, {
        builder.ensure(dist::Rustc { stage: self.stage, target: self.target });
        install_rustc(builder, self.stage, self.target);
    };
);
<|endoftext|>


<commit_before>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc_plugin;

use syntax::abi;
use syntax::ast;
use syntax::ast::Ty_::TyBareFn;
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, MultiModifier, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc_plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  MultiModifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          annot_item: Annotatable) -> Annotatable {
    let item = annot_item.expect_item();
    match item.node {
        ast::ItemFn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyBareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ast::ItemStatic(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    Annotatable::Item(item)
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> Annotatable {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, ast::DeclItem(prop.clone())));
    let inner_fn =
        P(codemap::respan(span, ast::StmtDecl(fn_decl, ast::DUMMY_NODE_ID)));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    Annotatable::Item(P(ast::Item {attrs: attrs, ..(*test).clone()}))
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::DefaultReturn(span),
        variadic: false,
    });
    let item = ast::ItemFn(decl,
                           ast::Unsafety::Normal,
                           ast::Constness::NotConst,
                           abi::Rust,
                           ast::Generics::default(),
                           body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<commit_msg>Update for Rust nightly 2016-02-12.<commit_after>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc_plugin;

use syntax::abi::Abi;
use syntax::ast;
use syntax::ast::{DeclKind, ItemKind, StmtKind, TyKind};
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, MultiModifier, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc_plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  MultiModifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          annot_item: Annotatable) -> Annotatable {
    let item = annot_item.expect_item();
    match item.node {
        ItemKind::Fn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyKind::BareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ItemKind::Static(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    Annotatable::Item(item)
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> Annotatable {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, DeclKind::Item(prop.clone())));
    let inner_fn = codemap::respan(
        span, StmtKind::Decl(fn_decl, ast::DUMMY_NODE_ID));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    Annotatable::Item(P(ast::Item {attrs: attrs, ..(*test).clone()}))
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::Default(span),
        variadic: false,
    });
    let item = ItemKind::Fn(
        decl, ast::Unsafety::Normal, ast::Constness::NotConst, Abi::Rust,
        ast::Generics::default(), body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<|endoftext|>


<commit_before>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc;

use syntax::abi;
use syntax::ast;
use syntax::ast::Ty_::TyBareFn;
use syntax::ast_util;
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, Modifier};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc::plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  Modifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          item: P<ast::Item>) -> P<ast::Item> {
    match item.node {
        ast::ItemFn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyBareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ast::ItemStatic(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    item
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> P<ast::Item> {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, ast::DeclItem(prop.clone())));
    let inner_fn =
        P(codemap::respan(span, ast::StmtDecl(fn_decl, ast::DUMMY_NODE_ID)));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    P(ast::Item {attrs: attrs, ..(*test).clone()})
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::DefaultReturn(span),
        variadic: false,
    });
    let item = ast::ItemFn(decl,
                           ast::Unsafety::Normal,
                           ast::Constness::NotConst,
                           abi::Rust,
                           ast_util::empty_generics(),
                           body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<commit_msg>Update to use MultiModifier<commit_after>//! This crate provides the `#[quickcheck]` attribute. Its use is
//! documented in the `quickcheck` crate.

#![crate_name = "quickcheck_macros"]
#![crate_type = "dylib"]
#![doc(html_root_url = "http://burntsushi.net/rustdoc/quickcheck")]

#![feature(plugin_registrar, rustc_private)]

extern crate syntax;
extern crate rustc;

use syntax::abi;
use syntax::ast;
use syntax::ast::Ty_::TyBareFn;
use syntax::ast_util;
use syntax::codemap;
use syntax::parse::token;
use syntax::ext::base::{ExtCtxt, MultiModifier, Annotatable};
use syntax::ext::build::AstBuilder;
use syntax::ptr::P;

use rustc::plugin::Registry;

/// For the `#[quickcheck]` attribute. Do not use.
#[plugin_registrar]
#[doc(hidden)]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(token::intern("quickcheck"),
                                  MultiModifier(Box::new(expand_meta_quickcheck)));
}

/// Expands the `#[quickcheck]` attribute.
///
/// Expands:
/// ```
/// #[quickcheck]
/// fn check_something(_: usize) -> bool {
///     true
/// }
/// ```
/// to:
/// ```
/// #[test]
/// fn check_something() {
///     fn check_something(_: usize) -> bool {
///         true
///     }
///     ::quickcheck::quickcheck(check_something as fn(usize) -> bool)
/// }
/// ```
fn expand_meta_quickcheck(cx: &mut ExtCtxt,
                          span: codemap::Span,
                          _: &ast::MetaItem,
                          annot_item: Annotatable) -> Annotatable {
    let item = annot_item.expect_item();
    match item.node {
        ast::ItemFn(ref decl, unsafety, _constness, abi, _, _) => {
            let prop_ident = cx.expr_ident(span, item.ident);
            let prop_ty = cx.ty(span, TyBareFn(P(ast::BareFnTy {
                unsafety: unsafety,
                abi: abi,
                lifetimes: vec![],
                decl: decl.clone(),
            })));
            let inner_ident = cx.expr_cast(span, prop_ident, prop_ty);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        ast::ItemStatic(..) => {
            let inner_ident = cx.expr_ident(span, item.ident);
            return wrap_item(cx, span, &*item, inner_ident);
        },
        _ => {
            cx.span_err(
                span, "#[quickcheck] only supported on statics and functions");
        }
    }
    Annotatable::Item(item)
}

fn wrap_item(cx: &mut ExtCtxt,
             span: codemap::Span,
             item: &ast::Item,
             inner_ident: P<ast::Expr>) -> Annotatable {
    // Copy original function without attributes
    let prop = P(ast::Item {attrs: Vec::new(), ..item.clone()});
    // ::quickcheck::quickcheck
    let check_ident = token::str_to_ident("quickcheck");
    let check_path = vec!(check_ident, check_ident);
    // Wrap original function in new outer function,
    // calling ::quickcheck::quickcheck()
    let fn_decl = P(codemap::respan(span, ast::DeclItem(prop.clone())));
    let inner_fn =
        P(codemap::respan(span, ast::StmtDecl(fn_decl, ast::DUMMY_NODE_ID)));
    let check_call = cx.expr_call_global(span, check_path, vec![inner_ident]);
    let body = cx.block(span, vec![inner_fn], Some(check_call));
    let test = item_fn(cx, span, item, body);

    // Copy attributes from original function
    let mut attrs = item.attrs.clone();
    // Add #[test] attribute
    attrs.push(cx.attribute(
        span, cx.meta_word(span, token::intern_and_get_ident("test"))));
    // Attach the attributes to the outer function
    Annotatable::Item(P(ast::Item {attrs: attrs, ..(*test).clone()}))
}

fn item_fn(cx: &mut ExtCtxt, span: codemap::Span,
           towrap_item: &ast::Item, body: P<ast::Block>) -> P<ast::Item> {
    let decl = P(ast::FnDecl {
        inputs: vec![],
        output: ast::FunctionRetTy::DefaultReturn(span),
        variadic: false,
    });
    let item = ast::ItemFn(decl,
                           ast::Unsafety::Normal,
                           ast::Constness::NotConst,
                           abi::Rust,
                           ast_util::empty_generics(),
                           body);
    cx.item(span, towrap_item.ident, vec![], item)
}
<|endoftext|>