Ashmem server for Android (#50)

* ashmem, initial commit

* ashmem

* ashmem_service: server side ready

* ashmem_service: client side ready. Ready for integration

* ashmem_service: changes to UnixShMem to make it 'threadable'

* ashmem_service: format

* ashmem_service: Undo changes to UnixShMem, make the thread own the AshmemService instead; Fix protocol bug

* ashmem_service: actually fix the protocol issue; clippy warnings

* no-std fixes

* fmt

Co-authored-by: Dominik Maier <domenukk@gmail.com>

libafl/src/bolts/llmp.rs

@@ -452,7 +452,7 @@ where
 {
     #[cfg(all(feature = "std", unix))]
     pub fn on_domain_socket(filename: &str) -> Result<Self, Error> {
-        match UnixListener::bind_unix_addr(&UnixSocketAddr::new(filename).unwrap()) {
+        match UnixListener::bind_unix_addr(&UnixSocketAddr::new(filename)?) {
             Ok(listener) => {
                 dbg!("We're the broker");
                 let mut broker = LlmpBroker::new()?;
@@ -732,7 +732,7 @@ where
                     #[cfg(feature = "std")]
                     return None;
                     #[cfg(not(feature = "std"))]
-                    panic!(&format!("Unexpected error allocing new msg {:?}", e));
+                    panic!("Unexpected error allocing new msg {:?}", e);
                 }
             };
             (*ret).message_id = (*last_msg).message_id + 1

libafl/src/bolts/os/ashmem_server.rs

@@ -0,0 +1,268 @@
+/*!
+On Android, we can only share maps between processes by serializing fds over sockets.
+Hence, the `ashmem_server` keeps track of existing maps, creates new maps for clients,
+and forwards them over unix domain sockets.
+*/
+
+use crate::{
+    bolts::shmem::{ShMem, ShMemDescription, UnixShMem},
+    Error,
+};
+use hashbrown::HashMap;
+use serde::{Deserialize, Serialize};
+use std::io::{Read, Write};
+
+#[cfg(all(feature = "std", unix))]
+use nix::poll::{poll, PollFd, PollFlags};
+
+#[cfg(all(feature = "std", unix))]
+use std::{
+    os::unix::{
+        net::{UnixListener, UnixStream},
+        {io::AsRawFd, prelude::RawFd},
+    },
+    thread,
+};
+
+#[cfg(all(unix, feature = "std"))]
+use uds::{UnixListenerExt, UnixSocketAddr, UnixStreamExt};
+
+#[derive(Debug)]
+/// The Sharedmem backed by a `ShmemService`a
+pub struct ServedShMem {
+    stream: UnixStream,
+    shmem: Option<UnixShMem>,
+    slice: Option<[u8; 20]>,
+    fd: Option<RawFd>,
+}
+const ASHMEM_SERVER_NAME: &str = "@ashmem_server";
+
+impl ServedShMem {
+    /// Create a new ServedShMem and connect to the ashmem server.
+    pub fn connect(name: &str) -> Self {
+        Self {
+            stream: UnixStream::connect_to_unix_addr(&UnixSocketAddr::from_abstract(name).unwrap())
+                .expect("Failed to connect to the ashmem server"),
+            shmem: None,
+            slice: None,
+            fd: None,
+        }
+    }
+
+    /// Send a request to the server, and wait for a response
+    fn send_receive(&mut self, request: AshmemRequest) -> ([u8; 20], RawFd) {
+        let body = postcard::to_allocvec(&request).unwrap();
+
+        let header = (body.len() as u32).to_be_bytes();
+        let mut message = header.to_vec();
+        message.extend(body);
+
+        self.stream
+            .write_all(&message)
+            .expect("Failed to send message");
+
+        let mut shm_slice = [0u8; 20];
+        let mut fd_buf = [-1; 1];
+        self.stream
+            .recv_fds(&mut shm_slice, &mut fd_buf)
+            .expect("Did not receive a response");
+        (shm_slice, fd_buf[0])
+    }
+}
+
+impl ShMem for ServedShMem {
+    fn new_map(map_size: usize) -> Result<Self, crate::Error> {
+        let mut res = Self::connect(ASHMEM_SERVER_NAME);
+        let (shm_slice, fd) = res.send_receive(AshmemRequest::NewMap(map_size));
+        if fd == -1 {
+            Err(Error::IllegalState(
+                "Could not allocate from the ashmem server".to_string(),
+            ))
+        } else {
+            res.slice = Some(shm_slice);
+            res.fd = Some(fd);
+            res.shmem = Some(
+                UnixShMem::existing_from_shm_slice(&shm_slice, map_size)
+                    .expect("Failed to create the UnixShMem"),
+            );
+            Ok(res)
+        }
+    }
+
+    fn existing_from_shm_slice(
+        map_str_bytes: &[u8; 20],
+        map_size: usize,
+    ) -> Result<Self, crate::Error> {
+        let mut res = Self::connect(ASHMEM_SERVER_NAME);
+        let (shm_slice, fd) = res.send_receive(AshmemRequest::ExistingMap(ShMemDescription {
+            size: map_size,
+            str_bytes: *map_str_bytes,
+        }));
+        if fd == -1 {
+            Err(Error::IllegalState(
+                "Could not allocate from the ashmem server".to_string(),
+            ))
+        } else {
+            res.slice = Some(shm_slice);
+            res.fd = Some(fd);
+            res.shmem = Some(
+                UnixShMem::existing_from_shm_slice(&shm_slice, map_size)
+                    .expect("Failed to create the UnixShMem"),
+            );
+            Ok(res)
+        }
+    }
+
+    fn shm_slice(&self) -> &[u8; 20] {
+        self.slice.as_ref().unwrap()
+    }
+
+    fn map(&self) -> &[u8] {
+        self.shmem.as_ref().unwrap().map()
+    }
+
+    fn map_mut(&mut self) -> &mut [u8] {
+        self.shmem.as_mut().unwrap().map_mut()
+    }
+}
+
+/// A request sent to the ShMem server to receive a fd to a shared map
+#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
+pub enum AshmemRequest {
+    /// Register a new map with a given size.
+    NewMap(usize),
+    /// Another client already has a map with this description mapped.
+    ExistingMap(ShMemDescription),
+    /// A client tells us it unregisters the previously allocated map
+    Deregister(u32),
+}
+
+#[derive(Debug)]
+pub struct AshmemClient {
+    unix_socket_file: String,
+}
+
+#[derive(Debug)]
+pub struct AshmemService {
+    maps: HashMap<[u8; 20], UnixShMem>,
+}
+
+impl AshmemService {
+    /// Create a new AshMem service
+    #[must_use]
+    fn new() -> Self {
+        AshmemService {
+            maps: HashMap::new(),
+        }
+    }
+
+    /// Read and handle the client request, send the answer over unix fd.
+    fn handle_client(&mut self, stream: &mut UnixStream) -> Result<(), Error> {
+        // Always receive one be u32 of size, then the command.
+        let mut size_bytes = [0u8; 4];
+        stream.read_exact(&mut size_bytes)?;
+        let size = u32::from_be_bytes(size_bytes);
+        let mut bytes = vec![];
+        bytes.resize(size as usize, 0u8);
+        stream
+            .read_exact(&mut bytes)
+            .expect("Failed to read message body");
+        let request: AshmemRequest = postcard::from_bytes(&bytes)?;
+
+        // Handle the client request
+        let (shmem_slice, fd): ([u8; 20], RawFd) = match request {
+            AshmemRequest::NewMap(map_size) => match UnixShMem::new(map_size) {
+                Err(e) => {
+                    println!("Error allocating shared map {:?}", e);
+                    ([0; 20], -1)
+                }
+                Ok(map) => {
+                    let res = (*map.shm_slice(), map.shm_id);
+                    self.maps.insert(*map.shm_slice(), map);
+                    res
+                }
+            },
+            AshmemRequest::ExistingMap(description) => {
+                match self.maps.get(&description.str_bytes) {
+                    None => {
+                        println!("Error finding shared map {:?}", description);
+                        ([0; 20], -1)
+                    }
+                    Some(map) => (*map.shm_slice(), map.shm_id),
+                }
+            }
+            AshmemRequest::Deregister(_) => {
+                return Ok(());
+            }
+        };
+
+        stream.send_fds(&shmem_slice, &[fd])?;
+        Ok(())
+    }
+
+    /// Create a new AshmemService, then listen and service incoming connections in a new thread.
+    pub fn start() -> Result<thread::JoinHandle<()>, Error> {
+        Ok(thread::spawn(move || {
+            Self::new().listen(ASHMEM_SERVER_NAME).unwrap()
+        }))
+    }
+
+    /// Listen on a filename (or abstract name) for new connections and serve them. This function
+    /// should not return.
+    fn listen(&mut self, filename: &str) -> Result<(), Error> {
+        let listener = UnixListener::bind_unix_addr(&UnixSocketAddr::new(filename)?)?;
+        let mut clients: HashMap<RawFd, (UnixStream, UnixSocketAddr)> = HashMap::new();
+        let mut poll_fds: HashMap<RawFd, PollFd> = HashMap::new();
+
+        poll_fds.insert(
+            listener.as_raw_fd(),
+            PollFd::new(listener.as_raw_fd(), PollFlags::POLLIN),
+        );
+
+        loop {
+            let mut fds_to_poll: Vec<PollFd> = poll_fds.values().copied().collect();
+            let fd = match poll(&mut fds_to_poll, -1) {
+                Ok(fd) => fd,
+                Err(e) => {
+                    println!("Error polling for activity: {:?}", e);
+                    continue;
+                }
+            };
+            if fd == listener.as_raw_fd() {
+                let (stream, addr) = match listener.accept_unix_addr() {
+                    Ok(stream_val) => stream_val,
+                    Err(e) => {
+                        println!("Error accepting client: {:?}", e);
+                        continue;
+                    }
+                };
+
+                println!("Recieved connection from {:?}", addr);
+                let pollfd = PollFd::new(stream.as_raw_fd(), PollFlags::POLLIN);
+                poll_fds.insert(stream.as_raw_fd(), pollfd);
+                clients
+                    .insert(stream.as_raw_fd(), (stream, addr))
+                    .as_ref()
+                    .unwrap();
+            } else if poll_fds
+                .get(&fd)
+                .unwrap()
+                .revents()
+                .unwrap()
+                .contains(PollFlags::POLLHUP)
+            {
+                poll_fds.remove(&fd);
+                clients.remove(&fd);
+            } else {
+                let (stream, _addr) = clients.get_mut(&fd).unwrap();
+                match self.handle_client(stream) {
+                    Ok(()) => (),
+                    Err(e) => {
+                        dbg!("Ignoring failed read from client", e);
+                        continue;
+                    }
+                };
+            }
+        }
+    }
+}

libafl/src/bolts/os/mod.rs

@@ -1,5 +1,8 @@
 //! Operating System specific abstractions
 
+#[cfg(all(unix, feature = "std"))]
+pub mod ashmem_server;
+
 #[cfg(unix)]
 pub mod unix_signals;
 #[cfg(windows)]

libafl/src/inputs/bytes.rs

@@ -1,4 +1,4 @@
-//! The BytesInput is the "normal" input, a map of bytes, that can be sent directly to the client
+//! The `BytesInput` is the "normal" input, a map of bytes, that can be sent directly to the client
 //! (As opposed to other, more abstract, imputs, like an Grammar-Based AST Input)
 
 use alloc::{borrow::ToOwned, rc::Rc, vec::Vec};

libafl/src/mutators/token_mutations.rs

@@ -1,5 +1,8 @@
 //! Tokens are what afl calls extras or dictionaries.
 //! They may be inserted as part of mutations during fuzzing.
+use alloc::vec::Vec;
+use core::marker::PhantomData;
+use serde::{Deserialize, Serialize};
 #[cfg(feature = "std")]
 use std::{
     fs::File,
@@ -9,18 +12,16 @@ use std::{
 
 use crate::{
     inputs::{HasBytesVec, Input},
-    mutators::{buffer_self_copy, mutations, str_decode, MutationResult, Mutator, Named},
+    mutators::{buffer_self_copy, mutations, MutationResult, Mutator, Named},
     state::{HasMaxSize, HasMetadata, HasRand},
     utils::Rand,
     Error,
 };
-use core::marker::PhantomData;
-
-use alloc::vec::Vec;
-use serde::{Deserialize, Serialize};
-
 use mutations::buffer_copy;
 
+#[cfg(feature = "std")]
+use crate::mutators::str_decode;
+
 /// A state metadata holding a list of tokens
 #[derive(Serialize, Deserialize)]
 pub struct Tokens {