voltlane_client/
lib.rs

1use tokio::net::{tcp::OwnedReadHalf, TcpStream};
2use tokio_util::codec::{FramedRead, LengthDelimitedCodec};
3
4pub use voltlane_enc as enc;
5pub use voltlane_net as net;
6
7pub struct Connection {
8    pub read: FramedRead<OwnedReadHalf, LengthDelimitedCodec>,
9    pub write: tokio::net::tcp::OwnedWriteHalf,
10    pub keys: enc::easy::Keys,
11    pub server_pubkey: enc::easy::PubKey,
12    pub client_id: u64,
13    pub addr: String,
14}
15
16impl Connection {
17    /// Creates a new connection to the server, with a new set of keys.
18    /// To re-connect an existing connection, use `Connection::reconnect` instead.
19    pub async fn new(addr: &str) -> anyhow::Result<Self> {
20        let mut stream = TcpStream::connect(addr).await?;
21        net::configure_performance_tcp_socket(&mut stream)?;
22        let (read, mut write) = stream.into_split();
23        let mut read = net::new_framed_reader(read);
24
25        // 0. send the current protocol version
26        net::send_size_prefixed(
27            &mut write,
28            &net::ClientServerPacket::ProtocolVersion(net::PROTOCOL_VERSION).into_vec()?,
29        )
30        .await?;
31
32        // 1. receive server public key
33        let buffer = net::recv_size_prefixed(&mut read).await?;
34        let pubkey_bytes = match net::ClientServerPacket::from_slice(&buffer) {
35            Ok(net::ClientServerPacket::PubKey(key)) => key,
36            Ok(_) => {
37                return Err(anyhow::format_err!("Expected server public key packet"));
38            }
39            Err(e) => return Err(anyhow::format_err!("Invalid server public key packet: {}", e)),
40        };
41        // validate it
42        let server_pubkey = match enc::easy::pubkey_from_bytes(&pubkey_bytes) {
43            Ok(val) => val,
44            Err(e) => return Err(anyhow::format_err!("Invalid server public key: {}", e)),
45        };
46
47        // 2. send my public key
48        let keys = enc::easy::Keys::new();
49        let pubkey = keys.pubkey_to_bytes();
50        let packet = net::ClientServerPacket::PubKey(pubkey);
51        net::send_size_prefixed(&mut write, &packet.into_vec()?).await?;
52
53        // 3. receive my client id
54        let buffer = net::recv_size_prefixed(&mut read).await?;
55        let client_id = match net::ClientServerPacket::from_slice(&buffer) {
56            Ok(net::ClientServerPacket::ClientId(id)) => id,
57            Ok(_) => return Err(anyhow::format_err!("Expected client ID packet")),
58            Err(e) => return Err(anyhow::format_err!("Invalid client ID packet: {}", e)),
59        };
60
61        Ok(Connection {
62            addr: addr.to_string(),
63            read,
64            write,
65            keys,
66            client_id,
67            server_pubkey,
68        })
69    }
70
71    pub async fn reconnect(&mut self) -> anyhow::Result<()> {
72        // to reconnect, we first send the protocol version,
73        // then we receive the server's public key and validate that it's the same we already know.
74        // then we send out client id, which signals to the server that we're trying to reconnect.
75        // then the server sends us a challenge, which we have to decrypt, and re-encrypt with the
76        // server's public key, and send it.
77        // if the server hasn't kicked us out by then, we're done :) we send a ping to confirm that it's
78        // working and we're done.
79        let mut stream = TcpStream::connect(&self.addr).await?;
80        net::configure_performance_tcp_socket(&mut stream)?;
81        let (read, mut write) = stream.into_split();
82        let mut read = net::new_framed_reader(read);
83
84        // 0. send the current protocol version
85        net::send_size_prefixed(
86            &mut write,
87            &net::ClientServerPacket::ProtocolVersion(net::PROTOCOL_VERSION).into_vec()?,
88        )
89        .await?;
90
91        let buffer = net::recv_size_prefixed(&mut read).await?;
92        let pubkey_bytes = match net::ClientServerPacket::from_slice(&buffer) {
93            Ok(net::ClientServerPacket::PubKey(key)) => key,
94            Ok(_) => {
95                return Err(anyhow::format_err!("Expected server public key packet"));
96            }
97            Err(e) => return Err(anyhow::format_err!("Invalid server public key packet: {}", e)),
98        };
99
100        // validate it
101        let server_pubkey = match enc::easy::pubkey_from_bytes(&pubkey_bytes) {
102            Ok(val) => val,
103            Err(e) => return Err(anyhow::format_err!("Invalid server public key: {}", e)),
104        };
105
106        // check if the server's public key is the same as the one we already know
107        if server_pubkey != self.server_pubkey {
108            return Err(anyhow::format_err!("Server public key does not match"));
109        }
110
111        // 2. send my client id
112        let packet = net::ClientServerPacket::ClientId(self.client_id);
113        net::send_size_prefixed(&mut write, &packet.into_vec()?).await?;
114
115        // 3. receive challenge
116        let buffer = net::recv_size_prefixed(&mut read).await?;
117        let challenge = match net::ClientServerPacket::from_slice(&buffer) {
118            Ok(net::ClientServerPacket::Challenge(challenge)) => challenge,
119            Ok(_) => return Err(anyhow::format_err!("Expected challenge packet")),
120            Err(e) => return Err(anyhow::format_err!("Invalid challenge packet: {}", e)),
121        };
122
123        // 4. decrypt challenge
124        let encryption = self.keys.create_encryption(&self.server_pubkey);
125        let decrypted_challenge = match encryption.decrypt(challenge) {
126            Ok(val) => val,
127            Err(e) => return Err(anyhow::format_err!("Failed to decrypt challenge: {}", e)),
128        };
129
130        // 5. re-encrypt challenge with server's public key
131        let encrypted_challenge = encryption.encrypt(decrypted_challenge);
132        let packet = net::ClientServerPacket::ChallengeResponse(encrypted_challenge);
133        net::send_size_prefixed(&mut write, &packet.into_vec()?).await?;
134
135        // 6. send ping
136        let packet = net::ClientServerPacket::Ping;
137        net::send_size_prefixed(&mut write, &packet.into_vec()?).await?;
138
139        // 7. receive pong
140        let buffer = net::recv_size_prefixed(&mut read).await?;
141        match net::ClientServerPacket::from_slice(&buffer) {
142            Ok(net::ClientServerPacket::Ping) => (),
143            Ok(_) => return Err(anyhow::format_err!("Expected pong packet")),
144            Err(e) => return Err(anyhow::format_err!("Invalid pong packet: {}", e)),
145        };
146
147        self.read = read;
148        self.write = write;
149        Ok(())
150    }
151}
voltlane_client/lib.rs

voltlane_client/
lib.rs
