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

//! Local state for keys in the snapshot in progress (the "index").

use std::borrow::Cow;

use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;

use std::sync::{Arc, Mutex, MutexGuard};
use time::Duration;

use blob;
use hash;
use periodic_timer::PeriodicTimer;
use util;

use super::schema;

error_type! {
    #[derive(Debug)]
    pub enum IndexError {
        SqlConnection(diesel::ConnectionError) {
            cause;
        },
        SqlMigration(diesel::migrations::MigrationError) {
            cause;
        },
        SqlRunMigration(diesel::migrations::RunMigrationsError) {
            cause;
        },
        SqlExecute(diesel::result::Error) {
            cause;
        },
        Message(Cow<'static, str>) {
            desc (e) &**e;
            from (s: &'static str) s.into();
            from (s: String) s.into();
        }
     }
}

#[derive(Clone, Debug)]
pub struct Entry {
    pub id: Option<u64>,
    pub parent_id: Option<u64>,

    pub name: Vec<u8>,

    pub created: Option<i64>,
    pub modified: Option<i64>,
    pub accessed: Option<i64>,

    pub permissions: Option<u64>,
    pub user_id: Option<u64>,
    pub group_id: Option<u64>,

    pub data_hash: Option<Vec<u8>>,
    pub data_length: Option<u64>,
}

#[derive(Clone)]
pub struct IndexProcess(Arc<Mutex<(Index, Option<i64>)>>);

pub enum Msg {
    /// Insert an entry in the key index.
    /// Returns `Id` with the new entry ID.
    Insert(Entry),

    /// Lookup an entry in the key index by its parent id and name.
    /// Returns either `Entry` with the found entry or `NotFound`.
    Lookup(Option<u64>, Vec<u8>),

    /// Update the `payload` and `persistent_ref` of an entry.
    /// Returns `UpdateOk`.
    UpdateDataHash(Entry, Option<hash::Hash>, Option<blob::ChunkRef>),

    /// List a directory (aka. `level`) in the index.
    /// Returns `ListResult` with all the entries under the given parent.
    ListDir(Option<u64>),

    /// Flush this key index.
    Flush,
}

pub enum Reply {
    Entry(Entry),
    NotFound,
    UpdateOk,
    ListResult(Vec<(Entry, Option<blob::ChunkRef>)>),
    FlushOk,
}

pub struct Index {
    path: String,
    conn: SqliteConnection,
    flush_timer: PeriodicTimer,
}

impl Index {
    pub fn new(path: String) -> Result<Index, IndexError> {
        let conn = try!(SqliteConnection::establish(&path));

        let ki = Index {
            path: path,
            conn: conn,
            flush_timer: PeriodicTimer::new(Duration::seconds(5)),
        };

        let dir = try!(diesel::migrations::find_migrations_directory());
        try!(diesel::migrations::run_pending_migrations_in_directory(&ki.conn,
                                                                     &dir,
                                                                     &mut util::InfoWriter));
        try!(ki.conn.begin_transaction());

        Ok(ki)
    }

    #[cfg(test)]
    pub fn new_for_testing() -> Result<Index, IndexError> {
        Index::new(":memory:".to_string())
    }

    fn last_insert_rowid(&self) -> Result<i64, IndexError> {
        let id = try!(diesel::select(diesel::expression::sql("last_insert_rowid()"))
                          .first::<i64>(&self.conn));
        Ok(id)
    }

    pub fn maybe_flush(&mut self) -> Result<(), IndexError> {
        if self.flush_timer.did_fire() {
            try!(self.flush());
        }

        Ok(())
    }

    pub fn flush(&mut self) -> Result<(), IndexError> {
        try!(self.conn.commit_transaction());
        try!(self.conn.begin_transaction());

        Ok(())
    }
}

impl Clone for Index {
    fn clone(&self) -> Index {
        // We succeeded setup once, and so should succeed again unless the underlying database has
        // disappeared, in which case it is acceptable to terminate the program.
        Index::new(self.path.clone()).expect("Could not clone index")
    }
}

impl Index {
    fn handle<F: Fn(Result<Reply, IndexError>)>
        (&mut self,
         msg: Msg,
         reply: F)-> Result<(), IndexError> {
        let reply_ok = |x| {
            reply(Ok(x));
            Ok(())
        };
        let reply_err = |e| {
            reply(Err(e));
            Ok(())
        };
        match msg {
            Msg::Insert(entry) => {
                use super::schema::keys::dsl::*;

                let entry = match entry.id {
                    Some(id_) => {
                        // Replace existing entry.
                        try!(diesel::update(keys.find(id_ as i64))
                                 .set((parent.eq(entry.parent_id.map(|x| x as i64)),
                                       name.eq(&entry.name[..]),
                                       created.eq(entry.created),
                                       modified.eq(entry.modified),
                                       accessed.eq(entry.accessed)))
                                 .execute(&self.conn));
                        entry
                    }
                    None => {
                        // Insert new entry.
                        {
                            let new = schema::NewKey {
                                parent: entry.parent_id.map(|x| x as i64),
                                name: &entry.name[..],
                                created: entry.created,
                                modified: entry.modified,
                                accessed: entry.accessed,
                                permissions: entry.permissions.map(|x| x as i64),
                                group_id: entry.group_id.map(|x| x as i64),
                                user_id: entry.user_id.map(|x| x as i64),
                                hash: None,
                                persistent_ref: None,
                            };

                            try!(diesel::insert(&new)
                                     .into(keys)
                                     .execute(&self.conn));
                        }
                        let mut entry = entry;
                        entry.id = Some(try!(self.last_insert_rowid()) as u64);
                        entry
                    }
                };

                reply_ok(Reply::Entry(entry))
            }
            Msg::Lookup(parent_, name_) => {
                use super::schema::keys::dsl::*;

                let row_opt = match parent_ {
                    Some(p) => {
                        try!(keys.filter(parent.eq(p as i64))
                                 .filter(name.eq(&name_[..]))
                                 .first::<schema::Key>(&self.conn)
                                 .optional())
                    }
                    None => {
                        try!(keys.filter(parent.is_null())
                                 .filter(name.eq(&name_[..]))
                                 .first::<schema::Key>(&self.conn)
                                 .optional())
                    }
                };

                if let Some(row) = row_opt {
                    return reply_ok(Reply::Entry(Entry {
                        id: Some(row.id as u64),
                        parent_id: parent_,
                        name: name_,
                        created: row.created,
                        modified: row.modified,
                        accessed: row.accessed,
                        permissions: row.permissions.map(|x| x as u64),
                        user_id: row.user_id.map(|x| x as u64),
                        group_id: row.group_id.map(|x| x as u64),
                        data_hash: row.hash,
                        data_length: None,
                    }));
                }

                reply_ok(Reply::NotFound)
            }
            Msg::UpdateDataHash(entry, hash_opt, persistent_ref_opt) => {
                use super::schema::keys::dsl::*;

                let id_ = match entry.id {
                    Some(i) => i as i64,
                    None => {
                        return reply_err(From::from("Tried to update data hash without an id"));
                    }
                };
                assert!(hash_opt.is_some() == persistent_ref_opt.is_some());

                let hash_bytes = hash_opt.map(|h| h.bytes);
                let persistent_ref_bytes = persistent_ref_opt.map(|p| p.as_bytes());

                if entry.modified.is_some() {
                    try!(diesel::update(keys.find(id_)
                                            .filter(modified.eq::<Option<i64>>(None)
                                                            .or(modified.le(entry.modified))))
                             .set((hash.eq(hash_bytes), persistent_ref.eq(persistent_ref_bytes)))
                             .execute(&self.conn));
                } else {
                    try!(diesel::update(keys.find(id_))
                             .set((hash.eq(hash_bytes), persistent_ref.eq(persistent_ref_bytes)))
                             .execute(&self.conn));
                }

                try!(self.maybe_flush());

                reply_ok(Reply::UpdateOk)
            }

            Msg::Flush => {
                try!(self.flush());

                reply_ok(Reply::FlushOk)
            }

            Msg::ListDir(parent_opt) => {
                use super::schema::keys::dsl::*;

                let rows = match parent_opt {
                    Some(p) => {
                        try!(keys.filter(parent.eq(p as i64))
                                 .load::<schema::Key>(&self.conn))
                    }
                    None => {
                        try!(keys.filter(parent.is_null())
                                 .load::<schema::Key>(&self.conn))
                    }
                };

                let res = rows.into_iter()
                              .map(|mut r| {
                                  (Entry {
                                      id: Some(r.id as u64),
                                      parent_id: r.parent.map(|x| x as u64),
                                      name: r.name,
                                      created: r.created,
                                      modified: r.modified,
                                      accessed: r.accessed,
                                      permissions: r.permissions
                                                    .map(|x| x as u64),
                                      user_id: r.user_id.map(|x| x as u64),
                                      group_id: r.group_id.map(|x| x as u64),
                                      data_hash: r.hash,
                                      data_length: None,
                                  },
                                   r.persistent_ref
                                    .as_mut()
                                    .map(|p| blob::ChunkRef::from_bytes(&mut &p[..]).unwrap()))
                              });

                reply_ok(Reply::ListResult(res.collect()))
            }
        }
    }
}

impl IndexProcess {
    pub fn new(index: Index) -> IndexProcess {
        IndexProcess::new_with_shutdown(index, None)
    }

    pub fn new_with_shutdown(index: Index, shutdown: Option<i64>) -> IndexProcess {
        IndexProcess(Arc::new(Mutex::new((index, shutdown))))
    }

    fn lock(&self) -> MutexGuard<(Index, Option<i64>)> {
        let mut guard = self.0.lock().expect("index-process has failed");

        match &mut guard.1 {
            &mut None => (),
            &mut Some(0) => panic!("No more requests for this index process"),
            &mut Some(ref mut n) => {
                *n -= 1;
            }
        }

        guard
    }

    pub fn send_reply(&self, msg: Msg) -> Result<Reply, IndexError>  {
        let mut guard = self.lock();
        let res = Mutex::new(None);

        let can_continue = guard.0.handle(msg, |r| {
            let mut guard = res.lock().unwrap();
            assert!((*guard).is_none());
            *guard = Some(r)
        }).is_ok();

        if !can_continue {
            guard.1 = Some(0);
        }

        let res: Option<Result<Reply, IndexError>> = res.into_inner().expect("Channel should not be poisened");
        res.expect("Callback must be called exactly once")
    }
}
<commit_msg>Closer to non-ugly code.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Local state for keys in the snapshot in progress (the "index").

use std::borrow::Cow;

use diesel;
use diesel::prelude::*;
use diesel::sqlite::SqliteConnection;

use std::sync::{Arc, Mutex, MutexGuard};
use time::Duration;

use blob;
use hash;
use periodic_timer::PeriodicTimer;
use util;

use super::schema;

error_type! {
    #[derive(Debug)]
    pub enum IndexError {
        SqlConnection(diesel::ConnectionError) {
            cause;
        },
        SqlMigration(diesel::migrations::MigrationError) {
            cause;
        },
        SqlRunMigration(diesel::migrations::RunMigrationsError) {
            cause;
        },
        SqlExecute(diesel::result::Error) {
            cause;
        },
        Message(Cow<'static, str>) {
            desc (e) &**e;
            from (s: &'static str) s.into();
            from (s: String) s.into();
        }
     }
}

#[derive(Clone, Debug)]
pub struct Entry {
    pub id: Option<u64>,
    pub parent_id: Option<u64>,

    pub name: Vec<u8>,

    pub created: Option<i64>,
    pub modified: Option<i64>,
    pub accessed: Option<i64>,

    pub permissions: Option<u64>,
    pub user_id: Option<u64>,
    pub group_id: Option<u64>,

    pub data_hash: Option<Vec<u8>>,
    pub data_length: Option<u64>,
}

#[derive(Clone)]
pub struct IndexProcess(Arc<Mutex<(Index, Option<i64>)>>);

pub enum Msg {
    /// Insert an entry in the key index.
    /// Returns `Id` with the new entry ID.
    Insert(Entry),

    /// Lookup an entry in the key index by its parent id and name.
    /// Returns either `Entry` with the found entry or `NotFound`.
    Lookup(Option<u64>, Vec<u8>),

    /// Update the `payload` and `persistent_ref` of an entry.
    /// Returns `UpdateOk`.
    UpdateDataHash(Entry, Option<hash::Hash>, Option<blob::ChunkRef>),

    /// List a directory (aka. `level`) in the index.
    /// Returns `ListResult` with all the entries under the given parent.
    ListDir(Option<u64>),

    /// Flush this key index.
    Flush,
}

pub enum Reply {
    Entry(Entry),
    NotFound,
    UpdateOk,
    ListResult(Vec<(Entry, Option<blob::ChunkRef>)>),
    FlushOk,
}

pub struct Index {
    path: String,
    conn: SqliteConnection,
    flush_timer: PeriodicTimer,
}

impl Index {
    pub fn new(path: String) -> Result<Index, IndexError> {
        let conn = try!(SqliteConnection::establish(&path));

        let ki = Index {
            path: path,
            conn: conn,
            flush_timer: PeriodicTimer::new(Duration::seconds(5)),
        };

        let dir = try!(diesel::migrations::find_migrations_directory());
        try!(diesel::migrations::run_pending_migrations_in_directory(&ki.conn,
                                                                     &dir,
                                                                     &mut util::InfoWriter));
        try!(ki.conn.begin_transaction());

        Ok(ki)
    }

    #[cfg(test)]
    pub fn new_for_testing() -> Result<Index, IndexError> {
        Index::new(":memory:".to_string())
    }

    fn last_insert_rowid(&self) -> Result<i64, IndexError> {
        let id = try!(diesel::select(diesel::expression::sql("last_insert_rowid()"))
                          .first::<i64>(&self.conn));
        Ok(id)
    }

    pub fn maybe_flush(&mut self) -> Result<(), IndexError> {
        if self.flush_timer.did_fire() {
            try!(self.flush());
        }

        Ok(())
    }

    pub fn flush(&mut self) -> Result<(), IndexError> {
        try!(self.conn.commit_transaction());
        try!(self.conn.begin_transaction());

        Ok(())
    }
}

impl Clone for Index {
    fn clone(&self) -> Index {
        // We succeeded setup once, and so should succeed again unless the underlying database has
        // disappeared, in which case it is acceptable to terminate the program.
        Index::new(self.path.clone()).expect("Could not clone index")
    }
}

impl Index {
    fn handle(&mut self,
              msg: Msg)-> Result<Reply, IndexError> {
        match msg {
            Msg::Insert(entry) => {
                use super::schema::keys::dsl::*;

                let entry = match entry.id {
                    Some(id_) => {
                        // Replace existing entry.
                        try!(diesel::update(keys.find(id_ as i64))
                                 .set((parent.eq(entry.parent_id.map(|x| x as i64)),
                                       name.eq(&entry.name[..]),
                                       created.eq(entry.created),
                                       modified.eq(entry.modified),
                                       accessed.eq(entry.accessed)))
                                 .execute(&self.conn));
                        entry
                    }
                    None => {
                        // Insert new entry.
                        {
                            let new = schema::NewKey {
                                parent: entry.parent_id.map(|x| x as i64),
                                name: &entry.name[..],
                                created: entry.created,
                                modified: entry.modified,
                                accessed: entry.accessed,
                                permissions: entry.permissions.map(|x| x as i64),
                                group_id: entry.group_id.map(|x| x as i64),
                                user_id: entry.user_id.map(|x| x as i64),
                                hash: None,
                                persistent_ref: None,
                            };

                            try!(diesel::insert(&new)
                                     .into(keys)
                                     .execute(&self.conn));
                        }
                        let mut entry = entry;
                        entry.id = Some(try!(self.last_insert_rowid()) as u64);
                        entry
                    }
                };

                Ok(Reply::Entry(entry))
            }
            Msg::Lookup(parent_, name_) => {
                use super::schema::keys::dsl::*;

                let row_opt = match parent_ {
                    Some(p) => {
                        try!(keys.filter(parent.eq(p as i64))
                                 .filter(name.eq(&name_[..]))
                                 .first::<schema::Key>(&self.conn)
                                 .optional())
                    }
                    None => {
                        try!(keys.filter(parent.is_null())
                                 .filter(name.eq(&name_[..]))
                                 .first::<schema::Key>(&self.conn)
                                 .optional())
                    }
                };

                if let Some(row) = row_opt {
                    Ok(Reply::Entry(Entry {
                        id: Some(row.id as u64),
                        parent_id: parent_,
                        name: name_,
                        created: row.created,
                        modified: row.modified,
                        accessed: row.accessed,
                        permissions: row.permissions.map(|x| x as u64),
                        user_id: row.user_id.map(|x| x as u64),
                        group_id: row.group_id.map(|x| x as u64),
                        data_hash: row.hash,
                        data_length: None,
                    }))
                } else {
                    Ok(Reply::NotFound)
                }
            }
            Msg::UpdateDataHash(entry, hash_opt, persistent_ref_opt) => {
                use super::schema::keys::dsl::*;

                let id_ = match entry.id {
                    Some(i) => i as i64,
                    None => {
                        return Err(From::from("Tried to update data hash without an id"))
                    }
                };
                assert!(hash_opt.is_some() == persistent_ref_opt.is_some());

                let hash_bytes = hash_opt.map(|h| h.bytes);
                let persistent_ref_bytes = persistent_ref_opt.map(|p| p.as_bytes());

                if entry.modified.is_some() {
                    try!(diesel::update(keys.find(id_)
                                            .filter(modified.eq::<Option<i64>>(None)
                                                            .or(modified.le(entry.modified))))
                             .set((hash.eq(hash_bytes), persistent_ref.eq(persistent_ref_bytes)))
                             .execute(&self.conn));
                } else {
                    try!(diesel::update(keys.find(id_))
                             .set((hash.eq(hash_bytes), persistent_ref.eq(persistent_ref_bytes)))
                             .execute(&self.conn));
                }

                try!(self.maybe_flush());

                Ok(Reply::UpdateOk)
            }

            Msg::Flush => {
                try!(self.flush());

                Ok(Reply::FlushOk)
            }

            Msg::ListDir(parent_opt) => {
                use super::schema::keys::dsl::*;

                let rows = match parent_opt {
                    Some(p) => {
                        try!(keys.filter(parent.eq(p as i64))
                                 .load::<schema::Key>(&self.conn))
                    }
                    None => {
                        try!(keys.filter(parent.is_null())
                                 .load::<schema::Key>(&self.conn))
                    }
                };

                let res = rows.into_iter()
                              .map(|mut r| {
                                  (Entry {
                                      id: Some(r.id as u64),
                                      parent_id: r.parent.map(|x| x as u64),
                                      name: r.name,
                                      created: r.created,
                                      modified: r.modified,
                                      accessed: r.accessed,
                                      permissions: r.permissions
                                                    .map(|x| x as u64),
                                      user_id: r.user_id.map(|x| x as u64),
                                      group_id: r.group_id.map(|x| x as u64),
                                      data_hash: r.hash,
                                      data_length: None,
                                  },
                                   r.persistent_ref
                                    .as_mut()
                                    .map(|p| blob::ChunkRef::from_bytes(&mut &p[..]).unwrap()))
                              });

                Ok(Reply::ListResult(res.collect()))
            }
        }
    }
}

impl IndexProcess {
    pub fn new(index: Index) -> IndexProcess {
        IndexProcess::new_with_shutdown(index, None)
    }

    pub fn new_with_shutdown(index: Index, shutdown: Option<i64>) -> IndexProcess {
        IndexProcess(Arc::new(Mutex::new((index, shutdown))))
    }

    fn lock(&self) -> MutexGuard<(Index, Option<i64>)> {
        let mut guard = self.0.lock().expect("index-process has failed");

        match &mut guard.1 {
            &mut None => (),
            &mut Some(0) => panic!("No more requests for this index process"),
            &mut Some(ref mut n) => {
                *n -= 1;
            }
        }

        guard
    }

    pub fn send_reply(&self, msg: Msg) -> Result<Reply, IndexError>  {
        self.lock().0.handle(msg)
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::RwLock;
use std::time::Instant;
use std::time::Duration;
use std::thread;
use std::collections::HashSet;

use futures::Future;
use futures::future::ok;
use futures::future::{loop_fn, Loop};

use grpc;
use grpc::futures_grpc::GrpcFutureSend;
use url::Url;
use rand::{self, Rng};

use kvproto::pdpb::{ResponseHeader, ErrorType, GetMembersRequest, GetMembersResponse, Member};
use kvproto::pdpb_grpc::{PDAsyncClient, PDAsync};

use util::HandyRwLock;

use super::super::{Result, Error, PdFuture};
use super::super::metrics::PD_SEND_MSG_HISTOGRAM;

pub struct Inner {
    client: PDAsyncClient,
    members: GetMembersResponse,
}

/// A leader client doing requests asynchronous.
pub struct LeaderClient {
    inner: Arc<RwLock<Inner>>,
}

impl LeaderClient {
    pub fn new(client: PDAsyncClient, members: GetMembersResponse) -> LeaderClient {
        LeaderClient {
            inner: Arc::new(RwLock::new(Inner {
                client: client,
                members: members,
            })),
        }
    }

    pub fn request<Req, Resp, F>(&self, req: Req, f: F, retry: usize) -> Request<Req, Resp, F>
        where Req: Clone + 'static,
              F: FnMut(&PDAsyncClient, Req) -> PdFuture<Resp> + Send + 'static
    {
        Request {
            reconnect_count: retry,
            request_sent: 0,
            client: LeaderClient { inner: self.inner.clone() },
            req: req,
            resp: None,
            func: f,
        }
    }

    pub fn get_leader(&self) -> Member {
        self.inner.rl().members.get_leader().clone()
    }

    // Re-establish connection with PD leader in synchronized fashion.
    pub fn reconnect(&self) {
        let start = Instant::now();
        let ret = try_connect_leader(&self.inner.rl().members);
        match ret {
            Ok((client, members)) => {
                let mut inner = self.inner.wl();
                inner.client = client;
                inner.members = members;
                warn!("updating PD client done, spent {:?}", start.elapsed());
            }

            Err(err) => {
                warn!("updating PD client spent {:?}, err {:?}",
                      start.elapsed(),
                      err);
                // FIXME: use tokio-timer instead.
                thread::sleep(Duration::from_secs(RETRY_INTERVAL));
            }
        }
    }
}

const RETRY_INTERVAL: u64 = 1;

/// The context of sending requets.
pub struct Request<Req, Resp, F> {
    reconnect_count: usize,
    request_sent: usize,

    client: LeaderClient,

    req: Req,
    resp: Option<Result<Resp>>,
    func: F,
}

const MAX_REQUEST_COUNT: usize = 5;

impl<Req, Resp, F> Request<Req, Resp, F>
    where Req: Clone + Send + 'static,
          Resp: Send + 'static,
          F: FnMut(&PDAsyncClient, Req) -> PdFuture<Resp> + Send + 'static
{
    fn reconnect_if_needed(mut self) -> PdFuture<Self> {
        debug!("reconnect remains: {}", self.reconnect_count);

        if self.request_sent == 0 {
            return ok(self).boxed();
        }

        if self.request_sent < MAX_REQUEST_COUNT {
            debug!("retry on the same client");
            return ok(self).boxed();
        }

        // Updating client.

        self.request_sent = 0;
        self.reconnect_count -= 1;

        // FIXME: should not block the core.
        warn!("updating PD client, block the tokio core");
        self.client.reconnect();

        ok(self).boxed()
    }

    fn send(mut self) -> PdFuture<Self> {
        self.request_sent += 1;
        debug!("request sent: {}", self.request_sent);
        let r = self.req.clone();
        let req = (self.func)(&self.client.inner.rl().client, r);
        req.then(|resp| {
                match resp {
                    Ok(resp) => self.resp = Some(Ok(resp)),
                    Err(err) => {
                        error!("request failed: {:?}", err);
                    }
                };
                Ok(self)
            })
            .boxed()
    }

    fn receive(self) -> Result<Loop<Self, Self>> {
        let done = self.reconnect_count == 0 || self.resp.is_some();
        if done {
            Ok(Loop::Break(self))
        } else {
            Ok(Loop::Continue(self))
        }
    }

    fn post_loop(ctx: Result<Self>) -> Result<Resp> {
        let ctx = ctx.expect("end loop with Ok(_)");
        ctx.resp.unwrap_or_else(|| Err(box_err!("fail to request")))
    }

    /// Returns a Future, it is resolves once a future returned by the closure
    /// is resolved successfully, otherwise it repeats `retry` times.
    pub fn execute(self) -> PdFuture<Resp> {
        let ctx = self;
        loop_fn(ctx, |ctx| {
                ctx.reconnect_if_needed()
                    .and_then(Self::send)
                    .and_then(Self::receive)
            })
            .then(Self::post_loop)
            .boxed()
    }
}

/// Do a request in synchronized fashion.
pub fn sync_request<F, R>(client: &LeaderClient, retry: usize, f: F) -> Result<R>
    where F: Fn(&PDAsyncClient) -> GrpcFutureSend<R>
{
    for _ in 0..retry {
        let r = {
            let timer = PD_SEND_MSG_HISTOGRAM.start_timer();
            let r = f(&client.inner.rl().client).wait();
            timer.observe_duration();
            r
        };

        match r {
            Ok(r) => {
                return Ok(r);
            }
            Err(e) => {
                error!("fail to request: {:?}", e);
                client.reconnect()
            }
        }
    }

    Err(box_err!("fail to request"))
}

pub fn validate_endpoints(endpoints: &[String]) -> Result<(PDAsyncClient, GetMembersResponse)> {
    if endpoints.is_empty() {
        return Err(box_err!("empty PD endpoints"));
    }

    let len = endpoints.len();
    let mut endpoints_set = HashSet::with_capacity(len);

    let mut members = None;
    let mut cluster_id = None;
    for ep in endpoints {
        if !endpoints_set.insert(ep) {
            return Err(box_err!("duplicate PD endpoint {}", ep));
        }

        let (_, resp) = match connect(ep) {
            Ok(resp) => resp,
            // Ignore failed PD node.
            Err(e) => {
                error!("PD endpoint {} failed to respond: {:?}", ep, e);
                continue;
            }
        };

        // Check cluster ID.
        let cid = resp.get_header().get_cluster_id();
        if let Some(sample) = cluster_id {
            if sample != cid {
                return Err(box_err!("PD response cluster_id mismatch, want {}, got {}",
                                    sample,
                                    cid));
            }
        } else {
            cluster_id = Some(cid);
        }
        // TODO: check all fields later?

        if members.is_none() {
            members = Some(resp);
        }
    }

    match members {
        Some(members) => {
            let (client, members) = try!(try_connect_leader(&members));
            info!("All PD endpoints are consistent: {:?}", endpoints);
            Ok((client, members))
        }
        _ => Err(box_err!("PD cluster failed to respond")),
    }
}

fn connect(addr: &str) -> Result<(PDAsyncClient, GetMembersResponse)> {
    debug!("connect to PD endpoint: {:?}", addr);
    let ep = box_try!(Url::parse(addr));
    let host = ep.host_str().unwrap();
    let port = ep.port().unwrap();

    let mut conf: grpc::client::GrpcClientConf = Default::default();
    conf.http.no_delay = Some(true);

    // TODO: It seems that `new` always return an Ok(_).
    PDAsyncClient::new(host, port, false, conf)
        .and_then(|client| {
            // try request.
            match client.GetMembers(GetMembersRequest::new()).wait() {
                Ok(resp) => Ok((client, resp)),
                Err(e) => Err(e),
            }
        })
        .map_err(Error::Grpc)
}

pub fn try_connect_leader(previous: &GetMembersResponse)
                          -> Result<(PDAsyncClient, GetMembersResponse)> {
    // Try to connect other members.
    // Randomize endpoints.
    let members = previous.get_members();
    let mut indexes: Vec<usize> = (0..members.len()).collect();
    rand::thread_rng().shuffle(&mut indexes);

    let mut resp = None;
    'outer: for i in indexes {
        for ep in members[i].get_client_urls() {
            match connect(ep.as_str()) {
                Ok((_, r)) => {
                    resp = Some(r);
                    break 'outer;
                }
                Err(e) => {
                    error!("failed to connect to {}, {:?}", ep, e);
                    continue;
                }
            }
        }
    }

    // Then try to connect the PD cluster leader.
    if let Some(resp) = resp {
        let leader = resp.get_leader().clone();
        for ep in leader.get_client_urls() {
            if let Ok((client, _)) = connect(ep.as_str()) {
                info!("connect to PD leader {:?}", ep);
                return Ok((client, resp));
            }
        }
    }

    Err(box_err!("failed to connect to {:?}", members))
}

pub fn check_resp_header(header: &ResponseHeader) -> Result<()> {
    if !header.has_error() {
        return Ok(());
    }
    // TODO: translate more error types
    let err = header.get_error();
    match err.get_field_type() {
        ErrorType::ALREADY_BOOTSTRAPPED => Err(Error::ClusterBootstrapped(header.get_cluster_id())),
        ErrorType::NOT_BOOTSTRAPPED => Err(Error::ClusterNotBootstrapped(header.get_cluster_id())),
        _ => Err(box_err!(err.get_message())),
    }
}
<commit_msg>pd: reduce RwLockReadGuard lifetime (#1781)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;
use std::sync::RwLock;
use std::time::Instant;
use std::time::Duration;
use std::thread;
use std::collections::HashSet;

use futures::Future;
use futures::future::ok;
use futures::future::{loop_fn, Loop};

use grpc;
use grpc::futures_grpc::GrpcFutureSend;
use url::Url;
use rand::{self, Rng};

use kvproto::pdpb::{ResponseHeader, ErrorType, GetMembersRequest, GetMembersResponse, Member};
use kvproto::pdpb_grpc::{PDAsyncClient, PDAsync};

use util::HandyRwLock;

use super::super::{Result, Error, PdFuture};
use super::super::metrics::PD_SEND_MSG_HISTOGRAM;

pub struct Inner {
    client: PDAsyncClient,
    members: GetMembersResponse,
}

/// A leader client doing requests asynchronous.
pub struct LeaderClient {
    inner: Arc<RwLock<Inner>>,
}

impl LeaderClient {
    pub fn new(client: PDAsyncClient, members: GetMembersResponse) -> LeaderClient {
        LeaderClient {
            inner: Arc::new(RwLock::new(Inner {
                client: client,
                members: members,
            })),
        }
    }

    pub fn request<Req, Resp, F>(&self, req: Req, f: F, retry: usize) -> Request<Req, Resp, F>
        where Req: Clone + 'static,
              F: FnMut(&PDAsyncClient, Req) -> PdFuture<Resp> + Send + 'static
    {
        Request {
            reconnect_count: retry,
            request_sent: 0,
            client: LeaderClient { inner: self.inner.clone() },
            req: req,
            resp: None,
            func: f,
        }
    }

    pub fn get_leader(&self) -> Member {
        self.inner.rl().members.get_leader().clone()
    }

    // Re-establish connection with PD leader in synchronized fashion.
    pub fn reconnect(&self) {
        let start = Instant::now();
        let ret = try_connect_leader(&self.inner.rl().members);
        match ret {
            Ok((client, members)) => {
                let mut inner = self.inner.wl();
                inner.client = client;
                inner.members = members;
                warn!("updating PD client done, spent {:?}", start.elapsed());
            }

            Err(err) => {
                warn!("updating PD client spent {:?}, err {:?}",
                      start.elapsed(),
                      err);
                // FIXME: use tokio-timer instead.
                thread::sleep(Duration::from_secs(RETRY_INTERVAL));
            }
        }
    }
}

const RETRY_INTERVAL: u64 = 1;

/// The context of sending requets.
pub struct Request<Req, Resp, F> {
    reconnect_count: usize,
    request_sent: usize,

    client: LeaderClient,

    req: Req,
    resp: Option<Result<Resp>>,
    func: F,
}

const MAX_REQUEST_COUNT: usize = 5;

impl<Req, Resp, F> Request<Req, Resp, F>
    where Req: Clone + Send + 'static,
          Resp: Send + 'static,
          F: FnMut(&PDAsyncClient, Req) -> PdFuture<Resp> + Send + 'static
{
    fn reconnect_if_needed(mut self) -> PdFuture<Self> {
        debug!("reconnect remains: {}", self.reconnect_count);

        if self.request_sent == 0 {
            return ok(self).boxed();
        }

        if self.request_sent < MAX_REQUEST_COUNT {
            debug!("retry on the same client");
            return ok(self).boxed();
        }

        // Updating client.

        self.request_sent = 0;
        self.reconnect_count -= 1;

        // FIXME: should not block the core.
        warn!("updating PD client, block the tokio core");
        self.client.reconnect();

        ok(self).boxed()
    }

    fn send(mut self) -> PdFuture<Self> {
        self.request_sent += 1;
        debug!("request sent: {}", self.request_sent);
        let r = self.req.clone();
        let req = (self.func)(&self.client.inner.rl().client, r);
        req.then(|resp| {
                match resp {
                    Ok(resp) => self.resp = Some(Ok(resp)),
                    Err(err) => {
                        error!("request failed: {:?}", err);
                    }
                };
                Ok(self)
            })
            .boxed()
    }

    fn receive(self) -> Result<Loop<Self, Self>> {
        let done = self.reconnect_count == 0 || self.resp.is_some();
        if done {
            Ok(Loop::Break(self))
        } else {
            Ok(Loop::Continue(self))
        }
    }

    fn post_loop(ctx: Result<Self>) -> Result<Resp> {
        let ctx = ctx.expect("end loop with Ok(_)");
        ctx.resp.unwrap_or_else(|| Err(box_err!("fail to request")))
    }

    /// Returns a Future, it is resolves once a future returned by the closure
    /// is resolved successfully, otherwise it repeats `retry` times.
    pub fn execute(self) -> PdFuture<Resp> {
        let ctx = self;
        loop_fn(ctx, |ctx| {
                ctx.reconnect_if_needed()
                    .and_then(Self::send)
                    .and_then(Self::receive)
            })
            .then(Self::post_loop)
            .boxed()
    }
}

/// Do a request in synchronized fashion.
pub fn sync_request<F, R>(client: &LeaderClient, retry: usize, func: F) -> Result<R>
    where F: Fn(&PDAsyncClient) -> GrpcFutureSend<R>
{
    for _ in 0..retry {
        let r = {
            let _timer = PD_SEND_MSG_HISTOGRAM.start_timer(); // observe on dropping.
            let f = func(&client.inner.rl().client);
            f.wait()
        };

        match r {
            Ok(r) => {
                return Ok(r);
            }
            Err(e) => {
                error!("fail to request: {:?}", e);
                client.reconnect()
            }
        }
    }

    Err(box_err!("fail to request"))
}

pub fn validate_endpoints(endpoints: &[String]) -> Result<(PDAsyncClient, GetMembersResponse)> {
    if endpoints.is_empty() {
        return Err(box_err!("empty PD endpoints"));
    }

    let len = endpoints.len();
    let mut endpoints_set = HashSet::with_capacity(len);

    let mut members = None;
    let mut cluster_id = None;
    for ep in endpoints {
        if !endpoints_set.insert(ep) {
            return Err(box_err!("duplicate PD endpoint {}", ep));
        }

        let (_, resp) = match connect(ep) {
            Ok(resp) => resp,
            // Ignore failed PD node.
            Err(e) => {
                error!("PD endpoint {} failed to respond: {:?}", ep, e);
                continue;
            }
        };

        // Check cluster ID.
        let cid = resp.get_header().get_cluster_id();
        if let Some(sample) = cluster_id {
            if sample != cid {
                return Err(box_err!("PD response cluster_id mismatch, want {}, got {}",
                                    sample,
                                    cid));
            }
        } else {
            cluster_id = Some(cid);
        }
        // TODO: check all fields later?

        if members.is_none() {
            members = Some(resp);
        }
    }

    match members {
        Some(members) => {
            let (client, members) = try!(try_connect_leader(&members));
            info!("All PD endpoints are consistent: {:?}", endpoints);
            Ok((client, members))
        }
        _ => Err(box_err!("PD cluster failed to respond")),
    }
}

fn connect(addr: &str) -> Result<(PDAsyncClient, GetMembersResponse)> {
    debug!("connect to PD endpoint: {:?}", addr);
    let ep = box_try!(Url::parse(addr));
    let host = ep.host_str().unwrap();
    let port = ep.port().unwrap();

    let mut conf: grpc::client::GrpcClientConf = Default::default();
    conf.http.no_delay = Some(true);

    // TODO: It seems that `new` always return an Ok(_).
    PDAsyncClient::new(host, port, false, conf)
        .and_then(|client| {
            // try request.
            match client.GetMembers(GetMembersRequest::new()).wait() {
                Ok(resp) => Ok((client, resp)),
                Err(e) => Err(e),
            }
        })
        .map_err(Error::Grpc)
}

pub fn try_connect_leader(previous: &GetMembersResponse)
                          -> Result<(PDAsyncClient, GetMembersResponse)> {
    // Try to connect other members.
    // Randomize endpoints.
    let members = previous.get_members();
    let mut indexes: Vec<usize> = (0..members.len()).collect();
    rand::thread_rng().shuffle(&mut indexes);

    let mut resp = None;
    'outer: for i in indexes {
        for ep in members[i].get_client_urls() {
            match connect(ep.as_str()) {
                Ok((_, r)) => {
                    resp = Some(r);
                    break 'outer;
                }
                Err(e) => {
                    error!("failed to connect to {}, {:?}", ep, e);
                    continue;
                }
            }
        }
    }

    // Then try to connect the PD cluster leader.
    if let Some(resp) = resp {
        let leader = resp.get_leader().clone();
        for ep in leader.get_client_urls() {
            if let Ok((client, _)) = connect(ep.as_str()) {
                info!("connect to PD leader {:?}", ep);
                return Ok((client, resp));
            }
        }
    }

    Err(box_err!("failed to connect to {:?}", members))
}

pub fn check_resp_header(header: &ResponseHeader) -> Result<()> {
    if !header.has_error() {
        return Ok(());
    }
    // TODO: translate more error types
    let err = header.get_error();
    match err.get_field_type() {
        ErrorType::ALREADY_BOOTSTRAPPED => Err(Error::ClusterBootstrapped(header.get_cluster_id())),
        ErrorType::NOT_BOOTSTRAPPED => Err(Error::ClusterNotBootstrapped(header.get_cluster_id())),
        _ => Err(box_err!(err.get_message())),
    }
}
<|endoftext|>


<commit_before>/// Windows implementation of dynamic library loading.
///
/// This module should be expanded with more Windows-specific functionality in the future.

extern crate winapi;
extern crate kernel32;

use std::ffi::{CStr, OsStr, OsString};
use std::marker;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::os::windows::ffi::{OsStrExt, OsStringExt};

/// A platform-specific equivalent of the cross-platform `Library`.
pub struct Library(winapi::HMODULE);

impl Library {
    /// Find and load a shared library (module).
    ///
    /// Locations where library is searched for is platform specific and can’t be adjusted
    /// portably.
    ///
    /// Corresponds to `LoadLibraryW(filename)`.
    #[inline]
    pub fn new<P: AsRef<OsStr>>(filename: P) -> ::Result<Library> {
        let mut wide_filename: Vec<u16> = filename.as_ref().encode_wide().collect();
        wide_filename.push(0);
        wide_filename.shrink_to_fit();
        let _guard = ErrorModeGuard::new();

        let ret = with_get_last_error(|| {
            // Make sure no winapi calls as a result of drop happen inside this closure, because
            // otherwise that might change the return value of the GetLastError.
            let handle = unsafe { kernel32::LoadLibraryW(wide_filename.as_ptr()) };
            if handle.is_null()  {
                None
            } else {
                Some(Library(handle))
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("LoadLibraryW failed but GetLastError did not report the error")
        ));

        drop(wide_filename); // Drop wide_filename here to ensure it doesn’t get moved and dropped
                             // inside the closure by mistake.
        ret
    }

    /// Get a symbol by name.
    ///
    /// Mangling or symbol rustification is not done: trying to `get` something like `x::y`
    /// will not work.
    ///
    /// # Unsafety
    ///
    /// The pointer to a symbol of arbitrary type or kind is returned. Requesting for function
    /// pointer while the symbol is not one and vice versa is not memory safe.
    ///
    /// The return value does not ensure the symbol does not outlive the library.
    pub unsafe fn get<T>(&self, symbol: &CStr) -> ::Result<Symbol<T>> {
        with_get_last_error(|| {
            let symbol = kernel32::GetProcAddress(self.0, symbol.as_ptr());
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }
}

impl Drop for Library {
    fn drop(&mut self) {
        with_get_last_error(|| {
            if unsafe { kernel32::FreeLibrary(self.0) == 0 } {
                None
            } else {
                Some(())
            }
        }).unwrap()
    }
}

impl ::std::fmt::Debug for Library {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        unsafe {
            let mut buf: [winapi::WCHAR; 1024] = mem::uninitialized();
            let len = kernel32::GetModuleFileNameW(self.0,
                                                   (&mut buf[..]).as_mut_ptr(), 1024) as usize;
            if len == 0 {
                f.write_str(&format!("Library@{:p}", self.0))
            } else {
                let string: OsString = OsString::from_wide(&buf[..len]);
                f.write_str(&format!("Library@{:p} from {:?}", self.0, string))
            }
        }
    }
}

pub struct Symbol<T> {
    pointer: winapi::FARPROC,
    pd: marker::PhantomData<T>
}

impl<T> ::std::ops::Deref for Symbol<T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe {
            &*(&self.pointer as *const _ as *const T)
        }
    }
}

impl<T> ::std::fmt::Debug for Symbol<T> {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        f.write_str(&format!("Symbol@{:p}", self.pointer))
    }
}

static USE_ERRORMODE: AtomicBool = ATOMIC_BOOL_INIT;
struct ErrorModeGuard(winapi::DWORD);

impl ErrorModeGuard {
    fn new() -> ErrorModeGuard {
        let mut ret = ErrorModeGuard(0);

        if !USE_ERRORMODE.load(Ordering::Acquire) {
            if unsafe { kernel32::SetThreadErrorMode(1, &mut ret.0) == 0
                        && kernel32::GetLastError() == winapi::ERROR_CALL_NOT_IMPLEMENTED } {
                USE_ERRORMODE.store(true, Ordering::Release);
            } else {
                return ret;
            }
        }
        ret.0 = unsafe { kernel32::SetErrorMode(1) };
        ret
    }
}

impl Drop for ErrorModeGuard {
    fn drop(&mut self) {
        unsafe {
            if !USE_ERRORMODE.load(Ordering::Relaxed) {
                kernel32::SetThreadErrorMode(self.0, ptr::null_mut());
            } else {
                kernel32::SetErrorMode(self.0);
            }
        }
    }
}

fn with_get_last_error<T, F>(closure: F) -> Result<T, Option<String>>
where F: FnOnce() -> Option<T> {
    closure().ok_or_else(|| {
        let error = unsafe { kernel32::GetLastError() };
        if error == 0 {
            None
        } else {
            // TODO: Possibly use FormatMessage (lots of work here)
            Some(format!("Error {}", error))
        }
    })
}

#[test]
fn works_new_kernel32() {
    let that = Library::new("kernel32.dll").unwrap();
    unsafe {
        that.get::<*mut usize>(&::std::ffi::CString::new("GetLastError").unwrap()).unwrap();
    }
}

#[test]
fn fails_new_kernel23() {
    Library::new("kernel23").err().unwrap();
}
<commit_msg>Fix windows import<commit_after>/// Windows implementation of dynamic library loading.
///
/// This module should be expanded with more Windows-specific functionality in the future.

extern crate winapi;
extern crate kernel32;

use std::ffi::{CStr, OsStr, OsString};
use std::marker;
use std::mem;
use std::ptr;
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
use std::os::windows::ffi::{OsStrExt, OsStringExt};

/// A platform-specific equivalent of the cross-platform `Library`.
pub struct Library(winapi::HMODULE);

impl Library {
    /// Find and load a shared library (module).
    ///
    /// Locations where library is searched for is platform specific and can’t be adjusted
    /// portably.
    ///
    /// Corresponds to `LoadLibraryW(filename)`.
    #[inline]
    pub fn new<P: AsRef<OsStr>>(filename: P) -> ::Result<Library> {
        let mut wide_filename: Vec<u16> = filename.as_ref().encode_wide().collect();
        wide_filename.push(0);
        wide_filename.shrink_to_fit();
        let _guard = ErrorModeGuard::new();

        let ret = with_get_last_error(|| {
            // Make sure no winapi calls as a result of drop happen inside this closure, because
            // otherwise that might change the return value of the GetLastError.
            let handle = unsafe { kernel32::LoadLibraryW(wide_filename.as_ptr()) };
            if handle.is_null()  {
                None
            } else {
                Some(Library(handle))
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("LoadLibraryW failed but GetLastError did not report the error")
        ));

        drop(wide_filename); // Drop wide_filename here to ensure it doesn’t get moved and dropped
                             // inside the closure by mistake.
        ret
    }

    /// Get a symbol by name.
    ///
    /// Mangling or symbol rustification is not done: trying to `get` something like `x::y`
    /// will not work.
    ///
    /// # Unsafety
    ///
    /// The pointer to a symbol of arbitrary type or kind is returned. Requesting for function
    /// pointer while the symbol is not one and vice versa is not memory safe.
    ///
    /// The return value does not ensure the symbol does not outlive the library.
    pub unsafe fn get<T>(&self, symbol: &CStr) -> ::Result<Symbol<T>> {
        with_get_last_error(|| {
            let symbol = kernel32::GetProcAddress(self.0, symbol.as_ptr());
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }
}

impl Drop for Library {
    fn drop(&mut self) {
        with_get_last_error(|| {
            if unsafe { kernel32::FreeLibrary(self.0) == 0 } {
                None
            } else {
                Some(())
            }
        }).unwrap()
    }
}

impl ::std::fmt::Debug for Library {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        unsafe {
            let mut buf: [winapi::WCHAR; 1024] = mem::uninitialized();
            let len = kernel32::GetModuleFileNameW(self.0,
                                                   (&mut buf[..]).as_mut_ptr(), 1024) as usize;
            if len == 0 {
                f.write_str(&format!("Library@{:p}", self.0))
            } else {
                let string: OsString = OsString::from_wide(&buf[..len]);
                f.write_str(&format!("Library@{:p} from {:?}", self.0, string))
            }
        }
    }
}

pub struct Symbol<T> {
    pointer: winapi::FARPROC,
    pd: marker::PhantomData<T>
}

impl<T> ::std::ops::Deref for Symbol<T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe {
            &*(&self.pointer as *const _ as *const T)
        }
    }
}

impl<T> ::std::fmt::Debug for Symbol<T> {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        f.write_str(&format!("Symbol@{:p}", self.pointer))
    }
}

static USE_ERRORMODE: AtomicBool = ATOMIC_BOOL_INIT;
struct ErrorModeGuard(winapi::DWORD);

impl ErrorModeGuard {
    fn new() -> ErrorModeGuard {
        let mut ret = ErrorModeGuard(0);

        if !USE_ERRORMODE.load(Ordering::Acquire) {
            if unsafe { kernel32::SetThreadErrorMode(1, &mut ret.0) == 0
                        && kernel32::GetLastError() == winapi::ERROR_CALL_NOT_IMPLEMENTED } {
                USE_ERRORMODE.store(true, Ordering::Release);
            } else {
                return ret;
            }
        }
        ret.0 = unsafe { kernel32::SetErrorMode(1) };
        ret
    }
}

impl Drop for ErrorModeGuard {
    fn drop(&mut self) {
        unsafe {
            if !USE_ERRORMODE.load(Ordering::Relaxed) {
                kernel32::SetThreadErrorMode(self.0, ptr::null_mut());
            } else {
                kernel32::SetErrorMode(self.0);
            }
        }
    }
}

fn with_get_last_error<T, F>(closure: F) -> Result<T, Option<String>>
where F: FnOnce() -> Option<T> {
    closure().ok_or_else(|| {
        let error = unsafe { kernel32::GetLastError() };
        if error == 0 {
            None
        } else {
            // TODO: Possibly use FormatMessage (lots of work here)
            Some(format!("Error {}", error))
        }
    })
}

#[test]
fn works_new_kernel32() {
    let that = Library::new("kernel32.dll").unwrap();
    unsafe {
        that.get::<*mut usize>(&::std::ffi::CString::new("GetLastError").unwrap()).unwrap();
    }
}

#[test]
fn fails_new_kernel23() {
    Library::new("kernel23").err().unwrap();
}
<|endoftext|>


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

//! Manages jailing and connecting virtio devices to the system bus.

use std::fmt;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::{Arc, Mutex};

use libc::STDERR_FILENO;

use io_jail::Minijail;
use kvm::{Vm, IoeventAddress};
use sys_util::{GuestMemory, syslog};
use sys_util;

use devices;
use kernel_cmdline;

/// Errors for device manager.
#[derive(Debug)]
pub enum Error {
    /// Could not create the mmio device to wrap a VirtioDevice.
    CreateMmioDevice(sys_util::Error),
    /// Failed to register ioevent with VM.
    RegisterIoevent(sys_util::Error),
    /// Failed to register irq eventfd with VM.
    RegisterIrqfd(sys_util::Error),
    /// Failed to initialize proxy device for jailed device.
    ProxyDeviceCreation(devices::ProxyError),
    /// Appending to kernel command line failed.
    Cmdline(kernel_cmdline::Error),
    /// No more IRQs are available.
    IrqsExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::CreateMmioDevice(ref e) => write!(f, "failed to create mmio device: {:?}", e),
            &Error::RegisterIoevent(ref e) => {
                write!(f, "failed to register ioevent to VM: {:?}", e)
            }
            &Error::RegisterIrqfd(ref e) => {
                write!(f, "failed to register irq eventfd to VM: {:?}", e)
            }
            &Error::ProxyDeviceCreation(ref e) => write!(f, "failed to create proxy device: {}", e),
            &Error::Cmdline(ref e) => {
                write!(f, "unable to add device to kernel command line: {}", e)
            }
            &Error::IrqsExhausted => write!(f, "no more IRQs are available"),
        }
    }
}

type Result<T> = ::std::result::Result<T, Error>;

const MAX_IRQ: u32 = 15;

/// Manages the complexities of adding a device.
pub struct DeviceManager<'a> {
    pub bus: devices::Bus,
    vm: &'a mut Vm,
    guest_mem: GuestMemory,
    mmio_len: u64,
    mmio_base: u64,
    irq: u32,
}

impl<'a> DeviceManager<'a> {
    /// Create a new DeviceManager.
    pub fn new(vm: &mut Vm,
               guest_mem: GuestMemory,
               mmio_len: u64,
               mmio_base: u64,
               irq_base: u32)
               -> DeviceManager {
        DeviceManager {
            bus: devices::Bus::new(),
            vm,
            guest_mem,
            mmio_len,
            mmio_base,
            irq: irq_base,
        }
    }

    /// Register a device to be used via MMIO transport.
    pub fn register_mmio(&mut self,
                         device: Box<devices::virtio::VirtioDevice>,
                         jail: Option<Minijail>,
                         cmdline: &mut kernel_cmdline::Cmdline)
                         -> Result<()> {
        if self.irq > MAX_IRQ {
            return Err(Error::IrqsExhausted);
        }

        // List of FDs to keep open in the child after it forks.
        let mut keep_fds: Vec<RawFd> = device.keep_fds();
        keep_fds.push(STDERR_FILENO);
        syslog::push_fds(&mut keep_fds);

        let mmio_device = devices::virtio::MmioDevice::new(self.guest_mem.clone(), device)
            .map_err(Error::CreateMmioDevice)?;
        for (i, queue_evt) in mmio_device.queue_evts().iter().enumerate() {
            let io_addr = IoeventAddress::Mmio(self.mmio_base +
                                               devices::virtio::NOITFY_REG_OFFSET as u64);
            self.vm
                .register_ioevent(&queue_evt, io_addr, i as u32)
                .map_err(Error::RegisterIoevent)?;
            keep_fds.push(queue_evt.as_raw_fd());
        }

        if let Some(interrupt_evt) = mmio_device.interrupt_evt() {
            self.vm
                .register_irqfd(&interrupt_evt, self.irq)
                .map_err(Error::RegisterIrqfd)?;
            keep_fds.push(interrupt_evt.as_raw_fd());
        }

        if let Some(jail) = jail {
            let proxy_dev = devices::ProxyDevice::new(mmio_device, &jail, keep_fds)
                .map_err(Error::ProxyDeviceCreation)?;

            self.bus
                .insert(Arc::new(Mutex::new(proxy_dev)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        } else {
            self.bus
                .insert(Arc::new(Mutex::new(mmio_device)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        }

        cmdline
            .insert("virtio_mmio.device",
                    &format!("4K@0x{:08x}:{}", self.mmio_base, self.irq))
            .map_err(Error::Cmdline)?;
        self.mmio_base += self.mmio_len;
        self.irq += 1;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;
    use std::os::unix::io::RawFd;
    use sys_util::{EventFd, GuestAddress, GuestMemory};
    use kvm::*;
    use device_manager;
    use kernel_cmdline;
    use devices;

    const QUEUE_SIZES: &'static [u16] = &[64];

    #[allow(dead_code)]
    struct DummyDevice {
        dummy: u32,
    }

    impl devices::virtio::VirtioDevice for DummyDevice {
        fn keep_fds(&self) -> Vec<RawFd> {
            Vec::new()
        }

        fn device_type(&self) -> u32 {
            0
        }

        fn queue_max_sizes(&self) -> &[u16] {
            QUEUE_SIZES
        }

        #[allow(unused_variables)]
        #[allow(unused_mut)]
        fn activate(&mut self,
                    mem: GuestMemory,
                    interrupt_evt: EventFd,
                    status: Arc<AtomicUsize>,
                    queues: Vec<devices::virtio::Queue>,
                    mut queue_evts: Vec<EventFd>) {
        }
    }

    #[test]
    #[ignore] // no access to /dev/kvm
    fn register_device() {
        let start_addr1 = GuestAddress(0x0);
        let start_addr2 = GuestAddress(0x1000);
        let guest_mem = GuestMemory::new(&vec![(start_addr1, 0x1000), (start_addr2, 0x1000)])
            .unwrap();
        let mut vm = Vm::new(&Kvm::new().unwrap(), guest_mem.clone()).unwrap();
        let mut device_manager =
            device_manager::DeviceManager::new(&mut vm, guest_mem, 0x1000, 0xd0000000, 5);

        let mut cmdline = kernel_cmdline::Cmdline::new(4096);
        let dummy_box = Box::new(DummyDevice { dummy: 0 });
        device_manager
            .register_mmio(dummy_box, None, &mut cmdline)
            .unwrap();
    }
}
<commit_msg>crosvm: remove stderr from preserved FDs<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Manages jailing and connecting virtio devices to the system bus.

use std::fmt;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::{Arc, Mutex};

use io_jail::Minijail;
use kvm::{Vm, IoeventAddress};
use sys_util::{GuestMemory, syslog};
use sys_util;

use devices;
use kernel_cmdline;

/// Errors for device manager.
#[derive(Debug)]
pub enum Error {
    /// Could not create the mmio device to wrap a VirtioDevice.
    CreateMmioDevice(sys_util::Error),
    /// Failed to register ioevent with VM.
    RegisterIoevent(sys_util::Error),
    /// Failed to register irq eventfd with VM.
    RegisterIrqfd(sys_util::Error),
    /// Failed to initialize proxy device for jailed device.
    ProxyDeviceCreation(devices::ProxyError),
    /// Appending to kernel command line failed.
    Cmdline(kernel_cmdline::Error),
    /// No more IRQs are available.
    IrqsExhausted,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::CreateMmioDevice(ref e) => write!(f, "failed to create mmio device: {:?}", e),
            &Error::RegisterIoevent(ref e) => {
                write!(f, "failed to register ioevent to VM: {:?}", e)
            }
            &Error::RegisterIrqfd(ref e) => {
                write!(f, "failed to register irq eventfd to VM: {:?}", e)
            }
            &Error::ProxyDeviceCreation(ref e) => write!(f, "failed to create proxy device: {}", e),
            &Error::Cmdline(ref e) => {
                write!(f, "unable to add device to kernel command line: {}", e)
            }
            &Error::IrqsExhausted => write!(f, "no more IRQs are available"),
        }
    }
}

type Result<T> = ::std::result::Result<T, Error>;

const MAX_IRQ: u32 = 15;

/// Manages the complexities of adding a device.
pub struct DeviceManager<'a> {
    pub bus: devices::Bus,
    vm: &'a mut Vm,
    guest_mem: GuestMemory,
    mmio_len: u64,
    mmio_base: u64,
    irq: u32,
}

impl<'a> DeviceManager<'a> {
    /// Create a new DeviceManager.
    pub fn new(vm: &mut Vm,
               guest_mem: GuestMemory,
               mmio_len: u64,
               mmio_base: u64,
               irq_base: u32)
               -> DeviceManager {
        DeviceManager {
            bus: devices::Bus::new(),
            vm,
            guest_mem,
            mmio_len,
            mmio_base,
            irq: irq_base,
        }
    }

    /// Register a device to be used via MMIO transport.
    pub fn register_mmio(&mut self,
                         device: Box<devices::virtio::VirtioDevice>,
                         jail: Option<Minijail>,
                         cmdline: &mut kernel_cmdline::Cmdline)
                         -> Result<()> {
        if self.irq > MAX_IRQ {
            return Err(Error::IrqsExhausted);
        }

        // List of FDs to keep open in the child after it forks.
        let mut keep_fds: Vec<RawFd> = device.keep_fds();
        syslog::push_fds(&mut keep_fds);

        let mmio_device = devices::virtio::MmioDevice::new(self.guest_mem.clone(), device)
            .map_err(Error::CreateMmioDevice)?;
        for (i, queue_evt) in mmio_device.queue_evts().iter().enumerate() {
            let io_addr = IoeventAddress::Mmio(self.mmio_base +
                                               devices::virtio::NOITFY_REG_OFFSET as u64);
            self.vm
                .register_ioevent(&queue_evt, io_addr, i as u32)
                .map_err(Error::RegisterIoevent)?;
            keep_fds.push(queue_evt.as_raw_fd());
        }

        if let Some(interrupt_evt) = mmio_device.interrupt_evt() {
            self.vm
                .register_irqfd(&interrupt_evt, self.irq)
                .map_err(Error::RegisterIrqfd)?;
            keep_fds.push(interrupt_evt.as_raw_fd());
        }

        if let Some(jail) = jail {
            let proxy_dev = devices::ProxyDevice::new(mmio_device, &jail, keep_fds)
                .map_err(Error::ProxyDeviceCreation)?;

            self.bus
                .insert(Arc::new(Mutex::new(proxy_dev)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        } else {
            self.bus
                .insert(Arc::new(Mutex::new(mmio_device)),
                        self.mmio_base,
                        self.mmio_len)
                .unwrap();
        }

        cmdline
            .insert("virtio_mmio.device",
                    &format!("4K@0x{:08x}:{}", self.mmio_base, self.irq))
            .map_err(Error::Cmdline)?;
        self.mmio_base += self.mmio_len;
        self.irq += 1;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;
    use std::os::unix::io::RawFd;
    use sys_util::{EventFd, GuestAddress, GuestMemory};
    use kvm::*;
    use device_manager;
    use kernel_cmdline;
    use devices;

    const QUEUE_SIZES: &'static [u16] = &[64];

    #[allow(dead_code)]
    struct DummyDevice {
        dummy: u32,
    }

    impl devices::virtio::VirtioDevice for DummyDevice {
        fn keep_fds(&self) -> Vec<RawFd> {
            Vec::new()
        }

        fn device_type(&self) -> u32 {
            0
        }

        fn queue_max_sizes(&self) -> &[u16] {
            QUEUE_SIZES
        }

        #[allow(unused_variables)]
        #[allow(unused_mut)]
        fn activate(&mut self,
                    mem: GuestMemory,
                    interrupt_evt: EventFd,
                    status: Arc<AtomicUsize>,
                    queues: Vec<devices::virtio::Queue>,
                    mut queue_evts: Vec<EventFd>) {
        }
    }

    #[test]
    #[ignore] // no access to /dev/kvm
    fn register_device() {
        let start_addr1 = GuestAddress(0x0);
        let start_addr2 = GuestAddress(0x1000);
        let guest_mem = GuestMemory::new(&vec![(start_addr1, 0x1000), (start_addr2, 0x1000)])
            .unwrap();
        let mut vm = Vm::new(&Kvm::new().unwrap(), guest_mem.clone()).unwrap();
        let mut device_manager =
            device_manager::DeviceManager::new(&mut vm, guest_mem, 0x1000, 0xd0000000, 5);

        let mut cmdline = kernel_cmdline::Cmdline::new(4096);
        let dummy_box = Box::new(DummyDevice { dummy: 0 });
        device_manager
            .register_mmio(dummy_box, None, &mut cmdline)
            .unwrap();
    }
}
<|endoftext|>


<commit_before>//  Copyright (C) 2020 Éloïs SANCHEZ.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::*;

fn gen_start_args(args: &DuniterStartArgs, duniter_ts_args: &mut Vec<String>) {
    if let Some(ref keyfile) = args.keyfile {
        duniter_ts_args.push("--keyfile".to_owned());
        duniter_ts_args.push(
            keyfile
                .to_str()
                .expect("keyfile path is invalid")
                .to_owned(),
        );
    }
}

fn gen_webstart_args(args: &DuniterWebstartArgs, duniter_ts_args: &mut Vec<String>) {
    if let Some(ref web_ui_host) = args.web_ui_host {
        duniter_ts_args.push("--webmhost".to_owned());
        duniter_ts_args.push(web_ui_host.into());
    }
    if let Some(ref web_ui_port) = args.web_ui_port {
        duniter_ts_args.push("--webmport".to_owned());
        duniter_ts_args.push(web_ui_port.into());
    }
}

pub(crate) fn gen_duniter_ts_args(args: &DuniterArgs, duniter_js_exe: String) -> Vec<String> {
    let mut duniter_ts_args = vec![duniter_js_exe];
    if let Some(ref home) = args.home {
        duniter_ts_args.push("--home".to_owned());
        duniter_ts_args.push(home.to_str().expect("invalid home path").to_owned());
    }
    if let Some(ref log_level) = args.log {
        duniter_ts_args.push("--loglevel".to_owned());
        duniter_ts_args.push(log_level.to_string().to_lowercase());
    }
    if let Some(ref profile) = args.profile {
        duniter_ts_args.push("--mdb".to_owned());
        duniter_ts_args.push(profile.clone());
    }
    match args.command {
        DuniterCommand::Completions { .. } => unreachable!(),
        DuniterCommand::DirectStart {
            keyprompt,
            ref start_args,
        } => {
            duniter_ts_args.push("direct_start".to_owned());
            if keyprompt {
                duniter_ts_args.push("--keyprompt".to_owned());
            }
            gen_start_args(start_args, &mut duniter_ts_args);
        }
        DuniterCommand::DirectWebstart {
            keyprompt,
            ref start_args,
            ref webstart_args,
        } => {
            duniter_ts_args.push("direct_webstart".to_owned());
            if keyprompt {
                duniter_ts_args.push("--keyprompt".to_owned());
            }
            gen_start_args(start_args, &mut duniter_ts_args);
            gen_webstart_args(webstart_args, &mut duniter_ts_args);
        }
        #[cfg(feature = "gva")]
        DuniterCommand::Gva(_) => unreachable!(),
        DuniterCommand::Start(ref start_args) => {
            duniter_ts_args.push("direct_start".to_owned());
            gen_start_args(start_args, &mut duniter_ts_args);
        }
        DuniterCommand::Webstart {
            ref start_args,
            ref webstart_args,
        } => {
            duniter_ts_args.push("direct_webstart".to_owned());
            gen_start_args(start_args, &mut duniter_ts_args);
            gen_webstart_args(webstart_args, &mut duniter_ts_args);
        }
        DuniterCommand::Stop => duniter_ts_args.push("stop".to_owned()),
        DuniterCommand::Sync(ref sync_args) => {
            duniter_ts_args.push("sync".to_owned());
            sync::gen_args(sync_args, &mut duniter_ts_args);
        }
        DuniterCommand::Config(ref config_args) => {
            duniter_ts_args.push("config".to_owned());
            config::gen_args(config_args, &mut duniter_ts_args);
        }
        DuniterCommand::Reset(ref reset_command) => {
            duniter_ts_args.push("reset".to_owned());
            match reset_command {
                ResetCommand::Config => duniter_ts_args.push("config".to_owned()),
                ResetCommand::Data => duniter_ts_args.push("data".to_owned()),
                ResetCommand::Peers => duniter_ts_args.push("peers".to_owned()),
                ResetCommand::Stats => duniter_ts_args.push("stats".to_owned()),
                ResetCommand::All => duniter_ts_args.push("all".to_owned()),
            }
        }
        DuniterCommand::Wizard(ref wizard_command) => {
            duniter_ts_args.push("wizard".to_owned());
            match wizard_command {
                WizardCommand::Bma => duniter_ts_args.push("network".to_owned()),
                WizardCommand::Key { n, r, p } => {
                    duniter_ts_args.push("key".to_owned());
                    if let Some(n) = n {
                        duniter_ts_args.push("--keyN".to_owned());
                        duniter_ts_args.push(n.to_string());
                    }
                    if let Some(r) = r {
                        duniter_ts_args.push("--keyr".to_owned());
                        duniter_ts_args.push(r.to_string());
                    }
                    if let Some(p) = p {
                        duniter_ts_args.push("--keyp".to_owned());
                        duniter_ts_args.push(p.to_string());
                    }
                }
            }
        }
        DuniterCommand::WS2P(ref ws2p_command) => {
            duniter_ts_args.push("ws2p".to_owned());
            match ws2p_command {
                WS2PCommand::ListNodes => duniter_ts_args.push("list-nodes".to_owned()),
                WS2PCommand::ListPrefered => duniter_ts_args.push("list-prefered".to_owned()),
                WS2PCommand::ListPrivileged => duniter_ts_args.push("list-privileged".to_owned()),
                WS2PCommand::ShowConf => duniter_ts_args.push("show-conf".to_owned()),
            }
        }
        DuniterCommand::Logs | DuniterCommand::Restart | DuniterCommand::Status => {}
    }
    duniter_ts_args
}
<commit_msg>fix: node v8 out of memory<commit_after>//  Copyright (C) 2020 Éloïs SANCHEZ.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

use crate::*;

fn gen_start_args(args: &DuniterStartArgs, duniter_ts_args: &mut Vec<String>) {
    if let Some(ref keyfile) = args.keyfile {
        duniter_ts_args.push("--keyfile".to_owned());
        duniter_ts_args.push(
            keyfile
                .to_str()
                .expect("keyfile path is invalid")
                .to_owned(),
        );
    }
}

fn gen_webstart_args(args: &DuniterWebstartArgs, duniter_ts_args: &mut Vec<String>) {
    if let Some(ref web_ui_host) = args.web_ui_host {
        duniter_ts_args.push("--webmhost".to_owned());
        duniter_ts_args.push(web_ui_host.into());
    }
    if let Some(ref web_ui_port) = args.web_ui_port {
        duniter_ts_args.push("--webmport".to_owned());
        duniter_ts_args.push(web_ui_port.into());
    }
}

pub(crate) fn gen_duniter_ts_args(args: &DuniterArgs, duniter_js_exe: String) -> Vec<String> {
    let mut duniter_ts_args = vec![
        "--max-old-space-size".to_owned(),
        "4096".to_owned(),
        duniter_js_exe,
    ];
    if let Some(ref home) = args.home {
        duniter_ts_args.push("--home".to_owned());
        duniter_ts_args.push(home.to_str().expect("invalid home path").to_owned());
    }
    if let Some(ref log_level) = args.log {
        duniter_ts_args.push("--loglevel".to_owned());
        duniter_ts_args.push(log_level.to_string().to_lowercase());
    }
    if let Some(ref profile) = args.profile {
        duniter_ts_args.push("--mdb".to_owned());
        duniter_ts_args.push(profile.clone());
    }
    match args.command {
        DuniterCommand::Completions { .. } => unreachable!(),
        DuniterCommand::DirectStart {
            keyprompt,
            ref start_args,
        } => {
            duniter_ts_args.push("direct_start".to_owned());
            if keyprompt {
                duniter_ts_args.push("--keyprompt".to_owned());
            }
            gen_start_args(start_args, &mut duniter_ts_args);
        }
        DuniterCommand::DirectWebstart {
            keyprompt,
            ref start_args,
            ref webstart_args,
        } => {
            duniter_ts_args.push("direct_webstart".to_owned());
            if keyprompt {
                duniter_ts_args.push("--keyprompt".to_owned());
            }
            gen_start_args(start_args, &mut duniter_ts_args);
            gen_webstart_args(webstart_args, &mut duniter_ts_args);
        }
        #[cfg(feature = "gva")]
        DuniterCommand::Gva(_) => unreachable!(),
        DuniterCommand::Start(ref start_args) => {
            duniter_ts_args.push("direct_start".to_owned());
            gen_start_args(start_args, &mut duniter_ts_args);
        }
        DuniterCommand::Webstart {
            ref start_args,
            ref webstart_args,
        } => {
            duniter_ts_args.push("direct_webstart".to_owned());
            gen_start_args(start_args, &mut duniter_ts_args);
            gen_webstart_args(webstart_args, &mut duniter_ts_args);
        }
        DuniterCommand::Stop => duniter_ts_args.push("stop".to_owned()),
        DuniterCommand::Sync(ref sync_args) => {
            duniter_ts_args.push("sync".to_owned());
            sync::gen_args(sync_args, &mut duniter_ts_args);
        }
        DuniterCommand::Config(ref config_args) => {
            duniter_ts_args.push("config".to_owned());
            config::gen_args(config_args, &mut duniter_ts_args);
        }
        DuniterCommand::Reset(ref reset_command) => {
            duniter_ts_args.push("reset".to_owned());
            match reset_command {
                ResetCommand::Config => duniter_ts_args.push("config".to_owned()),
                ResetCommand::Data => duniter_ts_args.push("data".to_owned()),
                ResetCommand::Peers => duniter_ts_args.push("peers".to_owned()),
                ResetCommand::Stats => duniter_ts_args.push("stats".to_owned()),
                ResetCommand::All => duniter_ts_args.push("all".to_owned()),
            }
        }
        DuniterCommand::Wizard(ref wizard_command) => {
            duniter_ts_args.push("wizard".to_owned());
            match wizard_command {
                WizardCommand::Bma => duniter_ts_args.push("network".to_owned()),
                WizardCommand::Key { n, r, p } => {
                    duniter_ts_args.push("key".to_owned());
                    if let Some(n) = n {
                        duniter_ts_args.push("--keyN".to_owned());
                        duniter_ts_args.push(n.to_string());
                    }
                    if let Some(r) = r {
                        duniter_ts_args.push("--keyr".to_owned());
                        duniter_ts_args.push(r.to_string());
                    }
                    if let Some(p) = p {
                        duniter_ts_args.push("--keyp".to_owned());
                        duniter_ts_args.push(p.to_string());
                    }
                }
            }
        }
        DuniterCommand::WS2P(ref ws2p_command) => {
            duniter_ts_args.push("ws2p".to_owned());
            match ws2p_command {
                WS2PCommand::ListNodes => duniter_ts_args.push("list-nodes".to_owned()),
                WS2PCommand::ListPrefered => duniter_ts_args.push("list-prefered".to_owned()),
                WS2PCommand::ListPrivileged => duniter_ts_args.push("list-privileged".to_owned()),
                WS2PCommand::ShowConf => duniter_ts_args.push("show-conf".to_owned()),
            }
        }
        DuniterCommand::Logs | DuniterCommand::Restart | DuniterCommand::Status => {}
    }
    duniter_ts_args
}
<|endoftext|>


<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! Base64 binary-to-text encoding

use core::iter;
use core::str;
use core::vec;

pub trait ToBase64 {
    fn to_base64(&self) -> ~str;
}

static CHARS: [char, ..64] = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/'
];

/**
 * Turn a vector of `u8` bytes into a string representing them in base64.
 *
 * *Example*:
 *
 * ~~~~
 * extern mod std;
 * use std::base64::ToBase64;
 *
 * fn main () {
 *   let str = [52,32].to_base64();
 *   println(fmt!("%s", str));
 * }
 * ~~~~
 */
impl<'self> ToBase64 for &'self [u8] {
    fn to_base64(&self) -> ~str {
        let mut s = ~"";
        unsafe {
            let len = self.len();
            str::reserve(&mut s, ((len + 3u) / 4u) * 3u);

            let mut i = 0u;

            while i < len - (len % 3u) {
                let n = (self[i] as uint) << 16u |
                        (self[i + 1u] as uint) << 8u |
                        (self[i + 2u] as uint);

                // This 24-bit number gets separated into four 6-bit numbers.
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, CHARS[n & 63u]);

                i += 3u;
            }

            // Heh, would be cool if we knew this was exhaustive
            // (the dream of bounded integer types)
            match len % 3 {
              0 => (),
              1 => {
                let n = (self[i] as uint) << 16u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, '=');
                str::push_char(&mut s, '=');
              }
              2 => {
                let n = (self[i] as uint) << 16u |
                    (self[i + 1u] as uint) << 8u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, '=');
              }
              _ => fail!(~"Algebra is broken, please alert the math police")
            }
        }
        s
    }
}

/**
 * Convert any string (literal, `@`, `&`, `~`) to base64 encoding.
 *
 *
 * *Example*:
 *
 * ~~~~
 * extern mod std;
 * use std::base64::ToBase64;
 *
 * fn main () {
 *   let str = "Hello, World".to_base64();
 *   println(fmt!("%s",str));
 * }
 * ~~~~
 *
 */
impl<'self> ToBase64 for &'self str {
    fn to_base64(&self) -> ~str {
        str::to_bytes(*self).to_base64()
    }
}

pub trait FromBase64 {
    fn from_base64(&self) -> ~[u8];
}

/**
 * Turn a vector of `u8`s representing characters
 * encoding byte values in base64 into the vector of `u8` byte values.
 *
 * *Example*:
 *
 * ~~~~
 * extern mod std;
 * use std::base64::ToBase64;
 * use std::base64::FromBase64;
 *
 * fn main () {
 *   let str = [52,32].to_base64();
 *   println(fmt!("%s", str));
 *   let bytes = str.from_base64();
 *   println(fmt!("%?",bytes));
 * }
 * ~~~~
 */
impl FromBase64 for ~[u8] {
    fn from_base64(&self) -> ~[u8] {
        if self.len() % 4u != 0u { fail!(~"invalid base64 length"); }

        let len = self.len();
        let mut padding = 0u;

        if len != 0u {
            if self[len - 1u] == '=' as u8 { padding += 1u; }
            if self[len - 2u] == '=' as u8 { padding += 1u; }
        }

        let mut r = vec::with_capacity((len / 4u) * 3u - padding);

        unsafe {
            let mut i = 0u;
            while i < len {
                let mut n = 0u;

                for iter::repeat(4u) {
                    let ch = self[i] as char;
                    n <<= 6u;

                    if ch >= 'A' && ch <= 'Z' {
                        n |= (ch as uint) - 0x41u;
                    } else if ch >= 'a' && ch <= 'z' {
                        n |= (ch as uint) - 0x47u;
                    } else if ch >= '0' && ch <= '9' {
                        n |= (ch as uint) + 0x04u;
                    } else if ch == '+' {
                        n |= 0x3Eu;
                    } else if ch == '/' {
                        n |= 0x3Fu;
                    } else if ch == '=' {
                        match len - i {
                          1u => {
                            r.push(((n >> 16u) & 0xFFu) as u8);
                            r.push(((n >> 8u ) & 0xFFu) as u8);
                            return copy r;
                          }
                          2u => {
                            r.push(((n >> 10u) & 0xFFu) as u8);
                            return copy r;
                          }
                          _ => fail!(~"invalid base64 padding")
                        }
                    } else {
                        fail!(~"invalid base64 character");
                    }

                    i += 1u;
                };

                r.push(((n >> 16u) & 0xFFu) as u8);
                r.push(((n >> 8u ) & 0xFFu) as u8);
                r.push(((n       ) & 0xFFu) as u8);
            }
        }
        r
    }
}

/**
 * Convert any string (literal, `@`, `&`, `~`)
 * that contains a base64 encoded value, to the byte values it encodes.
 *
 * You can use the `from_bytes` function in `core::str`
 * to turn a `[u8]` into a string with characters corresponding to those values.
 *
 * *Example*:
 *
 * This is an example of going from a string literal to the base64 encoding
 * and back to the same string.
 *
 * ~~~~
 * extern mod std;
 * use std::base64::ToBase64;
 * use std::base64::FromBase64;
 * use core::str;
 *
 * fn main () {
 *   let hello_str = "Hello, World".to_base64();
 *   println(fmt!("%s",hello_str));
 *   let bytes = hello_str.from_base64();
 *   println(fmt!("%?",bytes));
 *   let result_str = str::from_bytes(bytes);
 *   println(fmt!("%s",result_str));
 * }
 * ~~~~
 */
impl FromBase64 for ~str {
    fn from_base64(&self) -> ~[u8] {
        str::to_bytes(*self).from_base64()
    }
}

#[cfg(test)]
mod tests {
    use core::str;

    #[test]
    pub fn test_to_base64() {
        assert!((~"").to_base64()       == ~"");
        assert!((~"f").to_base64()      == ~"Zg==");
        assert!((~"fo").to_base64()     == ~"Zm8=");
        assert!((~"foo").to_base64()    == ~"Zm9v");
        assert!((~"foob").to_base64()   == ~"Zm9vYg==");
        assert!((~"fooba").to_base64()  == ~"Zm9vYmE=");
        assert!((~"foobar").to_base64() == ~"Zm9vYmFy");
    }

    #[test]
    pub fn test_from_base64() {
        assert!((~"").from_base64() == str::to_bytes(~""));
        assert!((~"Zg==").from_base64() == str::to_bytes(~"f"));
        assert!((~"Zm8=").from_base64() == str::to_bytes(~"fo"));
        assert!((~"Zm9v").from_base64() == str::to_bytes(~"foo"));
        assert!((~"Zm9vYg==").from_base64() == str::to_bytes(~"foob"));
        assert!((~"Zm9vYmE=").from_base64() == str::to_bytes(~"fooba"));
        assert!((~"Zm9vYmFy").from_base64() == str::to_bytes(~"foobar"));
    }
}
<commit_msg>move rustdoc comments so that they attach to the functions rather than the impl blocks.<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

//! Base64 binary-to-text encoding

use core::iter;
use core::str;
use core::vec;

pub trait ToBase64 {
    fn to_base64(&self) -> ~str;
}

static CHARS: [char, ..64] = [
    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
    '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/'
];

impl<'self> ToBase64 for &'self [u8] {
    /**
     * Turn a vector of `u8` bytes into a string representing them in base64.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *   let str = [52,32].to_base64();
     *   println(fmt!("%s", str));
     * }
     * ~~~~
     */
    fn to_base64(&self) -> ~str {
        let mut s = ~"";
        unsafe {
            let len = self.len();
            str::reserve(&mut s, ((len + 3u) / 4u) * 3u);

            let mut i = 0u;

            while i < len - (len % 3u) {
                let n = (self[i] as uint) << 16u |
                        (self[i + 1u] as uint) << 8u |
                        (self[i + 2u] as uint);

                // This 24-bit number gets separated into four 6-bit numbers.
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, CHARS[n & 63u]);

                i += 3u;
            }

            // Heh, would be cool if we knew this was exhaustive
            // (the dream of bounded integer types)
            match len % 3 {
              0 => (),
              1 => {
                let n = (self[i] as uint) << 16u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, '=');
                str::push_char(&mut s, '=');
              }
              2 => {
                let n = (self[i] as uint) << 16u |
                    (self[i + 1u] as uint) << 8u;
                str::push_char(&mut s, CHARS[(n >> 18u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 12u) & 63u]);
                str::push_char(&mut s, CHARS[(n >> 6u) & 63u]);
                str::push_char(&mut s, '=');
              }
              _ => fail!(~"Algebra is broken, please alert the math police")
            }
        }
        s
    }
}

impl<'self> ToBase64 for &'self str {
    /**
     * Convert any string (literal, `@`, `&`, `~`) to base64 encoding.
     *
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     *
     * fn main () {
     *   let str = "Hello, World".to_base64();
     *   println(fmt!("%s",str));
     * }
     * ~~~~
     *
     */
    fn to_base64(&self) -> ~str {
        str::to_bytes(*self).to_base64()
    }
}

pub trait FromBase64 {
    fn from_base64(&self) -> ~[u8];
}

impl FromBase64 for ~[u8] {
    /**
     * Turn a vector of `u8`s representing characters
     * encoding byte values in base64 into the vector of `u8` byte values.
     *
     * *Example*:
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     *
     * fn main () {
     *   let str = [52,32].to_base64();
     *   println(fmt!("%s", str));
     *   let bytes = str.from_base64();
     *   println(fmt!("%?",bytes));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        if self.len() % 4u != 0u { fail!(~"invalid base64 length"); }

        let len = self.len();
        let mut padding = 0u;

        if len != 0u {
            if self[len - 1u] == '=' as u8 { padding += 1u; }
            if self[len - 2u] == '=' as u8 { padding += 1u; }
        }

        let mut r = vec::with_capacity((len / 4u) * 3u - padding);

        unsafe {
            let mut i = 0u;
            while i < len {
                let mut n = 0u;

                for iter::repeat(4u) {
                    let ch = self[i] as char;
                    n <<= 6u;

                    if ch >= 'A' && ch <= 'Z' {
                        n |= (ch as uint) - 0x41u;
                    } else if ch >= 'a' && ch <= 'z' {
                        n |= (ch as uint) - 0x47u;
                    } else if ch >= '0' && ch <= '9' {
                        n |= (ch as uint) + 0x04u;
                    } else if ch == '+' {
                        n |= 0x3Eu;
                    } else if ch == '/' {
                        n |= 0x3Fu;
                    } else if ch == '=' {
                        match len - i {
                          1u => {
                            r.push(((n >> 16u) & 0xFFu) as u8);
                            r.push(((n >> 8u ) & 0xFFu) as u8);
                            return copy r;
                          }
                          2u => {
                            r.push(((n >> 10u) & 0xFFu) as u8);
                            return copy r;
                          }
                          _ => fail!(~"invalid base64 padding")
                        }
                    } else {
                        fail!(~"invalid base64 character");
                    }

                    i += 1u;
                };

                r.push(((n >> 16u) & 0xFFu) as u8);
                r.push(((n >> 8u ) & 0xFFu) as u8);
                r.push(((n       ) & 0xFFu) as u8);
            }
        }
        r
    }
}

impl FromBase64 for ~str {
    /**
     * Convert any string (literal, `@`, `&`, `~`)
     * that contains a base64 encoded value, to the byte values it encodes.
     *
     * You can use the `from_bytes` function in `core::str`
     * to turn a `[u8]` into a string with characters corresponding to those values.
     *
     * *Example*:
     *
     * This is an example of going from a string literal to the base64 encoding
     * and back to the same string.
     *
     * ~~~~
     * extern mod std;
     * use std::base64::ToBase64;
     * use std::base64::FromBase64;
     * use core::str;
     *
     * fn main () {
     *   let hello_str = "Hello, World".to_base64();
     *   println(fmt!("%s",hello_str));
     *   let bytes = hello_str.from_base64();
     *   println(fmt!("%?",bytes));
     *   let result_str = str::from_bytes(bytes);
     *   println(fmt!("%s",result_str));
     * }
     * ~~~~
     */
    fn from_base64(&self) -> ~[u8] {
        str::to_bytes(*self).from_base64()
    }
}

#[cfg(test)]
mod tests {
    use core::str;

    #[test]
    pub fn test_to_base64() {
        assert!((~"").to_base64()       == ~"");
        assert!((~"f").to_base64()      == ~"Zg==");
        assert!((~"fo").to_base64()     == ~"Zm8=");
        assert!((~"foo").to_base64()    == ~"Zm9v");
        assert!((~"foob").to_base64()   == ~"Zm9vYg==");
        assert!((~"fooba").to_base64()  == ~"Zm9vYmE=");
        assert!((~"foobar").to_base64() == ~"Zm9vYmFy");
    }

    #[test]
    pub fn test_from_base64() {
        assert!((~"").from_base64() == str::to_bytes(~""));
        assert!((~"Zg==").from_base64() == str::to_bytes(~"f"));
        assert!((~"Zm8=").from_base64() == str::to_bytes(~"fo"));
        assert!((~"Zm9v").from_base64() == str::to_bytes(~"foo"));
        assert!((~"Zm9vYg==").from_base64() == str::to_bytes(~"foob"));
        assert!((~"Zm9vYmE=").from_base64() == str::to_bytes(~"fooba"));
        assert!((~"Zm9vYmFy").from_base64() == str::to_bytes(~"foobar"));
    }
}
<|endoftext|>