<commit_before>#![warn(unused_variables)]

use std::cell::RefCell;
use std::collections::HashMap;
use std::error::Error;
use std::rc::Rc;

use errors::sovrin::SovrinError;
use errors::common::CommonError;
use services::agent::AgentService;
use services::pool::PoolService;
use services::signus::types::{MyDid, TheirDid};
use services::wallet::WalletService;
use utils::json::JsonDecodable;

pub enum AgentCommand {
    Connect(
        i32, // wallet handle
        String, // sender did
        String, // receiver did
        Box<Fn(Result<i32, SovrinError>) + Send>, // connect cb
        Box<Fn(Result<(i32, String), SovrinError>) + Send>, // message cb
    ),
    ConnectAck(
        i32, // cmd handle (eq conn handle)
        Result<i32, SovrinError> // conn handle or error
    ),
    CloseConnection,
    Listen(
        i32, // wallet handle
        String, // endpoint
        Box<Fn(Result<i32, SovrinError>) + Send>, // listen cb
        Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>, // connect cb
        Box<Fn(Result<(i32, String), SovrinError>) + Send>, // message cb
    ),
    ListenAck(
        i32, // cmd handle (eq listener handle)
        Result<i32, CommonError> // listener handle or error
    ),
    ListenerOnConnect(
        i32, // listener handle
        Result<(i32, i32, String, String), SovrinError>, // (listener handle, new connection handle, sender and receiver did) or error
    ),
    CloseListener,
    Send(
        i32, // connection handle
        Option<String>, // message
        Box<Fn(Result<(), SovrinError>) + Send>, // send cb
    ),
    SendAck(
        i32, // send cmd handle
        Result<(), CommonError>,
    )
}

pub struct AgentCommandExecutor {
    agent_service: Rc<AgentService>,
    pool_service: Rc<PoolService>,
    wallet_service: Rc<WalletService>,

    listeners: RefCell<HashMap<i32, Listener>>,

    listen_callbacks: RefCell<HashMap<i32, (
        Box<Fn(Result<i32, SovrinError>) + Send>, // listen cb
        Listener
    )>>,
    open_callbacks: RefCell<HashMap<i32, Box<Fn(Result<i32, SovrinError>)>>>,
    send_callbacks: RefCell<HashMap<i32, Box<Fn(Result<(), SovrinError>)>>>,
}

struct Listener {
    on_connect: Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>,
    on_msg: Box<Fn(Result<(i32, String), SovrinError>) + Send>,
}

impl AgentCommandExecutor {
    pub fn new(agent_service: Rc<AgentService>, pool_service: Rc<PoolService>, wallet_service: Rc<WalletService>) -> AgentCommandExecutor {
        AgentCommandExecutor {
            agent_service: agent_service,
            pool_service: pool_service,
            wallet_service: wallet_service,
            listeners: RefCell::new(HashMap::new()),
            listen_callbacks: RefCell::new(HashMap::new()),
            open_callbacks: RefCell::new(HashMap::new()),
            send_callbacks: RefCell::new(HashMap::new()),
        }
    }

    pub fn execute(&self, agent_cmd: AgentCommand) {
        match agent_cmd {
            AgentCommand::Connect(wallet_handle, sender_did, receiver_did, connect_cb, message_cb) => {
                info!(target: "agent_command_executor", "Connect command received");
                self.connect(wallet_handle, sender_did, receiver_did, connect_cb, message_cb)
            }
            AgentCommand::ConnectAck(cmd_id, res) => {
                info!(target: "agent_command_executor", "ConnectAck command received");
                self.open_callbacks.borrow_mut().remove(&cmd_id).unwrap()(res) //TODO extract method
            }
            AgentCommand::Listen(wallet_handle, endpoint, listen_cb, connect_cb, message_cb) => {
                info!(target: "agent_command_executor", "Listen command received");
                self.listen(wallet_handle, endpoint, listen_cb, connect_cb, message_cb);
            }
            AgentCommand::ListenAck(cmd_id, res) => {
                info!(target: "agent_command_executor", "ListenAck command received");
                let cbs = self.listen_callbacks.borrow_mut().remove(&cmd_id).unwrap();
                self.listeners.borrow_mut().insert(cmd_id, cbs.1);
                cbs.0(res.map_err(From::from)) //TODO extract method
            }
            AgentCommand::ListenerOnConnect(listener_id, res) => {
                info!(target: "agent_command_executor", "ListenerOnConnect command received");
                (self.listeners.borrow().get(&listener_id).unwrap().on_connect)(res);
            }
            AgentCommand::Send(connection_id, msg, cb) => {
                info!(target: "agent_command_executor", "Send command received");
                self.send(connection_id, msg, cb)
            }
            _ => unimplemented!(),
        }
    }

    fn connect(&self, wallet_handle: i32, sender_did: String, receiver_did: String,
               connect_cb: Box<Fn(Result<i32, SovrinError>) + Send>,
               /* FIXME message_cb */ _: Box<Fn(Result<(i32, String), SovrinError>) + Send>,
    ) {
        let result = self._get_connection_info(wallet_handle, &sender_did, &receiver_did)
            .and_then(|info: ConnectInfo| {
                debug!("AgentCommandExecutor::connect try to service.connect with {:?}", info);
                self.agent_service
                    .connect(sender_did.as_str(),
                             info.secret_key.as_str(), info.public_key.as_str(),
                             info.endpoint.as_str(), info.server_key.as_str())
                    .map_err(From::from)
            })
            .and_then(|conn_handle| {
                match self.open_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, conn_handle)),
                    Err(err) => Err(SovrinError::CommonError(CommonError::InvalidState(err.description().to_string()))),
                }
            });
        match result {
            Err(err) => { connect_cb(Err(err)); }
            Ok((mut cbs, handle)) => { cbs.insert(handle, connect_cb); /* TODO check if map contains same key */ }
        };
    }

    fn _get_connection_info(&self, wallet_handle: i32, sender_did: &String, receiver_did: &String)
                            -> Result<ConnectInfo, SovrinError> {
        let my_did_json = self.wallet_service.get(wallet_handle, &format!("my_did::{}", sender_did))?;
        let my_did: MyDid = MyDid::from_json(&my_did_json).map_err(|_| CommonError::InvalidState((format!("Invalid my did json"))))?;

        let their_did_json = self.wallet_service.get_not_expired(wallet_handle, &format!("their_did::{}", receiver_did))?; //FIXME implement: get from ledger
        let their_did: TheirDid = TheirDid::from_json(&their_did_json).map_err(|_| CommonError::InvalidState((format!("Invalid their did json"))))?;
        assert!(their_did.endpoint.is_some()); //FIXME implement: get from ledger
        Ok(ConnectInfo {
            secret_key: my_did.sk,
            public_key: my_did.pk,
            endpoint: their_did.endpoint.unwrap(),
            server_key: their_did.pk.unwrap()
        })
    }

    fn listen(&self, wallet_handle: i32, endpoint: String,
              listen_cb: Box<Fn(Result<i32, SovrinError>) + Send>,
              connect_cb: Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>,
              message_cb: Box<Fn(Result<(i32, String), SovrinError>) + Send>) {
        let my_did_json: String = self.wallet_service.list(wallet_handle, "my_did::").as_ref().unwrap().get(0).as_ref().unwrap().1.clone();
        let my_did: MyDid = MyDid::from_json(my_did_json.as_str()).map_err(|_| CommonError::InvalidState((format!("Invalid my did json")))).unwrap();

        let result = self.agent_service
            .listen(endpoint.as_str(), my_did.pk.as_str(), my_did.sk.as_str())
            .and_then(|cmd_id| {
                match self.listen_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, cmd_id)),
                    Err(err) => Err(CommonError::InvalidState(err.description().to_string())),
                }
            });
        match result {
            Err(err) => listen_cb(Err(From::from(err))),
            Ok((mut cbs, handle)) => {
                cbs.insert(handle, (listen_cb,
                                    Listener { on_connect: connect_cb, on_msg: message_cb })); /* TODO check if map contains same key */
            }
        };
    }

    fn send(&self, conn_id: i32, msg: Option<String>, cb: Box<Fn(Result<(), SovrinError>)>) {
        let result = self.agent_service
            .send(conn_id, msg.as_ref().map(String::as_str))
            .and_then(|cmd_id| {
                match self.send_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, cmd_id)),
                    Err(err) => Err(CommonError::InvalidState(err.description().to_string())),
                }
            });
        match result {
            Ok((mut cbs, cmd_id)) => { cbs.insert(cmd_id, cb); /* TODO check if map contains same key */ }
            Err(err) => cb(Err(From::from(err))),
        }
    }
}

#[derive(Debug)]
struct ConnectInfo {
    //TODO push to public service structure and use in service calls?
    secret_key: String,
    public_key: String,
    server_key: String,
    endpoint: String,
}
<commit_msg>Implement handling AgentCommand::SendAck.<commit_after>#![warn(unused_variables)]

use std::cell::RefCell;
use std::collections::HashMap;
use std::error::Error;
use std::rc::Rc;

use errors::sovrin::SovrinError;
use errors::common::CommonError;
use services::agent::AgentService;
use services::pool::PoolService;
use services::signus::types::{MyDid, TheirDid};
use services::wallet::WalletService;
use utils::json::JsonDecodable;

pub enum AgentCommand {
    Connect(
        i32, // wallet handle
        String, // sender did
        String, // receiver did
        Box<Fn(Result<i32, SovrinError>) + Send>, // connect cb
        Box<Fn(Result<(i32, String), SovrinError>) + Send>, // message cb
    ),
    ConnectAck(
        i32, // cmd handle (eq conn handle)
        Result<i32, SovrinError> // conn handle or error
    ),
    CloseConnection,
    Listen(
        i32, // wallet handle
        String, // endpoint
        Box<Fn(Result<i32, SovrinError>) + Send>, // listen cb
        Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>, // connect cb
        Box<Fn(Result<(i32, String), SovrinError>) + Send>, // message cb
    ),
    ListenAck(
        i32, // cmd handle (eq listener handle)
        Result<i32, CommonError> // listener handle or error
    ),
    ListenerOnConnect(
        i32, // listener handle
        Result<(i32, i32, String, String), SovrinError>, // (listener handle, new connection handle, sender and receiver did) or error
    ),
    CloseListener,
    Send(
        i32, // connection handle
        Option<String>, // message
        Box<Fn(Result<(), SovrinError>) + Send>, // send cb
    ),
    SendAck(
        i32, // send cmd handle
        Result<(), CommonError>,
    )
}

pub struct AgentCommandExecutor {
    agent_service: Rc<AgentService>,
    pool_service: Rc<PoolService>,
    wallet_service: Rc<WalletService>,

    listeners: RefCell<HashMap<i32, Listener>>,

    listen_callbacks: RefCell<HashMap<i32, (
        Box<Fn(Result<i32, SovrinError>) + Send>, // listen cb
        Listener
    )>>,
    open_callbacks: RefCell<HashMap<i32, Box<Fn(Result<i32, SovrinError>)>>>,
    send_callbacks: RefCell<HashMap<i32, Box<Fn(Result<(), SovrinError>)>>>,
}

struct Listener {
    on_connect: Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>,
    on_msg: Box<Fn(Result<(i32, String), SovrinError>) + Send>,
}

impl AgentCommandExecutor {
    pub fn new(agent_service: Rc<AgentService>, pool_service: Rc<PoolService>, wallet_service: Rc<WalletService>) -> AgentCommandExecutor {
        AgentCommandExecutor {
            agent_service: agent_service,
            pool_service: pool_service,
            wallet_service: wallet_service,
            listeners: RefCell::new(HashMap::new()),
            listen_callbacks: RefCell::new(HashMap::new()),
            open_callbacks: RefCell::new(HashMap::new()),
            send_callbacks: RefCell::new(HashMap::new()),
        }
    }

    pub fn execute(&self, agent_cmd: AgentCommand) {
        match agent_cmd {
            AgentCommand::Connect(wallet_handle, sender_did, receiver_did, connect_cb, message_cb) => {
                info!(target: "agent_command_executor", "Connect command received");
                self.connect(wallet_handle, sender_did, receiver_did, connect_cb, message_cb)
            }
            AgentCommand::ConnectAck(cmd_id, res) => {
                info!(target: "agent_command_executor", "ConnectAck command received");
                self.open_callbacks.borrow_mut().remove(&cmd_id).unwrap()(res) //TODO extract method
            }
            AgentCommand::Listen(wallet_handle, endpoint, listen_cb, connect_cb, message_cb) => {
                info!(target: "agent_command_executor", "Listen command received");
                self.listen(wallet_handle, endpoint, listen_cb, connect_cb, message_cb);
            }
            AgentCommand::ListenAck(cmd_id, res) => {
                info!(target: "agent_command_executor", "ListenAck command received");
                let cbs = self.listen_callbacks.borrow_mut().remove(&cmd_id).unwrap();
                self.listeners.borrow_mut().insert(cmd_id, cbs.1);
                cbs.0(res.map_err(From::from)) //TODO extract method
            }
            AgentCommand::ListenerOnConnect(listener_id, res) => {
                info!(target: "agent_command_executor", "ListenerOnConnect command received");
                (self.listeners.borrow().get(&listener_id).unwrap().on_connect)(res);
            }
            AgentCommand::Send(connection_id, msg, cb) => {
                info!(target: "agent_command_executor", "Send command received");
                self.send(connection_id, msg, cb)
            }
            AgentCommand::SendAck(cmd_id, res) => {
                info!(target: "agent_command_executor", "SendAck command received");
                self.send_callbacks.borrow_mut().remove(&cmd_id).unwrap()(res.map_err(From::from));
            }
            _ => unimplemented!(),
        }
    }

    fn connect(&self, wallet_handle: i32, sender_did: String, receiver_did: String,
               connect_cb: Box<Fn(Result<i32, SovrinError>) + Send>,
               /* FIXME message_cb */ _: Box<Fn(Result<(i32, String), SovrinError>) + Send>,
    ) {
        let result = self._get_connection_info(wallet_handle, &sender_did, &receiver_did)
            .and_then(|info: ConnectInfo| {
                debug!("AgentCommandExecutor::connect try to service.connect with {:?}", info);
                self.agent_service
                    .connect(sender_did.as_str(),
                             info.secret_key.as_str(), info.public_key.as_str(),
                             info.endpoint.as_str(), info.server_key.as_str())
                    .map_err(From::from)
            })
            .and_then(|conn_handle| {
                match self.open_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, conn_handle)),
                    Err(err) => Err(SovrinError::CommonError(CommonError::InvalidState(err.description().to_string()))),
                }
            });
        match result {
            Err(err) => { connect_cb(Err(err)); }
            Ok((mut cbs, handle)) => { cbs.insert(handle, connect_cb); /* TODO check if map contains same key */ }
        };
    }

    fn _get_connection_info(&self, wallet_handle: i32, sender_did: &String, receiver_did: &String)
                            -> Result<ConnectInfo, SovrinError> {
        let my_did_json = self.wallet_service.get(wallet_handle, &format!("my_did::{}", sender_did))?;
        let my_did: MyDid = MyDid::from_json(&my_did_json).map_err(|_| CommonError::InvalidState((format!("Invalid my did json"))))?;

        let their_did_json = self.wallet_service.get_not_expired(wallet_handle, &format!("their_did::{}", receiver_did))?; //FIXME implement: get from ledger
        let their_did: TheirDid = TheirDid::from_json(&their_did_json).map_err(|_| CommonError::InvalidState((format!("Invalid their did json"))))?;
        assert!(their_did.endpoint.is_some()); //FIXME implement: get from ledger
        Ok(ConnectInfo {
            secret_key: my_did.sk,
            public_key: my_did.pk,
            endpoint: their_did.endpoint.unwrap(),
            server_key: their_did.pk.unwrap()
        })
    }

    fn listen(&self, wallet_handle: i32, endpoint: String,
              listen_cb: Box<Fn(Result<i32, SovrinError>) + Send>,
              connect_cb: Box<Fn(Result<(i32, i32, String, String), SovrinError>) + Send>,
              message_cb: Box<Fn(Result<(i32, String), SovrinError>) + Send>) {
        let my_did_json: String = self.wallet_service.list(wallet_handle, "my_did::").as_ref().unwrap().get(0).as_ref().unwrap().1.clone();
        let my_did: MyDid = MyDid::from_json(my_did_json.as_str()).map_err(|_| CommonError::InvalidState((format!("Invalid my did json")))).unwrap();

        let result = self.agent_service
            .listen(endpoint.as_str(), my_did.pk.as_str(), my_did.sk.as_str())
            .and_then(|cmd_id| {
                match self.listen_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, cmd_id)),
                    Err(err) => Err(CommonError::InvalidState(err.description().to_string())),
                }
            });
        match result {
            Err(err) => listen_cb(Err(From::from(err))),
            Ok((mut cbs, handle)) => {
                cbs.insert(handle, (listen_cb,
                                    Listener { on_connect: connect_cb, on_msg: message_cb })); /* TODO check if map contains same key */
            }
        };
    }

    fn send(&self, conn_id: i32, msg: Option<String>, cb: Box<Fn(Result<(), SovrinError>)>) {
        let result = self.agent_service
            .send(conn_id, msg.as_ref().map(String::as_str))
            .and_then(|cmd_id| {
                match self.send_callbacks.try_borrow_mut() {
                    Ok(cbs) => Ok((cbs, cmd_id)),
                    Err(err) => Err(CommonError::InvalidState(err.description().to_string())),
                }
            });
        match result {
            Ok((mut cbs, cmd_id)) => { cbs.insert(cmd_id, cb); /* TODO check if map contains same key */ }
            Err(err) => cb(Err(From::from(err))),
        }
    }
}

#[derive(Debug)]
struct ConnectInfo {
    //TODO push to public service structure and use in service calls?
    secret_key: String,
    public_key: String,
    server_key: String,
    endpoint: String,
}
<|endoftext|>


<commit_before>use std;
import std::comm;
import std::task;

fn start(c : comm::_chan[int], n: int) {
    let i: int = n;

    while i > 0 { comm::send(c, 0); i = i - 1; }
}

fn main() {
    let p = comm::mk_port[comm::_chan[int]]();
    // Spawn a task that sends us back messages. The parent task
    // is likely to terminate before the child completes, so from
    // the child's point of view the receiver may die. We should
    // drop messages on the floor in this case, and not crash!
    let child = task::_spawn(bind start(p.mk_chan(), 10));
    let c = p.recv();
}<commit_msg>XFAIL task-comm-15<commit_after>// xfail-stage1
// xfail-stage2
// xfail-stage3

use std;
import std::comm;
import std::task;

fn start(c : comm::_chan[int], n: int) {
    let i: int = n;

    while i > 0 { comm::send(c, 0); i = i - 1; }
}

fn main() {
    let p = comm::mk_port[comm::_chan[int]]();
    // Spawn a task that sends us back messages. The parent task
    // is likely to terminate before the child completes, so from
    // the child's point of view the receiver may die. We should
    // drop messages on the floor in this case, and not crash!
    let child = task::_spawn(bind start(p.mk_chan(), 10));
    let c = p.recv();
}<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![allow(bare_trait_objects)]

use ::libra_types::proto::*;
use libra_mempool_shared_proto::proto::mempool_status;
use tokio::runtime::{Builder, Runtime};

pub mod admission_control {
    tonic::include_proto!("admission_control");
}

pub use self::admission_control::{
    admission_control_client::AdmissionControlClient, AdmissionControlMsg,
    SubmitTransactionRequest, SubmitTransactionResponse,
};

pub struct AdmissionControlClientBlocking {
    // Currently the runtime but be ordered before the tonic client to ensure that the runtime is
    // dropped last when this struct is dropped.
    // See https://github.com/tokio-rs/tokio/issues/1948 for more info.
    rt: Runtime,
    addr: String,
    client: Option<AdmissionControlClient<tonic::transport::Channel>>,
}

impl AdmissionControlClientBlocking {
    pub fn new<A: AsRef<str>>(address: A, port: u16) -> Self {
        let rt = Builder::new()
            .basic_scheduler()
            .enable_all()
            .build()
            .unwrap();
        let addr = format!("http://{}:{}", address.as_ref(), port);

        Self {
            client: None,
            addr,
            rt,
        }
    }

    fn client(
        &mut self,
    ) -> Result<
        (
            &mut Runtime,
            &mut AdmissionControlClient<tonic::transport::Channel>,
        ),
        tonic::Status,
    > {
        if self.client.is_none() {
            self.client = Some(
                self.rt
                    .block_on(AdmissionControlClient::connect(self.addr.clone()))
                    .map_err(|e| tonic::Status::new(tonic::Code::Unavailable, e.to_string()))?,
            );
        }

        // client is guaranteed to be populated by the time we reach here
        Ok((&mut self.rt, self.client.as_mut().unwrap()))
    }

    pub fn submit_transaction(
        &mut self,
        request: SubmitTransactionRequest,
    ) -> Result<SubmitTransactionResponse, tonic::Status> {
        let (rt, client) = self.client()?;
        rt.block_on(client.submit_transaction(request))
            .map(tonic::Response::into_inner)
    }

    pub fn update_to_latest_ledger(
        &mut self,
        request: types::UpdateToLatestLedgerRequest,
    ) -> Result<types::UpdateToLatestLedgerResponse, tonic::Status> {
        let (rt, client) = self.client()?;
        rt.block_on(async {
            tokio::time::timeout(
                std::time::Duration::from_millis(5000),
                client.update_to_latest_ledger(request),
            )
            .await
        })
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }
}

// Allow for lazily creating a Client
#[derive(Clone)]
pub struct AdmissionControlClientAsync {
    addr: String,
    client: Option<AdmissionControlClient<tonic::transport::Channel>>,
}

impl AdmissionControlClientAsync {
    pub fn new<A: AsRef<str>>(address: A, port: u16) -> Self {
        let addr = format!("http://{}:{}", address.as_ref(), port);

        Self { client: None, addr }
    }

    async fn client(
        &mut self,
    ) -> Result<&mut AdmissionControlClient<tonic::transport::Channel>, tonic::Status> {
        if self.client.is_none() {
            self.client = Some(
                AdmissionControlClient::connect(self.addr.clone())
                    .await
                    .map_err(|e| tonic::Status::new(tonic::Code::Unavailable, e.to_string()))?,
            );
        }

        // client is guaranteed to be populated by the time we reach here
        Ok(self.client.as_mut().unwrap())
    }

    pub async fn submit_transaction(
        &mut self,
        request: SubmitTransactionRequest,
    ) -> Result<SubmitTransactionResponse, tonic::Status> {
        let client = self.client().await?;
        tokio::time::timeout(
            std::time::Duration::from_millis(5000),
            client.submit_transaction(request),
        )
        .await
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }

    pub async fn update_to_latest_ledger(
        &mut self,
        request: types::UpdateToLatestLedgerRequest,
    ) -> Result<types::UpdateToLatestLedgerResponse, tonic::Status> {
        let client = self.client().await?;
        tokio::time::timeout(
            std::time::Duration::from_millis(5000),
            client.update_to_latest_ledger(request),
        )
        .await
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }
}
<commit_msg>client: add timeout to submit_transaction<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![allow(bare_trait_objects)]

use ::libra_types::proto::*;
use libra_mempool_shared_proto::proto::mempool_status;
use tokio::runtime::{Builder, Runtime};

pub mod admission_control {
    tonic::include_proto!("admission_control");
}

pub use self::admission_control::{
    admission_control_client::AdmissionControlClient, AdmissionControlMsg,
    SubmitTransactionRequest, SubmitTransactionResponse,
};

pub struct AdmissionControlClientBlocking {
    // Currently the runtime but be ordered before the tonic client to ensure that the runtime is
    // dropped last when this struct is dropped.
    // See https://github.com/tokio-rs/tokio/issues/1948 for more info.
    rt: Runtime,
    addr: String,
    client: Option<AdmissionControlClient<tonic::transport::Channel>>,
}

impl AdmissionControlClientBlocking {
    pub fn new<A: AsRef<str>>(address: A, port: u16) -> Self {
        let rt = Builder::new()
            .basic_scheduler()
            .enable_all()
            .build()
            .unwrap();
        let addr = format!("http://{}:{}", address.as_ref(), port);

        Self {
            client: None,
            addr,
            rt,
        }
    }

    fn client(
        &mut self,
    ) -> Result<
        (
            &mut Runtime,
            &mut AdmissionControlClient<tonic::transport::Channel>,
        ),
        tonic::Status,
    > {
        if self.client.is_none() {
            self.client = Some(
                self.rt
                    .block_on(AdmissionControlClient::connect(self.addr.clone()))
                    .map_err(|e| tonic::Status::new(tonic::Code::Unavailable, e.to_string()))?,
            );
        }

        // client is guaranteed to be populated by the time we reach here
        Ok((&mut self.rt, self.client.as_mut().unwrap()))
    }

    pub fn submit_transaction(
        &mut self,
        request: SubmitTransactionRequest,
    ) -> Result<SubmitTransactionResponse, tonic::Status> {
        let (rt, client) = self.client()?;
        rt.block_on(async {
            tokio::time::timeout(
                std::time::Duration::from_millis(5000),
                client.submit_transaction(request),
            )
            .await
        })
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }

    pub fn update_to_latest_ledger(
        &mut self,
        request: types::UpdateToLatestLedgerRequest,
    ) -> Result<types::UpdateToLatestLedgerResponse, tonic::Status> {
        let (rt, client) = self.client()?;
        rt.block_on(async {
            tokio::time::timeout(
                std::time::Duration::from_millis(5000),
                client.update_to_latest_ledger(request),
            )
            .await
        })
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }
}

// Allow for lazily creating a Client
#[derive(Clone)]
pub struct AdmissionControlClientAsync {
    addr: String,
    client: Option<AdmissionControlClient<tonic::transport::Channel>>,
}

impl AdmissionControlClientAsync {
    pub fn new<A: AsRef<str>>(address: A, port: u16) -> Self {
        let addr = format!("http://{}:{}", address.as_ref(), port);

        Self { client: None, addr }
    }

    async fn client(
        &mut self,
    ) -> Result<&mut AdmissionControlClient<tonic::transport::Channel>, tonic::Status> {
        if self.client.is_none() {
            self.client = Some(
                AdmissionControlClient::connect(self.addr.clone())
                    .await
                    .map_err(|e| tonic::Status::new(tonic::Code::Unavailable, e.to_string()))?,
            );
        }

        // client is guaranteed to be populated by the time we reach here
        Ok(self.client.as_mut().unwrap())
    }

    pub async fn submit_transaction(
        &mut self,
        request: SubmitTransactionRequest,
    ) -> Result<SubmitTransactionResponse, tonic::Status> {
        let client = self.client().await?;
        tokio::time::timeout(
            std::time::Duration::from_millis(5000),
            client.submit_transaction(request),
        )
        .await
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }

    pub async fn update_to_latest_ledger(
        &mut self,
        request: types::UpdateToLatestLedgerRequest,
    ) -> Result<types::UpdateToLatestLedgerResponse, tonic::Status> {
        let client = self.client().await?;
        tokio::time::timeout(
            std::time::Duration::from_millis(5000),
            client.update_to_latest_ledger(request),
        )
        .await
        .map_err(|_| tonic::Status::new(tonic::Code::DeadlineExceeded, ""))?
        .map(tonic::Response::into_inner)
    }
}
<|endoftext|>


<commit_before>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id`, `guild_id` and `webhook_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivalent in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivalent for the two
//! requests (`10`).
//!
//! Major parameters are why some variants (i.e. all of the channel/guild
//! variants) have an associated u64 as data. This is the Id of the parameter,
//! differentiating between different ratelimits.
//!
//! [Taken from]: https://discordapp.com/developers/docs/topics/rate-limits#rate-limits

pub use super::routing::Route;

use chrono::Utc;
use reqwest::{
    Client,
    Response,
    header::HeaderMap,
    StatusCode,
};
use crate::internal::prelude::*;
use parking_lot::{Mutex, RwLock};
use std::{
    collections::HashMap,
    sync::Arc,
    str::{
        self,
        FromStr,
    },
    time::Duration,
    thread,
    i64,
    u64,
};
use super::{HttpError, Request};
use log::debug;

/// Ratelimiter for requests to the Discord API.
///
/// This keeps track of ratelimit data for known routes through the
/// [`Ratelimit`] implementation for each route: how many tickets are
/// [`remaining`] until the user needs to wait for the known [`reset`] time, and
/// the [`limit`] of requests that can be made within that time.
///
/// When no tickets are available for some time, then the thread sleeps until
/// that time passes. The mechanism is known as "pre-emptive ratelimiting".
///
/// Occasionally for very high traffic bots, a global ratelimit may be reached
/// which blocks all future requests until the global ratelimit is over,
/// regardless of route. The value of this global ratelimit is never given
/// through the API, so it can't be pre-emptively ratelimited. This only affects
/// the largest of bots.
pub struct Ratelimiter {
    client: Arc<Client>,
    global: Arc<Mutex<()>>,
    // When futures is implemented, make tasks clear out their respective entry
    // when the 'reset' passes.
    routes: Arc<RwLock<HashMap<Route, Arc<Mutex<Ratelimit>>>>>,
    token: String,
}

impl Ratelimiter {
    /// Creates a new ratelimiter, with a shared `reqwest` client and the
    /// bot's token.
    ///
    /// The bot token must be prefixed with `"Bot "`. The ratelimiter does not
    /// prefix it.
    pub fn new(client: Arc<Client>, token: impl Into<String>) -> Self {
        Self::_new(client, token.into())
    }

    fn _new(client: Arc<Client>, token: String) -> Self {
        Self {
            client,
            global: Default::default(),
            routes: Default::default(),
            token,
        }
    }

    /// The routes mutex is a HashMap of each [`Route`] and their respective
    /// ratelimit information.
    ///
    /// See the documentation for [`Ratelimit`] for more information on how the
    /// library handles ratelimiting.
    ///
    /// # Examples
    ///
    /// View the `reset` time of the route for `ChannelsId(7)`:
    ///
    /// ```rust,no_run
    /// use serenity::http::{ratelimiting::{Route}};
    /// # use serenity::http::Http;
    /// # let http = Http::default();
    /// let routes = http.ratelimiter.routes();
    /// let reader = routes.read();
    ///
    /// if let Some(route) = reader.get(&Route::ChannelsId(7)) {
    ///     println!("Reset time at: {}", route.lock().reset());
    /// }
    /// ```
    ///
    /// [`Ratelimit`]: struct.Ratelimit.html
    /// [`Route`]: ../routing/enum.Route.html
    pub fn routes(&self) -> Arc<RwLock<HashMap<Route, Arc<Mutex<Ratelimit>>>>> {
        Arc::clone(&self.routes)
    }

    pub fn perform(&self, req: RatelimitedRequest<'_>) -> Result<Response> {
        let RatelimitedRequest { req } = req;

        loop {
            // This will block if another thread hit the global ratelimit.
            let _ = self.global.lock();

            // Destructure the tuple instead of retrieving the third value to
            // take advantage of the type system. If `RouteInfo::deconstruct`
            // returns a different number of tuple elements in the future,
            // directly accessing a certain index
            // (e.g. `req.route.deconstruct().1`) would mean this code would not
            // indicate it might need to be updated for the new tuple element
            // amount.
            //
            // This isn't normally important, but might be for ratelimiting.
            let (_, route, _) = req.route.deconstruct();

            // Perform pre-checking here:
            //
            // - get the route's relevant rate
            // - sleep if that route's already rate-limited until the end of the
            //   'reset' time;
            // - get the global rate;
            // - sleep if there is 0 remaining
            // - then, perform the request
            let bucket = Arc::clone(&self.routes
                .write()
                .entry(route)
                .or_default());

            bucket.lock().pre_hook(&route);

            let request = req.build(&self.client, &self.token)?;
            let response = request.send()?;

            // Check if the request got ratelimited by checking for status 429,
            // and if so, sleep for the value of the header 'retry-after' -
            // which is in milliseconds - and then `continue` to try again
            //
            // If it didn't ratelimit, subtract one from the Ratelimit's
            // 'remaining'
            //
            // Update `reset` with the value of 'x-ratelimit-reset' header.
            // Similarly, update `reset-after` with the 'x-ratelimit-reset-after' header.
            //
            // It _may_ be possible for the limit to be raised at any time,
            // so check if it did from the value of the 'x-ratelimit-limit'
            // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
            if route == Route::None {
                return Ok(response);
            } else {
                let redo = if response.headers().get("x-ratelimit-global").is_some() {
                    let _ = self.global.lock();

                    Ok(
                        if let Some(retry_after) = parse_header::<u64>(&response.headers(), "retry-after")? {
                            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
                            thread::sleep(Duration::from_millis(retry_after));

                            true
                        } else {
                            false
                        },
                    )
                } else {
                    bucket.lock().post_hook(&response, &route)
                };

                if !redo.unwrap_or(true) {
                    return Ok(response);
                }
            }
        }
    }
}

/// A set of data containing information about the ratelimits for a particular
/// [`Route`], which is stored in [`Http`].
///
/// See the [Discord docs] on ratelimits for more information.
///
/// **Note**: You should _not_ mutate any of the fields, as this can help cause
/// 429s.
///
/// [`Http`]: ../client/struct.Http.html#structfield.routes
/// [`Route`]: ../routing/enum.Route.html
/// [Discord docs]: https://discordapp.com/developers/docs/topics/rate-limits
#[derive(Debug)]
pub struct Ratelimit {
    /// The total number of requests that can be made in a period of time.
    limit: i64,
    /// The number of requests remaining in the period of time.
    remaining: i64,
    /// The absolute time in milliseconds when the interval resets.
    reset: i64,
    /// The total time in milliseconds when the interval resets.
    reset_after: i64,
}

impl Ratelimit {
    #[cfg(feature = "absolute_ratelimits")]
    fn get_delay(&self) -> i64 {
        let now = Utc::now().timestamp_millis();
        self.reset - now
    }

    #[cfg(not(feature = "absolute_ratelimits"))]
    fn get_delay(&self) -> i64 {
        self.reset_after
    }

    pub fn pre_hook(&mut self, route: &Route) {
        if self.limit() == 0 {
            return;
        }

		let delay = self.get_delay();

		if delay < 0 {
			// We're probably in the past.
			self.remaining = self.limit;

			return;
		}

        if self.remaining() == 0 {
            let delay = delay as u64;

            debug!(
                "Pre-emptive ratelimit on route {:?} for {:?}ms",
                route,
                delay,
            );

            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub fn post_hook(&mut self, response: &Response, route: &Route) -> Result<bool> {
        if let Some(limit) = parse_header(&response.headers(), "x-ratelimit-limit")? {
            self.limit = limit;
        }

        if let Some(remaining) = parse_header(&response.headers(), "x-ratelimit-remaining")? {
            self.remaining = remaining;
        }

        if let Some(reset) = parse_header::<f64>(&response.headers(), "x-ratelimit-reset")? {
            self.reset = (reset * 1000f64) as i64;
        }

        if let Some(reset_after) = parse_header::<f64>(&response.headers(), "x-ratelimit-reset-after")? {
            self.reset_after = (reset_after * 1000f64) as i64;
        }

        Ok(if response.status() != StatusCode::TOO_MANY_REQUESTS {
            false
        } else if let Some(retry_after) = parse_header::<u64>(&response.headers(), "retry-after")? {
            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
            thread::sleep(Duration::from_millis(retry_after));

            true
        } else {
            false
        })
    }

    /// The total number of requests that can be made in a period of time.
    #[inline]
    pub fn limit(&self) -> i64 {
        self.limit
    }

    /// The number of requests remaining in the period of time.
    #[inline]
    pub fn remaining(&self) -> i64 {
        self.remaining
    }

    /// The absolute time in milliseconds when the interval resets.
    #[inline]
    pub fn reset(&self) -> i64 {
        self.reset
    }

    /// The total time in milliseconds when the interval resets.
    #[inline]
    pub fn reset_after(&self) -> i64 {
        self.reset_after
    }
}

impl Default for Ratelimit {
    fn default() -> Self {
        Self {
            limit: i64::MAX,
            remaining: i64::MAX,
            reset: i64::MAX,
            reset_after: i64::MAX,
        }
    }
}

/// Information about a request for the ratelimiter to perform.
///
/// This only contains the basic information needed by the ratelimiter to
/// perform a full cycle of making the request and returning the response.
///
/// Use the `From` implementations for making one of these.
pub struct RatelimitedRequest<'a> {
    req: Request<'a>,
}

impl<'a> From<Request<'a>> for RatelimitedRequest<'a> {
    fn from(req: Request<'a>) -> Self {
        Self { req }
    }
}

fn parse_header<T: FromStr>(headers: &HeaderMap, header: &str) -> Result<Option<T>> {
    let header = match headers.get(header) {
        Some(v) => v,
        None => return Ok(None),
    };

    let unicode = str::from_utf8(&header.as_bytes()).map_err(|_| {
        Error::from(HttpError::RateLimitUtf8)
    })?;

    let num = unicode.parse().map_err(|_| {
        Error::from(HttpError::RateLimitI64F64)
    })?;

    Ok(Some(num))
}

#[cfg(test)]
mod tests {
    use crate::{
        error::Error,
        http::HttpError,
    };
    use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
    use std::{
        error::Error as StdError,
        result::Result as StdResult,
    };
    use super::parse_header;

    type Result<T> = StdResult<T, Box<dyn StdError>>;

    fn headers() -> HeaderMap {
        let pairs = &[
            (
                HeaderName::from_static("x-ratelimit-limit"),
                HeaderValue::from_static("5"),
            ),
            (
                HeaderName::from_static("x-ratelimit-remaining"),
                HeaderValue::from_static("4"),
            ),
            (
                HeaderName::from_static("x-ratelimit-reset"),
                HeaderValue::from_static("1560704880.423"),
            ),
            (
                HeaderName::from_static("x-bad-num"),
                HeaderValue::from_static("abc"),
            ),
            (
                HeaderName::from_static("x-bad-unicode"),
                HeaderValue::from_bytes(&[255, 255, 255, 255]).unwrap(),
            ),
        ];

        let mut map = HeaderMap::with_capacity(pairs.len());

        for (name, val) in pairs.into_iter() {
            map.insert(name, val.to_owned());
        }

        map
    }

    #[test]
    fn test_parse_header_good() -> Result<()> {
        let headers = headers();

        assert_eq!(parse_header::<i64>(&headers, "x-ratelimit-limit")?.unwrap(), 5);
        assert_eq!(
            parse_header::<i64>(&headers, "x-ratelimit-remaining")?.unwrap(),
            4,
        );
        assert_eq!(
            parse_header::<f64>(&headers, "x-ratelimit-reset")?.unwrap(),
            1_560_704_880.423,
        );

        Ok(())
    }

    #[test]
    fn test_parse_header_errors() -> Result<()> {
        let headers = headers();

        match parse_header::<i64>(&headers, "x-bad-num").unwrap_err() {
            Error::Http(x) => match *x {
                HttpError::RateLimitI64F64 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        match parse_header::<i64>(&headers, "x-bad-unicode").unwrap_err() {
            Error::Http(http_err) => match *http_err {
                HttpError::RateLimitUtf8 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        Ok(())
    }
}
<commit_msg>Remove unused import<commit_after>//! Routes are used for ratelimiting. These are to differentiate between the
//! different _types_ of routes - such as getting the current user's channels -
//! for the most part, with the exception being major parameters.
//!
//! [Taken from] the Discord docs, major parameters are:
//!
//! > Additionally, rate limits take into account major parameters in the URL.
//! > For example, `/channels/:channel_id` and
//! > `/channels/:channel_id/messages/:message_id` both take `channel_id` into
//! > account when generating rate limits since it's the major parameter. The
//! only current major parameters are `channel_id`, `guild_id` and `webhook_id`.
//!
//! This results in the two URIs of `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` being rate limited _separately_. However, the
//! two URIs of `GET /channels/10/messages/11` and
//! `GET /channels/10/messages/12` will count towards the "same ratelimit", as
//! the major parameter - `10` is equivalent in both URIs' format.
//!
//! # Examples
//!
//! First: taking the first two URIs - `GET /channels/4/messages/7` and
//! `GET /channels/5/messages/8` - and assuming both buckets have a `limit` of
//! `10`, requesting the first URI will result in the response containing a
//! `remaining` of `9`. Immediately after - prior to buckets resetting -
//! performing a request to the _second_ URI will also contain a `remaining` of
//! `9` in the response, as the major parameter - `channel_id` - is different
//! in the two requests (`4` and `5`).
//!
//! Second: take for example the last two URIs. Assuming the bucket's `limit` is
//! `10`, requesting the first URI will return a `remaining` of `9` in the
//! response. Immediately after - prior to buckets resetting - performing a
//! request to the _second_ URI will return a `remaining` of `8` in the
//! response, as the major parameter - `channel_id` - is equivalent for the two
//! requests (`10`).
//!
//! Major parameters are why some variants (i.e. all of the channel/guild
//! variants) have an associated u64 as data. This is the Id of the parameter,
//! differentiating between different ratelimits.
//!
//! [Taken from]: https://discordapp.com/developers/docs/topics/rate-limits#rate-limits

pub use super::routing::Route;

use reqwest::{
    Client,
    Response,
    header::HeaderMap,
    StatusCode,
};
use crate::internal::prelude::*;
use parking_lot::{Mutex, RwLock};
use std::{
    collections::HashMap,
    sync::Arc,
    str::{
        self,
        FromStr,
    },
    time::Duration,
    thread,
    i64,
    u64,
};
use super::{HttpError, Request};
use log::debug;

/// Ratelimiter for requests to the Discord API.
///
/// This keeps track of ratelimit data for known routes through the
/// [`Ratelimit`] implementation for each route: how many tickets are
/// [`remaining`] until the user needs to wait for the known [`reset`] time, and
/// the [`limit`] of requests that can be made within that time.
///
/// When no tickets are available for some time, then the thread sleeps until
/// that time passes. The mechanism is known as "pre-emptive ratelimiting".
///
/// Occasionally for very high traffic bots, a global ratelimit may be reached
/// which blocks all future requests until the global ratelimit is over,
/// regardless of route. The value of this global ratelimit is never given
/// through the API, so it can't be pre-emptively ratelimited. This only affects
/// the largest of bots.
pub struct Ratelimiter {
    client: Arc<Client>,
    global: Arc<Mutex<()>>,
    // When futures is implemented, make tasks clear out their respective entry
    // when the 'reset' passes.
    routes: Arc<RwLock<HashMap<Route, Arc<Mutex<Ratelimit>>>>>,
    token: String,
}

impl Ratelimiter {
    /// Creates a new ratelimiter, with a shared `reqwest` client and the
    /// bot's token.
    ///
    /// The bot token must be prefixed with `"Bot "`. The ratelimiter does not
    /// prefix it.
    pub fn new(client: Arc<Client>, token: impl Into<String>) -> Self {
        Self::_new(client, token.into())
    }

    fn _new(client: Arc<Client>, token: String) -> Self {
        Self {
            client,
            global: Default::default(),
            routes: Default::default(),
            token,
        }
    }

    /// The routes mutex is a HashMap of each [`Route`] and their respective
    /// ratelimit information.
    ///
    /// See the documentation for [`Ratelimit`] for more information on how the
    /// library handles ratelimiting.
    ///
    /// # Examples
    ///
    /// View the `reset` time of the route for `ChannelsId(7)`:
    ///
    /// ```rust,no_run
    /// use serenity::http::{ratelimiting::{Route}};
    /// # use serenity::http::Http;
    /// # let http = Http::default();
    /// let routes = http.ratelimiter.routes();
    /// let reader = routes.read();
    ///
    /// if let Some(route) = reader.get(&Route::ChannelsId(7)) {
    ///     println!("Reset time at: {}", route.lock().reset());
    /// }
    /// ```
    ///
    /// [`Ratelimit`]: struct.Ratelimit.html
    /// [`Route`]: ../routing/enum.Route.html
    pub fn routes(&self) -> Arc<RwLock<HashMap<Route, Arc<Mutex<Ratelimit>>>>> {
        Arc::clone(&self.routes)
    }

    pub fn perform(&self, req: RatelimitedRequest<'_>) -> Result<Response> {
        let RatelimitedRequest { req } = req;

        loop {
            // This will block if another thread hit the global ratelimit.
            let _ = self.global.lock();

            // Destructure the tuple instead of retrieving the third value to
            // take advantage of the type system. If `RouteInfo::deconstruct`
            // returns a different number of tuple elements in the future,
            // directly accessing a certain index
            // (e.g. `req.route.deconstruct().1`) would mean this code would not
            // indicate it might need to be updated for the new tuple element
            // amount.
            //
            // This isn't normally important, but might be for ratelimiting.
            let (_, route, _) = req.route.deconstruct();

            // Perform pre-checking here:
            //
            // - get the route's relevant rate
            // - sleep if that route's already rate-limited until the end of the
            //   'reset' time;
            // - get the global rate;
            // - sleep if there is 0 remaining
            // - then, perform the request
            let bucket = Arc::clone(&self.routes
                .write()
                .entry(route)
                .or_default());

            bucket.lock().pre_hook(&route);

            let request = req.build(&self.client, &self.token)?;
            let response = request.send()?;

            // Check if the request got ratelimited by checking for status 429,
            // and if so, sleep for the value of the header 'retry-after' -
            // which is in milliseconds - and then `continue` to try again
            //
            // If it didn't ratelimit, subtract one from the Ratelimit's
            // 'remaining'
            //
            // Update `reset` with the value of 'x-ratelimit-reset' header.
            // Similarly, update `reset-after` with the 'x-ratelimit-reset-after' header.
            //
            // It _may_ be possible for the limit to be raised at any time,
            // so check if it did from the value of the 'x-ratelimit-limit'
            // header. If the limit was 5 and is now 7, add 2 to the 'remaining'
            if route == Route::None {
                return Ok(response);
            } else {
                let redo = if response.headers().get("x-ratelimit-global").is_some() {
                    let _ = self.global.lock();

                    Ok(
                        if let Some(retry_after) = parse_header::<u64>(&response.headers(), "retry-after")? {
                            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
                            thread::sleep(Duration::from_millis(retry_after));

                            true
                        } else {
                            false
                        },
                    )
                } else {
                    bucket.lock().post_hook(&response, &route)
                };

                if !redo.unwrap_or(true) {
                    return Ok(response);
                }
            }
        }
    }
}

/// A set of data containing information about the ratelimits for a particular
/// [`Route`], which is stored in [`Http`].
///
/// See the [Discord docs] on ratelimits for more information.
///
/// **Note**: You should _not_ mutate any of the fields, as this can help cause
/// 429s.
///
/// [`Http`]: ../client/struct.Http.html#structfield.routes
/// [`Route`]: ../routing/enum.Route.html
/// [Discord docs]: https://discordapp.com/developers/docs/topics/rate-limits
#[derive(Debug)]
pub struct Ratelimit {
    /// The total number of requests that can be made in a period of time.
    limit: i64,
    /// The number of requests remaining in the period of time.
    remaining: i64,
    /// The absolute time in milliseconds when the interval resets.
    reset: i64,
    /// The total time in milliseconds when the interval resets.
    reset_after: i64,
}

impl Ratelimit {
    #[cfg(feature = "absolute_ratelimits")]
    fn get_delay(&self) -> i64 {
        let now = Utc::now().timestamp_millis();
        self.reset - now
    }

    #[cfg(not(feature = "absolute_ratelimits"))]
    fn get_delay(&self) -> i64 {
        self.reset_after
    }

    pub fn pre_hook(&mut self, route: &Route) {
        if self.limit() == 0 {
            return;
        }

		let delay = self.get_delay();

		if delay < 0 {
			// We're probably in the past.
			self.remaining = self.limit;

			return;
		}

        if self.remaining() == 0 {
            let delay = delay as u64;

            debug!(
                "Pre-emptive ratelimit on route {:?} for {:?}ms",
                route,
                delay,
            );

            thread::sleep(Duration::from_millis(delay));

            return;
        }

        self.remaining -= 1;
    }

    pub fn post_hook(&mut self, response: &Response, route: &Route) -> Result<bool> {
        if let Some(limit) = parse_header(&response.headers(), "x-ratelimit-limit")? {
            self.limit = limit;
        }

        if let Some(remaining) = parse_header(&response.headers(), "x-ratelimit-remaining")? {
            self.remaining = remaining;
        }

        if let Some(reset) = parse_header::<f64>(&response.headers(), "x-ratelimit-reset")? {
            self.reset = (reset * 1000f64) as i64;
        }

        if let Some(reset_after) = parse_header::<f64>(&response.headers(), "x-ratelimit-reset-after")? {
            self.reset_after = (reset_after * 1000f64) as i64;
        }

        Ok(if response.status() != StatusCode::TOO_MANY_REQUESTS {
            false
        } else if let Some(retry_after) = parse_header::<u64>(&response.headers(), "retry-after")? {
            debug!("Ratelimited on route {:?} for {:?}ms", route, retry_after);
            thread::sleep(Duration::from_millis(retry_after));

            true
        } else {
            false
        })
    }

    /// The total number of requests that can be made in a period of time.
    #[inline]
    pub fn limit(&self) -> i64 {
        self.limit
    }

    /// The number of requests remaining in the period of time.
    #[inline]
    pub fn remaining(&self) -> i64 {
        self.remaining
    }

    /// The absolute time in milliseconds when the interval resets.
    #[inline]
    pub fn reset(&self) -> i64 {
        self.reset
    }

    /// The total time in milliseconds when the interval resets.
    #[inline]
    pub fn reset_after(&self) -> i64 {
        self.reset_after
    }
}

impl Default for Ratelimit {
    fn default() -> Self {
        Self {
            limit: i64::MAX,
            remaining: i64::MAX,
            reset: i64::MAX,
            reset_after: i64::MAX,
        }
    }
}

/// Information about a request for the ratelimiter to perform.
///
/// This only contains the basic information needed by the ratelimiter to
/// perform a full cycle of making the request and returning the response.
///
/// Use the `From` implementations for making one of these.
pub struct RatelimitedRequest<'a> {
    req: Request<'a>,
}

impl<'a> From<Request<'a>> for RatelimitedRequest<'a> {
    fn from(req: Request<'a>) -> Self {
        Self { req }
    }
}

fn parse_header<T: FromStr>(headers: &HeaderMap, header: &str) -> Result<Option<T>> {
    let header = match headers.get(header) {
        Some(v) => v,
        None => return Ok(None),
    };

    let unicode = str::from_utf8(&header.as_bytes()).map_err(|_| {
        Error::from(HttpError::RateLimitUtf8)
    })?;

    let num = unicode.parse().map_err(|_| {
        Error::from(HttpError::RateLimitI64F64)
    })?;

    Ok(Some(num))
}

#[cfg(test)]
mod tests {
    use crate::{
        error::Error,
        http::HttpError,
    };
    use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
    use std::{
        error::Error as StdError,
        result::Result as StdResult,
    };
    use super::parse_header;

    type Result<T> = StdResult<T, Box<dyn StdError>>;

    fn headers() -> HeaderMap {
        let pairs = &[
            (
                HeaderName::from_static("x-ratelimit-limit"),
                HeaderValue::from_static("5"),
            ),
            (
                HeaderName::from_static("x-ratelimit-remaining"),
                HeaderValue::from_static("4"),
            ),
            (
                HeaderName::from_static("x-ratelimit-reset"),
                HeaderValue::from_static("1560704880.423"),
            ),
            (
                HeaderName::from_static("x-bad-num"),
                HeaderValue::from_static("abc"),
            ),
            (
                HeaderName::from_static("x-bad-unicode"),
                HeaderValue::from_bytes(&[255, 255, 255, 255]).unwrap(),
            ),
        ];

        let mut map = HeaderMap::with_capacity(pairs.len());

        for (name, val) in pairs.into_iter() {
            map.insert(name, val.to_owned());
        }

        map
    }

    #[test]
    fn test_parse_header_good() -> Result<()> {
        let headers = headers();

        assert_eq!(parse_header::<i64>(&headers, "x-ratelimit-limit")?.unwrap(), 5);
        assert_eq!(
            parse_header::<i64>(&headers, "x-ratelimit-remaining")?.unwrap(),
            4,
        );
        assert_eq!(
            parse_header::<f64>(&headers, "x-ratelimit-reset")?.unwrap(),
            1_560_704_880.423,
        );

        Ok(())
    }

    #[test]
    fn test_parse_header_errors() -> Result<()> {
        let headers = headers();

        match parse_header::<i64>(&headers, "x-bad-num").unwrap_err() {
            Error::Http(x) => match *x {
                HttpError::RateLimitI64F64 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        match parse_header::<i64>(&headers, "x-bad-unicode").unwrap_err() {
            Error::Http(http_err) => match *http_err {
                HttpError::RateLimitUtf8 => assert!(true),
                _ => assert!(false),
            },
            _ => assert!(false),
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Encode `InlineContent` nodes to HTML

use super::{
    attr, attr_id, attr_itemprop, attr_itemtype, attr_itemtype_str, concat, elem, elem_empty, json,
    EncodeContext, ToHtml,
};
use codec_txt::ToTxt;
use html_escape::encode_safe;
use std::{fs, path::PathBuf};
use stencila_schema::*;

impl ToHtml for InlineContent {
    fn to_html(&self, context: &EncodeContext) -> String {
        match self {
            InlineContent::AudioObject(node) => node.to_html(context),
            InlineContent::Boolean(node) => node.to_html(context),
            InlineContent::Cite(node) => node.to_html(context),
            InlineContent::CiteGroup(node) => node.to_html(context),
            InlineContent::CodeExpression(node) => node.to_html(context),
            InlineContent::CodeFragment(node) => node.to_html(context),
            InlineContent::Delete(node) => node.to_html(context),
            InlineContent::Emphasis(node) => node.to_html(context),
            InlineContent::ImageObject(node) => node.to_html(context),
            InlineContent::Integer(node) => node.to_html(context),
            InlineContent::Link(node) => node.to_html(context),
            InlineContent::MathFragment(node) => node.to_html(context),
            InlineContent::NontextualAnnotation(node) => node.to_html(context),
            InlineContent::Note(node) => node.to_html(context),
            InlineContent::Null(node) => node.to_html(context),
            InlineContent::Number(node) => node.to_html(context),
            InlineContent::Parameter(node) => node.to_html(context),
            InlineContent::Quote(node) => node.to_html(context),
            // Unlike everything else, `String` instances are not wrapped in an element.
            // However, `InlineContent::String` (and `Node::String` instances) should be
            // because they are usually part of a vector (e.g. `Paragraph.content`).
            // Wrapping them indicates to the DOM patching algorithm that the container element
            // represents a vector and not a placeholder for an optional string property.
            // To reduce the size of the generated HTML a little, unlike for `Node::String`,
            // the `itemtype` attribute is not used.
            InlineContent::String(node) => elem("span", &[], &node.to_html(context)),
            InlineContent::Strong(node) => node.to_html(context),
            InlineContent::Subscript(node) => node.to_html(context),
            InlineContent::Superscript(node) => node.to_html(context),
            InlineContent::VideoObject(node) => node.to_html(context),
        }
    }
}

macro_rules! mark_to_html {
    ($type:ident, $tag:literal) => {
        impl ToHtml for $type {
            fn to_html(&self, context: &EncodeContext) -> String {
                elem(
                    $tag,
                    &[attr_itemtype::<Self>(), attr_id(&self.id)],
                    &self.content.to_html(context),
                )
            }
        }
    };
}

mark_to_html!(Delete, "del");
mark_to_html!(Emphasis, "em");
mark_to_html!(NontextualAnnotation, "u");
mark_to_html!(Strong, "strong");
mark_to_html!(Subscript, "sub");
mark_to_html!(Superscript, "sup");

/// Convert a file:// URL to a data:// URI
///
/// Note that this function assumes that paths are absolute.
/// File URLs are usually resolves elsewhere e.g. in the `compile` method
/// before encoding to HTML.
fn file_uri_to_data_uri(url: &str) -> String {
    let path = if let Some(path) = url.strip_prefix("file://") {
        PathBuf::from(path)
    } else {
        return url.into();
    };

    // Read the file, convert it to a dataURI, and record it as a dependency
    match fs::read(&path) {
        Ok(bytes) => {
            let mime = match mime_guess::from_path(&path).first() {
                Some(mime) => mime.to_string(),
                None => "image/png".to_string(),
            };
            let data = base64::encode(bytes);
            ["data:", &mime, ";base64,", &data].concat()
        }
        Err(error) => {
            tracing::warn!("Unable to read media file {}: {}", path.display(), error);
            url.into()
        }
    }
}

fn content_url_to_src_attr(content_url: &str, context: &EncodeContext) -> String {
    let url = match context.bundle {
        true => file_uri_to_data_uri(content_url),
        false => content_url.to_string(),
    };
    attr("src", &url)
}

impl ToHtml for AudioObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "audio",
            &[
                attr_itemtype_str("AudioObject"),
                attr_id(&self.id),
                "controls".to_string(),
                content_url_to_src_attr(&self.content_url, context),
            ],
            "",
        )
    }
}

impl ToHtml for ImageObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem_empty(
            "img",
            &[
                attr_itemtype_str("ImageObject"),
                attr_id(&self.id),
                content_url_to_src_attr(&self.content_url, context),
            ],
        )
    }
}

impl ToHtml for VideoObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        let src_attr = content_url_to_src_attr(&self.content_url, context);
        let type_attr = match &self.media_type {
            Some(media_type) => attr("type", media_type),
            None => "".to_string(),
        };
        elem(
            "video",
            &[
                attr_itemtype_str("VideoObject"),
                attr_id(&self.id),
                "controls".to_string(),
            ],
            &elem("source", &[src_attr, type_attr], ""),
        )
    }
}

impl ToHtml for Cite {
    fn to_html(&self, context: &EncodeContext) -> String {
        let content = match &self.content {
            Some(nodes) => nodes.to_html(context),
            None => {
                // Get the list of references from the root nodes
                let references = match context.root {
                    Node::Article(article) => article.references.clone(),
                    Node::CreativeWork(work) => work.references.clone(),
                    _ => {
                        tracing::warn!("Unhandled root document type");
                        None
                    }
                };
                if let Some(references) = references {
                    // Find the reference that matches the `target`
                    let reference = references
                        .iter()
                        .enumerate()
                        .find_map(|(index, reference)| match reference {
                            // A string reference so match against its index in the references list
                            CreativeWorkReferences::String(_) => {
                                if self.target == format!("ref{}", index + 1)
                                    || self.target == format!("ref{}", index + 1)
                                {
                                    Some((index, reference))
                                } else {
                                    None
                                }
                            }
                            // A `CreativeWork` reference so match against its `id`/
                            CreativeWorkReferences::CreativeWorkTypes(work) => match work {
                                CreativeWorkTypes::Article(Article { id, .. }) => {
                                    if let Some(id) = id.as_deref() {
                                        if self.target == *id {
                                            Some((index, reference))
                                        } else {
                                            None
                                        }
                                    } else {
                                        None
                                    }
                                }
                                _ => None,
                            },
                        });
                    // Create the content for the citation
                    match reference {
                        None => {
                            tracing::warn!("When encoding citation was unable to find reference '{}' in root document", self.target);
                            format!(r#"{}"#, self.target.clone())
                        }
                        Some((index, reference)) => {
                            // Always have a numeric citation
                            let mut content = format!(r#"{}"#, index + 1);
                            // If a `CreativeWorkType` then add authors and year
                            if let CreativeWorkReferences::CreativeWorkTypes(work) = reference {
                                let (authors, date) = match &work {
                                    CreativeWorkTypes::Article(article) => {
                                        (article.authors.clone(), article.date_published.clone())
                                    }
                                    CreativeWorkTypes::CreativeWork(work) => {
                                        (work.authors.clone(), work.date_published.clone())
                                    }
                                    _ => {
                                        tracing::warn!("Unhandled root document type");
                                        (None, None)
                                    }
                                };
                                if let Some(authors) = authors {
                                    let names: Vec<String> = authors
                                        .iter()
                                        .map(|author| match author {
                                            CreativeWorkAuthors::Person(person) => {
                                                match &person.family_names {
                                                    Some(family_names) => family_names.join(" "),
                                                    None => "Anonymous".to_string(),
                                                }
                                            }
                                            CreativeWorkAuthors::Organization(org) => {
                                                match &org.name {
                                                    Some(name) => *name.clone(),
                                                    None => "Anonymous".to_string(),
                                                }
                                            }
                                        })
                                        .collect();
                                    let names = if authors.len() == 1 {
                                        names.join("")
                                    } else if authors.len() == 2 {
                                        format!(r#"{} and {}"#, names[0], names[1])
                                    } else {
                                        format!(r#"{} et al"#, names[0])
                                    };
                                    content += &format!(r#"{}"#, names)
                                }
                                if let Some(date) = date {
                                    if date.value.len() >= 4 {
                                        let year = date.value[..4].to_string();
                                        content += &format!(r#"{}"#, year)
                                    }
                                }
                            }
                            content
                        }
                    }
                } else {
                    tracing::warn!("When encoding citation was unable to find references list in root document");
                    format!(r#"{}"#, self.target.clone())
                }
            }
        };
        elem(
            "cite",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &elem("a", &[attr("href", &self.target)], &content),
        )
    }
}

impl ToHtml for CiteGroup {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "span",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &concat(&self.items, |cite| cite.to_html(context)),
        )
    }
}

impl ToHtml for CodeExpression {
    fn to_html(&self, context: &EncodeContext) -> String {
        let output = match &self.output {
            Some(output) => output.to_html(context),
            None => "".to_string(),
        };
        elem(
            "stencila-code-expression",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("programming-language", &self.programming_language),
            ],
            &[
                elem("code", &[attr("slot", "text")], &encode_safe(&self.text)),
                elem("output", &[attr("slot", "output")], &output),
            ]
            .concat(),
        )
    }
}

impl ToHtml for CodeFragment {
    /// Encode a [`CodeFragment`] as HTML
    ///
    /// See `CodeBlock::to_html` for why `programming_language` is encoded
    /// as both a `class` attribute and a `<meta>` element.
    fn to_html(&self, _context: &EncodeContext) -> String {
        let (class, meta) = match &self.programming_language {
            Some(programming_language) => (
                attr("class", &["language-", programming_language].concat()),
                elem_empty(
                    "meta",
                    &[
                        attr_itemprop("programming_language"),
                        attr("content", programming_language),
                    ],
                ),
            ),
            None => ("".to_string(), "".to_string()),
        };

        let text = encode_safe(&self.text);

        elem(
            "code",
            &[attr_itemtype::<Self>(), attr_id(&self.id), class],
            &[meta, text.to_string()].concat(),
        )
    }
}

impl ToHtml for Link {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "a",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("href", &self.target),
            ],
            &self.content.to_html(context),
        )
    }
}

impl ToHtml for MathFragment {
    fn to_html(&self, _context: &EncodeContext) -> String {
        elem(
            "code",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &encode_safe(&self.text),
        )
    }
}

impl ToHtml for Note {
    fn to_html(&self, _context: &EncodeContext) -> String {
        elem(
            "code",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("class", "todo"),
            ],
            &json(self),
        )
    }
}

impl ToHtml for Parameter {
    fn to_html(&self, _context: &EncodeContext) -> String {
        let input_type = match self.validator.as_deref() {
            Some(ValidatorTypes::NumberValidator(..)) => "number",
            _ => "text",
        };
        let value_attr = match self.value.as_deref() {
            Some(node) => attr("value", &node.to_txt()),
            _ => "".to_string(),
        };
        elem_empty(
            "input",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("type", input_type),
                attr("name", &self.name),
                value_attr,
            ],
        )
    }
}

impl ToHtml for Quote {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "q",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &self.content.to_html(context),
        )
    }
}
<commit_msg>fix(HTML encoding): Encode Plotly and Vega images as custom components<commit_after>//! Encode `InlineContent` nodes to HTML

use super::{
    attr, attr_id, attr_itemprop, attr_itemtype, attr_itemtype_str, concat, elem, elem_empty, json,
    EncodeContext, ToHtml,
};
use codec_txt::ToTxt;
use html_escape::encode_safe;
use std::{fs, path::PathBuf};
use stencila_schema::*;

impl ToHtml for InlineContent {
    fn to_html(&self, context: &EncodeContext) -> String {
        match self {
            InlineContent::AudioObject(node) => node.to_html(context),
            InlineContent::Boolean(node) => node.to_html(context),
            InlineContent::Cite(node) => node.to_html(context),
            InlineContent::CiteGroup(node) => node.to_html(context),
            InlineContent::CodeExpression(node) => node.to_html(context),
            InlineContent::CodeFragment(node) => node.to_html(context),
            InlineContent::Delete(node) => node.to_html(context),
            InlineContent::Emphasis(node) => node.to_html(context),
            InlineContent::ImageObject(node) => node.to_html(context),
            InlineContent::Integer(node) => node.to_html(context),
            InlineContent::Link(node) => node.to_html(context),
            InlineContent::MathFragment(node) => node.to_html(context),
            InlineContent::NontextualAnnotation(node) => node.to_html(context),
            InlineContent::Note(node) => node.to_html(context),
            InlineContent::Null(node) => node.to_html(context),
            InlineContent::Number(node) => node.to_html(context),
            InlineContent::Parameter(node) => node.to_html(context),
            InlineContent::Quote(node) => node.to_html(context),
            // Unlike everything else, `String` instances are not wrapped in an element.
            // However, `InlineContent::String` (and `Node::String` instances) should be
            // because they are usually part of a vector (e.g. `Paragraph.content`).
            // Wrapping them indicates to the DOM patching algorithm that the container element
            // represents a vector and not a placeholder for an optional string property.
            // To reduce the size of the generated HTML a little, unlike for `Node::String`,
            // the `itemtype` attribute is not used.
            InlineContent::String(node) => elem("span", &[], &node.to_html(context)),
            InlineContent::Strong(node) => node.to_html(context),
            InlineContent::Subscript(node) => node.to_html(context),
            InlineContent::Superscript(node) => node.to_html(context),
            InlineContent::VideoObject(node) => node.to_html(context),
        }
    }
}

macro_rules! mark_to_html {
    ($type:ident, $tag:literal) => {
        impl ToHtml for $type {
            fn to_html(&self, context: &EncodeContext) -> String {
                elem(
                    $tag,
                    &[attr_itemtype::<Self>(), attr_id(&self.id)],
                    &self.content.to_html(context),
                )
            }
        }
    };
}

mark_to_html!(Delete, "del");
mark_to_html!(Emphasis, "em");
mark_to_html!(NontextualAnnotation, "u");
mark_to_html!(Strong, "strong");
mark_to_html!(Subscript, "sub");
mark_to_html!(Superscript, "sup");

/// Convert a file:// URL to a data:// URI
///
/// Note that this function assumes that paths are absolute.
/// File URLs are usually resolves elsewhere e.g. in the `compile` method
/// before encoding to HTML.
fn file_uri_to_data_uri(url: &str) -> String {
    let path = if let Some(path) = url.strip_prefix("file://") {
        PathBuf::from(path)
    } else {
        return url.into();
    };

    // Read the file, convert it to a dataURI, and record it as a dependency
    match fs::read(&path) {
        Ok(bytes) => {
            let mime = match mime_guess::from_path(&path).first() {
                Some(mime) => mime.to_string(),
                None => "image/png".to_string(),
            };
            let data = base64::encode(bytes);
            ["data:", &mime, ";base64,", &data].concat()
        }
        Err(error) => {
            tracing::warn!("Unable to read media file {}: {}", path.display(), error);
            url.into()
        }
    }
}

fn content_url_to_src_attr(content_url: &str, context: &EncodeContext) -> String {
    let url = match context.bundle {
        true => file_uri_to_data_uri(content_url),
        false => content_url.to_string(),
    };
    attr("src", &url)
}

impl ToHtml for AudioObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "audio",
            &[
                attr_itemtype_str("AudioObject"),
                attr_id(&self.id),
                "controls".to_string(),
                content_url_to_src_attr(&self.content_url, context),
            ],
            "",
        )
    }
}

impl ToHtml for ImageObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        fn decode_base64(data: &str) -> String {
            match base64::decode(data) {
                Ok(data) => String::from_utf8_lossy(&data).to_string(),
                Err(error) => {
                    tracing::error!("While decoding base64 image data: {}", error);
                    format!("\"{}\"", error)
                }
            }
        }

        if let Some(data) = self
            .content_url
            .strip_prefix("data:application/vnd.plotly.")
        {
            let mut parts = data.split(";base64,");

            let version = parts.next().unwrap_or("1");
            let media_type = match &self.media_type {
                Some(value) => *value.clone(),
                None => ["application/vnd.plotly.", version].concat(),
            };

            let data = decode_base64(parts.next().unwrap_or_default());

            elem(
                "stencila-image-plotly",
                &[],
                &elem(
                    "picture",
                    &[],
                    &elem("script", &[attr("type", &media_type)], &data),
                ),
            )
        } else if let Some(data) = self.content_url.strip_prefix("data:application/vnd.vega") {
            let mut parts = data.split(";base64,");
            let data = decode_base64(parts.nth(1).unwrap_or_default());

            elem(
                "stencila-image-vega",
                &[],
                &elem(
                    "picture",
                    &[],
                    &elem(
                        "script",
                        &[attr("type", "application/vnd.vega+json")],
                        &data,
                    ),
                ),
            )
        } else {
            elem_empty(
                "img",
                &[
                    attr_itemtype_str("ImageObject"),
                    attr_id(&self.id),
                    content_url_to_src_attr(&self.content_url, context),
                ],
            )
        }
    }
}

impl ToHtml for VideoObjectSimple {
    fn to_html(&self, context: &EncodeContext) -> String {
        let src_attr = content_url_to_src_attr(&self.content_url, context);
        let type_attr = match &self.media_type {
            Some(media_type) => attr("type", media_type),
            None => "".to_string(),
        };
        elem(
            "video",
            &[
                attr_itemtype_str("VideoObject"),
                attr_id(&self.id),
                "controls".to_string(),
            ],
            &elem("source", &[src_attr, type_attr], ""),
        )
    }
}

impl ToHtml for Cite {
    fn to_html(&self, context: &EncodeContext) -> String {
        let content = match &self.content {
            Some(nodes) => nodes.to_html(context),
            None => {
                // Get the list of references from the root nodes
                let references = match context.root {
                    Node::Article(article) => article.references.clone(),
                    Node::CreativeWork(work) => work.references.clone(),
                    _ => {
                        tracing::warn!("Unhandled root document type");
                        None
                    }
                };
                if let Some(references) = references {
                    // Find the reference that matches the `target`
                    let reference = references
                        .iter()
                        .enumerate()
                        .find_map(|(index, reference)| match reference {
                            // A string reference so match against its index in the references list
                            CreativeWorkReferences::String(_) => {
                                if self.target == format!("ref{}", index + 1)
                                    || self.target == format!("ref{}", index + 1)
                                {
                                    Some((index, reference))
                                } else {
                                    None
                                }
                            }
                            // A `CreativeWork` reference so match against its `id`/
                            CreativeWorkReferences::CreativeWorkTypes(work) => match work {
                                CreativeWorkTypes::Article(Article { id, .. }) => {
                                    if let Some(id) = id.as_deref() {
                                        if self.target == *id {
                                            Some((index, reference))
                                        } else {
                                            None
                                        }
                                    } else {
                                        None
                                    }
                                }
                                _ => None,
                            },
                        });
                    // Create the content for the citation
                    match reference {
                        None => {
                            tracing::warn!("When encoding citation was unable to find reference '{}' in root document", self.target);
                            format!(r#"{}"#, self.target.clone())
                        }
                        Some((index, reference)) => {
                            // Always have a numeric citation
                            let mut content = format!(r#"{}"#, index + 1);
                            // If a `CreativeWorkType` then add authors and year
                            if let CreativeWorkReferences::CreativeWorkTypes(work) = reference {
                                let (authors, date) = match &work {
                                    CreativeWorkTypes::Article(article) => {
                                        (article.authors.clone(), article.date_published.clone())
                                    }
                                    CreativeWorkTypes::CreativeWork(work) => {
                                        (work.authors.clone(), work.date_published.clone())
                                    }
                                    _ => {
                                        tracing::warn!("Unhandled root document type");
                                        (None, None)
                                    }
                                };
                                if let Some(authors) = authors {
                                    let names: Vec<String> = authors
                                        .iter()
                                        .map(|author| match author {
                                            CreativeWorkAuthors::Person(person) => {
                                                match &person.family_names {
                                                    Some(family_names) => family_names.join(" "),
                                                    None => "Anonymous".to_string(),
                                                }
                                            }
                                            CreativeWorkAuthors::Organization(org) => {
                                                match &org.name {
                                                    Some(name) => *name.clone(),
                                                    None => "Anonymous".to_string(),
                                                }
                                            }
                                        })
                                        .collect();
                                    let names = if authors.len() == 1 {
                                        names.join("")
                                    } else if authors.len() == 2 {
                                        format!(r#"{} and {}"#, names[0], names[1])
                                    } else {
                                        format!(r#"{} et al"#, names[0])
                                    };
                                    content += &format!(r#"{}"#, names)
                                }
                                if let Some(date) = date {
                                    if date.value.len() >= 4 {
                                        let year = date.value[..4].to_string();
                                        content += &format!(r#"{}"#, year)
                                    }
                                }
                            }
                            content
                        }
                    }
                } else {
                    tracing::warn!("When encoding citation was unable to find references list in root document");
                    format!(r#"{}"#, self.target.clone())
                }
            }
        };
        elem(
            "cite",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &elem("a", &[attr("href", &self.target)], &content),
        )
    }
}

impl ToHtml for CiteGroup {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "span",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &concat(&self.items, |cite| cite.to_html(context)),
        )
    }
}

impl ToHtml for CodeExpression {
    fn to_html(&self, context: &EncodeContext) -> String {
        let output = match &self.output {
            Some(output) => output.to_html(context),
            None => "".to_string(),
        };
        elem(
            "stencila-code-expression",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("programming-language", &self.programming_language),
            ],
            &[
                elem("code", &[attr("slot", "text")], &encode_safe(&self.text)),
                elem("output", &[attr("slot", "output")], &output),
            ]
            .concat(),
        )
    }
}

impl ToHtml for CodeFragment {
    /// Encode a [`CodeFragment`] as HTML
    ///
    /// See `CodeBlock::to_html` for why `programming_language` is encoded
    /// as both a `class` attribute and a `<meta>` element.
    fn to_html(&self, _context: &EncodeContext) -> String {
        let (class, meta) = match &self.programming_language {
            Some(programming_language) => (
                attr("class", &["language-", programming_language].concat()),
                elem_empty(
                    "meta",
                    &[
                        attr_itemprop("programming_language"),
                        attr("content", programming_language),
                    ],
                ),
            ),
            None => ("".to_string(), "".to_string()),
        };

        let text = encode_safe(&self.text);

        elem(
            "code",
            &[attr_itemtype::<Self>(), attr_id(&self.id), class],
            &[meta, text.to_string()].concat(),
        )
    }
}

impl ToHtml for Link {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "a",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("href", &self.target),
            ],
            &self.content.to_html(context),
        )
    }
}

impl ToHtml for MathFragment {
    fn to_html(&self, _context: &EncodeContext) -> String {
        elem(
            "code",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &encode_safe(&self.text),
        )
    }
}

impl ToHtml for Note {
    fn to_html(&self, _context: &EncodeContext) -> String {
        elem(
            "code",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("class", "todo"),
            ],
            &json(self),
        )
    }
}

impl ToHtml for Parameter {
    fn to_html(&self, _context: &EncodeContext) -> String {
        let input_type = match self.validator.as_deref() {
            Some(ValidatorTypes::NumberValidator(..)) => "number",
            _ => "text",
        };
        let value_attr = match self.value.as_deref() {
            Some(node) => attr("value", &node.to_txt()),
            _ => "".to_string(),
        };
        elem_empty(
            "input",
            &[
                attr_itemtype::<Self>(),
                attr_id(&self.id),
                attr("type", input_type),
                attr("name", &self.name),
                value_attr,
            ],
        )
    }
}

impl ToHtml for Quote {
    fn to_html(&self, context: &EncodeContext) -> String {
        elem(
            "q",
            &[attr_itemtype::<Self>(), attr_id(&self.id)],
            &self.content.to_html(context),
        )
    }
}
<|endoftext|>