sstp/client/
getter.rs

1use super::ws_conn::WSConn;
2use crate::messages::{ContentMessage, Message};
3use crate::utils::*;
4use std::error::Error;
5use std::io::Write;
6use std::net::SocketAddr;
7use std::str;
8
9pub struct Getter {
10  pub code: String,
11  pub relay_addr: SocketAddr,
12  ws_conn: WSConn,
13  pub filename: Option<String>,
14  pub size: Option<u64>,
15  pub peer_addr: Option<SocketAddr>,
16}
17
18impl Getter {
19  pub fn new(code: Option<&str>, relay_addr: Option<&str>) -> Self {
20    let code = code.expect("Room code must be provided").to_string();
21    let relay_addr = parse_relay_addr(relay_addr);
22    Getter {
23      code,
24      relay_addr,
25      ws_conn: WSConn::new_empty(),
26      filename: None,
27      size: None,
28      peer_addr: None,
29    }
30  }
31
32  pub async fn connect(&mut self) -> Result<(), Box<dyn Error>> {
33    let (tx, rx) = start_ws_conn(self.relay_addr).await?;
34    self.ws_conn.init(tx, rx);
35    Ok(())
36  }
37
38  pub async fn get_room(&mut self) -> Result<(), Box<dyn Error>> {
39    let msg = Message::new_get(&self.code);
40    self.ws_conn.send(&msg).await?;
41
42    let mut data = Vec::new();
43    self.ws_conn.recv(&mut data).await?;
44    let message: Message = serde_json::from_slice(&data)?;
45    match message {
46      Message::Error(error) => Err(Box::new(error)),
47      Message::ApproveReq(req) => {
48        self.filename = Some(req.filename.clone());
49        self.size = Some(req.size);
50        self.peer_addr = Some(req.addr);
51        Ok(())
52      }
53      _ => unreachable!(),
54    }
55  }
56
57  pub async fn send_approval(&mut self, approved: bool) -> Result<(), Box<dyn Error>> {
58    let res_message = Message::new_approve_res(approved);
59    self.ws_conn.send(&res_message).await?;
60    Ok(())
61  }
62
63  pub async fn start_transfer(
64    &mut self,
65    wtr: &mut impl Write,
66    f: impl Fn(u64) -> (),
67  ) -> Result<(), Box<dyn Error>> {
68    let chunks = calc_chunks(self.size.unwrap());
69    for _ in 0..chunks {
70      let mut data = Vec::new();
71      self.ws_conn.recv(&mut data).await?;
72      let mut msg: ContentMessage = serde_json::from_slice(&data)?;
73      f(msg.content.len().try_into()?);
74      wtr.write(&mut msg.content)?;
75    }
76    self.ws_conn.close().await?;
77    Ok(())
78  }
79}