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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455

use tokio::sync::mpsc;
use tokio::sync::oneshot;
use std::sync::{Arc, atomic::{AtomicBool, Ordering}};

use crate::error::private::{
	connection_aborted,
	InnerError,
	UnexpectedMessageType,
};
use crate::peer::Command;
use crate::{Error, Message};

pub(crate) enum RequestHandleCommand<Body> {
	Close,
	Message(Message<Body>),
}

/// A handle for a sent request.
///
/// The handle can be used to receive updates and the response from the remote peer,
/// and to send update messages to the remote peer.
pub struct SentRequestHandle<Body> {
	write_handle: SentRequestWriteHandle<Body>,
	incoming_rx: mpsc::UnboundedReceiver<RequestHandleCommand<Body>>,
	peek_buffer: Option<Message<Body>>,
}

/// A write handle for a sent request.
///
/// Unlike [`SentRequestHandle`], write handles can be cloned.
/// But unlike regular handles, they can not be used to receive updates or the response from the remote peer.
///
/// Write handles can be used to send updates even when the regular handle is mutably borrowed.
///
/// You can get more write handles using [`SentRequestHandle::write_handle()`] or by cloning an existing one.
pub struct SentRequestWriteHandle<Body> {
	request_id: u32,
	service_id: i32,
	closed: Arc<AtomicBool>,
	command_tx: mpsc::UnboundedSender<Command<Body>>,
}

/// A handle for a received request.
///
/// The handle can be used to receive updates from the remote peer,
/// and to send updates and the response to the remote peer.
pub struct ReceivedRequestHandle<Body> {
	write_handle: ReceivedRequestWriteHandle<Body>,
	incoming_rx: mpsc::UnboundedReceiver<RequestHandleCommand<Body>>,
}

/// A write handle for a received request.
///
/// Unlike [`ReceivedRequestHandle`], write handles can be cloned.
/// But unlike regular handles, they can not be used to receive updates or the response from the remote peer.
///
/// Write handles can be used to send updates even when the regular handle is mutably borrowed.
///
/// You can get more write handles using [`ReceivedRequestHandle::write_handle()`] or by cloning an existing one.
pub struct ReceivedRequestWriteHandle<Body> {
	request_id: u32,
	service_id: i32,
	closed: Arc<AtomicBool>,
	command_tx: mpsc::UnboundedSender<Command<Body>>,
}

/// An incoming request or stream message.
pub enum ReceivedMessage<Body> {
	/// An incoming request.
	Request(ReceivedRequestHandle<Body>, Body),

	/// An incoming stream message.
	Stream(Message<Body>),
}

impl<Body> SentRequestHandle<Body> {
	/// Create a new sent request.
	pub(crate) fn new(
		request_id: u32,
		service_id: i32,
		closed: Arc<AtomicBool>,
		incoming_rx: mpsc::UnboundedReceiver<RequestHandleCommand<Body>>,
		command_tx: mpsc::UnboundedSender<Command<Body>>,
	) -> Self {
		let write_handle = SentRequestWriteHandle {
			request_id,
			service_id,
			closed,
			command_tx,
		};
		Self {
			write_handle,
			incoming_rx,
			peek_buffer: None,
		}
	}

	/// Get the request ID of the sent request.
	pub fn request_id(&self) -> u32 {
		self.write_handle.request_id()
	}

	/// Get the service ID of the initial request message.
	pub fn service_id(&self) -> i32 {
		self.write_handle.service_id()
	}

	/// Create a write handle for this request.
	///
	/// The write handle can be cloned and used even while this handle is mutably borrowed.
	pub fn write_handle(&self) -> SentRequestWriteHandle<Body> {
		self.write_handle.clone()
	}

	/// Receive the next update message of the request from the remote peer.
	///
	/// This function returns `None` if the final response is received instead of an update message.
	/// If that happens, the response message can be read using [`Self::recv_response`].
	pub async fn recv_update(&mut self) -> Option<Message<Body>> {
		let message = self.recv_message().await?;
		if message.header.message_type.is_responder_update() {
			Some(message)
		} else {
			self.peek_buffer = Some(message);
			None
		}
	}

	/// Receive the final response of the request from the remote peer.
	///
	/// This function returns an error if the received message is an update message.
	/// You can detect this situation using [`Error::is_unexpected_message_type()`].
	/// Afterwards, the update message can be read using [`Self::recv_update`].
	/// To ensure that there are no update messages left, keep calling [`Self::recv_update`] untill it returns `Ok(None)`.
	pub async fn recv_response(&mut self) -> Result<Message<Body>, Error> {
		let message = self.recv_message()
			.await
			.ok_or_else(connection_aborted)?;
		let kind = message.header.message_type;
		if kind.is_response() {
			Ok(message)
		} else {
			self.peek_buffer = Some(message);
			Err(
				InnerError::from(
					UnexpectedMessageType {
						value: kind,
						expected: crate::MessageType::Response,
					}
				).into()
			)
		}
	}

	/// Receive the next message of the request from the remote peer.
	///
	/// This could be an update message or a response message.
	async fn recv_message(&mut self) -> Option<Message<Body>> {
		if let Some(message) = self.peek_buffer.take() {
			Some(message)
		} else {
			match self.incoming_rx.recv().await? {
				RequestHandleCommand::Message(message) => {
					// Close the channel when reading a response message.
					if message.header.message_type.is_response() {
						self.incoming_rx.close();
					}
					Some(message)
				},
				// Close the channel when instructed to do so.
				// This is sent by the request tracker when unregistering the request.
				RequestHandleCommand::Close => {
					self.incoming_rx.close();
					None
				},
			}
		}
	}

	/// Send an update for the request to the remote peer.
	pub async fn send_update(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		self.write_handle.send_update(service_id, body).await
	}
}

impl<Body> SentRequestWriteHandle<Body> {
	/// Get the request ID of the sent request.
	pub fn request_id(&self) -> u32 {
		self.request_id
	}

	/// Get the service ID of the initial request message.
	pub fn service_id(&self) -> i32 {
		self.service_id
	}

	/// Send an update for the request to the remote peer.
	pub async fn send_update(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		use crate::peer::SendRawMessage;

		// If the response has already arrived, we're not allowed to send messages anymore.
		// The request ID could have been re-used already.
		if self.closed.load(Ordering::Acquire) {
			return Err(InnerError::RequestClosed.into())
		}

		let body = body.into();
		let (result_tx, result_rx) = oneshot::channel();
		let message = Message::requester_update(self.request_id, service_id, body);
		self.command_tx
			.send(SendRawMessage { message, result_tx }.into())
			.map_err(|_| connection_aborted())?;
		result_rx.await.map_err(|_| connection_aborted())??;
		Ok(())
	}
}

impl<Body> ReceivedRequestHandle<Body> {
	/// Create a new received request.
	pub(crate) fn new(
		request_id: u32,
		service_id: i32,
		closed: Arc<AtomicBool>,
		incoming_rx: mpsc::UnboundedReceiver<RequestHandleCommand<Body>>,
		command_tx: mpsc::UnboundedSender<Command<Body>>,
	) -> Self {
		let write_handle = ReceivedRequestWriteHandle {
			request_id,
			service_id,
			closed,
			command_tx,
		};
		Self {
			write_handle,
			incoming_rx,
		}
	}

	/// Get the request ID of the received request.
	pub fn request_id(&self) -> u32 {
		self.write_handle.request_id()
	}

	/// Get the service ID of the received request message.
	pub fn service_id(&self) -> i32 {
		self.write_handle.service_id()
	}

	/// Create a write handle for this request.
	///
	/// The write handle can be cloned and used even while this handle is mutably borrowed.
	pub fn write_handle(&self) -> ReceivedRequestWriteHandle<Body> {
		self.write_handle.clone()
	}

	/// Receive the next update message of the request from the remote peer.
	pub async fn recv_update(&mut self) -> Option<Message<Body>> {
		match self.incoming_rx.recv().await? {
			RequestHandleCommand::Message(x) => Some(x),
			// Close the channel when instructed to do so.
			// This is sent by the request tracker when unregistering the request.
			RequestHandleCommand::Close => {
				self.incoming_rx.close();
				None
			},
		}
	}

	/// Send an update for the request to the remote peer.
	pub async fn send_update(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		self.write_handle.send_update(service_id, body).await
	}

	/// Send the final response for the request to the remote peer.
	pub async fn send_response(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		self.write_handle.send_response(service_id, body).await
	}

	/// Send the final response with an error message.
	pub async fn send_error_response(&self, message: &str) -> Result<(), Error>
	where
		Body: crate::Body,
	{
		self.write_handle.send_error_response(message).await
	}
}

impl<Body> ReceivedRequestWriteHandle<Body> {
	/// Get the request ID of the sent request.
	pub fn request_id(&self) -> u32 {
		self.request_id
	}

	/// Get the service ID of the initial request message.
	pub fn service_id(&self) -> i32 {
		self.service_id
	}

	/// Send an update for the request to the remote peer.
	pub async fn send_update(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		let body = body.into();
		self.send_raw_message(Message::responder_update(self.request_id, service_id, body)).await
	}

	/// Send the final response for the request to the remote peer.
	pub async fn send_response(&self, service_id: i32, body: impl Into<Body>) -> Result<(), Error> {
		let body = body.into();
		self.send_raw_message(Message::response(self.request_id, service_id, body)).await
	}

	/// Send the final response with an error message.
	pub async fn send_error_response(&self, message: &str) -> Result<(), Error>
	where
		Body: crate::Body,
	{
		self.send_raw_message(Message::error_response(self.request_id, message)).await
	}

	/// Send a raw message.
	async fn send_raw_message(&self, message: Message<Body>) -> Result<(), Error> {
		use crate::peer::SendRawMessage;

		// If the response has already arrived, we're not allowed to send messages anymore.
		// The request ID could have been re-used already.
		if self.closed.load(Ordering::Acquire) {
			return Err(InnerError::RequestClosed.into())
		}

		let (result_tx, result_rx) = oneshot::channel();
		self.command_tx
			.send(SendRawMessage { message, result_tx }.into())
			.map_err(|_| connection_aborted())?;
		result_rx.await.map_err(|_| connection_aborted())??;
		Ok(())
	}
}

impl<Body> std::fmt::Debug for SentRequestHandle<Body> {
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
		f.debug_struct("SentRequestHandle")
			.field("request_id", &self.request_id())
			.field("service_id", &self.service_id())
			.finish_non_exhaustive()
	}
}

impl<Body> std::fmt::Debug for SentRequestWriteHandle<Body> {
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
		f.debug_struct("SentRequestWriteHandle")
			.field("request_id", &self.request_id())
			.field("service_id", &self.service_id())
			.finish_non_exhaustive()
	}
}

impl<Body> std::fmt::Debug for ReceivedRequestHandle<Body> {
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
		f.debug_struct("ReceivedRequestHandle")
			.field("request_id", &self.request_id())
			.field("service_id", &self.service_id())
			.finish_non_exhaustive()
	}
}

impl<Body> std::fmt::Debug for ReceivedRequestWriteHandle<Body> {
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
		f.debug_struct("ReceivedRequestWriteHandle")
			.field("request_id", &self.request_id())
			.field("service_id", &self.service_id())
			.finish_non_exhaustive()
	}
}

impl<Body> std::fmt::Debug for ReceivedMessage<Body> {
	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
		match self {
			Self::Request(x, _body) => write!(f, "Request({:?})", x),
			Self::Stream(x) => write!(f, "Stream({:?})", x),
		}
	}
}

impl<Body> Clone for SentRequestWriteHandle<Body> {
	fn clone(&self) -> Self {
		Self {
			request_id: self.request_id,
			service_id: self.service_id,
			closed: self.closed.clone(),
			command_tx: self.command_tx.clone(),
		}
	}
}

impl<Body> Clone for ReceivedRequestWriteHandle<Body> {
	fn clone(&self) -> Self {
		Self {
			request_id: self.request_id,
			service_id: self.service_id,
			closed: self.closed.clone(),
			command_tx: self.command_tx.clone(),
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;
	use crate::{Peer, UnixStreamTransport};
	use tokio::net::UnixStream;
	use assert2::{assert, let_assert};

	/// Test that request handles can not be used for sending messages after they are closed.
	///
	/// They should be closed automatically by the request tracker when the response is sent or received.
	#[tokio::test]
	async fn test_closed_after_respone() {
		let_assert!(Ok((peer_a, peer_b)) = UnixStream::pair());

		let (peer_a, handle_a) = Peer::new(UnixStreamTransport::new(peer_a, Default::default()));
		let (peer_b, mut handle_b) = Peer::new(UnixStreamTransport::new(peer_b, Default::default()));

		let task_a = tokio::spawn(peer_a.run());
		let task_b = tokio::spawn(peer_b.run());

		// Send a request from A.
		let_assert!(Ok(mut sent_request) = handle_a.send_request(1, &[2][..]).await);

		// Receive the request on B.
		let_assert!(Ok(ReceivedMessage::Request(mut received_request, _body)) = handle_b.recv_message().await);

		// Check that sending requests works.
		assert!(let Ok(()) = sent_request.send_update(1, vec![]).await);
		assert!(let Some(_) = received_request.recv_update().await);

		assert!(let Ok(()) = received_request.send_update(1, vec![]).await);
		assert!(let Some(_) = sent_request.recv_update().await);

		// Now we send and receive a response.
		// After that, sending responses should be impossible.
		assert!(let Ok(()) = received_request.send_response(1, vec![]).await);
		assert!(let Err(_) = received_request.send_update(1, vec![]).await);
		assert!(let Err(_) = received_request.send_response(1, vec![]).await);

		assert!(let Ok(_) = sent_request.recv_response().await);
		assert!(let Err(_) = sent_request.send_update(1, vec![]).await);

		drop(handle_a);
		drop(handle_b);
		drop(sent_request);
		drop(received_request);

		assert!(let Ok(()) = task_a.await);
		assert!(let Ok(()) = task_b.await);
	}
}