crosvm: add msg_sock.

MsgSock wraps UnixDatagram and provides simple macro to define Messages
that could be send through sock easily.

TEST=cargo test
BUG=None

Change-Id: I296fabc41893ad6a3ec42ef82dd29c3b752be8b8
Reviewed-on: https://chromium-review.googlesource.com/1255548
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Jingkui Wang <jkwang@google.com>
Reviewed-by: Zach Reizner <zachr@chromium.org>

1 files changed, 288 insertions, 0 deletions

diff --git a/msg_socket/src/lib.rs b/msg_socket/src/lib.rs
new file mode 100644
index 0000000..f868b19
--- /dev/null
+++ b/msg_socket/src/lib.rs
@@ -0,0 +1,288 @@
+// Copyright 2018 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#[allow(unused_imports)]
+#[macro_use]
+extern crate msg_on_socket_derive;
+extern crate data_model;
+extern crate sys_util;
+
+mod msg_on_socket;
+
+use std::marker::PhantomData;
+use std::os::unix::io::RawFd;
+use std::os::unix::net::UnixDatagram;
+use sys_util::{ScmSocket, UnlinkUnixDatagram};
+
+pub use msg_on_socket::*;
+pub use msg_on_socket_derive::*;
+
+/// Create a pair of socket. Request is send in one direction while response is in the other
+/// direction.
+pub fn pair<Request: MsgOnSocket, Response: MsgOnSocket>(
+) -> Option<(MsgSocket<Request, Response>, MsgSocket<Response, Request>)> {
+    let (sock1, sock2) = match UnixDatagram::pair() {
+        Ok((sock1, sock2)) => (sock1, sock2),
+        _ => {
+            return None;
+        }
+    };
+    let requester = MsgSocket {
+        sock: sock1,
+        _i: PhantomData,
+        _o: PhantomData,
+    };
+    let responder = MsgSocket {
+        sock: sock2,
+        _i: PhantomData,
+        _o: PhantomData,
+    };
+    Some((requester, responder))
+}
+
+/// Bidirection sock that support both send and recv.
+pub struct MsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
+    sock: UnixDatagram,
+    _i: PhantomData<I>,
+    _o: PhantomData<O>,
+}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> MsgSocket<I, O> {
+    // Create a new MsgSocket.
+    pub fn new(s: UnixDatagram) -> MsgSocket<I, O> {
+        MsgSocket {
+            sock: s,
+            _i: PhantomData,
+            _o: PhantomData,
+        }
+    }
+}
+
+/// Bidirection sock that support both send and recv.
+pub struct UnlinkMsgSocket<I: MsgOnSocket, O: MsgOnSocket> {
+    sock: UnlinkUnixDatagram,
+    _i: PhantomData<I>,
+    _o: PhantomData<O>,
+}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> UnlinkMsgSocket<I, O> {
+    // Create a new MsgSocket.
+    pub fn new(s: UnlinkUnixDatagram) -> UnlinkMsgSocket<I, O> {
+        UnlinkMsgSocket {
+            sock: s,
+            _i: PhantomData,
+            _o: PhantomData,
+        }
+    }
+}
+
+/// One direction socket that only supports sending.
+pub struct Sender<M: MsgOnSocket> {
+    sock: UnixDatagram,
+    _m: PhantomData<M>,
+}
+
+impl<M: MsgOnSocket> Sender<M> {
+    /// Create a new sender sock.
+    pub fn new(s: UnixDatagram) -> Sender<M> {
+        Sender {
+            sock: s,
+            _m: PhantomData,
+        }
+    }
+}
+
+/// One direction socket that only supports receiving.
+pub struct Receiver<M: MsgOnSocket> {
+    sock: UnixDatagram,
+    _m: PhantomData<M>,
+}
+
+impl<M: MsgOnSocket> Receiver<M> {
+    /// Create a new receiver sock.
+    pub fn new(s: UnixDatagram) -> Receiver<M> {
+        Receiver {
+            sock: s,
+            _m: PhantomData,
+        }
+    }
+}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for MsgSocket<I, O> {
+    fn as_ref(&self) -> &UnixDatagram {
+        &self.sock
+    }
+}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> AsRef<UnixDatagram> for UnlinkMsgSocket<I, O> {
+    fn as_ref(&self) -> &UnixDatagram {
+        self.sock.as_ref()
+    }
+}
+
+impl<M: MsgOnSocket> AsRef<UnixDatagram> for Sender<M> {
+    fn as_ref(&self) -> &UnixDatagram {
+        &self.sock
+    }
+}
+
+impl<M: MsgOnSocket> AsRef<UnixDatagram> for Receiver<M> {
+    fn as_ref(&self) -> &UnixDatagram {
+        &self.sock
+    }
+}
+
+/// Types that could send a message.
+pub trait MsgSender<M: MsgOnSocket>: AsRef<UnixDatagram> {
+    fn send(&self, msg: &M) -> MsgResult<()> {
+        let msg_size = M::msg_size();
+        let fd_size = M::max_fd_count();
+        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
+        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];
+
+        let fd_size = msg.write_to_buffer(&mut msg_buffer, &mut fd_buffer)?;
+        let sock: &UnixDatagram = self.as_ref();
+        sock.send_with_fds(&msg_buffer[..], &fd_buffer[0..fd_size])
+            .map_err(|e| MsgError::Send(e))?;
+        Ok(())
+    }
+}
+
+/// Types that could receive a message.
+pub trait MsgReceiver<M: MsgOnSocket>: AsRef<UnixDatagram> {
+    fn recv(&self) -> MsgResult<M> {
+        let msg_size = M::msg_size();
+        let fd_size = M::max_fd_count();
+        let mut msg_buffer: Vec<u8> = vec![0; msg_size];
+        let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];
+
+        let sock: &UnixDatagram = self.as_ref();
+        let (recv_msg_size, recv_fd_size) = sock
+            .recv_with_fds(&mut msg_buffer, &mut fd_buffer)
+            .map_err(|e| MsgError::Recv(e))?;
+        if msg_size != recv_msg_size {
+            return Err(MsgError::BadRecvSize(msg_size));
+        }
+        // Safe because fd buffer is read from socket.
+        let (v, read_fd_size) = unsafe {
+            M::read_from_buffer(&msg_buffer[0..recv_msg_size], &fd_buffer[0..recv_fd_size])?
+        };
+        if recv_fd_size != read_fd_size {
+            return Err(MsgError::NotExpectFd);
+        }
+        Ok(v)
+    }
+}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for MsgSocket<I, O> {}
+impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for MsgSocket<I, O> {}
+
+impl<I: MsgOnSocket, O: MsgOnSocket> MsgSender<I> for UnlinkMsgSocket<I, O> {}
+impl<I: MsgOnSocket, O: MsgOnSocket> MsgReceiver<O> for UnlinkMsgSocket<I, O> {}
+
+impl<M: MsgOnSocket> MsgSender<M> for Sender<M> {}
+impl<M: MsgOnSocket> MsgReceiver<M> for Receiver<M> {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use sys_util::EventFd;
+
+    #[derive(MsgOnSocket)]
+    struct Request {
+        field0: u8,
+        field1: EventFd,
+        field2: u32,
+    }
+
+    #[derive(MsgOnSocket)]
+    enum Response {
+        A(u8),
+        B,
+        C(u32, EventFd),
+        D([u8; 4]),
+        E { f0: u8, f1: u32 },
+    }
+
+    #[derive(MsgOnSocket)]
+    struct Message(u8, u16, EventFd);
+
+    #[test]
+    fn sock_send_recv_struct() {
+        let (req, res) = pair::<Request, Response>().unwrap();
+        let e0 = EventFd::new().unwrap();
+        let e1 = e0.try_clone().unwrap();
+        req.send(&Request {
+            field0: 2,
+            field1: e0,
+            field2: 0xf0f0,
+        }).unwrap();
+        let r = res.recv().unwrap();
+        assert_eq!(r.field0, 2);
+        assert_eq!(r.field2, 0xf0f0);
+        r.field1.write(0x0f0f).unwrap();
+        assert_eq!(e1.read().unwrap(), 0x0f0f);
+    }
+
+    #[test]
+    fn sock_send_recv_enum() {
+        let (req, res) = pair::<Request, Response>().unwrap();
+        let e0 = EventFd::new().unwrap();
+        let e1 = e0.try_clone().unwrap();
+        res.send(&Response::C(0xf0f0, e0)).unwrap();
+        let r = req.recv().unwrap();
+        match r {
+            Response::C(v, efd) => {
+                assert_eq!(v, 0xf0f0);
+                efd.write(0x0f0f).unwrap();
+            }
+            _ => panic!("wrong type"),
+        };
+        assert_eq!(e1.read().unwrap(), 0x0f0f);
+
+        res.send(&Response::B).unwrap();
+        match req.recv().unwrap() {
+            Response::B => {}
+            _ => panic!("Wrong enum type"),
+        };
+
+        res.send(&Response::A(0x3)).unwrap();
+        match req.recv().unwrap() {
+            Response::A(v) => assert_eq!(v, 0x3),
+            _ => panic!("Wrong enum type"),
+        };
+
+        res.send(&Response::D([0, 1, 2, 3])).unwrap();
+        match req.recv().unwrap() {
+            Response::D(v) => assert_eq!(v, [0, 1, 2, 3]),
+            _ => panic!("Wrong enum type"),
+        };
+
+        res.send(&Response::E {
+            f0: 0x12,
+            f1: 0x0f0f,
+        }).unwrap();
+        match req.recv().unwrap() {
+            Response::E { f0, f1 } => {
+                assert_eq!(f0, 0x12);
+                assert_eq!(f1, 0x0f0f);
+            }
+            _ => panic!("Wrong enum type"),
+        };
+    }
+
+    #[test]
+    fn sock_send_recv_tuple() {
+        let (req, res) = pair::<Message, Message>().unwrap();
+        let e0 = EventFd::new().unwrap();
+        let e1 = e0.try_clone().unwrap();
+        req.send(&Message(1, 0x12, e0)).unwrap();
+        let r = res.recv().unwrap();
+        assert_eq!(r.0, 1);
+        assert_eq!(r.1, 0x12);
+        r.2.write(0x0f0f).unwrap();
+        assert_eq!(e1.read().unwrap(), 0x0f0f);
+    }
+
+}