src/proc/bin/starnix/types/user_buffer.rs

// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use zerocopy::{AsBytes, FromBytes};

use crate::types::*;

/// Matches iovec_t.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, AsBytes, FromBytes)]
#[repr(C)]
pub struct UserBuffer {
    pub address: UserAddress,
    pub length: usize,
}

impl UserBuffer {
    pub fn get_total_length(buffers: &[UserBuffer]) -> Result<usize, Errno> {
        let mut total: usize = 0;
        for buffer in buffers {
            total = total.checked_add(buffer.length).ok_or_else(|| errno!(EINVAL))?;
        }
        // The docs say we should return EINVAL if "the sum of the iov_len
        // values overflows an ssize_t value."
        if total > isize::MAX as usize {
            return error!(EINVAL);
        }
        Ok(total)
    }
}

pub struct UserBufferIterator<'a> {
    buffers: &'a [UserBuffer],
    index: usize,
    offset: usize,
}

impl<'a> UserBufferIterator<'a> {
    pub fn new(buffers: &'a [UserBuffer]) -> UserBufferIterator<'a> {
        UserBufferIterator { buffers, index: 0, offset: 0 }
    }

    pub fn remaining(&self) -> usize {
        let remaining =
            self.buffers[self.index..].iter().fold(0, |sum, buffer| sum + buffer.length);
        remaining - self.offset
    }

    pub fn next(&mut self, limit: usize) -> Option<UserBuffer> {
        if self.index >= self.buffers.len() || limit == 0 {
            return None;
        }
        let buffer = &self.buffers[self.index];
        let chunk_size = std::cmp::min(limit, buffer.length - self.offset);
        let result = UserBuffer { address: buffer.address + self.offset, length: chunk_size };
        self.offset += chunk_size;
        while self.index < self.buffers.len() && self.offset == self.buffers[self.index].length {
            self.index += 1;
            self.offset = 0;
        }
        Some(result)
    }

    pub fn drain_to_vec(&mut self) -> Vec<UserBuffer> {
        let mut buffers = Vec::<UserBuffer>::new();
        while let Some(buffer) = self.next(usize::MAX) {
            buffers.push(buffer);
        }
        buffers
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[::fuchsia::test]
    fn test_get_total_length_overflow() {
        let buffers = vec![
            UserBuffer { address: UserAddress::from_ptr(0x20), length: usize::MAX - 10 },
            UserBuffer { address: UserAddress::from_ptr(0x820), length: usize::MAX - 10 },
            UserBuffer { address: UserAddress::from_ptr(0x3820), length: usize::MAX - 10 },
        ];
        assert!(UserBuffer::get_total_length(&buffers).is_err());
        let buffers = vec![
            UserBuffer { address: UserAddress::from_ptr(0x20), length: (isize::MAX - 10) as usize },
            UserBuffer {
                address: UserAddress::from_ptr(0x820),
                length: (isize::MAX - 10) as usize,
            },
            UserBuffer {
                address: UserAddress::from_ptr(0x3820),
                length: (isize::MAX - 10) as usize,
            },
        ];
        assert!(UserBuffer::get_total_length(&buffers).is_err());
    }

    #[::fuchsia::test]
    fn test_user_buffer_iterator() {
        let buffers = vec![
            UserBuffer { address: UserAddress::from_ptr(0x20), length: 7 },
            UserBuffer { address: UserAddress::from_ptr(0x820), length: 0 },
            UserBuffer { address: UserAddress::from_ptr(0x3820), length: 42 },
        ];
        let data = &buffers[..];
        let mut user_buffers = UserBufferIterator::new(data);
        assert_eq!(49, user_buffers.remaining());
        assert_eq!(None, user_buffers.next(0));
        assert_eq!(
            Some(UserBuffer { address: UserAddress::from_ptr(0x20), length: 1 }),
            user_buffers.next(1)
        );
        assert_eq!(48, user_buffers.remaining());
        assert_eq!(
            Some(UserBuffer { address: UserAddress::from_ptr(0x21), length: 2 }),
            user_buffers.next(2)
        );
        assert_eq!(46, user_buffers.remaining());
        assert_eq!(
            Some(UserBuffer { address: UserAddress::from_ptr(0x23), length: 4 }),
            user_buffers.next(31)
        );
        assert_eq!(42, user_buffers.remaining());
        assert_eq!(
            Some(UserBuffer { address: UserAddress::from_ptr(0x3820), length: 9 }),
            user_buffers.next(9)
        );
        assert_eq!(33, user_buffers.remaining());
        assert_eq!(
            Some(UserBuffer { address: UserAddress::from_ptr(0x3829), length: 33 }),
            user_buffers.next(40)
        );
        assert_eq!(0, user_buffers.remaining());
        assert_eq!(None, user_buffers.next(40));
        assert_eq!(0, user_buffers.remaining());
    }
}