1#![recursion_limit = "256"]
7
8use perfetto_protos::perfetto::protos::{
9 DisableTracingRequest, EnableTracingRequest, FreeBuffersRequest, GetAsyncCommandRequest,
10 GetAsyncCommandResponse, InitializeConnectionRequest, InitializeConnectionResponse, IpcFrame,
11 ReadBuffersRequest, RegisterDataSourceRequest, RegisterDataSourceResponse, ipc_frame,
12};
13use prost::Message;
14use starnix_core::task::{CurrentTask, EventHandler, Waiter};
15use starnix_core::vfs::buffers::{VecInputBuffer, VecOutputBuffer};
16use starnix_core::vfs::socket::{
17 SocketDomain, SocketFile, SocketPeer, SocketProtocol, SocketType, resolve_unix_socket_address,
18};
19use starnix_core::vfs::{FileHandle, FsStr};
20
21use starnix_uapi::errors::Errno;
22use starnix_uapi::open_flags::OpenFlags;
23use starnix_uapi::vfs::FdEvents;
24use std::collections::VecDeque;
25use thiserror::Error;
26
27pub struct IpcConnection {
33 file: FileHandle,
35 request_id: u64,
38}
39
40#[derive(Error, Debug)]
41pub enum IpcWriteError {
42 #[error(transparent)]
43 Encode(#[from] prost::EncodeError),
44 #[error(transparent)]
45 Write(#[from] Errno),
46 #[error("TooLong: {0} exceeds max for u32")]
47 TooLong(usize),
48}
49
50#[derive(Error, Debug)]
51pub enum IpcReadError {
52 #[error(transparent)]
53 Decode(#[from] prost::DecodeError),
54 #[error(transparent)]
55 Read(#[from] Errno),
56}
57
58#[derive(Error, Debug)]
59pub enum InvokeMethodError {
60 #[error("could not not look up method name: {0}")]
61 InvalidMethod(String),
62 #[error(transparent)]
63 IpcWrite(#[from] IpcWriteError),
64 #[error(transparent)]
65 IpcRead(#[from] IpcReadError),
66 #[error("unexpected response: {0}")]
67 InvalidResponse(String),
68}
69
70#[derive(Error, Debug)]
71pub enum ProducerError {
72 #[error(transparent)]
73 InvokeMethod(#[from] InvokeMethodError),
74 #[error(transparent)]
75 IpcWrite(#[from] IpcWriteError),
76 #[error(transparent)]
77 IpcRead(#[from] IpcReadError),
78 #[error("unexpected response: {0}")]
79 InvalidResponse(String),
80 #[error(transparent)]
81 Decode(#[from] prost::DecodeError),
82}
83
84impl IpcConnection {
85 pub fn new(file: FileHandle) -> Self {
86 Self { file, request_id: 0 }
87 }
88
89 pub fn bind_service(
90 &mut self,
91 service_name: &str,
92 current_task: &CurrentTask,
93 ) -> Result<(), IpcWriteError> {
94 let bind_service_message = IpcFrame {
97 request_id: Some(1),
98 msg: Some(ipc_frame::Msg::MsgBindService(ipc_frame::BindService {
99 service_name: Some(service_name.to_string()),
100 })),
101 ..Default::default()
102 };
103 self.write_frame(bind_service_message, current_task)
104 }
105
106 pub fn invoke_method(
107 &mut self,
108 service_id: u32,
109 method_id: u32,
110 arguments: Option<Vec<u8>>,
111 current_task: &CurrentTask,
112 ) -> Result<(), IpcWriteError> {
113 let msg = IpcFrame {
114 request_id: Some(self.allocate_request_id()),
115 msg: Some(ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
116 service_id: Some(service_id),
117 method_id: Some(method_id),
118 args_proto: arguments,
119 drop_reply: None,
120 })),
121 ..Default::default()
122 };
123 self.write_frame(msg, current_task)
124 }
125
126 fn write_frame(
127 &mut self,
128 frame: IpcFrame,
129 current_task: &CurrentTask,
130 ) -> Result<(), IpcWriteError> {
131 let frame_len = u32::try_from(frame.encoded_len())
133 .map_err(|_| IpcWriteError::TooLong(frame.encoded_len()))?;
134 let mut bind_service_bytes =
135 Vec::with_capacity(frame.encoded_len() + std::mem::size_of::<u32>());
136 bind_service_bytes.extend_from_slice(&frame_len.to_le_bytes());
137 frame.encode(&mut bind_service_bytes)?;
138 let mut bind_service_buffer: VecInputBuffer = bind_service_bytes.into();
139 self.file.write(current_task, &mut bind_service_buffer)?;
140 Ok(())
141 }
142
143 fn allocate_request_id(&mut self) -> u64 {
145 let id = self.request_id;
146 self.request_id += 1;
147 id
148 }
149}
150
151pub struct FrameReader {
157 file: FileHandle,
159 read_buffer: VecOutputBuffer,
164 data: VecDeque<u8>,
166 next_message_size: Option<usize>,
169}
170
171impl FrameReader {
172 pub fn new(file: FileHandle) -> Self {
173 Self {
174 file,
175 read_buffer: VecOutputBuffer::new(4096),
176 data: VecDeque::with_capacity(4096),
177 next_message_size: None,
178 }
179 }
180
181 pub fn next_frame_blocking(
183 &mut self,
184 current_task: &CurrentTask,
185 ) -> Result<IpcFrame, IpcReadError> {
186 loop {
187 if self.next_message_size.is_none() && self.data.len() >= 4 {
188 let len_bytes: [u8; 4] = self
189 .data
190 .drain(..4)
191 .collect::<Vec<_>>()
192 .try_into()
193 .expect("self.data has at least 4 elements");
194 self.next_message_size = Some(u32::from_le_bytes(len_bytes) as usize);
195 }
196 if let Some(message_size) = self.next_message_size {
197 if self.data.len() >= message_size {
198 let message: Vec<u8> = self.data.drain(..message_size).collect();
199 self.next_message_size = None;
200 return Ok(IpcFrame::decode(message.as_slice())?);
201 }
202 }
203
204 let waiter = Waiter::new();
205 self.file.wait_async(current_task, &waiter, FdEvents::POLLIN, EventHandler::None);
206 while self.file.query_events(current_task)? & FdEvents::POLLIN != FdEvents::POLLIN {
207 waiter.wait(current_task)?;
208 }
209 self.file.read(current_task, &mut self.read_buffer)?;
210 self.data.extend(self.read_buffer.data());
211 self.read_buffer.reset();
212 }
213 }
214}
215
216pub struct Consumer {
218 conn_file: FileHandle,
221 frame_reader: FrameReader,
223 bind_service_reply: ipc_frame::BindServiceReply,
226 request_id: u64,
228}
229
230impl Consumer {
231 pub fn new(current_task: &CurrentTask, socket_path: &FsStr) -> Result<Self, anyhow::Error> {
234 let conn_file = SocketFile::new_socket(
235 current_task,
236 SocketDomain::Unix,
237 SocketType::Stream,
238 OpenFlags::RDWR,
239 SocketProtocol::from_raw(0),
240 false,
241 )?;
242 let conn = SocketFile::get_from_file(&conn_file)?;
243 let peer = SocketPeer::Handle(resolve_unix_socket_address(current_task, socket_path)?);
244 conn.connect(current_task, peer)?;
245 let mut frame_reader = FrameReader::new(conn_file.clone());
246 let mut request_id = 1;
247
248 let bind_service_message = IpcFrame {
249 request_id: Some(request_id),
250 data_for_testing: Vec::new(),
251 msg: Some(ipc_frame::Msg::MsgBindService(ipc_frame::BindService {
252 service_name: Some("ConsumerPort".to_string()),
253 })),
254 };
255 request_id += 1;
256 let mut bind_service_bytes =
257 Vec::with_capacity(bind_service_message.encoded_len() + std::mem::size_of::<u32>());
258 bind_service_bytes.extend_from_slice(
259 &u32::try_from(bind_service_message.encoded_len()).unwrap().to_le_bytes(),
260 );
261 bind_service_message.encode(&mut bind_service_bytes)?;
262 let mut bind_service_buffer: VecInputBuffer = bind_service_bytes.into();
263 conn.file().write(current_task, &mut bind_service_buffer)?;
264
265 let reply_frame = frame_reader.next_frame_blocking(current_task)?;
266
267 let bind_service_reply = match reply_frame.msg {
268 Some(ipc_frame::Msg::MsgBindServiceReply(reply)) => reply,
269 m => return Err(anyhow::anyhow!("Got unexpected reply message: {:?}", m)),
270 };
271
272 Ok(Self { conn_file, frame_reader, bind_service_reply, request_id })
273 }
274
275 fn send_message(
276 &mut self,
277 current_task: &CurrentTask,
278 msg: ipc_frame::Msg,
279 ) -> Result<u64, anyhow::Error> {
280 let request_id = self.request_id;
281 let frame =
282 IpcFrame { request_id: Some(request_id), data_for_testing: Vec::new(), msg: Some(msg) };
283
284 self.request_id += 1;
285
286 let mut frame_bytes = Vec::with_capacity(frame.encoded_len() + std::mem::size_of::<u32>());
287 frame_bytes.extend_from_slice(&u32::try_from(frame.encoded_len())?.to_le_bytes());
288 frame.encode(&mut frame_bytes)?;
289 let mut buffer: VecInputBuffer = frame_bytes.into();
290 self.conn_file.write(current_task, &mut buffer)?;
291
292 Ok(request_id)
293 }
294
295 fn method_id(&self, name: &str) -> Result<u32, anyhow::Error> {
296 for method in &self.bind_service_reply.methods {
297 if let Some(method_name) = method.name.as_ref() {
298 if method_name == name {
299 if let Some(id) = method.id {
300 return Ok(id);
301 } else {
302 return Err(anyhow::anyhow!(
303 "Matched method name {} but found no id",
304 method_name
305 ));
306 }
307 }
308 }
309 }
310 Err(anyhow::anyhow!("Did not find method {}", name))
311 }
312
313 pub fn enable_tracing(
314 &mut self,
315 current_task: &CurrentTask,
316 req: EnableTracingRequest,
317 ) -> Result<u64, anyhow::Error> {
318 let method_id = self.method_id("EnableTracing")?;
319 let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
320 req.encode(&mut encoded_args)?;
321
322 self.send_message(
323 current_task,
324 ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
325 service_id: self.bind_service_reply.service_id,
326 method_id: Some(method_id),
327 args_proto: Some(encoded_args),
328 drop_reply: None,
329 }),
330 )
331 }
332
333 pub fn disable_tracing(
334 &mut self,
335 current_task: &CurrentTask,
336 req: DisableTracingRequest,
337 ) -> Result<u64, anyhow::Error> {
338 let method_id = self.method_id("DisableTracing")?;
339 let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
340 req.encode(&mut encoded_args)?;
341
342 self.send_message(
343 current_task,
344 ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
345 service_id: self.bind_service_reply.service_id,
346 method_id: Some(method_id),
347 args_proto: Some(encoded_args),
348 drop_reply: None,
349 }),
350 )
351 }
352
353 pub fn read_buffers(
354 &mut self,
355 current_task: &CurrentTask,
356 req: ReadBuffersRequest,
357 ) -> Result<u64, anyhow::Error> {
358 let method_id = self.method_id("ReadBuffers")?;
359 let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
360 req.encode(&mut encoded_args)?;
361
362 self.send_message(
363 current_task,
364 ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
365 service_id: self.bind_service_reply.service_id,
366 method_id: Some(method_id),
367 args_proto: Some(encoded_args),
368 drop_reply: None,
369 }),
370 )
371 }
372
373 pub fn free_buffers(
374 &mut self,
375 current_task: &CurrentTask,
376 req: FreeBuffersRequest,
377 ) -> Result<u64, anyhow::Error> {
378 let method_id = self.method_id("FreeBuffers")?;
379 let mut encoded_args: Vec<u8> = Vec::with_capacity(req.encoded_len());
380 req.encode(&mut encoded_args)?;
381
382 self.send_message(
383 current_task,
384 ipc_frame::Msg::MsgInvokeMethod(ipc_frame::InvokeMethod {
385 service_id: self.bind_service_reply.service_id,
386 method_id: Some(method_id),
387 args_proto: Some(encoded_args),
388 drop_reply: None,
389 }),
390 )
391 }
392
393 pub fn next_frame_blocking(
394 &mut self,
395 current_task: &CurrentTask,
396 ) -> Result<IpcFrame, IpcReadError> {
397 self.frame_reader.next_frame_blocking(current_task)
398 }
399}
400
401pub struct Producer {
408 frame_reader: FrameReader,
410
411 ipc_connection: IpcConnection,
413
414 service_id: u32,
417
418 method_map: std::collections::HashMap<String, u32>,
421}
422
423impl Producer {
424 pub fn new(current_task: &CurrentTask, socket: FileHandle) -> Result<Self, ProducerError> {
427 let mut producer = Self {
428 frame_reader: FrameReader::new(socket.clone()),
429 ipc_connection: IpcConnection::new(socket),
430 service_id: 0,
431 method_map: std::collections::HashMap::new(),
432 };
433
434 producer.ipc_connection.bind_service("ProducerPort", current_task)?;
437
438 let reply_frame = producer.frame_reader.next_frame_blocking(current_task)?;
442
443 let ipc_frame::BindServiceReply { success, service_id, methods } = match reply_frame.msg {
444 Some(ipc_frame::Msg::MsgBindServiceReply(reply)) => reply,
445 m => {
446 return Err(ProducerError::InvalidResponse(format!(
447 "Got unexpected reply message: {:?}",
448 m
449 )));
450 }
451 };
452
453 if !success.unwrap_or(false) {
454 return Err(ProducerError::InvalidResponse("Bind to socket failed".into()));
455 }
456
457 producer.method_map = methods
459 .into_iter()
460 .flat_map(|ipc_frame::bind_service_reply::MethodInfo { id, name }| match (id, name) {
461 (Some(id), Some(name)) => Some((name, id)),
462 _ => None,
463 })
464 .collect();
465 if let Some(service_id) = service_id {
466 producer.service_id = service_id
467 } else {
468 return Err(ProducerError::InvalidResponse(
469 "BindServiceReply did not include service_id".into(),
470 ));
471 }
472
473 Ok(producer)
474 }
475
476 pub fn initialize_connection(
479 &mut self,
480 request: InitializeConnectionRequest,
481 current_task: &CurrentTask,
482 ) -> Result<InitializeConnectionResponse, ProducerError> {
483 let (Some(reply), has_more) = self.invoke_method(
484 "InitializeConnection",
485 Some(request.encode_to_vec()),
486 current_task,
487 )?
488 else {
489 return Err(ProducerError::InvalidResponse("expected a response but got none".into()));
490 };
491 if has_more {
492 return Err(ProducerError::InvalidResponse(
493 "InitializeConnection should not stream but got a streaming response".into(),
494 ));
495 }
496 Ok(InitializeConnectionResponse::decode(reply.as_ref())?)
497 }
498
499 pub fn register_data_source(
501 &mut self,
502 request: RegisterDataSourceRequest,
503 current_task: &CurrentTask,
504 ) -> Result<RegisterDataSourceResponse, ProducerError> {
505 let (Some(reply), has_more) =
506 self.invoke_method("RegisterDataSource", Some(request.encode_to_vec()), current_task)?
507 else {
508 return Err(ProducerError::InvalidResponse(
509 "RegisterDataSource expected a response but got none".into(),
510 ));
511 };
512 if has_more {
513 return Err(ProducerError::InvalidResponse(
514 "RegisterDataSource should not stream but got a streaming response".into(),
515 ));
516 }
517 Ok(RegisterDataSourceResponse::decode(reply.as_ref())?)
518 }
519
520 fn invoke_method(
522 &mut self,
523 method_name: &str,
524 arguments: Option<Vec<u8>>,
525 current_task: &CurrentTask,
526 ) -> Result<(Option<Vec<u8>>, bool), InvokeMethodError> {
527 self.invoke_method_inner(method_name, arguments, current_task)?;
528
529 let reply_frame = self.frame_reader.next_frame_blocking(current_task)?;
530
531 let ipc_frame::InvokeMethodReply { success, has_more, reply_proto } = match reply_frame.msg
532 {
533 Some(ipc_frame::Msg::MsgInvokeMethodReply(reply)) => reply,
534 m => {
535 return Err(InvokeMethodError::InvalidResponse(format!(
536 "unexpected reply message: {:?}",
537 m
538 )));
539 }
540 };
541 match success {
542 Some(true) => Ok((reply_proto, has_more.unwrap_or(false))),
543 _ => {
544 return Err(InvokeMethodError::InvalidResponse(format!(
545 "InvokeMethod Reply did not succeed. Reply: success: {:?}, has_more: {:?}, proto: {:?}",
546 success, has_more, reply_proto,
547 )));
548 }
549 }
550 }
551
552 fn invoke_method_inner(
554 &mut self,
555 method_name: &str,
556 arguments: Option<Vec<u8>>,
557 current_task: &CurrentTask,
558 ) -> Result<(), InvokeMethodError> {
559 let Some(method_id) = self.method_map.get(method_name).copied() else {
560 return Err(InvokeMethodError::InvalidMethod(method_name.into()));
561 };
562 self.ipc_connection.invoke_method(self.service_id, method_id, arguments, current_task)?;
563 Ok(())
564 }
565
566 pub fn get_command_request(&mut self, current_task: &CurrentTask) -> Result<(), ProducerError> {
569 Ok(self.invoke_method_inner(
570 "GetAsyncCommand",
571 Some(GetAsyncCommandRequest {}.encode_to_vec()),
572 current_task,
573 )?)
574 }
575
576 pub fn get_command_response(
578 &mut self,
579 current_task: &CurrentTask,
580 ) -> Result<(Option<GetAsyncCommandResponse>, bool), ProducerError> {
581 let reply_frame = self.frame_reader.next_frame_blocking(current_task)?;
582
583 let ipc_frame::InvokeMethodReply { success, has_more, reply_proto } = match reply_frame.msg
584 {
585 Some(ipc_frame::Msg::MsgInvokeMethodReply(reply)) => reply,
586 m => {
587 return Err(ProducerError::InvalidResponse(format!(
588 "Got unexpected reply message: {:?}",
589 m
590 )));
591 }
592 };
593 if !success.unwrap_or(false) {
594 return Err(ProducerError::InvalidResponse(format!(
595 "InvokeMethod Reply did not include success. Reply: success: {:?}, has_more: {:?}, proto: {:?}",
596 success, has_more, reply_proto
597 )));
598 }
599
600 let Some(reply_proto) = reply_proto else {
601 return Err(ProducerError::InvalidResponse(
602 "InvokeMethod reply didn't include a proto".into(),
603 ));
604 };
605 Ok((
606 Some(GetAsyncCommandResponse::decode(reply_proto.as_ref())?),
607 has_more.unwrap_or(false),
608 ))
609 }
610}