1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! Low-level building blocks for the command server protocol.

use bytes::Bytes;
use futures::SinkExt;
use std::io;
use tokio::stream::{self, StreamExt};

use crate::codec::{BlockMessage, ChannelMessage};
use crate::connection::Connection;

/// Connection wrapper to process channel requests and responses.
///
/// This provides a low-level interface to communicate with the Mercurial
/// command server.
#[derive(Debug)]
pub struct Protocol<C> {
    conn: C,
}

impl<C> Protocol<C> {
    /// Creates a new Protocol that wraps the given connection.
    pub fn new(conn: C) -> Self {
        Protocol { conn }
    }

    /// Unwraps the underlying `Connection`.
    pub fn into_connection(self) -> C {
        self.conn
    }
}

#[cfg(unix)]
mod unix {
    use super::*;
    use std::os::unix::io::{AsRawFd, RawFd};

    impl<C> AsRawFd for Protocol<C>
    where
        C: AsRawFd,
    {
        fn as_raw_fd(&self) -> RawFd {
            self.conn.as_raw_fd()
        }
    }
}

impl<C> Protocol<C>
where
    C: Connection,
{
    /// Sends the command of no argument without waiting for response.
    ///
    /// This is equivalent to `OneShotRequest::start()` of tokio-hglib 0.2.
    /// For `MessageLoop::start()`, call this function and fetch responses.
    pub async fn send_command(&mut self, cmd: impl Into<Bytes>) -> io::Result<()> {
        self.conn
            .get_tx_mut()
            .send(BlockMessage::Command(cmd.into()))
            .await?;
        Ok(())
    }

    /// Sends the command and arguments without waiting for response.
    ///
    /// This is equivalent to `OneShotRequest::start_with_args()` of tokio-hglib 0.2.
    /// For `MessageLoop::start_with_args()`, call this function and fetch responses.
    pub async fn send_command_with_args(
        &mut self,
        cmd: impl Into<Bytes>,
        packed_args: impl Into<Bytes>,
    ) -> io::Result<()> {
        let blocks = vec![
            Ok(BlockMessage::Command(cmd.into())),
            Ok(BlockMessage::Data(packed_args.into())),
        ];
        self.conn
            .get_tx_mut()
            .send_all(&mut stream::iter(blocks))
            .await?;
        Ok(())
    }

    /// Sends the given data back to the server.
    ///
    /// For `MessageLoop::resume_with_data()` of tokio-hglib 0.2, call this function
    /// and fetch responses.
    pub async fn send_data(&mut self, data: impl Into<Bytes>) -> io::Result<()> {
        self.conn
            .get_tx_mut()
            .send(BlockMessage::Data(data.into()))
            .await?;
        Ok(())
    }

    /// Sends the command of no argument, and fetches the result data.
    ///
    /// This is equivalent to `OneShotQuery::start()` of tokio-hglib 0.2.
    pub async fn query(&mut self, cmd: impl Into<Bytes>) -> io::Result<Bytes> {
        self.send_command(cmd).await?;
        self.fetch_result().await
    }

    /// Sends the command and arguments, and fetches the result data.
    ///
    /// This is equivalent to `OneShotQuery::start_with_args()` of tokio-hglib 0.2.
    pub async fn query_with_args(
        &mut self,
        cmd: impl Into<Bytes>,
        packed_args: impl Into<Bytes>,
    ) -> io::Result<Bytes> {
        self.send_command_with_args(cmd, packed_args).await?;
        self.fetch_result().await
    }

    /// Fetches response message from the server.
    ///
    /// This is equivalent to `MessageLoop::resume()` of tokio-hglib 0.2.
    pub async fn fetch_response(&mut self) -> io::Result<ChannelMessage> {
        let v = self.conn.get_rx_mut().try_next().await?;
        expect_msg(v)
    }

    async fn fetch_result(&mut self) -> io::Result<Bytes> {
        loop {
            match self.fetch_response().await? {
                ChannelMessage::Data(b'r', data) => {
                    return Ok(data);
                }
                ChannelMessage::Data(..) => {
                    // just ignore data sent to uninteresting (optional) channel
                }
                ChannelMessage::InputRequest(..)
                | ChannelMessage::LineRequest(..)
                | ChannelMessage::SystemRequest(..) => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "unsupported request while querying",
                    ));
                }
            }
        }
    }
}

fn expect_msg(v: Option<ChannelMessage>) -> Result<ChannelMessage, io::Error> {
    v.ok_or(io::Error::new(
        io::ErrorKind::UnexpectedEof,
        "no result code received",
    ))
}