1pub use crate::{
2 api::proto::{
3 get_events::MaybeCount,
4 get_identity_command::Id,
5 get_messages_cmd::{MaybeApplication, TimeRange, TimeSlice},
6 import_identity_command::MaybeHandle,
7 import_identity_response::{FinalResponse, State},
8 send_message_cmd::SignIdentity,
9 Ack, CryptoMessage, GetEvents, GetIdentityCommand, GetMessagesCmd, IdentityResponse,
10 ImportIdentityCommand, ImportIdentityResponse, MessageResponse, PairingAck,
11 PairingInitiate, PairingRequest, SbEvent, SbEvents, SendMessageCmd, UnitResponse,
12 },
13 crypto::{CryptoMessageWrapper, Session, SessionState},
14 error::{Error, IntoRemoteErr, SbResult},
15 mdns::HostRecord,
16 response::{Identity, Message},
17 serialize::{ProtoStream, ToUuid},
18 types::ImportIdentityState,
19};
20use crate::{
21 api::{
22 proto::get_messages_cmd::time_range::{EndPoint, StartPoint},
23 proto::PairingSynAck,
24 },
25 types::DartFuture,
26};
27use chrono::NaiveDateTime;
28use dryoc::generichash::Key;
29
30use bip39::Mnemonic;
31use dryoc::{
32 constants::CRYPTO_GENERICHASH_BYTES_MIN, generichash::GenericHash, kx::PublicKey,
33 types::StackByteArray,
34};
35pub use std::{future::Future, net::SocketAddr};
36use tokio::io::{AsyncReadExt, AsyncWriteExt};
37pub use tokio::net::TcpStream;
38use uuid::Uuid;
39
40impl From<SocketAddr> for HostRecord {
41 fn from(value: SocketAddr) -> Self {
42 Self {
43 name: value.to_string(),
44 addr: [value.ip()].into_iter().collect(),
45 port: value.port(),
46 }
47 }
48}
49
50impl HostRecord {
51 #[cfg(not(feature = "flutter"))]
52 pub async fn connect(self) -> SbResult<ProtoStream<TcpStream>> {
53 self.connect_impl().await
54 }
55
56 pub(crate) async fn connect_impl(self) -> SbResult<ProtoStream<TcpStream>> {
57 for addr in self.addr {
58 println!("attempting to connect to {}", addr);
59 match TcpStream::connect((addr, self.port)).await {
60 Ok(c) => return Ok(ProtoStream::new(c)),
61 Err(err) => log::warn!("Failed to connect to {}: {}", addr, err),
62 }
63 }
64 Err(Error::NoAddr)
65 }
66}
67
68impl<A> ProtoStream<A>
69where
70 A: Unpin + Send + AsyncReadExt + AsyncWriteExt + 'static,
71{
72 pub async fn key_exchange(mut self, state: SessionState) -> SbResult<Option<Session<A>>> {
73 let i = PairingInitiate {
74 pubkey: state.kp.public_key.iter().copied().collect(),
75 };
76 self.write_message(&i).await?;
77
78 let v: PairingAck = self.read_message().await?;
79 let session_id = v
80 .session
81 .ok_or_else(|| Error::CorruptHeader)?
82 .session
83 .ok_or_else(|| Error::CorruptHeader)?;
84
85 let ack_remote_key: PublicKey = v.pubkey.as_slice().try_into()?;
86 let s = dryoc::kx::Session::new_client(&state.kp, &ack_remote_key)?;
87
88 log::debug!("tx key {:?}", s.tx_as_slice());
89 log::debug!("rx key {:?}", s.rx_as_slice());
90
91 if let Some(remotekey) = state.remotekey {
92 if remotekey != ack_remote_key {
93 return Err(Error::MitmDetected);
94 }
95 Ok(Some(Session {
96 session: session_id.as_uuid(),
97 session_keys: s,
98 state: SessionState {
99 kp: state.kp,
100 remotekey: Some(remotekey),
101 },
102 stream: self,
103 }))
104 } else {
105 Ok(None)
106 }
107 }
108
109 pub async fn pair<F, Fut>(
110 mut self,
111 state: SessionState,
112 app_name: String,
113 cb: F,
114 ) -> SbResult<Session<A>>
115 where
116 F: FnOnce(Mnemonic) -> Fut,
117 Fut: Future<Output = std::result::Result<bool, Box<dyn std::error::Error + Send + Sync>>>,
118 {
119 let i = PairingInitiate {
120 pubkey: state.kp.public_key.iter().copied().collect(),
121 };
122 self.write_message(&i).await?;
123
124 let v: PairingAck = self.read_message().await?;
125 let session_id = v
126 .session
127 .ok_or_else(|| Error::CorruptHeader)?
128 .session
129 .ok_or_else(|| Error::CorruptHeader)?;
130
131 let ack_remote_key: PublicKey = v.pubkey.as_slice().try_into()?;
132
133 let s = dryoc::kx::Session::new_client(&state.kp, &ack_remote_key)?;
134
135 log::debug!("tx key {:?}", s.tx_as_slice());
136 log::debug!("rx key {:?}", s.rx_as_slice());
137
138 if let Some(remotekey) = state.remotekey {
139 if remotekey != ack_remote_key {
140 return Err(Error::MitmDetected);
141 }
142 Ok(Session {
143 session: session_id.as_uuid(),
144 session_keys: s,
145 state: SessionState {
146 kp: state.kp,
147 remotekey: Some(remotekey),
148 },
149 stream: self,
150 })
151 } else {
152 let mut pr = PairingRequest::default();
153 pr.name = app_name;
154 pr.session = v.session;
155 let pr = CryptoMessageWrapper::new_message(&pr, s.rx_as_array())?;
156 self.write_message(pr.message()).await?;
157 let fingerprint: StackByteArray<{ CRYPTO_GENERICHASH_BYTES_MIN }> =
158 GenericHash::hash(&i.pubkey, None::<&Key>).unwrap();
159 let words = Mnemonic::from_entropy(fingerprint.as_ref())?;
160 let confirmed = cb(words).await?; let v: CryptoMessage = self.read_message().await?;
163
164 let v = CryptoMessageWrapper::new(v);
165
166 let ack: Ack = v.decrypt(s.tx_as_array())?;
167
168 log::info!("got ack {}", ack.success);
169
170 let mut synack = PairingSynAck::default();
171 if !ack.success || !confirmed {
172 synack.success = false;
173 synack.message = "Pairing request rejected".to_owned();
174 self.write_message(
175 CryptoMessageWrapper::new_message(&ack, s.rx_as_array())?.message(),
176 )
177 .await?;
178 return Err(Error::PairingFailed);
179 }
180
181 synack.success = true;
182 self.write_message(
183 CryptoMessageWrapper::new_message(&synack, s.rx_as_array())?.message(),
184 )
185 .await?;
186
187 Ok(Session {
188 session: session_id.as_uuid(),
189 session_keys: s,
190 state: SessionState {
191 kp: state.kp,
192 remotekey: Some(ack_remote_key),
193 },
194 stream: self,
195 })
196 }
197 }
198}
199
200impl<A> SessionTrait for Session<A>
201where
202 A: Unpin + Send + AsyncReadExt + AsyncWriteExt,
203{
204 #[cfg(feature = "flutter")]
205 fn set_on_connect(
206 &mut self,
207 on_connect: Box<
208 dyn Fn(Option<crate::api::api::SbSession>) -> flutter_rust_bridge::DartFnFuture<()>
209 + Send
210 + Sync
211 + 'static,
212 >,
213 ) {
214 self.stream.on_connect = Some(Box::new(on_connect));
215 }
216
217 #[cfg(feature = "flutter")]
218 fn on_connect<'a>(
219 &'a self,
220 ) -> Option<
221 &'a Box<
222 dyn Fn(Option<crate::api::api::SbSession>) -> flutter_rust_bridge::DartFnFuture<()>
223 + Send
224 + Sync
225 + 'static,
226 >,
227 > {
228 self.stream.on_connect.as_ref()
229 }
230
231 fn get_identity<'a>(&'a mut self, id: Option<Uuid>) -> DartFuture<'a, SbResult<Vec<Identity>>> {
232 Box::pin(async move {
233 let cmd = GetIdentityCommand {
234 header: Some(self.get_header()),
235 id: id.map(|v| Id::Identity(v.as_proto())),
236 owned: false,
237 };
238 self.write_crypto(cmd).await?;
239 let id: IdentityResponse = self.read_crypto().await?;
240 id.try_into()
241 })
242 }
243
244 fn is_closed<'a>(&'a mut self) -> DartFuture<'a, SbResult<bool>> {
245 Box::pin(async move { Ok(self.is_disconnected()) })
246 }
247
248 fn get_events<'a>(
249 &'a mut self,
250 block: bool,
251 count: Option<u32>,
252 ) -> DartFuture<'a, SbResult<Vec<SbEvent>>> {
253 Box::pin(async move {
254 let cmd = GetEvents {
255 header: Some(self.get_header()),
256 block,
257 maybe_count: count.map(|v| MaybeCount::Count(v)),
258 };
259 self.write_crypto(cmd).await?;
260 let resp: SbEvents = self.read_crypto().await?;
261 Ok(resp.events)
262 })
263 }
264
265 fn get_messages<'a>(
266 &'a mut self,
267 application: String,
268 limit: Option<i32>,
269 ) -> DartFuture<'a, SbResult<Vec<Message>>> {
270 Box::pin(async move {
271 let cmd = GetMessagesCmd {
272 header: Some(self.get_header()),
273 time_slice: None,
274 maybe_application: Some(MaybeApplication::Application(application)),
275 limit: limit.unwrap_or(-1),
276 };
277 self.write_crypto(cmd).await?;
278 let m: MessageResponse = self.read_crypto().await?;
279 m.try_into()
280 })
281 }
282
283 fn send_messages<'a>(
284 &'a mut self,
285 messages: Vec<Message>,
286 sign_identity: Option<Uuid>,
287 ) -> DartFuture<'a, SbResult<()>> {
288 Box::pin(async move {
289 let cmd = SendMessageCmd {
290 header: Some(self.get_header()),
291 messages: messages.into_iter().map(|v| v.into()).collect(),
292 sign_identity: sign_identity.map(|v| SignIdentity::Identity(v.as_proto())),
293 };
294 self.write_crypto(cmd).await?;
295 let m: UnitResponse = self.read_crypto().await?;
296 m.into_remote_err()?;
297 Ok(())
298 })
299 }
300
301 fn initiate_identity_import<'a>(
302 &'a mut self,
303 id: Option<Uuid>,
304 ) -> DartFuture<'a, SbResult<ImportIdentityState>> {
305 Box::pin(async move {
306 let cmd = ImportIdentityCommand {
307 header: Some(self.get_header()),
308 maybe_handle: id.map(|v| MaybeHandle::Handle(v.as_proto())),
309 };
310 self.write_crypto(cmd).await?;
311 let resp: ImportIdentityResponse = self.read_crypto().await?;
312 let state = resp
313 .state
314 .ok_or_else(|| Error::RemoteError("Missing state field".to_owned()))?;
315 let res = match state {
316 State::Handle(uuid) => ImportIdentityState::Initiated(uuid.as_uuid()),
317 State::Final(FinalResponse { identity, .. }) => ImportIdentityState::Complete(
318 identity
319 .ok_or_else(|| Error::RemoteError("missing identity uuid".to_owned()))?
320 .as_uuid(),
321 ),
322 };
323 Ok(res)
324 })
325 }
326
327 fn get_messages_send_date<'a>(
328 &'a mut self,
329 application: String,
330 limit: Option<i32>,
331 start_date: Option<NaiveDateTime>,
332 end_date: Option<NaiveDateTime>,
333 ) -> DartFuture<'a, SbResult<Vec<Message>>> {
334 Box::pin(async move {
335 let cmd = GetMessagesCmd {
336 header: Some(self.get_header()),
337 time_slice: Some(TimeSlice::SendDate(TimeRange {
338 start_point: start_date.map(|v| StartPoint::Start(v.and_utc().timestamp())),
339 end_point: end_date.map(|v| EndPoint::End(v.and_utc().timestamp())),
340 })),
341 maybe_application: Some(MaybeApplication::Application(application)),
342 limit: limit.unwrap_or(-1),
343 };
344 self.write_crypto(cmd).await?;
345 let m: MessageResponse = self.read_crypto().await?;
346 m.try_into()
347 })
348 }
349
350 fn get_messages_recieve_date<'a>(
351 &'a mut self,
352 application: String,
353 limit: Option<i32>,
354 start_date: Option<NaiveDateTime>,
355 end_date: Option<NaiveDateTime>,
356 ) -> DartFuture<'a, SbResult<Vec<Message>>> {
357 Box::pin(async move {
358 let cmd = GetMessagesCmd {
359 header: Some(self.get_header()),
360 time_slice: Some(TimeSlice::SendDate(TimeRange {
361 start_point: start_date.map(|v| StartPoint::Start(v.and_utc().timestamp())),
362 end_point: end_date.map(|v| EndPoint::End(v.and_utc().timestamp())),
363 })),
364 maybe_application: Some(MaybeApplication::Application(application)),
365 limit: limit.unwrap_or(-1),
366 };
367 self.write_crypto(cmd).await?;
368 let m: MessageResponse = self.read_crypto().await?;
369 m.try_into()
370 })
371 }
372}
373
374pub trait SessionTrait {
375 fn get_identity<'a>(&'a mut self, id: Option<Uuid>) -> DartFuture<'a, SbResult<Vec<Identity>>>;
376 #[cfg(feature = "flutter")]
377
378 fn set_on_connect(
379 &mut self,
380 on_connect: Box<
381 dyn Fn(Option<crate::api::api::SbSession>) -> flutter_rust_bridge::DartFnFuture<()>
382 + Send
383 + Sync
384 + 'static,
385 >,
386 );
387
388 #[cfg(feature = "flutter")]
389 fn on_connect<'a>(
390 &'a self,
391 ) -> Option<
392 &'a Box<
393 dyn Fn(Option<crate::api::api::SbSession>) -> flutter_rust_bridge::DartFnFuture<()>
394 + Send
395 + Sync
396 + 'static,
397 >,
398 >;
399
400 fn get_events<'a>(
401 &'a mut self,
402 block: bool,
403 count: Option<u32>,
404 ) -> DartFuture<'a, SbResult<Vec<SbEvent>>>;
405
406 fn get_messages<'a>(
407 &'a mut self,
408 application: String,
409 limit: Option<i32>,
410 ) -> DartFuture<'a, SbResult<Vec<Message>>>;
411
412 fn send_messages<'a>(
413 &'a mut self,
414 messages: Vec<Message>,
415 sign_identity: Option<Uuid>,
416 ) -> DartFuture<'a, SbResult<()>>;
417
418 fn initiate_identity_import<'a>(
419 &'a mut self,
420 id: Option<Uuid>,
421 ) -> DartFuture<'a, SbResult<ImportIdentityState>>;
422
423 fn get_messages_send_date<'a>(
424 &'a mut self,
425 application: String,
426 limit: Option<i32>,
427 start_date: Option<NaiveDateTime>,
428 end_date: Option<NaiveDateTime>,
429 ) -> DartFuture<'a, SbResult<Vec<Message>>>;
430
431 fn get_messages_recieve_date<'a>(
432 &'a mut self,
433 application: String,
434 limit: Option<i32>,
435 start_date: Option<NaiveDateTime>,
436 end_date: Option<NaiveDateTime>,
437 ) -> DartFuture<'a, SbResult<Vec<Message>>>;
438
439 fn is_closed<'a>(&'a mut self) -> DartFuture<'a, SbResult<bool>>;
440}