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
use crate::{
    syscall::{
        self,
        CreateChannelError,
        GetMessageError,
        RegisterServiceError,
        SendMessageError,
        SubscribeToServiceError,
        CHANNEL_MAX_NUM_HANDLES,
    },
    Handle,
};
use core::{future::Future, marker::PhantomData, mem, task::Poll};
use ptah::{DeserializeOwned, Serialize};

#[derive(Debug)]
pub enum ChannelSendError {
    FailedToSerialize(ptah::ser::Error),
    SendError(SendMessageError),
}

#[derive(Debug)]
pub enum ChannelReceiveError {
    FailedToDeserialize(ptah::de::Error),
    ReceiveError(GetMessageError),
}

pub struct Channel<S, R>(Handle, PhantomData<(S, R)>)
where
    S: Serialize + DeserializeOwned,
    R: Serialize + DeserializeOwned;

impl<S, R> Channel<S, R>
where
    S: Serialize + DeserializeOwned,
    R: Serialize + DeserializeOwned,
{
    pub fn new_from_handle(handle: Handle) -> Channel<S, R> {
        Channel(handle, PhantomData)
    }

    /// Create a new channel. Returns one end as a `Channel`, and a `Handle` for the other end.
    /// Generally, the handle
    pub fn create() -> Result<(Channel<S, R>, Handle), CreateChannelError> {
        let (this_end, other_end) = syscall::create_channel()?;
        Ok((Self::new_from_handle(this_end), other_end))
    }

    pub fn subscribe_to_service(name: &str) -> Result<Channel<S, R>, SubscribeToServiceError> {
        let handle = syscall::subscribe_to_service(name)?;
        Ok(Self::new_from_handle(handle))
    }

    pub fn send(&self, message: &S) -> Result<(), ChannelSendError> {
        let mut writer = ChannelWriter::new();
        ptah::to_wire(message, &mut writer).map_err(|err| ChannelSendError::FailedToSerialize(err))?;
        syscall::send_message(self.0, writer.bytes(), writer.handles())
            .map_err(|err| ChannelSendError::SendError(err))
    }

    /// Receive a message from the channel, if there's one waiting. Returns `Ok(None)` if there are no pending
    /// messages to be received.
    pub fn try_receive(&self) -> Result<Option<R>, ChannelReceiveError> {
        let mut byte_buffer = [0u8; BYTES_BUFFER_SIZE];
        let mut handle_buffer = [Handle::ZERO; CHANNEL_MAX_NUM_HANDLES];

        match syscall::get_message(self.0, &mut byte_buffer, &mut handle_buffer) {
            Ok((bytes, handles)) => {
                // TODO: this looks really bad, but is actually fine (since Handle is just a transparent wrapper
                // around a `u32`). There might be a better way.
                let ptah_handles: &[u32] = unsafe { mem::transmute(handles) };

                let message: R = ptah::from_wire(bytes, ptah_handles)
                    .map_err(|err| ChannelReceiveError::FailedToDeserialize(err))?;
                Ok(Some(message))
            }
            Err(GetMessageError::NoMessage) => Ok(None),
            Err(err) => Err(ChannelReceiveError::ReceiveError(err)),
        }
    }

    pub fn receive(&self) -> impl Future<Output = Result<R, ChannelReceiveError>> + '_ {
        core::future::poll_fn(|context| {
            let mut byte_buffer = [0u8; BYTES_BUFFER_SIZE];
            let mut handle_buffer = [Handle::ZERO; CHANNEL_MAX_NUM_HANDLES];

            match syscall::get_message(self.0, &mut byte_buffer, &mut handle_buffer) {
                Ok((bytes, handles)) => {
                    // TODO: this looks really bad, but is actually fine (since Handle is just a transparent wrapper
                    // around a `u32`). There might be a better way.
                    let ptah_handles: &[u32] = unsafe { mem::transmute(handles) };

                    let message: R = ptah::from_wire(bytes, ptah_handles)
                        .map_err(|err| ChannelReceiveError::FailedToDeserialize(err))?;
                    Poll::Ready(Ok(message))
                }
                Err(GetMessageError::NoMessage) => {
                    crate::rt::RUNTIME.get().reactor.lock().register(self.0, context.waker().clone());
                    Poll::Pending
                }
                Err(err) => Poll::Ready(Err(ChannelReceiveError::ReceiveError(err))),
            }
        })
    }
}

impl Channel<!, Handle> {
    pub fn register_service(name: &str) -> Result<Channel<!, Handle>, RegisterServiceError> {
        Ok(Self::new_from_handle(syscall::register_service(name)?))
    }
}

const BYTES_BUFFER_SIZE: usize = 512;

struct ChannelWriter {
    byte_buffer: [u8; BYTES_BUFFER_SIZE],
    handle_buffer: [Handle; CHANNEL_MAX_NUM_HANDLES],
    num_bytes: usize,
    num_handles: u8,
}

impl ChannelWriter {
    pub fn new() -> ChannelWriter {
        ChannelWriter {
            byte_buffer: [0u8; BYTES_BUFFER_SIZE],
            handle_buffer: [Handle::ZERO; CHANNEL_MAX_NUM_HANDLES],
            num_bytes: 0,
            num_handles: 0,
        }
    }

    pub fn bytes(&self) -> &[u8] {
        &self.byte_buffer[0..self.num_bytes]
    }

    pub fn handles(&self) -> &[Handle] {
        &self.handle_buffer[0..(self.num_handles as usize)]
    }
}

impl<'a> ptah::Writer for &'a mut ChannelWriter {
    fn write(&mut self, buf: &[u8]) -> ptah::ser::Result<()> {
        /*
         * Detect if the write will overflow the buffer.
         */
        if (self.num_bytes + buf.len()) > BYTES_BUFFER_SIZE {
            return Err(ptah::ser::Error::WriterFullOfBytes);
        }

        self.byte_buffer[self.num_bytes..(self.num_bytes + buf.len())].copy_from_slice(buf);
        self.num_bytes += buf.len();
        Ok(())
    }

    fn push_handle(&mut self, handle: ptah::Handle) -> ptah::ser::Result<ptah::HandleSlot> {
        /*
         * Check if we're full of handles yet.
         */
        if (self.num_handles as usize + 1) > CHANNEL_MAX_NUM_HANDLES {
            return Err(ptah::ser::Error::WriterFullOfHandles);
        }

        self.handle_buffer[self.num_handles as usize] = Handle(handle);

        let slot = ptah::make_handle_slot(self.num_handles);
        self.num_handles += 1;
        Ok(slot)
    }
}