<commit_before>//! Module for User-Interactive Authentication API types.

use std::{
    collections::BTreeMap,
    fmt::{self, Display, Formatter},
};

use ruma_api::{error::ResponseDeserializationError, EndpointError};
use ruma_serde::Outgoing;
use serde::{Deserialize, Serialize};
use serde_json::{
    from_slice as from_json_slice, to_vec as to_json_vec, value::RawValue as RawJsonValue,
    Value as JsonValue,
};

use crate::error::{Error as MatrixError, ErrorBody};

/// Additional authentication information for the user-interactive authentication API.
#[derive(Clone, Debug, Outgoing, Serialize)]
#[serde(untagged)]
pub enum AuthData<'a> {
    /// Used for sending UIAA authentication requests to the homeserver directly
    /// from the client.
    DirectRequest {
        /// The login type that the client is attempting to complete.
        #[serde(rename = "type")]
        kind: &'a str,

        /// The value of the session key given by the homeserver.
        #[serde(skip_serializing_if = "Option::is_none")]
        session: Option<&'a str>,

        /// Parameters submitted for a particular authentication stage.
        // FIXME: RawJsonValue doesn't work here, is that a bug?
        #[serde(flatten)]
        auth_parameters: BTreeMap<String, JsonValue>,
    },

    /// Used by the client to acknowledge that the user has completed a UIAA
    /// stage through the fallback method.
    FallbackAcknowledgement {
        /// The value of the session key given by the homeserver.
        session: &'a str,
    },
}

/// Information about available authentication flows and status for
/// User-Interactive Authenticiation API.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct UiaaInfo {
    /// List of authentication flows available for this endpoint.
    pub flows: Vec<AuthFlow>,

    /// List of stages in the current flow completed by the client.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub completed: Vec<String>,

    /// Authentication parameters required for the client to complete
    /// authentication.
    ///
    /// To create a `Box<RawJsonValue>`, use `serde_json::value::to_raw_value`.
    pub params: Box<RawJsonValue>,

    /// Session key for client to use to complete authentication.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session: Option<String>,

    /// Authentication-related errors for previous request returned by homeserver.
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub auth_error: Option<ErrorBody>,
}

/// Description of steps required to authenticate via the User-Interactive
/// Authentication API.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[cfg_attr(test, derive(PartialEq))]
pub struct AuthFlow {
    /// Ordered list of stages required to complete authentication.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub stages: Vec<String>,
}

/// Contains either a User-Interactive Authentication API response body or a
/// Matrix error.
#[derive(Clone, Debug)]
pub enum UiaaResponse {
    /// User-Interactive Authentication API response
    AuthResponse(UiaaInfo),

    /// Matrix error response
    MatrixError(MatrixError),
}

impl Display for UiaaResponse {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::AuthResponse(_) => write!(f, "User-Interactive Authentication required."),
            Self::MatrixError(err) => write!(f, "{}", err),
        }
    }
}

impl From<MatrixError> for UiaaResponse {
    fn from(error: MatrixError) -> Self {
        Self::MatrixError(error)
    }
}

impl EndpointError for UiaaResponse {
    fn try_from_response(
        response: http::Response<Vec<u8>>,
    ) -> Result<Self, ResponseDeserializationError> {
        if response.status() == http::StatusCode::UNAUTHORIZED {
            if let Ok(authentication_info) = from_json_slice::<UiaaInfo>(response.body()) {
                return Ok(UiaaResponse::AuthResponse(authentication_info));
            }
        }

        MatrixError::try_from_response(response).map(From::from)
    }
}

impl std::error::Error for UiaaResponse {}

impl From<UiaaResponse> for http::Response<Vec<u8>> {
    fn from(uiaa_response: UiaaResponse) -> http::Response<Vec<u8>> {
        match uiaa_response {
            UiaaResponse::AuthResponse(authentication_info) => http::Response::builder()
                .header(http::header::CONTENT_TYPE, "application/json")
                .status(&http::StatusCode::UNAUTHORIZED)
                .body(to_json_vec(&authentication_info).unwrap())
                .unwrap(),
            UiaaResponse::MatrixError(error) => http::Response::from(error),
        }
    }
}

#[cfg(test)]
mod tests {
    use maplit::btreemap;
    use matches::assert_matches;
    use ruma_api::EndpointError;
    use serde_json::{
        from_slice as from_json_slice, from_str as from_json_str, from_value as from_json_value,
        json, to_value as to_json_value, value::to_raw_value as to_raw_json_value,
        Value as JsonValue,
    };

    use super::{AuthData, AuthFlow, IncomingAuthData, UiaaInfo, UiaaResponse};
    use crate::error::{ErrorBody, ErrorKind};

    #[test]
    fn serialize_authentication_data_direct_request() {
        let authentication_data = AuthData::DirectRequest {
            kind: "example.type.foo",
            session: Some("ZXY000"),
            auth_parameters: btreemap! {
                "example_credential".to_owned() => json!("verypoorsharedsecret")
            },
        };

        assert_eq!(
            json!({
                "type": "example.type.foo",
                "session": "ZXY000",
                "example_credential": "verypoorsharedsecret",
            }),
            to_json_value(authentication_data).unwrap()
        );
    }

    #[test]
    fn deserialize_authentication_data_direct_request() {
        let json = json!({
            "type": "example.type.foo",
            "session": "opaque_session_id",
            "example_credential": "verypoorsharedsecret",
        });

        assert_matches!(
            from_json_value(json).unwrap(),
            IncomingAuthData::DirectRequest { kind, session: Some(session), auth_parameters }
            if kind == "example.type.foo"
                && session == "opaque_session_id"
                && auth_parameters == btreemap!{
                    "example_credential".to_owned() => json!("verypoorsharedsecret")
                }
        );
    }

    #[test]
    fn serialize_authentication_data_fallback() {
        let authentication_data = AuthData::FallbackAcknowledgement { session: "ZXY000" };

        assert_eq!(json!({ "session": "ZXY000" }), to_json_value(authentication_data).unwrap());
    }

    #[test]
    fn deserialize_authentication_data_fallback() {
        let json = json!({ "session": "opaque_session_id" });

        assert_matches!(
            from_json_value(json).unwrap(),
            IncomingAuthData::FallbackAcknowledgement { session }
            if session == "opaque_session_id"
        );
    }

    #[test]
    fn serialize_uiaa_info() {
        let uiaa_info = UiaaInfo {
            flows: vec![AuthFlow {
                stages: vec!["m.login.password".into(), "m.login.dummy".into()],
            }],
            completed: vec!["m.login.password".into()],
            params: to_raw_json_value(&json!({
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }))
            .unwrap(),
            session: None,
            auth_error: None,
        };

        let json = json!({
            "flows": [{ "stages": ["m.login.password", "m.login.dummy"] }],
            "completed": ["m.login.password"],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }
        });
        assert_eq!(to_json_value(uiaa_info).unwrap(), json);
    }

    #[test]
    fn deserialize_uiaa_info() {
        let json = json!({
            "errcode": "M_FORBIDDEN",
            "error": "Invalid password",
            "completed": ["example.type.foo"],
            "flows": [
                {
                    "stages": ["example.type.foo", "example.type.bar"]
                },
                {
                    "stages": ["example.type.foo", "example.type.baz"]
                }
            ],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            },
            "session": "xxxxxx"
        });

        assert_matches!(
            from_json_value::<UiaaInfo>(json).unwrap(),
            UiaaInfo {
                auth_error: Some(ErrorBody {
                    kind: ErrorKind::Forbidden,
                    message: error_message,
                }),
                completed,
                flows,
                params,
                session: Some(session),
            } if error_message == "Invalid password"
                && completed == vec!["example.type.foo".to_owned()]
                && flows == vec![
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.bar".into(),
                        ],
                    },
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.baz".into(),
                        ],
                    },
                ]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
                && session == "xxxxxx"
        );
    }

    #[test]
    fn try_uiaa_response_into_http_response() {
        let uiaa_info = UiaaInfo {
            flows: vec![AuthFlow {
                stages: vec!["m.login.password".into(), "m.login.dummy".into()],
            }],
            completed: vec!["m.login.password".into()],
            params: to_raw_json_value(&json!({
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }))
            .unwrap(),
            session: None,
            auth_error: None,
        };
        let uiaa_response: http::Response<Vec<u8>> = UiaaResponse::AuthResponse(uiaa_info).into();

        assert_matches!(
            from_json_slice::<UiaaInfo>(uiaa_response.body()).unwrap(),
            UiaaInfo {
                flows,
                completed,
                params,
                session: None,
                auth_error: None,
            } if flows == vec![AuthFlow {
                    stages: vec!["m.login.password".into(), "m.login.dummy".into()],
                }]
                && completed == vec!["m.login.password".to_owned()]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
        );
        assert_eq!(uiaa_response.status(), http::status::StatusCode::UNAUTHORIZED);
    }

    #[test]
    fn try_uiaa_response_from_http_response() {
        let json = serde_json::to_string(&json!({
            "errcode": "M_FORBIDDEN",
            "error": "Invalid password",
            "completed": [ "example.type.foo" ],
            "flows": [
                {
                    "stages": [ "example.type.foo", "example.type.bar" ]
                },
                {
                    "stages": [ "example.type.foo", "example.type.baz" ]
                }
            ],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            },
            "session": "xxxxxx"
        }))
        .unwrap();

        let http_response = http::Response::builder()
            .status(http::StatusCode::UNAUTHORIZED)
            .body(json.into())
            .unwrap();

        let parsed_uiaa_info = match UiaaResponse::try_from_response(http_response).unwrap() {
            UiaaResponse::AuthResponse(uiaa_info) => uiaa_info,
            _ => panic!("Expected UiaaResponse::AuthResponse"),
        };

        assert_matches!(
            parsed_uiaa_info,
            UiaaInfo {
                auth_error: Some(ErrorBody {
                    kind: ErrorKind::Forbidden,
                    message: error_message,
                }),
                completed,
                flows,
                params,
                session: Some(session),
            } if error_message == "Invalid password"
                && completed == vec!["example.type.foo".to_owned()]
                && flows == vec![
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.bar".into(),
                        ],
                    },
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.baz".into(),
                        ],
                    },
                ]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
                && session == "xxxxxx"
        );
    }
}
<commit_msg>client-api: Rewrap comments in uiaa<commit_after>//! Module for User-Interactive Authentication API types.

use std::{
    collections::BTreeMap,
    fmt::{self, Display, Formatter},
};

use ruma_api::{error::ResponseDeserializationError, EndpointError};
use ruma_serde::Outgoing;
use serde::{Deserialize, Serialize};
use serde_json::{
    from_slice as from_json_slice, to_vec as to_json_vec, value::RawValue as RawJsonValue,
    Value as JsonValue,
};

use crate::error::{Error as MatrixError, ErrorBody};

/// Additional authentication information for the user-interactive authentication API.
#[derive(Clone, Debug, Outgoing, Serialize)]
#[serde(untagged)]
pub enum AuthData<'a> {
    /// Used for sending UIAA authentication requests to the homeserver directly from the client.
    DirectRequest {
        /// The login type that the client is attempting to complete.
        #[serde(rename = "type")]
        kind: &'a str,

        /// The value of the session key given by the homeserver.
        #[serde(skip_serializing_if = "Option::is_none")]
        session: Option<&'a str>,

        /// Parameters submitted for a particular authentication stage.
        // FIXME: RawJsonValue doesn't work here, is that a bug?
        #[serde(flatten)]
        auth_parameters: BTreeMap<String, JsonValue>,
    },

    /// Used by the client to acknowledge that the user has completed a UIAA stage through the
    /// fallback method.
    FallbackAcknowledgement {
        /// The value of the session key given by the homeserver.
        session: &'a str,
    },
}

/// Information about available authentication flows and status for User-Interactive Authenticiation
/// API.
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct UiaaInfo {
    /// List of authentication flows available for this endpoint.
    pub flows: Vec<AuthFlow>,

    /// List of stages in the current flow completed by the client.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub completed: Vec<String>,

    /// Authentication parameters required for the client to complete authentication.
    ///
    /// To create a `Box<RawJsonValue>`, use `serde_json::value::to_raw_value`.
    pub params: Box<RawJsonValue>,

    /// Session key for client to use to complete authentication.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session: Option<String>,

    /// Authentication-related errors for previous request returned by homeserver.
    #[serde(flatten, skip_serializing_if = "Option::is_none")]
    pub auth_error: Option<ErrorBody>,
}

/// Description of steps required to authenticate via the User-Interactive Authentication API.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[cfg_attr(test, derive(PartialEq))]
pub struct AuthFlow {
    /// Ordered list of stages required to complete authentication.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub stages: Vec<String>,
}

/// Contains either a User-Interactive Authentication API response body or a Matrix error.
#[derive(Clone, Debug)]
pub enum UiaaResponse {
    /// User-Interactive Authentication API response
    AuthResponse(UiaaInfo),

    /// Matrix error response
    MatrixError(MatrixError),
}

impl Display for UiaaResponse {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            Self::AuthResponse(_) => write!(f, "User-Interactive Authentication required."),
            Self::MatrixError(err) => write!(f, "{}", err),
        }
    }
}

impl From<MatrixError> for UiaaResponse {
    fn from(error: MatrixError) -> Self {
        Self::MatrixError(error)
    }
}

impl EndpointError for UiaaResponse {
    fn try_from_response(
        response: http::Response<Vec<u8>>,
    ) -> Result<Self, ResponseDeserializationError> {
        if response.status() == http::StatusCode::UNAUTHORIZED {
            if let Ok(authentication_info) = from_json_slice::<UiaaInfo>(response.body()) {
                return Ok(UiaaResponse::AuthResponse(authentication_info));
            }
        }

        MatrixError::try_from_response(response).map(From::from)
    }
}

impl std::error::Error for UiaaResponse {}

impl From<UiaaResponse> for http::Response<Vec<u8>> {
    fn from(uiaa_response: UiaaResponse) -> http::Response<Vec<u8>> {
        match uiaa_response {
            UiaaResponse::AuthResponse(authentication_info) => http::Response::builder()
                .header(http::header::CONTENT_TYPE, "application/json")
                .status(&http::StatusCode::UNAUTHORIZED)
                .body(to_json_vec(&authentication_info).unwrap())
                .unwrap(),
            UiaaResponse::MatrixError(error) => http::Response::from(error),
        }
    }
}

#[cfg(test)]
mod tests {
    use maplit::btreemap;
    use matches::assert_matches;
    use ruma_api::EndpointError;
    use serde_json::{
        from_slice as from_json_slice, from_str as from_json_str, from_value as from_json_value,
        json, to_value as to_json_value, value::to_raw_value as to_raw_json_value,
        Value as JsonValue,
    };

    use super::{AuthData, AuthFlow, IncomingAuthData, UiaaInfo, UiaaResponse};
    use crate::error::{ErrorBody, ErrorKind};

    #[test]
    fn serialize_authentication_data_direct_request() {
        let authentication_data = AuthData::DirectRequest {
            kind: "example.type.foo",
            session: Some("ZXY000"),
            auth_parameters: btreemap! {
                "example_credential".to_owned() => json!("verypoorsharedsecret")
            },
        };

        assert_eq!(
            json!({
                "type": "example.type.foo",
                "session": "ZXY000",
                "example_credential": "verypoorsharedsecret",
            }),
            to_json_value(authentication_data).unwrap()
        );
    }

    #[test]
    fn deserialize_authentication_data_direct_request() {
        let json = json!({
            "type": "example.type.foo",
            "session": "opaque_session_id",
            "example_credential": "verypoorsharedsecret",
        });

        assert_matches!(
            from_json_value(json).unwrap(),
            IncomingAuthData::DirectRequest { kind, session: Some(session), auth_parameters }
            if kind == "example.type.foo"
                && session == "opaque_session_id"
                && auth_parameters == btreemap!{
                    "example_credential".to_owned() => json!("verypoorsharedsecret")
                }
        );
    }

    #[test]
    fn serialize_authentication_data_fallback() {
        let authentication_data = AuthData::FallbackAcknowledgement { session: "ZXY000" };

        assert_eq!(json!({ "session": "ZXY000" }), to_json_value(authentication_data).unwrap());
    }

    #[test]
    fn deserialize_authentication_data_fallback() {
        let json = json!({ "session": "opaque_session_id" });

        assert_matches!(
            from_json_value(json).unwrap(),
            IncomingAuthData::FallbackAcknowledgement { session }
            if session == "opaque_session_id"
        );
    }

    #[test]
    fn serialize_uiaa_info() {
        let uiaa_info = UiaaInfo {
            flows: vec![AuthFlow {
                stages: vec!["m.login.password".into(), "m.login.dummy".into()],
            }],
            completed: vec!["m.login.password".into()],
            params: to_raw_json_value(&json!({
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }))
            .unwrap(),
            session: None,
            auth_error: None,
        };

        let json = json!({
            "flows": [{ "stages": ["m.login.password", "m.login.dummy"] }],
            "completed": ["m.login.password"],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }
        });
        assert_eq!(to_json_value(uiaa_info).unwrap(), json);
    }

    #[test]
    fn deserialize_uiaa_info() {
        let json = json!({
            "errcode": "M_FORBIDDEN",
            "error": "Invalid password",
            "completed": ["example.type.foo"],
            "flows": [
                {
                    "stages": ["example.type.foo", "example.type.bar"]
                },
                {
                    "stages": ["example.type.foo", "example.type.baz"]
                }
            ],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            },
            "session": "xxxxxx"
        });

        assert_matches!(
            from_json_value::<UiaaInfo>(json).unwrap(),
            UiaaInfo {
                auth_error: Some(ErrorBody {
                    kind: ErrorKind::Forbidden,
                    message: error_message,
                }),
                completed,
                flows,
                params,
                session: Some(session),
            } if error_message == "Invalid password"
                && completed == vec!["example.type.foo".to_owned()]
                && flows == vec![
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.bar".into(),
                        ],
                    },
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.baz".into(),
                        ],
                    },
                ]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
                && session == "xxxxxx"
        );
    }

    #[test]
    fn try_uiaa_response_into_http_response() {
        let uiaa_info = UiaaInfo {
            flows: vec![AuthFlow {
                stages: vec!["m.login.password".into(), "m.login.dummy".into()],
            }],
            completed: vec!["m.login.password".into()],
            params: to_raw_json_value(&json!({
                "example.type.baz": {
                    "example_key": "foobar"
                }
            }))
            .unwrap(),
            session: None,
            auth_error: None,
        };
        let uiaa_response: http::Response<Vec<u8>> = UiaaResponse::AuthResponse(uiaa_info).into();

        assert_matches!(
            from_json_slice::<UiaaInfo>(uiaa_response.body()).unwrap(),
            UiaaInfo {
                flows,
                completed,
                params,
                session: None,
                auth_error: None,
            } if flows == vec![AuthFlow {
                    stages: vec!["m.login.password".into(), "m.login.dummy".into()],
                }]
                && completed == vec!["m.login.password".to_owned()]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
        );
        assert_eq!(uiaa_response.status(), http::status::StatusCode::UNAUTHORIZED);
    }

    #[test]
    fn try_uiaa_response_from_http_response() {
        let json = serde_json::to_string(&json!({
            "errcode": "M_FORBIDDEN",
            "error": "Invalid password",
            "completed": [ "example.type.foo" ],
            "flows": [
                {
                    "stages": [ "example.type.foo", "example.type.bar" ]
                },
                {
                    "stages": [ "example.type.foo", "example.type.baz" ]
                }
            ],
            "params": {
                "example.type.baz": {
                    "example_key": "foobar"
                }
            },
            "session": "xxxxxx"
        }))
        .unwrap();

        let http_response = http::Response::builder()
            .status(http::StatusCode::UNAUTHORIZED)
            .body(json.into())
            .unwrap();

        let parsed_uiaa_info = match UiaaResponse::try_from_response(http_response).unwrap() {
            UiaaResponse::AuthResponse(uiaa_info) => uiaa_info,
            _ => panic!("Expected UiaaResponse::AuthResponse"),
        };

        assert_matches!(
            parsed_uiaa_info,
            UiaaInfo {
                auth_error: Some(ErrorBody {
                    kind: ErrorKind::Forbidden,
                    message: error_message,
                }),
                completed,
                flows,
                params,
                session: Some(session),
            } if error_message == "Invalid password"
                && completed == vec!["example.type.foo".to_owned()]
                && flows == vec![
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.bar".into(),
                        ],
                    },
                    AuthFlow {
                        stages: vec![
                            "example.type.foo".into(),
                            "example.type.baz".into(),
                        ],
                    },
                ]
                && from_json_str::<JsonValue>(params.get()).unwrap() == json!({
                    "example.type.baz": {
                        "example_key": "foobar"
                    }
                })
                && session == "xxxxxx"
        );
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

fn main() {
    let arr = [1, 6, 3, 3, 5, 8, 7, 2, 7, 7, 4];

    let res = get_mmm(&arr);

    println!("{:?}", res);
}

fn get_mmm(lst: &[i32]) -> Vec<f64> {
    vec![mean(lst), mid(lst), mode(lst) as f64]
}

fn mean(lst: &[i32]) -> f64 {
    let sum = lst.iter().fold(0, |p, &q| p + q);
    sum as f64 / lst.len() as f64
}

fn mid(lst: &[i32]) -> f64 {
    let mut vec = lst.to_vec();
    vec.sort_by(|x, y| x.cmp(y));

    let len = vec.len();
    if len % 2 == 0 {
        (vec[len/2] + vec[len/2 - 1]) as f64 / 2.0
    } else {
        vec[len/2] as f64
    }
}

fn mode(lst: &[i32]) -> i32 {
    let mut count_map = HashMap::new();
    for num in lst {
        let count = count_map.entry(num).or_insert(0);
        *count += 1;
    }

    let mut max_num = 0;
    let mut max_count = 0;
    for (num, count) in count_map {
        if count > max_count {
            max_num = *num;
            max_count = count;
        }
    }

    max_num
}
<commit_msg>Add cha 8 exer 2.<commit_after>use std::collections::HashMap;

fn main() {
    let arr = [1, 6, 3, 3, 5, 8, 7, 2, 7, 7, 4];

    let res = get_mmm(&arr);

    println!("{:?}", res);

    let cons_word = String::from("first");
    println!("{} -> {}", cons_word, to_pig_latin(&cons_word));
    let vow_word = String::from("apple");
    println!("{} -> {}", vow_word, to_pig_latin(&vow_word));
}

fn get_mmm(lst: &[i32]) -> Vec<f64> {
    vec![mean(lst), mid(lst), mode(lst) as f64]
}

fn mean(lst: &[i32]) -> f64 {
    let sum = lst.iter().fold(0, |p, &q| p + q);
    sum as f64 / lst.len() as f64
}

fn mid(lst: &[i32]) -> f64 {
    let mut vec = lst.to_vec();
    vec.sort_by(|x, y| x.cmp(y));

    let len = vec.len();
    if len % 2 == 0 {
        (vec[len/2] + vec[len/2 - 1]) as f64 / 2.0
    } else {
        vec[len/2] as f64
    }
}

fn mode(lst: &[i32]) -> i32 {
    let mut count_map = HashMap::new();
    for num in lst {
        let count = count_map.entry(num).or_insert(0);
        *count += 1;
    }

    let mut max_num = 0;
    let mut max_count = 0;
    for (num, count) in count_map {
        if count > max_count {
            max_num = *num;
            max_count = count;
        }
    }

    max_num
}

fn to_pig_latin(s: &str) -> String {
    match s.chars().nth(0).unwrap() {
        'a' | 'e' | 'i' | 'o' | 'u' => s.to_owned() + "-hay",
        _ => s[1..].to_owned() + "-" + s[0..1].to_owned().as_str() + "ay"
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Runs hardware devices in child processes.

use std::ffi::CString;
use std::fmt::{self, Display};
use std::time::Duration;

use base::{error, AsRawDescriptor, RawDescriptor, Tube, TubeError};
use libc::{self, pid_t};
use minijail::{self, Minijail};
use serde::{Deserialize, Serialize};

use crate::bus::ConfigWriteResult;
use crate::{BusAccessInfo, BusDevice};

/// Errors for proxy devices.
#[derive(Debug)]
pub enum Error {
    ForkingJail(minijail::Error),
    Tube(TubeError),
}
pub type Result<T> = std::result::Result<T, Error>;

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            ForkingJail(e) => write!(f, "Failed to fork jail process: {}", e),
            Tube(e) => write!(f, "Failed to configure tube: {}.", e),
        }
    }
}

const SOCKET_TIMEOUT_MS: u64 = 2000;

#[derive(Debug, Serialize, Deserialize)]
enum Command {
    Read {
        len: u32,
        info: BusAccessInfo,
    },
    Write {
        len: u32,
        info: BusAccessInfo,
        data: [u8; 8],
    },
    ReadConfig(u32),
    WriteConfig {
        reg_idx: u32,
        offset: u32,
        len: u32,
        data: [u8; 4],
    },
    Shutdown,
}
#[derive(Debug, Serialize, Deserialize)]
enum CommandResult {
    Ok,
    ReadResult([u8; 8]),
    ReadConfigResult(u32),
    WriteConfigResult {
        mem_bus_new_state: Option<bool>,
        io_bus_new_state: Option<bool>,
    },
}

fn child_proc<D: BusDevice>(tube: Tube, device: &mut D) {
    let mut running = true;

    while running {
        let cmd = match tube.recv() {
            Ok(cmd) => cmd,
            Err(err) => {
                error!("child device process failed recv: {}", err);
                break;
            }
        };

        let res = match cmd {
            Command::Read { len, info } => {
                let mut buffer = [0u8; 8];
                device.read(info, &mut buffer[0..len as usize]);
                tube.send(&CommandResult::ReadResult(buffer))
            }
            Command::Write { len, info, data } => {
                let len = len as usize;
                device.write(info, &data[0..len]);
                // Command::Write does not have a result.
                Ok(())
            }
            Command::ReadConfig(idx) => {
                let val = device.config_register_read(idx as usize);
                tube.send(&CommandResult::ReadConfigResult(val))
            }
            Command::WriteConfig {
                reg_idx,
                offset,
                len,
                data,
            } => {
                let len = len as usize;
                let res =
                    device.config_register_write(reg_idx as usize, offset as u64, &data[0..len]);
                tube.send(&CommandResult::WriteConfigResult {
                    mem_bus_new_state: res.mem_bus_new_state,
                    io_bus_new_state: res.io_bus_new_state,
                })
            }
            Command::Shutdown => {
                running = false;
                tube.send(&CommandResult::Ok)
            }
        };
        if let Err(e) = res {
            error!("child device process failed send: {}", e);
        }
    }
}

/// Wraps an inner `BusDevice` that is run inside a child process via fork.
///
/// Because forks are very unfriendly to destructors and all memory mappings and file descriptors
/// are inherited, this should be used as early as possible in the main process.
pub struct ProxyDevice {
    tube: Tube,
    pid: pid_t,
    debug_label: String,
}

impl ProxyDevice {
    /// Takes the given device and isolates it into another process via fork before returning.
    ///
    /// The forked process will automatically be terminated when this is dropped, so be sure to keep
    /// a reference.
    ///
    /// # Arguments
    /// * `device` - The device to isolate to another process.
    /// * `jail` - The jail to use for isolating the given device.
    /// * `keep_rds` - File descriptors that will be kept open in the child.
    pub fn new<D: BusDevice>(
        mut device: D,
        jail: &Minijail,
        mut keep_rds: Vec<RawDescriptor>,
    ) -> Result<ProxyDevice> {
        let debug_label = device.debug_label();
        let (child_tube, parent_tube) = Tube::pair().map_err(Error::Tube)?;

        keep_rds.push(child_tube.as_raw_descriptor());
        // Forking here is safe as long as the program is still single threaded.
        let pid = unsafe {
            match jail.fork(Some(&keep_rds)).map_err(Error::ForkingJail)? {
                0 => {
                    let max_len = 15; // pthread_setname_np() limit on Linux
                    let debug_label_trimmed =
                        &debug_label.as_bytes()[..std::cmp::min(max_len, debug_label.len())];
                    let thread_name = CString::new(debug_label_trimmed).unwrap();
                    // TODO(crbug.com/1199487): remove this once libc provides the wrapper for all
                    // targets
                    #[cfg(all(target_os = "linux", target_env = "gnu"))]
                    let _ = libc::pthread_setname_np(libc::pthread_self(), thread_name.as_ptr());
                    device.on_sandboxed();
                    child_proc(child_tube, &mut device);

                    // We're explicitly not using std::process::exit here to avoid the cleanup of
                    // stdout/stderr globals. This can cause cascading panics and SIGILL if a worker
                    // thread attempts to log to stderr after at_exit handlers have been run.
                    // TODO(crbug.com/992494): Remove this once device shutdown ordering is clearly
                    // defined.
                    //
                    // exit() is trivially safe.
                    // ! Never returns
                    libc::exit(0);
                }
                p => p,
            }
        };

        parent_tube
            .set_send_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        parent_tube
            .set_recv_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        Ok(ProxyDevice {
            tube: parent_tube,
            pid,
            debug_label,
        })
    }

    pub fn pid(&self) -> pid_t {
        self.pid
    }

    /// Send a command that does not expect a response from the child device process.
    fn send_no_result(&self, cmd: &Command) {
        let res = self.tube.send(cmd);
        if let Err(e) = res {
            error!(
                "failed write to child device process {}: {}",
                self.debug_label, e,
            );
        }
    }

    /// Send a command and read its response from the child device process.
    fn sync_send(&self, cmd: &Command) -> Option<CommandResult> {
        self.send_no_result(cmd);
        match self.tube.recv() {
            Err(e) => {
                error!(
                    "failed to read result of {:?} from child device process {}: {}",
                    cmd, self.debug_label, e,
                );
                None
            }
            Ok(r) => Some(r),
        }
    }
}

impl BusDevice for ProxyDevice {
    fn debug_label(&self) -> String {
        self.debug_label.clone()
    }

    fn config_register_write(
        &mut self,
        reg_idx: usize,
        offset: u64,
        data: &[u8],
    ) -> ConfigWriteResult {
        let len = data.len() as u32;
        let mut buffer = [0u8; 4];
        buffer[0..data.len()].clone_from_slice(data);
        let reg_idx = reg_idx as u32;
        let offset = offset as u32;
        if let Some(CommandResult::WriteConfigResult {
            mem_bus_new_state,
            io_bus_new_state,
        }) = self.sync_send(&Command::WriteConfig {
            reg_idx,
            offset,
            len,
            data: buffer,
        }) {
            ConfigWriteResult {
                mem_bus_new_state,
                io_bus_new_state,
            }
        } else {
            Default::default()
        }
    }

    fn config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadConfig(reg_idx as u32));
        if let Some(CommandResult::ReadConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn read(&mut self, info: BusAccessInfo, data: &mut [u8]) {
        let len = data.len() as u32;
        if let Some(CommandResult::ReadResult(buffer)) =
            self.sync_send(&Command::Read { len, info })
        {
            let len = data.len();
            data.clone_from_slice(&buffer[0..len]);
        }
    }

    fn write(&mut self, info: BusAccessInfo, data: &[u8]) {
        let mut buffer = [0u8; 8];
        let len = data.len() as u32;
        buffer[0..data.len()].clone_from_slice(data);
        self.send_no_result(&Command::Write {
            len,
            info,
            data: buffer,
        });
    }
}

impl Drop for ProxyDevice {
    fn drop(&mut self) {
        self.sync_send(&Command::Shutdown);
    }
}

/// Note: These tests must be run with --test-threads=1 to allow minijail to fork
/// the process.
#[cfg(test)]
mod tests {
    use super::*;

    /// A simple test echo device that outputs the same u8 that was written to it.
    struct EchoDevice {
        data: u8,
        config: u8,
    }
    impl EchoDevice {
        fn new() -> EchoDevice {
            EchoDevice { data: 0, config: 0 }
        }
    }
    impl BusDevice for EchoDevice {
        fn debug_label(&self) -> String {
            "EchoDevice".to_owned()
        }

        fn write(&mut self, _info: BusAccessInfo, data: &[u8]) {
            assert!(data.len() == 1);
            self.data = data[0];
        }

        fn read(&mut self, _info: BusAccessInfo, data: &mut [u8]) {
            assert!(data.len() == 1);
            data[0] = self.data;
        }

        fn config_register_write(
            &mut self,
            _reg_idx: usize,
            _offset: u64,
            data: &[u8],
        ) -> ConfigWriteResult {
            let result = ConfigWriteResult {
                ..Default::default()
            };
            assert!(data.len() == 1);
            self.config = data[0];
            result
        }

        fn config_register_read(&self, _reg_idx: usize) -> u32 {
            self.config as u32
        }
    }

    fn new_proxied_echo_device() -> ProxyDevice {
        let device = EchoDevice::new();
        let keep_fds: Vec<RawDescriptor> = Vec::new();
        let minijail = Minijail::new().unwrap();
        ProxyDevice::new(device, &minijail, keep_fds).unwrap()
    }

    // TODO(b/173833661): Find a way to ensure these tests are run single-threaded.
    #[test]
    #[ignore]
    fn test_debug_label() {
        let proxy_device = new_proxied_echo_device();
        assert_eq!(proxy_device.debug_label(), "EchoDevice");
    }

    #[test]
    #[ignore]
    fn test_proxied_read_write() {
        let mut proxy_device = new_proxied_echo_device();
        let address = BusAccessInfo {
            offset: 0,
            address: 0,
            id: 0,
        };
        proxy_device.write(address, &[42]);
        let mut read_buffer = [0];
        proxy_device.read(address, &mut read_buffer);
        assert_eq!(read_buffer, [42]);
    }

    #[test]
    #[ignore]
    fn test_proxied_config() {
        let mut proxy_device = new_proxied_echo_device();
        proxy_device.config_register_write(0, 0, &[42]);
        assert_eq!(proxy_device.config_register_read(0), 42);
    }
}
<commit_msg>devices: Deduplicate mapped file descriptors.<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Runs hardware devices in child processes.

use std::ffi::CString;
use std::fmt::{self, Display};
use std::time::Duration;

use base::{error, AsRawDescriptor, RawDescriptor, Tube, TubeError};
use libc::{self, pid_t};
use minijail::{self, Minijail};
use serde::{Deserialize, Serialize};

use crate::bus::ConfigWriteResult;
use crate::{BusAccessInfo, BusDevice};

/// Errors for proxy devices.
#[derive(Debug)]
pub enum Error {
    ForkingJail(minijail::Error),
    Tube(TubeError),
}
pub type Result<T> = std::result::Result<T, Error>;

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            ForkingJail(e) => write!(f, "Failed to fork jail process: {}", e),
            Tube(e) => write!(f, "Failed to configure tube: {}.", e),
        }
    }
}

const SOCKET_TIMEOUT_MS: u64 = 2000;

#[derive(Debug, Serialize, Deserialize)]
enum Command {
    Read {
        len: u32,
        info: BusAccessInfo,
    },
    Write {
        len: u32,
        info: BusAccessInfo,
        data: [u8; 8],
    },
    ReadConfig(u32),
    WriteConfig {
        reg_idx: u32,
        offset: u32,
        len: u32,
        data: [u8; 4],
    },
    Shutdown,
}
#[derive(Debug, Serialize, Deserialize)]
enum CommandResult {
    Ok,
    ReadResult([u8; 8]),
    ReadConfigResult(u32),
    WriteConfigResult {
        mem_bus_new_state: Option<bool>,
        io_bus_new_state: Option<bool>,
    },
}

fn child_proc<D: BusDevice>(tube: Tube, device: &mut D) {
    let mut running = true;

    while running {
        let cmd = match tube.recv() {
            Ok(cmd) => cmd,
            Err(err) => {
                error!("child device process failed recv: {}", err);
                break;
            }
        };

        let res = match cmd {
            Command::Read { len, info } => {
                let mut buffer = [0u8; 8];
                device.read(info, &mut buffer[0..len as usize]);
                tube.send(&CommandResult::ReadResult(buffer))
            }
            Command::Write { len, info, data } => {
                let len = len as usize;
                device.write(info, &data[0..len]);
                // Command::Write does not have a result.
                Ok(())
            }
            Command::ReadConfig(idx) => {
                let val = device.config_register_read(idx as usize);
                tube.send(&CommandResult::ReadConfigResult(val))
            }
            Command::WriteConfig {
                reg_idx,
                offset,
                len,
                data,
            } => {
                let len = len as usize;
                let res =
                    device.config_register_write(reg_idx as usize, offset as u64, &data[0..len]);
                tube.send(&CommandResult::WriteConfigResult {
                    mem_bus_new_state: res.mem_bus_new_state,
                    io_bus_new_state: res.io_bus_new_state,
                })
            }
            Command::Shutdown => {
                running = false;
                tube.send(&CommandResult::Ok)
            }
        };
        if let Err(e) = res {
            error!("child device process failed send: {}", e);
        }
    }
}

/// Wraps an inner `BusDevice` that is run inside a child process via fork.
///
/// Because forks are very unfriendly to destructors and all memory mappings and file descriptors
/// are inherited, this should be used as early as possible in the main process.
pub struct ProxyDevice {
    tube: Tube,
    pid: pid_t,
    debug_label: String,
}

impl ProxyDevice {
    /// Takes the given device and isolates it into another process via fork before returning.
    ///
    /// The forked process will automatically be terminated when this is dropped, so be sure to keep
    /// a reference.
    ///
    /// # Arguments
    /// * `device` - The device to isolate to another process.
    /// * `jail` - The jail to use for isolating the given device.
    /// * `keep_rds` - File descriptors that will be kept open in the child.
    pub fn new<D: BusDevice>(
        mut device: D,
        jail: &Minijail,
        mut keep_rds: Vec<RawDescriptor>,
    ) -> Result<ProxyDevice> {
        let debug_label = device.debug_label();
        let (child_tube, parent_tube) = Tube::pair().map_err(Error::Tube)?;

        keep_rds.push(child_tube.as_raw_descriptor());

        // Deduplicate the FDs since minijail expects this.
        keep_rds.sort_unstable();
        keep_rds.dedup();

        // Forking here is safe as long as the program is still single threaded.
        let pid = unsafe {
            match jail.fork(Some(&keep_rds)).map_err(Error::ForkingJail)? {
                0 => {
                    let max_len = 15; // pthread_setname_np() limit on Linux
                    let debug_label_trimmed =
                        &debug_label.as_bytes()[..std::cmp::min(max_len, debug_label.len())];
                    let thread_name = CString::new(debug_label_trimmed).unwrap();
                    // TODO(crbug.com/1199487): remove this once libc provides the wrapper for all
                    // targets
                    #[cfg(all(target_os = "linux", target_env = "gnu"))]
                    let _ = libc::pthread_setname_np(libc::pthread_self(), thread_name.as_ptr());
                    device.on_sandboxed();
                    child_proc(child_tube, &mut device);

                    // We're explicitly not using std::process::exit here to avoid the cleanup of
                    // stdout/stderr globals. This can cause cascading panics and SIGILL if a worker
                    // thread attempts to log to stderr after at_exit handlers have been run.
                    // TODO(crbug.com/992494): Remove this once device shutdown ordering is clearly
                    // defined.
                    //
                    // exit() is trivially safe.
                    // ! Never returns
                    libc::exit(0);
                }
                p => p,
            }
        };

        parent_tube
            .set_send_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        parent_tube
            .set_recv_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        Ok(ProxyDevice {
            tube: parent_tube,
            pid,
            debug_label,
        })
    }

    pub fn pid(&self) -> pid_t {
        self.pid
    }

    /// Send a command that does not expect a response from the child device process.
    fn send_no_result(&self, cmd: &Command) {
        let res = self.tube.send(cmd);
        if let Err(e) = res {
            error!(
                "failed write to child device process {}: {}",
                self.debug_label, e,
            );
        }
    }

    /// Send a command and read its response from the child device process.
    fn sync_send(&self, cmd: &Command) -> Option<CommandResult> {
        self.send_no_result(cmd);
        match self.tube.recv() {
            Err(e) => {
                error!(
                    "failed to read result of {:?} from child device process {}: {}",
                    cmd, self.debug_label, e,
                );
                None
            }
            Ok(r) => Some(r),
        }
    }
}

impl BusDevice for ProxyDevice {
    fn debug_label(&self) -> String {
        self.debug_label.clone()
    }

    fn config_register_write(
        &mut self,
        reg_idx: usize,
        offset: u64,
        data: &[u8],
    ) -> ConfigWriteResult {
        let len = data.len() as u32;
        let mut buffer = [0u8; 4];
        buffer[0..data.len()].clone_from_slice(data);
        let reg_idx = reg_idx as u32;
        let offset = offset as u32;
        if let Some(CommandResult::WriteConfigResult {
            mem_bus_new_state,
            io_bus_new_state,
        }) = self.sync_send(&Command::WriteConfig {
            reg_idx,
            offset,
            len,
            data: buffer,
        }) {
            ConfigWriteResult {
                mem_bus_new_state,
                io_bus_new_state,
            }
        } else {
            Default::default()
        }
    }

    fn config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadConfig(reg_idx as u32));
        if let Some(CommandResult::ReadConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn read(&mut self, info: BusAccessInfo, data: &mut [u8]) {
        let len = data.len() as u32;
        if let Some(CommandResult::ReadResult(buffer)) =
            self.sync_send(&Command::Read { len, info })
        {
            let len = data.len();
            data.clone_from_slice(&buffer[0..len]);
        }
    }

    fn write(&mut self, info: BusAccessInfo, data: &[u8]) {
        let mut buffer = [0u8; 8];
        let len = data.len() as u32;
        buffer[0..data.len()].clone_from_slice(data);
        self.send_no_result(&Command::Write {
            len,
            info,
            data: buffer,
        });
    }
}

impl Drop for ProxyDevice {
    fn drop(&mut self) {
        self.sync_send(&Command::Shutdown);
    }
}

/// Note: These tests must be run with --test-threads=1 to allow minijail to fork
/// the process.
#[cfg(test)]
mod tests {
    use super::*;

    /// A simple test echo device that outputs the same u8 that was written to it.
    struct EchoDevice {
        data: u8,
        config: u8,
    }
    impl EchoDevice {
        fn new() -> EchoDevice {
            EchoDevice { data: 0, config: 0 }
        }
    }
    impl BusDevice for EchoDevice {
        fn debug_label(&self) -> String {
            "EchoDevice".to_owned()
        }

        fn write(&mut self, _info: BusAccessInfo, data: &[u8]) {
            assert!(data.len() == 1);
            self.data = data[0];
        }

        fn read(&mut self, _info: BusAccessInfo, data: &mut [u8]) {
            assert!(data.len() == 1);
            data[0] = self.data;
        }

        fn config_register_write(
            &mut self,
            _reg_idx: usize,
            _offset: u64,
            data: &[u8],
        ) -> ConfigWriteResult {
            let result = ConfigWriteResult {
                ..Default::default()
            };
            assert!(data.len() == 1);
            self.config = data[0];
            result
        }

        fn config_register_read(&self, _reg_idx: usize) -> u32 {
            self.config as u32
        }
    }

    fn new_proxied_echo_device() -> ProxyDevice {
        let device = EchoDevice::new();
        let keep_fds: Vec<RawDescriptor> = Vec::new();
        let minijail = Minijail::new().unwrap();
        ProxyDevice::new(device, &minijail, keep_fds).unwrap()
    }

    // TODO(b/173833661): Find a way to ensure these tests are run single-threaded.
    #[test]
    #[ignore]
    fn test_debug_label() {
        let proxy_device = new_proxied_echo_device();
        assert_eq!(proxy_device.debug_label(), "EchoDevice");
    }

    #[test]
    #[ignore]
    fn test_proxied_read_write() {
        let mut proxy_device = new_proxied_echo_device();
        let address = BusAccessInfo {
            offset: 0,
            address: 0,
            id: 0,
        };
        proxy_device.write(address, &[42]);
        let mut read_buffer = [0];
        proxy_device.read(address, &mut read_buffer);
        assert_eq!(read_buffer, [42]);
    }

    #[test]
    #[ignore]
    fn test_proxied_config() {
        let mut proxy_device = new_proxied_echo_device();
        proxy_device.config_register_write(0, 0, &[42]);
        assert_eq!(proxy_device.config_register_read(0), 42);
    }
}
<|endoftext|>


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

use notify::{Watcher, RecursiveMode, watcher, DebouncedEvent};
use std::path::Path;
use std::sync::mpsc::channel;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;
use std::collections::VecDeque;

use xi_rpc::RpcPeer;

/// xi_rpc idle Token for watcher related idle scheduling.
pub const WATCH_IDLE_TOKEN: usize = 1002;

/// Token provided to `FsWatcher`, to associate events with registrees.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EventToken(pub usize);

/// Wrapper around `notify::Watcher`, adding support for the xi_rpc runloop.
#[derive(Debug, Clone, Default)]
pub struct FsWatcher {
    pub events: Arc<Mutex<VecDeque<(EventToken, DebouncedEvent)>>>,
}

impl FsWatcher {
    /// Begin watching `path`. As `DebouncedEvent`s (documented in the [notify](https://docs.rs/notify/4.0.2/notify/) crate)
    /// arrive, they are stored with the associated `token` and a task is
    /// added to the runloop's idle queue.
    ///
    /// Delivery of events then requires that the runloop's handler
    /// correctly forward the `handle_idle` call to the interested party.
    pub fn watch<P>(&mut self, path: P, recursive_mode: RecursiveMode,
                token: EventToken, peer: &RpcPeer)
        where P: AsRef<Path>,
    {
        let path = path.as_ref().to_owned();
        let peer = peer.clone();
        let events = self.events.clone();
        thread::spawn(move || {

            let (tx, rx) = channel();
            let mut watcher = watcher(tx, Duration::from_secs(1)).unwrap();

            watcher.watch(&path, recursive_mode).unwrap();

            loop {
                match rx.recv() {
                    Ok(event) =>  {
                        events.lock().unwrap().push_back((token, event));
                        peer.schedule_idle(WATCH_IDLE_TOKEN);
                    },
                    Err(e) => {
                        //TODO: how do we handle unexpected disconnects?
                        eprintln!("watcher returned error {:?} for path {:?}, \
                        token {:?}", e, &path, token);
                        break
                    }
                }
            }
        });
    }
    //TODO impl unwatch, when we add in watching of opened files
}
<commit_msg>Add FsWatcher::watch_filtered<commit_after>// Copyright 2017 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use notify::{Watcher, RecursiveMode, watcher, DebouncedEvent};
use std::path::Path;
use std::sync::mpsc::channel;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;
use std::collections::VecDeque;

use xi_rpc::RpcPeer;

/// xi_rpc idle Token for watcher related idle scheduling.
pub const WATCH_IDLE_TOKEN: usize = 1002;

/// Token provided to `FsWatcher`, to associate events with registrees.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EventToken(pub usize);

/// Wrapper around `notify::Watcher`, adding support for the xi_rpc runloop.
#[derive(Debug, Clone, Default)]
pub struct FsWatcher {
    pub events: Arc<Mutex<VecDeque<(EventToken, DebouncedEvent)>>>,
}

impl FsWatcher {
    /// Begin watching `path`. As `DebouncedEvent`s (documented in the [notify](https://docs.rs/notify/4.0.2/notify/) crate)
    /// arrive, they are stored with the associated `token` and a task is
    /// added to the runloop's idle queue.
    ///
    /// Delivery of events then requires that the runloop's handler
    /// correctly forward the `handle_idle` call to the interested party.
    pub fn watch<P>(&mut self, path: P, recursive_mode: RecursiveMode,
                token: EventToken, peer: &RpcPeer)
        where P: AsRef<Path>,
    {
        self.watch_filtered(path, recursive_mode, token, peer, |_| { true });
    }

    /// Like `watch`, but taking a predicate function that filters delivery
    /// of events based on their path.
    pub fn watch_filtered<P, F>(&mut self, path: P, recursive_mode: RecursiveMode,
                                token: EventToken, peer: &RpcPeer, predicate: F)
        where P: AsRef<Path>,
              F: Fn(&Path) -> bool + Send + 'static,
    {
        let path = path.as_ref().to_owned();
        let peer = peer.clone();
        let events = self.events.clone();
        thread::spawn(move || {

            let (tx, rx) = channel();
            let mut watcher = watcher(tx, Duration::from_secs(1)).unwrap();

            watcher.watch(&path, recursive_mode).unwrap();

            loop {
                match rx.recv() {
                    Ok(event) =>  {
                        if apply_filter(&predicate, &event) {
                            events.lock().unwrap().push_back((token, event));
                            peer.schedule_idle(WATCH_IDLE_TOKEN);
                        }
                    },
                    Err(e) => {
                        //TODO: how do we handle unexpected disconnects?
                        eprintln!("watcher returned error {:?} for path {:?}, \
                        token {:?}", e, &path, token);
                        break
                    }
                }
            }
        });
    }
    //TODO impl unwatch, when we add in watching of opened files
}

/// Checks the predicate against the various event cases
fn apply_filter<F>(filter: &F, event: &DebouncedEvent) -> bool
    where F: Fn(&Path) -> bool,
{
    use self::DebouncedEvent::*;
    match *event {
        NoticeWrite(ref p) | NoticeRemove(ref p) | Create(ref p) |
            Write(ref p) | Chmod(ref p) | Remove(ref p) => {
                filter(p)
            }
        Rename(ref p1, ref p2) => {
            filter(p1) || filter(p2)
        }
        Rescan => false,
        Error(_, ref opt_p) => {
            opt_p.as_ref().map(|p| filter(p))
                .unwrap_or(false)
        }
    }
}

<|endoftext|>


<commit_before>//! An asynchronous, pipelined, PostgreSQL client.
//!
//! # Example
//!
//! ```no_run
//! use tokio_postgres::{NoTls, Error, Row};
//!
//! # #[cfg(not(feature = "runtime"))] fn main() {}
//! # #[cfg(feature = "runtime")]
//! #[tokio::main] // By default, tokio_postgres uses the tokio crate as its runtime.
//! async fn main() -> Result<(), Error> {
//!     // Connect to the database.
//!     let (mut client, connection) =
//!         tokio_postgres::connect("host=localhost user=postgres", NoTls).await?;
//!
//!     // The connection object performs the actual communication with the database,
//!     // so spawn it off to run on its own.
//!     tokio::spawn(async move {
//!         if let Err(e) = connection.await {
//!             eprintln!("connection error: {}", e);
//!         }
//!     });
//!
//!     // Now we can prepare a simple statement that just returns its parameter.
//!     let stmt = client.prepare("SELECT $1::TEXT").await?;
//!
//!     // And then execute it, returning a list of the resulting rows.
//!     let rows = client
//!         .query(&stmt, &[&"hello world"])
//!         .await?;
//!
//!     // Now we can check that we got back the same string we sent over.
//!     let value: &str = rows[0].get(0);
//!     assert_eq!(value, "hello world");
//!
//!     Ok(())
//! }
//! ```
//!
//! # Behavior
//!
//! Calling a method like `Client::query` on its own does nothing. The associated request is not sent to the database
//! until the future returned by the method is first polled. Requests are executed in the order that they are first
//! polled, not in the order that their futures are created.
//!
//! # Pipelining
//!
//! The client supports *pipelined* requests. Pipelining can improve performance in use cases in which multiple,
//! independent queries need to be executed. In a traditional workflow, each query is sent to the server after the
//! previous query completes. In contrast, pipelining allows the client to send all of the queries to the server up
//! front, minimizing time spent by one side waiting for the other to finish sending data:
//!
//! ```not_rust
//!             Sequential                              Pipelined
//! | Client         | Server          |    | Client         | Server          |
//! |----------------|-----------------|    |----------------|-----------------|
//! | send query 1   |                 |    | send query 1   |                 |
//! |                | process query 1 |    | send query 2   | process query 1 |
//! | receive rows 1 |                 |    | send query 3   | process query 2 |
//! | send query 2   |                 |    | receive rows 1 | process query 3 |
//! |                | process query 2 |    | receive rows 2 |                 |
//! | receive rows 2 |                 |    | receive rows 3 |                 |
//! | send query 3   |                 |
//! |                | process query 3 |
//! | receive rows 3 |                 |
//! ```
//!
//! In both cases, the PostgreSQL server is executing the queries sequentially - pipelining just allows both sides of
//! the connection to work concurrently when possible.
//!
//! Pipelining happens automatically when futures are polled concurrently (for example, by using the futures `join`
//! combinator):
//!
//! ```rust
//! use futures::future;
//! use std::future::Future;
//! use tokio_postgres::{Client, Error, Statement};
//!
//! async fn pipelined_prepare(
//!     client: &Client,
//! ) -> Result<(Statement, Statement), Error>
//! {
//!     future::try_join(
//!         client.prepare("SELECT * FROM foo"),
//!         client.prepare("INSERT INTO bar (id, name) VALUES ($1, $2)")
//!     ).await
//! }
//! ```
//!
//! # Runtime
//!
//! The client works with arbitrary `AsyncRead + AsyncWrite` streams. Convenience APIs are provided to handle the
//! connection process, but these are gated by the `runtime` Cargo feature, which is enabled by default. If disabled,
//! all dependence on the tokio runtime is removed.
//!
//! # SSL/TLS support
//!
//! TLS support is implemented via external libraries. `Client::connect` and `Config::connect` take a TLS implementation
//! as an argument. The `NoTls` type in this crate can be used when TLS is not required. Otherwise, the
//! `postgres-openssl` and `postgres-native-tls` crates provide implementations backed by the `openssl` and `native-tls`
//! crates, respectively.
#![doc(html_root_url = "https://docs.rs/tokio-postgres/0.5")]
#![warn(rust_2018_idioms, clippy::all, missing_docs)]

pub use crate::client::Client;
pub use crate::config::Config;
pub use crate::connection::Connection;
use crate::error::DbError;
pub use crate::error::Error;
pub use crate::portal::Portal;
pub use crate::query::RowStream;
pub use crate::row::{Row, SimpleQueryRow};
pub use crate::simple_query::SimpleQueryStream;
#[cfg(feature = "runtime")]
pub use crate::socket::Socket;
pub use crate::statement::{Column, Statement};
#[cfg(feature = "runtime")]
use crate::tls::MakeTlsConnect;
pub use crate::tls::NoTls;
pub use crate::to_statement::ToStatement;
pub use crate::transaction::Transaction;
use crate::types::ToSql;

mod bind;
#[cfg(feature = "runtime")]
mod cancel_query;
mod cancel_query_raw;
mod client;
mod codec;
pub mod config;
#[cfg(feature = "runtime")]
mod connect;
mod connect_raw;
#[cfg(feature = "runtime")]
mod connect_socket;
mod connect_tls;
mod connection;
mod copy_in;
mod copy_out;
pub mod error;
mod maybe_tls_stream;
mod portal;
mod prepare;
mod query;
pub mod row;
mod simple_query;
#[cfg(feature = "runtime")]
mod socket;
mod statement;
pub mod tls;
mod to_statement;
mod transaction;
pub mod types;

/// A convenience function which parses a connection string and connects to the database.
///
/// See the documentation for [`Config`] for details on the connection string format.
///
/// Requires the `runtime` Cargo feature (enabled by default).
///
/// [`Config`]: config/struct.Config.html
#[cfg(feature = "runtime")]
pub async fn connect<T>(
    config: &str,
    tls: T,
) -> Result<(Client, Connection<Socket, T::Stream>), Error>
where
    T: MakeTlsConnect<Socket>,
{
    let config = config.parse::<Config>()?;
    config.connect(tls).await
}

/// An asynchronous notification.
#[derive(Clone, Debug)]
pub struct Notification {
    process_id: i32,
    channel: String,
    payload: String,
}

impl Notification {
    /// The process ID of the notifying backend process.
    pub fn process_id(&self) -> i32 {
        self.process_id
    }

    /// The name of the channel that the notify has been raised on.
    pub fn channel(&self) -> &str {
        &self.channel
    }

    /// The "payload" string passed from the notifying process.
    pub fn payload(&self) -> &str {
        &self.payload
    }
}

/// An asynchronous message from the server.
#[allow(clippy::large_enum_variant)]
pub enum AsyncMessage {
    /// A notice.
    ///
    /// Notices use the same format as errors, but aren't "errors" per-se.
    Notice(DbError),
    /// A notification.
    ///
    /// Connections can subscribe to notifications with the `LISTEN` command.
    Notification(Notification),
    #[doc(hidden)]
    __NonExhaustive,
}

/// Message returned by the `SimpleQuery` stream.
pub enum SimpleQueryMessage {
    /// A row of data.
    Row(SimpleQueryRow),
    /// A statement in the query has completed.
    ///
    /// The number of rows modified or selected is returned.
    CommandComplete(u64),
    #[doc(hidden)]
    __NonExhaustive,
}

fn slice_iter<'a>(
    s: &'a [&'a (dyn ToSql + Sync)],
) -> impl ExactSizeIterator<Item = &'a dyn ToSql> + 'a {
    s.iter().map(|s| *s as _)
}
<commit_msg>more example fixes<commit_after>//! An asynchronous, pipelined, PostgreSQL client.
//!
//! # Example
//!
//! ```no_run
//! use tokio_postgres::{NoTls, Error};
//!
//! # #[cfg(not(feature = "runtime"))] fn main() {}
//! # #[cfg(feature = "runtime")]
//! #[tokio::main] // By default, tokio_postgres uses the tokio crate as its runtime.
//! async fn main() -> Result<(), Error> {
//!     // Connect to the database.
//!     let (client, connection) =
//!         tokio_postgres::connect("host=localhost user=postgres", NoTls).await?;
//!
//!     // The connection object performs the actual communication with the database,
//!     // so spawn it off to run on its own.
//!     tokio::spawn(async move {
//!         if let Err(e) = connection.await {
//!             eprintln!("connection error: {}", e);
//!         }
//!     });
//!
//!     // Now we can prepare a simple statement that just returns its parameter.
//!     let stmt = client.prepare("SELECT $1::TEXT").await?;
//!
//!     // And then execute it, returning a list of the resulting rows.
//!     let rows = client
//!         .query(&stmt, &[&"hello world"])
//!         .await?;
//!
//!     // Now we can check that we got back the same string we sent over.
//!     let value: &str = rows[0].get(0);
//!     assert_eq!(value, "hello world");
//!
//!     Ok(())
//! }
//! ```
//!
//! # Behavior
//!
//! Calling a method like `Client::query` on its own does nothing. The associated request is not sent to the database
//! until the future returned by the method is first polled. Requests are executed in the order that they are first
//! polled, not in the order that their futures are created.
//!
//! # Pipelining
//!
//! The client supports *pipelined* requests. Pipelining can improve performance in use cases in which multiple,
//! independent queries need to be executed. In a traditional workflow, each query is sent to the server after the
//! previous query completes. In contrast, pipelining allows the client to send all of the queries to the server up
//! front, minimizing time spent by one side waiting for the other to finish sending data:
//!
//! ```not_rust
//!             Sequential                              Pipelined
//! | Client         | Server          |    | Client         | Server          |
//! |----------------|-----------------|    |----------------|-----------------|
//! | send query 1   |                 |    | send query 1   |                 |
//! |                | process query 1 |    | send query 2   | process query 1 |
//! | receive rows 1 |                 |    | send query 3   | process query 2 |
//! | send query 2   |                 |    | receive rows 1 | process query 3 |
//! |                | process query 2 |    | receive rows 2 |                 |
//! | receive rows 2 |                 |    | receive rows 3 |                 |
//! | send query 3   |                 |
//! |                | process query 3 |
//! | receive rows 3 |                 |
//! ```
//!
//! In both cases, the PostgreSQL server is executing the queries sequentially - pipelining just allows both sides of
//! the connection to work concurrently when possible.
//!
//! Pipelining happens automatically when futures are polled concurrently (for example, by using the futures `join`
//! combinator):
//!
//! ```rust
//! use futures::future;
//! use std::future::Future;
//! use tokio_postgres::{Client, Error, Statement};
//!
//! async fn pipelined_prepare(
//!     client: &Client,
//! ) -> Result<(Statement, Statement), Error>
//! {
//!     future::try_join(
//!         client.prepare("SELECT * FROM foo"),
//!         client.prepare("INSERT INTO bar (id, name) VALUES ($1, $2)")
//!     ).await
//! }
//! ```
//!
//! # Runtime
//!
//! The client works with arbitrary `AsyncRead + AsyncWrite` streams. Convenience APIs are provided to handle the
//! connection process, but these are gated by the `runtime` Cargo feature, which is enabled by default. If disabled,
//! all dependence on the tokio runtime is removed.
//!
//! # SSL/TLS support
//!
//! TLS support is implemented via external libraries. `Client::connect` and `Config::connect` take a TLS implementation
//! as an argument. The `NoTls` type in this crate can be used when TLS is not required. Otherwise, the
//! `postgres-openssl` and `postgres-native-tls` crates provide implementations backed by the `openssl` and `native-tls`
//! crates, respectively.
#![doc(html_root_url = "https://docs.rs/tokio-postgres/0.5")]
#![warn(rust_2018_idioms, clippy::all, missing_docs)]

pub use crate::client::Client;
pub use crate::config::Config;
pub use crate::connection::Connection;
use crate::error::DbError;
pub use crate::error::Error;
pub use crate::portal::Portal;
pub use crate::query::RowStream;
pub use crate::row::{Row, SimpleQueryRow};
pub use crate::simple_query::SimpleQueryStream;
#[cfg(feature = "runtime")]
pub use crate::socket::Socket;
pub use crate::statement::{Column, Statement};
#[cfg(feature = "runtime")]
use crate::tls::MakeTlsConnect;
pub use crate::tls::NoTls;
pub use crate::to_statement::ToStatement;
pub use crate::transaction::Transaction;
use crate::types::ToSql;

mod bind;
#[cfg(feature = "runtime")]
mod cancel_query;
mod cancel_query_raw;
mod client;
mod codec;
pub mod config;
#[cfg(feature = "runtime")]
mod connect;
mod connect_raw;
#[cfg(feature = "runtime")]
mod connect_socket;
mod connect_tls;
mod connection;
mod copy_in;
mod copy_out;
pub mod error;
mod maybe_tls_stream;
mod portal;
mod prepare;
mod query;
pub mod row;
mod simple_query;
#[cfg(feature = "runtime")]
mod socket;
mod statement;
pub mod tls;
mod to_statement;
mod transaction;
pub mod types;

/// A convenience function which parses a connection string and connects to the database.
///
/// See the documentation for [`Config`] for details on the connection string format.
///
/// Requires the `runtime` Cargo feature (enabled by default).
///
/// [`Config`]: config/struct.Config.html
#[cfg(feature = "runtime")]
pub async fn connect<T>(
    config: &str,
    tls: T,
) -> Result<(Client, Connection<Socket, T::Stream>), Error>
where
    T: MakeTlsConnect<Socket>,
{
    let config = config.parse::<Config>()?;
    config.connect(tls).await
}

/// An asynchronous notification.
#[derive(Clone, Debug)]
pub struct Notification {
    process_id: i32,
    channel: String,
    payload: String,
}

impl Notification {
    /// The process ID of the notifying backend process.
    pub fn process_id(&self) -> i32 {
        self.process_id
    }

    /// The name of the channel that the notify has been raised on.
    pub fn channel(&self) -> &str {
        &self.channel
    }

    /// The "payload" string passed from the notifying process.
    pub fn payload(&self) -> &str {
        &self.payload
    }
}

/// An asynchronous message from the server.
#[allow(clippy::large_enum_variant)]
pub enum AsyncMessage {
    /// A notice.
    ///
    /// Notices use the same format as errors, but aren't "errors" per-se.
    Notice(DbError),
    /// A notification.
    ///
    /// Connections can subscribe to notifications with the `LISTEN` command.
    Notification(Notification),
    #[doc(hidden)]
    __NonExhaustive,
}

/// Message returned by the `SimpleQuery` stream.
pub enum SimpleQueryMessage {
    /// A row of data.
    Row(SimpleQueryRow),
    /// A statement in the query has completed.
    ///
    /// The number of rows modified or selected is returned.
    CommandComplete(u64),
    #[doc(hidden)]
    __NonExhaustive,
}

fn slice_iter<'a>(
    s: &'a [&'a (dyn ToSql + Sync)],
) -> impl ExactSizeIterator<Item = &'a dyn ToSql> + 'a {
    s.iter().map(|s| *s as _)
}
<|endoftext|>


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

//! Common storage-related types
use memmap;

use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, BufWriter, ErrorKind, Read, Write};
use std::os::unix::io::AsRawFd;
use std::path::Path;

#[cfg(not(any(target_os = "linux", target_os = "android")))]
use libc::{ftruncate as ftruncate64, off_t as off64_t};
#[cfg(any(target_os = "linux"))]
use libc::{ftruncate64, off64_t};

use core::core::hash::Hash;
use core::ser::{self, FixedLength};

/// A no-op function for doing nothing with some pruned data.
pub fn prune_noop(_pruned_data: &[u8]) {}

/// Hash file (MMR) wrapper around an append only file.
pub struct HashFile {
	file: AppendOnlyFile,
}

impl HashFile {
	/// Open (or create) a hash file at the provided path on disk.
	pub fn open(path: &str) -> io::Result<HashFile> {
		let file = AppendOnlyFile::open(path)?;
		Ok(HashFile { file })
	}

	/// Append a hash to this hash file.
	/// Will not be written to disk until flush() is subsequently called.
	/// Alternatively discard() may be called to discard any pending changes.
	pub fn append(&mut self, hash: &Hash) -> io::Result<()> {
		let mut bytes = ser::ser_vec(hash).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
		self.file.append(&mut bytes);
		Ok(())
	}

	/// Read a hash from the hash file by position.
	pub fn read(&self, position: u64) -> Option<Hash> {
		// The MMR starts at 1, our binary backend starts at 0.
		let pos = position - 1;

		// Must be on disk, doing a read at the correct position
		let file_offset = (pos as usize) * Hash::LEN;
		let data = self.file.read(file_offset, Hash::LEN);
		match ser::deserialize(&mut &data[..]) {
			Ok(h) => Some(h),
			Err(e) => {
				error!(
					"Corrupted storage, could not read an entry from hash file: {:?}",
					e
				);
				None
			}
		}
	}

	/// Rewind the backend file to the specified position.
	pub fn rewind(&mut self, position: u64) -> io::Result<()> {
		self.file.rewind(position * Hash::LEN as u64);
		Ok(())
	}

	/// Flush unsynced changes to the hash file to disk.
	pub fn flush(&mut self) -> io::Result<()> {
		self.file.flush()
	}

	/// Discard any unsynced changes to the hash file.
	pub fn discard(&mut self) {
		self.file.discard()
	}

	/// Size of the hash file in number of hashes (not bytes).
	pub fn size(&self) -> u64 {
		self.file.size() / Hash::LEN as u64
	}

	/// Size of the unsync'd hash file, in hashes (not bytes).
	pub fn size_unsync(&self) -> u64 {
		self.file.size_unsync() / Hash::LEN as u64
	}

	/// Rewrite the hash file out to disk, pruning removed hashes.
	pub fn save_prune<T>(&self, target: String, prune_offs: &[u64], prune_cb: T) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		let prune_offs = prune_offs
			.iter()
			.map(|x| x * Hash::LEN as u64)
			.collect::<Vec<_>>();
		self.file
			.save_prune(target, prune_offs.as_slice(), Hash::LEN as u64, prune_cb)
	}
}

/// Wrapper for a file that can be read at any position (random read) but for
/// which writes are append only. Reads are backed by a memory map (mmap(2)),
/// relying on the operating system for fast access and caching. The memory
/// map is reallocated to expand it when new writes are flushed.
///
/// Despite being append-only, the file can still be pruned and truncated. The
/// former simply happens by rewriting it, ignoring some of the data. The
/// latter by truncating the underlying file and re-creating the mmap.
pub struct AppendOnlyFile {
	path: String,
	file: File,
	mmap: Option<memmap::Mmap>,
	buffer_start: usize,
	buffer: Vec<u8>,
	buffer_start_bak: usize,
}

impl AppendOnlyFile {
	/// Open a file (existing or not) as append-only, backed by a mmap.
	pub fn open(path: &str) -> io::Result<AppendOnlyFile> {
		let file = OpenOptions::new()
			.read(true)
			.append(true)
			.create(true)
			.open(&path)?;
		let mut aof = AppendOnlyFile {
			file,
			path: path.to_string(),
			mmap: None,
			buffer_start: 0,
			buffer: vec![],
			buffer_start_bak: 0,
		};
		// If we have a non-empty file then mmap it.
		let sz = aof.size();
		if sz > 0 {
			aof.buffer_start = sz as usize;
			aof.mmap = Some(unsafe { memmap::Mmap::map(&aof.file)? });
		}
		Ok(aof)
	}

	/// Append data to the file. Until the append-only file is synced, data is
	/// only written to memory.
	pub fn append(&mut self, buf: &mut Vec<u8>) {
		self.buffer.append(buf);
	}

	/// Rewinds the data file back to a lower position. The new position needs
	/// to be the one of the first byte the next time data is appended.
	/// Supports two scenarios currently -
	///   * rewind from a clean state (rewinding to handle a forked block)
	///   * rewind within the buffer itself (raw_tx fails to validate)
	/// Note: we do not currently support a rewind() that
	/// crosses the buffer boundary.
	pub fn rewind(&mut self, file_pos: u64) {
		if self.buffer.is_empty() {
			// rewinding from clean state, no buffer, not already rewound anything
			if self.buffer_start_bak == 0 {
				self.buffer_start_bak = self.buffer_start;
			}
			self.buffer_start = file_pos as usize;
		} else {
			// rewinding (within) the buffer
			if self.buffer_start as u64 > file_pos {
				panic!("cannot rewind buffer beyond buffer_start");
			} else {
				let buffer_len = file_pos - self.buffer_start as u64;
				self.buffer.truncate(buffer_len as usize);
			}
		}
	}

	/// Syncs all writes (fsync), reallocating the memory map to make the newly
	/// written data accessible.
	pub fn flush(&mut self) -> io::Result<()> {
		if self.buffer_start_bak > 0 {
			// flushing a rewound state, we need to truncate before applying
			self.truncate(self.buffer_start)?;
			self.buffer_start_bak = 0;
		}

		self.buffer_start += self.buffer.len();
		self.file.write_all(&self.buffer[..])?;
		self.file.sync_all()?;

		self.buffer = vec![];

		// Note: file must be non-empty to memory map it
		if self.file.metadata()?.len() == 0 {
			self.mmap = None;
		} else {
			self.mmap = Some(unsafe { memmap::Mmap::map(&self.file)? });
		}

		Ok(())
	}

	/// Returns the last position (in bytes), taking into account whether data
	/// has been rewound
	pub fn last_buffer_pos(&self) -> usize {
		self.buffer_start
	}

	/// Discard the current non-flushed data.
	pub fn discard(&mut self) {
		if self.buffer_start_bak > 0 {
			// discarding a rewound state, restore the buffer start
			self.buffer_start = self.buffer_start_bak;
			self.buffer_start_bak = 0;
		}
		self.buffer = vec![];
	}

	/// Read length bytes of data at offset from the file.
	/// Leverages the memory map.
	pub fn read(&self, offset: usize, length: usize) -> Vec<u8> {
		if offset >= self.buffer_start {
			let buffer_offset = offset - self.buffer_start;
			return self.read_from_buffer(buffer_offset, length);
		}
		if self.mmap.is_none() {
			return vec![];
		}
		let mmap = self.mmap.as_ref().unwrap();

		if mmap.len() < (offset + length) {
			return vec![];
		}

		(&mmap[offset..(offset + length)]).to_vec()
	}

	// Read length bytes from the buffer, from offset.
	// Return empty vec if we do not have enough bytes in the buffer to read a full
	// vec.
	fn read_from_buffer(&self, offset: usize, length: usize) -> Vec<u8> {
		if self.buffer.len() < (offset + length) {
			vec![]
		} else {
			self.buffer[offset..(offset + length)].to_vec()
		}
	}

	/// Truncates the underlying file to the provided offset
	pub fn truncate(&self, offs: usize) -> io::Result<()> {
		let fd = self.file.as_raw_fd();
		let res = unsafe { ftruncate64(fd, offs as off64_t) };
		if res == -1 {
			Err(io::Error::last_os_error())
		} else {
			Ok(())
		}
	}

	/// Saves a copy of the current file content, skipping data at the provided
	/// prune indices. The prune Vec must be ordered.
	pub fn save_prune<T>(
		&self,
		target: String,
		prune_offs: &[u64],
		prune_len: u64,
		prune_cb: T,
	) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		if prune_offs.is_empty() {
			fs::copy(&self.path, &target)?;
			Ok(())
		} else {
			let mut reader = File::open(&self.path)?;
			let mut writer = BufWriter::new(File::create(&target)?);

			// align the buffer on prune_len to avoid misalignments
			let mut buf = vec![0; (prune_len * 256) as usize];
			let mut read = 0;
			let mut prune_pos = 0;
			loop {
				// fill our buffer
				let len = match reader.read(&mut buf) {
					Ok(0) => return Ok(()),
					Ok(len) => len,
					Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
					Err(e) => return Err(e),
				} as u64;

				// write the buffer, except if we prune offsets in the current span,
				// in which case we skip
				let mut buf_start = 0;
				while prune_offs[prune_pos] >= read && prune_offs[prune_pos] < read + len {
					let prune_at = (prune_offs[prune_pos] - read) as usize;
					if prune_at != buf_start {
						writer.write_all(&buf[buf_start..prune_at])?;
					} else {
						prune_cb(&buf[buf_start..prune_at]);
					}
					buf_start = prune_at + (prune_len as usize);
					if prune_offs.len() > prune_pos + 1 {
						prune_pos += 1;
					} else {
						break;
					}
				}
				writer.write_all(&buf[buf_start..(len as usize)])?;
				read += len;
			}
		}
	}

	/// Current size of the file in bytes.
	pub fn size(&self) -> u64 {
		fs::metadata(&self.path).map(|md| md.len()).unwrap_or(0)
	}

	/// Current size of the (unsynced) file in bytes.
	pub fn size_unsync(&self) -> u64 {
		(self.buffer_start + self.buffer.len()) as u64
	}

	/// Path of the underlying file
	pub fn path(&self) -> String {
		self.path.clone()
	}
}

/// Read an ordered vector of scalars from a file.
pub fn read_ordered_vec<T>(path: &str, elmt_len: usize) -> io::Result<Vec<T>>
where
	T: ser::Readable + cmp::Ord,
{
	let file_path = Path::new(&path);
	let mut ovec = Vec::with_capacity(1000);
	if file_path.exists() {
		let mut file = BufReader::with_capacity(elmt_len * 1000, File::open(&path)?);
		loop {
			// need a block to end mutable borrow before consume
			let buf_len = {
				let buf = file.fill_buf()?;
				if buf.len() == 0 {
					break;
				}
				let elmts_res: Result<Vec<T>, ser::Error> = ser::deserialize(&mut &buf[..]);
				match elmts_res {
					Ok(elmts) => for elmt in elmts {
						if let Err(idx) = ovec.binary_search(&elmt) {
							ovec.insert(idx, elmt);
						}
					},
					Err(_) => {
						return Err(io::Error::new(
							io::ErrorKind::InvalidData,
							format!("Corrupted storage, could not read file at {}", path),
						));
					}
				}
				buf.len()
			};
			file.consume(buf_len);
		}
	}
	Ok(ovec)
}

/// Writes an ordered vector to a file
pub fn write_vec<T>(path: &str, v: &Vec<T>) -> io::Result<()>
where
	T: ser::Writeable,
{
	let mut file_path = File::create(&path)?;
	ser::serialize(&mut file_path, v).map_err(|_| {
		io::Error::new(
			io::ErrorKind::InvalidInput,
			format!("Failed to serialize data when writing to {}", path),
		)
	})?;
	Ok(())
}
<commit_msg>replace ftruncate/ftruncate64 with a simple set_len() call<commit_after>// Copyright 2018 The Grin Developers
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Common storage-related types
use memmap;

use std::cmp;
use std::fs::{self, File, OpenOptions};
use std::io::{self, BufRead, BufReader, BufWriter, ErrorKind, Read, Write};
use std::path::Path;

use core::core::hash::Hash;
use core::ser::{self, FixedLength};

/// A no-op function for doing nothing with some pruned data.
pub fn prune_noop(_pruned_data: &[u8]) {}

/// Hash file (MMR) wrapper around an append only file.
pub struct HashFile {
	file: AppendOnlyFile,
}

impl HashFile {
	/// Open (or create) a hash file at the provided path on disk.
	pub fn open(path: &str) -> io::Result<HashFile> {
		let file = AppendOnlyFile::open(path)?;
		Ok(HashFile { file })
	}

	/// Append a hash to this hash file.
	/// Will not be written to disk until flush() is subsequently called.
	/// Alternatively discard() may be called to discard any pending changes.
	pub fn append(&mut self, hash: &Hash) -> io::Result<()> {
		let mut bytes = ser::ser_vec(hash).map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
		self.file.append(&mut bytes);
		Ok(())
	}

	/// Read a hash from the hash file by position.
	pub fn read(&self, position: u64) -> Option<Hash> {
		// The MMR starts at 1, our binary backend starts at 0.
		let pos = position - 1;

		// Must be on disk, doing a read at the correct position
		let file_offset = (pos as usize) * Hash::LEN;
		let data = self.file.read(file_offset, Hash::LEN);
		match ser::deserialize(&mut &data[..]) {
			Ok(h) => Some(h),
			Err(e) => {
				error!(
					"Corrupted storage, could not read an entry from hash file: {:?}",
					e
				);
				None
			}
		}
	}

	/// Rewind the backend file to the specified position.
	pub fn rewind(&mut self, position: u64) -> io::Result<()> {
		self.file.rewind(position * Hash::LEN as u64);
		Ok(())
	}

	/// Flush unsynced changes to the hash file to disk.
	pub fn flush(&mut self) -> io::Result<()> {
		self.file.flush()
	}

	/// Discard any unsynced changes to the hash file.
	pub fn discard(&mut self) {
		self.file.discard()
	}

	/// Size of the hash file in number of hashes (not bytes).
	pub fn size(&self) -> u64 {
		self.file.size() / Hash::LEN as u64
	}

	/// Size of the unsync'd hash file, in hashes (not bytes).
	pub fn size_unsync(&self) -> u64 {
		self.file.size_unsync() / Hash::LEN as u64
	}

	/// Rewrite the hash file out to disk, pruning removed hashes.
	pub fn save_prune<T>(&self, target: String, prune_offs: &[u64], prune_cb: T) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		let prune_offs = prune_offs
			.iter()
			.map(|x| x * Hash::LEN as u64)
			.collect::<Vec<_>>();
		self.file
			.save_prune(target, prune_offs.as_slice(), Hash::LEN as u64, prune_cb)
	}
}

/// Wrapper for a file that can be read at any position (random read) but for
/// which writes are append only. Reads are backed by a memory map (mmap(2)),
/// relying on the operating system for fast access and caching. The memory
/// map is reallocated to expand it when new writes are flushed.
///
/// Despite being append-only, the file can still be pruned and truncated. The
/// former simply happens by rewriting it, ignoring some of the data. The
/// latter by truncating the underlying file and re-creating the mmap.
pub struct AppendOnlyFile {
	path: String,
	file: File,
	mmap: Option<memmap::Mmap>,
	buffer_start: usize,
	buffer: Vec<u8>,
	buffer_start_bak: usize,
}

impl AppendOnlyFile {
	/// Open a file (existing or not) as append-only, backed by a mmap.
	pub fn open(path: &str) -> io::Result<AppendOnlyFile> {
		let file = OpenOptions::new()
			.read(true)
			.append(true)
			.create(true)
			.open(&path)?;
		let mut aof = AppendOnlyFile {
			file,
			path: path.to_string(),
			mmap: None,
			buffer_start: 0,
			buffer: vec![],
			buffer_start_bak: 0,
		};
		// If we have a non-empty file then mmap it.
		let sz = aof.size();
		if sz > 0 {
			aof.buffer_start = sz as usize;
			aof.mmap = Some(unsafe { memmap::Mmap::map(&aof.file)? });
		}
		Ok(aof)
	}

	/// Append data to the file. Until the append-only file is synced, data is
	/// only written to memory.
	pub fn append(&mut self, buf: &mut Vec<u8>) {
		self.buffer.append(buf);
	}

	/// Rewinds the data file back to a lower position. The new position needs
	/// to be the one of the first byte the next time data is appended.
	/// Supports two scenarios currently -
	///   * rewind from a clean state (rewinding to handle a forked block)
	///   * rewind within the buffer itself (raw_tx fails to validate)
	/// Note: we do not currently support a rewind() that
	/// crosses the buffer boundary.
	pub fn rewind(&mut self, file_pos: u64) {
		if self.buffer.is_empty() {
			// rewinding from clean state, no buffer, not already rewound anything
			if self.buffer_start_bak == 0 {
				self.buffer_start_bak = self.buffer_start;
			}
			self.buffer_start = file_pos as usize;
		} else {
			// rewinding (within) the buffer
			if self.buffer_start as u64 > file_pos {
				panic!("cannot rewind buffer beyond buffer_start");
			} else {
				let buffer_len = file_pos - self.buffer_start as u64;
				self.buffer.truncate(buffer_len as usize);
			}
		}
	}

	/// Syncs all writes (fsync), reallocating the memory map to make the newly
	/// written data accessible.
	pub fn flush(&mut self) -> io::Result<()> {
		if self.buffer_start_bak > 0 {
			// Flushing a rewound state, we need to truncate via set_len() before applying.
			self.file.set_len(self.buffer_start as u64)?;
			self.buffer_start_bak = 0;
		}

		self.buffer_start += self.buffer.len();
		self.file.write_all(&self.buffer[..])?;
		self.file.sync_all()?;

		self.buffer = vec![];

		// Note: file must be non-empty to memory map it
		if self.file.metadata()?.len() == 0 {
			self.mmap = None;
		} else {
			self.mmap = Some(unsafe { memmap::Mmap::map(&self.file)? });
		}

		Ok(())
	}

	/// Returns the last position (in bytes), taking into account whether data
	/// has been rewound
	pub fn last_buffer_pos(&self) -> usize {
		self.buffer_start
	}

	/// Discard the current non-flushed data.
	pub fn discard(&mut self) {
		if self.buffer_start_bak > 0 {
			// discarding a rewound state, restore the buffer start
			self.buffer_start = self.buffer_start_bak;
			self.buffer_start_bak = 0;
		}
		self.buffer = vec![];
	}

	/// Read length bytes of data at offset from the file.
	/// Leverages the memory map.
	pub fn read(&self, offset: usize, length: usize) -> Vec<u8> {
		if offset >= self.buffer_start {
			let buffer_offset = offset - self.buffer_start;
			return self.read_from_buffer(buffer_offset, length);
		}
		if self.mmap.is_none() {
			return vec![];
		}
		let mmap = self.mmap.as_ref().unwrap();

		if mmap.len() < (offset + length) {
			return vec![];
		}

		(&mmap[offset..(offset + length)]).to_vec()
	}

	// Read length bytes from the buffer, from offset.
	// Return empty vec if we do not have enough bytes in the buffer to read a full
	// vec.
	fn read_from_buffer(&self, offset: usize, length: usize) -> Vec<u8> {
		if self.buffer.len() < (offset + length) {
			vec![]
		} else {
			self.buffer[offset..(offset + length)].to_vec()
		}
	}

	// /// Truncates the underlying file to the provided offset
	// pub fn truncate(&self, offs: usize) -> io::Result<()> {
	// 	let fd = self.file.as_raw_fd();
	// 	let res = unsafe { ftruncate64(fd, offs as off64_t) };
	// 	if res == -1 {
	// 		Err(io::Error::last_os_error())
	// 	} else {
	// 		Ok(())
	// 	}
	// }

	/// Saves a copy of the current file content, skipping data at the provided
	/// prune indices. The prune Vec must be ordered.
	pub fn save_prune<T>(
		&self,
		target: String,
		prune_offs: &[u64],
		prune_len: u64,
		prune_cb: T,
	) -> io::Result<()>
	where
		T: Fn(&[u8]),
	{
		if prune_offs.is_empty() {
			fs::copy(&self.path, &target)?;
			Ok(())
		} else {
			let mut reader = File::open(&self.path)?;
			let mut writer = BufWriter::new(File::create(&target)?);

			// align the buffer on prune_len to avoid misalignments
			let mut buf = vec![0; (prune_len * 256) as usize];
			let mut read = 0;
			let mut prune_pos = 0;
			loop {
				// fill our buffer
				let len = match reader.read(&mut buf) {
					Ok(0) => return Ok(()),
					Ok(len) => len,
					Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
					Err(e) => return Err(e),
				} as u64;

				// write the buffer, except if we prune offsets in the current span,
				// in which case we skip
				let mut buf_start = 0;
				while prune_offs[prune_pos] >= read && prune_offs[prune_pos] < read + len {
					let prune_at = (prune_offs[prune_pos] - read) as usize;
					if prune_at != buf_start {
						writer.write_all(&buf[buf_start..prune_at])?;
					} else {
						prune_cb(&buf[buf_start..prune_at]);
					}
					buf_start = prune_at + (prune_len as usize);
					if prune_offs.len() > prune_pos + 1 {
						prune_pos += 1;
					} else {
						break;
					}
				}
				writer.write_all(&buf[buf_start..(len as usize)])?;
				read += len;
			}
		}
	}

	/// Current size of the file in bytes.
	pub fn size(&self) -> u64 {
		fs::metadata(&self.path).map(|md| md.len()).unwrap_or(0)
	}

	/// Current size of the (unsynced) file in bytes.
	pub fn size_unsync(&self) -> u64 {
		(self.buffer_start + self.buffer.len()) as u64
	}

	/// Path of the underlying file
	pub fn path(&self) -> String {
		self.path.clone()
	}
}

/// Read an ordered vector of scalars from a file.
pub fn read_ordered_vec<T>(path: &str, elmt_len: usize) -> io::Result<Vec<T>>
where
	T: ser::Readable + cmp::Ord,
{
	let file_path = Path::new(&path);
	let mut ovec = Vec::with_capacity(1000);
	if file_path.exists() {
		let mut file = BufReader::with_capacity(elmt_len * 1000, File::open(&path)?);
		loop {
			// need a block to end mutable borrow before consume
			let buf_len = {
				let buf = file.fill_buf()?;
				if buf.len() == 0 {
					break;
				}
				let elmts_res: Result<Vec<T>, ser::Error> = ser::deserialize(&mut &buf[..]);
				match elmts_res {
					Ok(elmts) => for elmt in elmts {
						if let Err(idx) = ovec.binary_search(&elmt) {
							ovec.insert(idx, elmt);
						}
					},
					Err(_) => {
						return Err(io::Error::new(
							io::ErrorKind::InvalidData,
							format!("Corrupted storage, could not read file at {}", path),
						));
					}
				}
				buf.len()
			};
			file.consume(buf_len);
		}
	}
	Ok(ovec)
}

/// Writes an ordered vector to a file
pub fn write_vec<T>(path: &str, v: &Vec<T>) -> io::Result<()>
where
	T: ser::Writeable,
{
	let mut file_path = File::create(&path)?;
	ser::serialize(&mut file_path, v).map_err(|_| {
		io::Error::new(
			io::ErrorKind::InvalidInput,
			format!("Failed to serialize data when writing to {}", path),
		)
	})?;
	Ok(())
}
<|endoftext|>