<commit_before>// Copyright Istio 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.

use crate::identity;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::path::PathBuf;
use std::time::Duration;
use tokio::time;

#[derive(Clone, Debug)]
pub struct Config {
    pub window_size: u32,
    pub connection_window_size: u32,
    pub frame_size: u32,

    pub socks5_addr: SocketAddr,
    pub admin_addr: SocketAddr,
    pub inbound_addr: SocketAddr,
    pub inbound_plaintext_addr: SocketAddr,
    pub outbound_addr: SocketAddr,

    /// The name of the node this ztunnel is running as.
    pub local_node: Option<String>,

    /// XDS address to use. If unset, XDS will not be used.
    pub xds_address: Option<String>,
    /// Filepath to a local xds file for workloads, as YAML.
    pub local_xds_path: Option<String>,
    /// If true, on-demand XDS will be used
    pub xds_on_demand: bool,

    pub auth: identity::AuthSource,

    pub termination_grace_period: time::Duration,

    /// Specify the number of worker threads the Tokio Runtime will use.
    pub num_worker_threads: usize,
}

const DEFAULT_WORKER_THREADS: usize = 2;

impl Default for Config {
    fn default() -> Config {
        // TODO: copy JWT auth logic from CA client and use TLS here (port 15012)
        let xds_address = Some(if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012".to_string()
        } else {
            "https://localhost:15012".to_string()
        });
        Config {
            window_size: 4 * 1024 * 1024,
            connection_window_size: 4 * 1024 * 1024,
            frame_size: 1024 * 1024,

            termination_grace_period: Duration::from_secs(5),

            admin_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15021),
            socks5_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15080),
            inbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15008),
            inbound_plaintext_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15006),
            outbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15001),

            local_node: std::env::var("NODE_NAME").ok(),

            xds_address,
            local_xds_path: std::env::var("LOCAL_XDS_PATH").ok(),
            xds_on_demand: std::env::var("XDS_ON_DEMAND").ok().as_deref() == Some("on"),

            auth: identity::AuthSource::Token(PathBuf::from(
                r"./var/run/secrets/tokens/istio-token",
            )),

            num_worker_threads: std::env::var("ZTUNNEL_WORKER_THREADS")
                .ok()
                .map(|v| {
                    v.parse::<usize>()
                        .ok()
                        .filter(|n| *n > 0)
                        .unwrap_or(DEFAULT_WORKER_THREADS)
                })
                .unwrap_or(DEFAULT_WORKER_THREADS),
        }
    }
}
<commit_msg>Add option to configure address of the control plane. (#198)<commit_after>// Copyright Istio 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.

use crate::identity;
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::path::PathBuf;
use std::time::Duration;
use tokio::time;

#[derive(Clone, Debug)]
pub struct Config {
    pub window_size: u32,
    pub connection_window_size: u32,
    pub frame_size: u32,

    pub socks5_addr: SocketAddr,
    pub admin_addr: SocketAddr,
    pub inbound_addr: SocketAddr,
    pub inbound_plaintext_addr: SocketAddr,
    pub outbound_addr: SocketAddr,

    /// The name of the node this ztunnel is running as.
    pub local_node: Option<String>,

    /// XDS address to use. If unset, XDS will not be used.
    pub xds_address: Option<String>,
    /// Filepath to a local xds file for workloads, as YAML.
    pub local_xds_path: Option<String>,
    /// If true, on-demand XDS will be used
    pub xds_on_demand: bool,

    pub auth: identity::AuthSource,

    pub termination_grace_period: time::Duration,

    /// Specify the number of worker threads the Tokio Runtime will use.
    pub num_worker_threads: usize,
}

const DEFAULT_WORKER_THREADS: usize = 2;

impl Default for Config {
    fn default() -> Config {
        // TODO: copy JWT auth logic from CA client and use TLS here (port 15012)
        let xds_address = match std::env::var("XDS_ADDRESS").ok() {
            Some(xds) if xds.as_str() == "" => None,
            Some(xds) => Some(xds),
            None => {
                if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
                    Some("https://istiod.istio-system:15012".to_string())
                } else {
                    Some("https://localhost:15012".to_string())
                }
            }
        };
        Config {
            window_size: 4 * 1024 * 1024,
            connection_window_size: 4 * 1024 * 1024,
            frame_size: 1024 * 1024,

            termination_grace_period: Duration::from_secs(5),

            admin_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15021),
            socks5_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15080),
            inbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15008),
            inbound_plaintext_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15006),
            outbound_addr: SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 15001),

            local_node: std::env::var("NODE_NAME").ok(),

            xds_address,
            local_xds_path: std::env::var("LOCAL_XDS_PATH").ok(),
            xds_on_demand: std::env::var("XDS_ON_DEMAND").ok().as_deref() == Some("on"),

            auth: identity::AuthSource::Token(PathBuf::from(
                r"./var/run/secrets/tokens/istio-token",
            )),

            num_worker_threads: std::env::var("ZTUNNEL_WORKER_THREADS")
                .ok()
                .map(|v| {
                    v.parse::<usize>()
                        .ok()
                        .filter(|n| *n > 0)
                        .unwrap_or(DEFAULT_WORKER_THREADS)
                })
                .unwrap_or(DEFAULT_WORKER_THREADS),
        }
    }
}
<|endoftext|>


<commit_before>use base64;
use bridges::oidc::ProviderKey;
use config::Config;
use email_address::EmailAddress;
use openssl::bn::{BigNum, BigNumRef};
use openssl::error::ErrorStack as SslErrorStack;
use openssl::hash::{Hasher, MessageDigest};
use openssl::rsa::Rsa;
use openssl::pkey::{PKey, Public, Private};
use openssl::sign::{Signer, Verifier};
use rand::{OsRng, Rng, random};
use serde_json as json;
use std::fs::File;
use std::io::{Read, Error as IoError};
use std::iter::Iterator;
use time::now_utc;

/// Union of all possible error types seen while parsing.
#[derive(Debug)]
pub enum CryptoError {
    Custom(&'static str),
    Io(IoError),
    Ssl(SslErrorStack),
}

impl From<&'static str> for CryptoError {
    fn from(err: &'static str) -> CryptoError {
        CryptoError::Custom(err)
    }
}

impl From<IoError> for CryptoError {
    fn from(err: IoError) -> CryptoError {
        CryptoError::Io(err)
    }
}

impl From<SslErrorStack> for CryptoError {
    fn from(err: SslErrorStack) -> CryptoError {
        CryptoError::Ssl(err)
    }
}

/// A named key pair, for use in JWS signing.
pub struct NamedKey {
    id: String,
    key: PKey<Private>,
}

impl NamedKey {
    /// Creates a NamedKey by reading a `file` path and generating an `id`.
    pub fn from_file(filename: &str) -> Result<NamedKey, CryptoError> {
        let mut file = File::open(filename)?;
        let mut file_contents = String::new();
        file.read_to_string(&mut file_contents)?;

        NamedKey::from_pem_str(&file_contents)
    }

    /// Creates a NamedKey from a PEM-encoded str.
    pub fn from_pem_str(pem: &str) -> Result<NamedKey, CryptoError> {
        let rsa = Rsa::private_key_from_pem(pem.as_bytes())?;

        NamedKey::from_rsa(rsa)
    }

    /// Creates a NamedKey from an Rsa
    pub fn from_rsa(rsa: Rsa<Private>) -> Result<NamedKey, CryptoError> {
        let id = {
            let (e, n) = (rsa.e(), rsa.n());
            let mut hasher = Hasher::new(MessageDigest::sha256())?;
            let hash = hasher.update(&e.to_vec())
                .and_then(|_| hasher.update(b"."))
                .and_then(|_| hasher.update(&n.to_vec()))
                .and_then(|_| hasher.finish())?;
            base64url_encode(&hash)
        };
        let key = PKey::from_rsa(rsa)?;
        Ok(NamedKey { id, key })
    }

    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    pub fn sign_jws(&self, payload: &json::Value) -> String {
        let header = json!({
            "kid": &self.id,
            "alg": "RS256",
        }).to_string();

        let payload = payload.to_string();
        let mut input = Vec::<u8>::new();
        input.extend(base64url_encode(&header).into_bytes());
        input.push(b'.');
        input.extend(base64url_encode(&payload).into_bytes());

        let mut signer = Signer::new(MessageDigest::sha256(), &self.key)
            .expect("could not initialize signer");
        let sig = signer.update(&input)
            .and_then(|_| signer.sign_to_vec())
            .expect("failed to sign jwt");

        input.push(b'.');
        input.extend(base64url_encode(&sig).into_bytes());
        String::from_utf8(input).expect("unable to coerce jwt into string")
    }

    /// Return JSON represenation of the public key for use in JWK key sets.
    pub fn public_jwk(&self) -> json::Value {
        fn json_big_num(n: &BigNumRef) -> String {
            base64url_encode(&n.to_vec())
        }

        let rsa = self.key.rsa().expect("unable to retrieve rsa key");
        let (n, e) = (rsa.n(), rsa.e());
        json!({
            "kty": "RSA",
            "alg": "RS256",
            "use": "sig",
            "kid": &self.id,
            "n": json_big_num(n),
            "e": json_big_num(e),
        })
    }
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
pub fn session_id(email: &EmailAddress, client_id: &str) -> String {
    let mut rng = OsRng::new().expect("unable to create rng");
    let rand_bytes: Vec<u8> = (0..16).map(|_| rng.gen()).collect();

    let mut hasher = Hasher::new(MessageDigest::sha256())
        .expect("couldn't initialize SHA256 hasher");
    let hash = hasher.update(email.as_str().as_bytes())
        .and_then(|_| hasher.update(client_id.as_bytes()))
        .and_then(|_| hasher.update(&rand_bytes))
        .and_then(|_| hasher.finish())
        .expect("session hashing failed");
    base64url_encode(&hash)
}

/// Helper function to create a secure nonce.
pub fn nonce() -> String {
    let mut rng = OsRng::new().expect("unable to create rng");
    let rand_bytes: Vec<u8> = (0..16).map(|_| rng.gen()).collect();
    base64url_encode(&rand_bytes)
}

/// Helper function to create a random string consisting of
/// characters from the z-base-32 set.
pub fn random_zbase32(len: usize) -> String {
    const CHARSET: &[u8] = b"13456789abcdefghijkmnopqrstuwxyz";
    String::from_utf8((0..len).map(|_| {
        CHARSET[random::<usize>() % CHARSET.len()]
    }).collect()).expect("failed to build one-time pad")
}

/// Helper function to deserialize key from JWK Key Set.
///
/// Searches the provided JWK Key Set Value for the key matching the given
/// id. Returns a usable public key if exactly one key is found.
pub fn jwk_key_set_find(key_set: &[ProviderKey], kid: &str) -> Result<PKey<Public>, ()> {
    let matching: Vec<&ProviderKey> = key_set.iter()
        .filter(|key| key.use_ == "sig" && key.kid == kid)
        .collect();

    // Verify that we found exactly one key matching the key ID.
    if matching.len() != 1 {
        return Err(());
    }

    // Then, use the data to build a public key object for verification.
    let key = matching.first().expect("expected one key");
    let n = base64url_decode(&key.n).and_then(|data| BigNum::from_slice(&data).map_err(|_| ()))?;
    let e = base64url_decode(&key.e).and_then(|data| BigNum::from_slice(&data).map_err(|_| ()))?;
    let rsa = Rsa::from_public_components(n, e).map_err(|_| ())?;
    Ok(PKey::from_rsa(rsa).map_err(|_| ())?)
}

/// Verify a JWS signature, returning the payload as Value if successful.
pub fn verify_jws(jws: &str, key_set: &[ProviderKey]) -> Result<json::Value, ()> {
    // Extract the header from the JWT structure. Determine what key was used
    // to sign the token, so we can then verify the signature.
    let parts: Vec<&str> = jws.split('.').collect();
    if parts.len() != 3 {
        return Err(());
    }
    let decoded = parts.iter().map(|s| base64url_decode(s))
        .collect::<Result<Vec<_>, _>>()?;
    let jwt_header: json::Value = json::from_slice(&decoded[0]).map_err(|_| ())?;
    let kid = jwt_header.get("kid").and_then(|v| v.as_str()).ok_or(())?;
    let pub_key = jwk_key_set_find(key_set, kid)?;

    // Verify the identity token's signature.
    let message_len = parts[0].len() + parts[1].len() + 1;
    let mut verifier = Verifier::new(MessageDigest::sha256(), &pub_key).map_err(|_| ())?;
    verifier.update(jws[..message_len].as_bytes())
        .and_then(|_| verifier.verify(&decoded[2]))
        .map_err(|_| ())
        .and_then(|ok| {
            if ok {
                Ok(json::from_slice(&decoded[1]).map_err(|_| ())?)
            } else {
                Err(())
            }
        })
}

/// Helper method to create a JWT for a given email address and audience.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
pub fn create_jwt(app: &Config, email: &str, email_addr: &EmailAddress, aud: &str, nonce: &str) -> String {
    let now = now_utc().to_timespec().sec;
    let payload = json!({
        "aud": aud,
        "email": email_addr.as_str(),
        "email_verified": email_addr.as_str(),
        "email_original": email,
        "exp": now + i64::from(app.token_ttl),
        "iat": now,
        "iss": &app.public_url,
        "sub": email_addr.as_str(),
        "nonce": nonce,
    });
    let key = app.keys.last().expect("unable to locate signing key");
    key.sign_jws(&payload)
}

#[inline]
fn base64url_encode<T: ?Sized + AsRef<[u8]>>(data: &T) -> String {
    base64::encode_config(data, base64::URL_SAFE_NO_PAD)
}

#[inline]
fn base64url_decode<T: ?Sized + AsRef<[u8]>>(data: &T) -> Result<Vec<u8>, ()> {
    base64::decode_config(data, base64::URL_SAFE_NO_PAD).map_err(|_| ())
}
<commit_msg>return verified_email as a bool (true) per oidc recommendations<commit_after>use base64;
use bridges::oidc::ProviderKey;
use config::Config;
use email_address::EmailAddress;
use openssl::bn::{BigNum, BigNumRef};
use openssl::error::ErrorStack as SslErrorStack;
use openssl::hash::{Hasher, MessageDigest};
use openssl::rsa::Rsa;
use openssl::pkey::{PKey, Public, Private};
use openssl::sign::{Signer, Verifier};
use rand::{OsRng, Rng, random};
use serde_json as json;
use std::fs::File;
use std::io::{Read, Error as IoError};
use std::iter::Iterator;
use time::now_utc;

/// Union of all possible error types seen while parsing.
#[derive(Debug)]
pub enum CryptoError {
    Custom(&'static str),
    Io(IoError),
    Ssl(SslErrorStack),
}

impl From<&'static str> for CryptoError {
    fn from(err: &'static str) -> CryptoError {
        CryptoError::Custom(err)
    }
}

impl From<IoError> for CryptoError {
    fn from(err: IoError) -> CryptoError {
        CryptoError::Io(err)
    }
}

impl From<SslErrorStack> for CryptoError {
    fn from(err: SslErrorStack) -> CryptoError {
        CryptoError::Ssl(err)
    }
}

/// A named key pair, for use in JWS signing.
pub struct NamedKey {
    id: String,
    key: PKey<Private>,
}

impl NamedKey {
    /// Creates a NamedKey by reading a `file` path and generating an `id`.
    pub fn from_file(filename: &str) -> Result<NamedKey, CryptoError> {
        let mut file = File::open(filename)?;
        let mut file_contents = String::new();
        file.read_to_string(&mut file_contents)?;

        NamedKey::from_pem_str(&file_contents)
    }

    /// Creates a NamedKey from a PEM-encoded str.
    pub fn from_pem_str(pem: &str) -> Result<NamedKey, CryptoError> {
        let rsa = Rsa::private_key_from_pem(pem.as_bytes())?;

        NamedKey::from_rsa(rsa)
    }

    /// Creates a NamedKey from an Rsa
    pub fn from_rsa(rsa: Rsa<Private>) -> Result<NamedKey, CryptoError> {
        let id = {
            let (e, n) = (rsa.e(), rsa.n());
            let mut hasher = Hasher::new(MessageDigest::sha256())?;
            let hash = hasher.update(&e.to_vec())
                .and_then(|_| hasher.update(b"."))
                .and_then(|_| hasher.update(&n.to_vec()))
                .and_then(|_| hasher.finish())?;
            base64url_encode(&hash)
        };
        let key = PKey::from_rsa(rsa)?;
        Ok(NamedKey { id, key })
    }

    /// Create a JSON Web Signature (JWS) for the given JSON structure.
    pub fn sign_jws(&self, payload: &json::Value) -> String {
        let header = json!({
            "kid": &self.id,
            "alg": "RS256",
        }).to_string();

        let payload = payload.to_string();
        let mut input = Vec::<u8>::new();
        input.extend(base64url_encode(&header).into_bytes());
        input.push(b'.');
        input.extend(base64url_encode(&payload).into_bytes());

        let mut signer = Signer::new(MessageDigest::sha256(), &self.key)
            .expect("could not initialize signer");
        let sig = signer.update(&input)
            .and_then(|_| signer.sign_to_vec())
            .expect("failed to sign jwt");

        input.push(b'.');
        input.extend(base64url_encode(&sig).into_bytes());
        String::from_utf8(input).expect("unable to coerce jwt into string")
    }

    /// Return JSON represenation of the public key for use in JWK key sets.
    pub fn public_jwk(&self) -> json::Value {
        fn json_big_num(n: &BigNumRef) -> String {
            base64url_encode(&n.to_vec())
        }

        let rsa = self.key.rsa().expect("unable to retrieve rsa key");
        let (n, e) = (rsa.n(), rsa.e());
        json!({
            "kty": "RSA",
            "alg": "RS256",
            "use": "sig",
            "kid": &self.id,
            "n": json_big_num(n),
            "e": json_big_num(e),
        })
    }
}

/// Helper function to build a session ID for a login attempt.
///
/// Put the email address, the client ID (RP origin) and some randomness into
/// a SHA256 hash, and encode it with URL-safe bas64 encoding. This is used
/// as the key in Redis, as well as the state for OAuth authentication.
pub fn session_id(email: &EmailAddress, client_id: &str) -> String {
    let mut rng = OsRng::new().expect("unable to create rng");
    let rand_bytes: Vec<u8> = (0..16).map(|_| rng.gen()).collect();

    let mut hasher = Hasher::new(MessageDigest::sha256())
        .expect("couldn't initialize SHA256 hasher");
    let hash = hasher.update(email.as_str().as_bytes())
        .and_then(|_| hasher.update(client_id.as_bytes()))
        .and_then(|_| hasher.update(&rand_bytes))
        .and_then(|_| hasher.finish())
        .expect("session hashing failed");
    base64url_encode(&hash)
}

/// Helper function to create a secure nonce.
pub fn nonce() -> String {
    let mut rng = OsRng::new().expect("unable to create rng");
    let rand_bytes: Vec<u8> = (0..16).map(|_| rng.gen()).collect();
    base64url_encode(&rand_bytes)
}

/// Helper function to create a random string consisting of
/// characters from the z-base-32 set.
pub fn random_zbase32(len: usize) -> String {
    const CHARSET: &[u8] = b"13456789abcdefghijkmnopqrstuwxyz";
    String::from_utf8((0..len).map(|_| {
        CHARSET[random::<usize>() % CHARSET.len()]
    }).collect()).expect("failed to build one-time pad")
}

/// Helper function to deserialize key from JWK Key Set.
///
/// Searches the provided JWK Key Set Value for the key matching the given
/// id. Returns a usable public key if exactly one key is found.
pub fn jwk_key_set_find(key_set: &[ProviderKey], kid: &str) -> Result<PKey<Public>, ()> {
    let matching: Vec<&ProviderKey> = key_set.iter()
        .filter(|key| key.use_ == "sig" && key.kid == kid)
        .collect();

    // Verify that we found exactly one key matching the key ID.
    if matching.len() != 1 {
        return Err(());
    }

    // Then, use the data to build a public key object for verification.
    let key = matching.first().expect("expected one key");
    let n = base64url_decode(&key.n).and_then(|data| BigNum::from_slice(&data).map_err(|_| ()))?;
    let e = base64url_decode(&key.e).and_then(|data| BigNum::from_slice(&data).map_err(|_| ()))?;
    let rsa = Rsa::from_public_components(n, e).map_err(|_| ())?;
    Ok(PKey::from_rsa(rsa).map_err(|_| ())?)
}

/// Verify a JWS signature, returning the payload as Value if successful.
pub fn verify_jws(jws: &str, key_set: &[ProviderKey]) -> Result<json::Value, ()> {
    // Extract the header from the JWT structure. Determine what key was used
    // to sign the token, so we can then verify the signature.
    let parts: Vec<&str> = jws.split('.').collect();
    if parts.len() != 3 {
        return Err(());
    }
    let decoded = parts.iter().map(|s| base64url_decode(s))
        .collect::<Result<Vec<_>, _>>()?;
    let jwt_header: json::Value = json::from_slice(&decoded[0]).map_err(|_| ())?;
    let kid = jwt_header.get("kid").and_then(|v| v.as_str()).ok_or(())?;
    let pub_key = jwk_key_set_find(key_set, kid)?;

    // Verify the identity token's signature.
    let message_len = parts[0].len() + parts[1].len() + 1;
    let mut verifier = Verifier::new(MessageDigest::sha256(), &pub_key).map_err(|_| ())?;
    verifier.update(jws[..message_len].as_bytes())
        .and_then(|_| verifier.verify(&decoded[2]))
        .map_err(|_| ())
        .and_then(|ok| {
            if ok {
                Ok(json::from_slice(&decoded[1]).map_err(|_| ())?)
            } else {
                Err(())
            }
        })
}

/// Helper method to create a JWT for a given email address and audience.
///
/// Builds the JSON payload, then signs it using the last key provided in
/// the configuration object.
pub fn create_jwt(app: &Config, email: &str, email_addr: &EmailAddress, aud: &str, nonce: &str) -> String {
    let now = now_utc().to_timespec().sec;
    let payload = json!({
        "aud": aud,
        "email": email_addr.as_str(),
        "email_verified": true,
        "email_original": email,
        "exp": now + i64::from(app.token_ttl),
        "iat": now,
        "iss": &app.public_url,
        "sub": email_addr.as_str(),
        "nonce": nonce,
    });
    let key = app.keys.last().expect("unable to locate signing key");
    key.sign_jws(&payload)
}

#[inline]
fn base64url_encode<T: ?Sized + AsRef<[u8]>>(data: &T) -> String {
    base64::encode_config(data, base64::URL_SAFE_NO_PAD)
}

#[inline]
fn base64url_decode<T: ?Sized + AsRef<[u8]>>(data: &T) -> Result<Vec<u8>, ()> {
    base64::decode_config(data, base64::URL_SAFE_NO_PAD).map_err(|_| ())
}
<|endoftext|>


<commit_before>use std::thread;
use std::time::Duration;
use time::precise_time_ns;

use bootstrap::{Action, Context, Key, Keyboard, Mouse, MouseButton, MouseMove, Scroll, Window};
use camera::{Camera, Freefly};
use transform::Translation;

/// All common stuff goes here.
pub struct App {
  /// Some kind of epoch start the application started at.
  start_time: u64,
  /// Keyboard receiver.
  kbd: Keyboard,
  /// Mouse receiver.
  mouse: Mouse,
  /// Cursor receiver.
  cursor: MouseMove,
  /// Scroll receiver.
  scroll: Scroll,
  /// Window.
  window: Window,
  /// Last known cursor position.
  last_cursor: [f64; 2],
  /// Is the mouse’s left button pressed down?
  left_down: bool,
  /// Is the mouse’s right button pressed down?
  right_down: bool,
}

impl App {
  pub fn init(kbd: Keyboard, mouse: Mouse, cursor: MouseMove, scroll: Scroll, window: Window) -> Self {
    App {
      start_time: precise_time_ns(),
      kbd: kbd,
      mouse: mouse,
      cursor: cursor,
      scroll: scroll,
      window: window,
      last_cursor: [0., 0.],
      left_down: false,
      right_down: false,
    }
  }

  pub fn time(&self) -> f32 {
    (precise_time_ns() - self.start_time) as f32 * 1e-9
  }

  pub fn handle_events<'a>(&mut self, camera: &mut Camera<Freefly>) -> bool {
    while let Ok((key, action)) = self.kbd.try_recv() {
      if action == Action::Release && key == Key::Escape{
        return false;
      } else {
        match key {
          Key::W => camera.mv(Translation::new(0., 0., 1.)),
          Key::S => camera.mv(Translation::new(0., 0., -1.)),
          Key::A => camera.mv(Translation::new(1., 0., 0.)),
          Key::D => camera.mv(Translation::new(-1., 0., 0.)),
          Key::Q => camera.mv(Translation::new(0., -1., 0.)),
          Key::E => camera.mv(Translation::new(0., 1., 0.)),
          _ => {}
        }
      }
    }

    while let Ok((button, action)) = self.mouse.try_recv() {
      match (button, action) {
        (MouseButton::Button1, Action::Press) => {
          self.left_down = true;
        },
        (MouseButton::Button1, Action::Release) => {
          self.left_down = false;
        },
        (MouseButton::Button2, Action::Press) => {
          self.right_down = true;
        },
        (MouseButton::Button2, Action::Release) => {
          self.right_down = false;
        },
        _ => {}
      }
    }

    while let Ok(xy) = self.cursor.try_recv() {
      if self.left_down {
        let (dx, dy) = (xy[0] - self.last_cursor[0], xy[1] - self.last_cursor[1]);
        camera.look_around(Translation::new(dy as f32, dx as f32, 0.));
      } else if self.right_down {
        let (dx, _) = (xy[0] - self.last_cursor[0], xy[1] - self.last_cursor[1]);
        camera.look_around(Translation::new(0., 0., dx as f32));
      }

      self.last_cursor = xy;
    }

    true
  }

  pub fn step<R>(&mut self, mut draw_frame: R) -> bool where R: FnMut(f32) {
    let loop_start_time = precise_time_ns();

    if self.window.should_close() {
      return false;
    }

    let t = self.time();

    draw_frame(t);

    self.window.swap_buffers();

    // wait for next frame according to the wished FPS
    let fps = 60.;
    let elapsed_time_ms = ((precise_time_ns() - loop_start_time) as f64 * 1e-6) as i64;
    let sleep_time_ms = (1. / fps * 1e3) as i64 - elapsed_time_ms;

    if sleep_time_ms > 0 {
      thread::sleep(Duration::from_millis(sleep_time_ms as u64));
    }

    true
  }
}
<commit_msg>Removed useless lifetime.<commit_after>use std::thread;
use std::time::Duration;
use time::precise_time_ns;

use bootstrap::{Action, Context, Key, Keyboard, Mouse, MouseButton, MouseMove, Scroll, Window};
use camera::{Camera, Freefly};
use transform::Translation;

/// All common stuff goes here.
pub struct App {
  /// Some kind of epoch start the application started at.
  start_time: u64,
  /// Keyboard receiver.
  kbd: Keyboard,
  /// Mouse receiver.
  mouse: Mouse,
  /// Cursor receiver.
  cursor: MouseMove,
  /// Scroll receiver.
  scroll: Scroll,
  /// Window.
  window: Window,
  /// Last known cursor position.
  last_cursor: [f64; 2],
  /// Is the mouse’s left button pressed down?
  left_down: bool,
  /// Is the mouse’s right button pressed down?
  right_down: bool,
}

impl App {
  pub fn init(kbd: Keyboard, mouse: Mouse, cursor: MouseMove, scroll: Scroll, window: Window) -> Self {
    App {
      start_time: precise_time_ns(),
      kbd: kbd,
      mouse: mouse,
      cursor: cursor,
      scroll: scroll,
      window: window,
      last_cursor: [0., 0.],
      left_down: false,
      right_down: false,
    }
  }

  pub fn time(&self) -> f32 {
    (precise_time_ns() - self.start_time) as f32 * 1e-9
  }

  pub fn handle_events(&mut self, camera: &mut Camera<Freefly>) -> bool {
    while let Ok((key, action)) = self.kbd.try_recv() {
      if action == Action::Release && key == Key::Escape{
        return false;
      } else {
        match key {
          Key::W => camera.mv(Translation::new(0., 0., 1.)),
          Key::S => camera.mv(Translation::new(0., 0., -1.)),
          Key::A => camera.mv(Translation::new(1., 0., 0.)),
          Key::D => camera.mv(Translation::new(-1., 0., 0.)),
          Key::Q => camera.mv(Translation::new(0., -1., 0.)),
          Key::E => camera.mv(Translation::new(0., 1., 0.)),
          _ => {}
        }
      }
    }

    while let Ok((button, action)) = self.mouse.try_recv() {
      match (button, action) {
        (MouseButton::Button1, Action::Press) => {
          self.left_down = true;
        },
        (MouseButton::Button1, Action::Release) => {
          self.left_down = false;
        },
        (MouseButton::Button2, Action::Press) => {
          self.right_down = true;
        },
        (MouseButton::Button2, Action::Release) => {
          self.right_down = false;
        },
        _ => {}
      }
    }

    while let Ok(xy) = self.cursor.try_recv() {
      if self.left_down {
        let (dx, dy) = (xy[0] - self.last_cursor[0], xy[1] - self.last_cursor[1]);
        camera.look_around(Translation::new(dy as f32, dx as f32, 0.));
      } else if self.right_down {
        let (dx, _) = (xy[0] - self.last_cursor[0], xy[1] - self.last_cursor[1]);
        camera.look_around(Translation::new(0., 0., dx as f32));
      }

      self.last_cursor = xy;
    }

    true
  }

  pub fn step<R>(&mut self, mut draw_frame: R) -> bool where R: FnMut(f32) {
    let loop_start_time = precise_time_ns();

    if self.window.should_close() {
      return false;
    }

    let t = self.time();

    draw_frame(t);

    self.window.swap_buffers();

    // wait for next frame according to the wished FPS
    let fps = 60.;
    let elapsed_time_ms = ((precise_time_ns() - loop_start_time) as f64 * 1e-6) as i64;
    let sleep_time_ms = (1. / fps * 1e3) as i64 - elapsed_time_ms;

    if sleep_time_ms > 0 {
      thread::sleep(Duration::from_millis(sleep_time_ms as u64));
    }

    true
  }
}
<|endoftext|>


<commit_before>use std::sync::*;
use rand::*;
use types::*;
use player::*;
use zobrist::*;
use field::*;
use uct::*;

pub struct Bot {
  rng: XorShiftRng,
  zobrist: Arc<Zobrist>,
  field: Field,
  uct: UctRoot
}

impl Bot {
  pub fn new(width: Coord, height: Coord) -> Bot {
    let length = length(width, height);
    let mut rng = XorShiftRng::new_unseeded();
    let zobrist = Arc::new(Zobrist::new(length, &mut rng));
    let field_zobrist = zobrist.clone();
    Bot {
      rng: rng,
      zobrist: zobrist,
      field: Field::new(width, height, field_zobrist),
      uct: UctRoot::new(length)
    }
  }

  pub fn best_move(&mut self, player: Player, time: Time) -> Option<(Coord, Coord)> {
    self.uct.best_move(&self.field, player, &mut self.rng, time).map(|pos| (self.field.to_x(pos), self.field.to_y(pos)))
  }

  pub fn put_point(&mut self, x: Coord, y: Coord, player: Player) -> bool {
    let pos = self.field.to_pos(x, y);
    self.field.put_point(pos, player)
  }

  pub fn undo(&mut self) -> bool {
    self.field.undo()
  }
}
<commit_msg>Bugfix.<commit_after>use std::sync::*;
use rand::*;
use types::*;
use player::*;
use zobrist::*;
use field::*;
use uct::*;

pub struct Bot {
  rng: XorShiftRng,
  zobrist: Arc<Zobrist>,
  field: Field,
  uct: UctRoot
}

impl Bot {
  pub fn new(width: Coord, height: Coord) -> Bot {
    let length = length(width, height);
    let mut rng = XorShiftRng::new_unseeded();
    let zobrist = Arc::new(Zobrist::new(length * 2, &mut rng));
    let field_zobrist = zobrist.clone();
    Bot {
      rng: rng,
      zobrist: zobrist,
      field: Field::new(width, height, field_zobrist),
      uct: UctRoot::new(length)
    }
  }

  pub fn best_move(&mut self, player: Player, time: Time) -> Option<(Coord, Coord)> {
    self.uct.best_move(&self.field, player, &mut self.rng, time).map(|pos| (self.field.to_x(pos), self.field.to_y(pos)))
  }

  pub fn put_point(&mut self, x: Coord, y: Coord, player: Player) -> bool {
    let pos = self.field.to_pos(x, y);
    self.field.put_point(pos, player)
  }

  pub fn undo(&mut self) -> bool {
    self.field.undo()
  }
}
<|endoftext|>


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

// FIXME: Try the efficient algorithm (takes only O(n) time with additional O(n) space)
// https://de.wikipedia.org/wiki/Suffix-Array-Induced-Sorting
pub fn bwt(mut input: Vec<u8>) -> Vec<u8> {
    input.push(0);

    // get the sorted suffix array
    let mut rotations = (0..input.len()).map(|i| &input[i..]).collect::<Vec<_>>();
    rotations.sort();
    let mut rots = vec![0; input.len()];

    // BWT[i] = S[SA[i] - 1]
    for i in 0..input.len() {
        if rotations[i].len() == input.len() {
            rots[i] = 0;
        } else {
            rots[i] = input[input.len() - rotations[i].len() - 1];
        }
    }

    rots
}

pub fn occurrence_index(input: &[u8]) -> HashMap<u8, usize> {
    // get the byte distribution
    let mut map = HashMap::new();
    for i in input {
        let mut count = map.entry(*i).or_insert(0);
        *count += 1;
    }

    // sort the bytes and locate the index of their first occurrences
    let mut chars = map.keys().collect::<Vec<_>>();
    chars.sort();

    let mut idx = 0;
    let mut occ_map = HashMap::new();
    for i in chars {
        occ_map.insert(*i, idx);
        idx += *map.get(i).unwrap();
    }

    occ_map
}

pub fn ibwt(input: Vec<u8>) -> Vec<u8> {
    let mut occ_map = occurrence_index(&input);

    // generate the LF vector
    let mut lf = vec![0; input.len()];
    for (i, c) in input.iter().enumerate() {
        let mut val = occ_map.get_mut(&c).unwrap();
        lf[i] = *val;
        *val += 1;
    }

    let mut idx = 0;
    // construct the sequence by traversing through the LF vector
    let mut output = vec![0; input.len()];
    for i in (0..(input.len() - 1)).rev() {
        output[i] = input[idx];
        idx = lf[idx];
    }

    output.pop();
    output
}

pub struct FMIndex {
    data: Vec<u8>,
    occ_map: HashMap<u8, usize>,
}

impl FMIndex {
    pub fn new(data: Vec<u8>) -> FMIndex {
        let bwt_data = bwt(data);
        FMIndex {
            occ_map: occurrence_index(&bwt_data),
            data: bwt_data,
        }
    }

    fn nearest_lf(&self, idx: usize, ch: u8) -> usize {
        match self.occ_map.get(&ch) {
            Some(occ) => occ + self.data[0..idx].iter().filter(|&i| *i == ch).count(),
            None => 0,
        }
    }

    // FIXME: We need "checkpointed" index!
    pub fn search(&mut self, query: &str) -> Vec<usize> {
        let (mut top, mut bottom) = (0, self.data.len());
        for ch in query.as_bytes().iter().rev() {
            top = self.nearest_lf(top, *ch);
            bottom = self.nearest_lf(bottom, *ch);
        }

        let mut results = vec![];

        for idx in top..bottom {
            let mut pos = 0;
            let mut i = idx;
            while self.data[i] != 0 {
                pos += 1;
                i = self.nearest_lf(i, self.data[i]);
            }

            results.push(pos);
        }

        return results
    }
}
<commit_msg>Cache the occurrence indices when building FM index<commit_after>use std::collections::HashMap;

// FIXME: Try the efficient algorithm (takes only O(n) time with additional O(n) space)
// https://de.wikipedia.org/wiki/Suffix-Array-Induced-Sorting
pub fn bwt<F: FnMut(u8)>(mut input: Vec<u8>, mut f: F) -> Vec<u8> {
    input.push(0);

    // get the sorted suffix array
    let mut rotations = (0..input.len()).map(|i| &input[i..]).collect::<Vec<_>>();
    rotations.sort();
    let mut rots = vec![0; input.len()];

    // BWT[i] = S[SA[i] - 1]
    for i in 0..input.len() {
        if rotations[i].len() == input.len() {
            rots[i] = 0;
        } else {
            rots[i] = input[input.len() - rotations[i].len() - 1];
        }

        f(rots[i]);     // call the function with the final value
    }

    rots
}

pub fn occurrence_index(map: HashMap<u8, usize>) -> HashMap<u8, usize> {
    // sort the bytes and locate the index of their first occurrences
    let mut chars = map.keys().collect::<Vec<_>>();
    chars.sort();

    let mut idx = 0;
    let mut occ_map = HashMap::new();
    for i in chars {
        occ_map.insert(*i, idx);
        idx += *map.get(i).unwrap();
    }

    occ_map
}

pub fn ibwt(input: Vec<u8>) -> Vec<u8> {
    // get the byte distribution
    let mut map = HashMap::new();
    for i in &input {
        let mut count = map.entry(*i).or_insert(0);
        *count += 1;
    }

    let mut occ_map = occurrence_index(map);

    // generate the LF vector
    let mut lf = vec![0; input.len()];
    for (i, c) in input.iter().enumerate() {
        let mut val = occ_map.get_mut(&c).unwrap();
        lf[i] = *val;
        *val += 1;
    }

    let mut idx = 0;
    // construct the sequence by traversing through the LF vector
    let mut output = vec![0; input.len()];
    for i in (0..(input.len() - 1)).rev() {
        output[i] = input[idx];
        idx = lf[idx];
    }

    output.pop();
    output
}

pub struct FMIndex {
    data: Vec<u8>,
    count_cache: Vec<usize>,    // remember the occurrence of each char (caching)
    occ_map: HashMap<u8, usize>,
}

impl FMIndex {
    pub fn new(data: Vec<u8>) -> FMIndex {
        let mut idx = 0;
        let mut map = HashMap::new();
        let mut count = vec![0; data.len() + 1];
        let bwt_data = bwt(data, |i| {
            let mut c = map.entry(i).or_insert(0);
            *c += 1;
            count[idx] = *c;
            idx += 1;
        });

        FMIndex {
            occ_map: occurrence_index(map),
            count_cache: count,
            data: bwt_data,
        }
    }

    fn nearest(&self, idx: usize, ch: u8) -> usize {
        match self.occ_map.get(&ch) {
            Some(occ) => {
                occ + (0..idx).rev()
                              .find(|&i| self.data[i] == ch)
                              .map(|i| self.count_cache[i])
                              .unwrap_or(0)
            },
            None => 0,
        }
    }

    // FIXME: We need "checkpointed" index!
    pub fn search(&mut self, query: &str) -> Option<usize> {
        let (mut top, mut bottom) = (0, self.data.len());
        for ch in query.as_bytes().iter().rev() {
            top = self.nearest(top, *ch);
            bottom = self.nearest(bottom, *ch);
        }

        for idx in top..bottom {
            let mut pos = 0;
            let mut i = idx;
            while self.data[i] != 0 {
                pos += 1;
                i = self.nearest(i, self.data[i]);
            }

            return Some(pos)    // break on first find
        }

        None
    }
}
<|endoftext|>


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

use address::Address;
use basicblock::{BasicBlock, BasicBlockEdge, EdgeType};
use instruction::Instruction;
use petgraph::graph::{Graph, NodeIndex};
use std::collections::BTreeMap;

/// A [control flow graph].
///
/// [control flow graph]: https://en.wikipedia.org/wiki/Control_flow_graph
pub struct CFG<'f> {
    /// The [`Graph`] that stores the actual CFG
    ///
    /// [`Graph`]: ../petgraph/graph/struct.Graph.html
    pub graph: Graph<BasicBlock<'f>, BasicBlockEdge>,
    /// The [`NodeIndex`] for the entry [`BasicBlock`] for this function.
    ///
    /// [`BasicBlock`]: struct.BasicBlock.html
    /// [`NodeIndex`]: ../petgraph/graph/struct.NodeIndex.html
    pub entry_block: Option<NodeIndex>,
    /// Map address to the corresponding basic block.
    pub block_finder: BTreeMap<Address, NodeIndex>,
}

impl<'f> CFG<'f> {
    /// Construct a new CFG
    pub fn new() -> Self {
        CFG {
            graph: Graph::new(),
            entry_block: None,
            block_finder: BTreeMap::new(),
        }
    }

    /// Build the CFG from the [`instructions`].
    ///
    /// This is conducted in a 2 step process:
    ///
    /// * First, each instruction is examined to identify block boundaries.
    /// * Then, we go through each instruction again, looking for the
    ///   previously identified block boundaries and build the edges.
    ///
    /// This two step process prevents us from having to construct and then
    /// subsequently split blocks as we find backward edges.
    ///
    /// [`instructions`]: trait.Instruction.html
    pub fn build(&mut self, instructions: &[Box<Instruction>]) {
        if instructions.is_empty() {
            return;
        }
        self.identify_blocks(instructions);
        self.build_edges(instructions);
    }

    /// Identify basic blocks by their boundaries.
    ///
    /// Inspect each instruction to see if it is a 'leader' or the start of
    /// a new basic block. An instruction is a leader if:
    ///
    /// * It is the first instruction in the basic block.
    /// * It comes after a branch of any sort. We identify these as things
    ///   for which `Instruction::is_block_terminator` returns `true`.
    /// * It is the target of a jump (conditional or unconditional) within
    ///   the function.
    fn identify_blocks(&mut self, instructions: &[Box<Instruction>]) {
        let start_addr = instructions[0].address();
        let end_addr = instructions.last().map(|i| i.address()).unwrap();
        let mut next_is_leader: bool = true;
        for inst in instructions {
            if next_is_leader {
                let idx = self.graph.add_node(BasicBlock::new(inst.address()));
                self.block_finder.insert(inst.address(), idx);
                next_is_leader = false;
            }
            if inst.is_block_terminator() {
                if let Some(target_addr) = inst.target_address() {
                    if !self.block_finder.contains_key(&target_addr) &&
                       target_addr >= start_addr && target_addr <= end_addr {
                        let idx = self.graph.add_node(BasicBlock::new(target_addr));
                        self.block_finder.insert(target_addr, idx);
                    }
                }
                // The next instruction, if any, will be the start of a new block.
                next_is_leader = true;
            }
        }
        self.entry_block = Some(self.block_finder.get(&instructions[0].address()).unwrap().clone());
    }

    /// Build an edge between 2 basic blocks.
    fn build_edge(&mut self,
                  current_block_idx: NodeIndex,
                  next_block_idx: NodeIndex,
                  current_inst: &Box<Instruction>) {
        if current_inst.is_local_conditional_jump() {
            // We have one edge for the jump target and one for the fallthrough.
            // We jump through some hoops here to keep the borrow checker happy.
            if let Some(target_addr) = current_inst.target_address() {
                if let Some(target_block_idx) = self.block_finder.get(&target_addr) {
                    let edge = BasicBlockEdge { edge_type: EdgeType::ConditionalTaken };
                    self.graph.add_edge(current_block_idx, target_block_idx.clone(), edge);
                }
                let edge = BasicBlockEdge { edge_type: EdgeType::ConditionalFallthrough };
                self.graph.add_edge(current_block_idx, next_block_idx, edge);
            }
        } else if current_inst.is_local_jump() || current_inst.is_call() {
            let edge = BasicBlockEdge { edge_type: EdgeType::Unconditional };
            self.graph.add_edge(current_block_idx, next_block_idx, edge);
        } else if current_inst.is_return() {
            // Do we want to record this exit anywhere?
        }
    }

    /// Build edges after having identified the basic blocks.
    ///
    /// After the basic blocks have been identified, we can go back and
    /// connect the basic blocks with their appropriate edges.
    ///
    /// We do this by iterating through the instructions looking for
    /// boundaries between the basic blocks and then setting up the
    /// new edges.
    fn build_edges(&mut self, instructions: &[Box<Instruction>]) {
        // Here, we're going to walk through the instructions again in
        // pairs of (i, Some(i + 1)):
        // [1,2..9] -> [Some(1),Some(2)..Some(9),None]
        let current_inst_iter = instructions.iter();
        let mut next_inst_iter = current_inst_iter.clone().map(Some).chain(Some(None).into_iter());
        // Skip the first one so that we're actually working with the next one.
        next_inst_iter.next();
        let mut current_block_idx = self.entry_block.unwrap();
        for ij in current_inst_iter.zip(next_inst_iter) {
            if let (current_inst, Some(next_inst)) = ij {
                // Does the next instruction begin a basic block?
                let next_block_idx = self.block_finder
                    .get(&next_inst.address())
                    .unwrap_or(&current_block_idx)
                    .clone();
                // If we're at a block boundary, create an edge between the
                // current and next blocks. The type of the edge is determined
                // by looking at the current instruction.
                if next_block_idx != current_block_idx {
                    self.build_edge(current_block_idx, next_block_idx, current_inst);
                    current_block_idx = next_block_idx;
                }
            } else {
                // No next instruction, so we're at the end.
                // Do we want to record this exit anywhere?
            }
        }
    }
}

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

    #[test]
    fn construct() {
        let cfg = CFG::new();
        assert!(cfg.entry_block.is_none());
        assert_eq!(cfg.graph.node_count(), 0);
    }

    #[test]
    fn build_one_basic_block() {
        let mut cfg = CFG::new();
        let insts: Vec<Box<Instruction>> = vec![
            TestInstruction::new(0, Opcode::Add),
            TestInstruction::new(1, Opcode::Add),
            TestInstruction::new(2, Opcode::Ret),
        ];
        cfg.build(&insts);
        assert!(cfg.entry_block.is_some());
        assert_eq!(cfg.graph.node_count(), 1);
    }
}
<commit_msg>Test the CFG's externals.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use address::Address;
use basicblock::{BasicBlock, BasicBlockEdge, EdgeType};
use instruction::Instruction;
use petgraph::graph::{Graph, NodeIndex};
use std::collections::BTreeMap;

/// A [control flow graph].
///
/// [control flow graph]: https://en.wikipedia.org/wiki/Control_flow_graph
pub struct CFG<'f> {
    /// The [`Graph`] that stores the actual CFG
    ///
    /// [`Graph`]: ../petgraph/graph/struct.Graph.html
    pub graph: Graph<BasicBlock<'f>, BasicBlockEdge>,
    /// The [`NodeIndex`] for the entry [`BasicBlock`] for this function.
    ///
    /// [`BasicBlock`]: struct.BasicBlock.html
    /// [`NodeIndex`]: ../petgraph/graph/struct.NodeIndex.html
    pub entry_block: Option<NodeIndex>,
    /// Map address to the corresponding basic block.
    pub block_finder: BTreeMap<Address, NodeIndex>,
}

impl<'f> CFG<'f> {
    /// Construct a new CFG
    pub fn new() -> Self {
        CFG {
            graph: Graph::new(),
            entry_block: None,
            block_finder: BTreeMap::new(),
        }
    }

    /// Build the CFG from the [`instructions`].
    ///
    /// This is conducted in a 2 step process:
    ///
    /// * First, each instruction is examined to identify block boundaries.
    /// * Then, we go through each instruction again, looking for the
    ///   previously identified block boundaries and build the edges.
    ///
    /// This two step process prevents us from having to construct and then
    /// subsequently split blocks as we find backward edges.
    ///
    /// [`instructions`]: trait.Instruction.html
    pub fn build(&mut self, instructions: &[Box<Instruction>]) {
        if instructions.is_empty() {
            return;
        }
        self.identify_blocks(instructions);
        self.build_edges(instructions);
    }

    /// Identify basic blocks by their boundaries.
    ///
    /// Inspect each instruction to see if it is a 'leader' or the start of
    /// a new basic block. An instruction is a leader if:
    ///
    /// * It is the first instruction in the basic block.
    /// * It comes after a branch of any sort. We identify these as things
    ///   for which `Instruction::is_block_terminator` returns `true`.
    /// * It is the target of a jump (conditional or unconditional) within
    ///   the function.
    fn identify_blocks(&mut self, instructions: &[Box<Instruction>]) {
        let start_addr = instructions[0].address();
        let end_addr = instructions.last().map(|i| i.address()).unwrap();
        let mut next_is_leader: bool = true;
        for inst in instructions {
            if next_is_leader {
                let idx = self.graph.add_node(BasicBlock::new(inst.address()));
                self.block_finder.insert(inst.address(), idx);
                next_is_leader = false;
            }
            if inst.is_block_terminator() {
                if let Some(target_addr) = inst.target_address() {
                    if !self.block_finder.contains_key(&target_addr) &&
                       target_addr >= start_addr && target_addr <= end_addr {
                        let idx = self.graph.add_node(BasicBlock::new(target_addr));
                        self.block_finder.insert(target_addr, idx);
                    }
                }
                // The next instruction, if any, will be the start of a new block.
                next_is_leader = true;
            }
        }
        self.entry_block = Some(self.block_finder.get(&instructions[0].address()).unwrap().clone());
    }

    /// Build an edge between 2 basic blocks.
    fn build_edge(&mut self,
                  current_block_idx: NodeIndex,
                  next_block_idx: NodeIndex,
                  current_inst: &Box<Instruction>) {
        if current_inst.is_local_conditional_jump() {
            // We have one edge for the jump target and one for the fallthrough.
            // We jump through some hoops here to keep the borrow checker happy.
            if let Some(target_addr) = current_inst.target_address() {
                if let Some(target_block_idx) = self.block_finder.get(&target_addr) {
                    let edge = BasicBlockEdge { edge_type: EdgeType::ConditionalTaken };
                    self.graph.add_edge(current_block_idx, target_block_idx.clone(), edge);
                }
                let edge = BasicBlockEdge { edge_type: EdgeType::ConditionalFallthrough };
                self.graph.add_edge(current_block_idx, next_block_idx, edge);
            }
        } else if current_inst.is_local_jump() || current_inst.is_call() {
            let edge = BasicBlockEdge { edge_type: EdgeType::Unconditional };
            self.graph.add_edge(current_block_idx, next_block_idx, edge);
        } else if current_inst.is_return() {
            // Do we want to record this exit anywhere?
        }
    }

    /// Build edges after having identified the basic blocks.
    ///
    /// After the basic blocks have been identified, we can go back and
    /// connect the basic blocks with their appropriate edges.
    ///
    /// We do this by iterating through the instructions looking for
    /// boundaries between the basic blocks and then setting up the
    /// new edges.
    fn build_edges(&mut self, instructions: &[Box<Instruction>]) {
        // Here, we're going to walk through the instructions again in
        // pairs of (i, Some(i + 1)):
        // [1,2..9] -> [Some(1),Some(2)..Some(9),None]
        let current_inst_iter = instructions.iter();
        let mut next_inst_iter = current_inst_iter.clone().map(Some).chain(Some(None).into_iter());
        // Skip the first one so that we're actually working with the next one.
        next_inst_iter.next();
        let mut current_block_idx = self.entry_block.unwrap();
        for ij in current_inst_iter.zip(next_inst_iter) {
            if let (current_inst, Some(next_inst)) = ij {
                // Does the next instruction begin a basic block?
                let next_block_idx = self.block_finder
                    .get(&next_inst.address())
                    .unwrap_or(&current_block_idx)
                    .clone();
                // If we're at a block boundary, create an edge between the
                // current and next blocks. The type of the edge is determined
                // by looking at the current instruction.
                if next_block_idx != current_block_idx {
                    self.build_edge(current_block_idx, next_block_idx, current_inst);
                    current_block_idx = next_block_idx;
                }
            } else {
                // No next instruction, so we're at the end.
                // Do we want to record this exit anywhere?
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use instruction::Instruction;
    use petgraph::EdgeDirection;
    use super::CFG;
    use tests::*;

    #[test]
    fn construct() {
        let cfg = CFG::new();
        assert!(cfg.entry_block.is_none());
        assert_eq!(cfg.graph.node_count(), 0);
    }

    #[test]
    fn build_one_basic_block() {
        let mut cfg = CFG::new();
        let insts: Vec<Box<Instruction>> = vec![
            TestInstruction::new(0, Opcode::Add),
            TestInstruction::new(1, Opcode::Add),
            TestInstruction::new(2, Opcode::Ret),
        ];
        cfg.build(&insts);
        assert!(cfg.entry_block.is_some());
        assert_eq!(cfg.graph.node_count(), 1);

        let inbound = cfg.graph.externals(EdgeDirection::Incoming);
        assert_eq!(inbound.count(), 1);

        let outbound = cfg.graph.externals(EdgeDirection::Outgoing);
        assert_eq!(outbound.count(), 1);
    }
}
<|endoftext|>


<commit_before>//!
//! The main `Desync` struct
//!

use super::scheduler::*;

use std::sync::{Arc};
use std::marker::{Unpin};
use futures::channel::oneshot;
use futures::future::{Future, BoxFuture};

use std::mem;
use std::result::{Result};

///
/// A data storage structure used to govern synchronous and asynchronous access to an underlying object.
///
pub struct Desync<T: Send+Unpin> {
    /// Queue used for scheduling runtime for this object
    queue:  Arc<JobQueue>,

    /// Data for this object. Boxed so the pointer remains the same through the lifetime of the object.
    data:   *mut T
}

unsafe impl<T: Send+Unpin> Send for Desync<T> {}

// True iff queue: Sync
unsafe impl<T: Send+Unpin> Sync for Desync<T> {}

///
/// Used for passing the data pointer through to the queue
///
/// 'Safe' because the queue is synchronised during drop, so we can never use the pointer
/// if the object does not exist.
///
struct DataRef<T: Send>(*mut T);
unsafe impl<T: Send> Send for DataRef<T> {}

// TODO: we can change DataRef to Shared (https://doc.rust-lang.org/std/ptr/struct.Shared.html in the future)

// TODO: T does not need to be static as we know that its lifetime is at least the lifetime of Desync<T> and hence the queue
impl<T: 'static+Send+Unpin> Desync<T> {
    ///
    /// Creates a new Desync object
    ///
    pub fn new(data: T) -> Desync<T> {
        let queue = queue();

        Desync {
            queue:  queue,
            data:   Box::into_raw(Box::new(data))
        }
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    #[inline]
    #[deprecated(since="0.3.0", note="please use `desync` instead")]
    pub fn r#async<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        self.desync(job)
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    pub fn desync<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        // As drop() is the last thing called, we know that this object will still exist at the point where the queue makes the asynchronous callback
        let data = DataRef::<T>(self.data);

        desync(&self.queue, move || {
            let data = data.0;
            job(unsafe { &mut *data });
        })
    }

    ///
    /// Performs an operation synchronously on this item. This will be queued with any other
    /// jobs that this item may be performing, and this function will not return until the
    /// job is complete and the result is available.
    ///
    pub fn sync<TFn, Result>(&self, job: TFn) -> Result
    where TFn: Send+FnOnce(&mut T) -> Result, Result: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            // Exclusivity is guaranteed because the queue executes only one task at a time
            let data = DataRef::<T>(self.data);

            sync(&self.queue, move || {
                let data = data.0;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Performs an operation synchronously on this item,
    ///
    pub fn try_sync<TFn, FnResult>(&self, job: TFn) -> Result<FnResult, TrySyncError>
    where TFn: Send+FnOnce(&mut T) -> FnResult, FnResult: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            let data = DataRef::<T>(self.data);

            try_sync(&self.queue, move || {
                let data = data.0;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Deprecated name for `future_desync`
    ///
    #[inline]
    #[deprecated(since="0.7.0", note="please use either `future_desync` or `future_sync` to schedule futures")]
    pub fn future<TFn, TOutput>(&self, job: TFn) -> impl Future<Output=Result<TOutput, oneshot::Canceled>>+Send
    where   TFn:        'static+Send+for<'a> FnOnce(&'a mut T) -> BoxFuture<'a, TOutput>,
            TOutput:    'static+Send {
        self.future_desync(job)
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning the
    /// result via a future.
    ///
    /// The future will be scheduled in the background, so it will make progress even if the current scheduler is
    /// blocked for any reason. Additionally, it's not necessary to await the returned future, which can be discarded
    /// if necessary. As soon as this function returns, the order of the operations in the future are fixed: that is,
    /// any other operation queued on the same `Desync` will be performed after the future completes.
    ///
    /// Call `detach()` on the returned future to fully put it in the background. If awaited the scheduler may try to
    /// 'steal' runtime on the current thread to complete the future, but it will also run the future in the background
    /// if the current thread is not available or a background thread is available faster. The future here has static
    /// lifetime so can be passed around in case something needs to monitor when it completes.
    ///
    /// This can also be used to run the future synchronously on the current thread, useful for cases where an async
    /// function needs to be called from a non-async context. Call `sync()` on the returned future to wait for its
    /// result synchronously.
    ///
    /// Note that this has a companion function, `future_sync()` which is more useful for cases where the future is
    /// awaited, as the 'inner' future is able to borrow values from the 'outer' future due to the shorter lifetime
    /// of the returned value.
    ///
    /// The normal `sync()` and `desync()` methods will schedule their operations in order around any `future_sync()`
    /// or `future_desync()` calls, so it's possible to easily mix async and traditional threaded code in one program.
    ///
    pub fn future_desync<'a, TFn, TFuture>(&self, job: TFn) -> SchedulerFuture<TFuture::Output>
    where
        TFn:                'static + Send + FnOnce(&'a mut T) -> TFuture,
        TFuture:            'a + Send + Future,
        TFuture::Output:    'static + Send,
    {
        // The future will have a lifetime shorter than the lifetime of this structure, and exclusivity is guaranteed
        // because queues only execute one task at a time
        let data = DataRef::<T>(self.data);

        scheduler().future_desync(&self.queue, move || {
            let data        = data.0;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning a future that must be awaited
    /// before it is dropped.
    ///
    /// The future will be scheduled in the current execution context, so it will only make progress if the current
    /// scheduler is running. The task will be cancelled and will not complete execution if the future is dropped before
    /// completion, so it's usually necessary to await the result of this function for the task to behave correctly.
    ///
    /// This is most useful for 'foreground' futures: note that the future can safely borrow values from the outer context
    /// to use while it's running (in much the same way `sync()` can in synchronous contexts). The need to dispatch to
    /// the awaiting context can make this future slightly slower than `future_desync()`
    ///
    /// Desync has a feature that makes it possible to wait synchronously for futures to complete: slightly
    /// counter-intuitively this is found in the future returned by `future_desync()` function rather than this function,
    /// which is synchronous with the current async context instead of with the current thread.
    ///
    /// The normal `sync()` and `desync()` methods will schedule their operations in order around any `future_sync()`
    /// or `future_desync()` calls, so it's possible to easily mix async and traditional threaded code in one program.
    ///
    pub fn future_sync<'a, 'b, TFn, TFuture>(&'a self, job: TFn) -> impl 'a + Future<Output=Result<TFuture::Output, oneshot::Canceled>> + Send
    where
        TFn:                'a + Send + FnOnce(&'b mut T) -> TFuture,
        TFuture:            'b + Send + Future,
        TFuture::Output:    'a + Send,
        'b:                 'a,
    {
        // The future will have a lifetime shorter than the lifetime of this structure
        let data = DataRef::<T>(self.data);

        scheduler().future_sync(&self.queue, move || {
            let data        = data.0;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// After the pending operations for this item are performed, waits for the
    /// supplied future to complete and then calls the specified function
    ///
    pub fn after<'a, TFn, Res, Fut>(&self, after: Fut, job: TFn) -> impl 'static + Future<Output=Result<Res, oneshot::Canceled>> + Send
    where
        Res: 'static + Send,
        Fut: 'static + Future + Send,
        TFn: 'static + Send + FnOnce(&mut T, Fut::Output) -> Res
    {
        self.future_desync(move |data| {
            async move {
                let future_result = after.await;
                job(data, future_result)
            }
        })
    }
}

impl<T: Send+Unpin> Drop for Desync<T> {
    fn drop(&mut self) {
        use std::thread;

        // Take the data we're about to drop from the object
        let data = DataRef::<T>(self.data);

        // Ensure that everything on the queue has committed by queueing a last synchronous event
        // (Not synchronising the queue would make this unsafe as we would hold on to a pointer to
        // the internal data structure)
        if thread::panicking() {
            // If the thread is already panicking when we're dropped, do not panic again
            scheduler().sync_no_panic(&self.queue, move || {
                let data = data.0;
                mem::drop(unsafe { Box::from_raw(data) });
            });
        } else {
            // Thread is not panicking
            sync(&self.queue, move || {
                let data = data.0;
                mem::drop(unsafe { Box::from_raw(data) });
            });
        }
    }
}
<commit_msg>Mark Desync as !RefUnwindSafe, it has no poisoning<commit_after>//!
//! The main `Desync` struct
//!

use super::scheduler::*;

use std::sync::{Arc};
use std::marker::{PhantomData, Unpin};
use futures::channel::oneshot;
use futures::future::{Future, BoxFuture};

use std::cell::{UnsafeCell};
use std::mem;
use std::panic::{UnwindSafe};
use std::result::{Result};

///
/// A data storage structure used to govern synchronous and asynchronous access to an underlying object.
///
pub struct Desync<T: Send+Unpin> {
    /// Queue used for scheduling runtime for this object
    queue:   Arc<JobQueue>,

    /// Data for this object. Boxed so the pointer remains the same through the lifetime of the object.
    data:    *mut T,

    /// Mark this type as !RefUnwindSafe.
    _marker: PhantomData<UnsafeCell<T>>
}

unsafe impl<T: Send+Unpin> Send for Desync<T> {}

// True iff queue: Sync
unsafe impl<T: Send+Unpin> Sync for Desync<T> {}

impl<T: Send+Unpin+UnwindSafe> UnwindSafe for Desync<T> {}

///
/// Used for passing the data pointer through to the queue
///
/// 'Safe' because the queue is synchronised during drop, so we can never use the pointer
/// if the object does not exist.
///
struct DataRef<T: Send>(*mut T);
unsafe impl<T: Send> Send for DataRef<T> {}

// TODO: we can change DataRef to Shared (https://doc.rust-lang.org/std/ptr/struct.Shared.html in the future)

// TODO: T does not need to be static as we know that its lifetime is at least the lifetime of Desync<T> and hence the queue
impl<T: 'static+Send+Unpin> Desync<T> {
    ///
    /// Creates a new Desync object
    ///
    pub fn new(data: T) -> Desync<T> {
        let queue = queue();

        Desync {
            queue:   queue,
            data:    Box::into_raw(Box::new(data)),
            _marker: PhantomData
        }
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    #[inline]
    #[deprecated(since="0.3.0", note="please use `desync` instead")]
    pub fn r#async<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        self.desync(job)
    }

    ///
    /// Performs an operation asynchronously on this item. This function will return
    /// immediately and the job will happen on a separate thread at some time in the
    /// future (generally fairly soon).
    ///
    /// Jobs are always performed in the order that they are queued and are always
    /// performed synchronously with respect to this object.
    ///
    pub fn desync<TFn>(&self, job: TFn)
    where TFn: 'static+Send+FnOnce(&mut T) -> () {
        // As drop() is the last thing called, we know that this object will still exist at the point where the queue makes the asynchronous callback
        let data = DataRef::<T>(self.data);

        desync(&self.queue, move || {
            let data = data.0;
            job(unsafe { &mut *data });
        })
    }

    ///
    /// Performs an operation synchronously on this item. This will be queued with any other
    /// jobs that this item may be performing, and this function will not return until the
    /// job is complete and the result is available.
    ///
    pub fn sync<TFn, Result>(&self, job: TFn) -> Result
    where TFn: Send+FnOnce(&mut T) -> Result, Result: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            // Exclusivity is guaranteed because the queue executes only one task at a time
            let data = DataRef::<T>(self.data);

            sync(&self.queue, move || {
                let data = data.0;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Performs an operation synchronously on this item,
    ///
    pub fn try_sync<TFn, FnResult>(&self, job: TFn) -> Result<FnResult, TrySyncError>
    where TFn: Send+FnOnce(&mut T) -> FnResult, FnResult: Send {
        let result = {
            // As drop() is the last thing called, we know that this object will still exist at the point where the callback occurs
            let data = DataRef::<T>(self.data);

            try_sync(&self.queue, move || {
                let data = data.0;
                job(unsafe { &mut *data })
            })
        };

        result
    }

    ///
    /// Deprecated name for `future_desync`
    ///
    #[inline]
    #[deprecated(since="0.7.0", note="please use either `future_desync` or `future_sync` to schedule futures")]
    pub fn future<TFn, TOutput>(&self, job: TFn) -> impl Future<Output=Result<TOutput, oneshot::Canceled>>+Send
    where   TFn:        'static+Send+for<'a> FnOnce(&'a mut T) -> BoxFuture<'a, TOutput>,
            TOutput:    'static+Send {
        self.future_desync(job)
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning the
    /// result via a future.
    ///
    /// The future will be scheduled in the background, so it will make progress even if the current scheduler is
    /// blocked for any reason. Additionally, it's not necessary to await the returned future, which can be discarded
    /// if necessary. As soon as this function returns, the order of the operations in the future are fixed: that is,
    /// any other operation queued on the same `Desync` will be performed after the future completes.
    ///
    /// Call `detach()` on the returned future to fully put it in the background. If awaited the scheduler may try to
    /// 'steal' runtime on the current thread to complete the future, but it will also run the future in the background
    /// if the current thread is not available or a background thread is available faster. The future here has static
    /// lifetime so can be passed around in case something needs to monitor when it completes.
    ///
    /// This can also be used to run the future synchronously on the current thread, useful for cases where an async
    /// function needs to be called from a non-async context. Call `sync()` on the returned future to wait for its
    /// result synchronously.
    ///
    /// Note that this has a companion function, `future_sync()` which is more useful for cases where the future is
    /// awaited, as the 'inner' future is able to borrow values from the 'outer' future due to the shorter lifetime
    /// of the returned value.
    ///
    /// The normal `sync()` and `desync()` methods will schedule their operations in order around any `future_sync()`
    /// or `future_desync()` calls, so it's possible to easily mix async and traditional threaded code in one program.
    ///
    pub fn future_desync<'a, TFn, TFuture>(&self, job: TFn) -> SchedulerFuture<TFuture::Output>
    where
        TFn:                'static + Send + FnOnce(&'a mut T) -> TFuture,
        TFuture:            'a + Send + Future,
        TFuture::Output:    'static + Send,
    {
        // The future will have a lifetime shorter than the lifetime of this structure, and exclusivity is guaranteed
        // because queues only execute one task at a time
        let data = DataRef::<T>(self.data);

        scheduler().future_desync(&self.queue, move || {
            let data        = data.0;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// Performs an operation asynchronously on the contents of this item, returning a future that must be awaited
    /// before it is dropped.
    ///
    /// The future will be scheduled in the current execution context, so it will only make progress if the current
    /// scheduler is running. The task will be cancelled and will not complete execution if the future is dropped before
    /// completion, so it's usually necessary to await the result of this function for the task to behave correctly.
    ///
    /// This is most useful for 'foreground' futures: note that the future can safely borrow values from the outer context
    /// to use while it's running (in much the same way `sync()` can in synchronous contexts). The need to dispatch to
    /// the awaiting context can make this future slightly slower than `future_desync()`
    ///
    /// Desync has a feature that makes it possible to wait synchronously for futures to complete: slightly
    /// counter-intuitively this is found in the future returned by `future_desync()` function rather than this function,
    /// which is synchronous with the current async context instead of with the current thread.
    ///
    /// The normal `sync()` and `desync()` methods will schedule their operations in order around any `future_sync()`
    /// or `future_desync()` calls, so it's possible to easily mix async and traditional threaded code in one program.
    ///
    pub fn future_sync<'a, 'b, TFn, TFuture>(&'a self, job: TFn) -> impl 'a + Future<Output=Result<TFuture::Output, oneshot::Canceled>> + Send
    where
        TFn:                'a + Send + FnOnce(&'b mut T) -> TFuture,
        TFuture:            'b + Send + Future,
        TFuture::Output:    'a + Send,
        'b:                 'a,
    {
        // The future will have a lifetime shorter than the lifetime of this structure
        let data = DataRef::<T>(self.data);

        scheduler().future_sync(&self.queue, move || {
            let data        = data.0;
            let job         = job(unsafe { &mut *data });

            async {
                job.await
            }
        })
    }

    ///
    /// After the pending operations for this item are performed, waits for the
    /// supplied future to complete and then calls the specified function
    ///
    pub fn after<'a, TFn, Res, Fut>(&self, after: Fut, job: TFn) -> impl 'static + Future<Output=Result<Res, oneshot::Canceled>> + Send
    where
        Res: 'static + Send,
        Fut: 'static + Future + Send,
        TFn: 'static + Send + FnOnce(&mut T, Fut::Output) -> Res
    {
        self.future_desync(move |data| {
            async move {
                let future_result = after.await;
                job(data, future_result)
            }
        })
    }
}

impl<T: Send+Unpin> Drop for Desync<T> {
    fn drop(&mut self) {
        use std::thread;

        // Take the data we're about to drop from the object
        let data = DataRef::<T>(self.data);

        // Ensure that everything on the queue has committed by queueing a last synchronous event
        // (Not synchronising the queue would make this unsafe as we would hold on to a pointer to
        // the internal data structure)
        if thread::panicking() {
            // If the thread is already panicking when we're dropped, do not panic again
            scheduler().sync_no_panic(&self.queue, move || {
                let data = data.0;
                mem::drop(unsafe { Box::from_raw(data) });
            });
        } else {
            // Thread is not panicking
            sync(&self.queue, move || {
                let data = data.0;
                mem::drop(unsafe { Box::from_raw(data) });
            });
        }
    }
}
<|endoftext|>