<commit_before>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::sync::{atomic::AtomicBool, mpsc, Arc};
use std::thread;
use std::time::Duration;

use raft::eraftpb::MessageType;
use tikv_util::HandyRwLock;

use test_raftstore::*;

enum FailureType<'a> {
    Partition(&'a [u64], &'a [u64]),
    Reboot(&'a [u64]),
}

fn attach_prevote_notifiers<T: Simulator>(cluster: &Cluster<T>, peer: u64) -> mpsc::Receiver<()> {
    // Setup a notifier
    let (tx, rx) = mpsc::channel();
    let response_notifier = Box::new(MessageTypeNotifier::new(
        MessageType::MsgRequestPreVoteResponse,
        tx.clone(),
        Arc::from(AtomicBool::new(true)),
    ));
    let request_notifier = Box::new(MessageTypeNotifier::new(
        MessageType::MsgRequestPreVote,
        tx.clone(),
        Arc::from(AtomicBool::new(true)),
    ));

    cluster.sim.wl().add_send_filter(peer, response_notifier);
    cluster.sim.wl().add_send_filter(peer, request_notifier);

    rx
}

// Validate that prevote is used in elections after partition or reboot of some nodes.
fn test_prevote<T: Simulator>(
    cluster: &mut Cluster<T>,
    failure_type: FailureType<'_>,
    leader_after_failure_id: impl Into<Option<u64>>,
    detect_during_failure: impl Into<Option<(u64, bool)>>,
    detect_during_recovery: impl Into<Option<(u64, bool)>>,
) {
    cluster.cfg.raft_store.prevote = true;
    // To stable the test, we use a large election timeout to make
    // leader's readiness get handle within an election timeout
    configure_for_lease_read(cluster, Some(10), Some(50));

    let leader_id = 1;
    let detect_during_failure = detect_during_failure.into();
    let detect_during_recovery = detect_during_recovery.into();

    // We must start the cluster before adding send filters, otherwise it panics.
    cluster.run();

    cluster.must_transfer_leader(1, new_peer(leader_id, 1));
    cluster.must_put(b"k1", b"v1");

    // Determine how to fail.
    let rx = if let Some((id, _)) = detect_during_failure {
        let rx = attach_prevote_notifiers(cluster, id);
        debug!("Attached failure prevote notifier.");
        Some(rx)
    } else {
        None
    };

    match failure_type {
        FailureType::Partition(majority, minority) => {
            cluster.partition(majority.to_vec(), minority.to_vec());
        }
        FailureType::Reboot(peers) => {
            peers.iter().for_each(|&peer| cluster.stop_node(peer));
        }
    };

    if let (Some(rx), Some((_, should_detect))) = (rx, detect_during_failure) {
        // Once we see a response on the wire we know a prevote round is happening.
        let received = rx.recv_timeout(Duration::from_secs(5));
        debug!("Done with failure prevote notifier, got {:?}", received);
        assert_eq!(
            received.is_ok(),
            should_detect,
            "Sends a PreVote or PreVoteResponse during failure.",
        );
    }

    // Let the cluster recover.
    match failure_type {
        FailureType::Partition(_, _) => {
            cluster.clear_send_filters();
        }
        FailureType::Reboot(peers) => {
            cluster.clear_send_filters();
            peers.iter().for_each(|&peer| {
                cluster.run_node(peer).unwrap();
            });
        }
    };

    // Prepare to listen.
    let rx = if let Some((id, _)) = detect_during_recovery {
        let rx = attach_prevote_notifiers(cluster, id);
        debug!("Attached recovery prevote notifier.");
        Some(rx)
    } else {
        None
    };

    if let Some(leader_id) = leader_after_failure_id.into() {
        cluster.must_transfer_leader(1, new_peer(leader_id, 1));
    };

    // Once we see a response on the wire we know a prevote round is happening.
    if let (Some(rx), Some((_, should_detect))) = (rx, detect_during_failure) {
        let received = rx.recv_timeout(Duration::from_secs(5));
        debug!("Done with recovery prevote notifier, got {:?}", received);

        assert_eq!(
            received.is_ok(),
            should_detect,
            "Sends a PreVote or PreVoteResponse during recovery.",
        );
    };

    cluster.must_put(b"k3", b"v3");
    assert_eq!(cluster.must_get(b"k1"), Some(b"v1".to_vec()));
}

#[test]
fn test_prevote_partition_leader_in_majority_detect_in_majority() {
    let mut cluster = new_server_cluster(0, 5);
    // Since the leader is in the majority and not rebooted, it sees no prevote.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[4, 5]),
        None,
        (1, false),
        (1, false),
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_majority_detect_in_minority() {
    let mut cluster = new_server_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[4, 5]),
        None,
        (4, true),
        None,
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_minority_detect_in_majority() {
    let mut cluster = new_server_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2], &[3, 4, 5]),
        None,
        (4, true),
        None,
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_minority_detect_in_minority() {
    let mut cluster = new_server_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[3, 4, 5]),
        None,
        (4, true),
        None,
    );
}

#[test]
fn test_prevote_reboot_majority_followers() {
    let mut cluster = new_server_cluster(0, 5);
    // A prevote round will start, but nothing will succeed.
    test_prevote(
        &mut cluster,
        FailureType::Reboot(&[3, 4, 5]),
        1,
        (1, true),
        None,
    );
}

#[test]
fn test_prevote_reboot_minority_followers() {
    let mut cluster = new_server_cluster(0, 5);
    // A prevote round will start, but nothing will succeed until recovery.
    test_prevote(
        &mut cluster,
        FailureType::Reboot(&[4, 5]),
        None,
        (2, false),
        (2, false),
    );
}

// Test isolating a minority of the cluster and make sure that the remove themselves.
#[cfg(feature = "protobuf-codec")]
fn test_pair_isolated<T: Simulator>(cluster: &mut Cluster<T>) {
    let region = 1;
    let pd_client = Arc::clone(&cluster.pd_client);

    // Given some nodes A, B, C, D, E, we partition the cluster such that D, E are isolated from the rest.
    cluster.run();
    // Choose a predictable leader so we don't accidentally partition the leader.
    cluster.must_transfer_leader(region, new_peer(1, 1));
    cluster.partition(vec![1, 2, 3], vec![4, 5]);

    // Then, add a policy to PD that it should ask the Raft leader to remove the peer from the group.
    pd_client.must_remove_peer(region, new_peer(4, 4));
    pd_client.must_remove_peer(region, new_peer(5, 5));

    // Verify the nodes have self removed.
    cluster.must_remove_region(4, region);
    cluster.must_remove_region(5, region);
}

// FIXME(nrc) failing on CI only
#[cfg(feature = "protobuf-codec")]
#[test]
fn test_server_pair_isolated() {
    let mut cluster = new_server_cluster(0, 5);
    test_pair_isolated(&mut cluster);
}

fn test_isolated_follower_leader_does_not_change<T: Simulator>(cluster: &mut Cluster<T>) {
    cluster.run();
    cluster.must_transfer_leader(1, new_peer(1, 1));
    cluster.must_put(b"k1", b"v1");
    let region_status = new_status_request(1, new_peer(1, 1), new_region_leader_cmd());
    let resp = cluster
        .call_command(region_status.clone(), Duration::from_secs(5))
        .unwrap();
    let term = resp.get_header().get_current_term();
    // Isolate peer5.
    cluster.partition(vec![1, 2, 3, 4], vec![5]);
    let election_timeout = cluster.cfg.raft_store.raft_base_tick_interval.0
        * cluster.cfg.raft_store.raft_election_timeout_ticks as u32;
    // Peer5 should not increase its term.
    thread::sleep(election_timeout * 2);
    // Now peer5 can send messages to others
    cluster.clear_send_filters();
    thread::sleep(election_timeout * 2);
    cluster.must_put(b"k1", b"v1");
    // Peer1 is still the leader.
    let leader = cluster.leader_of_region(1);
    assert_eq!(leader, Some(new_peer(1, 1)));
    // And the term is not changed.
    let resp = cluster
        .call_command(region_status, Duration::from_secs(5))
        .unwrap();
    let current_term = resp.get_header().get_current_term();
    assert_eq!(term, current_term);
}

#[test]
fn test_server_isolated_follower_leader_does_not_change() {
    let mut cluster = new_server_cluster(0, 5);
    test_isolated_follower_leader_does_not_change(&mut cluster);
}

fn test_create_peer_from_pre_vote<T: Simulator>(cluster: &mut Cluster<T>) {
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    let r1 = cluster.run_conf_change();
    cluster.must_put(b"k1", b"v1");

    let rx = attach_prevote_notifiers(cluster, 1);
    cluster.add_send_filter(IsolationFilterFactory::new(2));
    pd_client.must_add_peer(r1, new_peer(2, 2));

    if rx.recv_timeout(Duration::from_secs(3)).is_err() {
        panic!("peer 1 should send pre vote");
    }

    // The peer 2 should be created.
    cluster.clear_send_filters();
    must_get_equal(&cluster.get_engine(2), b"k1", b"v1");
}

#[test]
fn test_node_create_peer_from_pre_vote() {
    let mut cluster = new_node_cluster(0, 2);
    cluster.cfg.raft_store.prevote = true;
    test_create_peer_from_pre_vote(&mut cluster);
}
<commit_msg>tests: stabilize integration cases about prevote (#6320)<commit_after>// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.

use std::sync::{atomic::AtomicBool, mpsc, Arc};
use std::thread;
use std::time::Duration;

use raft::eraftpb::MessageType;
use tikv_util::HandyRwLock;

use test_raftstore::*;

enum FailureType<'a> {
    Partition(&'a [u64], &'a [u64]),
    Reboot(&'a [u64]),
}

fn attach_prevote_notifiers<T: Simulator>(cluster: &Cluster<T>, peer: u64) -> mpsc::Receiver<()> {
    // Setup a notifier
    let (tx, rx) = mpsc::channel();
    let response_notifier = Box::new(MessageTypeNotifier::new(
        MessageType::MsgRequestPreVoteResponse,
        tx.clone(),
        Arc::from(AtomicBool::new(true)),
    ));
    let request_notifier = Box::new(MessageTypeNotifier::new(
        MessageType::MsgRequestPreVote,
        tx.clone(),
        Arc::from(AtomicBool::new(true)),
    ));

    cluster.sim.wl().add_send_filter(peer, response_notifier);
    cluster.sim.wl().add_send_filter(peer, request_notifier);

    rx
}

// Validate that prevote is used in elections after partition or reboot of some nodes.
fn test_prevote<T: Simulator>(
    cluster: &mut Cluster<T>,
    failure_type: FailureType<'_>,
    leader_after_failure_id: impl Into<Option<u64>>,
    detect_during_failure: impl Into<Option<(u64, bool)>>,
    detect_during_recovery: impl Into<Option<(u64, bool)>>,
) {
    cluster.cfg.raft_store.prevote = true;
    // To stable the test, we use a large election timeout to make
    // leader's readiness get handle within an election timeout
    configure_for_lease_read(cluster, Some(20), Some(10));

    let leader_id = 1;
    let detect_during_failure = detect_during_failure.into();
    let detect_during_recovery = detect_during_recovery.into();

    // We must start the cluster before adding send filters, otherwise it panics.
    cluster.run();

    cluster.must_transfer_leader(1, new_peer(leader_id, 1));
    cluster.must_put(b"k1", b"v1");

    // Determine how to fail.
    let rx = if let Some((id, _)) = detect_during_failure {
        let rx = attach_prevote_notifiers(cluster, id);
        debug!("Attached failure prevote notifier.");
        Some(rx)
    } else {
        None
    };

    match failure_type {
        FailureType::Partition(majority, minority) => {
            cluster.partition(majority.to_vec(), minority.to_vec());
        }
        FailureType::Reboot(peers) => {
            peers.iter().for_each(|&peer| cluster.stop_node(peer));
        }
    };

    if let (Some(rx), Some((_, should_detect))) = (rx, detect_during_failure) {
        // Once we see a response on the wire we know a prevote round is happening.
        let received = rx.recv_timeout(Duration::from_secs(5));
        debug!("Done with failure prevote notifier, got {:?}", received);
        assert_eq!(
            received.is_ok(),
            should_detect,
            "Sends a PreVote or PreVoteResponse during failure.",
        );
    }

    // Let the cluster recover.
    match failure_type {
        FailureType::Partition(_, _) => {
            cluster.clear_send_filters();
        }
        FailureType::Reboot(peers) => {
            cluster.clear_send_filters();
            peers.iter().for_each(|&peer| {
                cluster.run_node(peer).unwrap();
            });
        }
    };

    // Prepare to listen.
    let rx = if let Some((id, _)) = detect_during_recovery {
        let rx = attach_prevote_notifiers(cluster, id);
        debug!("Attached recovery prevote notifier.");
        Some(rx)
    } else {
        None
    };

    if let Some(leader_id) = leader_after_failure_id.into() {
        cluster.must_transfer_leader(1, new_peer(leader_id, 1));
    };

    // Once we see a response on the wire we know a prevote round is happening.
    if let (Some(rx), Some((_, should_detect))) = (rx, detect_during_failure) {
        let received = rx.recv_timeout(Duration::from_secs(5));
        debug!("Done with recovery prevote notifier, got {:?}", received);

        assert_eq!(
            received.is_ok(),
            should_detect,
            "Sends a PreVote or PreVoteResponse during recovery.",
        );
    };

    cluster.must_put(b"k3", b"v3");
    assert_eq!(cluster.must_get(b"k1"), Some(b"v1".to_vec()));
}

#[test]
fn test_prevote_partition_leader_in_majority_detect_in_majority() {
    let mut cluster = new_node_cluster(0, 5);
    // Since the leader is in the majority and not rebooted, it sees no prevote.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[4, 5]),
        None,
        (1, false),
        (1, false),
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_majority_detect_in_minority() {
    let mut cluster = new_node_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[4, 5]),
        None,
        (4, true),
        None,
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_minority_detect_in_majority() {
    let mut cluster = new_node_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2], &[3, 4, 5]),
        None,
        (4, true),
        None,
    );
}

// TODO: Enable detect after failure when we can reliably capture the prevote.
#[test]
fn test_prevote_partition_leader_in_minority_detect_in_minority() {
    let mut cluster = new_node_cluster(0, 5);
    // The follower is in the minority and is part of a prevote process. On rejoin it adopts the
    // old leader.
    test_prevote(
        &mut cluster,
        FailureType::Partition(&[1, 2, 3], &[3, 4, 5]),
        None,
        (4, true),
        None,
    );
}

#[test]
fn test_prevote_reboot_majority_followers() {
    let mut cluster = new_node_cluster(0, 5);
    // A prevote round will start, but nothing will succeed.
    test_prevote(
        &mut cluster,
        FailureType::Reboot(&[3, 4, 5]),
        1,
        (1, true),
        None,
    );
}

#[test]
fn test_prevote_reboot_minority_followers() {
    let mut cluster = new_node_cluster(0, 5);
    // A prevote round will start, but nothing will succeed until recovery.
    test_prevote(
        &mut cluster,
        FailureType::Reboot(&[4, 5]),
        None,
        (2, false),
        (2, false),
    );
}

// Test isolating a minority of the cluster and make sure that the remove themselves.
#[cfg(feature = "protobuf-codec")]
fn test_pair_isolated<T: Simulator>(cluster: &mut Cluster<T>) {
    let region = 1;
    let pd_client = Arc::clone(&cluster.pd_client);

    // Given some nodes A, B, C, D, E, we partition the cluster such that D, E are isolated from the rest.
    cluster.run();
    // Choose a predictable leader so we don't accidentally partition the leader.
    cluster.must_transfer_leader(region, new_peer(1, 1));
    cluster.partition(vec![1, 2, 3], vec![4, 5]);

    // Then, add a policy to PD that it should ask the Raft leader to remove the peer from the group.
    pd_client.must_remove_peer(region, new_peer(4, 4));
    pd_client.must_remove_peer(region, new_peer(5, 5));

    // Verify the nodes have self removed.
    cluster.must_remove_region(4, region);
    cluster.must_remove_region(5, region);
}

// FIXME(nrc) failing on CI only
#[cfg(feature = "protobuf-codec")]
#[test]
fn test_server_pair_isolated() {
    let mut cluster = new_server_cluster(0, 5);
    test_pair_isolated(&mut cluster);
}

fn test_isolated_follower_leader_does_not_change<T: Simulator>(cluster: &mut Cluster<T>) {
    cluster.run();
    cluster.must_transfer_leader(1, new_peer(1, 1));
    cluster.must_put(b"k1", b"v1");
    let region_status = new_status_request(1, new_peer(1, 1), new_region_leader_cmd());
    let resp = cluster
        .call_command(region_status.clone(), Duration::from_secs(5))
        .unwrap();
    let term = resp.get_header().get_current_term();
    // Isolate peer5.
    cluster.partition(vec![1, 2, 3, 4], vec![5]);
    let election_timeout = cluster.cfg.raft_store.raft_base_tick_interval.0
        * cluster.cfg.raft_store.raft_election_timeout_ticks as u32;
    // Peer5 should not increase its term.
    thread::sleep(election_timeout * 2);
    // Now peer5 can send messages to others
    cluster.clear_send_filters();
    thread::sleep(election_timeout * 2);
    cluster.must_put(b"k1", b"v1");
    // Peer1 is still the leader.
    let leader = cluster.leader_of_region(1);
    assert_eq!(leader, Some(new_peer(1, 1)));
    // And the term is not changed.
    let resp = cluster
        .call_command(region_status, Duration::from_secs(5))
        .unwrap();
    let current_term = resp.get_header().get_current_term();
    assert_eq!(term, current_term);
}

#[test]
fn test_server_isolated_follower_leader_does_not_change() {
    let mut cluster = new_server_cluster(0, 5);
    test_isolated_follower_leader_does_not_change(&mut cluster);
}

fn test_create_peer_from_pre_vote<T: Simulator>(cluster: &mut Cluster<T>) {
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    let r1 = cluster.run_conf_change();
    cluster.must_put(b"k1", b"v1");

    let rx = attach_prevote_notifiers(cluster, 1);
    cluster.add_send_filter(IsolationFilterFactory::new(2));
    pd_client.must_add_peer(r1, new_peer(2, 2));

    if rx.recv_timeout(Duration::from_secs(3)).is_err() {
        panic!("peer 1 should send pre vote");
    }

    // The peer 2 should be created.
    cluster.clear_send_filters();
    must_get_equal(&cluster.get_engine(2), b"k1", b"v1");
}

#[test]
fn test_node_create_peer_from_pre_vote() {
    let mut cluster = new_node_cluster(0, 2);
    cluster.cfg.raft_store.prevote = true;
    test_create_peer_from_pre_vote(&mut cluster);
}
<|endoftext|>


<commit_before>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_http::AwsErrorRetryPolicy;
use aws_sdk_dynamodb::error::DescribeTableError;
use aws_sdk_dynamodb::input::DescribeTableInput;
use aws_sdk_dynamodb::middleware::DefaultMiddleware;
use aws_sdk_dynamodb::model::{
    AttributeDefinition, AttributeValue, KeySchemaElement, KeyType, ProvisionedThroughput,
    ScalarAttributeType, Select, TableStatus,
};
use aws_sdk_dynamodb::operation::DescribeTable;
use aws_sdk_dynamodb::output::DescribeTableOutput;
use aws_sdk_dynamodb::{Client, Config, Error, Region, PKG_VERSION};
use aws_smithy_client::erase::DynConnector;
use aws_smithy_http::result::{SdkError, SdkSuccess};

use aws_smithy_http::operation::Operation;
use aws_smithy_http::retry::ClassifyResponse;
use aws_smithy_types::retry::RetryKind;
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};
use std::io::{stdin, Read};
use std::time::Duration;
use std::{iter, process};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// Whether to run in interactive mode (you have to press return between operations)
    #[structopt(short, long)]
    interactive: bool,

    /// The AWS Region
    #[structopt(short, long)]
    region: Option<String>,

    /// Activate verbose mode
    #[structopt(short, long)]
    verbose: bool,
}

/// Create a random, n-length string
fn random_string(n: usize) -> String {
    let mut rng = thread_rng();
    iter::repeat(())
        .map(|()| rng.sample(Alphanumeric))
        .map(char::from)
        .take(n)
        .collect()
}

/// Create a new table.
// snippet-start:[dynamodb.rust.crud-make_table]
async fn make_table(
    client: &Client,
    table: &str,
    key: &str,
) -> Result<(), SdkError<aws_sdk_dynamodb::error::CreateTableError>> {
    let ad = AttributeDefinition::builder()
        .attribute_name(key)
        .attribute_type(ScalarAttributeType::S)
        .build();

    let ks = KeySchemaElement::builder()
        .attribute_name(key)
        .key_type(KeyType::Hash)
        .build();

    let pt = ProvisionedThroughput::builder()
        .read_capacity_units(10)
        .write_capacity_units(5)
        .build();

    match client
        .create_table()
        .table_name(table)
        .key_schema(ks)
        .attribute_definitions(ad)
        .provisioned_throughput(pt)
        .send()
        .await
    {
        Ok(_) => Ok(()),
        Err(e) => Err(e),
    }
}
// snippet-end:[dynamodb.rust.crud-make_table]

/// For add_item and query_item
#[derive(Clone)]
struct Item {
    table: String,
    key: String,
    value: String,
    first_name: String,
    last_name: String,
    age: String,
    utype: String,
}

/// Add an item to the table.
// snippet-start:[dynamodb.rust.crud-add_item]
async fn add_item(
    client: &Client,
    item: Item,
) -> Result<(), SdkError<aws_sdk_dynamodb::error::PutItemError>> {
    let user_av = AttributeValue::S(item.value);
    let type_av = AttributeValue::S(item.utype);
    let age_av = AttributeValue::S(item.age);
    let first_av = AttributeValue::S(item.first_name);
    let last_av = AttributeValue::S(item.last_name);

    match client
        .put_item()
        .table_name(item.table)
        .item(item.key, user_av)
        .item("account_type", type_av)
        .item("age", age_av)
        .item("first_name", first_av)
        .item("last_name", last_av)
        .send()
        .await
    {
        Ok(_) => Ok(()),
        Err(e) => Err(e),
    }
}
// snippet-start:[dynamodb.rust.crud-add_item]

/// Query the table for an item matching the input values.
/// Returns true if the item is found; otherwise false.
// snippet-start:[dynamodb.rust.crud-query_item]
async fn query_item(client: &Client, item: Item) -> bool {
    let value = &item.value;
    let key = &item.key;
    let user_av = AttributeValue::S(value.to_string());

    match client
        .query()
        .table_name(item.table)
        .key_condition_expression("#key = :value".to_string())
        .expression_attribute_names("#key".to_string(), key.to_string())
        .expression_attribute_values(":value".to_string(), user_av)
        .select(Select::AllAttributes)
        .send()
        .await
    {
        Ok(resp) => {
            if resp.count > 0 {
                println!("Found a matching entry in the table:");
                println!("{:?}", resp.items.unwrap_or_default().pop());
                true
            } else {
                println!("Did not find a match.");
                false
            }
        }
        Err(e) => {
            println!("Got an error querying table:");
            println!("{}", e);
            process::exit(1);
        }
    }
}
// snippet-end:[dynamodb.rust.crud-query_item]

// Deletes an item from a table.
// snippet-start:[dynamodb.rust.crud-remove_item]
async fn remove_item(client: &Client, table: &str, key: &str, value: String) -> Result<(), Error> {
    let user_av = AttributeValue::S(value);

    client
        .delete_item()
        .table_name(table)
        .key(key, user_av)
        .send()
        .await?;

    println!("Deleted item.");

    Ok(())
}
// snippet-end:[dynamodb.rust.crud-remove_item]

// Deletes a table.
// snippet-start:[dynamodb.rust.crud-remove_table]
async fn remove_table(client: &Client, table: &str) -> Result<(), Error> {
    client.delete_table().table_name(table).send().await?;

    Ok(())
}
// snippet-end:[dynamodb.rust.crud-remove_table]

/// Hand-written waiter to retry every second until the table is out of `Creating` state
#[derive(Clone)]
struct WaitForReadyTable<R> {
    inner: R,
}

impl<R> ClassifyResponse<SdkSuccess<DescribeTableOutput>, SdkError<DescribeTableError>>
    for WaitForReadyTable<R>
where
    R: ClassifyResponse<SdkSuccess<DescribeTableOutput>, SdkError<DescribeTableError>>,
{
    fn classify(
        &self,
        response: Result<&SdkSuccess<DescribeTableOutput>, &SdkError<DescribeTableError>>,
    ) -> RetryKind {
        match self.inner.classify(response) {
            RetryKind::NotRetryable => (),
            other => return other,
        };
        match response {
            Ok(SdkSuccess { parsed, .. }) => {
                if parsed
                    .table
                    .as_ref()
                    .unwrap()
                    .table_status
                    .as_ref()
                    .unwrap()
                    == &TableStatus::Creating
                {
                    RetryKind::Explicit(Duration::from_secs(1))
                } else {
                    RetryKind::NotRetryable
                }
            }
            _ => RetryKind::NotRetryable,
        }
    }
}

/// Wait for the user to press Enter.
fn pause() {
    println!("Press Enter to continue.");
    stdin().read_exact(&mut [0]).unwrap();
}

/// Performs CRUD (create, read, update, delete) operations on a DynamoDB table and table item.
/// It creates a table, adds an item to the table, updates the item, deletes the item, and deletes the table.
/// The table name, primary key, and primary key value are all created as random strings.
///
/// # Arguments
///
/// * `[-i]` - Whether to pause between operations.
/// * `[-r REGION]` - The region in which the client is created.
///    If not supplied, uses the value of the **AWS_REGION** environment variable.
///    If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();
    let Opt {
        interactive,
        region,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    let shared_config = aws_config::from_env().region(region_provider).load().await;

    // Create 10-character random table name
    let table = random_string(10);

    // Create a 6-character random key name
    let key = random_string(6);

    // Create a 12-character random key value
    let value = random_string(12);

    // Specify first name, last name, age, and type
    let first_name = "DummyFirstName";
    let last_name = "DummyLastName";
    let age = "33";
    let utype = "standard_user";

    println!();

    if verbose {
        println!("DynamoDB client version: {}", PKG_VERSION);
        println!(
            "Region:                  {}",
            shared_config.region().unwrap()
        );
        println!("Table:                   {}", table);
        println!("Key:                     {}", key);
        println!("Value:                   {}", value);
        println!("First name:              {}", first_name);
        println!("Last name:               {}", last_name);
        println!("Age:                     {}", age);
        println!("User type:               {}", utype);

        println!();
    }

    let client = Client::new(&shared_config);

    /* Create table */
    println!("Creating the table.");
    match make_table(&client, &table, &key).await {
        Err(e) => {
            println!("Got an error creating the table:");
            println!("{}", e);
            process::exit(1);
        }
        Ok(_) => {
            println!("Created the table.");
        }
    }

    println!("Waiting for table to be ready.");

    let raw_client = aws_smithy_client::Client::<DynConnector, DefaultMiddleware>::dyn_https();

    raw_client
        .call(wait_for_ready_table(&table, client.conf()).await)
        .await
        .expect("table should become ready.");

    println!("Table is now ready to use.");

    if interactive {
        pause();
    }

    println!("Adding item to table.");

    let mut item = Item {
        table: table.clone(),
        key: key.clone(),
        value: value.clone(),
        first_name: first_name.to_string(),
        last_name: last_name.to_string(),
        age: age.to_string(),
        utype: utype.to_string(),
    };

    add_item(&client, item.clone()).await?;
    println!("Added item to table.");

    if interactive {
        pause();
    }

    item.age = "44".to_string();

    /* Update the item */
    println!("Modifying table item to change age to 44.");

    add_item(&client, item.clone()).await?;

    println!("Modified table item.");

    if interactive {
        pause();
    }

    /* Get item and compare it with the one we added */
    println!("Comparing table item to original value.");

    query_item(&client, item).await;

    if interactive {
        pause();
    }

    /* Delete item */
    println!("Deleting item.");
    remove_item(&client, &table, &key, value.clone()).await?;

    if interactive {
        pause();
    }

    /* Delete table */
    println!("Deleting table.");
    remove_table(&client, &table).await?;
    println!();

    Ok(())
}

/// Construct a `DescribeTable` request with a policy to retry every second until the table
/// is ready
async fn wait_for_ready_table(
    table_name: &str,
    conf: &Config,
) -> Operation<DescribeTable, WaitForReadyTable<AwsErrorRetryPolicy>> {
    let operation = DescribeTableInput::builder()
        .table_name(table_name)
        .build()
        .expect("valid input")
        .make_operation(conf)
        .await
        .expect("valid operation");
    let waiting_policy = WaitForReadyTable {
        inner: operation.retry_policy().clone(),
    };
    operation.with_retry_policy(waiting_policy)
}
<commit_msg>Fixed snippet tag in DynamoDB code example in Rust<commit_after>/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */

use aws_config::meta::region::RegionProviderChain;
use aws_http::AwsErrorRetryPolicy;
use aws_sdk_dynamodb::error::DescribeTableError;
use aws_sdk_dynamodb::input::DescribeTableInput;
use aws_sdk_dynamodb::middleware::DefaultMiddleware;
use aws_sdk_dynamodb::model::{
    AttributeDefinition, AttributeValue, KeySchemaElement, KeyType, ProvisionedThroughput,
    ScalarAttributeType, Select, TableStatus,
};
use aws_sdk_dynamodb::operation::DescribeTable;
use aws_sdk_dynamodb::output::DescribeTableOutput;
use aws_sdk_dynamodb::{Client, Config, Error, Region, PKG_VERSION};
use aws_smithy_client::erase::DynConnector;
use aws_smithy_http::result::{SdkError, SdkSuccess};

use aws_smithy_http::operation::Operation;
use aws_smithy_http::retry::ClassifyResponse;
use aws_smithy_types::retry::RetryKind;
use rand::distributions::Alphanumeric;
use rand::{thread_rng, Rng};
use std::io::{stdin, Read};
use std::time::Duration;
use std::{iter, process};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    /// Whether to run in interactive mode (you have to press return between operations)
    #[structopt(short, long)]
    interactive: bool,

    /// The AWS Region
    #[structopt(short, long)]
    region: Option<String>,

    /// Activate verbose mode
    #[structopt(short, long)]
    verbose: bool,
}

/// Create a random, n-length string
fn random_string(n: usize) -> String {
    let mut rng = thread_rng();
    iter::repeat(())
        .map(|()| rng.sample(Alphanumeric))
        .map(char::from)
        .take(n)
        .collect()
}

/// Create a new table.
// snippet-start:[dynamodb.rust.crud-make_table]
async fn make_table(
    client: &Client,
    table: &str,
    key: &str,
) -> Result<(), SdkError<aws_sdk_dynamodb::error::CreateTableError>> {
    let ad = AttributeDefinition::builder()
        .attribute_name(key)
        .attribute_type(ScalarAttributeType::S)
        .build();

    let ks = KeySchemaElement::builder()
        .attribute_name(key)
        .key_type(KeyType::Hash)
        .build();

    let pt = ProvisionedThroughput::builder()
        .read_capacity_units(10)
        .write_capacity_units(5)
        .build();

    match client
        .create_table()
        .table_name(table)
        .key_schema(ks)
        .attribute_definitions(ad)
        .provisioned_throughput(pt)
        .send()
        .await
    {
        Ok(_) => Ok(()),
        Err(e) => Err(e),
    }
}
// snippet-end:[dynamodb.rust.crud-make_table]

/// For add_item and query_item
#[derive(Clone)]
struct Item {
    table: String,
    key: String,
    value: String,
    first_name: String,
    last_name: String,
    age: String,
    utype: String,
}

/// Add an item to the table.
// snippet-start:[dynamodb.rust.crud-add_item]
async fn add_item(
    client: &Client,
    item: Item,
) -> Result<(), SdkError<aws_sdk_dynamodb::error::PutItemError>> {
    let user_av = AttributeValue::S(item.value);
    let type_av = AttributeValue::S(item.utype);
    let age_av = AttributeValue::S(item.age);
    let first_av = AttributeValue::S(item.first_name);
    let last_av = AttributeValue::S(item.last_name);

    match client
        .put_item()
        .table_name(item.table)
        .item(item.key, user_av)
        .item("account_type", type_av)
        .item("age", age_av)
        .item("first_name", first_av)
        .item("last_name", last_av)
        .send()
        .await
    {
        Ok(_) => Ok(()),
        Err(e) => Err(e),
    }
}
// snippet-end:[dynamodb.rust.crud-add_item]

/// Query the table for an item matching the input values.
/// Returns true if the item is found; otherwise false.
// snippet-start:[dynamodb.rust.crud-query_item]
async fn query_item(client: &Client, item: Item) -> bool {
    let value = &item.value;
    let key = &item.key;
    let user_av = AttributeValue::S(value.to_string());

    match client
        .query()
        .table_name(item.table)
        .key_condition_expression("#key = :value".to_string())
        .expression_attribute_names("#key".to_string(), key.to_string())
        .expression_attribute_values(":value".to_string(), user_av)
        .select(Select::AllAttributes)
        .send()
        .await
    {
        Ok(resp) => {
            if resp.count > 0 {
                println!("Found a matching entry in the table:");
                println!("{:?}", resp.items.unwrap_or_default().pop());
                true
            } else {
                println!("Did not find a match.");
                false
            }
        }
        Err(e) => {
            println!("Got an error querying table:");
            println!("{}", e);
            process::exit(1);
        }
    }
}
// snippet-end:[dynamodb.rust.crud-query_item]

// Deletes an item from a table.
// snippet-start:[dynamodb.rust.crud-remove_item]
async fn remove_item(client: &Client, table: &str, key: &str, value: String) -> Result<(), Error> {
    let user_av = AttributeValue::S(value);

    client
        .delete_item()
        .table_name(table)
        .key(key, user_av)
        .send()
        .await?;

    println!("Deleted item.");

    Ok(())
}
// snippet-end:[dynamodb.rust.crud-remove_item]

// Deletes a table.
// snippet-start:[dynamodb.rust.crud-remove_table]
async fn remove_table(client: &Client, table: &str) -> Result<(), Error> {
    client.delete_table().table_name(table).send().await?;

    Ok(())
}
// snippet-end:[dynamodb.rust.crud-remove_table]

/// Hand-written waiter to retry every second until the table is out of `Creating` state
#[derive(Clone)]
struct WaitForReadyTable<R> {
    inner: R,
}

impl<R> ClassifyResponse<SdkSuccess<DescribeTableOutput>, SdkError<DescribeTableError>>
    for WaitForReadyTable<R>
where
    R: ClassifyResponse<SdkSuccess<DescribeTableOutput>, SdkError<DescribeTableError>>,
{
    fn classify(
        &self,
        response: Result<&SdkSuccess<DescribeTableOutput>, &SdkError<DescribeTableError>>,
    ) -> RetryKind {
        match self.inner.classify(response) {
            RetryKind::NotRetryable => (),
            other => return other,
        };
        match response {
            Ok(SdkSuccess { parsed, .. }) => {
                if parsed
                    .table
                    .as_ref()
                    .unwrap()
                    .table_status
                    .as_ref()
                    .unwrap()
                    == &TableStatus::Creating
                {
                    RetryKind::Explicit(Duration::from_secs(1))
                } else {
                    RetryKind::NotRetryable
                }
            }
            _ => RetryKind::NotRetryable,
        }
    }
}

/// Wait for the user to press Enter.
fn pause() {
    println!("Press Enter to continue.");
    stdin().read_exact(&mut [0]).unwrap();
}

/// Performs CRUD (create, read, update, delete) operations on a DynamoDB table and table item.
/// It creates a table, adds an item to the table, updates the item, deletes the item, and deletes the table.
/// The table name, primary key, and primary key value are all created as random strings.
///
/// # Arguments
///
/// * `[-i]` - Whether to pause between operations.
/// * `[-r REGION]` - The region in which the client is created.
///    If not supplied, uses the value of the **AWS_REGION** environment variable.
///    If the environment variable is not set, defaults to **us-west-2**.
/// * `[-v]` - Whether to display additional information.
#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt::init();
    let Opt {
        interactive,
        region,
        verbose,
    } = Opt::from_args();

    let region_provider = RegionProviderChain::first_try(region.map(Region::new))
        .or_default_provider()
        .or_else(Region::new("us-west-2"));
    let shared_config = aws_config::from_env().region(region_provider).load().await;

    // Create 10-character random table name
    let table = random_string(10);

    // Create a 6-character random key name
    let key = random_string(6);

    // Create a 12-character random key value
    let value = random_string(12);

    // Specify first name, last name, age, and type
    let first_name = "DummyFirstName";
    let last_name = "DummyLastName";
    let age = "33";
    let utype = "standard_user";

    println!();

    if verbose {
        println!("DynamoDB client version: {}", PKG_VERSION);
        println!(
            "Region:                  {}",
            shared_config.region().unwrap()
        );
        println!("Table:                   {}", table);
        println!("Key:                     {}", key);
        println!("Value:                   {}", value);
        println!("First name:              {}", first_name);
        println!("Last name:               {}", last_name);
        println!("Age:                     {}", age);
        println!("User type:               {}", utype);

        println!();
    }

    let client = Client::new(&shared_config);

    /* Create table */
    println!("Creating the table.");
    match make_table(&client, &table, &key).await {
        Err(e) => {
            println!("Got an error creating the table:");
            println!("{}", e);
            process::exit(1);
        }
        Ok(_) => {
            println!("Created the table.");
        }
    }

    println!("Waiting for table to be ready.");

    let raw_client = aws_smithy_client::Client::<DynConnector, DefaultMiddleware>::dyn_https();

    raw_client
        .call(wait_for_ready_table(&table, client.conf()).await)
        .await
        .expect("table should become ready.");

    println!("Table is now ready to use.");

    if interactive {
        pause();
    }

    println!("Adding item to table.");

    let mut item = Item {
        table: table.clone(),
        key: key.clone(),
        value: value.clone(),
        first_name: first_name.to_string(),
        last_name: last_name.to_string(),
        age: age.to_string(),
        utype: utype.to_string(),
    };

    add_item(&client, item.clone()).await?;
    println!("Added item to table.");

    if interactive {
        pause();
    }

    item.age = "44".to_string();

    /* Update the item */
    println!("Modifying table item to change age to 44.");

    add_item(&client, item.clone()).await?;

    println!("Modified table item.");

    if interactive {
        pause();
    }

    /* Get item and compare it with the one we added */
    println!("Comparing table item to original value.");

    query_item(&client, item).await;

    if interactive {
        pause();
    }

    /* Delete item */
    println!("Deleting item.");
    remove_item(&client, &table, &key, value.clone()).await?;

    if interactive {
        pause();
    }

    /* Delete table */
    println!("Deleting table.");
    remove_table(&client, &table).await?;
    println!();

    Ok(())
}

/// Construct a `DescribeTable` request with a policy to retry every second until the table
/// is ready
async fn wait_for_ready_table(
    table_name: &str,
    conf: &Config,
) -> Operation<DescribeTable, WaitForReadyTable<AwsErrorRetryPolicy>> {
    let operation = DescribeTableInput::builder()
        .table_name(table_name)
        .build()
        .expect("valid input")
        .make_operation(conf)
        .await
        .expect("valid operation");
    let waiting_policy = WaitForReadyTable {
        inner: operation.retry_policy().clone(),
    };
    operation.with_retry_policy(waiting_policy)
}
<|endoftext|>


<commit_before>use bytes::{Buf, BufMut, IntoBuf};
use futures::{Async, Future, Poll};
use hyper::client::{HttpConnector, Service};
use hyper::Uri;
use hyper_tls::{HttpsConnector, MaybeHttpsStream};
use native_tls::TlsConnector;
use tokio_core::reactor::Handle;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_tls::{TlsConnectorExt, TlsStream};

use std::io::{self, Cursor, Read, Write};
use std::sync::Arc;

use {proxy, Proxy};

// pub(crate)

pub struct Connector {
    https: HttpsConnector<HttpConnector>,
    proxies: Arc<Vec<Proxy>>,
    tls: TlsConnector,
}

impl Connector {
    pub fn new(tls: TlsConnector, proxies: Arc<Vec<Proxy>>, handle: &Handle) -> Connector {
        let mut http = HttpConnector::new(4, handle);
        http.enforce_http(false);
        let https = HttpsConnector::from((http, tls.clone()));

        Connector {
            https: https,
            proxies: proxies,
            tls: tls,
        }
    }

    pub fn danger_disable_hostname_verification(&mut self) {
        self.https.danger_disable_hostname_verification(true);
    }
}

impl Service for Connector {
    type Request = Uri;
    type Response = Conn;
    type Error = io::Error;
    type Future = Connecting;

    fn call(&self, uri: Uri) -> Self::Future {
        for prox in self.proxies.iter() {
            if let Some(puri) = proxy::proxies(prox, &uri) {
                if uri.scheme() == Some("https") {
                    let host = uri.authority().unwrap().to_owned();
                    let tls = self.tls.clone();
                    return Box::new(self.https.call(puri).and_then(move |conn| {
                        tunnel(conn, host.clone())
                            .and_then(move |tunneled| {
                                tls.connect_async(&host, tunneled)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .map(|io| Conn::Proxied(io))
                    }));
                }
                return Box::new(self.https.call(puri).map(|io| Conn::Normal(io)));
            }
        }
        Box::new(self.https.call(uri).map(|io| Conn::Normal(io)))
    }
}

type HttpStream = <HttpConnector as Service>::Response;
type HttpsStream = MaybeHttpsStream<HttpStream>;

pub type Connecting = Box<Future<Item=Conn, Error=io::Error>>;

pub enum Conn {
    Normal(HttpsStream),
    Proxied(TlsStream<MaybeHttpsStream<HttpStream>>),
}

impl Read for Conn {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Conn::Normal(ref mut s) => s.read(buf),
            Conn::Proxied(ref mut s) => s.read(buf),
        }
    }
}

impl Write for Conn {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Conn::Normal(ref mut s) => s.write(buf),
            Conn::Proxied(ref mut s) => s.write(buf),
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Conn::Normal(ref mut s) => s.flush(),
            Conn::Proxied(ref mut s) => s.flush(),
        }
    }
}

impl AsyncRead for Conn {
    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
        match *self {
            Conn::Normal(ref s) => s.prepare_uninitialized_buffer(buf),
            Conn::Proxied(ref s) => s.prepare_uninitialized_buffer(buf),
        }
    }

    fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.read_buf(buf),
            Conn::Proxied(ref mut s) => s.read_buf(buf),
        }
    }
}

impl AsyncWrite for Conn {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.shutdown(),
            Conn::Proxied(ref mut s) => s.shutdown(),
        }
    }

    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.write_buf(buf),
            Conn::Proxied(ref mut s) => s.write_buf(buf),
        }
    }
}

fn tunnel<T>(conn: T, host: String) -> Tunnel<T> {
     let buf = format!("\
        CONNECT {0} HTTP/1.1\r\n\
        Host: {0}\r\n\
        \r\n\
    ", host).into_bytes();

     Tunnel {
        buf: buf.into_buf(),
        conn: Some(conn),
        state: TunnelState::Writing,
     }
}

struct Tunnel<T> {
    buf: Cursor<Vec<u8>>,
    conn: Option<T>,
    state: TunnelState,
}

enum TunnelState {
    Writing,
    Reading
}

impl<T> Future for Tunnel<T>
where T: AsyncRead + AsyncWrite {
    type Item = T;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let TunnelState::Writing = self.state {
                let n = try_ready!(self.conn.as_mut().unwrap().write_buf(&mut self.buf));
                if !self.buf.has_remaining_mut() {
                    self.state = TunnelState::Reading;
                    self.buf.get_mut().truncate(0);
                } else if n == 0 {
                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected eof while tunneling"));
                }
            } else {
                let n = try_ready!(self.conn.as_mut().unwrap().read_buf(&mut self.buf.get_mut()));
                let read = &self.buf.get_ref()[..];
                if n == 0 {
                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected eof while tunneling"));
                } else if read.len() > 12 {
                    if read.starts_with(b"HTTP/1.1 200") {
                        if read.ends_with(b"\r\n\r\n") {
                            return Ok(Async::Ready(self.conn.take().unwrap()));
                        }
                        // else read more
                    } else {
                        return Err(io::Error::new(io::ErrorKind::Other, "unsuccessful tunnel"));
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Read, Write};
    use std::net::TcpListener;
    use std::thread;
    use futures::Future;
    use tokio_core::reactor::Core;
    use tokio_core::net::TcpStream;
    use super::tunnel;

    macro_rules! mock_tunnel {
        () => ({
            mock_tunnel!(b"\
                HTTP/1.1 200 OK\r\n\
                \r\n\
            ")
        });
        ($write:expr) => ({
            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
            let addr = listener.local_addr().unwrap();
            let connect_expected = format!("\
                CONNECT {0} HTTP/1.1\r\n\
                Host: {0}\r\n\
                \r\n\
            ", addr).into_bytes();

            thread::spawn(move || {
                let (mut sock, _) = listener.accept().unwrap();
                let mut buf = [0u8; 4096];
                let n = sock.read(&mut buf).unwrap();
                assert_eq!(&buf[..n], &connect_expected[..]);

                sock.write_all($write).unwrap();
            });
            addr
        })
    }

    #[test]
    fn test_tunnel() {
        let addr = mock_tunnel!();

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host)
        });

        core.run(work).unwrap();
    }

    #[test]
    fn test_tunnel_eof() {
        let addr = mock_tunnel!(b"HTTP/1.1 200 OK");

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host)
        });

        core.run(work).unwrap_err();
    }

    #[test]
    fn test_tunnel_bad_response() {
        let addr = mock_tunnel!(b"foo bar baz hallo");

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host)
        });

        core.run(work).unwrap_err();
    }
}
<commit_msg>format connect request with host:port and connect on  http/1.0 responses<commit_after>use bytes::{Buf, BufMut, IntoBuf};
use futures::{Async, Future, Poll};
use hyper::client::{HttpConnector, Service};
use hyper::Uri;
use hyper_tls::{HttpsConnector, MaybeHttpsStream};
use native_tls::TlsConnector;
use tokio_core::reactor::Handle;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_tls::{TlsConnectorExt, TlsStream};

use std::io::{self, Cursor, Read, Write};
use std::sync::Arc;

use {proxy, Proxy};

// pub(crate)

pub struct Connector {
    https: HttpsConnector<HttpConnector>,
    proxies: Arc<Vec<Proxy>>,
    tls: TlsConnector,
}

impl Connector {
    pub fn new(tls: TlsConnector, proxies: Arc<Vec<Proxy>>, handle: &Handle) -> Connector {
        let mut http = HttpConnector::new(4, handle);
        http.enforce_http(false);
        let https = HttpsConnector::from((http, tls.clone()));

        Connector {
            https: https,
            proxies: proxies,
            tls: tls,
        }
    }

    pub fn danger_disable_hostname_verification(&mut self) {
        self.https.danger_disable_hostname_verification(true);
    }
}

impl Service for Connector {
    type Request = Uri;
    type Response = Conn;
    type Error = io::Error;
    type Future = Connecting;

    fn call(&self, uri: Uri) -> Self::Future {
        for prox in self.proxies.iter() {
            if let Some(puri) = proxy::proxies(prox, &uri) {
                if uri.scheme() == Some("https") {
                    let host = uri.host().unwrap().to_owned();
                    let port = match uri.port(){
                        Some(port) => port.to_string(),
                        None => 443.to_string()
                    }.to_owned();
                    let tls = self.tls.clone();
                    return Box::new(self.https.call(puri).and_then(move |conn| {
                        tunnel(conn, host.clone(), port.clone())
                            .and_then(move |tunneled| {
                                tls.connect_async(&host, tunneled)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .map(|io| Conn::Proxied(io))
                    }));
                }
                return Box::new(self.https.call(puri).map(|io| Conn::Normal(io)));
            }
        }
        Box::new(self.https.call(uri).map(|io| Conn::Normal(io)))
    }
}

type HttpStream = <HttpConnector as Service>::Response;
type HttpsStream = MaybeHttpsStream<HttpStream>;

pub type Connecting = Box<Future<Item=Conn, Error=io::Error>>;

pub enum Conn {
    Normal(HttpsStream),
    Proxied(TlsStream<MaybeHttpsStream<HttpStream>>),
}

impl Read for Conn {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match *self {
            Conn::Normal(ref mut s) => s.read(buf),
            Conn::Proxied(ref mut s) => s.read(buf),
        }
    }
}

impl Write for Conn {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match *self {
            Conn::Normal(ref mut s) => s.write(buf),
            Conn::Proxied(ref mut s) => s.write(buf),
        }
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        match *self {
            Conn::Normal(ref mut s) => s.flush(),
            Conn::Proxied(ref mut s) => s.flush(),
        }
    }
}

impl AsyncRead for Conn {
    unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
        match *self {
            Conn::Normal(ref s) => s.prepare_uninitialized_buffer(buf),
            Conn::Proxied(ref s) => s.prepare_uninitialized_buffer(buf),
        }
    }

    fn read_buf<B: BufMut>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.read_buf(buf),
            Conn::Proxied(ref mut s) => s.read_buf(buf),
        }
    }
}

impl AsyncWrite for Conn {
    fn shutdown(&mut self) -> Poll<(), io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.shutdown(),
            Conn::Proxied(ref mut s) => s.shutdown(),
        }
    }

    fn write_buf<B: Buf>(&mut self, buf: &mut B) -> Poll<usize, io::Error> {
        match *self {
            Conn::Normal(ref mut s) => s.write_buf(buf),
            Conn::Proxied(ref mut s) => s.write_buf(buf),
        }
    }
}

fn tunnel<T>(conn: T, host: String, port: String) -> Tunnel<T> {
     let buf = format!("\
        CONNECT {0}:{1} HTTP/1.1\r\n\
        Host: {0}:{1}\r\n\
        \r\n\
    ", host, port).into_bytes();

     Tunnel {
        buf: buf.into_buf(),
        conn: Some(conn),
        state: TunnelState::Writing,
     }
}

struct Tunnel<T> {
    buf: Cursor<Vec<u8>>,
    conn: Option<T>,
    state: TunnelState,
}

enum TunnelState {
    Writing,
    Reading
}

impl<T> Future for Tunnel<T>
where T: AsyncRead + AsyncWrite {
    type Item = T;
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            if let TunnelState::Writing = self.state {
                let n = try_ready!(self.conn.as_mut().unwrap().write_buf(&mut self.buf));
                if !self.buf.has_remaining_mut() {
                    self.state = TunnelState::Reading;
                    self.buf.get_mut().truncate(0);
                } else if n == 0 {
                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected eof while tunneling"));
                }
            } else {
                let n = try_ready!(self.conn.as_mut().unwrap().read_buf(&mut self.buf.get_mut()));
                let read = &self.buf.get_ref()[..];
                if n == 0 {
                    return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected eof while tunneling"));
                } else if read.len() > 12 {
                    if read.starts_with(b"HTTP/1.1 200") || read.starts_with(b"HTTP/1.0 200"){
                        if read.ends_with(b"\r\n\r\n") {
                            return Ok(Async::Ready(self.conn.take().unwrap()));
                        }
                        // else read more
                    } else {
                        return Err(io::Error::new(io::ErrorKind::Other, "unsuccessful tunnel"));
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::io::{Read, Write};
    use std::net::TcpListener;
    use std::thread;
    use futures::Future;
    use tokio_core::reactor::Core;
    use tokio_core::net::TcpStream;
    use super::tunnel;

    macro_rules! mock_tunnel {
        () => ({
            mock_tunnel!(b"\
                HTTP/1.1 200 OK\r\n\
                \r\n\
            ")
        });
        ($write:expr) => ({
            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
            let addr = listener.local_addr().unwrap();
            let connect_expected = format!("\
                CONNECT {0}:{1} HTTP/1.1\r\n\
                Host: {0}:{1}\r\n\
                \r\n\
            ", addr.ip(), addr.port()).into_bytes();

            thread::spawn(move || {
                let (mut sock, _) = listener.accept().unwrap();
                let mut buf = [0u8; 4096];
                let n = sock.read(&mut buf).unwrap();
                assert_eq!(&buf[..n], &connect_expected[..]);

                sock.write_all($write).unwrap();
            });
            addr
        })
    }

    #[test]
    fn test_tunnel() {
        let addr = mock_tunnel!();

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.ip().to_string();
        let port = addr.port().to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host, port)
        });

        core.run(work).unwrap();
    }

    #[test]
    fn test_tunnel_eof() {
        let addr = mock_tunnel!(b"HTTP/1.1 200 OK");

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.ip().to_string();
        let port = addr.port().to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host, port)
        });

        core.run(work).unwrap_err();
    }

    #[test]
    fn test_tunnel_bad_response() {
        let addr = mock_tunnel!(b"foo bar baz hallo");

        let mut core = Core::new().unwrap();
        let work = TcpStream::connect(&addr, &core.handle());
        let host = addr.ip().to_string();
        let port = addr.port().to_string();
        let work = work.and_then(|tcp| {
            tunnel(tcp, host, port)
        });

        core.run(work).unwrap_err();
    }
}
<|endoftext|>


<commit_before>//! Kata OS global memory management support

extern crate alloc;
use core::ops::Range;
use kata_memory_interface::ObjDesc;
use kata_memory_interface::ObjDescBundle;
use kata_memory_interface::MemoryError;
use kata_memory_interface::MemoryManagerInterface;
use kata_memory_interface::MemoryManagerStats;
use kata_os_common::sel4_sys;
use log::{debug, error, trace};
use sel4_sys::seL4_CPtr;
use sel4_sys::seL4_CNode_Delete;
use sel4_sys::seL4_Error;
use sel4_sys::seL4_Result;
use sel4_sys::seL4_Word;
use sel4_sys::seL4_UntypedDesc;
use sel4_sys::seL4_Untyped_Retype;
use smallvec::SmallVec;

// SmallVec capacity for untyped memory slabs. There are two instances;
// one for anonymous memory and one for device-backed memory. The memory
// manager is expected to be setup as a static global so these data
// structures will land in .bss and only overflow to the heap if
// initialized with more than this count.
const UNTYPED_SLAB_CAPACITY: usize = 64;  // # slabs kept inline

// The MemoryManager supports allocating & freeing seL4 objects that are
// instantiated from UntypedMemory "slabs". Allocation causes untyped memory
// to be converted to concrete types. Freeing deletes the specified capabilities
// and updates the bookkeeping. Note that a free only releases the specified
// cap; if there are dups or derived objects the memory will not be returned
// to the untyped slab from which it was allocated and the bookkeeping done
// here will be out of sync with the kernel.
// TODO(sleffler): support device-backed memory objects
#[derive(Debug)]
struct UntypedSlab {
    pub size_bits: usize, // NB: only used to sort
    pub _base_paddr: seL4_Word,  // Physical address of slab start
    pub _last_paddr: seL4_Word,  // Physical address of slab end
    pub cptr: seL4_CPtr,    // seL4 untyped object
    pub _next_paddr: seL4_Word,  // Physical address of next available chunk
}
impl UntypedSlab {
    fn new(ut: &seL4_UntypedDesc, cptr: seL4_CPtr) -> Self {
        UntypedSlab {
            size_bits: ut.size_bits(),
            _base_paddr: ut.paddr,
            _last_paddr: ut.paddr + (1 << ut.size_bits()),
            cptr: cptr,
            _next_paddr: ut.paddr,
        }
    }
    fn _size(&self) -> usize { l2tob(self.size_bits) }
    fn size_bits(&self) -> usize { self.size_bits }
}
pub struct MemoryManager {
    untypeds: SmallVec<[UntypedSlab; UNTYPED_SLAB_CAPACITY]>,
    _device_untypeds: SmallVec<[UntypedSlab; UNTYPED_SLAB_CAPACITY]>,
    cur_untyped: usize,
    _cur_device_untyped: usize,

    total_bytes: usize,   // Total available space
    allocated_bytes: usize,   // Amount of space currently allocated
    requested_bytes: usize,  // Amount of space allocated over all time

    allocated_objs: usize,   // # seL4 objects currently allocated
    requested_objs: usize,  // # seL4 objects allocated over all time

    // Retype requests failed due to insufficient available memory.
    untyped_slab_too_small: usize,

    // Alloc requests failed due to lack of untyped memory (NB: may be
    // due to fragmentation of untyped slabs).
    out_of_memory: usize,
}

fn _howmany(value: usize, unit: usize) -> usize {
    value + (unit - 1) / unit
}
fn _round_up(value: usize, align: usize) -> usize {
    _howmany(value, align) * align
}

// Log2 bits to bytes.
fn l2tob(size_bits: usize) -> usize { 1 << size_bits }

impl MemoryManager {
    // Creates a new MemoryManager instance. The allocator is seeded
    // from the untyped memory descriptors.
    pub fn new(slots: Range<seL4_CPtr>, untypeds: &[seL4_UntypedDesc]) -> Self {
        assert!(untypeds.len() > 0);
        assert_eq!(slots.end - slots.start, untypeds.len());
        let mut m = MemoryManager {
            untypeds: SmallVec::new(),
            _device_untypeds: SmallVec::new(),
            cur_untyped: 0,
            _cur_device_untyped: 0,

            total_bytes: 0,
            allocated_bytes: 0,
            requested_bytes: 0,

            allocated_objs: 0,
            requested_objs: 0,

            untyped_slab_too_small: 0,
            out_of_memory: 0,
        };
        for (ut_index, ut) in untypeds.iter().enumerate() {
            if ut.is_device() {
                m._device_untypeds.push(UntypedSlab::new(ut, slots.start + ut_index));
            } else {
                m.untypeds.push(UntypedSlab::new(ut, slots.start + ut_index));
                m.total_bytes += l2tob(ut.size_bits());
            }
        }
        // Sort non-device slabs by descending size.
        m.untypeds.sort_unstable_by(|a, b| b.size_bits().cmp(&a.size_bits()));
        m
    }

    // Total available space.
    pub fn total_available_space(&self) -> usize {
        self.total_bytes
    }
    // Current allocated space
    pub fn allocated_space(&self) -> usize {
        self.allocated_bytes
    }
    // Current free space.
    pub fn free_space(&self) -> usize {
        self.total_bytes - self.allocated_bytes
    }
    // Total space allocated over time
    pub fn total_requested_space(&self) -> usize {
        self.requested_bytes
    }

    // Current allocated objects
    pub fn allocated_objs(&self) -> usize {
        self.allocated_objs
    }
    // Total objects allocated over time
    pub fn total_requested_objs(&self) -> usize {
        self.requested_objs
    }

    pub fn untyped_slab_too_small(&self) -> usize {
        self.untyped_slab_too_small
    }
    pub fn out_of_memory(&self) -> usize {
        self.out_of_memory
    }

    fn retype_untyped(
        free_untyped: seL4_CPtr,
        root: seL4_CPtr,
        obj: &ObjDesc,
    ) -> seL4_Result {
        unsafe {
            seL4_Untyped_Retype(
                free_untyped,
                /*type=*/ obj.type_.into(),
                /*size_bits=*/ obj.retype_size_bits().unwrap(),
                /*root=*/ root,
                /*node_index=*/ 0,  // Ignored 'cuz depth is zero
                /*node_depth=*/ 0,  // NB: store in cnode
                /*node_offset=*/ obj.cptr,
                /*num_objects=*/ obj.retype_count(),
            )
        }
    }

    fn delete_caps(root: seL4_CPtr, depth: u8, od: &ObjDesc) -> seL4_Result {
        for offset in 0..od.retype_count() {
            // TODO(sleffler) warn about errors?
            unsafe { seL4_CNode_Delete(root, od.cptr + offset, depth) }?;
        }
        Ok(())
    }
}

impl MemoryManagerInterface for MemoryManager {
    fn alloc(&mut self, bundle: &ObjDescBundle) -> Result<(), MemoryError> {
        trace!("alloc {:?}", bundle);

        // TODO(sleffler): split by device vs no-device (or allow mixing)
        let first_ut = self.cur_untyped;
        let mut ut_index = first_ut;

        let mut allocated_bytes: usize = 0;
        let mut allocated_objs: usize = 0;

        for od in &bundle.objs {
            // NB: we don't check slots are available (the kernel will tell us).
            // TODO(sleffler): maybe check size_bytes() against untyped slab?
            //    (we depend on the kernel for now)
            while let Err(e) =
                // NB: we don't allocate ASIDPool objects but if we did it
                //   would fail because it needs to map to an UntypedObject
                MemoryManager::retype_untyped(self.untypeds[ut_index].cptr, bundle.cnode, od)
            {
                if e != seL4_Error::seL4_NotEnoughMemory {
                    // Should not happen.
                    // TODO(sleffler): reclaim allocations
                    error!("Allocation request failed (retype returned {:?})", e);
                    return Err(MemoryError::UnknownMemoryError);
                }
                // This untyped does not have enough available space, try
                // the next slab until we exhaust all slabs. This is the best
                // we can do without per-slab bookkeeping.
                self.untyped_slab_too_small += 1;
                ut_index = (ut_index + 1) % self.untypeds.len();
                debug!("Advance to untyped slab {}", ut_index);
                if ut_index == first_ut {
                    // TODO(sleffler): reclaim allocations
                    self.out_of_memory += 1;
                    debug!("Allocation request failed (out of space)");
                    return Err(MemoryError::AllocFailed);
                }
            }
            allocated_objs += od.retype_count();
            allocated_bytes += od.size_bytes().unwrap();

        }
        self.cur_untyped = ut_index;

        self.allocated_bytes += allocated_bytes;
        self.allocated_objs += allocated_objs;

        // NB: does not include requests that fail
        self.requested_objs += allocated_objs;
        self.requested_bytes += allocated_bytes;

        Ok(())
    }
    fn free(&mut self, bundle: &ObjDescBundle) -> Result<(), MemoryError> {
        trace!("free {:?}", bundle);

        for od in &bundle.objs {
            // TODO(sleffler): support leaving objects so client can do bulk
            //   reclaim on exit (maybe require cptr != 0)
            if MemoryManager::delete_caps(bundle.cnode, bundle.depth, od).is_ok() {
                // NB: atm we do not do per-untyped bookkeeping so just track
                //   global stats.
                // TODO(sleffler): temp workaround for bad bookkeeping / client mis-handling
                let size_bytes = od.size_bytes().ok_or(MemoryError::ObjTypeInvalid)?;
                if size_bytes <= self.allocated_bytes {
                    self.allocated_bytes -= size_bytes;
                    self.allocated_objs -= od.retype_count();
                } else {
                    debug!("Undeflow on free of {:?}", od);
                }
            }
        }
        Ok(())
    }
    fn stats(&self) -> Result<MemoryManagerStats, MemoryError> {
        Ok(MemoryManagerStats {
            allocated_bytes: self.allocated_space(),
            free_bytes: self.free_space(),
            total_requested_bytes: self.total_requested_space(),
            // TODO(sleffler): track fragmentation
            // NB: assumes all heap usage is for the buddy allocator
            overhead_bytes: 0,

            allocated_objs: self.allocated_objs(),
            total_requested_objs: self.total_requested_objs(),

            untyped_slab_too_small: self.untyped_slab_too_small(),
            out_of_memory: self.out_of_memory(),
        })
    }
}
<commit_msg>MemoryManager: workaround release build problem<commit_after>//! Kata OS global memory management support

extern crate alloc;
use core::ops::Range;
use kata_memory_interface::ObjDesc;
use kata_memory_interface::ObjDescBundle;
use kata_memory_interface::MemoryError;
use kata_memory_interface::MemoryManagerInterface;
use kata_memory_interface::MemoryManagerStats;
use kata_os_common::sel4_sys;
use log::{debug, error, warn, trace};
use sel4_sys::seL4_CPtr;
use sel4_sys::seL4_CNode_Delete;
use sel4_sys::seL4_Error;
use sel4_sys::seL4_Result;
use sel4_sys::seL4_Word;
use sel4_sys::seL4_UntypedDesc;
use sel4_sys::seL4_Untyped_Retype;
use smallvec::SmallVec;

// SmallVec capacity for untyped memory slabs. There are two instances;
// one for anonymous memory and one for device-backed memory. The memory
// manager is expected to be setup as a static global so these data
// structures will land in .bss and only overflow to the heap if
// initialized with more than this count.
const UNTYPED_SLAB_CAPACITY: usize = 64;  // # slabs kept inline

// The MemoryManager supports allocating & freeing seL4 objects that are
// instantiated from UntypedMemory "slabs". Allocation causes untyped memory
// to be converted to concrete types. Freeing deletes the specified capabilities
// and updates the bookkeeping. Note that a free only releases the specified
// cap; if there are dups or derived objects the memory will not be returned
// to the untyped slab from which it was allocated and the bookkeeping done
// here will be out of sync with the kernel.
// TODO(sleffler): support device-backed memory objects
#[derive(Debug)]
struct UntypedSlab {
    pub size_bits: usize, // NB: only used to sort
    pub _base_paddr: seL4_Word,  // Physical address of slab start
    pub _last_paddr: seL4_Word,  // Physical address of slab end
    pub cptr: seL4_CPtr,    // seL4 untyped object
    pub _next_paddr: seL4_Word,  // Physical address of next available chunk
}
impl UntypedSlab {
    fn new(ut: &seL4_UntypedDesc, cptr: seL4_CPtr) -> Self {
        UntypedSlab {
            size_bits: ut.size_bits(),
            _base_paddr: ut.paddr,
            _last_paddr: ut.paddr + (1 << ut.size_bits()),
            cptr: cptr,
            _next_paddr: ut.paddr,
        }
    }
    fn _size(&self) -> usize { l2tob(self.size_bits) }
    fn size_bits(&self) -> usize { self.size_bits }
}
pub struct MemoryManager {
    untypeds: SmallVec<[UntypedSlab; UNTYPED_SLAB_CAPACITY]>,
    _device_untypeds: SmallVec<[UntypedSlab; UNTYPED_SLAB_CAPACITY]>,
    cur_untyped: usize,
    _cur_device_untyped: usize,

    total_bytes: usize,   // Total available space
    allocated_bytes: usize,   // Amount of space currently allocated
    requested_bytes: usize,  // Amount of space allocated over all time

    allocated_objs: usize,   // # seL4 objects currently allocated
    requested_objs: usize,  // # seL4 objects allocated over all time

    // Retype requests failed due to insufficient available memory.
    untyped_slab_too_small: usize,

    // Alloc requests failed due to lack of untyped memory (NB: may be
    // due to fragmentation of untyped slabs).
    out_of_memory: usize,
}

fn _howmany(value: usize, unit: usize) -> usize {
    value + (unit - 1) / unit
}
fn _round_up(value: usize, align: usize) -> usize {
    _howmany(value, align) * align
}

// Log2 bits to bytes.
fn l2tob(size_bits: usize) -> usize { 1 << size_bits }

impl MemoryManager {
    // Creates a new MemoryManager instance. The allocator is seeded
    // from the untyped memory descriptors.
    pub fn new(slots: Range<seL4_CPtr>, untypeds: &[seL4_UntypedDesc]) -> Self {
        assert!(untypeds.len() > 0);
        assert_eq!(slots.end - slots.start, untypeds.len());
        let mut m = MemoryManager {
            untypeds: SmallVec::new(),
            _device_untypeds: SmallVec::new(),
            cur_untyped: 0,
            _cur_device_untyped: 0,

            total_bytes: 0,
            allocated_bytes: 0,
            requested_bytes: 0,

            allocated_objs: 0,
            requested_objs: 0,

            untyped_slab_too_small: 0,
            out_of_memory: 0,
        };
        for (ut_index, ut) in untypeds.iter().enumerate() {
            if ut.is_device() {
                m._device_untypeds.push(UntypedSlab::new(ut, slots.start + ut_index));
            } else {
                m.untypeds.push(UntypedSlab::new(ut, slots.start + ut_index));
                m.total_bytes += l2tob(ut.size_bits());
            }
        }
        // Sort non-device slabs by descending size.
        m.untypeds.sort_unstable_by(|a, b| b.size_bits().cmp(&a.size_bits()));
        m
    }

    // Total available space.
    pub fn total_available_space(&self) -> usize {
        self.total_bytes
    }
    // Current allocated space
    pub fn allocated_space(&self) -> usize {
        self.allocated_bytes
    }
    // Current free space.
    pub fn free_space(&self) -> usize {
        self.total_bytes - self.allocated_bytes
    }
    // Total space allocated over time
    pub fn total_requested_space(&self) -> usize {
        self.requested_bytes
    }

    // Current allocated objects
    pub fn allocated_objs(&self) -> usize {
        self.allocated_objs
    }
    // Total objects allocated over time
    pub fn total_requested_objs(&self) -> usize {
        self.requested_objs
    }

    pub fn untyped_slab_too_small(&self) -> usize {
        self.untyped_slab_too_small
    }
    pub fn out_of_memory(&self) -> usize {
        self.out_of_memory
    }

    fn retype_untyped(
        free_untyped: seL4_CPtr,
        root: seL4_CPtr,
        obj: &ObjDesc,
    ) -> seL4_Result {
        unsafe {
            seL4_Untyped_Retype(
                free_untyped,
                /*type=*/ obj.type_.into(),
                /*size_bits=*/ obj.retype_size_bits().unwrap(),
                /*root=*/ root,
                /*node_index=*/ 0,  // Ignored 'cuz depth is zero
                /*node_depth=*/ 0,  // NB: store in cnode
                /*node_offset=*/ obj.cptr,
                /*num_objects=*/ obj.retype_count(),
            )
        }
    }

    fn delete_caps(root: seL4_CPtr, depth: u8, od: &ObjDesc) -> seL4_Result {
        for offset in 0..od.retype_count() {
            let path = (root, od.cptr + offset, depth);
            if let Err(e) = unsafe { seL4_CNode_Delete(path.0, path.1, path.2) } {
                warn!("DELETE {:?} failed: od {:?} error {:?}", &path, od, e);
            }
        }
        Ok(())
    }
}

impl MemoryManagerInterface for MemoryManager {
    fn alloc(&mut self, bundle: &ObjDescBundle) -> Result<(), MemoryError> {
        trace!("alloc {:?}", bundle);

        // TODO(sleffler): split by device vs no-device (or allow mixing)
        let first_ut = self.cur_untyped;
        let mut ut_index = first_ut;

        let mut allocated_bytes: usize = 0;
        let mut allocated_objs: usize = 0;

        for od in &bundle.objs {
            // NB: we don't check slots are available (the kernel will tell us).
            // TODO(sleffler): maybe check size_bytes() against untyped slab?
            //    (we depend on the kernel for now)
            while let Err(e) =
                // NB: we don't allocate ASIDPool objects but if we did it
                //   would fail because it needs to map to an UntypedObject
                MemoryManager::retype_untyped(self.untypeds[ut_index].cptr, bundle.cnode, od)
            {
                if e != seL4_Error::seL4_NotEnoughMemory {
                    // Should not happen.
                    // TODO(sleffler): reclaim allocations
                    error!("Allocation request failed (retype returned {:?})", e);
                    return Err(MemoryError::UnknownMemoryError);
                }
                // This untyped does not have enough available space, try
                // the next slab until we exhaust all slabs. This is the best
                // we can do without per-slab bookkeeping.
                self.untyped_slab_too_small += 1;
                ut_index = (ut_index + 1) % self.untypeds.len();
                debug!("Advance to untyped slab {}", ut_index);
                if ut_index == first_ut {
                    // TODO(sleffler): reclaim allocations
                    self.out_of_memory += 1;
                    debug!("Allocation request failed (out of space)");
                    return Err(MemoryError::AllocFailed);
                }
            }
            allocated_objs += od.retype_count();
            allocated_bytes += od.size_bytes().unwrap();

        }
        self.cur_untyped = ut_index;

        self.allocated_bytes += allocated_bytes;
        self.allocated_objs += allocated_objs;

        // NB: does not include requests that fail
        self.requested_objs += allocated_objs;
        self.requested_bytes += allocated_bytes;

        Ok(())
    }
    fn free(&mut self, bundle: &ObjDescBundle) -> Result<(), MemoryError> {
        trace!("free {:?}", bundle);

        for od in &bundle.objs {
            // TODO(sleffler): support leaving objects so client can do bulk
            //   reclaim on exit (maybe require cptr != 0)
            if MemoryManager::delete_caps(bundle.cnode, bundle.depth, od).is_ok() {
                // NB: atm we do not do per-untyped bookkeeping so just track
                //   global stats.
                // TODO(sleffler): temp workaround for bad bookkeeping / client mis-handling
                let size_bytes = od.size_bytes().ok_or(MemoryError::ObjTypeInvalid)?;
                if size_bytes <= self.allocated_bytes {
                    self.allocated_bytes -= size_bytes;
                    self.allocated_objs -= od.retype_count();
                } else {
                    debug!("Undeflow on free of {:?}", od);
                }
            }
        }
        Ok(())
    }
    fn stats(&self) -> Result<MemoryManagerStats, MemoryError> {
        Ok(MemoryManagerStats {
            allocated_bytes: self.allocated_space(),
            free_bytes: self.free_space(),
            total_requested_bytes: self.total_requested_space(),
            // TODO(sleffler): track fragmentation
            // NB: assumes all heap usage is for the buddy allocator
            overhead_bytes: 0,

            allocated_objs: self.allocated_objs(),
            total_requested_objs: self.total_requested_objs(),

            untyped_slab_too_small: self.untyped_slab_too_small(),
            out_of_memory: self.out_of_memory(),
        })
    }
}
<|endoftext|>