<commit_before>use std::collections::{BTreeMap, HashMap};
use std::hash::Hash;
use webcore::try_from::{TryFrom, TryInto};
use webcore::value::{Reference, Value, ConversionError, FromReferenceUnchecked};
use webcore::serialization::{JsSerializable, deserialize_object};

/// A type representing a JavaScript object.
#[derive(Clone, PartialEq, Debug)]
pub struct Object( Reference );

impl FromReferenceUnchecked for Object {
    unsafe fn from_reference_unchecked( reference: Reference ) -> Self {
        Object( reference )
    }
}

impl From< Object > for Reference {
    fn from( object: Object ) -> Self {
        object.0.clone()
    }
}

impl Object {
    #[inline]
    pub(crate) fn as_reference( &self ) -> &Reference {
        &self.0
    }

    pub fn len( &self ) -> usize {
        js!(
            return Object.keys( @{self} ).length;
        ).try_into().unwrap()
    }
}

impl From< Object > for BTreeMap< String, Value > {
    fn from( object: Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a Object > for BTreeMap< String, Value > {
    fn from( object: &'a Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a mut Object > for BTreeMap< String, Value > {
    fn from( object: &'a mut Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl From< Object > for HashMap< String, Value > {
    fn from( object: Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a Object > for HashMap< String, Value > {
    fn from( object: &'a Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a mut Object > for HashMap< String, Value > {
    fn from( object: &'a mut Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}
/*
impl< K, V > From< BTreeMap< K, V > > for Object where K: AsRef< str >, V: Into< Value > {
    #[inline]
    fn from( map: BTreeMap< K, V > ) -> Self {
        unimplemented!();
    }
}

impl< 'a, K, V > From< &'a BTreeMap< K, V > > for Object where K: AsRef< str >, &'a V: Into< Value > {
    #[inline]
    fn from( map: &'a BTreeMap< K, V > ) -> Self {
        unimplemented!();
    }
}

impl< 'a, K, V > From< &'a mut BTreeMap< K, V > > for Object where K: AsRef< str >, &'a V: Into< Value > {
    #[inline]
    fn from( map: &'a mut BTreeMap< K, V > ) -> Self {
        let map: &BTreeMap< _, _ > = map;
        map.into()
    }
}

impl< K, V > From< HashMap< K, V > > for Object where K: AsRef< str > + Eq + Hash, V: Into< Value > {
    #[inline]
    fn from( map: HashMap< K, V > ) -> Self {
        unimplemented!();
    }
}

impl< 'a, K, V > From< &'a HashMap< K, V > > for Object where K: AsRef< str > + Eq + Hash, &'a V: Into< Value > {
    #[inline]
    fn from( map: &'a HashMap< K, V > ) -> Self {
        unimplemented!();
    }
}

impl< 'a, K, V > From< &'a mut HashMap< K, V > > for Object where K: AsRef< str > + Eq + Hash, &'a V: Into< Value > {
    #[inline]
    fn from( map: &'a mut HashMap< K, V > ) -> Self {
        let map: &HashMap< _, _ > = map;
        map.into()
    }
}
*/

// TODO: It would be nice to specialize this for values which are already of type Value.
impl< K: AsRef< str >, V: JsSerializable > From< BTreeMap< K, V > > for Object {
    #[inline]
    fn from( value: BTreeMap< K, V > ) -> Self {
        (&value).into()
    }
}

impl< 'a, K, V > From< &'a BTreeMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a BTreeMap< K, V > ) -> Self {
        // TODO: Do something more efficient here?
        let value = js! {
            return @{value};
        };

        match value {
            Value::Object( object ) => return object,
            _ => unreachable!()
        }
    }
}

impl< 'a, K, V > From< &'a mut BTreeMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a mut BTreeMap< K, V > ) -> Self {
        let value: &BTreeMap< K, V > = value;
        value.into()
    }
}

impl< K, V > From< HashMap< K, V > > for Object where K: AsRef< str > + Hash + Eq, V: JsSerializable {
    #[inline]
    fn from( value: HashMap< K, V > ) -> Self {
        (&value).into()
    }
}

impl< 'a, K, V > From< &'a HashMap< K, V > > for Object where K: AsRef< str > + Hash + Eq, V: JsSerializable {
    #[inline]
    fn from( value: &'a HashMap< K, V > ) -> Self {
        // TODO: Do something more efficient here?
        let value = js! {
            return @{value};
        };

        match value {
            Value::Object( object ) => return object,
            _ => unreachable!()
        }
    }
}

impl< 'a, K: Hash + Eq, V > From< &'a mut HashMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a mut HashMap< K, V > ) -> Self {
        let value: &HashMap< K, V > = value;
        value.into()
    }
}

impl< E: Into< ConversionError >, V: TryFrom< Value, Error = E > > TryFrom< Object > for BTreeMap< String, V > {
    type Error = ConversionError;

    fn try_from( object: Object ) -> Result< Self, Self::Error > {
        deserialize_object( object.as_reference(), |deserializer| -> Result< BTreeMap< String, V >, E > {
            let mut output = BTreeMap::new();
            for (key, value) in deserializer {
                output.insert( key, value.try_into()? );
            }
            Ok( output )
        }).map_err( |err| err.into() )
    }
}

impl< E: Into< ConversionError >, V: TryFrom< Value, Error = E > > TryFrom< Object > for HashMap< String, V > {
    type Error = ConversionError;

    fn try_from( object: Object ) -> Result< Self, Self::Error > {
        deserialize_object( object.as_reference(), |deserializer| -> Result< HashMap< String, V >, E > {
            let mut output = HashMap::with_capacity( deserializer.len() );
            for (key, value) in deserializer {
                output.insert( key, value.try_into()? );
            }

            Ok( output )
        }).map_err( |err| err.into() )
    }
}
<commit_msg>Remove a bunch of commented out code that I've forgot to delete<commit_after>use std::collections::{BTreeMap, HashMap};
use std::hash::Hash;
use webcore::try_from::{TryFrom, TryInto};
use webcore::value::{Reference, Value, ConversionError, FromReferenceUnchecked};
use webcore::serialization::{JsSerializable, deserialize_object};

/// A type representing a JavaScript object.
#[derive(Clone, PartialEq, Debug)]
pub struct Object( Reference );

impl FromReferenceUnchecked for Object {
    unsafe fn from_reference_unchecked( reference: Reference ) -> Self {
        Object( reference )
    }
}

impl From< Object > for Reference {
    fn from( object: Object ) -> Self {
        object.0.clone()
    }
}

impl Object {
    #[inline]
    pub(crate) fn as_reference( &self ) -> &Reference {
        &self.0
    }

    pub fn len( &self ) -> usize {
        js!(
            return Object.keys( @{self} ).length;
        ).try_into().unwrap()
    }
}

impl From< Object > for BTreeMap< String, Value > {
    fn from( object: Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a Object > for BTreeMap< String, Value > {
    fn from( object: &'a Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a mut Object > for BTreeMap< String, Value > {
    fn from( object: &'a mut Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl From< Object > for HashMap< String, Value > {
    fn from( object: Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a Object > for HashMap< String, Value > {
    fn from( object: &'a Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

impl< 'a > From< &'a mut Object > for HashMap< String, Value > {
    fn from( object: &'a mut Object ) -> Self {
        deserialize_object( &object.0, |iter| iter.collect() )
    }
}

// TODO: It would be nice to specialize this for values which are already of type Value.
impl< K: AsRef< str >, V: JsSerializable > From< BTreeMap< K, V > > for Object {
    #[inline]
    fn from( value: BTreeMap< K, V > ) -> Self {
        (&value).into()
    }
}

impl< 'a, K, V > From< &'a BTreeMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a BTreeMap< K, V > ) -> Self {
        // TODO: Do something more efficient here?
        let value = js! {
            return @{value};
        };

        match value {
            Value::Object( object ) => return object,
            _ => unreachable!()
        }
    }
}

impl< 'a, K, V > From< &'a mut BTreeMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a mut BTreeMap< K, V > ) -> Self {
        let value: &BTreeMap< K, V > = value;
        value.into()
    }
}

impl< K, V > From< HashMap< K, V > > for Object where K: AsRef< str > + Hash + Eq, V: JsSerializable {
    #[inline]
    fn from( value: HashMap< K, V > ) -> Self {
        (&value).into()
    }
}

impl< 'a, K, V > From< &'a HashMap< K, V > > for Object where K: AsRef< str > + Hash + Eq, V: JsSerializable {
    #[inline]
    fn from( value: &'a HashMap< K, V > ) -> Self {
        // TODO: Do something more efficient here?
        let value = js! {
            return @{value};
        };

        match value {
            Value::Object( object ) => return object,
            _ => unreachable!()
        }
    }
}

impl< 'a, K: Hash + Eq, V > From< &'a mut HashMap< K, V > > for Object where K: AsRef< str >, V: JsSerializable {
    #[inline]
    fn from( value: &'a mut HashMap< K, V > ) -> Self {
        let value: &HashMap< K, V > = value;
        value.into()
    }
}

impl< E: Into< ConversionError >, V: TryFrom< Value, Error = E > > TryFrom< Object > for BTreeMap< String, V > {
    type Error = ConversionError;

    fn try_from( object: Object ) -> Result< Self, Self::Error > {
        deserialize_object( object.as_reference(), |deserializer| -> Result< BTreeMap< String, V >, E > {
            let mut output = BTreeMap::new();
            for (key, value) in deserializer {
                output.insert( key, value.try_into()? );
            }
            Ok( output )
        }).map_err( |err| err.into() )
    }
}

impl< E: Into< ConversionError >, V: TryFrom< Value, Error = E > > TryFrom< Object > for HashMap< String, V > {
    type Error = ConversionError;

    fn try_from( object: Object ) -> Result< Self, Self::Error > {
        deserialize_object( object.as_reference(), |deserializer| -> Result< HashMap< String, V >, E > {
            let mut output = HashMap::with_capacity( deserializer.len() );
            for (key, value) in deserializer {
                output.insert( key, value.try_into()? );
            }

            Ok( output )
        }).map_err( |err| err.into() )
    }
}
<|endoftext|>


<commit_before>//! HTTP response builder.

use std::{
    cell::{Ref, RefMut},
    fmt,
    future::Future,
    pin::Pin,
    str,
    task::{Context, Poll},
};

use bytes::Bytes;
use futures_core::Stream;

use crate::{
    body::{Body, BodyStream, ResponseBody},
    error::Error,
    extensions::Extensions,
    header::{IntoHeaderPair, IntoHeaderValue},
    http::{header, Error as HttpError, StatusCode},
    message::{BoxedResponseHead, ConnectionType, ResponseHead},
    Response,
};

/// An HTTP response builder.
///
/// Used to construct an instance of `Response` using a builder pattern. Response builders are often
/// created using [`Response::build`].
///
/// # Examples
/// ```
/// use actix_http::{Response, ResponseBuilder, body, http::StatusCode, http::header};
///
/// # actix_rt::System::new().block_on(async {
/// let mut res: Response<_> = Response::build(StatusCode::OK)
///     .content_type(mime::APPLICATION_JSON)
///     .insert_header((header::SERVER, "my-app/1.0"))
///     .append_header((header::SET_COOKIE, "a=1"))
///     .append_header((header::SET_COOKIE, "b=2"))
///     .body("1234");
///
/// assert_eq!(res.status(), StatusCode::OK);
/// assert_eq!(body::to_bytes(res.take_body()).await.unwrap(), &b"1234"[..]);
///
/// assert!(res.headers().contains_key("server"));
/// assert_eq!(res.headers().get_all("set-cookie").count(), 2);
/// # })
/// ```
pub struct ResponseBuilder {
    head: Option<BoxedResponseHead>,
    err: Option<HttpError>,
}

impl ResponseBuilder {
    /// Create response builder
    ///
    /// # Examples
    /// ```
    /// use actix_http::{Response, ResponseBuilder, http::StatusCode};
    ///
    /// let res: Response<_> = ResponseBuilder::default().finish();
    /// assert_eq!(res.status(), StatusCode::OK);
    /// ```
    #[inline]
    pub fn new(status: StatusCode) -> Self {
        ResponseBuilder {
            head: Some(BoxedResponseHead::new(status)),
            err: None,
        }
    }

    /// Set HTTP status code of this response.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::StatusCode};
    ///
    /// let res = ResponseBuilder::default().status(StatusCode::NOT_FOUND).finish();
    /// assert_eq!(res.status(), StatusCode::NOT_FOUND);
    /// ```
    #[inline]
    pub fn status(&mut self, status: StatusCode) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.status = status;
        }
        self
    }

    /// Insert a header, replacing any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .insert_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .insert_header(("X-TEST", "value"))
    ///     .finish();
    ///
    /// assert!(res.headers().contains_key("content-type"));
    /// assert!(res.headers().contains_key("x-test"));
    /// ```
    pub fn insert_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => {
                    parts.headers.insert(key, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Append a header, keeping any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .append_header(("X-TEST", "value1"))
    ///     .append_header(("X-TEST", "value2"))
    ///     .finish();
    ///
    /// assert_eq!(res.headers().get_all("content-type").count(), 1);
    /// assert_eq!(res.headers().get_all("x-test").count(), 2);
    /// ```
    pub fn append_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match header.try_into_header_pair() {
                Ok((key, value)) => parts.headers.append(key, value),
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Set the custom reason for the response.
    #[inline]
    pub fn reason(&mut self, reason: &'static str) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.reason = Some(reason);
        }
        self
    }

    /// Set connection type to KeepAlive
    #[inline]
    pub fn keep_alive(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::KeepAlive);
        }
        self
    }

    /// Set connection type to Upgrade
    #[inline]
    pub fn upgrade<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Upgrade);
        }

        if let Ok(value) = value.try_into_value() {
            self.insert_header((header::UPGRADE, value));
        }

        self
    }

    /// Force close connection, even if it is marked as keep-alive
    #[inline]
    pub fn force_close(&mut self) -> &mut Self {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.set_connection_type(ConnectionType::Close);
        }
        self
    }

    /// Disable chunked transfer encoding for HTTP/1.1 streaming responses.
    #[inline]
    pub fn no_chunking(&mut self, len: u64) -> &mut Self {
        let mut buf = itoa::Buffer::new();
        self.insert_header((header::CONTENT_LENGTH, buf.format(len)));

        if let Some(parts) = parts(&mut self.head, &self.err) {
            parts.no_chunking(true);
        }
        self
    }

    /// Set response content type.
    #[inline]
    pub fn content_type<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = parts(&mut self.head, &self.err) {
            match value.try_into_value() {
                Ok(value) => {
                    parts.headers.insert(header::CONTENT_TYPE, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }
        self
    }

    /// Responses extensions
    #[inline]
    pub fn extensions(&self) -> Ref<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow()
    }

    /// Mutable reference to a the response's extensions
    #[inline]
    pub fn extensions_mut(&mut self) -> RefMut<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow_mut()
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn body<B: Into<Body>>(&mut self, body: B) -> Response<Body> {
        self.message_body(body.into())
    }

    /// Set a body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    pub fn message_body<B>(&mut self, body: B) -> Response<B> {
        if let Some(e) = self.err.take() {
            return Response::from(Error::from(e)).into_body();
        }

        let response = self.head.take().expect("cannot reuse response builder");

        Response {
            head: response,
            body: ResponseBody::Body(body),
            error: None,
        }
    }

    /// Set a streaming body and generate `Response`.
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn streaming<S, E>(&mut self, stream: S) -> Response<Body>
    where
        S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
        E: Into<Error> + 'static,
    {
        self.body(Body::from_message(BodyStream::new(stream)))
    }

    /// Set an empty body and generate `Response`
    ///
    /// `ResponseBuilder` can not be used after this call.
    #[inline]
    pub fn finish(&mut self) -> Response<Body> {
        self.body(Body::Empty)
    }

    /// This method construct new `ResponseBuilder`
    pub fn take(&mut self) -> ResponseBuilder {
        ResponseBuilder {
            head: self.head.take(),
            err: self.err.take(),
        }
    }
}

#[inline]
fn parts<'a>(
    parts: &'a mut Option<BoxedResponseHead>,
    err: &Option<HttpError>,
) -> Option<&'a mut ResponseHead> {
    if err.is_some() {
        return None;
    }
    parts.as_mut().map(|r| &mut **r)
}

impl Default for ResponseBuilder {
    fn default() -> Self {
        Self::new(StatusCode::OK)
    }
}

/// Convert `Response` to a `ResponseBuilder`. Body get dropped.
impl<B> From<Response<B>> for ResponseBuilder {
    fn from(res: Response<B>) -> ResponseBuilder {
        ResponseBuilder {
            head: Some(res.head),
            err: None,
        }
    }
}

/// Convert `ResponseHead` to a `ResponseBuilder`
impl<'a> From<&'a ResponseHead> for ResponseBuilder {
    fn from(head: &'a ResponseHead) -> ResponseBuilder {
        let mut msg = BoxedResponseHead::new(head.status);
        msg.version = head.version;
        msg.reason = head.reason;

        for (k, v) in head.headers.iter() {
            msg.headers.append(k.clone(), v.clone());
        }

        msg.no_chunking(!head.chunked());

        ResponseBuilder {
            head: Some(msg),
            err: None,
        }
    }
}

impl Future for ResponseBuilder {
    type Output = Result<Response<Body>, Error>;

    fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(Ok(self.finish()))
    }
}

impl fmt::Debug for ResponseBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let head = self.head.as_ref().unwrap();

        let res = writeln!(
            f,
            "\nResponseBuilder {:?} {}{}",
            head.version,
            head.status,
            head.reason.unwrap_or(""),
        );
        let _ = writeln!(f, "  headers:");
        for (key, val) in head.headers.iter() {
            let _ = writeln!(f, "    {:?}: {:?}", key, val);
        }
        res
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::body::Body;
    use crate::http::header::{HeaderName, HeaderValue, CONTENT_TYPE};

    #[test]
    fn test_basic_builder() {
        let resp = Response::build(StatusCode::OK)
            .insert_header(("X-TEST", "value"))
            .finish();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn test_upgrade() {
        let resp = Response::build(StatusCode::OK)
            .upgrade("websocket")
            .finish();
        assert!(resp.upgrade());
        assert_eq!(
            resp.headers().get(header::UPGRADE).unwrap(),
            HeaderValue::from_static("websocket")
        );
    }

    #[test]
    fn test_force_close() {
        let resp = Response::build(StatusCode::OK).force_close().finish();
        assert!(!resp.keep_alive())
    }

    #[test]
    fn test_content_type() {
        let resp = Response::build(StatusCode::OK)
            .content_type("text/plain")
            .body(Body::Empty);
        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "text/plain")
    }

    #[test]
    fn test_into_builder() {
        let mut resp: Response<Body> = "test".into();
        assert_eq!(resp.status(), StatusCode::OK);

        resp.headers_mut().insert(
            HeaderName::from_static("cookie"),
            HeaderValue::from_static("cookie1=val100"),
        );

        let mut builder: ResponseBuilder = resp.into();
        let resp = builder.status(StatusCode::BAD_REQUEST).finish();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let cookie = resp.headers().get_all("Cookie").next().unwrap();
        assert_eq!(cookie.to_str().unwrap(), "cookie1=val100");
    }

    #[test]
    fn response_builder_header_insert_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header(("Content-Type", "application/octet-stream"));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_insert_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_append_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header(("Content-Type", "application/octet-stream"));
        res.append_header(("Content-Type", "application/json"));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }

    #[test]
    fn response_builder_header_append_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }
}
<commit_msg>implement parts as assoc method<commit_after>//! HTTP response builder.

use std::{
    cell::{Ref, RefMut},
    fmt,
    future::Future,
    pin::Pin,
    str,
    task::{Context, Poll},
};

use bytes::Bytes;
use futures_core::Stream;

use crate::{
    body::{Body, BodyStream, ResponseBody},
    error::Error,
    extensions::Extensions,
    header::{IntoHeaderPair, IntoHeaderValue},
    http::{header, Error as HttpError, StatusCode},
    message::{BoxedResponseHead, ConnectionType, ResponseHead},
    Response,
};

/// An HTTP response builder.
///
/// Used to construct an instance of `Response` using a builder pattern. Response builders are often
/// created using [`Response::build`].
///
/// # Examples
/// ```
/// use actix_http::{Response, ResponseBuilder, body, http::StatusCode, http::header};
///
/// # actix_rt::System::new().block_on(async {
/// let mut res: Response<_> = Response::build(StatusCode::OK)
///     .content_type(mime::APPLICATION_JSON)
///     .insert_header((header::SERVER, "my-app/1.0"))
///     .append_header((header::SET_COOKIE, "a=1"))
///     .append_header((header::SET_COOKIE, "b=2"))
///     .body("1234");
///
/// assert_eq!(res.status(), StatusCode::OK);
/// assert_eq!(body::to_bytes(res.take_body()).await.unwrap(), &b"1234"[..]);
///
/// assert!(res.headers().contains_key("server"));
/// assert_eq!(res.headers().get_all("set-cookie").count(), 2);
/// # })
/// ```
pub struct ResponseBuilder {
    head: Option<BoxedResponseHead>,
    err: Option<HttpError>,
}

impl ResponseBuilder {
    /// Create response builder
    ///
    /// # Examples
    /// ```
    /// use actix_http::{Response, ResponseBuilder, http::StatusCode};
    ///
    /// let res: Response<_> = ResponseBuilder::default().finish();
    /// assert_eq!(res.status(), StatusCode::OK);
    /// ```
    #[inline]
    pub fn new(status: StatusCode) -> Self {
        ResponseBuilder {
            head: Some(BoxedResponseHead::new(status)),
            err: None,
        }
    }

    /// Set HTTP status code of this response.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::StatusCode};
    ///
    /// let res = ResponseBuilder::default().status(StatusCode::NOT_FOUND).finish();
    /// assert_eq!(res.status(), StatusCode::NOT_FOUND);
    /// ```
    #[inline]
    pub fn status(&mut self, status: StatusCode) -> &mut Self {
        if let Some(parts) = self.inner() {
            parts.status = status;
        }
        self
    }

    /// Insert a header, replacing any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .insert_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .insert_header(("X-TEST", "value"))
    ///     .finish();
    ///
    /// assert!(res.headers().contains_key("content-type"));
    /// assert!(res.headers().contains_key("x-test"));
    /// ```
    pub fn insert_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = self.inner() {
            match header.try_into_header_pair() {
                Ok((key, value)) => {
                    parts.headers.insert(key, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Append a header, keeping any that were set with an equivalent field name.
    ///
    /// # Examples
    /// ```
    /// use actix_http::{ResponseBuilder, http::header};
    ///
    /// let res = ResponseBuilder::default()
    ///     .append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON))
    ///     .append_header(("X-TEST", "value1"))
    ///     .append_header(("X-TEST", "value2"))
    ///     .finish();
    ///
    /// assert_eq!(res.headers().get_all("content-type").count(), 1);
    /// assert_eq!(res.headers().get_all("x-test").count(), 2);
    /// ```
    pub fn append_header<H>(&mut self, header: H) -> &mut Self
    where
        H: IntoHeaderPair,
    {
        if let Some(parts) = self.inner() {
            match header.try_into_header_pair() {
                Ok((key, value)) => parts.headers.append(key, value),
                Err(e) => self.err = Some(e.into()),
            };
        }

        self
    }

    /// Set the custom reason for the response.
    #[inline]
    pub fn reason(&mut self, reason: &'static str) -> &mut Self {
        if let Some(parts) = self.inner() {
            parts.reason = Some(reason);
        }
        self
    }

    /// Set connection type to KeepAlive
    #[inline]
    pub fn keep_alive(&mut self) -> &mut Self {
        if let Some(parts) = self.inner() {
            parts.set_connection_type(ConnectionType::KeepAlive);
        }
        self
    }

    /// Set connection type to Upgrade
    #[inline]
    pub fn upgrade<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = self.inner() {
            parts.set_connection_type(ConnectionType::Upgrade);
        }

        if let Ok(value) = value.try_into_value() {
            self.insert_header((header::UPGRADE, value));
        }

        self
    }

    /// Force close connection, even if it is marked as keep-alive
    #[inline]
    pub fn force_close(&mut self) -> &mut Self {
        if let Some(parts) = self.inner() {
            parts.set_connection_type(ConnectionType::Close);
        }
        self
    }

    /// Disable chunked transfer encoding for HTTP/1.1 streaming responses.
    #[inline]
    pub fn no_chunking(&mut self, len: u64) -> &mut Self {
        let mut buf = itoa::Buffer::new();
        self.insert_header((header::CONTENT_LENGTH, buf.format(len)));

        if let Some(parts) = self.inner() {
            parts.no_chunking(true);
        }
        self
    }

    /// Set response content type.
    #[inline]
    pub fn content_type<V>(&mut self, value: V) -> &mut Self
    where
        V: IntoHeaderValue,
    {
        if let Some(parts) = self.inner() {
            match value.try_into_value() {
                Ok(value) => {
                    parts.headers.insert(header::CONTENT_TYPE, value);
                }
                Err(e) => self.err = Some(e.into()),
            };
        }
        self
    }

    /// Responses extensions
    #[inline]
    pub fn extensions(&self) -> Ref<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow()
    }

    /// Mutable reference to a the response's extensions
    #[inline]
    pub fn extensions_mut(&mut self) -> RefMut<'_, Extensions> {
        let head = self.head.as_ref().expect("cannot reuse response builder");
        head.extensions.borrow_mut()
    }

    /// Generate response with a wrapped body.
    ///
    /// This `ResponseBuilder` will be left in a useless state.
    #[inline]
    pub fn body<B: Into<Body>>(&mut self, body: B) -> Response<Body> {
        self.message_body(body.into())
    }

    /// Generate response with a body.
    ///
    /// This `ResponseBuilder` will be left in a useless state.
    pub fn message_body<B>(&mut self, body: B) -> Response<B> {
        if let Some(e) = self.err.take() {
            return Response::from(Error::from(e)).into_body();
        }

        let response = self.head.take().expect("cannot reuse response builder");

        Response {
            head: response,
            body: ResponseBody::Body(body),
            error: None,
        }
    }

    /// Generate response with a streaming body.
    ///
    /// This `ResponseBuilder` will be left in a useless state.
    #[inline]
    pub fn streaming<S, E>(&mut self, stream: S) -> Response<Body>
    where
        S: Stream<Item = Result<Bytes, E>> + Unpin + 'static,
        E: Into<Error> + 'static,
    {
        self.body(Body::from_message(BodyStream::new(stream)))
    }

    /// Generate response with an empty body.
    ///
    /// This `ResponseBuilder` will be left in a useless state.
    #[inline]
    pub fn finish(&mut self) -> Response<Body> {
        self.body(Body::Empty)
    }

    /// Create an owned `ResponseBuilder`, leaving the original in a useless state.
    pub fn take(&mut self) -> ResponseBuilder {
        ResponseBuilder {
            head: self.head.take(),
            err: self.err.take(),
        }
    }

    /// Get access to the inner response head if there has been no error.
    fn inner(&mut self) -> Option<&mut ResponseHead> {
        if self.err.is_some() {
            return None;
        }

        self.head.as_mut().map(|r| &mut **r)
    }
}

impl Default for ResponseBuilder {
    fn default() -> Self {
        Self::new(StatusCode::OK)
    }
}

/// Convert `Response` to a `ResponseBuilder`. Body get dropped.
impl<B> From<Response<B>> for ResponseBuilder {
    fn from(res: Response<B>) -> ResponseBuilder {
        ResponseBuilder {
            head: Some(res.head),
            err: None,
        }
    }
}

/// Convert `ResponseHead` to a `ResponseBuilder`
impl<'a> From<&'a ResponseHead> for ResponseBuilder {
    fn from(head: &'a ResponseHead) -> ResponseBuilder {
        let mut msg = BoxedResponseHead::new(head.status);
        msg.version = head.version;
        msg.reason = head.reason;

        for (k, v) in head.headers.iter() {
            msg.headers.append(k.clone(), v.clone());
        }

        msg.no_chunking(!head.chunked());

        ResponseBuilder {
            head: Some(msg),
            err: None,
        }
    }
}

impl Future for ResponseBuilder {
    type Output = Result<Response<Body>, Error>;

    fn poll(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(Ok(self.finish()))
    }
}

impl fmt::Debug for ResponseBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let head = self.head.as_ref().unwrap();

        let res = writeln!(
            f,
            "\nResponseBuilder {:?} {}{}",
            head.version,
            head.status,
            head.reason.unwrap_or(""),
        );
        let _ = writeln!(f, "  headers:");
        for (key, val) in head.headers.iter() {
            let _ = writeln!(f, "    {:?}: {:?}", key, val);
        }
        res
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::body::Body;
    use crate::http::header::{HeaderName, HeaderValue, CONTENT_TYPE};

    #[test]
    fn test_basic_builder() {
        let resp = Response::build(StatusCode::OK)
            .insert_header(("X-TEST", "value"))
            .finish();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[test]
    fn test_upgrade() {
        let resp = Response::build(StatusCode::OK)
            .upgrade("websocket")
            .finish();
        assert!(resp.upgrade());
        assert_eq!(
            resp.headers().get(header::UPGRADE).unwrap(),
            HeaderValue::from_static("websocket")
        );
    }

    #[test]
    fn test_force_close() {
        let resp = Response::build(StatusCode::OK).force_close().finish();
        assert!(!resp.keep_alive())
    }

    #[test]
    fn test_content_type() {
        let resp = Response::build(StatusCode::OK)
            .content_type("text/plain")
            .body(Body::Empty);
        assert_eq!(resp.headers().get(CONTENT_TYPE).unwrap(), "text/plain")
    }

    #[test]
    fn test_into_builder() {
        let mut resp: Response<Body> = "test".into();
        assert_eq!(resp.status(), StatusCode::OK);

        resp.headers_mut().insert(
            HeaderName::from_static("cookie"),
            HeaderValue::from_static("cookie1=val100"),
        );

        let mut builder: ResponseBuilder = resp.into();
        let resp = builder.status(StatusCode::BAD_REQUEST).finish();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let cookie = resp.headers().get_all("Cookie").next().unwrap();
        assert_eq!(cookie.to_str().unwrap(), "cookie1=val100");
    }

    #[test]
    fn response_builder_header_insert_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header(("Content-Type", "application/octet-stream"));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_insert_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.insert_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        let res = res.finish();

        assert_eq!(
            res.headers().get("Content-Type"),
            Some(&HeaderValue::from_static("application/octet-stream"))
        );
    }

    #[test]
    fn response_builder_header_append_kv() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header(("Content-Type", "application/octet-stream"));
        res.append_header(("Content-Type", "application/json"));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }

    #[test]
    fn response_builder_header_append_typed() {
        let mut res = Response::build(StatusCode::OK);
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_OCTET_STREAM));
        res.append_header((header::CONTENT_TYPE, mime::APPLICATION_JSON));
        let res = res.finish();

        let headers: Vec<_> = res.headers().get_all("Content-Type").cloned().collect();
        assert_eq!(headers.len(), 2);
        assert!(headers.contains(&HeaderValue::from_static("application/octet-stream")));
        assert!(headers.contains(&HeaderValue::from_static("application/json")));
    }
}
<|endoftext|>


<commit_before>// rust-xmpp
// Copyright (c) 2014 Florian Zeitz
//
// This project is MIT licensed.
// Please see the COPYING file for more information.

use xml::Element;
use std::string::CowString;
use std::borrow::IntoCow;

pub trait XmppSend {
    fn xmpp_str<'a>(&'a self) -> CowString<'a>;
}

impl XmppSend for str {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        self.into_cow()
    }
}

impl XmppSend for String {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        (&self[]).into_cow()
    }
}

impl XmppSend for Element {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        (format!("{}", *self)).into_cow()
    }
}

impl<'a, S: ?Sized + XmppSend> XmppSend for &'a S {
    fn xmpp_str<'b>(&'b self) -> CowString<'b> {
        self.xmpp_str()
    }
}
<commit_msg>xmpp_send: Fix infinite recursion<commit_after>// rust-xmpp
// Copyright (c) 2014 Florian Zeitz
//
// This project is MIT licensed.
// Please see the COPYING file for more information.

use xml::Element;
use std::string::CowString;
use std::borrow::IntoCow;

pub trait XmppSend {
    fn xmpp_str<'a>(&'a self) -> CowString<'a>;
}

impl XmppSend for str {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        self.into_cow()
    }
}

impl XmppSend for String {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        (&self[]).into_cow()
    }
}

impl XmppSend for Element {
    fn xmpp_str<'a>(&'a self) -> CowString<'a> {
        (format!("{}", *self)).into_cow()
    }
}

impl<'a, S: ?Sized + XmppSend> XmppSend for &'a S {
    fn xmpp_str<'b>(&'b self) -> CowString<'b> {
        (*self).xmpp_str()
    }
}
<|endoftext|>


<commit_before>use super::GbaMemory;
use arm::{AccessType, Waitstates};

impl GbaMemory {
    pub(super) fn load32_sram(&mut self, address: u32, access: AccessType) -> (u32, Waitstates) {
        let (value, waitstates) = self.load8_sram(address, access);
        (value as u32 * 0x01010101, waitstates) // return repeated byte
    }

    pub(super) fn load16_sram(&mut self, address: u32, access: AccessType) -> (u16, Waitstates) {
        let (value, waitstates) = self.load8_sram(address, access);
        (value as u16 * 0x0101, waitstates) // return repeated byte
    }

    pub(super) fn load8_sram(&mut self, address: u32, _access: AccessType) -> (u8, Waitstates) {
        // FIXME implement real Flash IDs

        // These are some default values that Pokemon works with
        const MANUFACTURER: u8 = 0xC2;
        const DEVICE: u8 = 0x09;

        let value = match address {
            0xE000000 => MANUFACTURER,
            0xE000001 => DEVICE,
            _ => 0,
        };
        log::warn!("attempted to load from  unimplemented SRAM; address=0x{address:08X}; value=0x{value:02X}");
        (value, self.sram_waitstates)
    }

    pub(super) fn store32_sram(
        &mut self,
        address: u32,
        value: u32,
        _access: AccessType,
    ) -> Waitstates {
        log::warn!("attempted to store to unimplemented SRAM; address=0x{address:08X}; value=0x{value:08X}",);
        self.sram_waitstates
    }

    pub(super) fn store16_sram(
        &mut self,
        address: u32,
        value: u16,
        _access: AccessType,
    ) -> Waitstates {
        log::warn!("attempted to store to unimplemented SRAM; address=0x{address:08X}; value=0x{value:04X}",);
        self.sram_waitstates
    }

    pub(super) fn store8_sram(
        &mut self,
        address: u32,
        value: u8,
        _access: AccessType,
    ) -> Waitstates {
        log::warn!("attempted to store to unimplemented SRAM; address=0x{address:08X}; value=0x{value:02X}",);
        self.sram_waitstates
    }
}
<commit_msg>Make SRAM debug messages less noisy.<commit_after>use super::GbaMemory;
use arm::{AccessType, Waitstates};

impl GbaMemory {
    pub(super) fn load32_sram(&mut self, address: u32, access: AccessType) -> (u32, Waitstates) {
        let (value, waitstates) = self.load8_sram(address, access);
        (value as u32 * 0x01010101, waitstates) // return repeated byte
    }

    pub(super) fn load16_sram(&mut self, address: u32, access: AccessType) -> (u16, Waitstates) {
        let (value, waitstates) = self.load8_sram(address, access);
        (value as u16 * 0x0101, waitstates) // return repeated byte
    }

    pub(super) fn load8_sram(&mut self, address: u32, _access: AccessType) -> (u8, Waitstates) {
        // FIXME implement real Flash IDs

        // These are some default values that Pokemon works with
        const MANUFACTURER: u8 = 0xC2;
        const DEVICE: u8 = 0x09;

        let value = match address {
            0xE000000 => MANUFACTURER,
            0xE000001 => DEVICE,
            _ => 0,
        };
        unimplemented_sram_load();
        (value, self.sram_waitstates)
    }

    pub(super) fn store32_sram(
        &mut self,
        _address: u32,
        _value: u32,
        _access: AccessType,
    ) -> Waitstates {
        unimplemented_sram_store();
        self.sram_waitstates
    }

    pub(super) fn store16_sram(
        &mut self,
        _address: u32,
        _value: u16,
        _access: AccessType,
    ) -> Waitstates {
        unimplemented_sram_store();
        self.sram_waitstates
    }

    pub(super) fn store8_sram(
        &mut self,
        _address: u32,
        _value: u8,
        _access: AccessType,
    ) -> Waitstates {
        unimplemented_sram_store();
        self.sram_waitstates
    }
}

fn unimplemented_sram_load() {
    static mut CALLED: bool = false;
    if unsafe { !CALLED } {
        unsafe { CALLED = true };
        log::warn!("unimplemented SRAM load");
    }
}

fn unimplemented_sram_store() {
    static mut CALLED: bool = false;
    if unsafe { !CALLED } {
        unsafe { CALLED = true };
        log::warn!("unimplemented SRAM store");
    }
}
<|endoftext|>


<commit_before>use gl;
use gl::types::*;
use gl_context::GLContext;
use num;
use std::ffi::CString;
use std::marker::PhantomData;
use std::mem;
use std::ptr;

use shader::*;

/// Gets the id number for a given input of the shader program.
#[allow(non_snake_case)]
pub fn glGetAttribLocation(shader_program: GLuint, name: &str) -> GLint {
  let c_str = CString::new(name.as_bytes()).unwrap();
  let ptr = c_str.as_ptr() as *const i8;
  unsafe {
    gl::GetAttribLocation(shader_program, ptr)
  }
}

pub struct BufferHandle<'a> {
  pub gl_id: GLuint,
  phantom: PhantomData<&'a ()>,
}

impl<'a> BufferHandle<'a> {
  pub fn new<'b:'a>(_gl: &'a GLContext) -> BufferHandle<'b> {
    let mut gl_id = 0;

    unsafe {
      gl::GenBuffers(1, &mut gl_id);
    }

    assert!(gl_id != 0);

    BufferHandle {
      gl_id: gl_id,
      phantom: PhantomData,
    }
  }
}

impl<'a> Drop for BufferHandle<'a> {
  fn drop(&mut self) {
    unsafe {
      gl::DeleteBuffers(1, &self.gl_id);
    }
  }
}

/// Fixed-size VRAM buffer for individual bytes.
pub struct GLByteBuffer<'a> {
  pub handle: BufferHandle<'a>,
  /// number of bytes in the buffer.
  pub length: usize,
  /// maximum number of bytes in the buffer.
  pub capacity: usize,
}

impl<'a> GLByteBuffer<'a> {
  /// Creates a new array of objects on the GPU.
  /// capacity is provided in units of size slice_span.
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    capacity: usize,
  ) -> GLByteBuffer<'b> {
    let handle = BufferHandle::new(gl);

    unsafe {
      gl::BindBuffer(gl::ARRAY_BUFFER, handle.gl_id);

      gl::BufferData(
        gl::ARRAY_BUFFER,
        capacity as GLsizeiptr,
        ptr::null(),
        gl::DYNAMIC_DRAW,
      );
    }

    match gl.get_error() {
      gl::NO_ERROR => {},
      gl::OUT_OF_MEMORY => panic!("Out of VRAM"),
      err => warn!("OpenGL error 0x{:x}", err),
    }

    GLByteBuffer {
      handle: handle,
      length: 0,
      capacity: capacity,
    }
  }

  pub fn bind(&self, _: &mut GLContext) {
    unsafe {
      gl::BindBuffer(gl::ARRAY_BUFFER, self.handle.gl_id);
    }
  }

  /// Add more data into this buffer.
  /// Returns false and does nothing if this would exceed the capacity of the buffer.
  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub unsafe fn push(&mut self, gl: &mut GLContext, vs: *const u8, count: usize) -> bool {
    if self.length + count > self.capacity {
      // This would overflow the buffer.
      return false;
    }

    self.update_inner(gl, self.length, vs, count);
    self.length += count;

    true
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn swap_remove(&mut self, _gl: &mut GLContext, i: usize, count: usize) {
    assert!(count <= self.length);
    self.length -= count;
    assert!(i <= self.length);

    // In the `i == self.length` case, we don't bother with the swap;
    // decreasing `self.length` is enough.

    if i < self.length {
      assert!(
        i <= self.length - count,
        "GLByteBuffer::swap_remove would cause copy in overlapping regions"
      );

      unsafe {
        gl::CopyBufferSubData(
          gl::ARRAY_BUFFER,
          gl::ARRAY_BUFFER,
          self.length as GLintptr,
          i as GLintptr,
          count as GLsizeiptr,
        );
      }
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub unsafe fn update(&self, gl: &mut GLContext, idx: usize, vs: *const u8, count: usize) {
    assert!(idx + count <= self.length);
    self.update_inner(gl, idx, vs, count);
  }

  unsafe fn update_inner(
    &self,
    _gl: &mut GLContext,
    idx: usize,
    vs: *const u8,
    count: usize,
  ) {
    assert!(idx + count <= self.capacity);

    gl::BufferSubData(
      gl::ARRAY_BUFFER,
      idx as GLintptr,
      count as GLsizeiptr,
      mem::transmute(vs),
    );
  }
}

/// Fixed-size typed VRAM buffer, optimized for bulk inserts.
pub struct GLBuffer<'a, T> {
  pub byte_buffer: GLByteBuffer<'a>,
  phantom: PhantomData<Vec<T>>,
}

impl<'a, T> GLBuffer<'a, T> {
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    capacity: usize,
  ) -> GLBuffer<'b, T> {
    GLBuffer {
      byte_buffer: GLByteBuffer::new(gl, capacity * mem::size_of::<T>()),
      phantom: PhantomData,
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn push(&mut self, gl: &mut GLContext, vs: &[T]) -> bool {
    unsafe {
      self.byte_buffer.push(
        gl,
        mem::transmute(vs.as_ptr()),
        mem::size_of::<T>() * vs.len()
      )
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn update(&mut self, gl: &mut GLContext, idx: usize, vs: &[T]) {
    unsafe {
      self.byte_buffer.update(
        gl,
        mem::size_of::<T>() * idx,
        mem::transmute(vs.as_ptr()),
        mem::size_of::<T>() * vs.len(),
      );
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn swap_remove(&mut self, gl: &mut GLContext, idx: usize, count: usize) {
    self.byte_buffer.swap_remove(
      gl,
      mem::size_of::<T>() * idx,
      mem::size_of::<T>() * count,
    );
  }
}

#[derive(Debug)]
#[derive(Copy, Clone)]
pub enum DrawMode {
  Lines,
  Triangles,
  Points,
}

impl DrawMode {
  fn to_enum(&self) -> GLenum {
    match *self {
      DrawMode::Lines     => gl::LINES,
      DrawMode::Triangles => gl::TRIANGLES,
      DrawMode::Points    => gl::POINTS,
    }
  }
}

#[derive(Debug)]
#[derive(Copy, Clone)]
pub enum GLType {
  Float,
  UInt,
  Int,
}

impl GLType {
  pub fn size(&self) -> u32 {
    num::NumCast::from(
      match *self {
        GLType::Float => mem::size_of::<GLfloat>(),
        GLType::UInt  => mem::size_of::<GLuint>(),
        GLType::Int   => mem::size_of::<GLint>(),
      })
    .unwrap()
  }

  pub fn gl_enum(&self) -> GLenum {
    match *self {
      GLType::Float => gl::FLOAT,
      GLType::UInt  => gl::UNSIGNED_INT,
      GLType::Int   => gl::INT,
    }
  }

  pub fn is_integral(&self) -> bool {
    match *self {
      GLType::Float => false,
      GLType::UInt  => true,
      GLType::Int   => true,
    }
  }
}

#[derive(Debug)]
/// Specifies how to pass data from OpenGL to the vertex shaders.
pub struct VertexAttribData<'a> {
  /// Cooresponds to the shader's `input variable`.
  pub name: &'a str,
  /// The size of this attribute, in the provided units.
  pub size: u32,
  pub unit: GLType,
}

pub struct ArrayHandle<'a> {
  pub gl_id: GLuint,
  phantom: PhantomData<&'a ()>,
}

impl<'a> ArrayHandle<'a> {
  pub fn new<'b:'a>(_gl: &'a GLContext) -> ArrayHandle<'b> {
    let mut gl_id = 0;
    unsafe {
      gl::GenVertexArrays(1, &mut gl_id);
    }

    ArrayHandle {
      gl_id: gl_id,
      phantom: PhantomData,
    }
  }
}

impl<'a> Drop for ArrayHandle<'a> {
  fn drop(&mut self) {
    unsafe {
      gl::DeleteVertexArrays(1, &self.gl_id);
    }
  }
}

/// A fixed-capacity array of bytes passed to OpenGL.
pub struct GLArray<'a, T> {
  pub buffer: GLBuffer<'a, T>,
  pub handle: ArrayHandle<'a>,
  /// How to draw this buffer. Ex: gl::LINES, gl::TRIANGLES, etc.
  pub mode: GLenum,
  /// length in `T`s.
  pub length: usize,
}

impl<'a, T> GLArray<'a, T> {
  /// Creates a new array of objects on the GPU.
  /// capacity is provided in units of size slice_span.
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    shader_program: &Shader<'b>,
    attribs: &[VertexAttribData],
    mode: DrawMode,
    buffer: GLBuffer<'b, T>,
  ) -> GLArray<'b, T> {
    let handle = ArrayHandle::new(gl);

    unsafe {
      gl::BindVertexArray(handle.gl_id);
    }

    let mut offset = 0;
    let attrib_span = {
      let mut attrib_span = 0;
      for attrib in attribs.iter() {
        attrib_span += attrib.size * attrib.unit.size();
      }
      attrib_span
    };
    for attrib in attribs.iter() {
      let shader_attrib =
        glGetAttribLocation(
          shader_program.handle.gl_id,
          attrib.name
        );
      assert!(shader_attrib != -1, "shader attribute \"{}\" not found", attrib.name);
      let shader_attrib = shader_attrib as GLuint;

      unsafe {
        gl::EnableVertexAttribArray(shader_attrib);

        if attrib.unit.is_integral() {
          gl::VertexAttribIPointer(
            shader_attrib,
            attrib.size as i32,
            attrib.unit.gl_enum(),
            attrib_span as i32,
            ptr::null().offset(offset),
          );
        } else {
          gl::VertexAttribPointer(
            shader_attrib,
            attrib.size as i32,
            attrib.unit.gl_enum(),
            gl::FALSE as GLboolean,
            attrib_span as i32,
            ptr::null().offset(offset),
          );
        }
      }
      offset += (attrib.size * attrib.unit.size()) as isize;
    }

    match unsafe { gl::GetError() } {
      gl::NO_ERROR => {},
      err => warn!("OpenGL error 0x{:x}", err),
    }

    if attrib_span as usize != mem::size_of::<T>() {
      panic!("GLArray attribs don't describe the right number of bytes");
    }

    let length = buffer.byte_buffer.length / mem::size_of::<T>();

    GLArray {
      buffer: buffer,
      handle: handle,
      mode: mode.to_enum(),
      length: length,
    }
  }

  pub fn bind(&self, _: &mut GLContext) {
    unsafe {
      gl::BindVertexArray(self.handle.gl_id);
      gl::BindBuffer(gl::ARRAY_BUFFER, self.buffer.byte_buffer.handle.gl_id);
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  /// The array doesn't need to be bound.
  pub fn push(&mut self, gl: &mut GLContext, vs: &[T]) -> bool {
    let r = self.buffer.push(gl, vs);
    if r {
      self.length += vs.len();
    }
    r
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  /// The array doesn't need to be bound.
  pub fn swap_remove(&mut self, gl: &mut GLContext, idx: usize, count: usize) {
    self.buffer.swap_remove(gl, idx, count);
    self.length -= count;
  }

  /// Draws all the queued triangles to the screen.
  /// N.B. For performance reasons, this does NOT bind the array.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn draw(&self, gl: &mut GLContext) {
    self.draw_slice(gl, 0, self.length);
  }

  /// Draw some subset of the triangle array.
  /// N.B. For performance reasons, this does NOT bind the array.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn draw_slice(&self, _gl: &mut GLContext, start: usize, len: usize) {
    assert!(start + len <= self.length);

    unsafe {
      gl::DrawArrays(self.mode, start as i32, len as i32);
    }
  }
}
<commit_msg>expose VertexAttribData::apply<commit_after>use gl;
use gl::types::*;
use gl_context::GLContext;
use num;
use std::ffi::CString;
use std::marker::PhantomData;
use std::mem;
use std::ptr;

use shader::*;

/// Gets the id number for a given input of the shader program.
#[allow(non_snake_case)]
pub fn glGetAttribLocation(shader_program: GLuint, name: &str) -> GLint {
  let c_str = CString::new(name.as_bytes()).unwrap();
  let ptr = c_str.as_ptr() as *const i8;
  unsafe {
    gl::GetAttribLocation(shader_program, ptr)
  }
}

pub struct BufferHandle<'a> {
  pub gl_id: GLuint,
  phantom: PhantomData<&'a ()>,
}

impl<'a> BufferHandle<'a> {
  pub fn new<'b:'a>(_gl: &'a GLContext) -> BufferHandle<'b> {
    let mut gl_id = 0;

    unsafe {
      gl::GenBuffers(1, &mut gl_id);
    }

    assert!(gl_id != 0);

    BufferHandle {
      gl_id: gl_id,
      phantom: PhantomData,
    }
  }
}

impl<'a> Drop for BufferHandle<'a> {
  fn drop(&mut self) {
    unsafe {
      gl::DeleteBuffers(1, &self.gl_id);
    }
  }
}

/// Fixed-size VRAM buffer for individual bytes.
pub struct GLByteBuffer<'a> {
  pub handle: BufferHandle<'a>,
  /// number of bytes in the buffer.
  pub length: usize,
  /// maximum number of bytes in the buffer.
  pub capacity: usize,
}

impl<'a> GLByteBuffer<'a> {
  /// Creates a new array of objects on the GPU.
  /// capacity is provided in units of size slice_span.
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    capacity: usize,
  ) -> GLByteBuffer<'b> {
    let handle = BufferHandle::new(gl);

    unsafe {
      gl::BindBuffer(gl::ARRAY_BUFFER, handle.gl_id);

      gl::BufferData(
        gl::ARRAY_BUFFER,
        capacity as GLsizeiptr,
        ptr::null(),
        gl::DYNAMIC_DRAW,
      );
    }

    match gl.get_error() {
      gl::NO_ERROR => {},
      gl::OUT_OF_MEMORY => panic!("Out of VRAM"),
      err => warn!("OpenGL error 0x{:x}", err),
    }

    GLByteBuffer {
      handle: handle,
      length: 0,
      capacity: capacity,
    }
  }

  pub fn bind(&self, _: &mut GLContext) {
    unsafe {
      gl::BindBuffer(gl::ARRAY_BUFFER, self.handle.gl_id);
    }
  }

  /// Add more data into this buffer.
  /// Returns false and does nothing if this would exceed the capacity of the buffer.
  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub unsafe fn push(&mut self, gl: &mut GLContext, vs: *const u8, count: usize) -> bool {
    if self.length + count > self.capacity {
      // This would overflow the buffer.
      return false;
    }

    self.update_inner(gl, self.length, vs, count);
    self.length += count;

    true
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn swap_remove(&mut self, _gl: &mut GLContext, i: usize, count: usize) {
    assert!(count <= self.length);
    self.length -= count;
    assert!(i <= self.length);

    // In the `i == self.length` case, we don't bother with the swap;
    // decreasing `self.length` is enough.

    if i < self.length {
      assert!(
        i <= self.length - count,
        "GLByteBuffer::swap_remove would cause copy in overlapping regions"
      );

      unsafe {
        gl::CopyBufferSubData(
          gl::ARRAY_BUFFER,
          gl::ARRAY_BUFFER,
          self.length as GLintptr,
          i as GLintptr,
          count as GLsizeiptr,
        );
      }
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub unsafe fn update(&self, gl: &mut GLContext, idx: usize, vs: *const u8, count: usize) {
    assert!(idx + count <= self.length);
    self.update_inner(gl, idx, vs, count);
  }

  unsafe fn update_inner(
    &self,
    _gl: &mut GLContext,
    idx: usize,
    vs: *const u8,
    count: usize,
  ) {
    assert!(idx + count <= self.capacity);

    gl::BufferSubData(
      gl::ARRAY_BUFFER,
      idx as GLintptr,
      count as GLsizeiptr,
      mem::transmute(vs),
    );
  }
}

/// Fixed-size typed VRAM buffer, optimized for bulk inserts.
pub struct GLBuffer<'a, T> {
  pub byte_buffer: GLByteBuffer<'a>,
  phantom: PhantomData<Vec<T>>,
}

impl<'a, T> GLBuffer<'a, T> {
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    capacity: usize,
  ) -> GLBuffer<'b, T> {
    GLBuffer {
      byte_buffer: GLByteBuffer::new(gl, capacity * mem::size_of::<T>()),
      phantom: PhantomData,
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn push(&mut self, gl: &mut GLContext, vs: &[T]) -> bool {
    unsafe {
      self.byte_buffer.push(
        gl,
        mem::transmute(vs.as_ptr()),
        mem::size_of::<T>() * vs.len()
      )
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn update(&mut self, gl: &mut GLContext, idx: usize, vs: &[T]) {
    unsafe {
      self.byte_buffer.update(
        gl,
        mem::size_of::<T>() * idx,
        mem::transmute(vs.as_ptr()),
        mem::size_of::<T>() * vs.len(),
      );
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn swap_remove(&mut self, gl: &mut GLContext, idx: usize, count: usize) {
    self.byte_buffer.swap_remove(
      gl,
      mem::size_of::<T>() * idx,
      mem::size_of::<T>() * count,
    );
  }
}

#[derive(Debug)]
#[derive(Copy, Clone)]
pub enum DrawMode {
  Lines,
  Triangles,
  Points,
}

impl DrawMode {
  fn to_enum(&self) -> GLenum {
    match *self {
      DrawMode::Lines     => gl::LINES,
      DrawMode::Triangles => gl::TRIANGLES,
      DrawMode::Points    => gl::POINTS,
    }
  }
}

#[derive(Debug)]
#[derive(Copy, Clone)]
pub enum GLType {
  Float,
  UInt,
  Int,
}

impl GLType {
  pub fn size(&self) -> u32 {
    num::NumCast::from(
      match *self {
        GLType::Float => mem::size_of::<GLfloat>(),
        GLType::UInt  => mem::size_of::<GLuint>(),
        GLType::Int   => mem::size_of::<GLint>(),
      })
    .unwrap()
  }

  pub fn gl_enum(&self) -> GLenum {
    match *self {
      GLType::Float => gl::FLOAT,
      GLType::UInt  => gl::UNSIGNED_INT,
      GLType::Int   => gl::INT,
    }
  }

  pub fn is_integral(&self) -> bool {
    match *self {
      GLType::Float => false,
      GLType::UInt  => true,
      GLType::Int   => true,
    }
  }
}

#[derive(Debug)]
/// Specifies how to pass data from OpenGL to the vertex shaders.
pub struct VertexAttribData<'a> {
  /// Cooresponds to the shader's `input variable`.
  pub name: &'a str,
  /// The size of this attribute, in the provided units.
  pub size: u32,
  pub unit: GLType,
}

impl<'a> VertexAttribData<'a> {
  pub fn apply(attribs: &[VertexAttribData<'a>], _gl: &mut GLContext, shader: &Shader) -> u32 {
    let mut offset = 0;
    let attrib_span = {
      let mut attrib_span = 0;
      for attrib in attribs.iter() {
        attrib_span += attrib.size * attrib.unit.size();
      }
      attrib_span
    };

    for attrib in attribs {
      let shader_attrib =
        glGetAttribLocation(
          shader.handle.gl_id,
          attrib.name
        );
      assert!(shader_attrib != -1, "shader attribute \"{}\" not found", attrib.name);
      let shader_attrib = shader_attrib as GLuint;

      unsafe {
        gl::EnableVertexAttribArray(shader_attrib);

        if attrib.unit.is_integral() {
          gl::VertexAttribIPointer(
            shader_attrib,
            attrib.size as i32,
            attrib.unit.gl_enum(),
            attrib_span as i32,
            ptr::null().offset(offset),
          );
        } else {
          gl::VertexAttribPointer(
            shader_attrib,
            attrib.size as i32,
            attrib.unit.gl_enum(),
            gl::FALSE as GLboolean,
            attrib_span as i32,
            ptr::null().offset(offset),
          );
        }
      }
      offset += (attrib.size * attrib.unit.size()) as isize;
    }

    attrib_span
  }
}

pub struct ArrayHandle<'a> {
  pub gl_id: GLuint,
  phantom: PhantomData<&'a ()>,
}

impl<'a> ArrayHandle<'a> {
  pub fn new<'b:'a>(_gl: &'a GLContext) -> ArrayHandle<'b> {
    let mut gl_id = 0;
    unsafe {
      gl::GenVertexArrays(1, &mut gl_id);
    }

    ArrayHandle {
      gl_id: gl_id,
      phantom: PhantomData,
    }
  }
}

impl<'a> Drop for ArrayHandle<'a> {
  fn drop(&mut self) {
    unsafe {
      gl::DeleteVertexArrays(1, &self.gl_id);
    }
  }
}

/// A fixed-capacity array of bytes passed to OpenGL.
pub struct GLArray<'a, T> {
  pub buffer: GLBuffer<'a, T>,
  pub handle: ArrayHandle<'a>,
  /// How to draw this buffer. Ex: gl::LINES, gl::TRIANGLES, etc.
  pub mode: GLenum,
  /// length in `T`s.
  pub length: usize,
}

impl<'a, T> GLArray<'a, T> {
  /// Creates a new array of objects on the GPU.
  /// capacity is provided in units of size slice_span.
  pub fn new<'b:'a>(
    gl: &'a mut GLContext,
    shader_program: &Shader<'b>,
    attribs: &[VertexAttribData],
    mode: DrawMode,
    buffer: GLBuffer<'b, T>,
  ) -> GLArray<'b, T> {
    let handle = ArrayHandle::new(gl);

    unsafe {
      gl::BindVertexArray(handle.gl_id);
    }

    let attrib_span = VertexAttribData::apply(attribs, gl, shader_program);
    if attrib_span as usize != mem::size_of::<T>() {
      panic!("GLArray attribs don't describe the right number of bytes");
    }

    match unsafe { gl::GetError() } {
      gl::NO_ERROR => {},
      err => warn!("OpenGL error 0x{:x}", err),
    }

    let length = buffer.byte_buffer.length / mem::size_of::<T>();

    GLArray {
      buffer: buffer,
      handle: handle,
      mode: mode.to_enum(),
      length: length,
    }
  }

  pub fn bind(&self, _: &mut GLContext) {
    unsafe {
      gl::BindVertexArray(self.handle.gl_id);
      gl::BindBuffer(gl::ARRAY_BUFFER, self.buffer.byte_buffer.handle.gl_id);
    }
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  /// The array doesn't need to be bound.
  pub fn push(&mut self, gl: &mut GLContext, vs: &[T]) -> bool {
    let r = self.buffer.push(gl, vs);
    if r {
      self.length += vs.len();
    }
    r
  }

  /// N.B. For performance reasons, this does NOT bind the buffer.
  /// It will do the wrong thing if `bind` has not been correctly called.
  /// The array doesn't need to be bound.
  pub fn swap_remove(&mut self, gl: &mut GLContext, idx: usize, count: usize) {
    self.buffer.swap_remove(gl, idx, count);
    self.length -= count;
  }

  /// Draws all the queued triangles to the screen.
  /// N.B. For performance reasons, this does NOT bind the array.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn draw(&self, gl: &mut GLContext) {
    self.draw_slice(gl, 0, self.length);
  }

  /// Draw some subset of the triangle array.
  /// N.B. For performance reasons, this does NOT bind the array.
  /// It will do the wrong thing if `bind` has not been correctly called.
  pub fn draw_slice(&self, _gl: &mut GLContext, start: usize, len: usize) {
    assert!(start + len <= self.length);

    unsafe {
      gl::DrawArrays(self.mode, start as i32, len as i32);
    }
  }
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::Write;
use std::str::FromStr;

use clap::ArgMatches;
use etcd::Client;

use decryption::Decryptor;

pub struct Agent {
    decryptor: Decryptor,
    etcd: Client,
    role: Role,
}

#[derive(Clone, Copy)]
enum Role {
    Master,
    Node,
}

impl Agent {
    pub fn new<'a>(matches: &'a ArgMatches) -> Result<Self, String> {
        Ok(Agent {
            decryptor: try!(Decryptor::new(
                matches.value_of("region").expect("clap should have required region")
            )),
            etcd: Client::new(&["http://127.0.0.1:2379"]).expect("Failed to create etcd client"),
            role: matches.value_of("role").expect("clap should have required role").parse().expect(
              "clap should have required a valid value for role"
            )
        })
    }

    pub fn run(mut self) -> Result<Option<String>, String> {
        println!("kaws-agent is starting...");
        println!("Checking initial modifiedIndex of /kaws key...");
        let mut modified_index = match self.etcd.get("/kaws", false, false, false) {
            Ok(key_space_info) => match key_space_info.node.unwrap().modified_index {
                Some(new_index) => new_index,
                None => return Err("WARNING: etcd node for /kaws had no initial modified index".to_owned()),
            },
            Err(errors) => return Err(errors[0].to_string()),
        };

        if let Err(error) = self.refresh() {
            println!("ERROR: {}", error);
        }

        loop {
            println!("Watching /kaws key...");
            match self.etcd.watch("/kaws", Some(modified_index + 1), true) {
                Ok(key_space_info) => {
                    if let Err(error) = self.refresh() {
                        println!("WARNING: {}", error);
                    }

                    match key_space_info.node.unwrap().modified_index {
                        Some(new_index) => modified_index = new_index,
                        None => println!("WARNING: etcd node for watched key had no modified index"),
                    }
                }
                Err(errors) => println!("ERROR: {}", errors[0]),
            }
        }
    }

    // private

    fn refresh(&mut self) -> Result<Option<String>, String> {
        println!("Refreshing files from data in etcd...");
        try!(self.refresh_item("/kaws/pki/ca.pem", "/etc/kubernetes/ssl/ca.pem", false));

        match &self.role {
            &Role::Master => {
                try!(self.refresh_item(
                    "/kaws/pki/master.pem",
                    "/etc/kubernetes/ssl/master.pem",
                    false,
                ));
                try!(self.refresh_item(
                    "/kaws/pki/master-key-encrypted.base64",
                    "/etc/kubernetes/ssl/master-key.pem",
                    true,
                ));
            }
            &Role::Node => {
                try!(self.refresh_item(
                    "/kaws/pki/node.pem",
                    "/etc/kubernetes/ssl/node.pem",
                    false,
                ));
                try!(self.refresh_item(
                    "/kaws/pki/node-key-encrypted.base64",
                    "/etc/kubernetes/ssl/node-key.pem",
                    true,
                ));
            }
        }

        Ok(None)
    }

    fn refresh_item(&mut self, key: &str, file_path: &str, encrypted: bool)
    -> Result<Option<String>, String> {
        println!("Refreshing file {} from etcd key {}", file_path, key);
        let key_space_info = match self.etcd.get(key, false, false, false) {
            Ok(key_space_info) => key_space_info,
            Err(errors) => return Err(errors[0].to_string()),
        };

        let value = match key_space_info.node.unwrap().value {
            Some(value) => if encrypted {
                match self.decryptor.decrypt(&value) {
                    Ok(plaintext) => plaintext,
                    Err(error) => return Err(error),
                }
            } else {
                value
            },
            None => return Err(format!("etcd key {} had no value", key)),
        };

        let mut file = match File::create(file_path) {
            Ok(file) => file,
            Err(error) => return Err(error.to_string()),
        };

        match file.write_all(value.as_bytes()) {
            Ok(_) => Ok(None),
            Err(error) => Err(error.to_string()),
        }
    }
}

impl FromStr for Role {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "master" => Ok(Role::Master),
            "node" => Ok(Role::Node),
            role => Err(format!("Unknown server role for agent: {}", role)),
        }
    }
}
<commit_msg>Handle the case where 1000 indexes may be missed.<commit_after>use std::fs::File;
use std::io::Write;
use std::str::FromStr;

use clap::ArgMatches;
use etcd::{Client, Error};

use decryption::Decryptor;

pub struct Agent {
    decryptor: Decryptor,
    etcd: Client,
    role: Role,
}

#[derive(Clone, Copy)]
enum Role {
    Master,
    Node,
}

impl Agent {
    pub fn new<'a>(matches: &'a ArgMatches) -> Result<Self, String> {
        Ok(Agent {
            decryptor: try!(Decryptor::new(
                matches.value_of("region").expect("clap should have required region")
            )),
            etcd: Client::new(&["http://127.0.0.1:2379"]).expect("Failed to create etcd client"),
            role: matches.value_of("role").expect("clap should have required role").parse().expect(
              "clap should have required a valid value for role"
            )
        })
    }

    pub fn run(mut self) -> Result<Option<String>, String> {
        println!("kaws-agent is starting...");
        let mut watch_index = 0;

        if let Err(error) = self.refresh() {
            println!("WARNING: {}", error);
        }

        loop {
            println!("Watching /kaws key from index {}...", watch_index + 1);
            match self.etcd.watch("/kaws", Some(watch_index + 1), true) {
                Ok(key_space_info) => {
                    if let Err(error) = self.refresh() {
                        println!("WARNING: {}", error);
                    }

                    match key_space_info.node.unwrap().modified_index {
                        Some(new_index) => watch_index = new_index,
                        None => println!("WARNING: etcd node for watched key had no modified index"),
                    }
                }
                Err(errors) => {
                    println!("WARNING: {}", errors[0]);

                    if let Error::Api(ref api_error) = errors[0] {
                        watch_index = api_error.index;
                    };
                }
            }
        }
    }

    // private

    fn refresh(&mut self) -> Result<Option<String>, String> {
        println!("Refreshing files from data in etcd...");
        try!(self.refresh_item("/kaws/pki/ca.pem", "/etc/kubernetes/ssl/ca.pem", false));

        match &self.role {
            &Role::Master => {
                try!(self.refresh_item(
                    "/kaws/pki/master.pem",
                    "/etc/kubernetes/ssl/master.pem",
                    false,
                ));
                try!(self.refresh_item(
                    "/kaws/pki/master-key-encrypted.base64",
                    "/etc/kubernetes/ssl/master-key.pem",
                    true,
                ));
            }
            &Role::Node => {
                try!(self.refresh_item(
                    "/kaws/pki/node.pem",
                    "/etc/kubernetes/ssl/node.pem",
                    false,
                ));
                try!(self.refresh_item(
                    "/kaws/pki/node-key-encrypted.base64",
                    "/etc/kubernetes/ssl/node-key.pem",
                    true,
                ));
            }
        }

        Ok(None)
    }

    fn refresh_item(&mut self, key: &str, file_path: &str, encrypted: bool)
    -> Result<Option<String>, String> {
        println!("Refreshing file {} from etcd key {}", file_path, key);
        let key_space_info = match self.etcd.get(key, false, false, false) {
            Ok(key_space_info) => key_space_info,
            Err(errors) => return Err(errors[0].to_string()),
        };

        let value = match key_space_info.node.unwrap().value {
            Some(value) => if encrypted {
                match self.decryptor.decrypt(&value) {
                    Ok(plaintext) => plaintext,
                    Err(error) => return Err(error),
                }
            } else {
                value
            },
            None => return Err(format!("etcd key {} had no value", key)),
        };

        let mut file = match File::create(file_path) {
            Ok(file) => file,
            Err(error) => return Err(error.to_string()),
        };

        match file.write_all(value.as_bytes()) {
            Ok(_) => Ok(None),
            Err(error) => Err(error.to_string()),
        }
    }
}

impl FromStr for Role {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "master" => Ok(Role::Master),
            "node" => Ok(Role::Node),
            role => Err(format!("Unknown server role for agent: {}", role)),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return unsafe { &*tls };
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state.queue_head().is_null() && spinwait.spin() {
                state = self.state.load(Ordering::Relaxed);
                continue;
            }

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            unsafe { thread_data.parker.prepare_park(); }

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            unsafe { thread_data.parker.park(); }

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    unsafe fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = (*current).queue_tail.get();
                if !queue_tail.is_null() {
                    break;
                }
                let next = (*current).next.get();
                (*next).prev.set(current);
                current = next;
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            (*queue_head).queue_tail.set(queue_tail);

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = (*queue_tail).prev.get();
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                (*queue_head).queue_tail.set(new_tail);
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            (*queue_tail).parker.unpark_lock().unpark();
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<commit_msg>Mark WordLock::unlock_slow as not unsafe<commit_after>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use spinwait::SpinWait;
use std::cell::Cell;
use std::mem;
#[cfg(not(has_localkey_try_with))]
use std::panic;
use std::ptr;
use std::sync::atomic::{fence, AtomicUsize, Ordering};
use std::thread::LocalKey;
use thread_parker::ThreadParker;

struct ThreadData {
    parker: ThreadParker,

    // Linked list of threads in the queue. The queue is split into two parts:
    // the processed part and the unprocessed part. When new nodes are added to
    // the list, they only have the next pointer set, and queue_tail is null.
    //
    // Nodes are processed with the queue lock held, which consists of setting
    // the prev pointer for each node and setting the queue_tail pointer on the
    // first processed node of the list.
    //
    // This setup allows nodes to be added to the queue without a lock, while
    // still allowing O(1) removal of nodes from the processed part of the list.
    // The only cost is the O(n) processing, but this only needs to be done
    // once for each node, and therefore isn't too expensive.
    queue_tail: Cell<*const ThreadData>,
    prev: Cell<*const ThreadData>,
    next: Cell<*const ThreadData>,
}

impl ThreadData {
    fn new() -> ThreadData {
        ThreadData {
            parker: ThreadParker::new(),
            queue_tail: Cell::new(ptr::null()),
            prev: Cell::new(ptr::null()),
            next: Cell::new(ptr::null()),
        }
    }
}

// Returns a ThreadData structure for the current thread
fn get_thread_data(local: &mut Option<ThreadData>) -> &ThreadData {
    // Try to read from thread-local storage, but return None if the TLS has
    // already been destroyed.
    #[cfg(has_localkey_try_with)]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        key.try_with(|x| x as *const ThreadData).ok()
    }
    #[cfg(not(has_localkey_try_with))]
    fn try_get_tls(key: &'static LocalKey<ThreadData>) -> Option<*const ThreadData> {
        panic::catch_unwind(|| key.with(|x| x as *const ThreadData)).ok()
    }

    // If ThreadData is expensive to construct, then we want to use a cached
    // version in thread-local storage if possible.
    if !cfg!(windows) && !cfg!(all(feature = "nightly", target_os = "linux")) {
        thread_local!(static THREAD_DATA: ThreadData = ThreadData::new());
        if let Some(tls) = try_get_tls(&THREAD_DATA) {
            return unsafe { &*tls };
        }
    }

    // Otherwise just create a ThreadData on the stack
    *local = Some(ThreadData::new());
    local.as_ref().unwrap()
}

const LOCKED_BIT: usize = 1;
const QUEUE_LOCKED_BIT: usize = 2;
const QUEUE_MASK: usize = !3;

// Word-sized lock that is used to implement the parking_lot API. Since this
// can't use parking_lot, it instead manages its own queue of waiting threads.
pub struct WordLock {
    state: AtomicUsize,
}

impl WordLock {
    #[inline]
    pub fn new() -> WordLock {
        WordLock {
            state: AtomicUsize::new(0),
        }
    }

    #[inline]
    pub fn lock(&self) {
        if self
            .state
            .compare_exchange_weak(0, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
        {
            return;
        }
        self.lock_slow();
    }

    #[inline]
    pub unsafe fn unlock(&self) {
        let state = self.state.fetch_sub(LOCKED_BIT, Ordering::Release);
        if state.is_queue_locked() || state.queue_head().is_null() {
            return;
        }
        self.unlock_slow();
    }

    #[cold]
    #[inline(never)]
    fn lock_slow(&self) {
        let mut spinwait = SpinWait::new();
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // Grab the lock if it isn't locked, even if there is a queue on it
            if !state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state | LOCKED_BIT,
                    Ordering::Acquire,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }
                continue;
            }

            // If there is no queue, try spinning a few times
            if state.queue_head().is_null() && spinwait.spin() {
                state = self.state.load(Ordering::Relaxed);
                continue;
            }

            // Get our thread data and prepare it for parking
            let mut thread_data = None;
            let thread_data = get_thread_data(&mut thread_data);
            assert!(mem::align_of_val(thread_data) > !QUEUE_MASK);
            unsafe { thread_data.parker.prepare_park(); }

            // Add our thread to the front of the queue
            let queue_head = state.queue_head();
            if queue_head.is_null() {
                thread_data.queue_tail.set(thread_data);
                thread_data.prev.set(ptr::null());
            } else {
                thread_data.queue_tail.set(ptr::null());
                thread_data.prev.set(ptr::null());
                thread_data.next.set(queue_head);
            }
            if let Err(x) = self.state.compare_exchange_weak(
                state,
                state.with_queue_head(thread_data),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                state = x;
                continue;
            }

            // Sleep until we are woken up by an unlock
            unsafe { thread_data.parker.park(); }

            // Loop back and try locking again
            spinwait.reset();
            self.state.load(Ordering::Relaxed);
        }
    }

    #[cold]
    #[inline(never)]
    fn unlock_slow(&self) {
        let mut state = self.state.load(Ordering::Relaxed);
        loop {
            // We just unlocked the WordLock. Just check if there is a thread
            // to wake up. If the queue is locked then another thread is already
            // taking care of waking up a thread.
            if state.is_queue_locked() || state.queue_head().is_null() {
                return;
            }

            // Try to grab the queue lock
            match self.state.compare_exchange_weak(
                state,
                state | QUEUE_LOCKED_BIT,
                Ordering::Acquire,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => state = x,
            }
        }

        // Now we have the queue lock and the queue is non-empty
        'outer: loop {
            // First, we need to fill in the prev pointers for any newly added
            // threads. We do this until we reach a node that we previously
            // processed, which has a non-null queue_tail pointer.
            let queue_head = state.queue_head();
            let mut queue_tail;
            let mut current = queue_head;
            loop {
                queue_tail = unsafe { (*current).queue_tail.get() };
                if !queue_tail.is_null() {
                    break;
                }
                unsafe {
                    let next = (*current).next.get();
                    (*next).prev.set(current);
                    current = next;
                }
            }

            // Set queue_tail on the queue head to indicate that the whole list
            // has prev pointers set correctly.
            unsafe {
                (*queue_head).queue_tail.set(queue_tail);
            }

            // If the WordLock is locked, then there is no point waking up a
            // thread now. Instead we let the next unlocker take care of waking
            // up a thread.
            if state.is_locked() {
                match self.state.compare_exchange_weak(
                    state,
                    state & !QUEUE_LOCKED_BIT,
                    Ordering::Release,
                    Ordering::Relaxed,
                ) {
                    Ok(_) => return,
                    Err(x) => state = x,
                }

                // Need an acquire fence before reading the new queue
                fence(Ordering::Acquire);
                continue;
            }

            // Remove the last thread from the queue and unlock the queue
            let new_tail = unsafe { (*queue_tail).prev.get() };
            if new_tail.is_null() {
                loop {
                    match self.state.compare_exchange_weak(
                        state,
                        state & LOCKED_BIT,
                        Ordering::Release,
                        Ordering::Relaxed,
                    ) {
                        Ok(_) => break,
                        Err(x) => state = x,
                    }

                    // If the compare_exchange failed because a new thread was
                    // added to the queue then we need to re-scan the queue to
                    // find the previous element.
                    if state.queue_head().is_null() {
                        continue;
                    } else {
                        // Need an acquire fence before reading the new queue
                        fence(Ordering::Acquire);
                        continue 'outer;
                    }
                }
            } else {
                unsafe {
                    (*queue_head).queue_tail.set(new_tail);
                }
                self.state.fetch_and(!QUEUE_LOCKED_BIT, Ordering::Release);
            }

            // Finally, wake up the thread we removed from the queue. Note that
            // we don't need to worry about any races here since the thread is
            // guaranteed to be sleeping right now and we are the only one who
            // can wake it up.
            unsafe {
                (*queue_tail).parker.unpark_lock().unpark();
            }
            break;
        }
    }
}

trait LockState {
    fn is_locked(self) -> bool;
    fn is_queue_locked(self) -> bool;
    fn queue_head(self) -> *const ThreadData;
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self;
}

impl LockState for usize {
    #[inline]
    fn is_locked(self) -> bool {
        self & LOCKED_BIT != 0
    }

    #[inline]
    fn is_queue_locked(self) -> bool {
        self & QUEUE_LOCKED_BIT != 0
    }

    #[inline]
    fn queue_head(self) -> *const ThreadData {
        (self & QUEUE_MASK) as *const ThreadData
    }

    #[inline]
    fn with_queue_head(self, thread_data: *const ThreadData) -> Self {
        (self & !QUEUE_MASK) | thread_data as *const _ as usize
    }
}
<|endoftext|>