<commit_before>//! Errors, type aliases, and functions related to working with `Result`.

use std::convert::From;
use std::error::Error as StdError;
use std::ffi::NulError;
use std::fmt::{self, Display};

#[derive(Debug)]
#[allow(clippy::enum_variant_names)]
/// Represents all the ways that a query can fail.
///
/// This type is not intended to be exhaustively matched, and new variants may
/// be added in the future without a major version bump.
#[non_exhaustive]
pub enum Error {
    /// The query contained a nul byte.
    ///
    /// This should never occur in normal usage.
    InvalidCString(NulError),

    /// The database returned an error.
    ///
    /// While Diesel prevents almost all sources of runtime errors at compile
    /// time, it does not attempt to prevent 100% of them. Typically this error
    /// will occur from insert or update statements due to a constraint
    /// violation.
    DatabaseError(
        DatabaseErrorKind,
        Box<dyn DatabaseErrorInformation + Send + Sync>,
    ),

    /// No rows were returned by a query expected to return at least one row.
    ///
    /// This variant is only returned by [`get_result`] and [`first`]. [`load`]
    /// does not treat 0 rows as an error. If you would like to allow either 0
    /// or 1 rows, call [`optional`] on the result.
    ///
    /// [`get_result`]: crate::query_dsl::RunQueryDsl::get_result()
    /// [`first`]: crate::query_dsl::RunQueryDsl::first()
    /// [`load`]: crate::query_dsl::RunQueryDsl::load()
    /// [`optional`]: OptionalExtension::optional
    NotFound,

    /// The query could not be constructed
    ///
    /// An example of when this error could occur is if you are attempting to
    /// construct an update statement with no changes (e.g. all fields on the
    /// struct are `None`).
    QueryBuilderError(Box<dyn StdError + Send + Sync>),

    /// An error occurred deserializing the data being sent to the database.
    ///
    /// Typically this error means that the stated type of the query is
    /// incorrect. An example of when this error might occur in normal usage is
    /// attempting to deserialize an infinite date into chrono.
    DeserializationError(Box<dyn StdError + Send + Sync>),

    /// An error occurred serializing the data being sent to the database.
    ///
    /// An example of when this error would be returned is if you attempted to
    /// serialize a `chrono::NaiveDate` earlier than the earliest date supported
    /// by PostgreSQL.
    SerializationError(Box<dyn StdError + Send + Sync>),

    /// An error occurred during the rollback of a transaction.
    ///
    /// An example of when this error would be returned is if a rollback has
    /// already be called on the current transaction.
    RollbackError(Box<Error>),

    /// Roll back the current transaction.
    ///
    /// You can return this variant inside of a transaction when you want to
    /// roll it back, but have no actual error to return. Diesel will never
    /// return this variant unless you gave it to us, and it can be safely
    /// ignored in error handling.
    RollbackTransaction,

    /// Attempted to perform an operation that cannot be done inside a transaction
    /// when a transaction was already open.
    AlreadyInTransaction,
}

#[derive(Debug, Clone, Copy)]
/// The kind of database error that occurred.
///
/// This is not meant to exhaustively cover all possible errors, but is used to
/// identify errors which are commonly recovered from programmatically. This enum
/// is not intended to be exhaustively matched, and new variants may be added in
/// the future without a major version bump.
#[non_exhaustive]
pub enum DatabaseErrorKind {
    /// A unique constraint was violated.
    UniqueViolation,

    /// A foreign key constraint was violated.
    ForeignKeyViolation,

    /// The query could not be sent to the database due to a protocol violation.
    ///
    /// An example of a case where this would occur is if you attempted to send
    /// a query with more than 65000 bind parameters using PostgreSQL.
    UnableToSendCommand,

    /// A serializable transaction failed to commit due to a read/write
    /// dependency on a concurrent transaction.
    ///
    /// Corresponds to SQLSTATE code 40001
    ///
    /// This error is only detected for PostgreSQL, as we do not yet support
    /// transaction isolation levels for other backends.
    SerializationFailure,

    /// The command could not be completed because the transaction was read
    /// only.
    ///
    /// This error will also be returned for `SELECT` statements which attempted
    /// to lock the rows.
    ReadOnlyTransaction,

    /// A not null constraint was violated.
    NotNullViolation,

    /// A check constraint was violated.
    CheckViolation,

    /// The connection to the server was unexpectedly closed.
    ClosedConnection,

    #[doc(hidden)]
    Unknown, // Match against _ instead, more variants may be added in the future
}

/// Information about an error that was returned by the database.
pub trait DatabaseErrorInformation {
    /// The primary human-readable error message. Typically one line.
    fn message(&self) -> &str;

    /// An optional secondary error message providing more details about the
    /// problem, if it was provided by the database. Might span multiple lines.
    fn details(&self) -> Option<&str>;

    /// An optional suggestion of what to do about the problem, if one was
    /// provided by the database.
    fn hint(&self) -> Option<&str>;

    /// The name of the table the error was associated with, if the error was
    /// associated with a specific table and the backend supports retrieving
    /// that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn table_name(&self) -> Option<&str>;

    /// The name of the column the error was associated with, if the error was
    /// associated with a specific column and the backend supports retrieving
    /// that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn column_name(&self) -> Option<&str>;

    /// The constraint that was violated if this error is a constraint violation
    /// and the backend supports retrieving that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn constraint_name(&self) -> Option<&str>;

    /// An optional integer indicating an error cursor position as an index into
    /// the original statement string.
    fn statement_position(&self) -> Option<i32>;
}

impl fmt::Debug for dyn DatabaseErrorInformation + Send + Sync {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.message(), f)
    }
}

impl DatabaseErrorInformation for String {
    fn message(&self) -> &str {
        self
    }
    fn details(&self) -> Option<&str> {
        None
    }
    fn hint(&self) -> Option<&str> {
        None
    }
    fn table_name(&self) -> Option<&str> {
        None
    }
    fn column_name(&self) -> Option<&str> {
        None
    }
    fn constraint_name(&self) -> Option<&str> {
        None
    }
    fn statement_position(&self) -> Option<i32> {
        None
    }
}

/// Errors which can occur during [`Connection::establish`]
///
/// [`Connection::establish`]: crate::connection::Connection::establish
#[derive(Debug, PartialEq)]
#[non_exhaustive]
pub enum ConnectionError {
    /// The connection URL contained a `NUL` byte.
    InvalidCString(NulError),
    /// The database returned an error.
    BadConnection(String),
    /// The connection URL could not be parsed.
    InvalidConnectionUrl(String),
    /// Diesel could not configure the database connection.
    ///
    /// Diesel may try to automatically set session specific configuration
    /// values, such as UTF8 encoding, or enabling the `||` operator on MySQL.
    /// This variant is returned if an error occurred executing the query to set
    /// those options. Diesel will never affect global configuration.
    CouldntSetupConfiguration(Error),
}

/// A specialized result type for queries.
///
/// This type is exported by `diesel::prelude`, and is generally used by any
/// code which is interacting with Diesel. This type exists to avoid writing out
/// `diesel::result::Error`, and is otherwise a direct mapping to `Result`.
pub type QueryResult<T> = Result<T, Error>;

/// A specialized result type for establishing connections.
///
/// This type exists to avoid writing out `diesel::result::ConnectionError`, and
/// is otherwise a direct mapping to `Result`.
pub type ConnectionResult<T> = Result<T, ConnectionError>;

/// See the [method documentation](OptionalExtension::optional).
pub trait OptionalExtension<T> {
    /// Converts a `QueryResult<T>` into a `QueryResult<Option<T>>`.
    ///
    /// By default, Diesel treats 0 rows being returned from a query that is expected to return 1
    /// row as an error (e.g. the return value of [`get_result`] or [`first`]). This method will
    /// handle that error, and give you back an `Option<T>` instead.
    ///
    /// [`get_result`]: crate::query_dsl::RunQueryDsl::get_result()
    /// [`first`]: crate::query_dsl::RunQueryDsl::first()
    ///
    /// # Example
    ///
    /// ```rust
    /// use diesel::{QueryResult, NotFound, OptionalExtension};
    ///
    /// let result: QueryResult<i32> = Ok(1);
    /// assert_eq!(Ok(Some(1)), result.optional());
    ///
    /// let result: QueryResult<i32> = Err(NotFound);
    /// assert_eq!(Ok(None), result.optional());
    /// ```
    fn optional(self) -> Result<Option<T>, Error>;
}

impl<T> OptionalExtension<T> for QueryResult<T> {
    fn optional(self) -> Result<Option<T>, Error> {
        match self {
            Ok(value) => Ok(Some(value)),
            Err(Error::NotFound) => Ok(None),
            Err(e) => Err(e),
        }
    }
}

impl From<NulError> for ConnectionError {
    fn from(e: NulError) -> Self {
        ConnectionError::InvalidCString(e)
    }
}

impl From<NulError> for Error {
    fn from(e: NulError) -> Self {
        Error::InvalidCString(e)
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::InvalidCString(ref nul_err) => write!(f, "{}", nul_err),
            Error::DatabaseError(_, ref e) => write!(f, "{}", e.message()),
            Error::NotFound => f.write_str("Record not found"),
            Error::QueryBuilderError(ref e) => e.fmt(f),
            Error::DeserializationError(ref e) => e.fmt(f),
            Error::SerializationError(ref e) => e.fmt(f),
            Error::RollbackError(ref e) => e.fmt(f),
            Error::RollbackTransaction => write!(f, "The current transaction was aborted"),
            Error::AlreadyInTransaction => write!(
                f,
                "Cannot perform this operation while a transaction is open",
            ),
        }
    }
}

impl StdError for Error {
    fn cause(&self) -> Option<&dyn StdError> {
        match *self {
            Error::InvalidCString(ref e) => Some(e),
            Error::QueryBuilderError(ref e) => Some(&**e),
            Error::DeserializationError(ref e) => Some(&**e),
            Error::SerializationError(ref e) => Some(&**e),
            _ => None,
        }
    }
}

impl Display for ConnectionError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConnectionError::InvalidCString(ref nul_err) => nul_err.fmt(f),
            ConnectionError::BadConnection(ref s) => write!(f, "{}", s),
            ConnectionError::InvalidConnectionUrl(ref s) => write!(f, "{}", s),
            ConnectionError::CouldntSetupConfiguration(ref e) => e.fmt(f),
        }
    }
}

impl StdError for ConnectionError {
    fn cause(&self) -> Option<&dyn StdError> {
        match *self {
            ConnectionError::InvalidCString(ref e) => Some(e),
            ConnectionError::CouldntSetupConfiguration(ref e) => Some(e),
            _ => None,
        }
    }
}

impl PartialEq for Error {
    fn eq(&self, other: &Error) -> bool {
        match (self, other) {
            (&Error::InvalidCString(ref a), &Error::InvalidCString(ref b)) => a == b,
            (&Error::DatabaseError(_, ref a), &Error::DatabaseError(_, ref b)) => {
                a.message() == b.message()
            }
            (&Error::NotFound, &Error::NotFound) => true,
            (&Error::RollbackTransaction, &Error::RollbackTransaction) => true,
            (&Error::AlreadyInTransaction, &Error::AlreadyInTransaction) => true,
            _ => false,
        }
    }
}

#[cfg(test)]
#[allow(warnings)]
fn error_impls_send() {
    let err: Error = unimplemented!();
    let x: &Send = &err;
}

pub(crate) fn first_or_not_found<T>(records: QueryResult<Vec<T>>) -> QueryResult<T> {
    records?.into_iter().next().ok_or(Error::NotFound)
}

/// An unexpected `NULL` was encountered during deserialization
#[derive(Debug, Clone, Copy)]
pub struct UnexpectedNullError;

impl fmt::Display for UnexpectedNullError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Unexpected null for non-null column")
    }
}

impl StdError for UnexpectedNullError {}

/// Expected more fields then present in the current row while deserialising results
#[derive(Debug, Clone, Copy)]
pub struct UnexpectedEndOfRow;

impl fmt::Display for UnexpectedEndOfRow {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Unexpected end of row")
    }
}

impl StdError for UnexpectedEndOfRow {}
<commit_msg>add more comments<commit_after>//! Errors, type aliases, and functions related to working with `Result`.

use std::convert::From;
use std::error::Error as StdError;
use std::ffi::NulError;
use std::fmt::{self, Display};

#[derive(Debug)]
#[allow(clippy::enum_variant_names)]
/// Represents all the ways that a query can fail.
///
/// This type is not intended to be exhaustively matched, and new variants may
/// be added in the future without a major version bump.
#[non_exhaustive]
pub enum Error {
    /// The query contained a nul byte.
    ///
    /// This should never occur in normal usage.
    InvalidCString(NulError),

    /// The database returned an error.
    ///
    /// While Diesel prevents almost all sources of runtime errors at compile
    /// time, it does not attempt to prevent 100% of them. Typically this error
    /// will occur from insert or update statements due to a constraint
    /// violation.
    DatabaseError(
        DatabaseErrorKind,
        Box<dyn DatabaseErrorInformation + Send + Sync>,
    ),

    /// No rows were returned by a query expected to return at least one row.
    ///
    /// This variant is only returned by [`get_result`] and [`first`]. [`load`]
    /// does not treat 0 rows as an error. If you would like to allow either 0
    /// or 1 rows, call [`optional`] on the result.
    ///
    /// [`get_result`]: crate::query_dsl::RunQueryDsl::get_result()
    /// [`first`]: crate::query_dsl::RunQueryDsl::first()
    /// [`load`]: crate::query_dsl::RunQueryDsl::load()
    /// [`optional`]: OptionalExtension::optional
    NotFound,

    /// The query could not be constructed
    ///
    /// An example of when this error could occur is if you are attempting to
    /// construct an update statement with no changes (e.g. all fields on the
    /// struct are `None`).
    QueryBuilderError(Box<dyn StdError + Send + Sync>),

    /// An error occurred deserializing the data being sent to the database.
    ///
    /// Typically this error means that the stated type of the query is
    /// incorrect. An example of when this error might occur in normal usage is
    /// attempting to deserialize an infinite date into chrono.
    DeserializationError(Box<dyn StdError + Send + Sync>),

    /// An error occurred serializing the data being sent to the database.
    ///
    /// An example of when this error would be returned is if you attempted to
    /// serialize a `chrono::NaiveDate` earlier than the earliest date supported
    /// by PostgreSQL.
    SerializationError(Box<dyn StdError + Send + Sync>),

    /// An error occurred during the rollback of a transaction.
    ///
    /// An example of when this error would be returned is if a rollback has
    /// already be called on the current transaction.
    RollbackError(Box<Error>),

    /// Roll back the current transaction.
    ///
    /// You can return this variant inside of a transaction when you want to
    /// roll it back, but have no actual error to return. Diesel will never
    /// return this variant unless you gave it to us, and it can be safely
    /// ignored in error handling.
    RollbackTransaction,

    /// Attempted to perform an operation that cannot be done inside a transaction
    /// when a transaction was already open.
    AlreadyInTransaction,
}

#[derive(Debug, Clone, Copy)]
/// The kind of database error that occurred.
///
/// This is not meant to exhaustively cover all possible errors, but is used to
/// identify errors which are commonly recovered from programmatically. This enum
/// is not intended to be exhaustively matched, and new variants may be added in
/// the future without a major version bump.
#[non_exhaustive]
pub enum DatabaseErrorKind {
    /// A unique constraint was violated.
    UniqueViolation,

    /// A foreign key constraint was violated.
    ForeignKeyViolation,

    /// The query could not be sent to the database due to a protocol violation.
    ///
    /// An example of a case where this would occur is if you attempted to send
    /// a query with more than 65000 bind parameters using PostgreSQL.
    UnableToSendCommand,

    /// A serializable transaction failed to commit due to a read/write
    /// dependency on a concurrent transaction.
    ///
    /// Corresponds to SQLSTATE code 40001
    ///
    /// This error is only detected for PostgreSQL, as we do not yet support
    /// transaction isolation levels for other backends.
    SerializationFailure,

    /// The command could not be completed because the transaction was read
    /// only.
    ///
    /// This error will also be returned for `SELECT` statements which attempted
    /// to lock the rows.
    ReadOnlyTransaction,

    /// A not null constraint was violated.
    NotNullViolation,

    /// A check constraint was violated.
    CheckViolation,

    /// The connection to the server was unexpectedly closed.
    ///
    /// This error is only detected for PostgreSQL and is emitted on a best-effort basis
    /// and may be missed.
    ClosedConnection,

    #[doc(hidden)]
    Unknown, // Match against _ instead, more variants may be added in the future
}

/// Information about an error that was returned by the database.
pub trait DatabaseErrorInformation {
    /// The primary human-readable error message. Typically one line.
    fn message(&self) -> &str;

    /// An optional secondary error message providing more details about the
    /// problem, if it was provided by the database. Might span multiple lines.
    fn details(&self) -> Option<&str>;

    /// An optional suggestion of what to do about the problem, if one was
    /// provided by the database.
    fn hint(&self) -> Option<&str>;

    /// The name of the table the error was associated with, if the error was
    /// associated with a specific table and the backend supports retrieving
    /// that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn table_name(&self) -> Option<&str>;

    /// The name of the column the error was associated with, if the error was
    /// associated with a specific column and the backend supports retrieving
    /// that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn column_name(&self) -> Option<&str>;

    /// The constraint that was violated if this error is a constraint violation
    /// and the backend supports retrieving that information.
    ///
    /// Currently this method will return `None` for all backends other than
    /// PostgreSQL.
    fn constraint_name(&self) -> Option<&str>;

    /// An optional integer indicating an error cursor position as an index into
    /// the original statement string.
    fn statement_position(&self) -> Option<i32>;
}

impl fmt::Debug for dyn DatabaseErrorInformation + Send + Sync {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.message(), f)
    }
}

impl DatabaseErrorInformation for String {
    fn message(&self) -> &str {
        self
    }
    fn details(&self) -> Option<&str> {
        None
    }
    fn hint(&self) -> Option<&str> {
        None
    }
    fn table_name(&self) -> Option<&str> {
        None
    }
    fn column_name(&self) -> Option<&str> {
        None
    }
    fn constraint_name(&self) -> Option<&str> {
        None
    }
    fn statement_position(&self) -> Option<i32> {
        None
    }
}

/// Errors which can occur during [`Connection::establish`]
///
/// [`Connection::establish`]: crate::connection::Connection::establish
#[derive(Debug, PartialEq)]
#[non_exhaustive]
pub enum ConnectionError {
    /// The connection URL contained a `NUL` byte.
    InvalidCString(NulError),
    /// The database returned an error.
    BadConnection(String),
    /// The connection URL could not be parsed.
    InvalidConnectionUrl(String),
    /// Diesel could not configure the database connection.
    ///
    /// Diesel may try to automatically set session specific configuration
    /// values, such as UTF8 encoding, or enabling the `||` operator on MySQL.
    /// This variant is returned if an error occurred executing the query to set
    /// those options. Diesel will never affect global configuration.
    CouldntSetupConfiguration(Error),
}

/// A specialized result type for queries.
///
/// This type is exported by `diesel::prelude`, and is generally used by any
/// code which is interacting with Diesel. This type exists to avoid writing out
/// `diesel::result::Error`, and is otherwise a direct mapping to `Result`.
pub type QueryResult<T> = Result<T, Error>;

/// A specialized result type for establishing connections.
///
/// This type exists to avoid writing out `diesel::result::ConnectionError`, and
/// is otherwise a direct mapping to `Result`.
pub type ConnectionResult<T> = Result<T, ConnectionError>;

/// See the [method documentation](OptionalExtension::optional).
pub trait OptionalExtension<T> {
    /// Converts a `QueryResult<T>` into a `QueryResult<Option<T>>`.
    ///
    /// By default, Diesel treats 0 rows being returned from a query that is expected to return 1
    /// row as an error (e.g. the return value of [`get_result`] or [`first`]). This method will
    /// handle that error, and give you back an `Option<T>` instead.
    ///
    /// [`get_result`]: crate::query_dsl::RunQueryDsl::get_result()
    /// [`first`]: crate::query_dsl::RunQueryDsl::first()
    ///
    /// # Example
    ///
    /// ```rust
    /// use diesel::{QueryResult, NotFound, OptionalExtension};
    ///
    /// let result: QueryResult<i32> = Ok(1);
    /// assert_eq!(Ok(Some(1)), result.optional());
    ///
    /// let result: QueryResult<i32> = Err(NotFound);
    /// assert_eq!(Ok(None), result.optional());
    /// ```
    fn optional(self) -> Result<Option<T>, Error>;
}

impl<T> OptionalExtension<T> for QueryResult<T> {
    fn optional(self) -> Result<Option<T>, Error> {
        match self {
            Ok(value) => Ok(Some(value)),
            Err(Error::NotFound) => Ok(None),
            Err(e) => Err(e),
        }
    }
}

impl From<NulError> for ConnectionError {
    fn from(e: NulError) -> Self {
        ConnectionError::InvalidCString(e)
    }
}

impl From<NulError> for Error {
    fn from(e: NulError) -> Self {
        Error::InvalidCString(e)
    }
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::InvalidCString(ref nul_err) => write!(f, "{}", nul_err),
            Error::DatabaseError(_, ref e) => write!(f, "{}", e.message()),
            Error::NotFound => f.write_str("Record not found"),
            Error::QueryBuilderError(ref e) => e.fmt(f),
            Error::DeserializationError(ref e) => e.fmt(f),
            Error::SerializationError(ref e) => e.fmt(f),
            Error::RollbackError(ref e) => e.fmt(f),
            Error::RollbackTransaction => write!(f, "The current transaction was aborted"),
            Error::AlreadyInTransaction => write!(
                f,
                "Cannot perform this operation while a transaction is open",
            ),
        }
    }
}

impl StdError for Error {
    fn cause(&self) -> Option<&dyn StdError> {
        match *self {
            Error::InvalidCString(ref e) => Some(e),
            Error::QueryBuilderError(ref e) => Some(&**e),
            Error::DeserializationError(ref e) => Some(&**e),
            Error::SerializationError(ref e) => Some(&**e),
            _ => None,
        }
    }
}

impl Display for ConnectionError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ConnectionError::InvalidCString(ref nul_err) => nul_err.fmt(f),
            ConnectionError::BadConnection(ref s) => write!(f, "{}", s),
            ConnectionError::InvalidConnectionUrl(ref s) => write!(f, "{}", s),
            ConnectionError::CouldntSetupConfiguration(ref e) => e.fmt(f),
        }
    }
}

impl StdError for ConnectionError {
    fn cause(&self) -> Option<&dyn StdError> {
        match *self {
            ConnectionError::InvalidCString(ref e) => Some(e),
            ConnectionError::CouldntSetupConfiguration(ref e) => Some(e),
            _ => None,
        }
    }
}

impl PartialEq for Error {
    fn eq(&self, other: &Error) -> bool {
        match (self, other) {
            (&Error::InvalidCString(ref a), &Error::InvalidCString(ref b)) => a == b,
            (&Error::DatabaseError(_, ref a), &Error::DatabaseError(_, ref b)) => {
                a.message() == b.message()
            }
            (&Error::NotFound, &Error::NotFound) => true,
            (&Error::RollbackTransaction, &Error::RollbackTransaction) => true,
            (&Error::AlreadyInTransaction, &Error::AlreadyInTransaction) => true,
            _ => false,
        }
    }
}

#[cfg(test)]
#[allow(warnings)]
fn error_impls_send() {
    let err: Error = unimplemented!();
    let x: &Send = &err;
}

pub(crate) fn first_or_not_found<T>(records: QueryResult<Vec<T>>) -> QueryResult<T> {
    records?.into_iter().next().ok_or(Error::NotFound)
}

/// An unexpected `NULL` was encountered during deserialization
#[derive(Debug, Clone, Copy)]
pub struct UnexpectedNullError;

impl fmt::Display for UnexpectedNullError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Unexpected null for non-null column")
    }
}

impl StdError for UnexpectedNullError {}

/// Expected more fields then present in the current row while deserialising results
#[derive(Debug, Clone, Copy)]
pub struct UnexpectedEndOfRow;

impl fmt::Display for UnexpectedEndOfRow {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Unexpected end of row")
    }
}

impl StdError for UnexpectedEndOfRow {}
<|endoftext|>


<commit_before>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::marker::PhantomData;
use std::mem;
use std::ptr::null_mut;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread::Builder as ThreadBuilder;
use std::thread::JoinHandle;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use base::error;
use base::info;
use base::warn;
use base::Event;
use base::Tube;
use euclid::size2;
use sync::Mutex;
#[cfg(feature = "kiwi")]
use vm_control::ServiceSendToGpu;
use win_util::win32_string;
use winapi::shared::minwindef::LPARAM;
use winapi::shared::minwindef::LRESULT;
use winapi::shared::minwindef::UINT;
use winapi::shared::minwindef::WPARAM;
use winapi::shared::windef::HWND;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::processthreadsapi::GetCurrentThreadId;
use winapi::um::winuser::*;

#[cfg(feature = "kiwi")]
use super::message_relay_thread::MessageRelayThread;
use super::message_relay_thread::MessageRelayThreadTrait;
use super::thread_message_util;
use super::window::MessagePacket;
use super::window::Window;
use super::window_message_dispatcher::WindowMessageDispatcher;
use super::window_message_dispatcher::DISPATCHER_PROPERTY_NAME;
use super::window_message_processor::*;

// The default app icon id, which is defined in crosvm-manifest.rc.
const APP_ICON_ID: u16 = 1;

// 0 is not a valid thread ID: https://devblogs.microsoft.com/oldnewthing/20040223-00/?p=40503
const INVALID_THREAD_ID: u32 = 0;

#[derive(Debug)]
enum MessageLoopState {
    /// The initial state.
    NotStarted = 0,
    /// The loop is running normally.
    Running,
    /// The loop has ended normally.
    ExitedNormally,
    /// The loop never started because errors occurred.
    EarlyTerminatedWithError,
    /// The loop has ended because errors occurred.
    ExitedWithError,
}

/// This class runs the WndProc thread, and provides helper functions for other threads to
/// communicate with it.
pub struct WindowProcedureThread<T: HandleWindowMessage> {
    thread: Option<JoinHandle<()>>,
    thread_id: u32,
    message_loop_state: Option<Arc<AtomicI32>>,
    thread_terminated_event: Event,
    _marker: PhantomData<T>,
}

impl<T: HandleWindowMessage> WindowProcedureThread<T> {
    pub fn start_thread(vm_tube: Option<Arc<Mutex<Tube>>>) -> Result<Self> {
        let thread_id = Arc::new(AtomicU32::new(INVALID_THREAD_ID));
        let message_loop_state = Arc::new(AtomicI32::new(MessageLoopState::NotStarted as i32));
        let thread_terminated_event = Event::new().unwrap();

        let thread_id_clone = Arc::clone(&thread_id);
        let message_loop_state_clone = Arc::clone(&message_loop_state);
        let thread_terminated_event_clone = thread_terminated_event
            .try_clone()
            .map_err(|e| anyhow!("Failed to clone thread_terminated_event: {}", e))?;

        match ThreadBuilder::new()
            .name("gpu_display_wndproc".into())
            .spawn(move || {
                // Safe because GetCurrentThreadId has no failure mode.
                let thread_id = unsafe { GetCurrentThreadId() };

                // Must be called before any other threads post messages to the WndProc thread.
                thread_message_util::force_create_message_queue();
                thread_id_clone.store(thread_id, Ordering::SeqCst);
                drop(thread_id_clone);

                // Safe because the message queue has been created, and the returned thread will go
                // out of scope and get dropped before the WndProc thread exits.
                let _message_relay_thread = unsafe {
                    vm_tube.and_then(|tube| Self::start_message_relay_thread(tube, thread_id))
                };

                Self::run_message_loop(message_loop_state_clone);

                if let Err(e) = thread_terminated_event_clone.write(1) {
                    error!("Failed to write to thread terminated event: {}", e);
                }
            }) {
            Ok(thread) => {
                // TODO(b/243184256): Use `Condvar` to avoid busy-waiting on the atomic variable.
                while message_loop_state.load(Ordering::SeqCst)
                    == MessageLoopState::NotStarted as i32
                {}

                let thread_id = Arc::try_unwrap(thread_id).unwrap().into_inner();
                if thread_id == INVALID_THREAD_ID {
                    bail!("Failed to retrieve thread ID when spawning WndProc thread!");
                }
                Ok(Self {
                    thread: Some(thread),
                    thread_id,
                    message_loop_state: Some(message_loop_state),
                    thread_terminated_event,
                    _marker: PhantomData,
                })
            }
            Err(e) => bail!("Failed to spawn WndProc thread: {:?}", e),
        }
    }

    pub fn try_clone_thread_terminated_event(&self) -> Result<Event> {
        self.thread_terminated_event
            .try_clone()
            .map_err(|e| anyhow!("Failed to clone thread_terminated_event: {}", e))
    }

    pub fn post_display_command(&self, message: DisplaySendToWndProc<T>) -> Result<()> {
        if !self.is_message_loop_running() {
            bail!("Host window has been destroyed!");
        }

        // Safe because the WndProc thread is still running the message loop.
        unsafe {
            thread_message_util::post_message_carrying_object(
                self.thread_id,
                WM_USER_HANDLE_DISPLAY_MESSAGE_INTERNAL,
                message,
            )
            .context("When posting DisplaySendToWndProc message")
        }
    }

    fn run_message_loop(message_loop_state: Arc<AtomicI32>) {
        // Safe because the dispatcher will take care of the lifetime of the `Window` object.
        let mut dispatcher = match unsafe {
            Self::create_window().and_then(|window| WindowMessageDispatcher::<T>::create(window))
        } {
            Ok(dispatcher) => dispatcher,
            Err(e) => {
                error!("WndProc thread didn't enter message loop: {:?}", e);
                message_loop_state.store(
                    MessageLoopState::EarlyTerminatedWithError as i32,
                    Ordering::SeqCst,
                );
                return;
            }
        };

        info!("WndProc thread entering message loop");
        message_loop_state.store(MessageLoopState::Running as i32, Ordering::SeqCst);
        loop {
            let mut message = mem::MaybeUninit::uninit();
            // Safe because we know the lifetime of `message`.
            match unsafe { GetMessageW(message.as_mut_ptr(), null_mut(), 0, 0) } {
                0 => {
                    info!("WndProc thread exiting message loop since WM_QUIT is received");
                    message_loop_state
                        .store(MessageLoopState::ExitedNormally as i32, Ordering::SeqCst);
                    break;
                }
                -1 => {
                    error!(
                        "WndProc thread exiting message loop because GetMessageW() failed with \
                        error code {}",
                        unsafe { GetLastError() }
                    );
                    message_loop_state
                        .store(MessageLoopState::ExitedWithError as i32, Ordering::SeqCst);
                    break;
                }
                _ => (),
            }

            // Safe because `GetMessageW()` will block until `message` is populated.
            let new_message = unsafe { message.assume_init() };
            if new_message.hwnd.is_null() {
                // Thread messages don't target a specific window and `DispatchMessageA()` won't
                // send them to `wnd_proc()` function, hence we need to handle it as a special case.
                dispatcher
                    .as_mut()
                    .process_thread_message(&new_message.into());
            } else {
                // Safe because `GetMessageW()` will block until `message` is populated.
                unsafe {
                    TranslateMessage(&new_message);
                    DispatchMessageW(&new_message);
                }
            }
        }
    }

    fn is_message_loop_running(&self) -> bool {
        self.message_loop_state.as_ref().map_or(false, |state| {
            state.load(Ordering::SeqCst) == MessageLoopState::Running as i32
        })
    }

    /// In the normal case, when all windows are closed by the user, the WndProc thread exits the
    /// message loop and terminates naturally. If we have to shutdown the VM before all windows are
    /// closed because of errors, this function will post a message to let the WndProc thread kill
    /// all windows and terminate.
    fn signal_exit_message_loop_if_needed(&self) {
        if !self.is_message_loop_running() {
            return;
        }

        info!("WndProc thread is still in message loop before dropping. Signaling killing windows");
        // Safe because the WndProc thread is still running the message loop.
        if let Err(e) = unsafe {
            thread_message_util::post_message(
                self.thread_id,
                WM_USER_WNDPROC_THREAD_DROP_KILL_WINDOW_INTERNAL,
                /* w_param */ 0,
                /* l_param */ 0,
            )
        } {
            error!("Failed to signal WndProc thread to kill windows: {:?}", e);
        }
    }

    /// Checks if the message loop has exited normally. This should be called after joining with the
    /// WndProc thread.
    fn check_message_loop_final_state(&mut self) {
        match Arc::try_unwrap(self.message_loop_state.take().unwrap()) {
            Ok(state) => {
                let state = state.into_inner();
                if state == MessageLoopState::ExitedNormally as i32 {
                    info!("WndProc thread exited gracefully");
                } else {
                    warn!("WndProc thread exited with message loop state: {:?}", state);
                }
            }
            Err(e) => error!(
                "WndProc thread exited but message loop state retrieval failed: {:?}",
                e
            ),
        }
    }

    /// # Safety
    /// The owner of the returned `Window` object is responsible for dropping it before we finish
    /// processing `WM_NCDESTROY`, because the window handle will become invalid afterwards.
    unsafe fn create_window() -> Result<Window> {
        // Gfxstream window is a child window of CrosVM window. Without WS_CLIPCHILDREN, the parent
        // window may use the background brush to clear the gfxstream window client area when
        // drawing occurs. This caused the screen flickering issue during resizing.
        // See b/197786842 for details.
        let dw_style = WS_POPUP | WS_CLIPCHILDREN;
        Window::new(
            Some(Self::wnd_proc),
            /* class_name */ "CROSVM",
            /* title */ "crosvm",
            APP_ICON_ID,
            dw_style,
            // The window size and style can be adjusted later when `Surface` is created.
            &size2(1, 1),
        )
    }

    /// # Safety
    /// The message queue must have been created on the WndProc thread before calling this, and the
    /// returned thread must not outlive the WndProc thread.
    unsafe fn start_message_relay_thread(
        #[allow(unused_variables)] vm_tube: Arc<Mutex<Tube>>,
        #[allow(unused_variables)] wndproc_thread_id: u32,
    ) -> Option<Box<dyn MessageRelayThreadTrait>> {
        #[cfg(feature = "kiwi")]
        match MessageRelayThread::<ServiceSendToGpu>::start_thread(
            vm_tube,
            wndproc_thread_id,
            WM_USER_HANDLE_SERVICE_MESSAGE_INTERNAL,
        ) {
            Ok(thread) => Some(thread),
            // We won't get messages from the service if we failed to spawn this thread. It may not
            // worth terminating the WndProc thread and crashing the emulator in that case, so we
            // just log the error.
            Err(e) => {
                error!("{:?}", e);
                None
            }
        }

        #[cfg(not(feature = "kiwi"))]
        None
    }

    unsafe extern "system" fn wnd_proc(
        hwnd: HWND,
        msg: UINT,
        w_param: WPARAM,
        l_param: LPARAM,
    ) -> LRESULT {
        let dispatcher_ptr = GetPropA(hwnd, win32_string(DISPATCHER_PROPERTY_NAME).as_ptr())
            as *mut WindowMessageDispatcher<T>;
        if let Some(dispatcher) = dispatcher_ptr.as_mut() {
            if let Some(ret) =
                dispatcher.dispatch_window_message(hwnd, &MessagePacket::new(msg, w_param, l_param))
            {
                return ret;
            }
        }
        DefWindowProcA(hwnd, msg, w_param, l_param)
    }
}

impl<T: HandleWindowMessage> Drop for WindowProcedureThread<T> {
    fn drop(&mut self) {
        self.signal_exit_message_loop_if_needed();
        match self.thread.take().unwrap().join() {
            Ok(_) => self.check_message_loop_final_state(),
            Err(e) => error!("Failed to join with WndProc thread: {:?}", e),
        }
    }
}

// `Send` may not be automatically inherited because of the `PhantomData`.
// Since `WindowProcedureThread` does not hold anything that cannot be transferred between threads,
// we can implement `Send` for it.
unsafe impl<T: HandleWindowMessage> Send for WindowProcedureThread<T> {}
<commit_msg>gpu_display: use mpsc channels to avoid busy-waiting (2).<commit_after>// Copyright 2022 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::marker::PhantomData;
use std::mem;
use std::pin::Pin;
use std::ptr::null_mut;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::Ordering;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::Arc;
use std::thread::Builder as ThreadBuilder;
use std::thread::JoinHandle;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use base::error;
use base::info;
use base::warn;
use base::Event;
use base::Tube;
use euclid::size2;
use sync::Mutex;
#[cfg(feature = "kiwi")]
use vm_control::ServiceSendToGpu;
use win_util::win32_string;
use winapi::shared::minwindef::LPARAM;
use winapi::shared::minwindef::LRESULT;
use winapi::shared::minwindef::UINT;
use winapi::shared::minwindef::WPARAM;
use winapi::shared::windef::HWND;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::processthreadsapi::GetCurrentThreadId;
use winapi::um::winuser::*;

#[cfg(feature = "kiwi")]
use super::message_relay_thread::MessageRelayThread;
use super::message_relay_thread::MessageRelayThreadTrait;
use super::thread_message_util;
use super::window::MessagePacket;
use super::window::Window;
use super::window_message_dispatcher::WindowMessageDispatcher;
use super::window_message_dispatcher::DISPATCHER_PROPERTY_NAME;
use super::window_message_processor::*;

// The default app icon id, which is defined in crosvm-manifest.rc.
const APP_ICON_ID: u16 = 1;

#[derive(Debug)]
enum MessageLoopState {
    /// The initial state.
    NotStarted = 0,
    /// The loop is running normally.
    Running,
    /// The loop has ended normally.
    ExitedNormally,
    /// The loop never started because errors occurred.
    EarlyTerminatedWithError,
    /// The loop has ended because errors occurred.
    ExitedWithError,
}

/// This class runs the WndProc thread, and provides helper functions for other threads to
/// communicate with it.
pub struct WindowProcedureThread<T: HandleWindowMessage> {
    thread: Option<JoinHandle<()>>,
    thread_id: u32,
    message_loop_state: Option<Arc<AtomicI32>>,
    thread_terminated_event: Event,
    _marker: PhantomData<T>,
}

impl<T: HandleWindowMessage> WindowProcedureThread<T> {
    pub fn start_thread(vm_tube: Option<Arc<Mutex<Tube>>>) -> Result<Self> {
        let (thread_id_sender, thread_id_receiver) = channel();
        let message_loop_state = Arc::new(AtomicI32::new(MessageLoopState::NotStarted as i32));
        let thread_terminated_event = Event::new().unwrap();

        let message_loop_state_clone = Arc::clone(&message_loop_state);
        let thread_terminated_event_clone = thread_terminated_event
            .try_clone()
            .map_err(|e| anyhow!("Failed to clone thread_terminated_event: {}", e))?;

        let thread = match ThreadBuilder::new()
            .name("gpu_display_wndproc".into())
            .spawn(move || {
                // Safe because GetCurrentThreadId has no failure mode.
                let thread_id = unsafe { GetCurrentThreadId() };

                // Must be called before any other threads post messages to the WndProc thread.
                thread_message_util::force_create_message_queue();

                // Safe because the message queue has been created, and the returned thread will go
                // out of scope and get dropped before the WndProc thread exits.
                let _message_relay_thread = unsafe {
                    vm_tube.and_then(|tube| Self::start_message_relay_thread(tube, thread_id))
                };

                Self::run_message_loop(thread_id_sender, message_loop_state_clone);

                if let Err(e) = thread_terminated_event_clone.write(1) {
                    error!("Failed to write to thread terminated event: {}", e);
                }
            }) {
            Ok(thread) => thread,
            Err(e) => bail!("Failed to spawn WndProc thread: {}", e),
        };

        match thread_id_receiver.recv() {
            Ok(thread_id_res) => match thread_id_res {
                Ok(thread_id) => Ok(Self {
                    thread: Some(thread),
                    thread_id,
                    message_loop_state: Some(message_loop_state),
                    thread_terminated_event,
                    _marker: PhantomData,
                }),
                Err(e) => bail!("WndProc internal failure: {:?}", e),
            },
            Err(e) => bail!("Failed to receive WndProc thread ID: {}", e),
        }
    }

    pub fn try_clone_thread_terminated_event(&self) -> Result<Event> {
        self.thread_terminated_event
            .try_clone()
            .map_err(|e| anyhow!("Failed to clone thread_terminated_event: {}", e))
    }

    pub fn post_display_command(&self, message: DisplaySendToWndProc<T>) -> Result<()> {
        if !self.is_message_loop_running() {
            bail!("Host window has been destroyed!");
        }

        // Safe because the WndProc thread is still running the message loop.
        unsafe {
            thread_message_util::post_message_carrying_object(
                self.thread_id,
                WM_USER_HANDLE_DISPLAY_MESSAGE_INTERNAL,
                message,
            )
            .context("When posting DisplaySendToWndProc message")
        }
    }

    fn run_message_loop(thread_id_sender: Sender<Result<u32>>, message_loop_state: Arc<AtomicI32>) {
        // Safe because the dispatcher will take care of the lifetime of the `Window` object.
        let create_window_res = unsafe { Self::create_window() };
        match create_window_res.and_then(|window| WindowMessageDispatcher::<T>::create(window)) {
            Ok(dispatcher) => {
                info!("WndProc thread entering message loop");
                message_loop_state.store(MessageLoopState::Running as i32, Ordering::SeqCst);
                if let Err(e) = thread_id_sender.send(Ok(unsafe { GetCurrentThreadId() })) {
                    error!("Failed to send WndProc thread ID: {}", e);
                }
                Self::run_message_loop_body(dispatcher, message_loop_state);
            }
            Err(e) => {
                error!("WndProc thread didn't enter message loop: {:?}", e);
                message_loop_state.store(
                    MessageLoopState::EarlyTerminatedWithError as i32,
                    Ordering::SeqCst,
                );
                if let Err(e) = thread_id_sender.send(Err(e)) {
                    error!("Failed to report message loop early termination: {}", e)
                }
            }
        }
    }

    fn run_message_loop_body(
        mut message_dispatcher: Pin<Box<WindowMessageDispatcher<T>>>,
        message_loop_state: Arc<AtomicI32>,
    ) {
        loop {
            let mut message = mem::MaybeUninit::uninit();
            // Safe because we know the lifetime of `message`.
            match unsafe { GetMessageW(message.as_mut_ptr(), null_mut(), 0, 0) } {
                0 => {
                    info!("WndProc thread exiting message loop since WM_QUIT is received");
                    message_loop_state
                        .store(MessageLoopState::ExitedNormally as i32, Ordering::SeqCst);
                    break;
                }
                -1 => {
                    error!(
                        "WndProc thread exiting message loop because GetMessageW() failed with \
                        error code {}",
                        unsafe { GetLastError() }
                    );
                    message_loop_state
                        .store(MessageLoopState::ExitedWithError as i32, Ordering::SeqCst);
                    break;
                }
                _ => (),
            }

            // Safe because `GetMessageW()` will block until `message` is populated.
            let new_message = unsafe { message.assume_init() };
            if new_message.hwnd.is_null() {
                // Thread messages don't target a specific window and `DispatchMessageA()` won't
                // send them to `wnd_proc()` function, hence we need to handle it as a special case.
                message_dispatcher
                    .as_mut()
                    .process_thread_message(&new_message.into());
            } else {
                // Safe because `GetMessageW()` will block until `message` is populated.
                unsafe {
                    TranslateMessage(&new_message);
                    DispatchMessageW(&new_message);
                }
            }
        }
    }

    fn is_message_loop_running(&self) -> bool {
        self.message_loop_state.as_ref().map_or(false, |state| {
            state.load(Ordering::SeqCst) == MessageLoopState::Running as i32
        })
    }

    /// In the normal case, when all windows are closed by the user, the WndProc thread exits the
    /// message loop and terminates naturally. If we have to shutdown the VM before all windows are
    /// closed because of errors, this function will post a message to let the WndProc thread kill
    /// all windows and terminate.
    fn signal_exit_message_loop_if_needed(&self) {
        if !self.is_message_loop_running() {
            return;
        }

        info!("WndProc thread is still in message loop before dropping. Signaling killing windows");
        // Safe because the WndProc thread is still running the message loop.
        if let Err(e) = unsafe {
            thread_message_util::post_message(
                self.thread_id,
                WM_USER_WNDPROC_THREAD_DROP_KILL_WINDOW_INTERNAL,
                /* w_param */ 0,
                /* l_param */ 0,
            )
        } {
            error!("Failed to signal WndProc thread to kill windows: {:?}", e);
        }
    }

    /// Checks if the message loop has exited normally. This should be called after joining with the
    /// WndProc thread.
    fn check_message_loop_final_state(&mut self) {
        match Arc::try_unwrap(self.message_loop_state.take().unwrap()) {
            Ok(state) => {
                let state = state.into_inner();
                if state == MessageLoopState::ExitedNormally as i32 {
                    info!("WndProc thread exited gracefully");
                } else {
                    warn!("WndProc thread exited with message loop state: {:?}", state);
                }
            }
            Err(e) => error!(
                "WndProc thread exited but message loop state retrieval failed: {:?}",
                e
            ),
        }
    }

    /// # Safety
    /// The owner of the returned `Window` object is responsible for dropping it before we finish
    /// processing `WM_NCDESTROY`, because the window handle will become invalid afterwards.
    unsafe fn create_window() -> Result<Window> {
        // Gfxstream window is a child window of CrosVM window. Without WS_CLIPCHILDREN, the parent
        // window may use the background brush to clear the gfxstream window client area when
        // drawing occurs. This caused the screen flickering issue during resizing.
        // See b/197786842 for details.
        let dw_style = WS_POPUP | WS_CLIPCHILDREN;
        Window::new(
            Some(Self::wnd_proc),
            /* class_name */ "CROSVM",
            /* title */ "crosvm",
            APP_ICON_ID,
            dw_style,
            // The window size and style can be adjusted later when `Surface` is created.
            &size2(1, 1),
        )
    }

    /// # Safety
    /// The message queue must have been created on the WndProc thread before calling this, and the
    /// returned thread must not outlive the WndProc thread.
    unsafe fn start_message_relay_thread(
        #[allow(unused_variables)] vm_tube: Arc<Mutex<Tube>>,
        #[allow(unused_variables)] wndproc_thread_id: u32,
    ) -> Option<Box<dyn MessageRelayThreadTrait>> {
        #[cfg(feature = "kiwi")]
        match MessageRelayThread::<ServiceSendToGpu>::start_thread(
            vm_tube,
            wndproc_thread_id,
            WM_USER_HANDLE_SERVICE_MESSAGE_INTERNAL,
        ) {
            Ok(thread) => Some(thread),
            // We won't get messages from the service if we failed to spawn this thread. It may not
            // worth terminating the WndProc thread and crashing the emulator in that case, so we
            // just log the error.
            Err(e) => {
                error!("{:?}", e);
                None
            }
        }

        #[cfg(not(feature = "kiwi"))]
        None
    }

    unsafe extern "system" fn wnd_proc(
        hwnd: HWND,
        msg: UINT,
        w_param: WPARAM,
        l_param: LPARAM,
    ) -> LRESULT {
        let dispatcher_ptr = GetPropA(hwnd, win32_string(DISPATCHER_PROPERTY_NAME).as_ptr())
            as *mut WindowMessageDispatcher<T>;
        if let Some(dispatcher) = dispatcher_ptr.as_mut() {
            if let Some(ret) =
                dispatcher.dispatch_window_message(hwnd, &MessagePacket::new(msg, w_param, l_param))
            {
                return ret;
            }
        }
        DefWindowProcA(hwnd, msg, w_param, l_param)
    }
}

impl<T: HandleWindowMessage> Drop for WindowProcedureThread<T> {
    fn drop(&mut self) {
        self.signal_exit_message_loop_if_needed();
        match self.thread.take().unwrap().join() {
            Ok(_) => self.check_message_loop_final_state(),
            Err(e) => error!("Failed to join with WndProc thread: {:?}", e),
        }
    }
}

// `Send` may not be automatically inherited because of the `PhantomData`.
// Since `WindowProcedureThread` does not hold anything that cannot be transferred between threads,
// we can implement `Send` for it.
unsafe impl<T: HandleWindowMessage> Send for WindowProcedureThread<T> {}
<|endoftext|>


<commit_before>//
// Copyright 2022 The Project Oak Authors
//
// 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.
//

//! Rust implementations of the AMD SEV-SNP GHCB MSR protocol.
//!
//! The protocol is implemented by writing to, and reading from, model-specific register (MSR)
//! 0xC001_0130. The guest makes a request by writing an appropriate value to the MSR and calling
//! VMGEXIT. After the hypervisor processes the request and resumes the guest, the guest can read
//! the response from the same MSR.
//!
//! The 12 least significant bits represent information about the operation being performed. The
//! remaining 52 bits represent the data for the operation. The format of the data is dependent on
//! the operation.
//!
//! See section 2.3.1 in <https://developer.amd.com/wp-content/resources/56421.pdf> for more detail.

use crate::instructions::vmgexit;
use anyhow::{bail, Result};
use strum::FromRepr;
use x86_64::registers::model_specific::Msr;

/// The version of the GHCB MSR protocol supported by this library. This represents the version
/// specific to AMD SEV-SNP.
pub const SUPPORTED_PROTOCOL_VERION: u16 = 2;

/// The identifier for the MSR used in the protocol.
const MSR_IDENTIFIER: u32 = 0xC001_0130;

/// Mask to extract the GHCB info from a u64 value.
const GHCB_INFO_MASK: u64 = 0xFFF;

/// Contains the guest-physical address of the GHCB page. The address must have been registered with
/// the hypervisor before using it.
pub struct GhcbGpa {
    gpa: u64,
}

impl GhcbGpa {
    pub fn new(gpa: usize) -> Result<Self> {
        let gpa = gpa as u64;
        if gpa & GHCB_INFO_MASK != 0 {
            bail!("GHCB must be 4KiB-aligned");
        }
        Ok(Self { gpa })
    }
}

impl From<GhcbGpa> for u64 {
    fn from(ghcb_gpa: GhcbGpa) -> Self {
        ghcb_gpa.gpa
    }
}

/// Sets the address of the GHCB page before exiting.
pub fn set_ghcb_address_and_exit(ghcb_gpa: GhcbGpa) {
    write_msr_and_exit(ghcb_gpa.into());
}

/// Response from the hypervisor about the encryption bit and supported GHCB protocol versions. The
/// encryption bit value is validated by the hardware before the guest is resumed.
pub struct SevInfoResponse {
    /// The minimum version of the GHCB MSR protocol supported by the hypervisor.
    pub min_protocol_version: u16,
    /// The maximum version of the GHCB MSR protocol supported by the hypervisor.
    pub max_protocol_version: u16,
    /// The page table bit used for inidicating that a page is encrypted.
    pub encryption_bit: u8,
}

impl TryFrom<u64> for SevInfoResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const SEV_INFO_RESPONSE_INFO: u64 = 0x001;
        if msr_value & GHCB_INFO_MASK != SEV_INFO_RESPONSE_INFO {
            bail!("Value is not a valid SEV information response.");
        }
        let max_protocol_version = (msr_value >> 48) as u16;
        let min_protocol_version = (msr_value >> 32) as u16;
        let encryption_bit = (msr_value >> 24) as u8;
        Ok(Self {
            min_protocol_version,
            max_protocol_version,
            encryption_bit,
        })
    }
}

/// A request for information about the supported GHCB MSR protocol version and the encryption bit.
pub struct SevInfoRequest;

impl From<SevInfoRequest> for u64 {
    fn from(_request: SevInfoRequest) -> Self {
        0x002
    }
}

/// Gets information about the supported GHCB MSR protocol verions and the location of the
/// encryption bit.
pub fn get_sev_info() -> Result<SevInfoResponse> {
    let request = SevInfoRequest;
    write_msr_and_exit(request.into());
    read_msr().try_into()
}

/// The register of interest from the result of executing CPUID.
#[derive(Debug, FromRepr)]
#[repr(u8)]
pub enum CpuidRegister {
    Eax = 0,
    Ebx = 1,
    Ecx = 2,
    Edx = 3,
}

/// A request to execute CPUID for a specific leaf and return one of the result registers.
pub struct CpuidRequest {
    /// The CPUID leaf to request. Sub-leafs are not supported byt this protocol.
    pub leaf: u32,
    /// The register to return from the result. This protocol only supports a single register at a
    /// time.
    pub register: CpuidRegister,
}

impl From<CpuidRequest> for u64 {
    fn from(request: CpuidRequest) -> Self {
        let mut result = 0x004;
        result |= (request.leaf as u64) << 32;
        result |= (request.register as u64) << 30;
        result
    }
}

/// A response from executing CPUID for a specific leaf. Only one register is returned at a time.
pub struct CpuidResponse {
    /// The value of the requested register after executing CPUID.
    pub value: u32,
    /// The register that the value represents.
    pub register: CpuidRegister,
}

impl TryFrom<u64> for CpuidResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const SEV_INFO_RESPONSE_INFO: u64 = 0x005;
        const RESERVED_MASK: u64 = 0x3FFFF000;
        if msr_value & GHCB_INFO_MASK != SEV_INFO_RESPONSE_INFO || msr_value & RESERVED_MASK != 0 {
            bail!("Value is not a valid CPUID response.");
        }
        let value = (msr_value >> 32) as u32;
        const REGISTER_MASK: u64 = 0xC0000000;
        let register = CpuidRegister::from_repr(((msr_value & REGISTER_MASK) >> 30) as u8)
            .ok_or_else(|| anyhow::anyhow!("Invalid register"))?;
        Ok(Self { value, register })
    }
}

/// Gets the value of the specified register that was returned when executing CPUID for the
/// specified leaf. Sub-leafs are not supported.
pub fn get_cpuid(request: CpuidRequest) -> Result<CpuidResponse> {
    write_msr_and_exit(request.into());
    read_msr().try_into()
}

/// Writes a value to the MSR and calls VMGEXIT to hand control to the hypervisor.
pub fn write_msr_and_exit(msr_value: u64) {
    // Safety: This operation is safe because this specific MSR is used only for communicating with
    // the hypervisor, and does not have any other side-effects within the guest.
    unsafe {
        Msr::new(MSR_IDENTIFIER).write(msr_value);
    }
    vmgexit();
}

/// Reads the value of the model-specific register.
///
/// This typically used to get the result of an operation when the hypervisor resumes the guest
/// following a call to VMGEGIT.
pub fn read_msr() -> u64 {
    // Safety: This operation is safe because this specific MSR is used only for communicating with
    // the hypervisor, and does not have any other side-effects within the guest.
    unsafe { Msr::new(MSR_IDENTIFIER).read() }
}
<commit_msg>Add more GHCB MSR protocol support (#3048)<commit_after>//
// Copyright 2022 The Project Oak Authors
//
// 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.
//

//! Rust implementations of the AMD SEV-SNP GHCB MSR protocol.
//!
//! The protocol is implemented by writing to, and reading from, model-specific register (MSR)
//! 0xC001_0130. The guest makes a request by writing an appropriate value to the MSR and calling
//! VMGEXIT. After the hypervisor processes the request and resumes the guest, the guest can read
//! the response from the same MSR.
//!
//! The 12 least significant bits represent information about the operation being performed. The
//! remaining 52 bits represent the data for the operation. The format of the data is dependent on
//! the operation.
//!
//! The MSR protocol is primarily used to set up the GHCB communications with the hypervisor.
//!
//! See section 2.3.1 in <https://developer.amd.com/wp-content/resources/56421.pdf> for more detail.
//!
//! Note: This implementation does not include the AP Reset Hold operations seeing that it is not
//! the preferred mechanism for starting a new application processor (AP) under AMD SEV-SNP. The
//! SEV-SNP AP creation NAE should be used instead. See section 4.3.2 of
//! <https://developer.amd.com/wp-content/resources/56421.pdf>.

use crate::instructions::vmgexit;
use anyhow::{bail, Result};
use strum::FromRepr;
use x86_64::registers::model_specific::Msr;

/// The version of the GHCB MSR protocol supported by this library. This represents the version
/// specific to AMD SEV-SNP.
pub const SUPPORTED_PROTOCOL_VERION: u16 = 2;

/// Value indicating that the hypervisor does not have a preferred location for the GHCB.
pub const NO_PREFERRED_GHCB_LOCATION: u64 = 0xFFFFFFFFFFFFF000;

/// Value indicating that the hypervisor did not accept the registered location for the GHCB.
const GHCB_LOCATION_NOT_ACCEPTED: u64 = 0xFFFFFFFFFFFFF000;

/// The identifier for the MSR used in the protocol.
const MSR_IDENTIFIER: u32 = 0xC001_0130;

/// Mask to extract the GHCB info from a u64 value.
const GHCB_INFO_MASK: u64 = 0xFFF;

/// Mask to extract the GHCB data from a u64 value.
const GCHP_DATA_MASK: u64 = !GHCB_INFO_MASK;

/// Contains the guest-physical address of the GHCB page. The address must have been registered with
/// the hypervisor before using it.
pub struct GhcbGpa {
    gpa: u64,
}

impl GhcbGpa {
    pub fn new(gpa: usize) -> Result<Self> {
        let gpa = gpa as u64;
        if gpa & GHCB_INFO_MASK != 0 {
            bail!("GHCB must be 4KiB-aligned");
        }
        Ok(Self { gpa })
    }
}

impl From<GhcbGpa> for u64 {
    fn from(ghcb_gpa: GhcbGpa) -> Self {
        ghcb_gpa.gpa
    }
}

/// Sets the address of the GHCB page before exiting to the hypervisor.
pub fn set_ghcb_address_and_exit(ghcb_gpa: GhcbGpa) {
    write_msr_and_exit(ghcb_gpa.into());
}

/// Response from the hypervisor about the encryption bit and supported GHCB protocol versions. The
/// encryption bit value is validated by the hardware before the guest is resumed.
pub struct SevInfoResponse {
    /// The minimum version of the GHCB MSR protocol supported by the hypervisor.
    pub min_protocol_version: u16,
    /// The maximum version of the GHCB MSR protocol supported by the hypervisor.
    pub max_protocol_version: u16,
    /// The page table bit used for inidicating that a page is encrypted.
    pub encryption_bit: u8,
}

impl TryFrom<u64> for SevInfoResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const SEV_INFO_RESPONSE_INFO: u64 = 0x001;
        if msr_value & GHCB_INFO_MASK != SEV_INFO_RESPONSE_INFO {
            bail!("Value is not a valid SEV information response");
        }
        let max_protocol_version = (msr_value >> 48) as u16;
        let min_protocol_version = (msr_value >> 32) as u16;
        let encryption_bit = (msr_value >> 24) as u8;
        Ok(Self {
            min_protocol_version,
            max_protocol_version,
            encryption_bit,
        })
    }
}

/// A request for information about the supported GHCB MSR protocol version and the encryption bit.
pub struct SevInfoRequest;

impl From<SevInfoRequest> for u64 {
    fn from(_request: SevInfoRequest) -> Self {
        0x002
    }
}

/// Gets information about the supported GHCB MSR protocol verions and the location of the
/// encryption bit.
pub fn get_sev_info() -> Result<SevInfoResponse> {
    let request = SevInfoRequest;
    write_msr_and_exit(request.into());
    read_msr().try_into()
}

/// The register of interest from the result of executing CPUID.
#[derive(Debug, FromRepr)]
#[repr(u8)]
pub enum CpuidRegister {
    Eax = 0,
    Ebx = 1,
    Ecx = 2,
    Edx = 3,
}

/// A request to execute CPUID for a specific leaf and return one of the result registers.
pub struct CpuidRequest {
    /// The CPUID leaf to request. Sub-leafs are not supported byt this protocol.
    pub leaf: u32,
    /// The register to return from the result. This protocol only supports a single register at a
    /// time.
    pub register: CpuidRegister,
}

impl From<CpuidRequest> for u64 {
    fn from(request: CpuidRequest) -> Self {
        let mut result = 0x004;
        result |= (request.leaf as u64) << 32;
        result |= (request.register as u64) << 30;
        result
    }
}

/// A response from executing CPUID for a specific leaf. Only one register is returned at a time.
pub struct CpuidResponse {
    /// The value of the requested register after executing CPUID.
    pub value: u32,
    /// The register that the value represents.
    pub register: CpuidRegister,
}

impl TryFrom<u64> for CpuidResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const SEV_INFO_RESPONSE_INFO: u64 = 0x005;
        const RESERVED_MASK: u64 = 0x3FFFF000;
        if msr_value & GHCB_INFO_MASK != SEV_INFO_RESPONSE_INFO || msr_value & RESERVED_MASK != 0 {
            bail!("Value is not a valid CPUID response");
        }
        let value = (msr_value >> 32) as u32;
        const REGISTER_MASK: u64 = 0xC0000000;
        let register = CpuidRegister::from_repr(((msr_value & REGISTER_MASK) >> 30) as u8)
            .ok_or_else(|| anyhow::anyhow!("Invalid register"))?;
        Ok(Self { value, register })
    }
}

/// Gets the value of the specified register that was returned when executing CPUID for the
/// specified leaf. Sub-leafs are not supported.
pub fn get_cpuid(request: CpuidRequest) -> Result<CpuidResponse> {
    write_msr_and_exit(request.into());
    read_msr().try_into()
}

/// A request for the hypervisor's preferred location for the GHCB page.
pub struct PreferredGhcbGpaRequest;

impl From<PreferredGhcbGpaRequest> for u64 {
    fn from(_request: PreferredGhcbGpaRequest) -> Self {
        0x010
    }
}

/// The response containing the preferred location of the GHCB.
pub struct PreferredGhcbGpaResponse {
    /// The preferred guest-physical address for the GHCB.
    pub ghcb_gpa: usize,
}

impl TryFrom<u64> for PreferredGhcbGpaResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const PREFERRED_GPA_RESPONSE_INFO: u64 = 0x011;
        if msr_value & GHCB_INFO_MASK != PREFERRED_GPA_RESPONSE_INFO {
            bail!("Value is not a valid preferred GHCP GPA response");
        }
        let ghcb_gpa = (msr_value & GCHP_DATA_MASK) as usize;
        Ok(Self { ghcb_gpa })
    }
}

/// Requests the hypervisor's preferred location for the GHCB page.
///
/// A response of `NO_PREFERRED_GHCB_LOCATION` indicates that the hypervisor does not have a
/// preferred location.
///
/// The guest must validate that the returned guest-physical address is not part of its known
/// guest-private memory.
pub fn get_preferred_ghcb_location() -> Result<PreferredGhcbGpaResponse> {
    let request = PreferredGhcbGpaRequest;
    write_msr_and_exit(request.into());
    read_msr().try_into()
}

/// Request to register a guest-physical address for the GHCB with the hypervisor.
pub struct RegisterGhcbGpaRequest {
    ghcb_gpa: u64,
}

impl RegisterGhcbGpaRequest {
    pub fn new(ghcb_gpa: usize) -> Result<Self> {
        let ghcb_gpa = ghcb_gpa as u64;
        if ghcb_gpa & GHCB_INFO_MASK != 0 {
            bail!("GHCB must be 4KiB-aligned");
        }
        Ok(Self { ghcb_gpa })
    }
}

impl From<RegisterGhcbGpaRequest> for u64 {
    fn from(register_ghcb_gpa_request: RegisterGhcbGpaRequest) -> Self {
        let mut result = 0x012;
        result |= register_ghcb_gpa_request.ghcb_gpa;
        result
    }
}

/// The response containing the result of the GHCB registration.
pub struct RegisterGhcbGpaResponse {
    /// The registered guest-physical address for the GHCB.
    ghcb_gpa: u64,
}

impl TryFrom<u64> for RegisterGhcbGpaResponse {
    type Error = anyhow::Error;
    fn try_from(msr_value: u64) -> Result<Self> {
        const REGISTER_GHCB_GPA_RESPONSE_INFO: u64 = 0x013;
        if msr_value & GHCB_INFO_MASK != REGISTER_GHCB_GPA_RESPONSE_INFO {
            bail!("Value is not a valid GHCP GPA registration response");
        }
        let ghcb_gpa = msr_value & GCHP_DATA_MASK;
        Ok(Self { ghcb_gpa })
    }
}

/// Registers the location of the GHCB page for the current vCPU with the hypervisor.
pub fn register_ghcb_location(request: RegisterGhcbGpaRequest) -> Result<()> {
    let request_ghcb_gpa: u64 = request.ghcb_gpa;
    write_msr_and_exit(request.into());
    let response: RegisterGhcbGpaResponse = read_msr().try_into()?;
    // Ensure that the registration was successful.
    if response.ghcb_gpa == GHCB_LOCATION_NOT_ACCEPTED {
        bail!("GHCB location registration not accepted");
    }
    if response.ghcb_gpa != request_ghcb_gpa {
        bail!("The registration response did not match the request");
    }
    Ok(())
}

/// Writes a value to the MSR and calls VMGEXIT to hand control to the hypervisor.
pub fn write_msr_and_exit(msr_value: u64) {
    // Safety: This operation is safe because this specific MSR is used only for communicating with
    // the hypervisor, and does not have any other side-effects within the guest.
    unsafe {
        Msr::new(MSR_IDENTIFIER).write(msr_value);
    }
    vmgexit();
}

/// Reads the value of the model-specific register.
///
/// This typically used to get the result of an operation when the hypervisor resumes the guest
/// following a call to VMGEGIT.
pub fn read_msr() -> u64 {
    // Safety: This operation is safe because this specific MSR is used only for communicating with
    // the hypervisor, and does not have any other side-effects within the guest.
    unsafe { Msr::new(MSR_IDENTIFIER).read() }
}
<|endoftext|>


<commit_before>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for its existence is to compare the efficiency Rust implementation against
//! the existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::action::timeline;
use rrg::session::Sink;

/// A binary for the timeline action.
#[derive(argh::FromArgs)]
struct Opts {
    /// A path to the root directory to timeline.
    #[argh(option,
           long = "root",
           arg_name = "FILE",
           default = "::std::path::PathBuf::from(\"/\")",
           description = "root directory to timeline")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[argh(positional,
           arg_name = "OUTPUT",
           description = "path to dump the results into")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // For now we are not interested in doing anything useful with chunk ids
        // since everything is dumped into one file and there is no need to
        // refer to a particular chunk.
        //
        // In the future they might be useful for printing some statistics about
        // the collected files.
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        // Just a sanity check in case the implementation of the timeline action
        // starts sending data to other sinks.
        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let opts: Opts = argh::from_env();

    timeline::handle(&mut Session::open(opts.output), timeline::Request {
        root: opts.root,
    }).expect("failed to execute the action");
}
<commit_msg>Rename the args struct in the timeline example.<commit_after>//! An example that timelines the specified folder.
//!
//! The timeline action is one of the filesystem tools that the GRR agent
//! exposes. Its description is pretty simple: given a folder, it returns a
//! stream of stat entries for all the files in it (including all subfolders).
//!
//! This example is a tiny wrapper around this action that allows to execute it
//! as a standalone binary without all the RRG setup required. A primary reason
//! for its existence is to compare the efficiency Rust implementation against
//! the existing Python one.

use std::fs::File;
use std::path::{Path, PathBuf};

use rrg::action::timeline;
use rrg::session::Sink;

/// A binary for the timeline action.
#[derive(argh::FromArgs)]
struct Args {
    /// A path to the root directory to timeline.
    #[argh(option,
           long = "root",
           arg_name = "FILE",
           default = "::std::path::PathBuf::from(\"/\")",
           description = "root directory to timeline")]
    root: PathBuf,

    /// A path to a file to dump the results into.
    #[argh(positional,
           arg_name = "OUTPUT",
           description = "path to dump the results into")]
    output: PathBuf,
}

/// A session type of the timeline example.
///
/// Every action handler needs a session object. "Real" sessions handle comms
/// with the GRR server, error reporting and alike. However, this is not needed
/// here and the example session simply dumps everything into the specified file
/// encoding data in the gzchunked format.
struct Session {
    /// A file into which the data chunks of the output are dumped.
    output: File,
}

impl Session {

    /// Creates a session object that will write to the given `output` path.
    fn open<P: AsRef<Path>>(output: P) -> Session {
        let output = File::create(output)
            .expect("failed to create the output file");

        Session {
            output: output,
        }
    }
}

impl rrg::session::Session for Session {

    fn reply<R>(&mut self, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        // For now we are not interested in doing anything useful with chunk ids
        // since everything is dumped into one file and there is no need to
        // refer to a particular chunk.
        //
        // In the future they might be useful for printing some statistics about
        // the collected files.
        drop(response);

        Ok(())
    }

    fn send<R>(&mut self, sink: Sink, response: R) -> rrg::session::Result<()>
    where
        R: rrg::action::Response + 'static,
    {
        use std::io::Write as _;
        use byteorder::{BigEndian, WriteBytesExt as _};

        // Just a sanity check in case the implementation of the timeline action
        // starts sending data to other sinks.
        assert_eq!(sink, Sink::TRANSFER_STORE);

        let response = (&response as &dyn std::any::Any)
            .downcast_ref::<timeline::Chunk>()
            .expect("unexpected response type");

        self.output.write_u64::<BigEndian>(response.data.len() as u64)
            .expect("failed to write chunk size tag");
        self.output.write_all(&response.data[..])
            .expect("failed to write chunk data");

        Ok(())
    }
}

fn main() {
    let args: Args = argh::from_env();

    timeline::handle(&mut Session::open(args.output), timeline::Request {
        root: args.root,
    }).expect("failed to execute the action");
}
<|endoftext|>


<commit_before>#![crate_type = "dylib"]
#![feature(plugin)]
#![plugin(interpolate_idents)]

#[macro_use] extern crate cpython;
extern crate yaml_rust;

use std::collections::BTreeMap;
use cpython::{PyString, Python, PyResult, PyObject, PyDict, PyList, PythonObject, ToPyObject};
use yaml_rust::{Yaml, YamlLoader};

fn convert_yaml_to_dict(py: Python, yaml: &BTreeMap<Yaml, Yaml>) -> PyDict {
    let dict = PyDict::new(py);

    for (k, v) in yaml.iter() {
        let key: PyObject = match from_yaml_to_python(py, k) {
            Some(o) => o,
            None => panic!("Error converting key {:?}", k)
        };
        let value: PyObject = match from_yaml_to_python(py, v) {
            Some(o) => o,
            None => panic!("Error converting value {:?}", v)
        };

        let _ = dict.set_item(py, key, value);
    }

    dict
}

fn convert_yaml_to_list(py: Python, yaml: &Vec<Yaml>) -> PyList {
    let vec: Vec<PyObject> = yaml.iter()
        .map(|e| from_yaml_to_python(py, e).unwrap())
        .collect();

    PyList::new(py, &vec)
}

fn from_yaml_to_python(py: Python, yaml: &Yaml) -> Option<PyObject> {
    // None
    if yaml.is_null() {
        return Some(py.None())
    }

    // Hash
    let hash = yaml.as_hash();
    if hash.is_some() {
        let dict = convert_yaml_to_dict(py, hash.unwrap());
        return Some(dict.into_object())
    }

    // String
    let string = yaml.as_str();
    if string.is_some() {
        let pystring = PyString::new(py, string.unwrap());
        return Some(pystring.into_object())
    }

    // Integer
    let int = yaml.as_i64();
    if int.is_some() {
        return Some(int.to_py_object(py).into_object())
    }

    // List
    let vec = yaml.as_vec();
    if vec.is_some() {
        let list = convert_yaml_to_list(py, vec.unwrap());
        return Some(list.into_object())
    }

    // Float
    let float = yaml.as_f64();
    if float.is_some() {
        return Some(float.to_py_object(py).into_object())
    }

    // Boolean
    let boolean = yaml.as_bool();
    if boolean.is_some() {
        match boolean.unwrap() {
            true => return Some(py.True().into_object()),
            false => return Some(py.False().into_object())
        }
    }

    // Default
    return None
}

fn safe_load(py: Python, stream: PyString) -> PyResult<PyObject> {
    // Convert stream into Rust string
    let native_stream = match stream.to_string(py) {
        Ok(s) => s,
        Err(e) => panic!("Cannot convert Python string into &str: {:?}", e)
    };

    // // Load first doc from stream
    let docs = match YamlLoader::load_from_str(&native_stream) {
        Ok(d) => d,
        Err(e) => panic!("{:?}", e)
    };
    let doc = &docs[0];

    // Convert into proper Python's objects
    let result: PyObject = match from_yaml_to_python(py, doc) {
        Some(o) => o,
        None => py.None()
    };

    Ok(result)
}

fn test(py: Python) -> PyResult<PyObject> {
    Ok(py.None())
}

py_module_initializer!(ryaml, initryaml, PyInit_ryaml, |py, m| {
    try!(m.add(py, "safe_load", py_fn!(py, safe_load(stream: PyString))));
    try!(m.add(py, "test", py_fn!(py, test())));
    Ok(())
});
<commit_msg>Removed test code<commit_after>#![crate_type = "dylib"]
#![feature(plugin)]
#![plugin(interpolate_idents)]

#[macro_use] extern crate cpython;
extern crate yaml_rust;

use std::collections::BTreeMap;
use cpython::{PyString, Python, PyResult, PyObject, PyDict, PyList, PythonObject, ToPyObject};
use yaml_rust::{Yaml, YamlLoader};

fn convert_yaml_to_dict(py: Python, yaml: &BTreeMap<Yaml, Yaml>) -> PyDict {
    let dict = PyDict::new(py);

    for (k, v) in yaml.iter() {
        let key: PyObject = match from_yaml_to_python(py, k) {
            Some(o) => o,
            None => panic!("Error converting key {:?}", k)
        };
        let value: PyObject = match from_yaml_to_python(py, v) {
            Some(o) => o,
            None => panic!("Error converting value {:?}", v)
        };

        let _ = dict.set_item(py, key, value);
    }

    dict
}

fn convert_yaml_to_list(py: Python, yaml: &Vec<Yaml>) -> PyList {
    let vec: Vec<PyObject> = yaml.iter()
        .map(|e| from_yaml_to_python(py, e).unwrap())
        .collect();

    PyList::new(py, &vec)
}

fn from_yaml_to_python(py: Python, yaml: &Yaml) -> Option<PyObject> {
    // None
    if yaml.is_null() {
        return Some(py.None())
    }

    // Hash
    let hash = yaml.as_hash();
    if hash.is_some() {
        let dict = convert_yaml_to_dict(py, hash.unwrap());
        return Some(dict.into_object())
    }

    // String
    let string = yaml.as_str();
    if string.is_some() {
        let pystring = PyString::new(py, string.unwrap());
        return Some(pystring.into_object())
    }

    // Integer
    let int = yaml.as_i64();
    if int.is_some() {
        return Some(int.to_py_object(py).into_object())
    }

    // List
    let vec = yaml.as_vec();
    if vec.is_some() {
        let list = convert_yaml_to_list(py, vec.unwrap());
        return Some(list.into_object())
    }

    // Float
    let float = yaml.as_f64();
    if float.is_some() {
        return Some(float.to_py_object(py).into_object())
    }

    // Boolean
    let boolean = yaml.as_bool();
    if boolean.is_some() {
        match boolean.unwrap() {
            true => return Some(py.True().into_object()),
            false => return Some(py.False().into_object())
        }
    }

    // Default
    return None
}

fn safe_load(py: Python, stream: PyString) -> PyResult<PyObject> {
    // Convert stream into Rust string
    let native_stream = match stream.to_string(py) {
        Ok(s) => s,
        Err(e) => panic!("Cannot convert Python string into &str: {:?}", e)
    };

    // // Load first doc from stream
    let docs = match YamlLoader::load_from_str(&native_stream) {
        Ok(d) => d,
        Err(e) => panic!("{:?}", e)
    };
    let doc = &docs[0];

    // Convert into proper Python's objects
    let result: PyObject = match from_yaml_to_python(py, doc) {
        Some(o) => o,
        None => py.None()
    };

    Ok(result)
}

py_module_initializer!(ryaml, initryaml, PyInit_ryaml, |py, m| {
    try!(m.add(py, "safe_load", py_fn!(py, safe_load(stream: PyString))));
    Ok(())
});
<|endoftext|>


<commit_before>//! Wrapper around a wl_output

use std::fmt;

use wlroots::{Area, Size, Origin};
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::{Proxy, NewProxy};

use lua::LUA;
use objects::screen::{self, Screen};

/// Wrapper around WlOutput.
#[derive(Clone, Eq, PartialEq)]
pub struct Output {
    proxy: Proxy<WlOutput>,
}

/// The cached state for the WlOutput.
#[derive(Debug, Default, Clone, Eq, PartialEq)]
struct OutputState {
    name: String,
    resolution: (i32, i32)
}

impl <'this> Into<&'this Proxy<WlOutput>> for &'this Output {
    fn into(self) -> &'this Proxy<WlOutput> {
        &self.proxy
    }
}

impl Output {
    pub fn new(new_proxy: Result<NewProxy<WlOutput>, u32>, _: ()) {
        let new_proxy = new_proxy.expect("Could not create WlOutput");
        let state = Box::new(OutputState::default());
        let proxy = new_proxy.implement(move |event, mut proxy: Proxy<WlOutput>| {
            use wayland_client::protocol::wl_output::Event;
            let output = Output { proxy: proxy.clone() };
            let state = unwrap_state_mut(&mut proxy);
            LUA.with(|lua| {
                let lua = lua.borrow();
                let lua = &*lua;
                let mut screen = screen::get_screen(lua, output)
                    .expect("Got screen state for screen that doesn't exist");
                match event {
                    Event::Geometry { make, model,  .. } => {
                        state.name = format!("{} ({})", make, model);
                    },
                    Event::Mode { width, height, .. } => {
                        let geometry = Area { origin: Origin { x: 0, y: 0 },
                                              size: Size { width, height }
                        };
                        screen.set_geometry(geometry)
                            .expect("could not set geometry");
                        screen.set_workarea(geometry)
                            .expect("could not set workarea ");
                        state.resolution = (width, height);
                    }
                    _ => {/* TODO */}
                }
            });
        });
        proxy.set_user_data(Box::into_raw(state) as _);
        let output = Output { proxy };
        LUA.with(|lua| {
            let lua = lua.borrow();
            let lua = &*lua;
            let mut screen = Screen::new(lua)
                .expect("Could not allocate new screen");
            screen.init_screens(output.clone(), vec![output])
                .expect("Could not initilize new output with a screen");
            screen::add_screen(lua, screen)
                .expect("Could not add screen to the list of screens");
        });
    }

    pub fn resolution(&self) -> (i32, i32) {
        unwrap_state(self).resolution
    }

    pub fn name(&self) -> &str {
        unwrap_state(self).name.as_str()
    }
}

impl fmt::Debug for Output {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.proxy.c_ptr())
    }
}

fn unwrap_state_mut<'this, I: Into<&'this mut Proxy<WlOutput>>>(proxy: I) -> &'this mut OutputState {
    unsafe {
        let user_data = proxy.into().get_user_data() as *mut OutputState;
        if user_data.is_null() {
            panic!("User data has not been set yet");
        }
        &mut *user_data
    }
}

fn unwrap_state<'this, I: Into<&'this Proxy<WlOutput>>>(proxy: I) -> &'this OutputState {
    unsafe {
        let user_data = proxy.into().get_user_data() as *const OutputState;
        if user_data.is_null() {
            panic!("User data has not been set yet");
        }
        &*user_data
    }
}
<commit_msg>Add the output to screen after Output::Event::Done<commit_after>//! Wrapper around a wl_output

use std::fmt;

use wlroots::{Area, Size, Origin};
use wayland_client::protocol::wl_output::WlOutput;
use wayland_client::{Proxy, NewProxy};

use lua::LUA;
use objects::screen::{self, Screen};

/// Wrapper around WlOutput.
#[derive(Clone, Eq, PartialEq)]
pub struct Output {
    proxy: Proxy<WlOutput>,
}

/// The cached state for the WlOutput.
#[derive(Debug, Default, Clone, Eq, PartialEq)]
struct OutputState {
    name: String,
    resolution: (i32, i32)
}

impl <'this> Into<&'this Proxy<WlOutput>> for &'this Output {
    fn into(self) -> &'this Proxy<WlOutput> {
        &self.proxy
    }
}

impl Output {
    pub fn new(new_proxy: Result<NewProxy<WlOutput>, u32>, _: ()) {
        let new_proxy = new_proxy.expect("Could not create WlOutput");
        let state = Box::new(OutputState::default());
        let proxy = new_proxy.implement(move |event, mut proxy: Proxy<WlOutput>| {
            use wayland_client::protocol::wl_output::Event;
            let output = Output { proxy: proxy.clone() };
            let state = unwrap_state_mut(&mut proxy);
            LUA.with(|lua| {
                let lua = lua.borrow();
                let lua = &*lua;
                match event {
                    Event::Geometry { make, model,  .. } => {
                        state.name = format!("{} ({})", make, model);
                    },
                    Event::Mode { width, height, .. } => {
                        state.resolution = (width, height);
                        let geometry = Area { origin: Origin { x: 0, y: 0 },
                                              size: Size { width, height }
                        };
                        if let Ok(mut screen) = screen::get_screen(lua, output) {
                            screen.set_geometry(geometry)
                                .expect("could not set geometry");
                            screen.set_workarea(geometry)
                                .expect("could not set workarea ");
                        }
                    }
                    Event::Done => {
                        let mut screen = Screen::new(lua)
                            .expect("Could not allocate new screen");
                        screen.init_screens(output.clone(), vec![output])
                            .expect("Could not initilize new output with a screen");
                        screen::add_screen(lua, screen)
                            .expect("Could not add screen to the list of screens");
                    }
                    _ => {/* TODO */}
                }
            });
        });
        proxy.set_user_data(Box::into_raw(state) as _);
        let output = Output { proxy };
        LUA.with(|lua| {
            let lua = lua.borrow();
            let lua = &*lua;
            let mut screen = Screen::new(lua)
                .expect("Could not allocate new screen");
            screen.init_screens(output.clone(), vec![output])
                .expect("Could not initilize new output with a screen");
            screen::add_screen(lua, screen)
                .expect("Could not add screen to the list of screens");
        });
    }

    pub fn resolution(&self) -> (i32, i32) {
        unwrap_state(self).resolution
    }

    pub fn name(&self) -> &str {
        unwrap_state(self).name.as_str()
    }
}

impl fmt::Debug for Output {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self.proxy.c_ptr())
    }
}

fn unwrap_state_mut<'this, I: Into<&'this mut Proxy<WlOutput>>>(proxy: I) -> &'this mut OutputState {
    unsafe {
        let user_data = proxy.into().get_user_data() as *mut OutputState;
        if user_data.is_null() {
            panic!("User data has not been set yet");
        }
        &mut *user_data
    }
}

fn unwrap_state<'this, I: Into<&'this Proxy<WlOutput>>>(proxy: I) -> &'this OutputState {
    unsafe {
        let user_data = proxy.into().get_user_data() as *const OutputState;
        if user_data.is_null() {
            panic!("User data has not been set yet");
        }
        &*user_data
    }
}
<|endoftext|>