<commit_before>//! Association requester module
//!
//! The module provides an abstraction for a DICOM association
//! in which this application entity is the one requesting the association.
//! See [`ClientAssociationOptions`](self::ClientAssociationOptions)
//! for details and examples on how to create an association.
use std::{
    borrow::Cow,
    net::{TcpStream, ToSocketAddrs},
};

use crate::pdu::{AbortRQSource, AssociationRJResult, AssociationRJSource, Pdu, PresentationContextProposed, PresentationContextResult, PresentationContextResultReason, reader::read_pdu, writer::write_pdu};
use snafu::{ensure, ResultExt, Snafu};

use super::pdata::PDataWriter;

#[derive(Debug, Snafu)]
#[non_exhaustive]
pub enum Error {
    /// missing abstract syntax to begin negotiation
    MissingAbstractSyntax,

    /// could not connect to server
    Connect { source: std::io::Error },

    /// failed to send association request
    SendRequest { source: crate::pdu::writer::Error },

    /// failed to receive association response
    ReceiveResponse { source: crate::pdu::reader::Error },

    #[snafu(display("unexpected response from server `{:?}`", pdu))]
    #[non_exhaustive]
    UnexpectedResponse {
        /// the PDU obtained from the server
        pdu: Pdu,
    },

    #[snafu(display("unknown response from server `{:?}`", pdu))]
    #[non_exhaustive]
    UnknownResponse {
        /// the PDU obtained from the server, of variant Unknown
        pdu: Pdu,
    },

    #[snafu(display("protocol version mismatch: expected {}, got {}", expected, got))]
    ProtocolVersionMismatch { expected: u16, got: u16 },

    /// the association was rejected by the server
    Rejected {
        association_result: AssociationRJResult,
        association_source: AssociationRJSource,
    },

    /// no presentation contexts accepted by the server
    NoAcceptedPresentationContexts,

    /// failed to send PDU message
    #[non_exhaustive]
    Send { source: crate::pdu::writer::Error },

    /// failed to receive PDU message
    #[non_exhaustive]
    Receive { source: crate::pdu::reader::Error },
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

/// A DICOM association builder for a client node.
/// The final outcome is a [`ClientAssociation`].
///
/// This is the standard way of requesting and establishing
/// an association with another DICOM node,
/// that one usually taking the role of a service class provider (SCP).
///
/// # Example
///
/// ```no_run
/// # use dicom_ul::association::client::ClientAssociationOptions;
/// # fn run() -> Result<(), Box<dyn std::error::Error>> {
/// let association = ClientAssociationOptions::new()
///    .with_abstract_syntax("1.2.840.10008.1.1")
///    .with_transfer_syntax("1.2.840.10008.1.2.1")
///    .establish("129.168.0.5:104")?;
/// # Ok(())
/// # }
/// ```
///
/// At least one abstract syntax must be specified,
/// using the method [`with_abstract_syntax`](Self::with_abstract_syntax).
/// The requester will admit by default the transfer syntaxes
/// _Implicit VR Little Endian_
/// and _Explicit VR Little Endian_.
/// Other transfer syntaxes can be requested in the association
/// via the method [`with_transfer_syntax`](Self::with_transfer_syntax).
#[derive(Debug, Clone)]
pub struct ClientAssociationOptions<'a> {
    /// the calling AE title
    calling_ae_title: Cow<'a, str>,
    /// the called AE title
    called_ae_title: Cow<'a, str>,
    /// the requested application context name
    application_context_name: Cow<'a, str>,
    /// the list of requested abstract syntaxes
    abstract_syntax_uids: Vec<Cow<'a, str>>,
    /// the list of requested transfer syntaxes
    transfer_syntax_uids: Vec<Cow<'a, str>>,
    /// the expected protocol version
    protocol_version: u16,
    /// the maximum PDU length
    max_pdu_length: u32,
}

impl<'a> Default for ClientAssociationOptions<'a> {
    fn default() -> Self {
        ClientAssociationOptions {
            /// the calling AE title
            calling_ae_title: "THIS-SCU".into(),
            /// the called AE title
            called_ae_title: "ANY-SCP".into(),
            /// the requested application context name
            application_context_name: "1.2.840.10008.3.1.1.1".into(),
            /// the list of requested abstract syntaxes
            abstract_syntax_uids: Vec::new(),
            /// the application context name
            transfer_syntax_uids: Vec::new(),
            protocol_version: 1,
            max_pdu_length: crate::pdu::reader::DEFAULT_MAX_PDU,
        }
    }
}

impl<'a> ClientAssociationOptions<'a> {
    /// Create a new set of options for establishing an association.
    pub fn new() -> Self {
        Self::default()
    }

    /// Define the calling application entity title for the association,
    /// which refers to this DICOM node.
    ///
    /// The default is `THIS-SCU`.
    pub fn calling_ae_title<T>(mut self, calling_ae_title: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.calling_ae_title = calling_ae_title.into();
        self
    }

    /// Define the called application entity title for the association,
    /// which refers to the target DICOM node.
    ///
    /// The default is `ANY-SCP`.
    pub fn called_ae_title<T>(mut self, called_ae_title: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.called_ae_title = called_ae_title.into();
        self
    }

    /// Include this abstract syntax
    /// in the list of proposed presentation contexts.
    pub fn with_abstract_syntax<T>(mut self, abstract_syntax_uid: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.abstract_syntax_uids.push(abstract_syntax_uid.into());
        self
    }

    /// Include this transfer syntax in each proposed presentation context.
    pub fn with_transfer_syntax<T>(mut self, transfer_syntax_uid: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.transfer_syntax_uids.push(transfer_syntax_uid.into());
        self
    }

    /// Override the maximum expected PDU length.
    pub fn max_pdu_length(mut self, value: u32) -> Self {
        self.max_pdu_length = value;
        self
    }

    /// Initiate the TCP connection to the given address
    /// and request a new DICOM association,
    /// negotiating the presentation contexts in the process.
    pub fn establish<A: ToSocketAddrs>(self, address: A) -> Result<ClientAssociation> {
        let ClientAssociationOptions {
            calling_ae_title,
            called_ae_title,
            application_context_name,
            abstract_syntax_uids,
            mut transfer_syntax_uids,
            protocol_version,
            max_pdu_length,
        } = self;

        // fail if no abstract syntaxes were provided: they represent intent,
        // should not be omitted by the user
        ensure!(!abstract_syntax_uids.is_empty(), MissingAbstractSyntax);

        // provide default transfer syntaxes
        if transfer_syntax_uids.is_empty() {
            // Explicit VR Little Endian
            transfer_syntax_uids.push("1.2.840.10008.1.2.1".into());
            // Implicit VR Little Endian
            transfer_syntax_uids.push("1.2.840.10008.1.2".into());
        }

        let presentation_contexts: Vec<_> = abstract_syntax_uids
            .into_iter()
            .enumerate()
            .map(|(i, abstract_syntax)| PresentationContextProposed {
                id: (i + 1) as u8,
                abstract_syntax: abstract_syntax.to_string(),
                transfer_syntaxes: transfer_syntax_uids
                    .iter()
                    .map(|uid| uid.to_string())
                    .collect(),
            })
            .collect();
        let msg = Pdu::AssociationRQ {
            protocol_version,
            calling_ae_title: calling_ae_title.to_string(),
            called_ae_title: called_ae_title.to_string(),
            application_context_name: application_context_name.to_string(),
            presentation_contexts: presentation_contexts.clone(),
            user_variables: vec![],
        };

        let mut socket = std::net::TcpStream::connect(address).context(Connect)?;

        // send request
        write_pdu(&mut socket, &msg).context(SendRequest)?;

        // receive response
        let msg = read_pdu(&mut socket, max_pdu_length).context(ReceiveResponse)?;

        match msg {
            Pdu::AssociationAC {
                protocol_version: protocol_version_scp,
                application_context_name: _,
                presentation_contexts: presentation_contexts_scp,
                user_variables: _,
            } => {
                ensure!(
                    protocol_version == protocol_version_scp,
                    ProtocolVersionMismatch {
                        expected: protocol_version,
                        got: protocol_version_scp,
                    }
                );

                let presentation_contexts: Vec<_> = presentation_contexts_scp
                    .into_iter()
                    .filter(|c| c.reason == PresentationContextResultReason::Acceptance)
                    .collect();
                ensure!(!presentation_contexts.is_empty(), NoAcceptedPresentationContexts);

                Ok(ClientAssociation {
                    presentation_contexts,
                    max_pdu_length,
                    socket,
                })
            }
            Pdu::AssociationRJ { result, source } => Rejected {
                association_result: result,
                association_source: source,
            }
            .fail(),
            pdu @ Pdu::AbortRQ { .. }
            | pdu @ Pdu::ReleaseRQ { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::PData { .. }
            | pdu @ Pdu::ReleaseRP { .. } => UnexpectedResponse { pdu }.fail(),
            pdu @ Pdu::Unknown { .. } => UnknownResponse { pdu }.fail(),
        }
    }
}

/// A DICOM upper level association from the perspective
/// of a requesting application entity.
///
/// The most common operations of an established association are
/// [`send`](Self::send)
/// and [`receive`](Self::receive).
/// Sending large P-Data fragments may be easier through the P-Data sender
/// abstraction (see [`send_pdata`](Self::send_pdata)).
///
/// When the value falls out of scope,
/// the program will automatically try to gracefully release the association
/// through a standard C-RELEASE message exchange,
/// then shut down the underlying TCP connection.
#[derive(Debug)]
pub struct ClientAssociation {
    /// The presentation contexts accorded with the acceptor application entity,
    /// without the rejected ones.
    presentation_contexts: Vec<PresentationContextResult>,
    /// The maximum PDU length
    max_pdu_length: u32,
    /// The TCP stream to the other DICOM node
    socket: TcpStream,
}

impl ClientAssociation {
    /// Retrieve the list of negotiated presentation contexts.
    pub fn presentation_contexts(&self) -> &[PresentationContextResult] {
        &self.presentation_contexts
    }

    /// Send a PDU message to the other intervenient.
    pub fn send(&mut self, msg: &Pdu) -> Result<()> {
        write_pdu(&mut self.socket, &msg).context(Send)
    }

    /// Read a PDU message from the other intervenient.
    pub fn receive(&mut self) -> Result<Pdu> {
        read_pdu(&mut self.socket, self.max_pdu_length).context(Receive)
    }

    /// Gracefully terminate the association by exchanging release messages
    /// and then shutting down the TCP connection.
    pub fn release(mut self) -> Result<()> {
        let out = self.release_impl();
        let _ = self.socket.shutdown(std::net::Shutdown::Both);
        out
    }

    /// Release implementation function,
    /// which tries to send a release request and receive a release response.
    /// This is in a separate private function because
    /// terminating a connection should still close the connection
    /// if the exchange fails.
    /// Send an abort message and shut down the TCP connection,
    /// terminating the association.
    pub fn abort(mut self) -> Result<()> {
        let out = write_pdu(
            &mut self.socket,
            &Pdu::AbortRQ {
                source: AbortRQSource::ServiceUser,
            },
        )
        .context(Send);

        let _ = self.socket.shutdown(std::net::Shutdown::Both);
        out
    }

    /// Obtain access to the inner TCP stream
    /// connected to the association acceptor.
    ///
    /// This can be used to send the PDU in semantic fragments of the message,
    /// thus using less memory.
    ///
    /// **Note:** reading and writing should be done with care
    /// to avoid inconsistencies in the association state.
    /// Do not call `send` and `receive` while not in a PDU boundary.
    pub fn inner_stream(&mut self) -> &mut TcpStream {
        &mut self.socket
    }

    /// Prepare a P-Data writer for sending
    /// one or more data items.
    ///
    /// Returns a writer which automatically
    /// splits the inner data into separate PDUs if necessary.
    pub fn send_pdata(&mut self, presentation_context_id: u8) -> PDataWriter<&mut TcpStream> {
        PDataWriter::new(
            &mut self.socket,
            presentation_context_id,
            self.max_pdu_length,
        )
    }

    fn release_impl(&mut self) -> Result<()> {
        write_pdu(&mut self.socket, &Pdu::ReleaseRQ).context(Send)?;

        let pdu = read_pdu(&mut self.socket, self.max_pdu_length).context(Receive)?;

        match pdu {
            Pdu::ReleaseRP => {}
            pdu @ Pdu::AbortRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::PData { .. }
            | pdu @ Pdu::ReleaseRQ { .. } => return UnexpectedResponse { pdu }.fail(),
            pdu @ Pdu::Unknown { .. } => return UnknownResponse { pdu }.fail(),
        }
        Ok(())
    }
}

/// Automatically release the association and shut down the connection.
impl Drop for ClientAssociation {
    fn drop(&mut self) {
        let _ = self.release_impl();
        let _ = self.socket.shutdown(std::net::Shutdown::Both);
    }
}
<commit_msg>[ul] Send connection aborts on client error<commit_after>//! Association requester module
//!
//! The module provides an abstraction for a DICOM association
//! in which this application entity is the one requesting the association.
//! See [`ClientAssociationOptions`](self::ClientAssociationOptions)
//! for details and examples on how to create an association.
use std::{
    borrow::Cow,
    net::{TcpStream, ToSocketAddrs},
};

use crate::pdu::{AbortRQSource, AssociationRJResult, AssociationRJSource, Pdu, PresentationContextProposed, PresentationContextResult, PresentationContextResultReason, reader::read_pdu, writer::write_pdu};
use snafu::{ensure, ResultExt, Snafu};

use super::pdata::PDataWriter;

#[derive(Debug, Snafu)]
#[non_exhaustive]
pub enum Error {
    /// missing abstract syntax to begin negotiation
    MissingAbstractSyntax,

    /// could not connect to server
    Connect { source: std::io::Error },

    /// failed to send association request
    SendRequest { source: crate::pdu::writer::Error },

    /// failed to receive association response
    ReceiveResponse { source: crate::pdu::reader::Error },

    #[snafu(display("unexpected response from server `{:?}`", pdu))]
    #[non_exhaustive]
    UnexpectedResponse {
        /// the PDU obtained from the server
        pdu: Pdu,
    },

    #[snafu(display("unknown response from server `{:?}`", pdu))]
    #[non_exhaustive]
    UnknownResponse {
        /// the PDU obtained from the server, of variant Unknown
        pdu: Pdu,
    },

    #[snafu(display("protocol version mismatch: expected {}, got {}", expected, got))]
    ProtocolVersionMismatch { expected: u16, got: u16 },

    /// the association was rejected by the server
    Rejected {
        association_result: AssociationRJResult,
        association_source: AssociationRJSource,
    },

    /// no presentation contexts accepted by the server
    NoAcceptedPresentationContexts,

    /// failed to send PDU message
    #[non_exhaustive]
    Send { source: crate::pdu::writer::Error },

    /// failed to receive PDU message
    #[non_exhaustive]
    Receive { source: crate::pdu::reader::Error },
}

pub type Result<T, E = Error> = std::result::Result<T, E>;

/// A DICOM association builder for a client node.
/// The final outcome is a [`ClientAssociation`].
///
/// This is the standard way of requesting and establishing
/// an association with another DICOM node,
/// that one usually taking the role of a service class provider (SCP).
///
/// # Example
///
/// ```no_run
/// # use dicom_ul::association::client::ClientAssociationOptions;
/// # fn run() -> Result<(), Box<dyn std::error::Error>> {
/// let association = ClientAssociationOptions::new()
///    .with_abstract_syntax("1.2.840.10008.1.1")
///    .with_transfer_syntax("1.2.840.10008.1.2.1")
///    .establish("129.168.0.5:104")?;
/// # Ok(())
/// # }
/// ```
///
/// At least one abstract syntax must be specified,
/// using the method [`with_abstract_syntax`](Self::with_abstract_syntax).
/// The requester will admit by default the transfer syntaxes
/// _Implicit VR Little Endian_
/// and _Explicit VR Little Endian_.
/// Other transfer syntaxes can be requested in the association
/// via the method [`with_transfer_syntax`](Self::with_transfer_syntax).
#[derive(Debug, Clone)]
pub struct ClientAssociationOptions<'a> {
    /// the calling AE title
    calling_ae_title: Cow<'a, str>,
    /// the called AE title
    called_ae_title: Cow<'a, str>,
    /// the requested application context name
    application_context_name: Cow<'a, str>,
    /// the list of requested abstract syntaxes
    abstract_syntax_uids: Vec<Cow<'a, str>>,
    /// the list of requested transfer syntaxes
    transfer_syntax_uids: Vec<Cow<'a, str>>,
    /// the expected protocol version
    protocol_version: u16,
    /// the maximum PDU length
    max_pdu_length: u32,
}

impl<'a> Default for ClientAssociationOptions<'a> {
    fn default() -> Self {
        ClientAssociationOptions {
            /// the calling AE title
            calling_ae_title: "THIS-SCU".into(),
            /// the called AE title
            called_ae_title: "ANY-SCP".into(),
            /// the requested application context name
            application_context_name: "1.2.840.10008.3.1.1.1".into(),
            /// the list of requested abstract syntaxes
            abstract_syntax_uids: Vec::new(),
            /// the application context name
            transfer_syntax_uids: Vec::new(),
            protocol_version: 1,
            max_pdu_length: crate::pdu::reader::DEFAULT_MAX_PDU,
        }
    }
}

impl<'a> ClientAssociationOptions<'a> {
    /// Create a new set of options for establishing an association.
    pub fn new() -> Self {
        Self::default()
    }

    /// Define the calling application entity title for the association,
    /// which refers to this DICOM node.
    ///
    /// The default is `THIS-SCU`.
    pub fn calling_ae_title<T>(mut self, calling_ae_title: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.calling_ae_title = calling_ae_title.into();
        self
    }

    /// Define the called application entity title for the association,
    /// which refers to the target DICOM node.
    ///
    /// The default is `ANY-SCP`.
    pub fn called_ae_title<T>(mut self, called_ae_title: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.called_ae_title = called_ae_title.into();
        self
    }

    /// Include this abstract syntax
    /// in the list of proposed presentation contexts.
    pub fn with_abstract_syntax<T>(mut self, abstract_syntax_uid: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.abstract_syntax_uids.push(abstract_syntax_uid.into());
        self
    }

    /// Include this transfer syntax in each proposed presentation context.
    pub fn with_transfer_syntax<T>(mut self, transfer_syntax_uid: T) -> Self
    where
        T: Into<Cow<'a, str>>,
    {
        self.transfer_syntax_uids.push(transfer_syntax_uid.into());
        self
    }

    /// Override the maximum expected PDU length.
    pub fn max_pdu_length(mut self, value: u32) -> Self {
        self.max_pdu_length = value;
        self
    }

    /// Initiate the TCP connection to the given address
    /// and request a new DICOM association,
    /// negotiating the presentation contexts in the process.
    pub fn establish<A: ToSocketAddrs>(self, address: A) -> Result<ClientAssociation> {
        let ClientAssociationOptions {
            calling_ae_title,
            called_ae_title,
            application_context_name,
            abstract_syntax_uids,
            mut transfer_syntax_uids,
            protocol_version,
            max_pdu_length,
        } = self;

        // fail if no abstract syntaxes were provided: they represent intent,
        // should not be omitted by the user
        ensure!(!abstract_syntax_uids.is_empty(), MissingAbstractSyntax);

        // provide default transfer syntaxes
        if transfer_syntax_uids.is_empty() {
            // Explicit VR Little Endian
            transfer_syntax_uids.push("1.2.840.10008.1.2.1".into());
            // Implicit VR Little Endian
            transfer_syntax_uids.push("1.2.840.10008.1.2".into());
        }

        let presentation_contexts: Vec<_> = abstract_syntax_uids
            .into_iter()
            .enumerate()
            .map(|(i, abstract_syntax)| PresentationContextProposed {
                id: (i + 1) as u8,
                abstract_syntax: abstract_syntax.to_string(),
                transfer_syntaxes: transfer_syntax_uids
                    .iter()
                    .map(|uid| uid.to_string())
                    .collect(),
            })
            .collect();
        let msg = Pdu::AssociationRQ {
            protocol_version,
            calling_ae_title: calling_ae_title.to_string(),
            called_ae_title: called_ae_title.to_string(),
            application_context_name: application_context_name.to_string(),
            presentation_contexts: presentation_contexts.clone(),
            user_variables: vec![],
        };

        let mut socket = std::net::TcpStream::connect(address).context(Connect)?;

        // send request
        write_pdu(&mut socket, &msg).context(SendRequest)?;

        // receive response
        let msg = read_pdu(&mut socket, max_pdu_length).context(ReceiveResponse)?;

        match msg {
            Pdu::AssociationAC {
                protocol_version: protocol_version_scp,
                application_context_name: _,
                presentation_contexts: presentation_contexts_scp,
                user_variables: _,
            } => {
                ensure!(
                    protocol_version == protocol_version_scp,
                    ProtocolVersionMismatch {
                        expected: protocol_version,
                        got: protocol_version_scp,
                    }
                );

                let presentation_contexts: Vec<_> = presentation_contexts_scp
                    .into_iter()
                    .filter(|c| c.reason == PresentationContextResultReason::Acceptance)
                    .collect();
                if presentation_contexts.is_empty() {
                    // abort connection
                    let _ = write_pdu(&mut socket, &Pdu::AbortRQ {
                        source: AbortRQSource::ServiceUser,
                    });
                    return NoAcceptedPresentationContexts.fail();
                }

                Ok(ClientAssociation {
                    presentation_contexts,
                    max_pdu_length,
                    socket,
                })
            }
            Pdu::AssociationRJ { result, source } => Rejected {
                association_result: result,
                association_source: source,
            }
            .fail(),
            pdu @ Pdu::AbortRQ { .. }
            | pdu @ Pdu::ReleaseRQ { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::PData { .. }
            | pdu @ Pdu::ReleaseRP { .. } => {
                // abort connection
                let _ = write_pdu(&mut socket, &Pdu::AbortRQ {
                    source: AbortRQSource::ServiceUser,
                });
                UnexpectedResponse { pdu }.fail()
            }
            pdu @ Pdu::Unknown { .. } => {
                // abort connection
                let _ = write_pdu(&mut socket, &Pdu::AbortRQ {
                    source: AbortRQSource::ServiceUser,
                });
                UnknownResponse { pdu }.fail()
            }
        }
    }
}

/// A DICOM upper level association from the perspective
/// of a requesting application entity.
///
/// The most common operations of an established association are
/// [`send`](Self::send)
/// and [`receive`](Self::receive).
/// Sending large P-Data fragments may be easier through the P-Data sender
/// abstraction (see [`send_pdata`](Self::send_pdata)).
///
/// When the value falls out of scope,
/// the program will automatically try to gracefully release the association
/// through a standard C-RELEASE message exchange,
/// then shut down the underlying TCP connection.
#[derive(Debug)]
pub struct ClientAssociation {
    /// The presentation contexts accorded with the acceptor application entity,
    /// without the rejected ones.
    presentation_contexts: Vec<PresentationContextResult>,
    /// The maximum PDU length
    max_pdu_length: u32,
    /// The TCP stream to the other DICOM node
    socket: TcpStream,
}

impl ClientAssociation {
    /// Retrieve the list of negotiated presentation contexts.
    pub fn presentation_contexts(&self) -> &[PresentationContextResult] {
        &self.presentation_contexts
    }

    /// Send a PDU message to the other intervenient.
    pub fn send(&mut self, msg: &Pdu) -> Result<()> {
        write_pdu(&mut self.socket, &msg).context(Send)
    }

    /// Read a PDU message from the other intervenient.
    pub fn receive(&mut self) -> Result<Pdu> {
        read_pdu(&mut self.socket, self.max_pdu_length).context(Receive)
    }

    /// Gracefully terminate the association by exchanging release messages
    /// and then shutting down the TCP connection.
    pub fn release(mut self) -> Result<()> {
        let out = self.release_impl();
        let _ = self.socket.shutdown(std::net::Shutdown::Both);
        out
    }

    /// Release implementation function,
    /// which tries to send a release request and receive a release response.
    /// This is in a separate private function because
    /// terminating a connection should still close the connection
    /// if the exchange fails.
    /// Send an abort message and shut down the TCP connection,
    /// terminating the association.
    pub fn abort(mut self) -> Result<()> {
        let out = write_pdu(
            &mut self.socket,
            &Pdu::AbortRQ {
                source: AbortRQSource::ServiceUser,
            },
        )
        .context(Send);

        let _ = self.socket.shutdown(std::net::Shutdown::Both);
        out
    }

    /// Obtain access to the inner TCP stream
    /// connected to the association acceptor.
    ///
    /// This can be used to send the PDU in semantic fragments of the message,
    /// thus using less memory.
    ///
    /// **Note:** reading and writing should be done with care
    /// to avoid inconsistencies in the association state.
    /// Do not call `send` and `receive` while not in a PDU boundary.
    pub fn inner_stream(&mut self) -> &mut TcpStream {
        &mut self.socket
    }

    /// Prepare a P-Data writer for sending
    /// one or more data items.
    ///
    /// Returns a writer which automatically
    /// splits the inner data into separate PDUs if necessary.
    pub fn send_pdata(&mut self, presentation_context_id: u8) -> PDataWriter<&mut TcpStream> {
        PDataWriter::new(
            &mut self.socket,
            presentation_context_id,
            self.max_pdu_length,
        )
    }

    fn release_impl(&mut self) -> Result<()> {
        write_pdu(&mut self.socket, &Pdu::ReleaseRQ).context(Send)?;

        let pdu = read_pdu(&mut self.socket, self.max_pdu_length).context(Receive)?;

        match pdu {
            Pdu::ReleaseRP => {}
            pdu @ Pdu::AbortRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::PData { .. }
            | pdu @ Pdu::ReleaseRQ { .. } => return UnexpectedResponse { pdu }.fail(),
            pdu @ Pdu::Unknown { .. } => return UnknownResponse { pdu }.fail(),
        }
        Ok(())
    }
}

/// Automatically release the association and shut down the connection.
impl Drop for ClientAssociation {
    fn drop(&mut self) {
        let _ = self.release_impl();
        let _ = self.socket.shutdown(std::net::Shutdown::Both);
    }
}
<|endoftext|>


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

#![feature(box_syntax)]
#![feature(collections)]
#![feature(rustc_private)]
#![feature(slice_patterns)]
#![feature(step_by)]

extern crate geom;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate png;
extern crate stb_image;
extern crate url;
extern crate util;
extern crate msg;

use hyper::header::{ContentType, Headers};
use hyper::http::RawStatus;
use hyper::method::Method;
use hyper::mime::{Mime, Attr};
use msg::constellation_msg::{PipelineId};
use url::Url;

use std::sync::mpsc::{channel, Receiver, Sender};

pub mod image_cache_task;
pub mod storage_task;

/// Image handling.
///
/// It may be surprising that this goes in the network crate as opposed to the graphics crate.
/// However, image handling is generally very integrated with the network stack (especially where
/// caching is involved) and as a result it must live in here.
pub mod image {
    pub mod base;
}

#[derive(Clone)]
pub struct LoadData {
    pub url: Url,
    pub method: Method,
    /// Headers that will apply to the initial request only
    pub headers: Headers,
    /// Headers that will apply to the initial request and any redirects
    pub preserved_headers: Headers,
    pub data: Option<Vec<u8>>,
    pub cors: Option<ResourceCORSData>,
    pub pipeline_id: Option<PipelineId>,
}

impl LoadData {
    pub fn new(url: Url, id: Option<PipelineId>) -> LoadData {
        LoadData {
            url: url,
            method: Method::Get,
            headers: Headers::new(),
            preserved_headers: Headers::new(),
            data: None,
            cors: None,
            pipeline_id: id,
        }
    }
}

/// A listener for asynchronous network events. Cancelling the underlying request is unsupported.
pub trait AsyncResponseListener {
    /// The response headers for a request have been received.
    fn headers_available(&self, metadata: Metadata);
    /// A portion of the response body has been received. This data is unavailable after
    /// this method returned, and must be stored accordingly.
    fn data_available(&self, payload: Vec<u8>);
    /// The response is complete. If the provided status is an Err value, there is no guarantee
    /// that the response body was completely read.
    fn response_complete(&self, status: Result<(), String>);
}

/// Data for passing between threads/processes to indicate a particular action to
/// take on a provided network listener.
pub enum ResponseAction {
    /// Invoke headers_available
    HeadersAvailable(Metadata),
    /// Invoke data_available
    DataAvailable(Vec<u8>),
    /// Invoke response_complete
    ResponseComplete(Result<(), String>)
}

impl ResponseAction {
    /// Execute the default action on a provided listener.
    pub fn process(self, listener: &AsyncResponseListener) {
        match self {
            ResponseAction::HeadersAvailable(m) => listener.headers_available(m),
            ResponseAction::DataAvailable(d) => listener.data_available(d),
            ResponseAction::ResponseComplete(r) => listener.response_complete(r),
        }
    }
}

/// A target for async networking events. Commonly used to dispatch a runnable event to another
/// thread storing the wrapped closure for later execution.
pub trait AsyncResponseTarget {
    fn invoke_with_listener(&self, action: ResponseAction);
}

/// A wrapper for a network load that can either be channel or event-based.
pub enum LoadConsumer {
    Channel(Sender<LoadResponse>),
    Listener(Box<AsyncResponseTarget + Send>),
}

/// Handle to a resource task
pub type ResourceTask = Sender<ControlMsg>;

pub enum ControlMsg {
    /// Request the data associated with a particular URL
    Load(LoadData, LoadConsumer),
    /// Store a set of cookies for a given originating URL
    SetCookiesForUrl(Url, String, CookieSource),
    /// Retrieve the stored cookies for a given URL
    GetCookiesForUrl(Url, Sender<Option<String>>, CookieSource),
    Exit
}

/// Initialized but unsent request. Encapsulates everything necessary to instruct
/// the resource task to make a new request. The `load` method *must* be called before
/// destruction or the task will panic.
pub struct PendingAsyncLoad {
    resource_task: ResourceTask,
    url: Url,
    pipeline: Option<PipelineId>,
    guard: PendingLoadGuard,
}

struct PendingLoadGuard {
    loaded: bool,
}

impl PendingLoadGuard {
    fn neuter(&mut self) {
        self.loaded = true;
    }
}

impl Drop for PendingLoadGuard {
    fn drop(&mut self) {
        assert!(self.loaded)
    }
}

impl PendingAsyncLoad {
    pub fn new(resource_task: ResourceTask, url: Url, pipeline: Option<PipelineId>)
               -> PendingAsyncLoad {
        PendingAsyncLoad {
            resource_task: resource_task,
            url: url,
            pipeline: pipeline,
            guard: PendingLoadGuard { loaded: false, },
        }
    }

    /// Initiate the network request associated with this pending load.
    pub fn load(mut self) -> Receiver<LoadResponse> {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let (sender, receiver) = channel();
        let consumer = LoadConsumer::Channel(sender);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
        receiver
    }

    /// Initiate the network request associated with this pending load, using the provided target.
    pub fn load_async(mut self, listener: Box<AsyncResponseTarget + Send>) {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let consumer = LoadConsumer::Listener(listener);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
    }
}

/// Message sent in response to `Load`.  Contains metadata, and a port
/// for receiving the data.
///
/// Even if loading fails immediately, we send one of these and the
/// progress_port will provide the error.
pub struct LoadResponse {
    /// Metadata, such as from HTTP headers.
    pub metadata: Metadata,
    /// Port for reading data.
    pub progress_port: Receiver<ProgressMsg>,
}

#[derive(Clone)]
pub struct ResourceCORSData {
    /// CORS Preflight flag
    pub preflight: bool,
    /// Origin of CORS Request
    pub origin: Url
}

/// Metadata about a loaded resource, such as is obtained from HTTP headers.
#[derive(Clone)]
pub struct Metadata {
    /// Final URL after redirects.
    pub final_url: Url,

    /// MIME type / subtype.
    pub content_type: Option<(ContentType)>,

    /// Character set.
    pub charset: Option<String>,

    /// Headers
    pub headers: Option<Headers>,

    /// HTTP Status
    pub status: Option<RawStatus>,
}

impl Metadata {
    /// Metadata with defaults for everything optional.
    pub fn default(url: Url) -> Self {
        Metadata {
            final_url:    url,
            content_type: None,
            charset:      None,
            headers: None,
            // https://fetch.spec.whatwg.org/#concept-response-status-message
            status: Some(RawStatus(200, "OK".into())),
        }
    }

    /// Extract the parts of a Mime that we care about.
    pub fn set_content_type(&mut self, content_type: Option<&Mime>) {
        match content_type {
            None => (),
            Some(mime) => {
                self.content_type = Some(ContentType(mime.clone()));
                let &Mime(_, _, ref parameters) = mime;
                for &(ref k, ref v) in parameters.iter() {
                    if &Attr::Charset == k {
                        self.charset = Some(v.to_string());
                    }
                }
            }
        }
    }
}

/// The creator of a given cookie
#[derive(PartialEq, Copy, Clone)]
pub enum CookieSource {
    /// An HTTP API
    HTTP,
    /// A non-HTTP API
    NonHTTP,
}

/// Messages sent in response to a `Load` message
#[derive(PartialEq,Debug)]
pub enum ProgressMsg {
    /// Binary data - there may be multiple of these
    Payload(Vec<u8>),
    /// Indicates loading is complete, either successfully or not
    Done(Result<(), String>)
}

/// Convenience function for synchronously loading a whole resource.
pub fn load_whole_resource(resource_task: &ResourceTask, url: Url)
        -> Result<(Metadata, Vec<u8>), String> {
    let (start_chan, start_port) = channel();
    resource_task.send(ControlMsg::Load(LoadData::new(url, None), LoadConsumer::Channel(start_chan))).unwrap();
    let response = start_port.recv().unwrap();

    let mut buf = vec!();
    loop {
        match response.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => buf.push_all(&data),
            ProgressMsg::Done(Ok(()))  => return Ok((response.metadata, buf)),
            ProgressMsg::Done(Err(e))  => return Err(e)
        }
    }
}

/// Load a URL asynchronously and iterate over chunks of bytes from the response.
pub fn load_bytes_iter(pending: PendingAsyncLoad) -> (Metadata, ProgressMsgPortIterator) {
    let input_port = pending.load();
    let response = input_port.recv().unwrap();
    let iter = ProgressMsgPortIterator { progress_port: response.progress_port };
    (response.metadata, iter)
}

/// Iterator that reads chunks of bytes from a ProgressMsg port
pub struct ProgressMsgPortIterator {
    progress_port: Receiver<ProgressMsg>
}

impl Iterator for ProgressMsgPortIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => Some(data),
            ProgressMsg::Done(Ok(()))  => None,
            ProgressMsg::Done(Err(e))  => {
                error!("error receiving bytes: {}", e);
                None
            }
        }
    }
}

<commit_msg>Auto merge of #6162 - r0e:testing, r=jdm<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(box_syntax)]
#![feature(collections)]
#![feature(rustc_private)]
#![feature(slice_patterns)]
#![feature(step_by)]

extern crate geom;
extern crate hyper;
#[macro_use]
extern crate log;
extern crate png;
extern crate stb_image;
extern crate url;
extern crate util;
extern crate msg;

use hyper::header::{ContentType, Headers};
use hyper::http::RawStatus;
use hyper::method::Method;
use hyper::mime::{Mime, Attr};
use msg::constellation_msg::{PipelineId};
use url::Url;

use std::sync::mpsc::{channel, Receiver, Sender};
use std::thread;

pub mod image_cache_task;
pub mod storage_task;

/// Image handling.
///
/// It may be surprising that this goes in the network crate as opposed to the graphics crate.
/// However, image handling is generally very integrated with the network stack (especially where
/// caching is involved) and as a result it must live in here.
pub mod image {
    pub mod base;
}

#[derive(Clone)]
pub struct LoadData {
    pub url: Url,
    pub method: Method,
    /// Headers that will apply to the initial request only
    pub headers: Headers,
    /// Headers that will apply to the initial request and any redirects
    pub preserved_headers: Headers,
    pub data: Option<Vec<u8>>,
    pub cors: Option<ResourceCORSData>,
    pub pipeline_id: Option<PipelineId>,
}

impl LoadData {
    pub fn new(url: Url, id: Option<PipelineId>) -> LoadData {
        LoadData {
            url: url,
            method: Method::Get,
            headers: Headers::new(),
            preserved_headers: Headers::new(),
            data: None,
            cors: None,
            pipeline_id: id,
        }
    }
}

/// A listener for asynchronous network events. Cancelling the underlying request is unsupported.
pub trait AsyncResponseListener {
    /// The response headers for a request have been received.
    fn headers_available(&self, metadata: Metadata);
    /// A portion of the response body has been received. This data is unavailable after
    /// this method returned, and must be stored accordingly.
    fn data_available(&self, payload: Vec<u8>);
    /// The response is complete. If the provided status is an Err value, there is no guarantee
    /// that the response body was completely read.
    fn response_complete(&self, status: Result<(), String>);
}

/// Data for passing between threads/processes to indicate a particular action to
/// take on a provided network listener.
pub enum ResponseAction {
    /// Invoke headers_available
    HeadersAvailable(Metadata),
    /// Invoke data_available
    DataAvailable(Vec<u8>),
    /// Invoke response_complete
    ResponseComplete(Result<(), String>)
}

impl ResponseAction {
    /// Execute the default action on a provided listener.
    pub fn process(self, listener: &AsyncResponseListener) {
        match self {
            ResponseAction::HeadersAvailable(m) => listener.headers_available(m),
            ResponseAction::DataAvailable(d) => listener.data_available(d),
            ResponseAction::ResponseComplete(r) => listener.response_complete(r),
        }
    }
}

/// A target for async networking events. Commonly used to dispatch a runnable event to another
/// thread storing the wrapped closure for later execution.
pub trait AsyncResponseTarget {
    fn invoke_with_listener(&self, action: ResponseAction);
}

/// A wrapper for a network load that can either be channel or event-based.
pub enum LoadConsumer {
    Channel(Sender<LoadResponse>),
    Listener(Box<AsyncResponseTarget + Send>),
}

/// Handle to a resource task
pub type ResourceTask = Sender<ControlMsg>;

pub enum ControlMsg {
    /// Request the data associated with a particular URL
    Load(LoadData, LoadConsumer),
    /// Store a set of cookies for a given originating URL
    SetCookiesForUrl(Url, String, CookieSource),
    /// Retrieve the stored cookies for a given URL
    GetCookiesForUrl(Url, Sender<Option<String>>, CookieSource),
    Exit
}

/// Initialized but unsent request. Encapsulates everything necessary to instruct
/// the resource task to make a new request. The `load` method *must* be called before
/// destruction or the task will panic.
pub struct PendingAsyncLoad {
    resource_task: ResourceTask,
    url: Url,
    pipeline: Option<PipelineId>,
    guard: PendingLoadGuard,
}

struct PendingLoadGuard {
    loaded: bool,
}

impl PendingLoadGuard {
    fn neuter(&mut self) {
        self.loaded = true;
    }
}

impl Drop for PendingLoadGuard {
    fn drop(&mut self) {
        if !thread::panicking() {
            assert!(self.loaded)
        }
    }
}

impl PendingAsyncLoad {
    pub fn new(resource_task: ResourceTask, url: Url, pipeline: Option<PipelineId>)
               -> PendingAsyncLoad {
        PendingAsyncLoad {
            resource_task: resource_task,
            url: url,
            pipeline: pipeline,
            guard: PendingLoadGuard { loaded: false, },
        }
    }

    /// Initiate the network request associated with this pending load.
    pub fn load(mut self) -> Receiver<LoadResponse> {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let (sender, receiver) = channel();
        let consumer = LoadConsumer::Channel(sender);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
        receiver
    }

    /// Initiate the network request associated with this pending load, using the provided target.
    pub fn load_async(mut self, listener: Box<AsyncResponseTarget + Send>) {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let consumer = LoadConsumer::Listener(listener);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
    }
}

/// Message sent in response to `Load`.  Contains metadata, and a port
/// for receiving the data.
///
/// Even if loading fails immediately, we send one of these and the
/// progress_port will provide the error.
pub struct LoadResponse {
    /// Metadata, such as from HTTP headers.
    pub metadata: Metadata,
    /// Port for reading data.
    pub progress_port: Receiver<ProgressMsg>,
}

#[derive(Clone)]
pub struct ResourceCORSData {
    /// CORS Preflight flag
    pub preflight: bool,
    /// Origin of CORS Request
    pub origin: Url
}

/// Metadata about a loaded resource, such as is obtained from HTTP headers.
#[derive(Clone)]
pub struct Metadata {
    /// Final URL after redirects.
    pub final_url: Url,

    /// MIME type / subtype.
    pub content_type: Option<(ContentType)>,

    /// Character set.
    pub charset: Option<String>,

    /// Headers
    pub headers: Option<Headers>,

    /// HTTP Status
    pub status: Option<RawStatus>,
}

impl Metadata {
    /// Metadata with defaults for everything optional.
    pub fn default(url: Url) -> Self {
        Metadata {
            final_url:    url,
            content_type: None,
            charset:      None,
            headers: None,
            // https://fetch.spec.whatwg.org/#concept-response-status-message
            status: Some(RawStatus(200, "OK".into())),
        }
    }

    /// Extract the parts of a Mime that we care about.
    pub fn set_content_type(&mut self, content_type: Option<&Mime>) {
        match content_type {
            None => (),
            Some(mime) => {
                self.content_type = Some(ContentType(mime.clone()));
                let &Mime(_, _, ref parameters) = mime;
                for &(ref k, ref v) in parameters.iter() {
                    if &Attr::Charset == k {
                        self.charset = Some(v.to_string());
                    }
                }
            }
        }
    }
}

/// The creator of a given cookie
#[derive(PartialEq, Copy, Clone)]
pub enum CookieSource {
    /// An HTTP API
    HTTP,
    /// A non-HTTP API
    NonHTTP,
}

/// Messages sent in response to a `Load` message
#[derive(PartialEq,Debug)]
pub enum ProgressMsg {
    /// Binary data - there may be multiple of these
    Payload(Vec<u8>),
    /// Indicates loading is complete, either successfully or not
    Done(Result<(), String>)
}

/// Convenience function for synchronously loading a whole resource.
pub fn load_whole_resource(resource_task: &ResourceTask, url: Url)
        -> Result<(Metadata, Vec<u8>), String> {
    let (start_chan, start_port) = channel();
    resource_task.send(ControlMsg::Load(LoadData::new(url, None), LoadConsumer::Channel(start_chan))).unwrap();
    let response = start_port.recv().unwrap();

    let mut buf = vec!();
    loop {
        match response.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => buf.push_all(&data),
            ProgressMsg::Done(Ok(()))  => return Ok((response.metadata, buf)),
            ProgressMsg::Done(Err(e))  => return Err(e)
        }
    }
}

/// Load a URL asynchronously and iterate over chunks of bytes from the response.
pub fn load_bytes_iter(pending: PendingAsyncLoad) -> (Metadata, ProgressMsgPortIterator) {
    let input_port = pending.load();
    let response = input_port.recv().unwrap();
    let iter = ProgressMsgPortIterator { progress_port: response.progress_port };
    (response.metadata, iter)
}

/// Iterator that reads chunks of bytes from a ProgressMsg port
pub struct ProgressMsgPortIterator {
    progress_port: Receiver<ProgressMsg>
}

impl Iterator for ProgressMsgPortIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => Some(data),
            ProgressMsg::Done(Ok(()))  => None,
            ProgressMsg::Done(Err(e))  => {
                error!("error receiving bytes: {}", e);
                None
            }
        }
    }
}

<|endoftext|>


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

//! Generic implementations of some DOM APIs so they can be shared between Servo
//! and Gecko.

use context::QuirksMode;
use selectors::{Element, NthIndexCache, SelectorList};
use selectors::matching::{self, MatchingContext, MatchingMode};

/// <https://dom.spec.whatwg.org/#dom-element-matches>
pub fn element_matches<E>(
    element: &E,
    selector_list: &SelectorList<E::Impl>,
    quirks_mode: QuirksMode,
) -> bool
where
    E: Element,
{
    let mut context = MatchingContext::new(
        MatchingMode::Normal,
        None,
        None,
        quirks_mode,
    );
    context.scope_element = Some(element.opaque());
    matching::matches_selector_list(selector_list, element, &mut context)
}

/// <https://dom.spec.whatwg.org/#dom-element-closest>
pub fn element_closest<E>(
    element: E,
    selector_list: &SelectorList<E::Impl>,
    quirks_mode: QuirksMode,
) -> Option<E>
where
    E: Element,
{
    let mut nth_index_cache = NthIndexCache::default();

    let mut context = MatchingContext::new(
        MatchingMode::Normal,
        None,
        Some(&mut nth_index_cache),
        quirks_mode,
    );
    context.scope_element = Some(element.opaque());

    let mut current = Some(element);
    while let Some(element) = current.take() {
        if matching::matches_selector_list(selector_list, &element, &mut context) {
            return Some(element);
        }
        current = element.parent_element();
    }

    return None;
}
<commit_msg>style: Add a very simple invalidation-based querySelector.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Generic implementations of some DOM APIs so they can be shared between Servo
//! and Gecko.

use context::QuirksMode;
use dom::{TElement, TNode};
use invalidation::element::invalidator::{Invalidation, InvalidationProcessor, InvalidationVector};
use selectors::{Element, NthIndexCache, SelectorList};
use selectors::matching::{self, MatchingContext, MatchingMode};
use smallvec::SmallVec;

/// <https://dom.spec.whatwg.org/#dom-element-matches>
pub fn element_matches<E>(
    element: &E,
    selector_list: &SelectorList<E::Impl>,
    quirks_mode: QuirksMode,
) -> bool
where
    E: Element,
{
    let mut context = MatchingContext::new(
        MatchingMode::Normal,
        None,
        None,
        quirks_mode,
    );
    context.scope_element = Some(element.opaque());
    matching::matches_selector_list(selector_list, element, &mut context)
}

/// <https://dom.spec.whatwg.org/#dom-element-closest>
pub fn element_closest<E>(
    element: E,
    selector_list: &SelectorList<E::Impl>,
    quirks_mode: QuirksMode,
) -> Option<E>
where
    E: Element,
{
    let mut nth_index_cache = NthIndexCache::default();

    let mut context = MatchingContext::new(
        MatchingMode::Normal,
        None,
        Some(&mut nth_index_cache),
        quirks_mode,
    );
    context.scope_element = Some(element.opaque());

    let mut current = Some(element);
    while let Some(element) = current.take() {
        if matching::matches_selector_list(selector_list, &element, &mut context) {
            return Some(element);
        }
        current = element.parent_element();
    }

    return None;
}

/// The result of a querySelector call.
pub type QuerySelectorResult<E> = SmallVec<[E; 128]>;

/// The query kind we're doing (either only the first descendant that matches or
/// all of them).
pub enum QuerySelectorKind {
    /// <https://dom.spec.whatwg.org/#dom-parentnode-queryselectorall>
    All,
    /// <https://dom.spec.whatwg.org/#dom-parentnode-queryselector>
    First,
}

struct QuerySelectorProcessor<'a, E: TElement + 'a> {
    kind: QuerySelectorKind,
    results: &'a mut QuerySelectorResult<E>,
    matching_context: MatchingContext<'a, E::Impl>,
    selector_list: &'a SelectorList<E::Impl>,
}

impl<'a, E> InvalidationProcessor<'a, E> for QuerySelectorProcessor<'a, E>
where
    E: TElement + 'a,
{
    fn collect_invalidations(
        &mut self,
        _element: E,
        self_invalidations: &mut InvalidationVector<'a>,
        descendant_invalidations: &mut InvalidationVector<'a>,
        _sibling_invalidations: &mut InvalidationVector<'a>,
    ) -> bool {
        // FIXME(emilio): If the element is not a root element, and
        // selector_list has any descendant combinator, we need to do extra work
        // in order to handle properly things like:
        //
        //

        //

        //

        //

        //

        //
        // b.querySelector('#a div'); // Should return "c".
        //
        let target_vector =
            if self.matching_context.scope_element.is_some() {
                descendant_invalidations
            } else {
                self_invalidations
            };

        for selector in self.selector_list.0.iter() {
            target_vector.push(Invalidation::new(selector, 0))
        }

        false
    }

    fn matching_context(&mut self) -> &mut MatchingContext<'a, E::Impl> {
        &mut self.matching_context
    }

    fn should_process_descendants(&mut self, _: E) -> bool {
        match self.kind {
            QuerySelectorKind::All => true,
            QuerySelectorKind::First => self.results.is_empty(),
        }
    }

    fn invalidated_self(&mut self, e: E) {
        self.results.push(e);
    }

    fn recursion_limit_exceeded(&mut self, _e: E) {}
    fn invalidated_descendants(&mut self, _e: E, _child: E) {}
}

/// <https://dom.spec.whatwg.org/#dom-parentnode-queryselector>
pub fn query_selector<E: TElement>(
    root: E::ConcreteNode,
    selector_list: &SelectorList<E::Impl>,
    results: &mut QuerySelectorResult<E>,
    kind: QuerySelectorKind,
    quirks_mode: QuirksMode,
) {
    use invalidation::element::invalidator::TreeStyleInvalidator;

    let mut nth_index_cache = NthIndexCache::default();
    let mut matching_context = MatchingContext::new(
        MatchingMode::Normal,
        None,
        Some(&mut nth_index_cache),
        quirks_mode,
    );

    let root_element = root.as_element();
    matching_context.scope_element = root_element.map(|e| e.opaque());

    let mut processor = QuerySelectorProcessor {
        kind,
        results,
        matching_context,
        selector_list,
    };

    match root_element {
        Some(e) => {
            TreeStyleInvalidator::new(
                e,
                /* stack_limit_checker = */ None,
                &mut processor,
            ).invalidate();
        }
        None => {
            for node in root.dom_children() {
                if let Some(e) = node.as_element() {
                    TreeStyleInvalidator::new(
                        e,
                        /* stack_limit_checker = */ None,
                        &mut processor,
                    ).invalidate();
                }
            }
        }
    }
}
<|endoftext|>


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

use tokio::sync::mpsc;

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

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

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

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

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

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

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut, if remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
            else => unreachable!(),
        }
    }
}

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

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

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

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

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

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

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

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

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

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

        let context = context.clone();
        let remote_upstream = remote_upstream.clone();
        let mut qtx = qtx.clone();

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

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), context.local_dns(), remote_upstream, question).await;

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

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

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

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

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

use tokio::sync::mpsc;

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

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

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

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

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

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

    let decider = async {
        let local_response = local_response_fut.await;
        if should_forward_by_response(acl, &local_response, query) {
            None
        } else {
            Some(local_response)
        }
    };
    tokio::pin!(remote_response_fut, decider);
    let mut use_remote = false;
    let mut remote_response = None;
    loop {
        tokio::select! {
            response = &mut remote_response_fut, if remote_response.is_none() => {
                if use_remote {
                    debug!("pick remote response (response): {:?}", response);
                    return (response, true);
                } else {
                    remote_response = Some(response);
                }
            }
            decision = &mut decider, if !use_remote => {
                if let Some(local_response) = decision {
                    debug!("pick local response (response): {:?}", local_response);
                    return (local_response, false);
                } else if let Some(remote_response) = remote_response {
                    debug!("pick remote response (response): {:?}", remote_response);
                    return (remote_response, true);
                } else {
                    use_remote = true;
                }
            }
            else => unreachable!(),
        }
    }
}

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

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

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

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

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

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

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

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

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

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

        let context = context.clone();
        let remote_upstream = remote_upstream.clone();
        let mut qtx = qtx.clone();

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

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), context.local_dns(), remote_upstream, question).await;

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

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

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

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

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


<commit_before>pub fn normalize(num: &str) -> String {
    let num = num.replace("IIIII", "V");
    let num = num.replace("IIII", "IV");
    let num = num.replace("VV", "X");
    let num = num.replace("VIV", "IX");
    let num = num.replace("XXXXX", "L");
    let num = num.replace("XXXX", "XL");
    let num = num.replace("LL", "C");
    let num = num.replace("LXL", "XC");
    let num = num.replace("CCCCC", "D");
    let num = num.replace("CCCC", "CD");
    num
}

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

    #[test]
    fn iiiii_becomes_v() {
        assert_eq!("V", normalize("IIIII"));
    }

    #[test]
    fn iiii_becomes_iv() {
        assert_eq!("IV", normalize("IIII"));
    }

    #[test]
    fn vv_becomes_x() {
        assert_eq!("X", normalize("VV"));
    }

    #[test]
    fn viiii_becomes_ix() {
        assert_eq!("IX", normalize("VIIII"));
    }

    #[test]
    fn viiiii_becomes_x() {
        assert_eq!("X", normalize("VIIIII"));
    }

    #[test]
    fn vviv_becomes_xiv() {
        assert_eq!("XIV", normalize("VVIV"));
    }

    #[test]
    fn xxxxx_becomes_l() {
        assert_eq!("L", normalize("XXXXX"));
    }

    #[test]
    fn xxxx_becomes_xl() {
        assert_eq!("XL", normalize("XXXX"));
    }

    #[test]
    fn lxxxx_becomes_xc() {
        assert_eq!("XC", normalize("LXXXX"));
    }

    #[test]
    fn lxxxxx_becomes_c() {
        assert_eq!("C", normalize("LXXXXX"));
    }

    #[test]
    fn llxl_becomes_cxl() {
        assert_eq!("CXL", normalize("LLXL"));
    }

    #[test]
    fn cccc_becomes_cd() {
        assert_eq!("CD", normalize("CCCC"));
    }

    #[test]
    fn ccccc_becomes_d() {
        assert_eq!("D", normalize("CCCCC"));
    }
}
<commit_msg>Initial support for M in rust<commit_after>pub fn normalize(num: &str) -> String {
    let num = num.replace("IIIII", "V");
    let num = num.replace("IIII", "IV");
    let num = num.replace("VV", "X");
    let num = num.replace("VIV", "IX");
    let num = num.replace("XXXXX", "L");
    let num = num.replace("XXXX", "XL");
    let num = num.replace("LL", "C");
    let num = num.replace("LXL", "XC");
    let num = num.replace("CCCCC", "D");
    let num = num.replace("CCCC", "CD");
    let num = num.replace("DD", "M");
    num
}

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

    #[test]
    fn iiiii_becomes_v() {
        assert_eq!("V", normalize("IIIII"));
    }

    #[test]
    fn iiii_becomes_iv() {
        assert_eq!("IV", normalize("IIII"));
    }

    #[test]
    fn vv_becomes_x() {
        assert_eq!("X", normalize("VV"));
    }

    #[test]
    fn viiii_becomes_ix() {
        assert_eq!("IX", normalize("VIIII"));
    }

    #[test]
    fn viiiii_becomes_x() {
        assert_eq!("X", normalize("VIIIII"));
    }

    #[test]
    fn vviv_becomes_xiv() {
        assert_eq!("XIV", normalize("VVIV"));
    }

    #[test]
    fn xxxxx_becomes_l() {
        assert_eq!("L", normalize("XXXXX"));
    }

    #[test]
    fn xxxx_becomes_xl() {
        assert_eq!("XL", normalize("XXXX"));
    }

    #[test]
    fn lxxxx_becomes_xc() {
        assert_eq!("XC", normalize("LXXXX"));
    }

    #[test]
    fn lxl_becomes_xc() {
        assert_eq!("XC", normalize("LXL"));
    }

    #[test]
    fn lxxxxx_becomes_c() {
        assert_eq!("C", normalize("LXXXXX"));
    }

    #[test]
    fn ll_becomes_c() {
        assert_eq!("C", normalize("LL"));
    }

    #[test]
    fn llxl_becomes_cxl() {
        assert_eq!("CXL", normalize("LLXL"));
    }

    #[test]
    fn cccc_becomes_cd() {
        assert_eq!("CD", normalize("CCCC"));
    }

    #[test]
    fn ccccc_becomes_d() {
        assert_eq!("D", normalize("CCCCC"));
    }

    #[test]
    fn dd_becomes_m() {
        assert_eq!("M", normalize("DD"));
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaList(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc::plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));
}
<commit_msg>fix(rustup): nightly rust moved Registry into rustc_plugin<commit_after>#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]
#![cfg_attr(not(feature = "with-syntex"), plugin(quasi_macros))]

extern crate aster;
extern crate quasi;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/lib.rs"));

#[cfg(not(feature = "with-syntex"))]
include!("lib.rs.in");

#[cfg(feature = "with-syntex")]
pub fn register(reg: &mut syntex::Registry) {
    use syntax::{ast, fold};

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", ser::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", de::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaList(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(ser::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(de::expand_derive_deserialize)));
}
<|endoftext|>