<commit_before>// Copyright 2015 The tiny-http Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ascii::AsciiExt;

use std::io::Error as IoError;
use std::io::{self, Cursor, Read, Write, ErrorKind};

use std::net::SocketAddr;
use std::fmt;
use std::str::FromStr;

use {Header, HTTPVersion, Method, Response, StatusCode};
use util::EqualReader;
use chunked_transfer::Decoder;

/// Represents an HTTP request made by a client.
///
/// A `Request` object is what is produced by the server, and is your what
/// your code must analyse and answer.
///
/// This object implements the `Send` trait, therefore you can dispatch your requests to
/// worker threads.
///
/// # Pipelining
///
/// If a client sends multiple requests in a row (without waiting for the response), then you will
/// get multiple `Request` objects simultaneously. This is called *requests pipelining*.
/// Tiny-http automatically reorders the responses so that you don't need to worry about the order
/// in which you call `respond` or `into_writer`.
///
/// This mechanic is disabled if:
///
///  - The body of a request is large enough (handling requires pipelining requires storing the
///    body of the request in a buffer ; if the body is too big, tiny-http will avoid doing that)
///  - A request sends a `Expect: 100-continue` header (which means that the client waits to
///    know whether its body will be processed before sending it)
///  - A request sends a `Connection: close` header or `Connection: upgrade` header (used for
///    websockets), which indicates that this is the last request that will be received on this
///    connection
///
/// # Automatic cleanup
///
/// If a `Request` object is destroyed without `into_writer` or `respond` being called,
/// an empty response with a 500 status code (internal server error) will automatically be
/// sent back to the client.
/// This means that if your code fails during the handling of a request, this "internal server
/// error" response will automatically be sent during the stack unwinding.
pub struct Request {
    // where to read the body from
    data_reader: Option<Box<Read + Send + 'static>>,

    // if this writer is empty, then the request has been answered
    response_writer: Option<Box<Write + Send + 'static>>,

    remote_addr: SocketAddr,

    // true if HTTPS, false if HTTP
    secure: bool,

    method: Method,

    path: String,

    http_version: HTTPVersion,

    headers: Vec<Header>,

    body_length: Option<usize>,

    // true if a `100 Continue` response must be sent when `as_reader()` is called
    must_send_continue: bool,
}

/// Error that can happen when building a `Request` object.
pub enum RequestCreationError {
    /// The client sent an `Expect` header that was not recognized by tiny-http.
    ExpectationFailed,

    /// Error while reading data from the socket during the creation of the `Request`.
    CreationIoError(IoError),
}

impl From<IoError> for RequestCreationError {
    fn from(err: IoError) -> RequestCreationError {
        RequestCreationError::CreationIoError(err)
    }
}

/// Builds a new request.
///
/// After the request line and headers have been read from the socket, a new `Request` object
/// is built.
///
/// You must pass a `Read` that will allow the `Request` object to read from the incoming data.
/// It is the responsibility of the `Request` to read only the data of the request and not further.
///
/// The `Write` object will be used by the `Request` to write the response.
pub fn new_request<R, W>(secure: bool, method: Method, path: String,
                         version: HTTPVersion, headers: Vec<Header>,
                         remote_addr: SocketAddr, mut source_data: R, writer: W)
                         -> Result<Request, RequestCreationError>
                         where R: Read + Send + 'static, W: Write + Send + 'static
{
    // finding the transfer-encoding header
    let transfer_encoding = headers.iter()
        .find(|h: &&Header| h.field.equiv(&"Transfer-Encoding"))
        .map(|h| h.value.clone());

    // finding the content-length header
    let content_length = if transfer_encoding.is_some() {
        // if transfer-encoding is specified, the Content-Length
        // header must be ignored (RFC2616 #4.4)
        None

    } else {
        headers.iter()
               .find(|h: &&Header| h.field.equiv(&"Content-Length"))
               .and_then(|h| FromStr::from_str(h.value.as_str()).ok())
    };

    // true if the client sent a `Expect: 100-continue` header
    let expects_continue = {
        match headers.iter().find(|h: &&Header| h.field.equiv(&"Expect")).map(|h| AsRef::<str>::as_ref(h.value.as_ref())) {
            None => false,
            Some(v) if v.eq_ignore_ascii_case("100-continue")
                => true,
            _ => return Err(RequestCreationError::ExpectationFailed)
        }
    };

    // true if the client sent a `Connection: upgrade` header
    let connection_upgrade = {
        match headers.iter().find(|h: &&Header| h.field.equiv(&"Connection")).map(|h| AsRef::<str>::as_ref(h.value.as_ref())) {
            None => false,
            Some(v) if v.eq_ignore_ascii_case("upgrade")
                => true,
            _ => false
        }
    };

    // we wrap `source_data` around a reading whose nature depends on the transfer-encoding and
    // content-length headers
    let reader =
        if connection_upgrade {
            // if we have a `Connection: upgrade`, always keeping the whole reader
            Box::new(source_data) as Box<Read + Send + 'static>

        } else if let Some(content_length) = content_length {
            if content_length == 0 {
                Box::new(io::empty()) as Box<Read + Send + 'static>

            } else if content_length <= 1024 && !expects_continue {
                // if the content-length is small enough, we just read everything into a buffer

                let mut buffer = vec![0; content_length];
                let mut offset = 0;

                while offset != content_length {
                    let read = try!(source_data.read(&mut buffer[offset..]));
                    if read == 0 {
                        // the socket returned EOF, but we were before the expected content-length
                        // aborting
                        let info = "Connection has been closed before we received enough data";
                        let err = IoError::new(ErrorKind::ConnectionAborted, info);
                        return Err(RequestCreationError::CreationIoError(err));
                    }

                    offset += read;
                }

                Box::new(Cursor::new(buffer)) as Box<Read + Send + 'static>

            } else {
                let (data_reader, _) = EqualReader::new(source_data, content_length);   // TODO:
                Box::new(data_reader) as Box<Read + Send + 'static>
            }

        } else if transfer_encoding.is_some() {
            // if a transfer-encoding was specified, then "chunked" is ALWAYS applied
            // over the message (RFC2616 #3.6)
            Box::new(Decoder::new(source_data)) as Box<Read + Send + 'static>

        } else {
            // if we have neither a Content-Length nor a Transfer-Encoding,
            // assuming that we have no data
            // TODO: could also be multipart/byteranges
            Box::new(io::empty()) as Box<Read + Send + 'static>
        };

    Ok(Request {
        data_reader: Some(reader),
        response_writer: Some(Box::new(writer) as Box<Write + Send + 'static>),
        remote_addr: remote_addr,
        secure: secure,
        method: method,
        path: path,
        http_version: version,
        headers: headers,
        body_length: content_length,
        must_send_continue: expects_continue,
    })
}

impl Request {
    /// Returns true if the request was made through HTTPS.
    #[inline]
    pub fn secure(&self) -> bool {
        self.secure
    }

    /// Returns the method requested by the client (eg. `GET`, `POST`, etc.).
    #[inline]
    pub fn method(&self) -> &Method {
        &self.method
    }

    /// Returns the resource requested by the client.
    #[inline]
    pub fn url(&self) -> &str {
        &self.path
    }

    /// Returns a list of all headers sent by the client.
    #[inline]
    pub fn headers(&self) -> &[Header] {
        &self.headers
    }

    /// Returns the HTTP version of the request.
    #[inline]
    pub fn http_version(&self) -> &HTTPVersion {
        &self.http_version
    }

    /// Returns the length of the body in bytes.
    ///
    /// Returns `None` if the length is unknown.
    #[inline]
    pub fn body_length(&self) -> Option<usize> {
        self.body_length
    }

    /// Returns the length of the body in bytes.
    #[inline]
    pub fn remote_addr(&self) -> &SocketAddr {
        &self.remote_addr
    }

/*      // FIXME: reimplement this
    /// Sends a response with a `Connection: upgrade` header, then turns the `Request` into a `Stream`.
    ///
    /// The main purpose of this function is to support websockets.
    /// If you detect that the request wants to use some kind of protocol upgrade, you can
    ///  call this function to obtain full control of the socket stream.
    ///
    /// If you call this on a non-websocket request, tiny-http will wait until this `Stream` object
    ///  is destroyed before continuing to read or write on the socket. Therefore you should always
    ///  destroy it as soon as possible.
    pub fn upgrade<R: Read>(mut self, protocol: &str, response: Response<R>) -> Box<Read + Write + Send> {
        use util::CustomStream;

        response.raw_print(self.response_writer.as_mut().unwrap().by_ref(), self.http_version,
                           self.headers, false, Some(protocol)).ok();   // TODO: unused result

        self.response_writer.as_mut().unwrap().flush().ok();    // TODO: unused result

        let stream = CustomStream::new(self.into_reader_impl(), self.into_writer_impl());
        Box::new(stream) as Box<Read + Write + Send>
    }*/

    /// Allows to read the body of the request.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # extern crate rustc_serialize;
    /// # extern crate tiny_http;
    /// # use rustc_serialize::json::Json;
    /// # use std::io::Read;
    /// # fn get_content_type(_: &tiny_http::Request) -> &'static str { "" }
    /// # fn main() {
    /// # let server = tiny_http::Server::http("0.0.0.0:0").unwrap();
    /// let mut request = server.recv().unwrap();
    ///
    /// if get_content_type(&request) == "application/json" {
    ///     let mut content = String::new();
    ///     request.as_reader().read_to_string(&mut content).unwrap();
    ///     let json: Json = content.parse().unwrap();
    /// }
    /// # }
    /// ```
    ///
    /// If the client sent a `Expect: 100-continue` header with the request, calling this
    ///  function will send back a `100 Continue` response.
    #[inline]
    pub fn as_reader(&mut self) -> &mut Read {
        if self.must_send_continue {
            let msg = Response::new_empty(StatusCode(100));
            msg.raw_print(self.response_writer.as_mut().unwrap().by_ref(),
                          self.http_version.clone(), &self.headers, true, None).ok();
            self.response_writer.as_mut().unwrap().flush().ok();
            self.must_send_continue = false;
        }

        self.data_reader.as_mut().unwrap()
    }

    /// Turns the `Request` into a writer.
    ///
    /// The writer has a raw access to the stream to the user.
    /// This function is useful for things like CGI.
    ///
    /// Note that the destruction of the `Writer` object may trigger
    /// some events. For exemple if a client has sent multiple requests and the requests
    /// have been processed in parallel, the destruction of a writer will trigger
    /// the writing of the next response.
    /// Therefore you should always destroy the `Writer` as soon as possible.
    #[inline]
    pub fn into_writer(mut self) -> Box<Write + Send + 'static> {
        self.into_writer_impl()
    }

    fn into_writer_impl(&mut self) -> Box<Write + Send + 'static> {
        use std::mem;

        assert!(self.response_writer.is_some());

        let mut writer = None;
        mem::swap(&mut self.response_writer, &mut writer);
        writer.unwrap()
    }

    // TODO: unused
    /*fn into_reader_impl(&mut self) -> AnyReader {
        use std::mem;

        assert!(self.data_reader.is_some());

        let mut reader = None;
        mem::swap(&mut self.data_reader, &mut reader);
        reader.unwrap()
    }*/

    /// Sends a response to this request.
    #[inline]
    pub fn respond<R>(mut self, response: Response<R>) -> Result<(), IoError>
        where R: Read
    {
        self.respond_impl(response)
    }

    fn respond_impl<R>(&mut self, response: Response<R>) -> Result<(), IoError>
        where R: Read
    {
        let mut writer = self.into_writer_impl();

        let do_not_send_body = self.method == Method::Head;

        match response.raw_print(writer.by_ref(),
                                 self.http_version.clone(), &self.headers,
                                 do_not_send_body, None)
        {
            Ok(_) => (),
            Err(ref err) if err.kind() == ErrorKind::BrokenPipe => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionAborted => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionRefused => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionReset => (),
            Err(err) => return Err(err)
        };

        writer.flush()
    }
}

impl fmt::Debug for Request {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Request({} {} from {})", self.method, self.path, self.remote_addr)
    }
}

impl Drop for Request {
    fn drop(&mut self) {
        if self.response_writer.is_some() {
            let response = Response::empty(500);
            let _ = self.respond_impl(response);        // ignoring any potential error
        }
    }
}

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

    #[test]
    fn must_be_send() {
        #![allow(dead_code)]
        fn f<T: Send>(_: &T) {}
        fn bar(rq: &Request) { f(rq); }
    }
}
<commit_msg>Update request.rs<commit_after>// Copyright 2015 The tiny-http Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ascii::AsciiExt;

use std::io::Error as IoError;
use std::io::{self, Cursor, Read, Write, ErrorKind};

use std::net::SocketAddr;
use std::fmt;
use std::str::FromStr;

use {Header, HTTPVersion, Method, Response, StatusCode};
use util::EqualReader;
use chunked_transfer::Decoder;

/// Represents an HTTP request made by a client.
///
/// A `Request` object is what is produced by the server, and is your what
/// your code must analyse and answer.
///
/// This object implements the `Send` trait, therefore you can dispatch your requests to
/// worker threads.
///
/// # Pipelining
///
/// If a client sends multiple requests in a row (without waiting for the response), then you will
/// get multiple `Request` objects simultaneously. This is called *requests pipelining*.
/// Tiny-http automatically reorders the responses so that you don't need to worry about the order
/// in which you call `respond` or `into_writer`.
///
/// This mechanic is disabled if:
///
///  - The body of a request is large enough (handling requires pipelining requires storing the
///    body of the request in a buffer ; if the body is too big, tiny-http will avoid doing that)
///  - A request sends a `Expect: 100-continue` header (which means that the client waits to
///    know whether its body will be processed before sending it)
///  - A request sends a `Connection: close` header or `Connection: upgrade` header (used for
///    websockets), which indicates that this is the last request that will be received on this
///    connection
///
/// # Automatic cleanup
///
/// If a `Request` object is destroyed without `into_writer` or `respond` being called,
/// an empty response with a 500 status code (internal server error) will automatically be
/// sent back to the client.
/// This means that if your code fails during the handling of a request, this "internal server
/// error" response will automatically be sent during the stack unwinding.
pub struct Request {
    // where to read the body from
    data_reader: Option<Box<Read + Send + 'static>>,

    // if this writer is empty, then the request has been answered
    response_writer: Option<Box<Write + Send + 'static>>,

    remote_addr: SocketAddr,

    // true if HTTPS, false if HTTP
    secure: bool,

    method: Method,

    path: String,

    http_version: HTTPVersion,

    headers: Vec<Header>,

    body_length: Option<usize>,

    // true if a `100 Continue` response must be sent when `as_reader()` is called
    must_send_continue: bool,
}

/// Error that can happen when building a `Request` object.
pub enum RequestCreationError {
    /// The client sent an `Expect` header that was not recognized by tiny-http.
    ExpectationFailed,

    /// Error while reading data from the socket during the creation of the `Request`.
    CreationIoError(IoError),
}

impl From<IoError> for RequestCreationError {
    fn from(err: IoError) -> RequestCreationError {
        RequestCreationError::CreationIoError(err)
    }
}

/// Builds a new request.
///
/// After the request line and headers have been read from the socket, a new `Request` object
/// is built.
///
/// You must pass a `Read` that will allow the `Request` object to read from the incoming data.
/// It is the responsibility of the `Request` to read only the data of the request and not further.
///
/// The `Write` object will be used by the `Request` to write the response.
pub fn new_request<R, W>(secure: bool, method: Method, path: String,
                         version: HTTPVersion, headers: Vec<Header>,
                         remote_addr: SocketAddr, mut source_data: R, writer: W)
                         -> Result<Request, RequestCreationError>
                         where R: Read + Send + 'static, W: Write + Send + 'static
{
    // finding the transfer-encoding header
    let transfer_encoding = headers.iter()
        .find(|h: &&Header| h.field.equiv(&"Transfer-Encoding"))
        .map(|h| h.value.clone());

    // finding the content-length header
    let content_length = if transfer_encoding.is_some() {
        // if transfer-encoding is specified, the Content-Length
        // header must be ignored (RFC2616 #4.4)
        None

    } else {
        headers.iter()
               .find(|h: &&Header| h.field.equiv(&"Content-Length"))
               .and_then(|h| FromStr::from_str(h.value.as_str()).ok())
    };

    // true if the client sent a `Expect: 100-continue` header
    let expects_continue = {
        match headers.iter().find(|h: &&Header| h.field.equiv(&"Expect")).map(|h| AsRef::<str>::as_ref(h.value.as_ref())) {
            None => false,
            Some(v) if v.eq_ignore_ascii_case("100-continue")
                => true,
            _ => return Err(RequestCreationError::ExpectationFailed)
        }
    };

    // true if the client sent a `Connection: upgrade` header
    let connection_upgrade = {
        match headers.iter().find(|h: &&Header| h.field.equiv(&"Connection")).map(|h| AsRef::<str>::as_ref(h.value.as_ref())) {
            None => false,
            Some(v) if v.eq_ignore_ascii_case("upgrade")
                => true,
            _ => false
        }
    };

    // we wrap `source_data` around a reading whose nature depends on the transfer-encoding and
    // content-length headers
    let reader =
        if connection_upgrade {
            // if we have a `Connection: upgrade`, always keeping the whole reader
            Box::new(source_data) as Box<Read + Send + 'static>

        } else if let Some(content_length) = content_length {
            if content_length == 0 {
                Box::new(io::empty()) as Box<Read + Send + 'static>

            } else if content_length <= 1024 && !expects_continue {
                // if the content-length is small enough, we just read everything into a buffer

                let mut buffer = vec![0; content_length];
                let mut offset = 0;

                while offset != content_length {
                    let read = try!(source_data.read(&mut buffer[offset..]));
                    if read == 0 {
                        // the socket returned EOF, but we were before the expected content-length
                        // aborting
                        let info = "Connection has been closed before we received enough data";
                        let err = IoError::new(ErrorKind::ConnectionAborted, info);
                        return Err(RequestCreationError::CreationIoError(err));
                    }

                    offset += read;
                }

                Box::new(Cursor::new(buffer)) as Box<Read + Send + 'static>

            } else {
                let (data_reader, _) = EqualReader::new(source_data, content_length);   // TODO:
                Box::new(data_reader) as Box<Read + Send + 'static>
            }

        } else if transfer_encoding.is_some() {
            // if a transfer-encoding was specified, then "chunked" is ALWAYS applied
            // over the message (RFC2616 #3.6)
            Box::new(Decoder::new(source_data)) as Box<Read + Send + 'static>

        } else {
            // if we have neither a Content-Length nor a Transfer-Encoding,
            // assuming that we have no data
            // TODO: could also be multipart/byteranges
            Box::new(io::empty()) as Box<Read + Send + 'static>
        };

    Ok(Request {
        data_reader: Some(reader),
        response_writer: Some(Box::new(writer) as Box<Write + Send + 'static>),
        remote_addr: remote_addr,
        secure: secure,
        method: method,
        path: path,
        http_version: version,
        headers: headers,
        body_length: content_length,
        must_send_continue: expects_continue,
    })
}

impl Request {
    /// Returns true if the request was made through HTTPS.
    #[inline]
    pub fn secure(&self) -> bool {
        self.secure
    }

    /// Returns the method requested by the client (eg. `GET`, `POST`, etc.).
    #[inline]
    pub fn method(&self) -> &Method {
        &self.method
    }

    /// Returns the resource requested by the client.
    #[inline]
    pub fn url(&self) -> &str {
        &self.path
    }

    /// Returns a list of all headers sent by the client.
    #[inline]
    pub fn headers(&self) -> &[Header] {
        &self.headers
    }

    /// Returns the HTTP version of the request.
    #[inline]
    pub fn http_version(&self) -> &HTTPVersion {
        &self.http_version
    }

    /// Returns the length of the body in bytes.
    ///
    /// Returns `None` if the length is unknown.
    #[inline]
    pub fn body_length(&self) -> Option<usize> {
        self.body_length
    }

    /// Returns the address of the client that sent this request.
    ///
    /// Note that this is gathered from the socket. If you receive the request from a proxy,
    /// this function will return the address of the proxy and not the address of the actual
    /// user.
    #[inline]
    pub fn remote_addr(&self) -> &SocketAddr {
        &self.remote_addr
    }

/*      // FIXME: reimplement this
    /// Sends a response with a `Connection: upgrade` header, then turns the `Request` into a `Stream`.
    ///
    /// The main purpose of this function is to support websockets.
    /// If you detect that the request wants to use some kind of protocol upgrade, you can
    ///  call this function to obtain full control of the socket stream.
    ///
    /// If you call this on a non-websocket request, tiny-http will wait until this `Stream` object
    ///  is destroyed before continuing to read or write on the socket. Therefore you should always
    ///  destroy it as soon as possible.
    pub fn upgrade<R: Read>(mut self, protocol: &str, response: Response<R>) -> Box<Read + Write + Send> {
        use util::CustomStream;

        response.raw_print(self.response_writer.as_mut().unwrap().by_ref(), self.http_version,
                           self.headers, false, Some(protocol)).ok();   // TODO: unused result

        self.response_writer.as_mut().unwrap().flush().ok();    // TODO: unused result

        let stream = CustomStream::new(self.into_reader_impl(), self.into_writer_impl());
        Box::new(stream) as Box<Read + Write + Send>
    }*/

    /// Allows to read the body of the request.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # extern crate rustc_serialize;
    /// # extern crate tiny_http;
    /// # use rustc_serialize::json::Json;
    /// # use std::io::Read;
    /// # fn get_content_type(_: &tiny_http::Request) -> &'static str { "" }
    /// # fn main() {
    /// # let server = tiny_http::Server::http("0.0.0.0:0").unwrap();
    /// let mut request = server.recv().unwrap();
    ///
    /// if get_content_type(&request) == "application/json" {
    ///     let mut content = String::new();
    ///     request.as_reader().read_to_string(&mut content).unwrap();
    ///     let json: Json = content.parse().unwrap();
    /// }
    /// # }
    /// ```
    ///
    /// If the client sent a `Expect: 100-continue` header with the request, calling this
    ///  function will send back a `100 Continue` response.
    #[inline]
    pub fn as_reader(&mut self) -> &mut Read {
        if self.must_send_continue {
            let msg = Response::new_empty(StatusCode(100));
            msg.raw_print(self.response_writer.as_mut().unwrap().by_ref(),
                          self.http_version.clone(), &self.headers, true, None).ok();
            self.response_writer.as_mut().unwrap().flush().ok();
            self.must_send_continue = false;
        }

        self.data_reader.as_mut().unwrap()
    }

    /// Turns the `Request` into a writer.
    ///
    /// The writer has a raw access to the stream to the user.
    /// This function is useful for things like CGI.
    ///
    /// Note that the destruction of the `Writer` object may trigger
    /// some events. For exemple if a client has sent multiple requests and the requests
    /// have been processed in parallel, the destruction of a writer will trigger
    /// the writing of the next response.
    /// Therefore you should always destroy the `Writer` as soon as possible.
    #[inline]
    pub fn into_writer(mut self) -> Box<Write + Send + 'static> {
        self.into_writer_impl()
    }

    fn into_writer_impl(&mut self) -> Box<Write + Send + 'static> {
        use std::mem;

        assert!(self.response_writer.is_some());

        let mut writer = None;
        mem::swap(&mut self.response_writer, &mut writer);
        writer.unwrap()
    }

    // TODO: unused
    /*fn into_reader_impl(&mut self) -> AnyReader {
        use std::mem;

        assert!(self.data_reader.is_some());

        let mut reader = None;
        mem::swap(&mut self.data_reader, &mut reader);
        reader.unwrap()
    }*/

    /// Sends a response to this request.
    #[inline]
    pub fn respond<R>(mut self, response: Response<R>) -> Result<(), IoError>
        where R: Read
    {
        self.respond_impl(response)
    }

    fn respond_impl<R>(&mut self, response: Response<R>) -> Result<(), IoError>
        where R: Read
    {
        let mut writer = self.into_writer_impl();

        let do_not_send_body = self.method == Method::Head;

        match response.raw_print(writer.by_ref(),
                                 self.http_version.clone(), &self.headers,
                                 do_not_send_body, None)
        {
            Ok(_) => (),
            Err(ref err) if err.kind() == ErrorKind::BrokenPipe => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionAborted => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionRefused => (),
            Err(ref err) if err.kind() == ErrorKind::ConnectionReset => (),
            Err(err) => return Err(err)
        };

        writer.flush()
    }
}

impl fmt::Debug for Request {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Request({} {} from {})", self.method, self.path, self.remote_addr)
    }
}

impl Drop for Request {
    fn drop(&mut self) {
        if self.response_writer.is_some() {
            let response = Response::empty(500);
            let _ = self.respond_impl(response);        // ignoring any potential error
        }
    }
}

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

    #[test]
    fn must_be_send() {
        #![allow(dead_code)]
        fn f<T: Send>(_: &T) {}
        fn bar(rq: &Request) { f(rq); }
    }
}
<|endoftext|>


<commit_before>//! Module for dealing with RFC2045 style headers.
use super::rfc5322::Rfc5322Parser;

use std::collections::HashMap;

/// Parser over RFC 2045 style headers.
///
/// Things of the style `value; param1=foo; param2="bar"`
pub struct Rfc2045Parser<'s> {
    parser: Rfc5322Parser<'s>,
}

impl<'s> Rfc2045Parser<'s> {
    /// Create a new parser over `s`
    pub fn new(s: &str) -> Rfc2045Parser {
        Rfc2045Parser {
            parser: Rfc5322Parser::new(s)
        }
    }

    fn consume_token(&mut self) -> Option<String> {
        let token = self.parser.consume_while(|c| {
            match c {
                // Not any tspecials
                '(' | ')' | '<' | '>' | '@' | ',' | ';' | ':' |
                    '\\' | '\"' | '/' | '[' | ']' | '?' | '=' => false,
                '!'...'~' => true,
                _ => false,
            }
        });

        if token.len() > 0 {
            Some(token)
        } else {
            None
        }
    }

    /// Consume up to all of the input into the value and a hashmap
    /// over parameters to values.
    pub fn consume_all(&mut self) -> (String, HashMap<String, String>) {
        let value = self.parser.consume_while(|c| { c != ';' });

        // Find the parameters
        let mut params = HashMap::new();
        while !self.parser.eof() {
            // Eat the ; and any whitespace
            assert_eq!(self.parser.consume_char(), Some(';'));
            self.parser.consume_linear_whitespace();

            let attribute = self.consume_token();
            assert_eq!(self.parser.consume_char(), Some('='));
            // Value can be token or quoted-string
            let value = if self.parser.peek() == '"' {
                self.parser.consume_quoted_string()
            } else {
                self.consume_token()
            };

            match (attribute, value) {
                (Some(attrib), Some(val)) => {
                    params.insert(attrib, val);
                },
                _ => {}
            }
        }

        (value, params)
    }
}

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

    use std::collections::HashMap;

    struct ParserTestCase<'s> {
        input: &'s str,
        output: (&'s str, Vec<(&'s str, &'s str)>),
        name: &'s str,
    }

    #[test]
    pub fn test_foo() {
        let tests = vec![
            ParserTestCase {
                input: "foo/bar",
                output: ("foo/bar", vec![]),
                name: "Basic value",
            },

            ParserTestCase {
                input: "foo/bar; foo=bar",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                ]),
                name: "Basic value with parameter",
            },

            ParserTestCase {
                input: "foo/bar; foo=\"bar\"",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                ]),
                name: "Basic value with quoted parameter",
            },

            ParserTestCase {
                input: "foo/bar; foo=\"bar\"; baz=qux",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                    ("baz", "qux"),
                ]),
                name: "Multiple values",
            },
        ];

        for test in tests.into_iter() {
            let (expected_value, expected_param_list) = test.output;
            let mut expected_params = HashMap::new();
            for &(param_name, param_value) in expected_param_list.iter() {
                expected_params.insert(param_name.to_string(), param_value.to_string());
            }

            let mut parser = Rfc2045Parser::new(test.input);
            let (value, parameters) = parser.consume_all();

            assert!(value == expected_value.to_string(), test.name);
            assert!(parameters == expected_params, test.name);
        }
    }
}
<commit_msg>Handle spaces in Content-Type header parameters<commit_after>//! Module for dealing with RFC2045 style headers.
use super::rfc5322::Rfc5322Parser;

use std::collections::HashMap;

/// Parser over RFC 2045 style headers.
///
/// Things of the style `value; param1=foo; param2="bar"`
pub struct Rfc2045Parser<'s> {
    parser: Rfc5322Parser<'s>,
}

impl<'s> Rfc2045Parser<'s> {
    /// Create a new parser over `s`
    pub fn new(s: &str) -> Rfc2045Parser {
        Rfc2045Parser {
            parser: Rfc5322Parser::new(s)
        }
    }

    fn consume_token(&mut self) -> Option<String> {
        let token = self.parser.consume_while(|c| {
            match c {
                // Not any tspecials
                '(' | ')' | '<' | '>' | '@' | ',' | ';' | ':' |
                    '\\' | '\"' | '/' | '[' | ']' | '?' | '=' => false,
                '!'...'~' => true,
                _ => false,
            }
        });

        if token.len() > 0 {
            Some(token)
        } else {
            None
        }
    }

    /// Consume up to all of the input into the value and a hashmap
    /// over parameters to values.
    pub fn consume_all(&mut self) -> (String, HashMap<String, String>) {
        let value = self.parser.consume_while(|c| { c != ';' });

        // Find the parameters
        let mut params = HashMap::new();
        while !self.parser.eof() {
            // Eat the ; and any whitespace
            assert_eq!(self.parser.consume_char(), Some(';'));
            self.parser.consume_linear_whitespace();

            let attribute = self.consume_token();
            self.parser.consume_linear_whitespace();
            assert_eq!(self.parser.consume_char(), Some('='));
            self.parser.consume_linear_whitespace();
            // Value can be token or quoted-string
            let value = if self.parser.peek() == '"' {
                self.parser.consume_quoted_string()
            } else {
                self.consume_token()
            };

            match (attribute, value) {
                (Some(attrib), Some(val)) => {
                    params.insert(attrib, val);
                },
                _ => {}
            }
        }

        (value, params)
    }
}

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

    use std::collections::HashMap;

    struct ParserTestCase<'s> {
        input: &'s str,
        output: (&'s str, Vec<(&'s str, &'s str)>),
        name: &'s str,
    }

    #[test]
    pub fn test_foo() {
        let tests = vec![
            ParserTestCase {
                input: "foo/bar",
                output: ("foo/bar", vec![]),
                name: "Basic value",
            },

            ParserTestCase {
                input: "foo/bar; foo=bar",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                ]),
                name: "Basic value with parameter",
            },

            ParserTestCase {
                input: "foo/bar; foo=\"bar\"",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                ]),
                name: "Basic value with quoted parameter",
            },

            ParserTestCase {
                input: "foo/bar; foo=\"bar\"; baz=qux",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                    ("baz", "qux"),
                ]),
                name: "Multiple values",
            },

            ParserTestCase {
                input: "foo/bar; foo = \"bar\"; baz=qux",
                output: ("foo/bar", vec![
                    ("foo", "bar"),
                    ("baz", "qux"),
                ]),
                name: "Parameter with space",
            },
        ];

        for test in tests.into_iter() {
            let (expected_value, expected_param_list) = test.output;
            let mut expected_params = HashMap::new();
            for &(param_name, param_value) in expected_param_list.iter() {
                expected_params.insert(param_name.to_string(), param_value.to_string());
            }

            let mut parser = Rfc2045Parser::new(test.input);
            let (value, parameters) = parser.consume_all();

            assert!(value == expected_value.to_string(), test.name);
            assert!(parameters == expected_params, test.name);
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>rustuv: Remove an old file<commit_after><|endoftext|>


<commit_before>use std::{
    fs::File,
    io::{
        BufRead,
        BufReader,
        Read,
    },
    path::Path,
};

use std::marker::PhantomData;

use anyhow::{
    Context,
    Result,
};

use serde::de::DeserializeOwned;

pub struct JsonlIterator<R: Read, T: DeserializeOwned> {
    reader: BufReader<R>,
    buf: String,
    phantom: PhantomData<T>,
}

impl<R: Read, T: DeserializeOwned> JsonlIterator<R, T> {
    pub fn from_read(reader: R) -> Self {
        Self {
            reader: BufReader::new(reader),
            buf: String::new(),
            phantom: PhantomData,
        }
    }
}

impl<T: DeserializeOwned> JsonlIterator<File, T> {
    pub fn from_path(path: &Path) -> Result<Self> {
        let file = File::open(path).with_context(|| format!("Could not open file {:?}", &path))?;
        Ok(JsonlIterator::from_read(file))
    }
}

impl<R: Read, T: DeserializeOwned> Iterator for JsonlIterator<R, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match self.reader.read_line(&mut self.buf).unwrap() {
            0 => None,
            _ => {
                let value: T = serde_json::from_str(&self.buf)
                    .with_context(|| format!("Could not parse line {}", &self.buf))
                    .unwrap();
                self.buf.clear();
                Some(value)
            },
        }
    }
}
<commit_msg>Try to handle escaped lines as well<commit_after>use std::{
    fs::File,
    io::{
        BufRead,
        BufReader,
        Read,
    },
    path::Path,
};

use std::marker::PhantomData;

use anyhow::{
    Context,
    Result,
};

use serde::de::DeserializeOwned;

pub struct JsonlIterator<R: Read, T: DeserializeOwned> {
    reader: BufReader<R>,
    buf: String,
    phantom: PhantomData<T>,
}

impl<R: Read, T: DeserializeOwned> JsonlIterator<R, T> {
    pub fn from_read(reader: R) -> Self {
        Self {
            reader: BufReader::new(reader),
            buf: String::new(),
            phantom: PhantomData,
        }
    }
}

impl<T: DeserializeOwned> JsonlIterator<File, T> {
    pub fn from_path(path: &Path) -> Result<Self> {
        let file = File::open(path).with_context(|| format!("Could not open file {:?}", &path))?;
        Ok(JsonlIterator::from_read(file))
    }
}

impl<R: Read, T: DeserializeOwned> Iterator for JsonlIterator<R, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        match self.reader.read_line(&mut self.buf).unwrap() {
            0 => None,
            _ => {
                let value: T = serde_json::from_str(&self.buf)
                    .or(serde_json::from_str(&self.buf.replace("////", "//")))
                    .with_context(|| format!("Could not parse line {}", &self.buf))
                    .unwrap();
                self.buf.clear();
                Some(value)
            },
        }
    }
}
<|endoftext|>


<commit_before>use std::io;

use serde::{ser, Serialize};

use super::{Error, Result, Serializer};

pub struct RawValueSerializer<'a, W: io::Write> {
    ser: &'a mut Serializer<W>,
}

impl<'a, W: io::Write> RawValueSerializer<'a, W> {
    pub fn new(ser: &'a mut Serializer<W>) -> Self {
        Self { ser }
    }
}

impl<'a, W: io::Write> ser::Serializer for RawValueSerializer<'a, W> {
    type Error = Error;
    type Ok = ();
    type SerializeMap = ser::Impossible<(), Error>;
    type SerializeSeq = ser::Impossible<(), Error>;
    type SerializeStruct = ser::Impossible<(), Error>;
    type SerializeStructVariant = ser::Impossible<(), Error>;
    type SerializeTuple = ser::Impossible<(), Error>;
    type SerializeTupleStruct = ser::Impossible<(), Error>;
    type SerializeTupleVariant = ser::Impossible<(), Error>;

    fn serialize_bool(self, _: bool) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i8(self, _: i8) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i16(self, _: i16) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i32(self, _: i32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i64(self, _: i64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    #[cfg(feature = "integer128")]
    fn serialize_i128(self, _: i128) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u8(self, _: u8) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u16(self, _: u16) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u32(self, _: u32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u64(self, _: u64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    #[cfg(feature = "integer128")]
    fn serialize_u128(self, _: u128) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_f32(self, _: f32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_f64(self, _: f64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_char(self, _: char) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_str(self, ron: &str) -> Result<()> {
        self.ser.output.write_all(ron.as_bytes())?;
        Ok(())
    }

    fn serialize_bytes(self, _: &[u8]) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_none(self) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_some<T: ?Sized + Serialize>(self, _: &T) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit(self) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit_struct(self, _: &'static str) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit_variant(self, _: &'static str, _: u32, _: &'static str) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_newtype_struct<T: ?Sized + Serialize>(self, _: &'static str, _: &T) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_newtype_variant<T: ?Sized + Serialize>(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: &T,
    ) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_seq(self, _: Option<usize>) -> Result<Self::SerializeSeq> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple(self, _: usize) -> Result<Self::SerializeTuple> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple_struct(
        self,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeTupleStruct> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple_variant(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeTupleVariant> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_map(self, _: Option<usize>) -> Result<Self::SerializeMap> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_struct(self, _: &'static str, _: usize) -> Result<Self::SerializeStruct> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_struct_variant(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeStructVariant> {
        Err(Error::ExpectedRawValue)
    }
}
<commit_msg>Add tests for raw value serialising (#419)<commit_after>use std::io;

use serde::{ser, Serialize};

use super::{Error, Result, Serializer};

pub struct RawValueSerializer<'a, W: io::Write> {
    ser: &'a mut Serializer<W>,
}

impl<'a, W: io::Write> RawValueSerializer<'a, W> {
    pub fn new(ser: &'a mut Serializer<W>) -> Self {
        Self { ser }
    }
}

impl<'a, W: io::Write> ser::Serializer for RawValueSerializer<'a, W> {
    type Error = Error;
    type Ok = ();
    type SerializeMap = ser::Impossible<(), Error>;
    type SerializeSeq = ser::Impossible<(), Error>;
    type SerializeStruct = ser::Impossible<(), Error>;
    type SerializeStructVariant = ser::Impossible<(), Error>;
    type SerializeTuple = ser::Impossible<(), Error>;
    type SerializeTupleStruct = ser::Impossible<(), Error>;
    type SerializeTupleVariant = ser::Impossible<(), Error>;

    fn serialize_bool(self, _: bool) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i8(self, _: i8) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i16(self, _: i16) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i32(self, _: i32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_i64(self, _: i64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    #[cfg(feature = "integer128")]
    fn serialize_i128(self, _: i128) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u8(self, _: u8) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u16(self, _: u16) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u32(self, _: u32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_u64(self, _: u64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    #[cfg(feature = "integer128")]
    fn serialize_u128(self, _: u128) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_f32(self, _: f32) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_f64(self, _: f64) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_char(self, _: char) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_str(self, ron: &str) -> Result<()> {
        self.ser.output.write_all(ron.as_bytes())?;
        Ok(())
    }

    fn serialize_bytes(self, _: &[u8]) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_none(self) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_some<T: ?Sized + Serialize>(self, _: &T) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit(self) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit_struct(self, _: &'static str) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_unit_variant(self, _: &'static str, _: u32, _: &'static str) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_newtype_struct<T: ?Sized + Serialize>(self, _: &'static str, _: &T) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_newtype_variant<T: ?Sized + Serialize>(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: &T,
    ) -> Result<()> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_seq(self, _: Option<usize>) -> Result<Self::SerializeSeq> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple(self, _: usize) -> Result<Self::SerializeTuple> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple_struct(
        self,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeTupleStruct> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_tuple_variant(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeTupleVariant> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_map(self, _: Option<usize>) -> Result<Self::SerializeMap> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_struct(self, _: &'static str, _: usize) -> Result<Self::SerializeStruct> {
        Err(Error::ExpectedRawValue)
    }

    fn serialize_struct_variant(
        self,
        _: &'static str,
        _: u32,
        _: &'static str,
        _: usize,
    ) -> Result<Self::SerializeStructVariant> {
        Err(Error::ExpectedRawValue)
    }
}

#[cfg(test)]
mod tests {
    macro_rules! test_non_raw_value {
        ($test_name:ident => $serialize_method:ident($($serialize_param:expr),*)) => {
            #[test]
            fn $test_name() {
                use serde::{Serialize, Serializer};

                struct Inner;

                impl Serialize for Inner {
                    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
                        serializer.$serialize_method($($serialize_param),*).map(|_| unreachable!())
                    }
                }

                #[derive(Debug)]
                struct Newtype;

                impl Serialize for Newtype {
                    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
                        serializer.serialize_newtype_struct(
                            crate::value::raw::RAW_VALUE_TOKEN, &Inner,
                        )
                    }
                }

                assert_eq!(
                    crate::to_string(&Newtype).unwrap_err(),
                    crate::Error::ExpectedRawValue
                )
            }
        };
    }

    test_non_raw_value! { test_bool => serialize_bool(false) }
    test_non_raw_value! { test_i8 => serialize_i8(0) }
    test_non_raw_value! { test_i16 => serialize_i16(0) }
    test_non_raw_value! { test_i32 => serialize_i32(0) }
    test_non_raw_value! { test_i64 => serialize_i64(0) }
    #[cfg(feature = "integer128")]
    test_non_raw_value! { test_i128 => serialize_i128(0) }
    test_non_raw_value! { test_u8 => serialize_u8(0) }
    test_non_raw_value! { test_u16 => serialize_u16(0) }
    test_non_raw_value! { test_u32 => serialize_u32(0) }
    test_non_raw_value! { test_u64 => serialize_u64(0) }
    #[cfg(feature = "integer128")]
    test_non_raw_value! { test_u128 => serialize_u128(0) }
    test_non_raw_value! { test_f32 => serialize_f32(0.0) }
    test_non_raw_value! { test_f64 => serialize_f64(0.0) }
    test_non_raw_value! { test_char => serialize_char('\0') }
    test_non_raw_value! { test_bytes => serialize_bytes(b"") }
    test_non_raw_value! { test_none => serialize_none() }
    test_non_raw_value! { test_some => serialize_some(&()) }
    test_non_raw_value! { test_unit => serialize_unit() }
    test_non_raw_value! { test_unit_struct => serialize_unit_struct("U") }
    test_non_raw_value! { test_unit_variant => serialize_unit_variant("E", 0, "U") }
    test_non_raw_value! { test_newtype_struct => serialize_newtype_struct("N", &()) }
    test_non_raw_value! { test_newtype_variant => serialize_newtype_variant("E", 0, "N", &()) }
    test_non_raw_value! { test_seq => serialize_seq(None) }
    test_non_raw_value! { test_tuple => serialize_tuple(0) }
    test_non_raw_value! { test_tuple_struct => serialize_tuple_struct("T", 0) }
    test_non_raw_value! { test_tuple_variant => serialize_tuple_variant("E", 0, "T", 0) }
    test_non_raw_value! { test_map => serialize_map(None) }
    test_non_raw_value! { test_struct => serialize_struct("S", 0) }
    test_non_raw_value! { test_struct_variant => serialize_struct_variant("E", 0, "S", 0) }
}
<|endoftext|>


<commit_before>use core::cell::Cell;
use core::mem;
use syscalls;
use alloc::heap::Heap;
use alloc::String;
use alloc::boxed::Box;
use alloc::raw_vec::RawVec;
use alloc::allocator::{Alloc, Layout};

const DRIVER_NUM: u32 = 0xff;

pub struct Client {
    pid: u32
}

impl Client {
    pub fn new(pkg_name: String) -> Result<Client, ()> {
        unsafe {
            let res = discover(pkg_name)?;
            Ok(Client {
                pid: res as u32
            })
        }
    }

    pub fn share(&mut self, mut len: usize) -> Result<Box<[u8]>, ()> {
        if len < 3 {
            len = 3;
        }
        unsafe {
            let shared_val = Heap.alloc_zeroed(
                    Layout::from_size_align(1 << len, 1 << len).unwrap());

            match shared_val {
                Ok(v) => {
                    share(self.pid, v, 1 << len)?;
                    Ok(RawVec::from_raw_parts(v, 1 << len).into_box())
                },
                _ => {
                    Err(())
                }
            }
        }
    }

    pub fn notify(&mut self) -> Result<(), ()> {
        let done = Cell::new(false);
        let ptr = &done as *const _ as usize;
        unsafe {
            if syscalls::subscribe(DRIVER_NUM, self.pid, Self::cb, ptr) < 0 {
                return Err(())
            }
            if syscalls::command(DRIVER_NUM, self.pid, 0) < 0 {
                return Err(())
            }
        }
        syscalls::yieldk_for(|| done.get());
        Ok(())
    }

    extern fn cb(_: usize, _: usize, _: usize, ptr: usize) {
        let done: &Cell<bool> = unsafe { mem::transmute(ptr) };
        done.set(true);

    }
}

unsafe fn discover(pkg_name: String) -> Result<isize, ()> {
    let res = syscalls::allow(DRIVER_NUM, 0, pkg_name.as_bytes());
    if res < 0 {
        Err(())
    } else {
        Ok(res)
    }
}

unsafe fn share(pid: u32, base: *mut u8, len: usize) -> Result<(), ()> {
    use core::slice::from_raw_parts;
    let res = syscalls::allow(DRIVER_NUM, pid, from_raw_parts(base, len));
    if res < 0 {
        Err(())
    } else {
        Ok(())
    }
}

<commit_msg>breakout ipc::notify into notify_async and subscribe<commit_after>use core::cell::Cell;
use core::mem;
use syscalls;
use alloc::heap::Heap;
use alloc::String;
use alloc::boxed::Box;
use alloc::raw_vec::RawVec;
use alloc::allocator::{Alloc, Layout};

const DRIVER_NUM: u32 = 0xff;

pub struct Client {
    pid: u32
}

impl Client {
    pub fn new(pkg_name: String) -> Result<Client, ()> {
        unsafe {
            let res = discover(pkg_name)?;
            Ok(Client {
                pid: res as u32
            })
        }
    }

    pub fn share(&mut self, mut len: usize) -> Result<Box<[u8]>, ()> {
        if len < 3 {
            len = 3;
        }
        unsafe {
            let shared_val = Heap.alloc_zeroed(
                    Layout::from_size_align(1 << len, 1 << len).unwrap());

            match shared_val {
                Ok(v) => {
                    share(self.pid, v, 1 << len)?;
                    Ok(RawVec::from_raw_parts(v, 1 << len).into_box())
                },
                _ => {
                    Err(())
                }
            }
        }
    }

    pub unsafe fn notify_async(&mut self) -> Result<(), ()> {
        if syscalls::command(DRIVER_NUM, self.pid, 0) < 0 {
            return Err(())
        }
        Ok(())
    }

    pub unsafe fn subscribe(&mut self,
            cb: extern fn(_: usize, _: usize, _: usize, ptr: usize), ud: usize)
                -> Result<(), ()> {
        if syscalls::subscribe(DRIVER_NUM, self.pid, cb, ud) < 0 {
            return Err(())
        }
        Ok(())
    }

    pub fn notify(&mut self) -> Result<(), ()> {
        let done = Cell::new(false);
        let ptr = &done as *const _ as usize;
        unsafe {
            self.subscribe(Self::cb, ptr)?;
            self.notify_async()?;
        }
        syscalls::yieldk_for(|| done.get());
        Ok(())
    }

    extern fn cb(_: usize, _: usize, _: usize, ptr: usize) {
        let done: &Cell<bool> = unsafe { mem::transmute(ptr) };
        done.set(true);

    }
}

unsafe fn discover(pkg_name: String) -> Result<isize, ()> {
    let res = syscalls::allow(DRIVER_NUM, 0, pkg_name.as_bytes());
    if res < 0 {
        Err(())
    } else {
        Ok(res)
    }
}

unsafe fn share(pid: u32, base: *mut u8, len: usize) -> Result<(), ()> {
    use core::slice::from_raw_parts;
    let res = syscalls::allow(DRIVER_NUM, pid, from_raw_parts(base, len));
    if res < 0 {
        Err(())
    } else {
        Ok(())
    }
}

<|endoftext|>


<commit_before>use std::{
    collections::{hash_map::Entry, HashMap},
    net::SocketAddr,
    ops::Deref,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};

use byte_string::ByteStr;
use kcp::KcpResult;
use log::{error, trace};
use spin::Mutex as SpinMutex;
use tokio::{
    net::UdpSocket,
    sync::{mpsc, Notify},
    time::{self, Instant},
};

use crate::{skcp::KcpSocket, KcpConfig};

pub struct KcpSession {
    socket: SpinMutex<KcpSocket>,
    closed: AtomicBool,
    session_expire: Duration,
    session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
    input_tx: mpsc::Sender<Vec<u8>>,
    notifier: Notify,
}

impl KcpSession {
    fn new(
        socket: KcpSocket,
        session_expire: Duration,
        session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
        input_tx: mpsc::Sender<Vec<u8>>,
    ) -> KcpSession {
        KcpSession {
            socket: SpinMutex::new(socket),
            closed: AtomicBool::new(false),
            session_expire,
            session_close_notifier,
            input_tx,
            notifier: Notify::new(),
        }
    }

    pub fn new_shared(
        socket: KcpSocket,
        session_expire: Duration,
        session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
    ) -> Arc<KcpSession> {
        let is_client = session_close_notifier.is_none();

        let (input_tx, mut input_rx) = mpsc::channel(64);

        let udp_socket = socket.udp_socket().clone();

        let session = Arc::new(KcpSession::new(
            socket,
            session_expire,
            session_close_notifier,
            input_tx,
        ));

        let io_task_handle = {
            let session = session.clone();
            tokio::spawn(async move {
                let mut input_buffer = [0u8; 65536];

                loop {
                    tokio::select! {
                        // recv() then input()
                        // Drives the KCP machine forward
                        recv_result = udp_socket.recv(&mut input_buffer), if is_client => {
                            match recv_result {
                                Err(err) => {
                                    error!("[SESSION] UDP recv failed, error: {}", err);
                                }
                                Ok(n) => {
                                    let input_buffer = &input_buffer[..n];
                                    let input_conv = kcp::get_conv(input_buffer);
                                    trace!("[SESSION] UDP recv {} bytes, conv: {}, going to input {:?}", n, input_conv, ByteStr::new(input_buffer));

                                    let mut socket = session.socket.lock();

                                    // Server may allocate another conv for this client.
                                    if !socket.waiting_conv() && socket.conv() != input_conv {
                                        trace!("[SESSION] UDP input conv: {} replaces session conv: {}", input_conv, socket.conv());
                                        socket.set_conv(input_conv);
                                    }

                                    match socket.input(input_buffer) {
                                        Ok(true) => {
                                            trace!("[SESSION] UDP input {} bytes and waked sender/receiver", n);
                                        }
                                        Ok(false) => {}
                                        Err(err) => {
                                            error!("[SESSION] UDP input {} bytes error: {}, input buffer {:?}", n, err, ByteStr::new(input_buffer));
                                        }
                                    }
                                }
                            }
                        }

                        // bytes received from listener socket
                        input_opt = input_rx.recv() => {
                            if let Some(input_buffer) = input_opt {
                                let mut socket = session.socket.lock();
                                match socket.input(&input_buffer) {
                                    Ok(..) => {
                                        trace!("[SESSION] UDP input {} bytes from channel {:?}", input_buffer.len(), ByteStr::new(&input_buffer));
                                    }
                                    Err(err) => {
                                        error!("[SESSION] UDP input {} bytes from channel failed, error: {}, input buffer {:?}",
                                               input_buffer.len(), err, ByteStr::new(&input_buffer));
                                    }
                                }
                            }
                        }
                    }
                }
            })
        };

        // Per-session updater
        {
            let session = session.clone();
            tokio::spawn(async move {
                while !session.closed.load(Ordering::Relaxed) {
                    let next = {
                        let mut socket = session.socket.lock();

                        let is_closed = session.closed.load(Ordering::Acquire);
                        if is_closed && socket.can_close() {
                            trace!("[SESSION] KCP session closing");
                            break;
                        }

                        // server socket expires
                        if !is_client {
                            // If this is a server stream, close it automatically after a period of time
                            let last_update_time = socket.last_update_time();
                            let elapsed = last_update_time.elapsed();

                            if elapsed > session.session_expire {
                                if elapsed > session.session_expire * 2 {
                                    // Force close. Client may have already gone.
                                    trace!(
                                        "[SESSION] force close inactive session, conv: {}, last_update: {}s ago",
                                        socket.conv(),
                                        elapsed.as_secs()
                                    );
                                    break;
                                }

                                if !is_closed {
                                    trace!(
                                        "[SESSION] closing inactive session, conv: {}, last_update: {}s ago",
                                        socket.conv(),
                                        elapsed.as_secs()
                                    );
                                    session.closed.store(true, Ordering::Release);
                                }
                            }
                        }

                        match socket.update() {
                            Ok(next_next) => Instant::from_std(next_next),
                            Err(err) => {
                                error!("[SESSION] KCP update failed, error: {}", err);
                                Instant::now() + Duration::from_millis(10)
                            }
                        }
                    };

                    tokio::select! {
                        _ = time::sleep_until(next) => {},
                        _ = session.notifier.notified() => {},
                    }
                }

                {
                    // Close the socket.
                    // Wake all pending tasks and let all send/recv return EOF

                    let mut socket = session.socket.lock();
                    socket.close();
                }

                if let Some((ref notifier, peer_addr)) = session.session_close_notifier {
                    let _ = notifier.send(peer_addr).await;
                }

                session.closed.store(true, Ordering::Release);
                io_task_handle.abort();

                trace!("[SESSION] KCP session closed");
            });
        }

        session
    }

    pub fn kcp_socket(&self) -> &SpinMutex<KcpSocket> {
        &self.socket
    }

    pub fn close(&self) {
        self.closed.store(true, Ordering::Release);
        self.notify();
    }

    pub async fn input(&self, buf: &[u8]) {
        self.input_tx.send(buf.to_owned()).await.expect("input channel closed")
    }

    pub async fn conv(&self) -> u32 {
        let socket = self.socket.lock();
        socket.conv()
    }

    pub fn notify(&self) {
        self.notifier.notify_one();
    }
}

struct KcpSessionUniq(Arc<KcpSession>);

impl Drop for KcpSessionUniq {
    fn drop(&mut self) {
        self.0.close();
    }
}

impl Deref for KcpSessionUniq {
    type Target = KcpSession;

    fn deref(&self) -> &KcpSession {
        &*self.0
    }
}

pub struct KcpSessionManager {
    sessions: HashMap<SocketAddr, KcpSessionUniq>,
}

impl KcpSessionManager {
    pub fn new() -> KcpSessionManager {
        KcpSessionManager {
            sessions: HashMap::new(),
        }
    }

    #[inline]
    pub fn alloc_conv(&mut self) -> u32 {
        rand::random()
    }

    pub fn close_peer(&mut self, peer_addr: SocketAddr) {
        self.sessions.remove(&peer_addr);
    }

    pub async fn get_or_create(
        &mut self,
        config: &KcpConfig,
        conv: u32,
        sn: u32,
        udp: &Arc<UdpSocket>,
        peer_addr: SocketAddr,
        session_close_notifier: &mpsc::Sender<SocketAddr>,
    ) -> KcpResult<(Arc<KcpSession>, bool)> {
        match self.sessions.entry(peer_addr) {
            Entry::Occupied(mut occ) => {
                let session = occ.get();

                if sn == 0 && session.conv().await != conv {
                    // This is the first packet received from this peer.
                    // Recreate a new session for this specific client.

                    let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
                    let session = KcpSession::new_shared(
                        socket,
                        config.session_expire,
                        Some((session_close_notifier.clone(), peer_addr)),
                    );

                    let old_session = occ.insert(KcpSessionUniq(session.clone()));
                    let old_conv = old_session.conv().await;
                    trace!(
                        "replaced session with conv: {} (old: {}), peer: {}",
                        conv,
                        old_conv,
                        peer_addr
                    );

                    Ok((session, true))
                } else {
                    Ok((session.0.clone(), false))
                }
            }
            Entry::Vacant(vac) => {
                let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
                let session = KcpSession::new_shared(
                    socket,
                    config.session_expire,
                    Some((session_close_notifier.clone(), peer_addr)),
                );
                trace!("created session for conv: {}, peer: {}", conv, peer_addr);
                vac.insert(KcpSessionUniq(session.clone()));
                Ok((session, true))
            }
        }
    }
}
<commit_msg>log trace when KcpSession was dropped<commit_after>use std::{
    collections::{hash_map::Entry, HashMap},
    net::SocketAddr,
    ops::Deref,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};

use byte_string::ByteStr;
use kcp::KcpResult;
use log::{error, trace};
use spin::Mutex as SpinMutex;
use tokio::{
    net::UdpSocket,
    sync::{mpsc, Notify},
    time::{self, Instant},
};

use crate::{skcp::KcpSocket, KcpConfig};

pub struct KcpSession {
    socket: SpinMutex<KcpSocket>,
    closed: AtomicBool,
    session_expire: Duration,
    session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
    input_tx: mpsc::Sender<Vec<u8>>,
    notifier: Notify,
}

impl Drop for KcpSession {
    fn drop(&mut self) {
        trace!(
            "[SESSION] KcpSession conv {} is dropping, closed? {}",
            self.socket.lock().conv(),
            self.closed.load(Ordering::Acquire),
        );
    }
}

impl KcpSession {
    fn new(
        socket: KcpSocket,
        session_expire: Duration,
        session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
        input_tx: mpsc::Sender<Vec<u8>>,
    ) -> KcpSession {
        KcpSession {
            socket: SpinMutex::new(socket),
            closed: AtomicBool::new(false),
            session_expire,
            session_close_notifier,
            input_tx,
            notifier: Notify::new(),
        }
    }

    pub fn new_shared(
        socket: KcpSocket,
        session_expire: Duration,
        session_close_notifier: Option<(mpsc::Sender<SocketAddr>, SocketAddr)>,
    ) -> Arc<KcpSession> {
        let is_client = session_close_notifier.is_none();

        let (input_tx, mut input_rx) = mpsc::channel(64);

        let udp_socket = socket.udp_socket().clone();

        let session = Arc::new(KcpSession::new(
            socket,
            session_expire,
            session_close_notifier,
            input_tx,
        ));

        let io_task_handle = {
            let session = session.clone();
            tokio::spawn(async move {
                let mut input_buffer = [0u8; 65536];

                loop {
                    tokio::select! {
                        // recv() then input()
                        // Drives the KCP machine forward
                        recv_result = udp_socket.recv(&mut input_buffer), if is_client => {
                            match recv_result {
                                Err(err) => {
                                    error!("[SESSION] UDP recv failed, error: {}", err);
                                }
                                Ok(n) => {
                                    let input_buffer = &input_buffer[..n];
                                    let input_conv = kcp::get_conv(input_buffer);
                                    trace!("[SESSION] UDP recv {} bytes, conv: {}, going to input {:?}", n, input_conv, ByteStr::new(input_buffer));

                                    let mut socket = session.socket.lock();

                                    // Server may allocate another conv for this client.
                                    if !socket.waiting_conv() && socket.conv() != input_conv {
                                        trace!("[SESSION] UDP input conv: {} replaces session conv: {}", input_conv, socket.conv());
                                        socket.set_conv(input_conv);
                                    }

                                    match socket.input(input_buffer) {
                                        Ok(true) => {
                                            trace!("[SESSION] UDP input {} bytes and waked sender/receiver", n);
                                        }
                                        Ok(false) => {}
                                        Err(err) => {
                                            error!("[SESSION] UDP input {} bytes error: {}, input buffer {:?}", n, err, ByteStr::new(input_buffer));
                                        }
                                    }
                                }
                            }
                        }

                        // bytes received from listener socket
                        input_opt = input_rx.recv() => {
                            if let Some(input_buffer) = input_opt {
                                let mut socket = session.socket.lock();
                                match socket.input(&input_buffer) {
                                    Ok(..) => {
                                        trace!("[SESSION] UDP input {} bytes from channel {:?}", input_buffer.len(), ByteStr::new(&input_buffer));
                                    }
                                    Err(err) => {
                                        error!("[SESSION] UDP input {} bytes from channel failed, error: {}, input buffer {:?}",
                                               input_buffer.len(), err, ByteStr::new(&input_buffer));
                                    }
                                }
                            }
                        }
                    }
                }
            })
        };

        // Per-session updater
        {
            let session = session.clone();
            tokio::spawn(async move {
                while !session.closed.load(Ordering::Relaxed) {
                    let next = {
                        let mut socket = session.socket.lock();

                        let is_closed = session.closed.load(Ordering::Acquire);
                        if is_closed && socket.can_close() {
                            trace!("[SESSION] KCP session closing");
                            break;
                        }

                        // server socket expires
                        if !is_client {
                            // If this is a server stream, close it automatically after a period of time
                            let last_update_time = socket.last_update_time();
                            let elapsed = last_update_time.elapsed();

                            if elapsed > session.session_expire {
                                if elapsed > session.session_expire * 2 {
                                    // Force close. Client may have already gone.
                                    trace!(
                                        "[SESSION] force close inactive session, conv: {}, last_update: {}s ago",
                                        socket.conv(),
                                        elapsed.as_secs()
                                    );
                                    break;
                                }

                                if !is_closed {
                                    trace!(
                                        "[SESSION] closing inactive session, conv: {}, last_update: {}s ago",
                                        socket.conv(),
                                        elapsed.as_secs()
                                    );
                                    session.closed.store(true, Ordering::Release);
                                }
                            }
                        }

                        match socket.update() {
                            Ok(next_next) => Instant::from_std(next_next),
                            Err(err) => {
                                error!("[SESSION] KCP update failed, error: {}", err);
                                Instant::now() + Duration::from_millis(10)
                            }
                        }
                    };

                    tokio::select! {
                        _ = time::sleep_until(next) => {},
                        _ = session.notifier.notified() => {},
                    }
                }

                {
                    // Close the socket.
                    // Wake all pending tasks and let all send/recv return EOF

                    let mut socket = session.socket.lock();
                    socket.close();
                }

                if let Some((ref notifier, peer_addr)) = session.session_close_notifier {
                    let _ = notifier.send(peer_addr).await;
                }

                session.closed.store(true, Ordering::Release);
                io_task_handle.abort();

                trace!("[SESSION] KCP session closed");
            });
        }

        session
    }

    pub fn kcp_socket(&self) -> &SpinMutex<KcpSocket> {
        &self.socket
    }

    pub fn close(&self) {
        self.closed.store(true, Ordering::Release);
        self.notify();
    }

    pub async fn input(&self, buf: &[u8]) {
        self.input_tx.send(buf.to_owned()).await.expect("input channel closed")
    }

    pub async fn conv(&self) -> u32 {
        let socket = self.socket.lock();
        socket.conv()
    }

    pub fn notify(&self) {
        self.notifier.notify_one();
    }
}

struct KcpSessionUniq(Arc<KcpSession>);

impl Drop for KcpSessionUniq {
    fn drop(&mut self) {
        self.0.close();
    }
}

impl Deref for KcpSessionUniq {
    type Target = KcpSession;

    fn deref(&self) -> &KcpSession {
        &*self.0
    }
}

pub struct KcpSessionManager {
    sessions: HashMap<SocketAddr, KcpSessionUniq>,
}

impl KcpSessionManager {
    pub fn new() -> KcpSessionManager {
        KcpSessionManager {
            sessions: HashMap::new(),
        }
    }

    #[inline]
    pub fn alloc_conv(&mut self) -> u32 {
        rand::random()
    }

    pub fn close_peer(&mut self, peer_addr: SocketAddr) {
        self.sessions.remove(&peer_addr);
    }

    pub async fn get_or_create(
        &mut self,
        config: &KcpConfig,
        conv: u32,
        sn: u32,
        udp: &Arc<UdpSocket>,
        peer_addr: SocketAddr,
        session_close_notifier: &mpsc::Sender<SocketAddr>,
    ) -> KcpResult<(Arc<KcpSession>, bool)> {
        match self.sessions.entry(peer_addr) {
            Entry::Occupied(mut occ) => {
                let session = occ.get();

                if sn == 0 && session.conv().await != conv {
                    // This is the first packet received from this peer.
                    // Recreate a new session for this specific client.

                    let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
                    let session = KcpSession::new_shared(
                        socket,
                        config.session_expire,
                        Some((session_close_notifier.clone(), peer_addr)),
                    );

                    let old_session = occ.insert(KcpSessionUniq(session.clone()));
                    let old_conv = old_session.conv().await;
                    trace!(
                        "replaced session with conv: {} (old: {}), peer: {}",
                        conv,
                        old_conv,
                        peer_addr
                    );

                    Ok((session, true))
                } else {
                    Ok((session.0.clone(), false))
                }
            }
            Entry::Vacant(vac) => {
                let socket = KcpSocket::new(config, conv, udp.clone(), peer_addr, config.stream)?;
                let session = KcpSession::new_shared(
                    socket,
                    config.session_expire,
                    Some((session_close_notifier.clone(), peer_addr)),
                );
                trace!("created session for conv: {}, peer: {}", conv, peer_addr);
                vac.insert(KcpSessionUniq(session.clone()));
                Ok((session, true))
            }
        }
    }
}
<|endoftext|>