thrussh_keys/agent/
server.rs

1use crate::encoding::{Encoding, Position, Reader};
2use crate::key;
3#[cfg(feature = "openssl")]
4use crate::key::SignatureHash;
5use byteorder::{BigEndian, ByteOrder};
6use cryptovec::CryptoVec;
7use futures::future::Future;
8use futures::stream::{Stream, StreamExt};
9use std;
10use std::collections::HashMap;
11use std::sync::{Arc, RwLock};
12use std::time::Duration;
13use std::time::SystemTime;
14use tokio;
15use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
16use tokio::time::sleep;
17
18use super::msg;
19use super::Constraint;
20use crate::Error;
21
22#[derive(Clone)]
23struct KeyStore(Arc<RwLock<HashMap<Vec<u8>, (Arc<key::KeyPair>, SystemTime, Vec<Constraint>)>>>);
24
25#[derive(Clone)]
26struct Lock(Arc<RwLock<CryptoVec>>);
27
28#[allow(missing_docs)]
29#[derive(Debug)]
30pub enum ServerError<E> {
31    E(E),
32    Error(Error),
33}
34
35pub trait Agent: Clone + Send + 'static {
36    fn confirm(
37        self,
38        _pk: Arc<key::KeyPair>,
39    ) -> Box<dyn Future<Output = (Self, bool)> + Unpin + Send> {
40        Box::new(futures::future::ready((self, true)))
41    }
42}
43
44pub async fn serve<S, L, A>(mut listener: L, agent: A) -> Result<(), Error>
45where
46    S: AsyncRead + AsyncWrite + Send + Sync + Unpin + 'static,
47    L: Stream<Item = tokio::io::Result<S>> + Unpin,
48    A: Agent + Send + Sync + 'static,
49{
50    let keys = KeyStore(Arc::new(RwLock::new(HashMap::new())));
51    let lock = Lock(Arc::new(RwLock::new(CryptoVec::new())));
52    while let Some(Ok(stream)) = listener.next().await {
53        let mut buf = CryptoVec::new();
54        buf.resize(4);
55        tokio::spawn(
56            (Connection {
57                lock: lock.clone(),
58                keys: keys.clone(),
59                agent: Some(agent.clone()),
60                s: stream,
61                buf: CryptoVec::new(),
62            })
63            .run(),
64        );
65    }
66    Ok(())
67}
68
69impl Agent for () {
70    fn confirm(
71        self,
72        _: Arc<key::KeyPair>,
73    ) -> Box<dyn Future<Output = (Self, bool)> + Unpin + Send> {
74        Box::new(futures::future::ready((self, true)))
75    }
76}
77
78struct Connection<S: AsyncRead + AsyncWrite + Send + 'static, A: Agent> {
79    lock: Lock,
80    keys: KeyStore,
81    agent: Option<A>,
82    s: S,
83    buf: CryptoVec,
84}
85
86impl<S: AsyncRead + AsyncWrite + Send + Unpin + 'static, A: Agent + Send + 'static>
87    Connection<S, A>
88{
89    async fn run(mut self) -> Result<(), Error> {
90        let mut writebuf = CryptoVec::new();
91        loop {
92            // Reading the length
93            self.buf.clear();
94            self.buf.resize(4);
95            self.s.read_exact(&mut self.buf).await?;
96            // Reading the rest of the buffer
97            let len = BigEndian::read_u32(&self.buf) as usize;
98            self.buf.clear();
99            self.buf.resize(len);
100            self.s.read_exact(&mut self.buf).await?;
101            // respond
102            writebuf.clear();
103            self.respond(&mut writebuf).await?;
104            self.s.write_all(&writebuf).await?;
105            self.s.flush().await?
106        }
107    }
108
109    async fn respond(&mut self, writebuf: &mut CryptoVec) -> Result<(), Error> {
110        let is_locked = {
111            if let Ok(password) = self.lock.0.read() {
112                !password.is_empty()
113            } else {
114                true
115            }
116        };
117        writebuf.extend(&[0, 0, 0, 0]);
118        let mut r = self.buf.reader(0);
119        match r.read_byte() {
120            Ok(11) if !is_locked => {
121                // request identities
122                if let Ok(keys) = self.keys.0.read() {
123                    writebuf.push(msg::IDENTITIES_ANSWER);
124                    writebuf.push_u32_be(keys.len() as u32);
125                    for (k, _) in keys.iter() {
126                        writebuf.extend_ssh_string(k);
127                        writebuf.extend_ssh_string(b"");
128                    }
129                } else {
130                    writebuf.push(msg::FAILURE)
131                }
132            }
133            Ok(13) if !is_locked => {
134                // sign request
135                let agent = self.agent.take().unwrap();
136                let (agent, signed) = self.try_sign(agent, r, writebuf).await?;
137                self.agent = Some(agent);
138                if signed {
139                    return Ok(());
140                } else {
141                    writebuf.resize(4);
142                    writebuf.push(msg::FAILURE)
143                }
144            }
145            Ok(17) if !is_locked => {
146                // add identity
147                if let Ok(true) = self.add_key(r, false, writebuf).await {
148                } else {
149                    writebuf.push(msg::FAILURE)
150                }
151            }
152            Ok(18) if !is_locked => {
153                // remove identity
154                if let Ok(true) = self.remove_identity(r) {
155                    writebuf.push(msg::SUCCESS)
156                } else {
157                    writebuf.push(msg::FAILURE)
158                }
159            }
160            Ok(19) if !is_locked => {
161                // remove all identities
162                if let Ok(mut keys) = self.keys.0.write() {
163                    keys.clear();
164                    writebuf.push(msg::SUCCESS)
165                } else {
166                    writebuf.push(msg::FAILURE)
167                }
168            }
169            Ok(22) if !is_locked => {
170                // lock
171                if let Ok(()) = self.lock(r) {
172                    writebuf.push(msg::SUCCESS)
173                } else {
174                    writebuf.push(msg::FAILURE)
175                }
176            }
177            Ok(23) if is_locked => {
178                // unlock
179                if let Ok(true) = self.unlock(r) {
180                    writebuf.push(msg::SUCCESS)
181                } else {
182                    writebuf.push(msg::FAILURE)
183                }
184            }
185            Ok(25) if !is_locked => {
186                // add identity constrained
187                if let Ok(true) = self.add_key(r, true, writebuf).await {
188                } else {
189                    writebuf.push(msg::FAILURE)
190                }
191            }
192            _ => {
193                // Message not understood
194                writebuf.push(msg::FAILURE)
195            }
196        }
197        let len = writebuf.len() - 4;
198        BigEndian::write_u32(&mut writebuf[0..], len as u32);
199        Ok(())
200    }
201
202    fn lock(&self, mut r: Position) -> Result<(), Error> {
203        let password = r.read_string()?;
204        let mut lock = self.lock.0.write().unwrap();
205        lock.extend(password);
206        Ok(())
207    }
208
209    fn unlock(&self, mut r: Position) -> Result<bool, Error> {
210        let password = r.read_string()?;
211        let mut lock = self.lock.0.write().unwrap();
212        if &lock[0..] == password {
213            lock.clear();
214            Ok(true)
215        } else {
216            Ok(false)
217        }
218    }
219
220    fn remove_identity(&self, mut r: Position) -> Result<bool, Error> {
221        if let Ok(mut keys) = self.keys.0.write() {
222            if keys.remove(r.read_string()?).is_some() {
223                Ok(true)
224            } else {
225                Ok(false)
226            }
227        } else {
228            Ok(false)
229        }
230    }
231
232    async fn add_key<'a>(
233        &self,
234        mut r: Position<'a>,
235        constrained: bool,
236        writebuf: &mut CryptoVec,
237    ) -> Result<bool, Error> {
238        let pos0 = r.position;
239        let t = r.read_string()?;
240        let (blob, key) = match t {
241            b"ssh-ed25519" => {
242                let public_ = r.read_string()?;
243                let pos1 = r.position;
244                let concat = r.read_string()?;
245                let _comment = r.read_string()?;
246                if &concat[32..64] != public_ {
247                    return Ok(false);
248                }
249                use key::ed25519::*;
250                let mut public = PublicKey::new_zeroed();
251                let mut secret = SecretKey::new_zeroed();
252                public.key.clone_from_slice(&public_[..32]);
253                secret.key.clone_from_slice(&concat[..]);
254                writebuf.push(msg::SUCCESS);
255                (self.buf[pos0..pos1].to_vec(), key::KeyPair::Ed25519(secret))
256            }
257            #[cfg(feature = "p256")]
258            b"ecdsa-sha2-nistp256" => {
259                let public_ = r.read_string()?;
260                let pos1 = r.position;
261                let priv_ = r.read_mpint()?;
262                let _comment = r.read_string()?;
263                let mut priv_bytes = [0u8; 32];
264                priv_.into_iter().rev().enumerate().for_each(|(i, b)| {
265                    if i < 32 {
266                        priv_bytes[31 - i] = *b;
267                    }
268                });
269                let key =
270                    crate::key::KeyPair::P256(p256::SecretKey::from_bytes(&priv_bytes.into())?);
271                if public_ != crate::PublicKeyBase64::public_key_bytes(&key) {
272                    return Ok(false);
273                }
274                (self.buf[pos0..pos1].to_vec(), key)
275            }
276            #[cfg(feature = "openssl")]
277            b"ssh-rsa" => {
278                use openssl::bn::{BigNum, BigNumContext};
279                use openssl::rsa::Rsa;
280                let n = r.read_mpint()?;
281                let e = r.read_mpint()?;
282                let d = BigNum::from_slice(r.read_mpint()?)?;
283                let q_inv = r.read_mpint()?;
284                let p = BigNum::from_slice(r.read_mpint()?)?;
285                let q = BigNum::from_slice(r.read_mpint()?)?;
286                let (dp, dq) = {
287                    let one = BigNum::from_u32(1)?;
288                    let p1 = p.as_ref() - one.as_ref();
289                    let q1 = q.as_ref() - one.as_ref();
290                    let mut context = BigNumContext::new()?;
291                    let mut dp = BigNum::new()?;
292                    let mut dq = BigNum::new()?;
293                    dp.checked_rem(&d, &p1, &mut context)?;
294                    dq.checked_rem(&d, &q1, &mut context)?;
295                    (dp, dq)
296                };
297                let _comment = r.read_string()?;
298                let key = Rsa::from_private_components(
299                    BigNum::from_slice(n)?,
300                    BigNum::from_slice(e)?,
301                    d,
302                    p,
303                    q,
304                    dp,
305                    dq,
306                    BigNum::from_slice(&q_inv)?,
307                )?;
308
309                let len0 = writebuf.len();
310                writebuf.extend_ssh_string(b"ssh-rsa");
311                writebuf.extend_ssh_mpint(&e);
312                writebuf.extend_ssh_mpint(&n);
313                let blob = writebuf[len0..].to_vec();
314                writebuf.resize(len0);
315                writebuf.push(msg::SUCCESS);
316                (
317                    blob,
318                    key::KeyPair::RSA {
319                        key,
320                        hash: SignatureHash::SHA2_256,
321                    },
322                )
323            }
324            _ => return Ok(false),
325        };
326        let mut w = self.keys.0.write().unwrap();
327        let now = SystemTime::now();
328        if constrained {
329            let n = r.read_u32()?;
330            let mut c = Vec::new();
331            for _ in 0..n {
332                let t = r.read_byte()?;
333                if t == msg::CONSTRAIN_LIFETIME {
334                    let seconds = r.read_u32()?;
335                    c.push(Constraint::KeyLifetime { seconds });
336                    let blob = blob.clone();
337                    let keys = self.keys.clone();
338                    tokio::spawn(async move {
339                        sleep(Duration::from_secs(seconds as u64)).await;
340                        let mut keys = keys.0.write().unwrap();
341                        let delete = if let Some(&(_, time, _)) = keys.get(&blob) {
342                            time == now
343                        } else {
344                            false
345                        };
346                        if delete {
347                            keys.remove(&blob);
348                        }
349                    });
350                } else if t == msg::CONSTRAIN_CONFIRM {
351                    c.push(Constraint::Confirm)
352                } else {
353                    return Ok(false);
354                }
355            }
356            w.insert(blob, (Arc::new(key), now, Vec::new()));
357        } else {
358            w.insert(blob, (Arc::new(key), now, Vec::new()));
359        }
360        Ok(true)
361    }
362
363    async fn try_sign<'a>(
364        &self,
365        agent: A,
366        mut r: Position<'a>,
367        writebuf: &mut CryptoVec,
368    ) -> Result<(A, bool), Error> {
369        let mut needs_confirm = false;
370        let key = {
371            let blob = r.read_string()?;
372            let k = self.keys.0.read().unwrap();
373            if let Some(&(ref key, _, ref constraints)) = k.get(blob) {
374                if constraints.iter().any(|c| *c == Constraint::Confirm) {
375                    needs_confirm = true;
376                }
377                key.clone()
378            } else {
379                return Ok((agent, false));
380            }
381        };
382        let agent = if needs_confirm {
383            let (agent, ok) = agent.confirm(key.clone()).await;
384            if !ok {
385                return Ok((agent, false));
386            }
387            agent
388        } else {
389            agent
390        };
391        writebuf.push(msg::SIGN_RESPONSE);
392        let data = r.read_string()?;
393        key.add_signature(writebuf, data)?;
394        let len = writebuf.len();
395        BigEndian::write_u32(writebuf, (len - 4) as u32);
396
397        Ok((agent, true))
398    }
399}
thrussh_keys/agent/server.rs

thrussh_keys/agent/
server.rs
