<commit_before>//! This handles the various supported encoding mechanism for the Eth 2.0 RPC.

use crate::rpc::methods::ErrorType;
use crate::rpc::{RPCCodedResponse, RPCRequest, RPCResponse};
use libp2p::bytes::BufMut;
use libp2p::bytes::BytesMut;
use std::marker::PhantomData;
use tokio_util::codec::{Decoder, Encoder};
use types::EthSpec;

pub trait OutboundCodec<TItem>: Encoder<TItem> + Decoder {
    type CodecErrorType;

    fn decode_error(
        &mut self,
        src: &mut BytesMut,
    ) -> Result<Option<Self::CodecErrorType>, <Self as Decoder>::Error>;
}

/* Global Inbound Codec */
// This deals with Decoding RPC Requests from other peers and encoding our responses

pub struct BaseInboundCodec<TCodec, TSpec>
where
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder,
    TSpec: EthSpec,
{
    /// Inner codec for handling various encodings
    inner: TCodec,
    phantom: PhantomData<TSpec>,
}

impl<TCodec, TSpec> BaseInboundCodec<TCodec, TSpec>
where
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder,
    TSpec: EthSpec,
{
    pub fn new(codec: TCodec) -> Self {
        BaseInboundCodec {
            inner: codec,
            phantom: PhantomData,
        }
    }
}

/* Global Outbound Codec */
// This deals with Decoding RPC Responses from other peers and encoding our requests
pub struct BaseOutboundCodec<TOutboundCodec, TSpec>
where
    TOutboundCodec: OutboundCodec<RPCRequest<TSpec>>,
    TSpec: EthSpec,
{
    /// Inner codec for handling various encodings.
    inner: TOutboundCodec,
    /// Keeps track of the current response code for a chunk.
    current_response_code: Option<u8>,
    phantom: PhantomData<TSpec>,
}

impl<TOutboundCodec, TSpec> BaseOutboundCodec<TOutboundCodec, TSpec>
where
    TSpec: EthSpec,
    TOutboundCodec: OutboundCodec<RPCRequest<TSpec>>,
{
    pub fn new(codec: TOutboundCodec) -> Self {
        BaseOutboundCodec {
            inner: codec,
            current_response_code: None,
            phantom: PhantomData,
        }
    }
}

/* Implementation of the Encoding/Decoding for the global codecs */

/* Base Inbound Codec */

// This Encodes RPC Responses sent to external peers
impl<TCodec, TSpec> Encoder<RPCCodedResponse<TSpec>> for BaseInboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: Decoder + Encoder<RPCCodedResponse<TSpec>>,
{
    type Error = <TCodec as Encoder<RPCCodedResponse<TSpec>>>::Error;

    fn encode(
        &mut self,
        item: RPCCodedResponse<TSpec>,
        dst: &mut BytesMut,
    ) -> Result<(), Self::Error> {
        dst.clear();
        dst.reserve(1);
        dst.put_u8(
            item.as_u8()
                .expect("Should never encode a stream termination"),
        );
        self.inner.encode(item, dst)
    }
}

// This Decodes RPC Requests from external peers
impl<TCodec, TSpec> Decoder for BaseInboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder<Item = RPCRequest<TSpec>>,
{
    type Item = RPCRequest<TSpec>;
    type Error = <TCodec as Decoder>::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        self.inner.decode(src)
    }
}

/* Base Outbound Codec */

// This Encodes RPC Requests sent to external peers
impl<TCodec, TSpec> Encoder<RPCRequest<TSpec>> for BaseOutboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: OutboundCodec<RPCRequest<TSpec>> + Encoder<RPCRequest<TSpec>>,
{
    type Error = <TCodec as Encoder<RPCRequest<TSpec>>>::Error;

    fn encode(&mut self, item: RPCRequest<TSpec>, dst: &mut BytesMut) -> Result<(), Self::Error> {
        self.inner.encode(item, dst)
    }
}

// This decodes RPC Responses received from external peers
impl<TCodec, TSpec> Decoder for BaseOutboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: OutboundCodec<RPCRequest<TSpec>, CodecErrorType = ErrorType>
        + Decoder<Item = RPCResponse<TSpec>>,
{
    type Item = RPCCodedResponse<TSpec>;
    type Error = <TCodec as Decoder>::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        // if we have only received the response code, wait for more bytes
        if src.len() <= 1 {
            return Ok(None);
        }
        // using the response code determine which kind of payload needs to be decoded.
        let response_code = self.current_response_code.unwrap_or_else(|| {
            let resp_code = src.split_to(1)[0];
            self.current_response_code = Some(resp_code);
            resp_code
        });

        let inner_result = {
            if RPCCodedResponse::<TSpec>::is_response(response_code) {
                // decode an actual response and mutates the buffer if enough bytes have been read
                // returning the result.
                self.inner
                    .decode(src)
                    .map(|r| r.map(RPCCodedResponse::Success))
            } else {
                // decode an error
                self.inner
                    .decode_error(src)
                    .map(|r| r.map(|resp| RPCCodedResponse::from_error(response_code, resp)))
            }
        };
        // if the inner decoder was capable of decoding a chunk, we need to reset the current
        // response code for the next chunk
        if let Ok(Some(_)) = inner_result {
            self.current_response_code = None;
        }
        // return the result
        inner_result
    }
}

#[cfg(test)]
mod tests {
    use super::super::ssz_snappy::*;
    use super::*;
    use crate::rpc::methods::StatusMessage;
    use crate::rpc::protocol::*;
    use snap::write::FrameEncoder;
    use ssz::Encode;
    use std::io::Write;
    use types::{Epoch, Hash256, Slot};
    use unsigned_varint::codec::Uvi;

    type Spec = types::MainnetEthSpec;

    #[test]
    fn test_decode_status_message() {
        let message = hex::decode("0054ff060000734e615070590032000006e71e7b54989925efd6c9cbcb8ceb9b5f71216f5137282bf6a1e3b50f64e42d6c7fb347abe07eb0db8200000005029e2800").unwrap();
        let mut buf = BytesMut::new();
        buf.extend_from_slice(&message);

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);

        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        // remove response code
        let mut snappy_buf = buf.clone();
        let _ = snappy_buf.split_to(1);

        // decode message just as snappy message
        let snappy_decoded_message = snappy_outbound_codec.decode(&mut snappy_buf).unwrap();

        // build codecs for entire chunk
        let mut snappy_base_outbound_codec = BaseOutboundCodec::new(snappy_outbound_codec);

        // decode message as ssz snappy chunk
        let snappy_decoded_chunk = snappy_base_outbound_codec.decode(&mut buf).unwrap();

        dbg!(snappy_decoded_message);
        dbg!(snappy_decoded_chunk);
    }

    #[test]
    fn test_invalid_length_prefix() {
        let mut uvi_codec: Uvi<u128> = Uvi::default();
        let mut dst = BytesMut::with_capacity(1024);

        // Smallest > 10 byte varint
        let len: u128 = 2u128.pow(70);

        // Insert length-prefix
        uvi_codec.encode(len, &mut dst).unwrap();

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);
        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        let snappy_decoded_message = snappy_outbound_codec.decode(&mut dst).unwrap_err();

        assert_eq!(
            snappy_decoded_message,
            RPCError::IoError("input bytes exceed maximum".to_string()),
            "length-prefix of > 10 bytes is invalid"
        );
    }

    #[test]
    fn test_length_limits() {
        fn encode_len(len: usize) -> BytesMut {
            let mut uvi_codec: Uvi<usize> = Uvi::default();
            let mut dst = BytesMut::with_capacity(1024);
            uvi_codec.encode(len, &mut dst).unwrap();
            dst
        }

        let protocol_id =
            ProtocolId::new(Protocol::BlocksByRange, Version::V1, Encoding::SSZSnappy);

        // Response limits
        let limit = protocol_id.rpc_response_limits::<Spec>();
        let mut max = encode_len(limit.max + 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut max).unwrap_err(), RPCError::InvalidData);

        let mut min = encode_len(limit.min - 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut min).unwrap_err(), RPCError::InvalidData);

        // Request limits
        let limit = protocol_id.rpc_request_limits();
        let mut max = encode_len(limit.max + 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut max).unwrap_err(), RPCError::InvalidData);

        let mut min = encode_len(limit.min - 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id, 1_048_576);
        assert_eq!(codec.decode(&mut min).unwrap_err(), RPCError::InvalidData);
    }

    #[test]
    fn test_decode_malicious_status_message() {
        // Snappy stream identifier
        let stream_identifier: &'static [u8] = b"\xFF\x06\x00\x00sNaPpY";

        // byte 0(0xFE) is padding chunk type identifier for snappy messages
        // byte 1,2,3 are chunk length (little endian)
        let malicious_padding: &'static [u8] = b"\xFE\x00\x00\x00";

        // Status message is 84 bytes uncompressed. `max_compressed_len` is 130.
        let status_message_bytes = StatusMessage {
            fork_digest: [0; 4],
            finalized_root: Hash256::from_low_u64_be(0),
            finalized_epoch: Epoch::new(1),
            head_root: Hash256::from_low_u64_be(0),
            head_slot: Slot::new(1),
        }
        .as_ssz_bytes();

        let mut uvi_codec: Uvi<usize> = Uvi::default();
        let mut dst = BytesMut::with_capacity(1024);

        // Insert length-prefix
        uvi_codec
            .encode(status_message_bytes.len(), &mut dst)
            .unwrap();

        // Insert snappy stream identifier
        dst.extend_from_slice(stream_identifier);

        // Insert malicious padding of 80 bytes.
        for _ in 0..20 {
            dst.extend_from_slice(malicious_padding);
        }

        // Insert payload (42 bytes compressed)
        let mut writer = FrameEncoder::new(Vec::new());
        writer.write_all(&status_message_bytes).unwrap();
        writer.flush().unwrap();
        dst.extend_from_slice(writer.get_ref());

        // 42 + 80 = 132 > max_compressed_len. Hence, decoding should fail with `InvalidData`.

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);

        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        let snappy_decoded_message = snappy_outbound_codec.decode(&mut dst.clone()).unwrap_err();
        assert_eq!(snappy_decoded_message, RPCError::InvalidData);
    }
}
<commit_msg>`test_decode_malicious_status_message` improvements (#2104)<commit_after>//! This handles the various supported encoding mechanism for the Eth 2.0 RPC.

use crate::rpc::methods::ErrorType;
use crate::rpc::{RPCCodedResponse, RPCRequest, RPCResponse};
use libp2p::bytes::BufMut;
use libp2p::bytes::BytesMut;
use std::marker::PhantomData;
use tokio_util::codec::{Decoder, Encoder};
use types::EthSpec;

pub trait OutboundCodec<TItem>: Encoder<TItem> + Decoder {
    type CodecErrorType;

    fn decode_error(
        &mut self,
        src: &mut BytesMut,
    ) -> Result<Option<Self::CodecErrorType>, <Self as Decoder>::Error>;
}

/* Global Inbound Codec */
// This deals with Decoding RPC Requests from other peers and encoding our responses

pub struct BaseInboundCodec<TCodec, TSpec>
where
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder,
    TSpec: EthSpec,
{
    /// Inner codec for handling various encodings
    inner: TCodec,
    phantom: PhantomData<TSpec>,
}

impl<TCodec, TSpec> BaseInboundCodec<TCodec, TSpec>
where
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder,
    TSpec: EthSpec,
{
    pub fn new(codec: TCodec) -> Self {
        BaseInboundCodec {
            inner: codec,
            phantom: PhantomData,
        }
    }
}

/* Global Outbound Codec */
// This deals with Decoding RPC Responses from other peers and encoding our requests
pub struct BaseOutboundCodec<TOutboundCodec, TSpec>
where
    TOutboundCodec: OutboundCodec<RPCRequest<TSpec>>,
    TSpec: EthSpec,
{
    /// Inner codec for handling various encodings.
    inner: TOutboundCodec,
    /// Keeps track of the current response code for a chunk.
    current_response_code: Option<u8>,
    phantom: PhantomData<TSpec>,
}

impl<TOutboundCodec, TSpec> BaseOutboundCodec<TOutboundCodec, TSpec>
where
    TSpec: EthSpec,
    TOutboundCodec: OutboundCodec<RPCRequest<TSpec>>,
{
    pub fn new(codec: TOutboundCodec) -> Self {
        BaseOutboundCodec {
            inner: codec,
            current_response_code: None,
            phantom: PhantomData,
        }
    }
}

/* Implementation of the Encoding/Decoding for the global codecs */

/* Base Inbound Codec */

// This Encodes RPC Responses sent to external peers
impl<TCodec, TSpec> Encoder<RPCCodedResponse<TSpec>> for BaseInboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: Decoder + Encoder<RPCCodedResponse<TSpec>>,
{
    type Error = <TCodec as Encoder<RPCCodedResponse<TSpec>>>::Error;

    fn encode(
        &mut self,
        item: RPCCodedResponse<TSpec>,
        dst: &mut BytesMut,
    ) -> Result<(), Self::Error> {
        dst.clear();
        dst.reserve(1);
        dst.put_u8(
            item.as_u8()
                .expect("Should never encode a stream termination"),
        );
        self.inner.encode(item, dst)
    }
}

// This Decodes RPC Requests from external peers
impl<TCodec, TSpec> Decoder for BaseInboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: Encoder<RPCCodedResponse<TSpec>> + Decoder<Item = RPCRequest<TSpec>>,
{
    type Item = RPCRequest<TSpec>;
    type Error = <TCodec as Decoder>::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        self.inner.decode(src)
    }
}

/* Base Outbound Codec */

// This Encodes RPC Requests sent to external peers
impl<TCodec, TSpec> Encoder<RPCRequest<TSpec>> for BaseOutboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: OutboundCodec<RPCRequest<TSpec>> + Encoder<RPCRequest<TSpec>>,
{
    type Error = <TCodec as Encoder<RPCRequest<TSpec>>>::Error;

    fn encode(&mut self, item: RPCRequest<TSpec>, dst: &mut BytesMut) -> Result<(), Self::Error> {
        self.inner.encode(item, dst)
    }
}

// This decodes RPC Responses received from external peers
impl<TCodec, TSpec> Decoder for BaseOutboundCodec<TCodec, TSpec>
where
    TSpec: EthSpec,
    TCodec: OutboundCodec<RPCRequest<TSpec>, CodecErrorType = ErrorType>
        + Decoder<Item = RPCResponse<TSpec>>,
{
    type Item = RPCCodedResponse<TSpec>;
    type Error = <TCodec as Decoder>::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        // if we have only received the response code, wait for more bytes
        if src.len() <= 1 {
            return Ok(None);
        }
        // using the response code determine which kind of payload needs to be decoded.
        let response_code = self.current_response_code.unwrap_or_else(|| {
            let resp_code = src.split_to(1)[0];
            self.current_response_code = Some(resp_code);
            resp_code
        });

        let inner_result = {
            if RPCCodedResponse::<TSpec>::is_response(response_code) {
                // decode an actual response and mutates the buffer if enough bytes have been read
                // returning the result.
                self.inner
                    .decode(src)
                    .map(|r| r.map(RPCCodedResponse::Success))
            } else {
                // decode an error
                self.inner
                    .decode_error(src)
                    .map(|r| r.map(|resp| RPCCodedResponse::from_error(response_code, resp)))
            }
        };
        // if the inner decoder was capable of decoding a chunk, we need to reset the current
        // response code for the next chunk
        if let Ok(Some(_)) = inner_result {
            self.current_response_code = None;
        }
        // return the result
        inner_result
    }
}

#[cfg(test)]
mod tests {
    use super::super::ssz_snappy::*;
    use super::*;
    use crate::rpc::methods::StatusMessage;
    use crate::rpc::protocol::*;
    use snap::write::FrameEncoder;
    use ssz::Encode;
    use std::io::Write;
    use types::{Epoch, Hash256, Slot};
    use unsigned_varint::codec::Uvi;

    type Spec = types::MainnetEthSpec;

    #[test]
    fn test_decode_status_message() {
        let message = hex::decode("0054ff060000734e615070590032000006e71e7b54989925efd6c9cbcb8ceb9b5f71216f5137282bf6a1e3b50f64e42d6c7fb347abe07eb0db8200000005029e2800").unwrap();
        let mut buf = BytesMut::new();
        buf.extend_from_slice(&message);

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);

        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        // remove response code
        let mut snappy_buf = buf.clone();
        let _ = snappy_buf.split_to(1);

        // decode message just as snappy message
        let snappy_decoded_message = snappy_outbound_codec.decode(&mut snappy_buf).unwrap();

        // build codecs for entire chunk
        let mut snappy_base_outbound_codec = BaseOutboundCodec::new(snappy_outbound_codec);

        // decode message as ssz snappy chunk
        let snappy_decoded_chunk = snappy_base_outbound_codec.decode(&mut buf).unwrap();

        dbg!(snappy_decoded_message);
        dbg!(snappy_decoded_chunk);
    }

    #[test]
    fn test_invalid_length_prefix() {
        let mut uvi_codec: Uvi<u128> = Uvi::default();
        let mut dst = BytesMut::with_capacity(1024);

        // Smallest > 10 byte varint
        let len: u128 = 2u128.pow(70);

        // Insert length-prefix
        uvi_codec.encode(len, &mut dst).unwrap();

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);
        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        let snappy_decoded_message = snappy_outbound_codec.decode(&mut dst).unwrap_err();

        assert_eq!(
            snappy_decoded_message,
            RPCError::IoError("input bytes exceed maximum".to_string()),
            "length-prefix of > 10 bytes is invalid"
        );
    }

    #[test]
    fn test_length_limits() {
        fn encode_len(len: usize) -> BytesMut {
            let mut uvi_codec: Uvi<usize> = Uvi::default();
            let mut dst = BytesMut::with_capacity(1024);
            uvi_codec.encode(len, &mut dst).unwrap();
            dst
        }

        let protocol_id =
            ProtocolId::new(Protocol::BlocksByRange, Version::V1, Encoding::SSZSnappy);

        // Response limits
        let limit = protocol_id.rpc_response_limits::<Spec>();
        let mut max = encode_len(limit.max + 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut max).unwrap_err(), RPCError::InvalidData);

        let mut min = encode_len(limit.min - 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut min).unwrap_err(), RPCError::InvalidData);

        // Request limits
        let limit = protocol_id.rpc_request_limits();
        let mut max = encode_len(limit.max + 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id.clone(), 1_048_576);
        assert_eq!(codec.decode(&mut max).unwrap_err(), RPCError::InvalidData);

        let mut min = encode_len(limit.min - 1);
        let mut codec = SSZSnappyOutboundCodec::<Spec>::new(protocol_id, 1_048_576);
        assert_eq!(codec.decode(&mut min).unwrap_err(), RPCError::InvalidData);
    }

    #[test]
    fn test_decode_malicious_status_message() {
        // 10 byte snappy stream identifier
        let stream_identifier: &'static [u8] = b"\xFF\x06\x00\x00sNaPpY";

        assert_eq!(stream_identifier.len(), 10);

        // byte 0(0xFE) is padding chunk type identifier for snappy messages
        // byte 1,2,3 are chunk length (little endian)
        let malicious_padding: &'static [u8] = b"\xFE\x00\x00\x00";

        // Status message is 84 bytes uncompressed. `max_compressed_len` is 32 + 84 + 84/6 = 130.
        let status_message_bytes = StatusMessage {
            fork_digest: [0; 4],
            finalized_root: Hash256::from_low_u64_be(0),
            finalized_epoch: Epoch::new(1),
            head_root: Hash256::from_low_u64_be(0),
            head_slot: Slot::new(1),
        }
        .as_ssz_bytes();

        assert_eq!(status_message_bytes.len(), 84);
        assert_eq!(snap::raw::max_compress_len(status_message_bytes.len()), 130);

        let mut uvi_codec: Uvi<usize> = Uvi::default();
        let mut dst = BytesMut::with_capacity(1024);

        // Insert length-prefix
        uvi_codec
            .encode(status_message_bytes.len(), &mut dst)
            .unwrap();

        // Insert snappy stream identifier
        dst.extend_from_slice(stream_identifier);

        // Insert malicious padding of 80 bytes.
        for _ in 0..20 {
            dst.extend_from_slice(malicious_padding);
        }

        // Insert payload (42 bytes compressed)
        let mut writer = FrameEncoder::new(Vec::new());
        writer.write_all(&status_message_bytes).unwrap();
        writer.flush().unwrap();
        assert_eq!(writer.get_ref().len(), 42);
        dst.extend_from_slice(writer.get_ref());

        // 10 (for stream identifier) + 80 + 42 = 132 > `max_compressed_len`. Hence, decoding should fail with `InvalidData`.

        let snappy_protocol_id =
            ProtocolId::new(Protocol::Status, Version::V1, Encoding::SSZSnappy);

        let mut snappy_outbound_codec =
            SSZSnappyOutboundCodec::<Spec>::new(snappy_protocol_id, 1_048_576);

        let snappy_decoded_message = snappy_outbound_codec.decode(&mut dst).unwrap_err();
        assert_eq!(snappy_decoded_message, RPCError::InvalidData);
    }
}
<|endoftext|>


<commit_before>use common::{HTTPDate, HTTPVersion, Header, StatusCode};

use std::cmp::Ordering;
use std::sync::mpsc::Receiver;

use std::io::Result as IoResult;
use std::io::{self, Cursor, Read, Write};

use std::fs::File;

use std::str::FromStr;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
/// of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///     will assume that the data from the `Read` object has the specified encoding
///     and will just pass-through.
///
///  - `Content-Length`: The length of the data should be set manually
///     using the `Reponse` object's API. Attempting to set the value of this
///     header will be equivalent to modifying the size of the data but the header
///     itself may not be present in the final result.
///
pub struct Response<R>
where
    R: Read,
{
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<usize>,
    chunked_threshold: Option<usize>,
}

/// A `Response` without a template parameter.
pub type ResponseBox = Response<Box<dyn Read + Send>>;

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
#[derive(Copy, Clone)]
enum TransferEncoding {
    Identity,
    Chunked,
}

impl FromStr for TransferEncoding {
    type Err = ();

    fn from_str(input: &str) -> Result<TransferEncoding, ()> {
        if input.eq_ignore_ascii_case("identity") {
            Ok(TransferEncoding::Identity)
        } else if input.eq_ignore_ascii_case("chunked") {
            Ok(TransferEncoding::Chunked)
        } else {
            Err(())
        }
    }
}

/// Builds a Date: header with the current date.
fn build_date_header() -> Header {
    let d = HTTPDate::new();
    Header::from_bytes(&b"Date"[..], &d.to_string().into_bytes()[..]).unwrap()
}

fn write_message_header<W>(
    mut writer: W,
    http_version: &HTTPVersion,
    status_code: &StatusCode,
    headers: &[Header],
) -> IoResult<()>
where
    W: Write,
{
    // writing status line
    write!(
        &mut writer,
        "HTTP/{}.{} {} {}\r\n",
        http_version.0,
        http_version.1,
        status_code.0,
        status_code.default_reason_phrase()
    )?;

    // writing headers
    for header in headers.iter() {
        writer.write_all(header.field.as_str().as_ref())?;
        write!(&mut writer, ": ")?;
        writer.write_all(header.value.as_str().as_ref())?;
        write!(&mut writer, "\r\n")?;
    }

    // separator between header and data
    write!(&mut writer, "\r\n")?;

    Ok(())
}

fn choose_transfer_encoding(
    request_headers: &[Header],
    http_version: &HTTPVersion,
    entity_length: &Option<usize>,
    has_additional_headers: bool,
    chunked_threshold: usize,
) -> TransferEncoding {
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= (1, 0) {
        return TransferEncoding::Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers
        .iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))
        // getting its value
        .map(|h| h.value.clone())
        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value.as_str()); // TODO: remove conversion

            // sorting elements by most priority
            parse.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                // q=0 are ignored
                if value.1 <= 0.0 {
                    continue;
                }

                if let Ok(te) = TransferEncoding::from_str(value.0) {
                    return Some(te);
                }
            }

            // encoding not found
            None
        });

    //
    if let Some(user_request) = user_request {
        return user_request;
    }

    // if we have additional headers, using chunked
    if has_additional_headers {
        return TransferEncoding::Chunked;
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    if entity_length
        .as_ref()
        .map_or(true, |val| *val >= chunked_threshold)
    {
        return TransferEncoding::Chunked;
    }

    // Identity by default
    TransferEncoding::Identity
}

impl<R> Response<R>
where
    R: Read,
{
    /// Creates a new Response object.
    ///
    /// The `additional_headers` argument is a receiver that
    ///  may provide headers even after the response has been sent.
    ///
    /// All the other arguments are straight-forward.
    pub fn new(
        status_code: StatusCode,
        headers: Vec<Header>,
        data: R,
        data_length: Option<usize>,
        additional_headers: Option<Receiver<Header>>,
    ) -> Response<R> {
        let mut response = Response {
            reader: data,
            status_code,
            headers: Vec::with_capacity(16),
            data_length,
            chunked_threshold: None,
        };

        for h in headers {
            response.add_header(h)
        }

        // dummy implementation
        if let Some(additional_headers) = additional_headers {
            for h in additional_headers.iter() {
                response.add_header(h)
            }
        }

        response
    }

    /// Set a threshold for `Content-Length` where we chose chunked
    /// transfer. Notice that chunked transfer might happen regardless of
    /// this threshold, for instance when the request headers indicate
    /// it is wanted or when there is no `Content-Length`.
    pub fn with_chunked_threshold(mut self, length: usize) -> Response<R> {
        self.chunked_threshold = Some(length);
        self
    }

    /// The current `Content-Length` threshold for switching over to
    /// chunked transfer. The default is 32768 bytes. Notice that
    /// chunked transfer is mutually exclusive with sending a
    /// `Content-Length` header as per the HTTP spec.
    pub fn chunked_threshold(&self) -> usize {
        self.chunked_threshold.unwrap_or(32768)
    }

    /// Adds a header to the list.
    /// Does all the checks.
    pub fn add_header<H>(&mut self, header: H)
    where
        H: Into<Header>,
    {
        let header = header.into();

        // ignoring forbidden headers
        if header.field.equiv(&"Accept-Ranges")
            || header.field.equiv(&"Connection")
            || header.field.equiv(&"Content-Range")
            || header.field.equiv(&"Trailer")
            || header.field.equiv(&"Transfer-Encoding")
            || header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            if let Ok(val) = usize::from_str(header.value.as_str()) {
                self.data_length = Some(val)
            }

            return;
        }

        self.headers.push(header);
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[inline]
    pub fn with_header<H>(mut self, header: H) -> Response<R>
    where
        H: Into<Header>,
    {
        self.add_header(header.into());
        self
    }

    /// Returns the same request, but with a different status code.
    #[inline]
    pub fn with_status_code<S>(mut self, code: S) -> Response<R>
    where
        S: Into<StatusCode>,
    {
        self.status_code = code.into();
        self
    }

    /// Returns the same request, but with different data.
    pub fn with_data<S>(self, reader: S, data_length: Option<usize>) -> Response<S>
    where
        S: Read,
    {
        Response {
            reader,
            headers: self.headers,
            status_code: self.status_code,
            data_length,
            chunked_threshold: None,
        }
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    ///
    /// Note: does not flush the writer.
    pub fn raw_print<W: Write>(
        mut self,
        mut writer: W,
        http_version: HTTPVersion,
        request_headers: &[Header],
        do_not_send_body: bool,
        upgrade: Option<&str>,
    ) -> IoResult<()> {
        let mut transfer_encoding = Some(choose_transfer_encoding(
            request_headers,
            &http_version,
            &self.data_length,
            false, /* TODO */
            self.chunked_threshold(),
        ));

        // add `Date` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Date"))
            .is_none()
        {
            self.headers.insert(0, build_date_header());
        }

        // add `Server` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Server"))
            .is_none()
        {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Server"[..], &b"tiny-http (Rust)"[..]).unwrap(),
            );
        }

        // handling upgrade
        if let Some(upgrade) = upgrade {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Upgrade"[..], upgrade.as_bytes()).unwrap(),
            );
            self.headers.insert(
                0,
                Header::from_bytes(&b"Connection"[..], &b"upgrade"[..]).unwrap(),
            );
            transfer_encoding = None;
        }

        // if the transfer encoding is identity, the content length must be known ; therefore if
        // we don't know it, we buffer the entire response first here
        // while this is an expensive operation, it is only ever needed for clients using HTTP 1.0
        let (mut reader, data_length) = match (self.data_length, transfer_encoding) {
            (Some(l), _) => (Box::new(self.reader) as Box<dyn Read>, Some(l)),
            (None, Some(TransferEncoding::Identity)) => {
                let mut buf = Vec::new();
                self.reader.read_to_end(&mut buf)?;
                let l = buf.len();
                (Box::new(Cursor::new(buf)) as Box<dyn Read>, Some(l))
            }
            _ => (Box::new(self.reader) as Box<dyn Read>, None),
        };

        // checking whether to ignore the body of the response
        let do_not_send_body = do_not_send_body
            || match self.status_code.0 {
                // sattus code 1xx, 204 and 304 MUST not include a body
                100..=199 | 204 | 304 => true,
                _ => false,
            };

        // preparing headers for transfer
        match transfer_encoding {
            Some(TransferEncoding::Chunked) => self
                .headers
                .push(Header::from_bytes(&b"Transfer-Encoding"[..], &b"chunked"[..]).unwrap()),

            Some(TransferEncoding::Identity) => {
                assert!(data_length.is_some());
                let data_length = data_length.unwrap();

                self.headers.push(
                    Header::from_bytes(
                        &b"Content-Length"[..],
                        format!("{}", data_length).as_bytes(),
                    )
                    .unwrap(),
                )
            }

            _ => (),
        };

        // sending headers
        write_message_header(
            writer.by_ref(),
            &http_version,
            &self.status_code,
            &self.headers,
        )?;

        // sending the body
        if !do_not_send_body {
            match transfer_encoding {
                Some(TransferEncoding::Chunked) => {
                    use chunked_transfer::Encoder;

                    let mut writer = Encoder::new(writer);
                    io::copy(&mut reader, &mut writer)?;
                }

                Some(TransferEncoding::Identity) => {
                    use util::EqualReader;

                    assert!(data_length.is_some());
                    let data_length = data_length.unwrap();

                    if data_length >= 1 {
                        let (mut equ_reader, _) = EqualReader::new(reader.by_ref(), data_length);
                        io::copy(&mut equ_reader, &mut writer)?;
                    }
                }

                _ => (),
            }
        }

        Ok(())
    }
}

impl<R> Response<R>
where
    R: Read + Send + 'static,
{
    /// Turns this response into a `Response<Box<Read + Send>>`.
    pub fn boxed(self) -> ResponseBox {
        Response {
            reader: Box::new(self.reader) as Box<dyn Read + Send>,
            status_code: self.status_code,
            headers: self.headers,
            data_length: self.data_length,
            chunked_threshold: None,
        }
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    pub fn from_file(file: File) -> Response<File> {
        let file_size = file.metadata().ok().map(|v| v.len() as usize);

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            file,
            file_size,
            None,
        )
    }
}

impl Response<Cursor<Vec<u8>>> {
    pub fn from_data<D>(data: D) -> Response<Cursor<Vec<u8>>>
    where
        D: Into<Vec<u8>>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            Cursor::new(data),
            Some(data_len),
            None,
        )
    }

    pub fn from_string<S>(data: S) -> Response<Cursor<Vec<u8>>>
    where
        S: Into<String>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec![
                Header::from_bytes(&b"Content-Type"[..], &b"text/plain; charset=UTF-8"[..])
                    .unwrap(),
            ],
            Cursor::new(data.into_bytes()),
            Some(data_len),
            None,
        )
    }
}

impl Response<io::Empty> {
    /// Builds an empty `Response` with the given status code.
    pub fn empty<S>(status_code: S) -> Response<io::Empty>
    where
        S: Into<StatusCode>,
    {
        Response::new(
            status_code.into(),
            Vec::with_capacity(0),
            io::empty(),
            Some(0),
            None,
        )
    }

    /// DEPRECATED. Use `empty` instead.
    pub fn new_empty(status_code: StatusCode) -> Response<io::Empty> {
        Response::empty(status_code)
    }
}

impl Clone for Response<io::Empty> {
    fn clone(&self) -> Response<io::Empty> {
        Response {
            reader: io::empty(),
            status_code: self.status_code.clone(),
            headers: self.headers.clone(),
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}
<commit_msg>Response: Don't forget `chunked_threshold`<commit_after>use common::{HTTPDate, HTTPVersion, Header, StatusCode};

use std::cmp::Ordering;
use std::sync::mpsc::Receiver;

use std::io::Result as IoResult;
use std::io::{self, Cursor, Read, Write};

use std::fs::File;

use std::str::FromStr;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
/// of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///     will assume that the data from the `Read` object has the specified encoding
///     and will just pass-through.
///
///  - `Content-Length`: The length of the data should be set manually
///     using the `Reponse` object's API. Attempting to set the value of this
///     header will be equivalent to modifying the size of the data but the header
///     itself may not be present in the final result.
///
pub struct Response<R>
where
    R: Read,
{
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<usize>,
    chunked_threshold: Option<usize>,
}

/// A `Response` without a template parameter.
pub type ResponseBox = Response<Box<dyn Read + Send>>;

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
#[derive(Copy, Clone)]
enum TransferEncoding {
    Identity,
    Chunked,
}

impl FromStr for TransferEncoding {
    type Err = ();

    fn from_str(input: &str) -> Result<TransferEncoding, ()> {
        if input.eq_ignore_ascii_case("identity") {
            Ok(TransferEncoding::Identity)
        } else if input.eq_ignore_ascii_case("chunked") {
            Ok(TransferEncoding::Chunked)
        } else {
            Err(())
        }
    }
}

/// Builds a Date: header with the current date.
fn build_date_header() -> Header {
    let d = HTTPDate::new();
    Header::from_bytes(&b"Date"[..], &d.to_string().into_bytes()[..]).unwrap()
}

fn write_message_header<W>(
    mut writer: W,
    http_version: &HTTPVersion,
    status_code: &StatusCode,
    headers: &[Header],
) -> IoResult<()>
where
    W: Write,
{
    // writing status line
    write!(
        &mut writer,
        "HTTP/{}.{} {} {}\r\n",
        http_version.0,
        http_version.1,
        status_code.0,
        status_code.default_reason_phrase()
    )?;

    // writing headers
    for header in headers.iter() {
        writer.write_all(header.field.as_str().as_ref())?;
        write!(&mut writer, ": ")?;
        writer.write_all(header.value.as_str().as_ref())?;
        write!(&mut writer, "\r\n")?;
    }

    // separator between header and data
    write!(&mut writer, "\r\n")?;

    Ok(())
}

fn choose_transfer_encoding(
    request_headers: &[Header],
    http_version: &HTTPVersion,
    entity_length: &Option<usize>,
    has_additional_headers: bool,
    chunked_threshold: usize,
) -> TransferEncoding {
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= (1, 0) {
        return TransferEncoding::Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers
        .iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))
        // getting its value
        .map(|h| h.value.clone())
        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value.as_str()); // TODO: remove conversion

            // sorting elements by most priority
            parse.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                // q=0 are ignored
                if value.1 <= 0.0 {
                    continue;
                }

                if let Ok(te) = TransferEncoding::from_str(value.0) {
                    return Some(te);
                }
            }

            // encoding not found
            None
        });

    //
    if let Some(user_request) = user_request {
        return user_request;
    }

    // if we have additional headers, using chunked
    if has_additional_headers {
        return TransferEncoding::Chunked;
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    if entity_length
        .as_ref()
        .map_or(true, |val| *val >= chunked_threshold)
    {
        return TransferEncoding::Chunked;
    }

    // Identity by default
    TransferEncoding::Identity
}

impl<R> Response<R>
where
    R: Read,
{
    /// Creates a new Response object.
    ///
    /// The `additional_headers` argument is a receiver that
    ///  may provide headers even after the response has been sent.
    ///
    /// All the other arguments are straight-forward.
    pub fn new(
        status_code: StatusCode,
        headers: Vec<Header>,
        data: R,
        data_length: Option<usize>,
        additional_headers: Option<Receiver<Header>>,
    ) -> Response<R> {
        let mut response = Response {
            reader: data,
            status_code,
            headers: Vec::with_capacity(16),
            data_length,
            chunked_threshold: None,
        };

        for h in headers {
            response.add_header(h)
        }

        // dummy implementation
        if let Some(additional_headers) = additional_headers {
            for h in additional_headers.iter() {
                response.add_header(h)
            }
        }

        response
    }

    /// Set a threshold for `Content-Length` where we chose chunked
    /// transfer. Notice that chunked transfer might happen regardless of
    /// this threshold, for instance when the request headers indicate
    /// it is wanted or when there is no `Content-Length`.
    pub fn with_chunked_threshold(mut self, length: usize) -> Response<R> {
        self.chunked_threshold = Some(length);
        self
    }

    /// The current `Content-Length` threshold for switching over to
    /// chunked transfer. The default is 32768 bytes. Notice that
    /// chunked transfer is mutually exclusive with sending a
    /// `Content-Length` header as per the HTTP spec.
    pub fn chunked_threshold(&self) -> usize {
        self.chunked_threshold.unwrap_or(32768)
    }

    /// Adds a header to the list.
    /// Does all the checks.
    pub fn add_header<H>(&mut self, header: H)
    where
        H: Into<Header>,
    {
        let header = header.into();

        // ignoring forbidden headers
        if header.field.equiv(&"Accept-Ranges")
            || header.field.equiv(&"Connection")
            || header.field.equiv(&"Content-Range")
            || header.field.equiv(&"Trailer")
            || header.field.equiv(&"Transfer-Encoding")
            || header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            if let Ok(val) = usize::from_str(header.value.as_str()) {
                self.data_length = Some(val)
            }

            return;
        }

        self.headers.push(header);
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[inline]
    pub fn with_header<H>(mut self, header: H) -> Response<R>
    where
        H: Into<Header>,
    {
        self.add_header(header.into());
        self
    }

    /// Returns the same request, but with a different status code.
    #[inline]
    pub fn with_status_code<S>(mut self, code: S) -> Response<R>
    where
        S: Into<StatusCode>,
    {
        self.status_code = code.into();
        self
    }

    /// Returns the same request, but with different data.
    pub fn with_data<S>(self, reader: S, data_length: Option<usize>) -> Response<S>
    where
        S: Read,
    {
        Response {
            reader,
            headers: self.headers,
            status_code: self.status_code,
            data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    ///
    /// Note: does not flush the writer.
    pub fn raw_print<W: Write>(
        mut self,
        mut writer: W,
        http_version: HTTPVersion,
        request_headers: &[Header],
        do_not_send_body: bool,
        upgrade: Option<&str>,
    ) -> IoResult<()> {
        let mut transfer_encoding = Some(choose_transfer_encoding(
            request_headers,
            &http_version,
            &self.data_length,
            false, /* TODO */
            self.chunked_threshold(),
        ));

        // add `Date` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Date"))
            .is_none()
        {
            self.headers.insert(0, build_date_header());
        }

        // add `Server` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Server"))
            .is_none()
        {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Server"[..], &b"tiny-http (Rust)"[..]).unwrap(),
            );
        }

        // handling upgrade
        if let Some(upgrade) = upgrade {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Upgrade"[..], upgrade.as_bytes()).unwrap(),
            );
            self.headers.insert(
                0,
                Header::from_bytes(&b"Connection"[..], &b"upgrade"[..]).unwrap(),
            );
            transfer_encoding = None;
        }

        // if the transfer encoding is identity, the content length must be known ; therefore if
        // we don't know it, we buffer the entire response first here
        // while this is an expensive operation, it is only ever needed for clients using HTTP 1.0
        let (mut reader, data_length) = match (self.data_length, transfer_encoding) {
            (Some(l), _) => (Box::new(self.reader) as Box<dyn Read>, Some(l)),
            (None, Some(TransferEncoding::Identity)) => {
                let mut buf = Vec::new();
                self.reader.read_to_end(&mut buf)?;
                let l = buf.len();
                (Box::new(Cursor::new(buf)) as Box<dyn Read>, Some(l))
            }
            _ => (Box::new(self.reader) as Box<dyn Read>, None),
        };

        // checking whether to ignore the body of the response
        let do_not_send_body = do_not_send_body
            || match self.status_code.0 {
                // sattus code 1xx, 204 and 304 MUST not include a body
                100..=199 | 204 | 304 => true,
                _ => false,
            };

        // preparing headers for transfer
        match transfer_encoding {
            Some(TransferEncoding::Chunked) => self
                .headers
                .push(Header::from_bytes(&b"Transfer-Encoding"[..], &b"chunked"[..]).unwrap()),

            Some(TransferEncoding::Identity) => {
                assert!(data_length.is_some());
                let data_length = data_length.unwrap();

                self.headers.push(
                    Header::from_bytes(
                        &b"Content-Length"[..],
                        format!("{}", data_length).as_bytes(),
                    )
                    .unwrap(),
                )
            }

            _ => (),
        };

        // sending headers
        write_message_header(
            writer.by_ref(),
            &http_version,
            &self.status_code,
            &self.headers,
        )?;

        // sending the body
        if !do_not_send_body {
            match transfer_encoding {
                Some(TransferEncoding::Chunked) => {
                    use chunked_transfer::Encoder;

                    let mut writer = Encoder::new(writer);
                    io::copy(&mut reader, &mut writer)?;
                }

                Some(TransferEncoding::Identity) => {
                    use util::EqualReader;

                    assert!(data_length.is_some());
                    let data_length = data_length.unwrap();

                    if data_length >= 1 {
                        let (mut equ_reader, _) = EqualReader::new(reader.by_ref(), data_length);
                        io::copy(&mut equ_reader, &mut writer)?;
                    }
                }

                _ => (),
            }
        }

        Ok(())
    }
}

impl<R> Response<R>
where
    R: Read + Send + 'static,
{
    /// Turns this response into a `Response<Box<Read + Send>>`.
    pub fn boxed(self) -> ResponseBox {
        Response {
            reader: Box::new(self.reader) as Box<dyn Read + Send>,
            status_code: self.status_code,
            headers: self.headers,
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    pub fn from_file(file: File) -> Response<File> {
        let file_size = file.metadata().ok().map(|v| v.len() as usize);

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            file,
            file_size,
            None,
        )
    }
}

impl Response<Cursor<Vec<u8>>> {
    pub fn from_data<D>(data: D) -> Response<Cursor<Vec<u8>>>
    where
        D: Into<Vec<u8>>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            Cursor::new(data),
            Some(data_len),
            None,
        )
    }

    pub fn from_string<S>(data: S) -> Response<Cursor<Vec<u8>>>
    where
        S: Into<String>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec![
                Header::from_bytes(&b"Content-Type"[..], &b"text/plain; charset=UTF-8"[..])
                    .unwrap(),
            ],
            Cursor::new(data.into_bytes()),
            Some(data_len),
            None,
        )
    }
}

impl Response<io::Empty> {
    /// Builds an empty `Response` with the given status code.
    pub fn empty<S>(status_code: S) -> Response<io::Empty>
    where
        S: Into<StatusCode>,
    {
        Response::new(
            status_code.into(),
            Vec::with_capacity(0),
            io::empty(),
            Some(0),
            None,
        )
    }

    /// DEPRECATED. Use `empty` instead.
    pub fn new_empty(status_code: StatusCode) -> Response<io::Empty> {
        Response::empty(status_code)
    }
}

impl Clone for Response<io::Empty> {
    fn clone(&self) -> Response<io::Empty> {
        Response {
            reader: io::empty(),
            status_code: self.status_code.clone(),
            headers: self.headers.clone(),
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}
<|endoftext|>


<commit_before>use common::{Header, HTTPVersion, StatusCode};
use std::io::{IoResult, MemReader};
use std::io::fs::File;
use std::io::util;
use std::io::util::NullReader;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
///  of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///      will assume that the data from the `Reader` has the specified encoding
///      and will just pass-through.
///  - `Content-Length`: If you define this header to `N`, only the first `N` bytes
///      of the `Reader` will be read. If the `Reader` reaches `EOF` before `N` bytes have
///      been read, `0`s will be sent. Also, this header may not be passed to the final
///      output.
///
#[experimental]
pub struct Response<R> {
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<uint>,
}

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
enum TransferEncoding {
    Identity,
    Chunked,
}

impl ::std::from_str::FromStr for TransferEncoding {
    fn from_str(input: &str) -> Option<TransferEncoding> {
        use std::ascii::StrAsciiExt;

        if input.eq_ignore_ascii_case("identity") {
            return Some(Identity);
        } else if input.eq_ignore_ascii_case("chunked") {
            return Some(Chunked);
        }

        None
    }
}

/// Builds a Date: header with the current date.
// TODO: this is optimisable
fn build_date_header() -> Header {
    use time;
    let date = time::now_utc().strftime("%a, %d %b %Y %H:%M:%S GMT");
    from_str(format!("Date: {}", date).as_slice()).unwrap()
}

fn write_message_header<W: Writer>(mut writer: W, http_version: &HTTPVersion,
                                   status_code: &StatusCode, headers: &[Header])
    -> IoResult<()>
{
    // writing status line
    try!(write!(writer, "HTTP/{} {} {}\r\n",
        http_version,
        status_code.as_uint(),
        status_code.get_default_reason_phrase()
    ));

    // writing headers
    for header in headers.iter() {
        try!(write!(writer, "{}: {}\r\n", header.field, header.value));
    }

    // separator between header and data
    try!(write!(writer, "\r\n"));

    Ok(())
}

fn choose_transfer_encoding(request_headers: &[Header], http_version: &HTTPVersion,
                            entity_length: &Option<uint>)
    -> TransferEncoding
{
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= HTTPVersion(1, 0) {
        return Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers.iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))

        // getting its value
        .map(|h| h.value.as_slice())

        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value);

            // sorting elements by most priority
            parse.sort_by(|a, b| b.val1().partial_cmp(&a.val1()).unwrap_or(Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                match from_str::<TransferEncoding>(value.val0()) {
                    Some(te) => return Some(te),
                    _ => ()     // unrecognized/unsupported encoding
                };
            }

            // encoding not found
            None
        });

    //
    if user_request.is_some() {
        return user_request.unwrap();
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    let chunks_threshold = 32768;
    if entity_length.as_ref().filtered(|val| **val < chunks_threshold).is_none() {
        return Chunked;
    }

    // Identity by default
    Identity
}

impl<R: Reader> Response<R> {
    #[experimental]
    pub fn new(status_code: StatusCode, headers: Vec<Header>,
               data: R, data_length: Option<uint>) -> Response<R>
    {
        let mut response = Response {
            reader: data,
            status_code: status_code,
            headers: Vec::new(),
            data_length: data_length,
        };

        for h in headers.move_iter() {
            response.add_header(h)
        }

        response
    }

    /// Adds a header to the list.
    /// Does all the checks.
    #[experimental]
    pub fn add_header(&mut self, header: Header) {
        // ignoring forbidden headers
        if header.field.equiv(&"Accept-Ranges") ||
           header.field.equiv(&"Connection") ||
           header.field.equiv(&"Content-Range") ||
           header.field.equiv(&"Trailer") ||
           header.field.equiv(&"Transfer-Encoding") ||
           header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            match from_str::<uint>(header.value.as_slice()) {
                Some(val) => self.data_length = Some(val),
                None => ()      // wrong value for content-length
            };

            return;
        }

        self.headers.push(header)
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[unstable]
    pub fn with_header(mut self, header: Header) -> Response<R> {
        self.add_header(header);
        self
    }

    /// Returns the same request, but with a different status code.
    #[unstable]
    pub fn with_status_code(mut self, code: StatusCode) -> Response<R> {
        self.status_code = code;
        self
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    // TODO: status codes 1xx, 204 and 304 must not include a response body
    #[unstable]
    pub fn raw_print<W: Writer>(mut self, mut writer: W, http_version: HTTPVersion,
                                request_headers: &[Header]) -> IoResult<()>
    {
        let transfer_encoding = choose_transfer_encoding(request_headers,
                                    &http_version, &self.data_length);

        // add `Date` if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Date")).is_none() {
            self.headers.unshift(build_date_header());
        }

        // add `Server` if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Server")).is_none() {
            self.headers.unshift(
                from_str("Server: tiny-http (Rust)").unwrap()
            );
        }

        // writing data
        match transfer_encoding {
            Chunked => {
                use util::ChunksEncoder;

                self.headers.push(
                    from_str("Transfer-Encoding: chunked").unwrap()
                );

                try!(write_message_header(writer.by_ref(), &http_version,
                    &self.status_code, self.headers.as_slice()));

                let mut writer = ChunksEncoder::new(writer);
                try!(util::copy(&mut self.reader, &mut writer));

                // flushing!
                try!(writer.flush());
            },

            Identity => {
                use util::EqualReader;

                assert!(self.data_length.is_some());
                let data_length = self.data_length.unwrap();

                self.headers.push(
                    from_str(format!("Content-Length: {}", data_length).as_slice()).unwrap()
                );

                try!(write_message_header(writer.by_ref(), &http_version,
                    &self.status_code, self.headers.as_slice()));

                if data_length >= 1 {
                    let (mut equ_reader, _) =
                        EqualReader::new(self.reader.by_ref(), data_length);
                    try!(util::copy(&mut equ_reader, &mut writer));
                }

                // flushing!
                try!(writer.flush());
            }
        };

        Ok(())
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    #[experimental]
    pub fn from_file(mut file: File) -> Response<File> {
        let file_size = file.stat().ok().map(|v| v.size as uint);

        Response::new(
            StatusCode(200),
            Vec::new(),
            file,
            file_size
        )
    }
}

impl Response<MemReader> {
    #[experimental]
    pub fn from_data(data: Vec<u8>) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::new(),
            MemReader::new(data),
            Some(data_len)
        )
    }

    #[experimental]
    pub fn from_string(data: String) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec!(
                from_str("Content-Type: text/plain; charset=UTF-8").unwrap()
            ),
            MemReader::new(data.into_bytes()),
            Some(data_len)
        )
    }
}

impl Response<NullReader> {
    /// Builds an empty `Response` with the given status code.
    #[experimental]
    pub fn new_empty(status_code: StatusCode) -> Response<NullReader> {
        Response::new(
            status_code,
            Vec::new(),
            NullReader,
            Some(0)
        )
    }
}
<commit_msg>Status codes 1xx, 204 and 304 will now always ignore body<commit_after>use common::{Header, HTTPVersion, StatusCode};
use std::io::{IoResult, MemReader};
use std::io::fs::File;
use std::io::util;
use std::io::util::NullReader;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
///  of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///      will assume that the data from the `Reader` has the specified encoding
///      and will just pass-through.
///  - `Content-Length`: If you define this header to `N`, only the first `N` bytes
///      of the `Reader` will be read. If the `Reader` reaches `EOF` before `N` bytes have
///      been read, `0`s will be sent. Also, this header may not be passed to the final
///      output.
///
#[experimental]
pub struct Response<R> {
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<uint>,
}

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
enum TransferEncoding {
    Identity,
    Chunked,
}

impl ::std::from_str::FromStr for TransferEncoding {
    fn from_str(input: &str) -> Option<TransferEncoding> {
        use std::ascii::StrAsciiExt;

        if input.eq_ignore_ascii_case("identity") {
            return Some(Identity);
        } else if input.eq_ignore_ascii_case("chunked") {
            return Some(Chunked);
        }

        None
    }
}

/// Builds a Date: header with the current date.
// TODO: this is optimisable
fn build_date_header() -> Header {
    use time;
    let date = time::now_utc().strftime("%a, %d %b %Y %H:%M:%S GMT");
    from_str(format!("Date: {}", date).as_slice()).unwrap()
}

fn write_message_header<W: Writer>(mut writer: W, http_version: &HTTPVersion,
                                   status_code: &StatusCode, headers: &[Header])
    -> IoResult<()>
{
    // writing status line
    try!(write!(writer, "HTTP/{} {} {}\r\n",
        http_version,
        status_code.as_uint(),
        status_code.get_default_reason_phrase()
    ));

    // writing headers
    for header in headers.iter() {
        try!(write!(writer, "{}: {}\r\n", header.field, header.value));
    }

    // separator between header and data
    try!(write!(writer, "\r\n"));

    Ok(())
}

fn choose_transfer_encoding(request_headers: &[Header], http_version: &HTTPVersion,
                            entity_length: &Option<uint>)
    -> TransferEncoding
{
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= HTTPVersion(1, 0) {
        return Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers.iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))

        // getting its value
        .map(|h| h.value.as_slice())

        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value);

            // sorting elements by most priority
            parse.sort_by(|a, b| b.val1().partial_cmp(&a.val1()).unwrap_or(Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                match from_str::<TransferEncoding>(value.val0()) {
                    Some(te) => return Some(te),
                    _ => ()     // unrecognized/unsupported encoding
                };
            }

            // encoding not found
            None
        });

    //
    if user_request.is_some() {
        return user_request.unwrap();
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    let chunks_threshold = 32768;
    if entity_length.as_ref().filtered(|val| **val < chunks_threshold).is_none() {
        return Chunked;
    }

    // Identity by default
    Identity
}

impl<R: Reader> Response<R> {
    #[experimental]
    pub fn new(status_code: StatusCode, headers: Vec<Header>,
               data: R, data_length: Option<uint>) -> Response<R>
    {
        let mut response = Response {
            reader: data,
            status_code: status_code,
            headers: Vec::new(),
            data_length: data_length,
        };

        for h in headers.move_iter() {
            response.add_header(h)
        }

        response
    }

    /// Adds a header to the list.
    /// Does all the checks.
    #[experimental]
    pub fn add_header(&mut self, header: Header) {
        // ignoring forbidden headers
        if header.field.equiv(&"Accept-Ranges") ||
           header.field.equiv(&"Connection") ||
           header.field.equiv(&"Content-Range") ||
           header.field.equiv(&"Trailer") ||
           header.field.equiv(&"Transfer-Encoding") ||
           header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            match from_str::<uint>(header.value.as_slice()) {
                Some(val) => self.data_length = Some(val),
                None => ()      // wrong value for content-length
            };

            return;
        }

        self.headers.push(header)
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[unstable]
    pub fn with_header(mut self, header: Header) -> Response<R> {
        self.add_header(header);
        self
    }

    /// Returns the same request, but with a different status code.
    #[unstable]
    pub fn with_status_code(mut self, code: StatusCode) -> Response<R> {
        self.status_code = code;
        self
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    #[unstable]
    pub fn raw_print<W: Writer>(mut self, mut writer: W, http_version: HTTPVersion,
                                request_headers: &[Header]) -> IoResult<()>
    {
        let transfer_encoding = choose_transfer_encoding(request_headers,
                                    &http_version, &self.data_length);

        // add `Date` if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Date")).is_none() {
            self.headers.unshift(build_date_header());
        }

        // add `Server` if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Server")).is_none() {
            self.headers.unshift(
                from_str("Server: tiny-http (Rust)").unwrap()
            );
        }

        // checking whether to ignore the body of the response
        let ignore_response_body = match self.status_code.as_uint() {
            // sattus code 1xx, 204 and 304 MUST not include a body
            100..199 | 204 | 304 => true,
            _ => false
        };

        // preparing headers for transfer
        match transfer_encoding {
            Chunked => {
                self.headers.push(
                    from_str("Transfer-Encoding: chunked").unwrap()
                )
            },

            Identity => {
                assert!(self.data_length.is_some());
                let data_length = self.data_length.unwrap();

                self.headers.push(
                    from_str(format!("Content-Length: {}", data_length).as_slice()).unwrap()
                )
            }
        };

        // sending headers
        try!(write_message_header(writer.by_ref(), &http_version,
            &self.status_code, self.headers.as_slice()));

        // sending the body
        if !ignore_response_body {
            match transfer_encoding {

                Chunked => {
                    use util::ChunksEncoder;

                    let mut writer = ChunksEncoder::new(writer);
                    try!(util::copy(&mut self.reader, &mut writer));
                },

                Identity => {
                    use util::EqualReader;

                    assert!(self.data_length.is_some());
                    let data_length = self.data_length.unwrap();

                    if data_length >= 1 {
                        let (mut equ_reader, _) =
                            EqualReader::new(self.reader.by_ref(), data_length);
                        try!(util::copy(&mut equ_reader, &mut writer));
                    }
                }

            }
        }

        Ok(())
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    #[experimental]
    pub fn from_file(mut file: File) -> Response<File> {
        let file_size = file.stat().ok().map(|v| v.size as uint);

        Response::new(
            StatusCode(200),
            Vec::new(),
            file,
            file_size
        )
    }
}

impl Response<MemReader> {
    #[experimental]
    pub fn from_data(data: Vec<u8>) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::new(),
            MemReader::new(data),
            Some(data_len)
        )
    }

    #[experimental]
    pub fn from_string(data: String) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec!(
                from_str("Content-Type: text/plain; charset=UTF-8").unwrap()
            ),
            MemReader::new(data.into_bytes()),
            Some(data_len)
        )
    }
}

impl Response<NullReader> {
    /// Builds an empty `Response` with the given status code.
    #[experimental]
    pub fn new_empty(status_code: StatusCode) -> Response<NullReader> {
        Response::new(
            status_code,
            Vec::new(),
            NullReader,
            Some(0)
        )
    }
}
<|endoftext|>


<commit_before>#![feature(thread_local)]

use std::thread;

#[thread_local]
static mut A: u8 = 0;
#[thread_local]
static mut B: u8 = 0;
static mut C: u8 = 0;

unsafe fn get_a_ref() -> *mut u8 {
    &mut A
}

struct Sender(*mut u8);

unsafe impl Send for Sender {}

fn main() {

    let ptr = unsafe {
        let x = get_a_ref();
        *x = 5;
        assert_eq!(A, 5);
        B = 15;
        C = 25;
        Sender(&mut A)
    };

    thread::spawn(move || {
        unsafe {
            assert_eq!(*ptr.0, 5);
            assert_eq!(A, 0);
            assert_eq!(B, 0);
            assert_eq!(C, 25);
            B = 14;
            C = 24;
            let y = get_a_ref();
            assert_eq!(*y, 0);
            *y = 4;
            assert_eq!(*ptr.0, 5);
            assert_eq!(A, 4);
            assert_eq!(*get_a_ref(), 4);

        }
    }).join().unwrap();

    unsafe {
        assert_eq!(*get_a_ref(), 5);
        assert_eq!(A, 5);
        assert_eq!(B, 15);
        assert_eq!(C, 24);
    }

}

<commit_msg>Rustfmt the test.<commit_after>#![feature(thread_local)]

use std::thread;

#[thread_local]
static mut A: u8 = 0;
#[thread_local]
static mut B: u8 = 0;
static mut C: u8 = 0;

unsafe fn get_a_ref() -> *mut u8 {
    &mut A
}

struct Sender(*mut u8);

unsafe impl Send for Sender {}

fn main() {
    let ptr = unsafe {
        let x = get_a_ref();
        *x = 5;
        assert_eq!(A, 5);
        B = 15;
        C = 25;
        Sender(&mut A)
    };

    thread::spawn(move || unsafe {
        assert_eq!(*ptr.0, 5);
        assert_eq!(A, 0);
        assert_eq!(B, 0);
        assert_eq!(C, 25);
        B = 14;
        C = 24;
        let y = get_a_ref();
        assert_eq!(*y, 0);
        *y = 4;
        assert_eq!(*ptr.0, 5);
        assert_eq!(A, 4);
        assert_eq!(*get_a_ref(), 4);
    })
    .join()
    .unwrap();

    unsafe {
        assert_eq!(*get_a_ref(), 5);
        assert_eq!(A, 5);
        assert_eq!(B, 15);
        assert_eq!(C, 24);
    }
}
<|endoftext|>


<commit_before>#[macro_export]
macro_rules! gc_ref_type {
    (pub struct $ref_type:ident / $storage_type:ident <'a> {
        $($field_name:ident / $field_setter_name:ident : $field_type: ty),*
    }) => {
        struct $storage_type <'a> {
            $($field_name: <$field_type as $crate::HeapInline<'a>>::Storage,)*
        }

        unsafe impl<'a> $crate::Mark<'a> for $storage_type<'a> {
            unsafe fn mark(ptr: *mut $storage_type<'a>) {
                if !$crate::Heap::get_mark_bit(ptr) {
                    $crate::Heap::set_mark_bit(ptr);
                    $(
                        $crate::Mark::mark(
                            &mut (*ptr).$field_name
                                as *mut <$field_type as $crate::HeapInline<'a>>::Storage);
                    )*
                }
            }
        }

        unsafe impl<'a> $crate::Mark<'a> for *mut $storage_type<'a> {
            unsafe fn mark(field_ptr: *mut *mut $storage_type<'a>) {
                let ptr = *field_ptr;
                if !ptr.is_null() {
                    Mark::mark(ptr);
                }
            }
        }

        #[allow(raw_pointer_derive)]
        #[derive(Clone, Debug, PartialEq)]
        pub struct $ref_type<'a>($crate::PinnedRef<'a, $storage_type<'a>>);

        impl<'a> $ref_type<'a> {
            $(
                pub fn $field_name(&self) -> $field_type {
                    let ptr = self.0.ptr;
                    unsafe {
                        HeapInline::from_heap(
                            &*self.0.heap,
                            &(*ptr).$field_name)
                    }
                }

                pub fn $field_setter_name(&self, v: $field_type) {
                    let ptr = self.0.ptr;
                    unsafe {
                        (*ptr).$field_name = HeapInline::to_heap(v);
                    }
                }
            )*
        }

        unsafe impl<'a> $crate::HeapInline<'a> for $ref_type<'a> {
            type Storage = *mut $storage_type<'a>;

            fn to_heap(self) -> Self::Storage {
                self.0.ptr
            }

            unsafe fn from_heap(heap: &$crate::Heap<'a>, v: &Self::Storage) -> Self {
                $ref_type($crate::PinnedRef::new(heap, *v))
            }
        }

        impl<'a> $crate::GCRef for $ref_type<'a> {
            #[cfg(test)]
            fn address(&self) -> usize {
                unsafe { ::std::mem::transmute(self.0.ptr) }
            }
        }
    }
}

#[macro_export]
macro_rules! gc_inline_enum {
    { AS_ITEM $x:item } => { $x };

    {
        PARSE_VARIANTS $helper_name:ident
        {}
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name DONE
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name NO_FIELDS
            {}
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident , $($more_variants:tt)* }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name NO_FIELDS
            { $($more_variants)* }
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident ( $($field_types:tt)* ) }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name ( $($field_types)* )
            { }
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident ( $($field_types:tt)* ), $($more_variants:tt)* }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name ( $($field_types)* )
            { $($more_variants)* }
            $($etc)*
        }
    };

    {
        DECLARE_STORAGE_TYPE DONE { $($accumulated_output:tt)* } $storage_type:ident
    } => {
        gc_inline_enum! {
            AS_ITEM
            enum $storage_type<'a> {
                $($accumulated_output)*
            }
        }
    };

    {
        DECLARE_STORAGE_TYPE VARIANT $variant_name:ident NO_FIELDS
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $more_variants {
                $($accumulated_output)*
                $variant_name,
            }
            $storage_type
        }
    };

    {
        DECLARE_STORAGE_TYPE VARIANT $variant_name:ident ( $($field_type:ty),* )
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $more_variants {
                $($accumulated_output)*
                $variant_name($(<$field_type as $crate::HeapInline<'a>>::Storage),*),
            }
            $storage_type
        }
    };

    {
        IMPL_MARK DONE { $($accumulated_output:tt)* } $storage_type:ident
    } => {
        gc_inline_enum! {
            AS_ITEM
            unsafe impl<'a> $crate::Mark<'a> for $storage_type<'a> {
                unsafe fn mark(ptr: *mut $storage_type<'a>) {
                    match *ptr {
                        $($accumulated_output)*
                    }
                }
            }
        }
    };

    {
        IMPL_MARK VARIANT $name:ident NO_FIELDS
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $more_variants {
                $($accumulated_output)*
                $storage_type::$name => (),
            }
            $storage_type
        }
    };

    {
        IMPL_MARK VARIANT $name:ident ( $($field_type:ty),* )
        $($etc:tt)*
    } => {
        gc_inline_enum! {
            TYPES_TO_IDENTS ( $($field_type),*, ) () (a b c d e f g h i j k l m n o p q r s t u v w x y z)
            (IMPL_MARK CONTINUE_VARIANT $name $($etc)*)
        }
    };

    {
        IMPL_MARK CONTINUE_VARIANT $name:ident
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
        ( $(($binding:ident : $field_type:ty))* )
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $more_variants {
                $($accumulated_output)*
                $storage_type::$name ( $(ref mut $binding),* ) => {
                    $( $crate::Mark::mark($binding as *mut <$field_type as HeapInline<'a>>::Storage); )*
                }
            }
            $storage_type
        }
    };

    { TYPES_TO_IDENTS () ($(($binding:ident : $btype:ty))*) $_leftovers:tt ($($ctn:tt)*) } => {
        gc_inline_enum! { $($ctn)* ($(($binding : $btype))*) }
    };
    {
        TYPES_TO_IDENTS
        ($t:ty, $($ts:ty),*)
        ($(($binding:ident : $btype:ty))*)
        ($id:ident $($ids:tt)*)
        ($($ctn:tt)*)
    } => {
        gc_inline_enum! {
            TYPES_TO_IDENTS
            ($($ts),*)
            ($(($binding : $btype))* ($id : $t))
            ($($ids)*)
            ($($ctn)*)
        }
    };

    {
        pub enum $stack_type:ident / $storage_type:ident <'a>
            $variants:tt
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $variants {}
            $storage_type
        }

        gc_inline_enum! {
            AS_ITEM
            #[derive(Debug, Clone, PartialEq)]
            pub enum $stack_type<'a>
                $variants
        }

        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $variants {}
            $storage_type
        }

        unsafe impl<'a> $crate::HeapInline<'a> for $stack_type<'a> {
            type Storage = $storage_type<'a>;

            fn to_heap(self) -> $storage_type<'a> {
                match self {
                    $stack_type::Null => $storage_type::Null,
                    $stack_type::Int(n) => $storage_type::Int(n),
                    $stack_type::Str(rcstr) => $storage_type::Str(rcstr),
                    $stack_type::Pair(pair) => $storage_type::Pair(HeapInline::<'a>::to_heap(pair))
                }
            }

            unsafe fn from_heap(heap: &Heap<'a>, v: &$storage_type<'a>) -> $stack_type<'a> {
                match v {
                    &$storage_type::Null => $stack_type::Null,
                    &$storage_type::Int(n) => $stack_type::Int(n),
                    &$storage_type::Str(ref rcstr) => $stack_type::Str(rcstr.clone()),
                    &$storage_type::Pair(ref ptr) => $stack_type::Pair(HeapInline::<'a>::from_heap(heap, ptr))
                }
            }
        }
    }
}
<commit_msg>gc_inline_enum: Working toward macroizing to_heap and from_heap.<commit_after>#[macro_export]
macro_rules! gc_ref_type {
    (pub struct $ref_type:ident / $storage_type:ident <'a> {
        $($field_name:ident / $field_setter_name:ident : $field_type: ty),*
    }) => {
        struct $storage_type <'a> {
            $($field_name: <$field_type as $crate::HeapInline<'a>>::Storage,)*
        }

        unsafe impl<'a> $crate::Mark<'a> for $storage_type<'a> {
            unsafe fn mark(ptr: *mut $storage_type<'a>) {
                if !$crate::Heap::get_mark_bit(ptr) {
                    $crate::Heap::set_mark_bit(ptr);
                    $(
                        $crate::Mark::mark(
                            &mut (*ptr).$field_name
                                as *mut <$field_type as $crate::HeapInline<'a>>::Storage);
                    )*
                }
            }
        }

        unsafe impl<'a> $crate::Mark<'a> for *mut $storage_type<'a> {
            unsafe fn mark(field_ptr: *mut *mut $storage_type<'a>) {
                let ptr = *field_ptr;
                if !ptr.is_null() {
                    Mark::mark(ptr);
                }
            }
        }

        #[allow(raw_pointer_derive)]
        #[derive(Clone, Debug, PartialEq)]
        pub struct $ref_type<'a>($crate::PinnedRef<'a, $storage_type<'a>>);

        impl<'a> $ref_type<'a> {
            $(
                pub fn $field_name(&self) -> $field_type {
                    let ptr = self.0.ptr;
                    unsafe {
                        HeapInline::from_heap(
                            &*self.0.heap,
                            &(*ptr).$field_name)
                    }
                }

                pub fn $field_setter_name(&self, v: $field_type) {
                    let ptr = self.0.ptr;
                    unsafe {
                        (*ptr).$field_name = HeapInline::to_heap(v);
                    }
                }
            )*
        }

        unsafe impl<'a> $crate::HeapInline<'a> for $ref_type<'a> {
            type Storage = *mut $storage_type<'a>;

            fn to_heap(self) -> Self::Storage {
                self.0.ptr
            }

            unsafe fn from_heap(heap: &$crate::Heap<'a>, v: &Self::Storage) -> Self {
                $ref_type($crate::PinnedRef::new(heap, *v))
            }
        }

        impl<'a> $crate::GCRef for $ref_type<'a> {
            #[cfg(test)]
            fn address(&self) -> usize {
                unsafe { ::std::mem::transmute(self.0.ptr) }
            }
        }
    }
}

#[macro_export]
macro_rules! gc_inline_enum {
    { AS_ITEM $x:item } => { $x };
    { AS_EXPR $x:expr } => { $x };

    {
        PARSE_VARIANTS $helper_name:ident
        {}
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name DONE
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name NO_FIELDS
            {}
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident , $($more_variants:tt)* }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name NO_FIELDS
            { $($more_variants)* }
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident ( $($field_types:tt)* ) }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name ( $($field_types)* )
            { }
            $($etc)*
        }
    };

    {
        PARSE_VARIANTS $helper_name:ident
        { $variant_name:ident ( $($field_types:tt)* ), $($more_variants:tt)* }
        $( $etc:tt )*
    } => {
        gc_inline_enum! {
            $helper_name VARIANT $variant_name ( $($field_types)* )
            { $($more_variants)* }
            $($etc)*
        }
    };

    {
        DECLARE_STORAGE_TYPE DONE { $($accumulated_output:tt)* } $storage_type:ident
    } => {
        gc_inline_enum! {
            AS_ITEM
            enum $storage_type<'a> {
                $($accumulated_output)*
            }
        }
    };

    {
        DECLARE_STORAGE_TYPE VARIANT $variant_name:ident NO_FIELDS
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $more_variants {
                $($accumulated_output)*
                $variant_name,
            }
            $storage_type
        }
    };

    {
        DECLARE_STORAGE_TYPE VARIANT $variant_name:ident ( $($field_type:ty),* )
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $more_variants {
                $($accumulated_output)*
                $variant_name($(<$field_type as $crate::HeapInline<'a>>::Storage),*),
            }
            $storage_type
        }
    };

    {
        IMPL_MARK DONE { $($accumulated_output:tt)* } $storage_type:ident
    } => {
        gc_inline_enum! {
            AS_ITEM
            unsafe impl<'a> $crate::Mark<'a> for $storage_type<'a> {
                unsafe fn mark(ptr: *mut $storage_type<'a>) {
                    match *ptr {
                        $($accumulated_output)*
                    }
                }
            }
        }
    };

    {
        IMPL_MARK VARIANT $name:ident NO_FIELDS
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $more_variants {
                $($accumulated_output)*
                $storage_type::$name => (),
            }
            $storage_type
        }
    };

    {
        IMPL_MARK VARIANT $name:ident ( $($field_type:ty),* )
        $($etc:tt)*
    } => {
        gc_inline_enum! {
            TYPES_TO_IDENTS ( $($field_type),*, ) () (a b c d e f g h i j k l m n o p q r s t u v w x y z)
            (IMPL_MARK CONTINUE_VARIANT $name $($etc)*)
        }
    };

    {
        IMPL_MARK CONTINUE_VARIANT $name:ident
        $more_variants:tt
        { $($accumulated_output:tt)* }
        $storage_type:ident
        ( $(($binding:ident : $field_type:ty))* )
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $more_variants {
                $($accumulated_output)*
                $storage_type::$name ( $(ref mut $binding),* ) => {
                    $( $crate::Mark::mark($binding as *mut <$field_type as HeapInline<'a>>::Storage); )*
                }
            }
            $storage_type
        }
    };

    { TYPES_TO_IDENTS () ($(($binding:ident : $btype:ty))*) $_leftovers:tt ($($ctn:tt)*) } => {
        gc_inline_enum! { $($ctn)* ($(($binding : $btype))*) }
    };
    {
        TYPES_TO_IDENTS
        ($t:ty, $($ts:ty),*)
        ($(($binding:ident : $btype:ty))*)
        ($id:ident $($ids:tt)*)
        ($($ctn:tt)*)
    } => {
        gc_inline_enum! {
            TYPES_TO_IDENTS
            ($($ts),*)
            ($(($binding : $btype))* ($id : $t))
            ($($ids)*)
            ($($ctn)*)
        }
    };

    {
        TO_HEAP DONE { $($accumulated_output:tt)* }
        $self_:expr, $_stack_type:ident / $_storage_type:ident
    } => {
        gc_inline_enum! {
            AS_EXPR
            match $self_ {
                $($accumulated_output)*
            }
        }
    };

    {
        FROM_HEAP DONE { $($accumulated_output:tt)* }
        $value:expr, $_stack_type:ident / $_storage_type:ident
    } => {
        gc_inline_enum! {
            AS_EXPR
            match $value {
                $($accumulated_output)*
            }
        }
    };

    {
        pub enum $stack_type:ident / $storage_type:ident <'a>
            $variants:tt
    } => {
        gc_inline_enum! {
            PARSE_VARIANTS DECLARE_STORAGE_TYPE $variants {}
            $storage_type
        }

        gc_inline_enum! {
            AS_ITEM
            #[derive(Debug, Clone, PartialEq)]
            pub enum $stack_type<'a>
                $variants
        }

        gc_inline_enum! {
            PARSE_VARIANTS IMPL_MARK $variants {}
            $storage_type
        }

        unsafe impl<'a> $crate::HeapInline<'a> for $stack_type<'a> {
            type Storage = $storage_type<'a>;

            fn to_heap(self) -> $storage_type<'a> {
                gc_inline_enum! {
                    TO_HEAP DONE {
                        $stack_type::Null => $storage_type::Null,
                        $stack_type::Int(n) => $storage_type::Int(n),
                        $stack_type::Str(rcstr) => $storage_type::Str(rcstr),
                        $stack_type::Pair(pair) => $storage_type::Pair(HeapInline::<'a>::to_heap(pair))
                    }
                    self, $stack_type / $storage_type
                }
            }

            unsafe fn from_heap(heap: &Heap<'a>, v: &$storage_type<'a>) -> $stack_type<'a> {
                gc_inline_enum! {
                    FROM_HEAP DONE {
                        &$storage_type::Null => $stack_type::Null,
                        &$storage_type::Int(n) => $stack_type::Int(n),
                        &$storage_type::Str(ref rcstr) => $stack_type::Str(rcstr.clone()),
                        &$storage_type::Pair(ref ptr) => $stack_type::Pair(HeapInline::<'a>::from_heap(heap, ptr))
                    }
                    v, $stack_type / $storage_type
                }
            }
        }
    }
}
<|endoftext|>