From 664cc3ca49cb58d5bf7d936686fd211d6dd728bf Mon Sep 17 00:00:00 2001
From: Charles William Dick <cwd@google.com>
Date: Fri, 10 Jan 2020 14:31:52 +0900
Subject: crosvm virtio balloon stats

Introduces the ability to request BalloonStats from a
BalloonControlCommand.

BUG=b:147334004
TEST=tast run <DUT> arc.Boot.vm, and the balance available changes based
on this.

Change-Id: I808c4024f8c644c9cc4e30cc455ceda5f477bff3
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2061517
Reviewed-by: Chirantan Ekbote <chirantan@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Charles Dueck <cwd@chromium.org>
---
 vm_control/src/lib.rs | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

(limited to 'vm_control/src')

diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index b2577d9..92af2c6 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -99,7 +99,33 @@ pub const USB_CONTROL_MAX_PORTS: usize = 16;
 #[derive(MsgOnSocket, Debug)]
 pub enum BalloonControlCommand {
     /// Set the size of the VM's balloon.
-    Adjust { num_bytes: u64 },
+    Adjust {
+        num_bytes: u64,
+    },
+    Stats,
+}
+
+// BalloonStats holds stats returned from the stats_queue.
+#[derive(Default, MsgOnSocket, Debug)]
+pub struct BalloonStats {
+    pub swap_in: Option<u64>,
+    pub swap_out: Option<u64>,
+    pub major_faults: Option<u64>,
+    pub minor_faults: Option<u64>,
+    pub free_memory: Option<u64>,
+    pub total_memory: Option<u64>,
+    pub available_memory: Option<u64>,
+    pub disk_caches: Option<u64>,
+    pub hugetlb_allocations: Option<u64>,
+    pub hugetlb_failures: Option<u64>,
+}
+
+#[derive(MsgOnSocket, Debug)]
+pub enum BalloonControlResult {
+    Stats {
+        stats: BalloonStats,
+        balloon_actual: u64,
+    },
 }
 
 #[derive(MsgOnSocket, Debug)]
@@ -379,8 +405,8 @@ pub enum VmIrqResponse {
     Err(SysError),
 }
 
-pub type BalloonControlRequestSocket = MsgSocket<BalloonControlCommand, ()>;
-pub type BalloonControlResponseSocket = MsgSocket<(), BalloonControlCommand>;
+pub type BalloonControlRequestSocket = MsgSocket<BalloonControlCommand, BalloonControlResult>;
+pub type BalloonControlResponseSocket = MsgSocket<BalloonControlResult, BalloonControlCommand>;
 
 pub type DiskControlRequestSocket = MsgSocket<DiskControlCommand, DiskControlResult>;
 pub type DiskControlResponseSocket = MsgSocket<DiskControlResult, DiskControlCommand>;
-- 
cgit 1.4.1


From 146450b4569e86657d1d8c4ffe17524781aae7e3 Mon Sep 17 00:00:00 2001
From: Zach Reizner <zachr@google.com>
Date: Tue, 28 Jan 2020 13:18:09 -0800
Subject: msg_socket: support dynamically sized types

This change is a major shift in how the MsgOnSocket trait works to allow
`self` to be used to determine the result `msg_size()`. This is to
support data structures with `Vec` or other dynamically sized type.

TEST=./build_test
     cargo test -p msg_socket
     tast run <DUT> crostini.CopyPaste.*
BUG=None

Cq-Depend: chromium:2025907
Change-Id: Ibdb51b377b2a2a77892f6c75e1a9f30b2f8b0240
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2029930
Tested-by: Zach Reizner <zachr@chromium.org>
Auto-Submit: Zach Reizner <zachr@chromium.org>
Reviewed-by: Zach Reizner <zachr@chromium.org>
Commit-Queue: Zach Reizner <zachr@chromium.org>
---
 .../msg_on_socket_derive/msg_on_socket_derive.rs   | 471 ++++++++++++---------
 msg_socket/src/lib.rs                              |  25 +-
 msg_socket/src/msg_on_socket.rs                    | 234 +++++++---
 vm_control/src/lib.rs                              |   9 +-
 4 files changed, 453 insertions(+), 286 deletions(-)

(limited to 'vm_control/src')

diff --git a/msg_socket/msg_on_socket_derive/msg_on_socket_derive.rs b/msg_socket/msg_on_socket_derive/msg_on_socket_derive.rs
index 3779357..fd6c19a 100644
--- a/msg_socket/msg_on_socket_derive/msg_on_socket_derive.rs
+++ b/msg_socket/msg_on_socket_derive/msg_on_socket_derive.rs
@@ -8,8 +8,10 @@ extern crate proc_macro;
 use std::vec::Vec;
 
 use proc_macro2::{Span, TokenStream};
-use quote::quote;
-use syn::{parse_macro_input, Data, DataEnum, DataStruct, DeriveInput, Fields, Ident};
+use quote::{format_ident, quote};
+use syn::{
+    parse_macro_input, Data, DataEnum, DataStruct, DeriveInput, Fields, Ident, Index, Member, Type,
+};
 
 /// The function that derives the recursive implementation for struct or enum.
 #[proc_macro_derive(MsgOnSocket)]
@@ -50,13 +52,14 @@ fn is_named_struct(ds: &DataStruct) -> bool {
 /************************** Named Struct Impls ********************************************/
 fn impl_for_named_struct(name: Ident, ds: DataStruct) -> TokenStream {
     let fields = get_struct_fields(ds);
-    let fields_types = get_types_from_fields_vec(&fields);
-    let buffer_sizes_impls = define_buffer_size_for_struct(&fields_types);
+    let uses_fd_impl = define_uses_fd_for_struct(&fields);
+    let buffer_sizes_impls = define_buffer_size_for_struct(&fields);
 
     let read_buffer = define_read_buffer_for_struct(&name, &fields);
     let write_buffer = define_write_buffer_for_struct(&name, &fields);
     quote! {
         impl msg_socket::MsgOnSocket for #name {
+            #uses_fd_impl
             #buffer_sizes_impls
             #read_buffer
             #write_buffer
@@ -64,17 +67,8 @@ fn impl_for_named_struct(name: Ident, ds: DataStruct) -> TokenStream {
     }
 }
 
-fn get_types_from_fields_vec(v: &[(Ident, syn::Type)]) -> Vec<syn::Type> {
-    let mut fields_types = Vec::new();
-    for (_i, t) in v {
-        fields_types.push(t.clone());
-    }
-    fields_types
-}
-
 // Flatten struct fields.
-// "myfield : Type" -> \(ident\("myfield"\), Token\(Type\)\)
-fn get_struct_fields(ds: DataStruct) -> Vec<(Ident, syn::Type)> {
+fn get_struct_fields(ds: DataStruct) -> Vec<(Member, Type)> {
     let fields = match ds.fields {
         Fields::Named(fields_named) => fields_named.named,
         _ => {
@@ -83,35 +77,52 @@ fn get_struct_fields(ds: DataStruct) -> Vec<(Ident, syn::Type)> {
     };
     let mut vec = Vec::new();
     for field in fields {
-        let ident = match field.ident {
-            Some(ident) => ident,
+        let member = match field.ident {
+            Some(ident) => Member::Named(ident),
             None => panic!("Unknown Error."),
         };
         let ty = field.ty;
-        vec.push((ident, ty));
+        vec.push((member, ty));
     }
     vec
 }
 
-fn define_buffer_size_for_struct(field_types: &[syn::Type]) -> TokenStream {
-    let (msg_size, max_fd_count) = get_fields_buffer_size_sum(field_types);
+fn define_uses_fd_for_struct(fields: &[(Member, Type)]) -> TokenStream {
+    if fields.len() == 0 {
+        return quote!();
+    }
+
+    let field_types = fields.iter().map(|(_, ty)| ty);
+    quote! {
+        fn uses_fd() -> bool {
+            #(<#field_types>::uses_fd())||*
+        }
+    }
+}
+
+fn define_buffer_size_for_struct(fields: &[(Member, Type)]) -> TokenStream {
+    let (msg_size, fd_count) = get_fields_buffer_size_sum(fields);
     quote! {
-        fn msg_size() -> usize {
+        fn msg_size(&self) -> usize {
             #msg_size
         }
-        fn max_fd_count() -> usize {
-            #max_fd_count
+        fn fd_count(&self) -> usize {
+            #fd_count
         }
     }
 }
 
-fn define_read_buffer_for_struct(_name: &Ident, fields: &[(Ident, syn::Type)]) -> TokenStream {
+fn define_read_buffer_for_struct(_name: &Ident, fields: &[(Member, Type)]) -> TokenStream {
     let mut read_fields = Vec::new();
     let mut init_fields = Vec::new();
-    for f in fields {
-        let read_field = read_from_buffer_and_move_offset(&f.0, &f.1);
+    for (field_member, field_ty) in fields {
+        let ident = match field_member {
+            Member::Named(ident) => ident,
+            Member::Unnamed(_) => unreachable!(),
+        };
+        let read_field = read_from_buffer_and_move_offset(&ident, &field_ty);
         read_fields.push(read_field);
-        let name = f.0.clone();
+        let name = ident.clone();
         init_fields.push(quote!(#name));
     }
     quote! {
@@ -132,10 +143,14 @@ fn define_read_buffer_for_struct(_name: &Ident, fields: &[(Ident, syn::Type)]) -
     }
 }
 
-fn define_write_buffer_for_struct(_name: &Ident, fields: &[(Ident, syn::Type)]) -> TokenStream {
+fn define_write_buffer_for_struct(_name: &Ident, fields: &[(Member, Type)]) -> TokenStream {
     let mut write_fields = Vec::new();
-    for f in fields {
-        let write_field = write_to_buffer_and_move_offset(&f.0, &f.1);
+    for (field_member, _) in fields {
+        let ident = match field_member {
+            Member::Named(ident) => ident,
+            Member::Unnamed(_) => unreachable!(),
+        };
+        let write_field = write_to_buffer_and_move_offset(&ident);
         write_fields.push(write_field);
     }
     quote! {
@@ -154,13 +169,13 @@ fn define_write_buffer_for_struct(_name: &Ident, fields: &[(Ident, syn::Type)])
 
 /************************** Enum Impls ********************************************/
 fn impl_for_enum(name: Ident, de: DataEnum) -> TokenStream {
-    let variants = get_enum_variant_types(&de);
-    let buffer_sizes_impls = define_buffer_size_for_enum(&variants);
-
+    let uses_fd_impl = define_uses_fd_for_enum(&de);
+    let buffer_sizes_impls = define_buffer_size_for_enum(&name, &de);
     let read_buffer = define_read_buffer_for_enum(&name, &de);
     let write_buffer = define_write_buffer_for_enum(&name, &de);
     quote! {
         impl msg_socket::MsgOnSocket for #name {
+            #uses_fd_impl
             #buffer_sizes_impls
             #read_buffer
             #write_buffer
@@ -168,74 +183,103 @@ fn impl_for_enum(name: Ident, de: DataEnum) -> TokenStream {
     }
 }
 
-fn define_buffer_size_for_enum(variants: &[(Ident, Vec<syn::Type>)]) -> TokenStream {
-    let mut variant_buffer_sizes = Vec::new();
-    let mut variant_fd_sizes = Vec::new();
-    for v in variants {
-        let (msg_size_impl, fd_count_impl) = get_fields_buffer_size_sum(&v.1);
-        variant_buffer_sizes.push(msg_size_impl);
-        variant_fd_sizes.push(fd_count_impl);
+fn define_uses_fd_for_enum(de: &DataEnum) -> TokenStream {
+    let mut variant_field_types = Vec::new();
+    for variant in &de.variants {
+        for variant_field_ty in variant.fields.iter().map(|f| &f.ty) {
+            variant_field_types.push(variant_field_ty);
+        }
     }
     quote! {
-        fn msg_size() -> usize {
-            // First byte is used for variant.
-            [#(#variant_buffer_sizes,)*].iter().max().unwrap().clone() as usize + 1
-        }
-        fn max_fd_count() -> usize {
-            [#(#variant_fd_sizes,)*].iter().max().unwrap().clone() as usize
+        fn uses_fd() -> bool {
+            #(<#variant_field_types>::uses_fd())||*
         }
     }
 }
 
-// Flatten enum variants. Return value = \[variant_name, \[types_of_this_variant\]\]
-fn get_enum_variant_types(de: &DataEnum) -> Vec<(Ident, Vec<syn::Type>)> {
-    let mut variants = Vec::new();
-    let de = de.clone();
-    for v in de.variants {
-        let name = v.ident;
-        match v.fields {
+fn define_buffer_size_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
+    let mut msg_size_match_variants = Vec::new();
+    let mut fd_count_match_variants = Vec::new();
+
+    for variant in &de.variants {
+        let variant_name = &variant.ident;
+        match &variant.fields {
+            Fields::Named(fields) => {
+                let mut tmp_names = Vec::new();
+                for field in &fields.named {
+                    tmp_names.push(field.ident.clone().unwrap());
+                }
+
+                let v = quote! {
+                    #name::#variant_name { #(#tmp_names),* } => #(#tmp_names.msg_size())+*,
+                };
+                msg_size_match_variants.push(v);
+
+                let v = quote! {
+                    #name::#variant_name { #(#tmp_names),* } => #(#tmp_names.fd_count())+*,
+                };
+                fd_count_match_variants.push(v);
+            }
             Fields::Unnamed(fields) => {
-                let mut vec = Vec::new();
-                for field in fields.unnamed {
-                    let ty = field.ty;
-                    vec.push(ty);
+                let mut tmp_names = Vec::new();
+                for idx in 0..fields.unnamed.len() {
+                    let tmp_name = format!("enum_field{}", idx);
+                    let tmp_name = Ident::new(&tmp_name, Span::call_site());
+                    tmp_names.push(tmp_name.clone());
                 }
-                variants.push((name, vec));
+
+                let v = quote! {
+                    #name::#variant_name(#(#tmp_names),*) => #(#tmp_names.msg_size())+*,
+                };
+                msg_size_match_variants.push(v);
+
+                let v = quote! {
+                    #name::#variant_name(#(#tmp_names),*) => #(#tmp_names.fd_count())+*,
+                };
+                fd_count_match_variants.push(v);
             }
             Fields::Unit => {
-                variants.push((name, Vec::new()));
-                continue;
+                let v = quote! {
+                    #name::#variant_name => 0,
+                };
+                msg_size_match_variants.push(v.clone());
+                fd_count_match_variants.push(v);
             }
-            Fields::Named(fields) => {
-                let mut vec = Vec::new();
-                for field in fields.named {
-                    let ty = field.ty;
-                    vec.push(ty);
-                }
-                variants.push((name, vec));
+        }
+    }
+
+    quote! {
+        fn msg_size(&self) -> usize {
+            1 + match self {
+                #(#msg_size_match_variants)*
             }
-        };
+        }
+        fn fd_count(&self) -> usize {
+            match self {
+                #(#fd_count_match_variants)*
+            }
+        }
     }
-    variants
 }
 
 fn define_read_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
     let mut match_variants = Vec::new();
     let de = de.clone();
-    let mut i = 0u8;
-    for v in de.variants {
-        let variant_name = v.ident;
-        match v.fields {
+    for (idx, variant) in de.variants.iter().enumerate() {
+        let idx = idx as u8;
+        let variant_name = &variant.ident;
+        match &variant.fields {
             Fields::Named(fields) => {
                 let mut tmp_names = Vec::new();
                 let mut read_tmps = Vec::new();
-                for f in fields.named {
+                for f in &fields.named {
                     tmp_names.push(f.ident.clone());
-                    let read_tmp = read_from_buffer_and_move_offset(&f.ident.unwrap(), &f.ty);
+                    let read_tmp =
+                        read_from_buffer_and_move_offset(f.ident.as_ref().unwrap(), &f.ty);
                     read_tmps.push(read_tmp);
                 }
                 let v = quote! {
-                    #i => {
+                    #idx => {
                         let mut __offset = 1usize;
                         let mut __fd_offset = 0usize;
                         #(#read_tmps)*
@@ -247,18 +291,15 @@ fn define_read_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
             Fields::Unnamed(fields) => {
                 let mut tmp_names = Vec::new();
                 let mut read_tmps = Vec::new();
-                let mut j = 0usize;
-                for f in fields.unnamed {
-                    let tmp_name = format!("enum_variant_tmp{}", j);
-                    let tmp_name = Ident::new(&tmp_name, Span::call_site());
+                for (idx, field) in fields.unnamed.iter().enumerate() {
+                    let tmp_name = format_ident!("enum_field{}", idx);
                     tmp_names.push(tmp_name.clone());
-                    let read_tmp = read_from_buffer_and_move_offset(&tmp_name, &f.ty);
+                    let read_tmp = read_from_buffer_and_move_offset(&tmp_name, &field.ty);
                     read_tmps.push(read_tmp);
-                    j += 1;
                 }
 
                 let v = quote! {
-                    #i => {
+                    #idx => {
                         let mut __offset = 1usize;
                         let mut __fd_offset = 0usize;
                         #(#read_tmps)*
@@ -269,19 +310,18 @@ fn define_read_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
             }
             Fields::Unit => {
                 let v = quote! {
-                    #i => Ok((#name::#variant_name, 0)),
+                    #idx => Ok((#name::#variant_name, 0)),
                 };
                 match_variants.push(v);
             }
         }
-        i += 1;
     }
     quote! {
         unsafe fn read_from_buffer(
             buffer: &[u8],
             fds: &[std::os::unix::io::RawFd],
         ) -> msg_socket::MsgResult<(Self, usize)> {
-            let v = buffer[0];
+            let v = buffer.get(0).ok_or(msg_socket::MsgError::WrongMsgBufferSize)?;
             match v {
                 #(#match_variants)*
                 _ => Err(msg_socket::MsgError::InvalidType),
@@ -292,23 +332,24 @@ fn define_read_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
 
 fn define_write_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
     let mut match_variants = Vec::new();
-    let mut i = 0u8;
     let de = de.clone();
-    for v in de.variants {
-        let variant_name = v.ident;
-        match v.fields {
+    for (idx, variant) in de.variants.iter().enumerate() {
+        let idx = idx as u8;
+        let variant_name = &variant.ident;
+        match &variant.fields {
             Fields::Named(fields) => {
                 let mut tmp_names = Vec::new();
                 let mut write_tmps = Vec::new();
-                for f in fields.named {
+                for f in &fields.named {
                     tmp_names.push(f.ident.clone().unwrap());
-                    let write_tmp = enum_write_to_buffer_and_move_offset(&f.ident.unwrap(), &f.ty);
+                    let write_tmp =
+                        enum_write_to_buffer_and_move_offset(&f.ident.as_ref().unwrap());
                     write_tmps.push(write_tmp);
                 }
 
                 let v = quote! {
                     #name::#variant_name { #(#tmp_names),* } => {
-                        buffer[0] = #i;
+                        buffer[0] = #idx;
                         let mut __offset = 1usize;
                         let mut __fd_offset = 0usize;
                         #(#write_tmps)*
@@ -320,19 +361,16 @@ fn define_write_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
             Fields::Unnamed(fields) => {
                 let mut tmp_names = Vec::new();
                 let mut write_tmps = Vec::new();
-                let mut j = 0usize;
-                for f in fields.unnamed {
-                    let tmp_name = format!("enum_variant_tmp{}", j);
-                    let tmp_name = Ident::new(&tmp_name, Span::call_site());
+                for idx in 0..fields.unnamed.len() {
+                    let tmp_name = format_ident!("enum_field{}", idx);
                     tmp_names.push(tmp_name.clone());
-                    let write_tmp = enum_write_to_buffer_and_move_offset(&tmp_name, &f.ty);
+                    let write_tmp = enum_write_to_buffer_and_move_offset(&tmp_name);
                     write_tmps.push(write_tmp);
-                    j += 1;
                 }
 
                 let v = quote! {
                     #name::#variant_name(#(#tmp_names),*) => {
-                        buffer[0] = #i;
+                        buffer[0] = #idx;
                         let mut __offset = 1usize;
                         let mut __fd_offset = 0usize;
                         #(#write_tmps)*
@@ -344,14 +382,13 @@ fn define_write_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
             Fields::Unit => {
                 let v = quote! {
                     #name::#variant_name => {
-                        buffer[0] = #i;
+                        buffer[0] = #idx;
                         Ok(0)
                     }
                 };
                 match_variants.push(v);
             }
         }
-        i += 1;
     }
 
     quote! {
@@ -360,6 +397,9 @@ fn define_write_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
             buffer: &mut [u8],
             fds: &mut [std::os::unix::io::RawFd],
         ) -> msg_socket::MsgResult<usize> {
+            if buffer.is_empty() {
+                return Err(msg_socket::MsgError::WrongMsgBufferSize)
+            }
             match self {
                 #(#match_variants)*
             }
@@ -367,24 +407,25 @@ fn define_write_buffer_for_enum(name: &Ident, de: &DataEnum) -> TokenStream {
     }
 }
 
-fn enum_write_to_buffer_and_move_offset(name: &Ident, ty: &syn::Type) -> TokenStream {
+fn enum_write_to_buffer_and_move_offset(name: &Ident) -> TokenStream {
     quote! {
         let o = #name.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-        __offset += <#ty>::msg_size();
+        __offset += #name.msg_size();
         __fd_offset += o;
     }
 }
 
 /************************** Tuple Impls ********************************************/
 fn impl_for_tuple_struct(name: Ident, ds: DataStruct) -> TokenStream {
-    let types = get_tuple_types(ds);
+    let fields = get_tuple_fields(ds);
 
-    let buffer_sizes_impls = define_buffer_size_for_struct(&types);
-
-    let read_buffer = define_read_buffer_for_tuples(&name, &types);
-    let write_buffer = define_write_buffer_for_tuples(&name, &types);
+    let uses_fd_impl = define_uses_fd_for_tuples(&fields);
+    let buffer_sizes_impls = define_buffer_size_for_struct(&fields);
+    let read_buffer = define_read_buffer_for_tuples(&name, &fields);
+    let write_buffer = define_write_buffer_for_tuples(&name, &fields);
     quote! {
         impl msg_socket::MsgOnSocket for #name {
+            #uses_fd_impl
             #buffer_sizes_impls
             #read_buffer
             #write_buffer
@@ -392,28 +433,42 @@ fn impl_for_tuple_struct(name: Ident, ds: DataStruct) -> TokenStream {
     }
 }
 
-fn get_tuple_types(ds: DataStruct) -> Vec<syn::Type> {
-    let mut types = Vec::new();
+fn get_tuple_fields(ds: DataStruct) -> Vec<(Member, Type)> {
+    let mut field_idents = Vec::new();
     let fields = match ds.fields {
         Fields::Unnamed(fields_unnamed) => fields_unnamed.unnamed,
         _ => {
             panic!("Tuple struct must have unnamed fields.");
         }
     };
-    for field in fields {
-        let ty = field.ty;
-        types.push(ty);
+    for (idx, field) in fields.iter().enumerate() {
+        let member = Member::Unnamed(Index::from(idx));
+        let ty = field.ty.clone();
+        field_idents.push((member, ty));
+    }
+    field_idents
+}
+
+fn define_uses_fd_for_tuples(fields: &[(Member, Type)]) -> TokenStream {
+    if fields.len() == 0 {
+        return quote!();
+    }
+
+    let field_types = fields.iter().map(|(_, ty)| ty);
+    quote! {
+        fn uses_fd() -> bool {
+            #(<#field_types>::uses_fd())||*
+        }
     }
-    types
 }
 
-fn define_read_buffer_for_tuples(name: &Ident, fields: &[syn::Type]) -> TokenStream {
+fn define_read_buffer_for_tuples(name: &Ident, fields: &[(Member, Type)]) -> TokenStream {
     let mut read_fields = Vec::new();
     let mut init_fields = Vec::new();
-    for i in 0..fields.len() {
-        let tmp_name = format!("tuple_tmp{}", i);
+    for (idx, (_, field_ty)) in fields.iter().enumerate() {
+        let tmp_name = format!("tuple_tmp{}", idx);
         let tmp_name = Ident::new(&tmp_name, Span::call_site());
-        let read_field = read_from_buffer_and_move_offset(&tmp_name, &fields[i]);
+        let read_field = read_from_buffer_and_move_offset(&tmp_name, field_ty);
         read_fields.push(read_field);
         init_fields.push(quote!(#tmp_name));
     }
@@ -436,13 +491,12 @@ fn define_read_buffer_for_tuples(name: &Ident, fields: &[syn::Type]) -> TokenStr
     }
 }
 
-fn define_write_buffer_for_tuples(name: &Ident, fields: &[syn::Type]) -> TokenStream {
+fn define_write_buffer_for_tuples(name: &Ident, fields: &[(Member, Type)]) -> TokenStream {
     let mut write_fields = Vec::new();
     let mut tmp_names = Vec::new();
-    for i in 0..fields.len() {
-        let tmp_name = format!("tuple_tmp{}", i);
-        let tmp_name = Ident::new(&tmp_name, Span::call_site());
-        let write_field = enum_write_to_buffer_and_move_offset(&tmp_name, &fields[i]);
+    for idx in 0..fields.len() {
+        let tmp_name = format_ident!("tuple_tmp{}", idx);
+        let write_field = enum_write_to_buffer_and_move_offset(&tmp_name);
         write_fields.push(write_field);
         tmp_names.push(tmp_name);
     }
@@ -461,14 +515,15 @@ fn define_write_buffer_for_tuples(name: &Ident, fields: &[syn::Type]) -> TokenSt
     }
 }
 /************************** Helpers ********************************************/
-fn get_fields_buffer_size_sum(field_types: &[syn::Type]) -> (TokenStream, TokenStream) {
-    if field_types.len() > 0 {
+fn get_fields_buffer_size_sum(fields: &[(Member, Type)]) -> (TokenStream, TokenStream) {
+    let fields: Vec<_> = fields.iter().map(|(m, _)| m).collect();
+    if fields.len() > 0 {
         (
             quote! {
-                #( <#field_types>::msg_size() as usize )+*
+                #( self.#fields.msg_size() as usize )+*
             },
             quote! {
-                #( <#field_types>::max_fd_count() as usize )+*
+                #( self.#fields.fd_count() as usize )+*
             },
         )
     } else {
@@ -476,19 +531,19 @@ fn get_fields_buffer_size_sum(field_types: &[syn::Type]) -> (TokenStream, TokenS
     }
 }
 
-fn read_from_buffer_and_move_offset(name: &Ident, ty: &syn::Type) -> TokenStream {
+fn read_from_buffer_and_move_offset(name: &Ident, ty: &Type) -> TokenStream {
     quote! {
         let t = <#ty>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-        __offset += <#ty>::msg_size();
+        __offset += t.0.msg_size();
         __fd_offset += t.1;
         let #name = t.0;
     }
 }
 
-fn write_to_buffer_and_move_offset(name: &Ident, ty: &syn::Type) -> TokenStream {
+fn write_to_buffer_and_move_offset(name: &Ident) -> TokenStream {
     quote! {
         let o = self.#name.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-        __offset += <#ty>::msg_size();
+        __offset += self.#name.msg_size();
         __fd_offset += o;
     }
 }
@@ -511,15 +566,18 @@ mod tests {
 
         let expected = quote! {
             impl msg_socket::MsgOnSocket for MyMsg {
-                fn msg_size() -> usize {
-                    <u8>::msg_size() as usize
-                        + <RawFd>::msg_size() as usize
-                        + <u32>::msg_size() as usize
+                fn uses_fd() -> bool {
+                    <u8>::uses_fd() || <RawFd>::uses_fd() || <u32>::uses_fd()
                 }
-                fn max_fd_count() -> usize {
-                    <u8>::max_fd_count() as usize
-                        + <RawFd>::max_fd_count() as usize
-                        + <u32>::max_fd_count() as usize
+                fn msg_size(&self) -> usize {
+                    self.a.msg_size() as usize
+                        + self.b.msg_size() as usize
+                        + self.c.msg_size() as usize
+                }
+                fn fd_count(&self) -> usize {
+                    self.a.fd_count() as usize
+                        + self.b.fd_count() as usize
+                        + self.c.fd_count() as usize
                 }
                 unsafe fn read_from_buffer(
                     buffer: &[u8],
@@ -528,15 +586,15 @@ mod tests {
                     let mut __offset = 0usize;
                     let mut __fd_offset = 0usize;
                     let t = <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <u8>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let a = t.0;
                     let t = <RawFd>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <RawFd>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let b = t.0;
                     let t = <u32>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <u32>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let c = t.0;
                     Ok((Self { a, b, c }, __fd_offset))
@@ -548,21 +606,25 @@ mod tests {
                 ) -> msg_socket::MsgResult<usize> {
                     let mut __offset = 0usize;
                     let mut __fd_offset = 0usize;
-                    let o = self.a
+                    let o = self
+                        .a
                         .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <u8>::msg_size();
+                    __offset += self.a.msg_size();
                     __fd_offset += o;
-                    let o = self.b
+                    let o = self
+                        .b
                         .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <RawFd>::msg_size();
+                    __offset += self.b.msg_size();
                     __fd_offset += o;
-                    let o = self.c
+                    let o = self
+                        .c
                         .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <u32>::msg_size();
+                    __offset += self.c.msg_size();
                     __fd_offset += o;
                     Ok(__fd_offset)
                 }
             }
+
         };
 
         assert_eq!(socket_msg_impl(input).to_string(), expected.to_string());
@@ -576,15 +638,17 @@ mod tests {
 
         let expected = quote! {
             impl msg_socket::MsgOnSocket for MyMsg {
-                fn msg_size() -> usize {
-                    <u8>::msg_size() as usize
-                        + <u32>::msg_size() as usize
-                        + <File>::msg_size() as usize
+                fn uses_fd() -> bool {
+                    <u8>::uses_fd() || <u32>::uses_fd() || <File>::uses_fd()
+                }
+                fn msg_size(&self) -> usize {
+                    self.0.msg_size() as usize
+                        + self.1.msg_size() as usize + self.2.msg_size() as usize
                 }
-                fn max_fd_count() -> usize {
-                    <u8>::max_fd_count() as usize
-                        + <u32>::max_fd_count() as usize
-                        + <File>::max_fd_count() as usize
+                fn fd_count(&self) -> usize {
+                    self.0.fd_count() as usize
+                        + self.1.fd_count() as usize
+                        + self.2.fd_count() as usize
                 }
                 unsafe fn read_from_buffer(
                     buffer: &[u8],
@@ -593,15 +657,15 @@ mod tests {
                     let mut __offset = 0usize;
                     let mut __fd_offset = 0usize;
                     let t = <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <u8>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let tuple_tmp0 = t.0;
                     let t = <u32>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <u32>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let tuple_tmp1 = t.0;
                     let t = <File>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                    __offset += <File>::msg_size();
+                    __offset += t.0.msg_size();
                     __fd_offset += t.1;
                     let tuple_tmp2 = t.0;
                     Ok((MyMsg(tuple_tmp0, tuple_tmp1, tuple_tmp2), __fd_offset))
@@ -614,17 +678,14 @@ mod tests {
                     let mut __offset = 0usize;
                     let mut __fd_offset = 0usize;
                     let MyMsg(tuple_tmp0, tuple_tmp1, tuple_tmp2) = self;
-                    let o = tuple_tmp0
-                        .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <u8>::msg_size();
+                    let o = tuple_tmp0.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
+                    __offset += tuple_tmp0.msg_size();
                     __fd_offset += o;
-                    let o = tuple_tmp1
-                        .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <u32>::msg_size();
+                    let o = tuple_tmp1.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
+                    __offset += tuple_tmp1.msg_size();
                     __fd_offset += o;
-                    let o = tuple_tmp2
-                        .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                    __offset += <File>::msg_size();
+                    let o = tuple_tmp2.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
+                    __offset += tuple_tmp2.msg_size();
                     __fd_offset += o;
                     Ok(__fd_offset)
                 }
@@ -649,52 +710,50 @@ mod tests {
 
         let expected = quote! {
             impl msg_socket::MsgOnSocket for MyMsg {
-                fn msg_size() -> usize {
-                    [
-                        <u8>::msg_size() as usize,
-                        0,
-                        <u8>::msg_size() as usize + <RawFd>::msg_size() as usize,
-                    ].iter()
-                        .max().unwrap().clone() as usize+ 1
+                fn uses_fd() -> bool {
+                    <u8>::uses_fd() || <u8>::uses_fd() || <RawFd>::uses_fd()
+                }
+                fn msg_size(&self) -> usize {
+                    1 + match self {
+                        MyMsg::A(enum_field0) => enum_field0.msg_size(),
+                        MyMsg::B => 0,
+                        MyMsg::C { f0, f1 } => f0.msg_size() + f1.msg_size(),
+                    }
                 }
-                fn max_fd_count() -> usize {
-                    [
-                        <u8>::max_fd_count() as usize,
-                        0,
-                        <u8>::max_fd_count() as usize + <RawFd>::max_fd_count() as usize,
-                    ].iter()
-                        .max().unwrap().clone() as usize
+                fn fd_count(&self) -> usize {
+                    match self {
+                        MyMsg::A(enum_field0) => enum_field0.fd_count(),
+                        MyMsg::B => 0,
+                        MyMsg::C { f0, f1 } => f0.fd_count() + f1.fd_count(),
+                    }
                 }
                 unsafe fn read_from_buffer(
                     buffer: &[u8],
                     fds: &[std::os::unix::io::RawFd],
                 ) -> msg_socket::MsgResult<(Self, usize)> {
-                    let v = buffer[0];
+                    let v = buffer
+                        .get(0)
+                        .ok_or(msg_socket::MsgError::WrongMsgBufferSize)?;
                     match v {
                         0u8 => {
                             let mut __offset = 1usize;
                             let mut __fd_offset = 0usize;
-                            let t =
-                                <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                            __offset += <u8>::msg_size();
+                            let t = <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
+                            __offset += t.0.msg_size();
                             __fd_offset += t.1;
-                            let enum_variant_tmp0 = t.0;
-                            Ok((MyMsg::A(enum_variant_tmp0), __fd_offset))
+                            let enum_field0 = t.0;
+                            Ok((MyMsg::A(enum_field0), __fd_offset))
                         }
                         1u8 => Ok((MyMsg::B, 0)),
                         2u8 => {
                             let mut __offset = 1usize;
                             let mut __fd_offset = 0usize;
-                            let t =
-                                <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
-                            __offset += <u8>::msg_size();
+                            let t = <u8>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
+                            __offset += t.0.msg_size();
                             __fd_offset += t.1;
                             let f0 = t.0;
-                            let t = <RawFd>::read_from_buffer(
-                                &buffer[__offset..],
-                                &fds[__fd_offset..]
-                            )?;
-                            __offset += <RawFd>::msg_size();
+                            let t = <RawFd>::read_from_buffer(&buffer[__offset..], &fds[__fd_offset..])?;
+                            __offset += t.0.msg_size();
                             __fd_offset += t.1;
                             let f1 = t.0;
                             Ok((MyMsg::C { f0, f1 }, __fd_offset))
@@ -707,14 +766,17 @@ mod tests {
                     buffer: &mut [u8],
                     fds: &mut [std::os::unix::io::RawFd],
                 ) -> msg_socket::MsgResult<usize> {
+                    if buffer.is_empty() {
+                        return Err(msg_socket::MsgError::WrongMsgBufferSize)
+                    }
                     match self {
-                        MyMsg::A(enum_variant_tmp0) => {
+                        MyMsg::A(enum_field0) => {
                             buffer[0] = 0u8;
                             let mut __offset = 1usize;
                             let mut __fd_offset = 0usize;
-                            let o = enum_variant_tmp0
+                            let o = enum_field0
                                 .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                            __offset += <u8>::msg_size();
+                            __offset += enum_field0.msg_size();
                             __fd_offset += o;
                             Ok(__fd_offset)
                         }
@@ -726,20 +788,17 @@ mod tests {
                             buffer[0] = 2u8;
                             let mut __offset = 1usize;
                             let mut __fd_offset = 0usize;
-                            let o = f0
-                                .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                            __offset += <u8>::msg_size();
+                            let o = f0.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
+                            __offset += f0.msg_size();
                             __fd_offset += o;
-                            let o = f1
-                                .write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
-                            __offset += <RawFd>::msg_size();
+                            let o = f1.write_to_buffer(&mut buffer[__offset..], &mut fds[__fd_offset..])?;
+                            __offset += f1.msg_size();
                             __fd_offset += o;
                             Ok(__fd_offset)
                         }
                     }
                 }
             }
-
         };
 
         assert_eq!(socket_msg_impl(input).to_string(), expected.to_string());
diff --git a/msg_socket/src/lib.rs b/msg_socket/src/lib.rs
index ea817f0..5871735 100644
--- a/msg_socket/src/lib.rs
+++ b/msg_socket/src/lib.rs
@@ -127,8 +127,8 @@ impl<M: MsgOnSocket> AsRawFd for Receiver<M> {
 pub trait MsgSender: AsRef<UnixSeqpacket> {
     type M: MsgOnSocket;
     fn send(&self, msg: &Self::M) -> MsgResult<()> {
-        let msg_size = Self::M::msg_size();
-        let fd_size = Self::M::max_fd_count();
+        let msg_size = msg.msg_size();
+        let fd_size = msg.fd_count();
         let mut msg_buffer: Vec<u8> = vec![0; msg_size];
         let mut fd_buffer: Vec<RawFd> = vec![0; fd_size];
 
@@ -149,34 +149,23 @@ pub trait MsgSender: AsRef<UnixSeqpacket> {
 pub trait MsgReceiver: AsRef<UnixSeqpacket> {
     type M: MsgOnSocket;
     fn recv(&self) -> MsgResult<Self::M> {
-        let msg_size = Self::M::msg_size();
-        let fd_size = Self::M::max_fd_count();
-
         let sock: &UnixSeqpacket = self.as_ref();
 
         let (msg_buffer, fd_buffer) = {
-            if fd_size == 0 {
+            if Self::M::uses_fd() {
+                sock.recv_as_vec_with_fds()
+                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?
+            } else {
                 (
                     sock.recv_as_vec().map_err(|e| {
                         MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0)))
                     })?,
                     vec![],
                 )
-            } else {
-                sock.recv_as_vec_with_fds()
-                    .map_err(|e| MsgError::Recv(SysError::new(e.raw_os_error().unwrap_or(0))))?
             }
         };
-
-        if msg_size != msg_buffer.len() {
-            return Err(MsgError::BadRecvSize {
-                expected: msg_size,
-                actual: msg_buffer.len(),
-            });
-        }
         // Safe because fd buffer is read from socket.
-        let (v, read_fd_size) =
-            unsafe { Self::M::read_from_buffer(&msg_buffer[..], &fd_buffer[..])? };
+        let (v, read_fd_size) = unsafe { Self::M::read_from_buffer(&msg_buffer, &fd_buffer)? };
         if fd_buffer.len() != read_fd_size {
             return Err(MsgError::NotExpectFd);
         }
diff --git a/msg_socket/src/msg_on_socket.rs b/msg_socket/src/msg_on_socket.rs
index f03c36f..3f34019 100644
--- a/msg_socket/src/msg_on_socket.rs
+++ b/msg_socket/src/msg_on_socket.rs
@@ -4,9 +4,11 @@
 
 use std::fmt::{self, Display};
 use std::fs::File;
+use std::mem::{size_of, transmute_copy, MaybeUninit};
 use std::net::{TcpListener, TcpStream, UdpSocket};
 use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
 use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
+use std::ptr::drop_in_place;
 use std::result;
 
 use data_model::*;
@@ -90,10 +92,24 @@ impl Display for MsgError {
 /// Thus, read/write functions always the return correct count of fds in this variant. There will be
 /// no padding in fd_buffer.
 pub trait MsgOnSocket: Sized {
+    // `true` if this structure can potentially serialize fds.
+    fn uses_fd() -> bool {
+        false
+    }
+
+    // Returns `Some(size)` if this structure always has a fixed size.
+    fn fixed_size() -> Option<usize> {
+        None
+    }
+
     /// Size of message in bytes.
-    fn msg_size() -> usize;
-    /// Max possible fd count in this type.
-    fn max_fd_count() -> usize {
+    fn msg_size(&self) -> usize {
+        Self::fixed_size().unwrap()
+    }
+
+    /// Number of FDs in this message. This must be overridden if `uses_fd()` returns true.
+    fn fd_count(&self) -> usize {
+        assert!(!Self::uses_fd());
         0
     }
     /// Returns (self, fd read count).
@@ -103,13 +119,14 @@ pub trait MsgOnSocket: Sized {
     ///     2. write_to_buffer is implemented correctly(put valid fds into the buffer, has no padding,
     ///        return correct count).
     unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)>;
+
     /// Serialize self to buffers.
     fn write_to_buffer(&self, buffer: &mut [u8], fds: &mut [RawFd]) -> MsgResult<usize>;
 }
 
 impl MsgOnSocket for SysError {
-    fn msg_size() -> usize {
-        u32::msg_size()
+    fn fixed_size() -> Option<usize> {
+        Some(size_of::<u32>())
     }
     unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
         let (v, size) = u32::read_from_buffer(buffer, fds)?;
@@ -122,12 +139,14 @@ impl MsgOnSocket for SysError {
 }
 
 impl MsgOnSocket for RawFd {
-    fn msg_size() -> usize {
-        0
+    fn fixed_size() -> Option<usize> {
+        Some(0)
     }
-    fn max_fd_count() -> usize {
+
+    fn fd_count(&self) -> usize {
         1
     }
+
     unsafe fn read_from_buffer(_buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
         if fds.is_empty() {
             return Err(MsgError::ExpectFd);
@@ -144,12 +163,22 @@ impl MsgOnSocket for RawFd {
 }
 
 impl<T: MsgOnSocket> MsgOnSocket for Option<T> {
-    fn msg_size() -> usize {
-        T::msg_size() + 1
+    fn uses_fd() -> bool {
+        T::uses_fd()
+    }
+
+    fn msg_size(&self) -> usize {
+        match self {
+            Some(v) => v.msg_size() + 1,
+            None => 0,
+        }
     }
 
-    fn max_fd_count() -> usize {
-        T::max_fd_count()
+    fn fd_count(&self) -> usize {
+        match self {
+            Some(v) => v.fd_count(),
+            None => 0,
+        }
     }
 
     unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
@@ -178,12 +207,8 @@ impl<T: MsgOnSocket> MsgOnSocket for Option<T> {
 }
 
 impl MsgOnSocket for () {
-    fn msg_size() -> usize {
-        0
-    }
-
-    fn max_fd_count() -> usize {
-        0
+    fn fixed_size() -> Option<usize> {
+        Some(0)
     }
 
     unsafe fn read_from_buffer(_buffer: &[u8], _fds: &[RawFd]) -> MsgResult<(Self, usize)> {
@@ -198,17 +223,20 @@ impl MsgOnSocket for () {
 macro_rules! rawfd_impl {
     ($type:ident) => {
         impl MsgOnSocket for $type {
-            fn msg_size() -> usize {
+            fn uses_fd() -> bool {
+                true
+            }
+            fn msg_size(&self) -> usize {
                 0
             }
-            fn max_fd_count() -> usize {
+            fn fd_count(&self) -> usize {
                 1
             }
             unsafe fn read_from_buffer(_buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
                 if fds.len() < 1 {
                     return Err(MsgError::ExpectFd);
                 }
-                Ok(($type::from_raw_fd(fds[0].clone()), 1))
+                Ok(($type::from_raw_fd(fds[0]), 1))
             }
             fn write_to_buffer(&self, _buffer: &mut [u8], fds: &mut [RawFd]) -> MsgResult<usize> {
                 if fds.len() < 1 {
@@ -246,11 +274,11 @@ where
 
 // usize could be different sizes on different targets. We always use u64.
 impl MsgOnSocket for usize {
-    fn msg_size() -> usize {
-        std::mem::size_of::<u64>()
+    fn msg_size(&self) -> usize {
+        size_of::<u64>()
     }
     unsafe fn read_from_buffer(buffer: &[u8], _fds: &[RawFd]) -> MsgResult<(Self, usize)> {
-        if buffer.len() < std::mem::size_of::<u64>() {
+        if buffer.len() < size_of::<u64>() {
             return Err(MsgError::WrongMsgBufferSize);
         }
         let t = u64::from_le_bytes(slice_to_array(buffer));
@@ -258,22 +286,22 @@ impl MsgOnSocket for usize {
     }
 
     fn write_to_buffer(&self, buffer: &mut [u8], _fds: &mut [RawFd]) -> MsgResult<usize> {
-        if buffer.len() < std::mem::size_of::<u64>() {
+        if buffer.len() < size_of::<u64>() {
             return Err(MsgError::WrongMsgBufferSize);
         }
         let t: Le64 = (*self as u64).into();
-        buffer[0..Self::msg_size()].copy_from_slice(t.as_slice());
+        buffer[0..self.msg_size()].copy_from_slice(t.as_slice());
         Ok(0)
     }
 }
 
 // Encode bool as a u8 of value 0 or 1
 impl MsgOnSocket for bool {
-    fn msg_size() -> usize {
-        std::mem::size_of::<u8>()
+    fn msg_size(&self) -> usize {
+        size_of::<u8>()
     }
     unsafe fn read_from_buffer(buffer: &[u8], _fds: &[RawFd]) -> MsgResult<(Self, usize)> {
-        if buffer.len() < std::mem::size_of::<u8>() {
+        if buffer.len() < size_of::<u8>() {
             return Err(MsgError::WrongMsgBufferSize);
         }
         let t: u8 = buffer[0];
@@ -284,7 +312,7 @@ impl MsgOnSocket for bool {
         }
     }
     fn write_to_buffer(&self, buffer: &mut [u8], _fds: &mut [RawFd]) -> MsgResult<usize> {
-        if buffer.len() < std::mem::size_of::<u8>() {
+        if buffer.len() < size_of::<u8>() {
             return Err(MsgError::WrongMsgBufferSize);
         }
         buffer[0] = *self as u8;
@@ -295,11 +323,12 @@ impl MsgOnSocket for bool {
 macro_rules! le_impl {
     ($type:ident, $native_type:ident) => {
         impl MsgOnSocket for $type {
-            fn msg_size() -> usize {
-                std::mem::size_of::<$native_type>()
+            fn fixed_size() -> Option<usize> {
+                Some(size_of::<$native_type>())
             }
+
             unsafe fn read_from_buffer(buffer: &[u8], _fds: &[RawFd]) -> MsgResult<(Self, usize)> {
-                if buffer.len() < std::mem::size_of::<$native_type>() {
+                if buffer.len() < size_of::<$native_type>() {
                     return Err(MsgError::WrongMsgBufferSize);
                 }
                 let t = $native_type::from_le_bytes(slice_to_array(buffer));
@@ -307,11 +336,11 @@ macro_rules! le_impl {
             }
 
             fn write_to_buffer(&self, buffer: &mut [u8], _fds: &mut [RawFd]) -> MsgResult<usize> {
-                if buffer.len() < std::mem::size_of::<$native_type>() {
+                if buffer.len() < size_of::<$native_type>() {
                     return Err(MsgError::WrongMsgBufferSize);
                 }
                 let t: $native_type = self.clone().into();
-                buffer[0..Self::msg_size()].copy_from_slice(&t.to_le_bytes());
+                buffer[0..self.msg_size()].copy_from_slice(&t.to_le_bytes());
                 Ok(0)
             }
         }
@@ -331,30 +360,85 @@ macro_rules! array_impls {
     ($N:expr, $t: ident $($ts:ident)*)
     => {
         impl<T: MsgOnSocket + Clone> MsgOnSocket for [T; $N] {
-            fn msg_size() -> usize {
-                T::msg_size() * $N
+            fn uses_fd() -> bool {
+                T::uses_fd()
             }
-            fn max_fd_count() -> usize {
-                T::max_fd_count() * $N
+
+            fn fixed_size() -> Option<usize> {
+                Some(T::fixed_size()? * $N)
             }
-            unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
-                if buffer.len() < Self::msg_size() {
-                    return Err(MsgError::WrongMsgBufferSize);
+
+            fn msg_size(&self) -> usize {
+                match T::fixed_size() {
+                    Some(s) => s * $N,
+                    None => self.iter().map(|i| i.msg_size()).sum::<usize>() + size_of::<u64>() * $N
                 }
+            }
+
+            fn fd_count(&self) -> usize {
+                if T::uses_fd() {
+                    self.iter().map(|i| i.fd_count()).sum()
+                } else {
+                    0
+                }
+            }
+
+            unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
+                // Taken from the canonical example of initializing an array, the `assume_init` can
+                // be assumed safe because the array elements (`MaybeUninit<T>` in this case)
+                // themselves don't require initializing.
+                let mut msgs: [MaybeUninit<T>; $N] =  MaybeUninit::uninit().assume_init();
+
                 let mut offset = 0usize;
                 let mut fd_offset = 0usize;
-                let ($t, fd_size) =
-                    T::read_from_buffer(&buffer[offset..], &fds[fd_offset..])?;
-                offset += T::msg_size();
-                fd_offset += fd_size;
-                $(
-                    let ($ts, fd_size) =
-                        T::read_from_buffer(&buffer[offset..], &fds[fd_offset..])?;
-                    offset += T::msg_size();
-                    fd_offset += fd_size;
-                    )*
-                assert_eq!(offset, Self::msg_size());
-                Ok(([$t, $($ts),*], fd_offset))
+
+                // In case of an error, we need to keep track of how many elements got initialized.
+                // In order to perform the necessary drops, the below loop is executed in a closure
+                // to capture errors without returning.
+                let mut last_index = 0;
+                let res = (|| {
+                    for msg in &mut msgs[..] {
+                        let element_size = match T::fixed_size() {
+                            Some(s) => s,
+                            None => {
+                                let (element_size, _) = u64::read_from_buffer(&buffer[offset..], &[])?;
+                                offset += element_size.msg_size();
+                                element_size as usize
+                            }
+                        };
+                        let (m, fd_size) =
+                            T::read_from_buffer(&buffer[offset..], &fds[fd_offset..])?;
+                        *msg = MaybeUninit::new(m);
+                        offset += element_size;
+                        fd_offset += fd_size;
+                        last_index += 1;
+                    }
+                    Ok(())
+                })();
+
+                // Because `MaybeUninit` will not automatically call drops, we have to drop the
+                // partially initialized array manually in the case of an error.
+                if let Err(e) = res {
+                    for msg in &mut msgs[..last_index] {
+                        // The call to `as_mut_ptr()` turns the `MaybeUninit` element of the array
+                        // into a pointer, which can be used with `drop_in_place` to call the
+                        // destructor without moving the element, which is impossible. This is safe
+                        // because `last_index` prevents this loop from traversing into the
+                        // uninitialized parts of the array.
+                        drop_in_place(msg.as_mut_ptr());
+                    }
+                    return Err(e)
+                }
+
+                // Also taken from the canonical example, we initialized every member of the array
+                // in the first loop of this function, so it is safe to `transmute_copy` the array
+                // of `MaybeUninit` data to plain data. Although `transmute`, which checks the
+                // types' sizes, would have been preferred in this code, the compiler complains with
+                // "cannot transmute between types of different sizes, or dependently-sized types."
+                // Because this function operates on generic data, the type is "dependently-sized"
+                // and so the compiler will not check that the size of the input and output match.
+                // See this issue for details: https://github.com/rust-lang/rust/issues/61956
+                Ok((transmute_copy::<_, [T; $N]>(&msgs), fd_offset))
             }
 
             fn write_to_buffer(
@@ -362,21 +446,53 @@ macro_rules! array_impls {
                 buffer: &mut [u8],
                 fds: &mut [RawFd],
                 ) -> MsgResult<usize> {
-                if buffer.len() < Self::msg_size() {
-                    return Err(MsgError::WrongMsgBufferSize);
-                }
                 let mut offset = 0usize;
                 let mut fd_offset = 0usize;
                 for idx in 0..$N {
-                    let fd_size = self[idx].clone().write_to_buffer(&mut buffer[offset..],
+                    let element_size = match T::fixed_size() {
+                        Some(s) => s,
+                        None => {
+                            let element_size = self[idx].msg_size() as u64;
+                            element_size.write_to_buffer(&mut buffer[offset..], &mut [])?;
+                            offset += element_size.msg_size();
+                            element_size as usize
+                        }
+                    };
+                    let fd_size = self[idx].write_to_buffer(&mut buffer[offset..],
                                                             &mut fds[fd_offset..])?;
-                    offset += T::msg_size();
+                    offset += element_size;
                     fd_offset += fd_size;
                 }
 
                 Ok(fd_offset)
             }
         }
+        #[cfg(test)]
+        mod $t {
+            use super::MsgOnSocket;
+
+            #[test]
+            fn read_write_option_array() {
+                type ArrayType = [Option<u32>; $N];
+                let array = [Some($N); $N];
+                let mut buffer = vec![0; array.msg_size()];
+                array.write_to_buffer(&mut buffer, &mut []).unwrap();
+                let read_array = unsafe { ArrayType::read_from_buffer(&buffer, &[]) }.unwrap().0;
+
+                assert_eq!(array, read_array);
+            }
+
+            #[test]
+            fn read_write_fixed() {
+                type ArrayType = [u32; $N];
+                let mut buffer = vec![0; <ArrayType>::fixed_size().unwrap()];
+                let array = [$N as u32; $N];
+                array.write_to_buffer(&mut buffer, &mut []).unwrap();
+                let read_array = unsafe { ArrayType::read_from_buffer(&buffer, &[]) }.unwrap().0;
+
+                assert_eq!(array, read_array);
+            }
+        }
         array_impls!(($N - 1), $($ts)*);
     };
     {$N:expr, } => {};
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index 92af2c6..cb836be 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -43,10 +43,13 @@ impl AsRawFd for MaybeOwnedFd {
 
 // When sent, it could be owned or borrowed. On the receiver end, it always owned.
 impl MsgOnSocket for MaybeOwnedFd {
-    fn msg_size() -> usize {
-        0usize
+    fn uses_fd() -> bool {
+        true
     }
-    fn max_fd_count() -> usize {
+    fn fixed_size() -> Option<usize> {
+        Some(0)
+    }
+    fn fd_count(&self) -> usize {
         1usize
     }
     unsafe fn read_from_buffer(buffer: &[u8], fds: &[RawFd]) -> MsgResult<(Self, usize)> {
-- 
cgit 1.4.1


From 252d5b3cf3fd7a48fe9d610b59e3d6da9f2c6fe9 Mon Sep 17 00:00:00 2001
From: Dylan Reid <dgreid@chromium.org>
Date: Fri, 27 Mar 2020 19:07:59 +0000
Subject: handle mmap of large offsets on 32 bit systems

While only 32 bits of address can be mapped, that 32 bits can be offset
by further than 32 bits in to a large file. As chirantan points out, the
try_mmap call was already casting the usize to u64 on all architectures.

Convert the usize offset in mmap to u64 and address users of the API as
well.

Change-Id: I67aed928ea521049fb51eb7aa61ea4de8b4d096c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2124879
Reviewed-by: Dylan Reid <dgreid@chromium.org>
Tested-by: Dylan Reid <dgreid@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Dylan Reid <dgreid@chromium.org>
---
 devices/src/pci/ac97_bus_master.rs | 10 ++++++++--
 devices/src/pci/vfio_pci.rs        |  2 +-
 src/plugin/process.rs              |  2 +-
 sys_util/src/guest_memory.rs       | 10 +++++-----
 sys_util/src/mmap.rs               | 22 +++++++++++-----------
 vm_control/src/lib.rs              |  4 ++--
 6 files changed, 28 insertions(+), 22 deletions(-)

(limited to 'vm_control/src')

diff --git a/devices/src/pci/ac97_bus_master.rs b/devices/src/pci/ac97_bus_master.rs
index 5f4ca75..809f31f 100644
--- a/devices/src/pci/ac97_bus_master.rs
+++ b/devices/src/pci/ac97_bus_master.rs
@@ -5,6 +5,7 @@
 use std;
 use std::collections::VecDeque;
 use std::convert::AsRef;
+use std::convert::TryInto;
 use std::error::Error;
 use std::fmt::{self, Display};
 use std::os::unix::io::{AsRawFd, RawFd};
@@ -106,6 +107,8 @@ type GuestMemoryResult<T> = std::result::Result<T, GuestMemoryError>;
 enum AudioError {
     // Failed to create a new stream.
     CreateStream(Box<dyn Error>),
+    // Invalid buffer offset received from the audio server.
+    InvalidBufferOffset,
     // Guest did not provide a buffer when needed.
     NoBufferAvailable,
     // Failure to read guest memory.
@@ -124,6 +127,7 @@ impl Display for AudioError {
 
         match self {
             CreateStream(e) => write!(f, "Failed to create audio stream: {}.", e),
+            InvalidBufferOffset => write!(f, "Offset > max usize"),
             NoBufferAvailable => write!(f, "No buffer was available from the Guest"),
             ReadingGuestError(e) => write!(f, "Failed to read guest memory: {}.", e),
             RespondRequest(e) => write!(f, "Failed to respond to the ServerRequest: {}", e),
@@ -610,7 +614,7 @@ fn get_buffer_offset(
     func_regs: &Ac97FunctionRegs,
     mem: &GuestMemory,
     index: u8,
-) -> GuestMemoryResult<usize> {
+) -> GuestMemoryResult<u64> {
     let descriptor_addr = func_regs.bdbar + u32::from(index) * DESCRIPTOR_LENGTH as u32;
     let buffer_addr_reg: u32 = mem
         .read_obj_from_addr(GuestAddress(u64::from(descriptor_addr)))
@@ -673,7 +677,9 @@ fn next_guest_buffer<'a>(
         return Ok(None);
     }
 
-    let offset = get_buffer_offset(func_regs, mem, index)?;
+    let offset = get_buffer_offset(func_regs, mem, index)?
+        .try_into()
+        .map_err(|_| AudioError::InvalidBufferOffset)?;
     let frames = get_buffer_samples(func_regs, mem, index)? / DEVICE_CHANNEL_COUNT;
 
     Ok(Some(GuestBuffer {
diff --git a/devices/src/pci/vfio_pci.rs b/devices/src/pci/vfio_pci.rs
index 4d5ad9e..d1b6f2b 100644
--- a/devices/src/pci/vfio_pci.rs
+++ b/devices/src/pci/vfio_pci.rs
@@ -709,7 +709,7 @@ impl VfioPciDevice {
                 let mmap_size = mmap.size;
                 let guest_map_start = bar_addr + mmap_offset;
                 let region_offset = self.device.get_region_offset(index);
-                let offset: usize = (region_offset + mmap_offset) as usize;
+                let offset = region_offset + mmap_offset;
                 if self
                     .vm_socket_mem
                     .send(&VmMemoryRequest::RegisterMmapMemory {
diff --git a/src/plugin/process.rs b/src/plugin/process.rs
index ea7a78c..51fc892 100644
--- a/src/plugin/process.rs
+++ b/src/plugin/process.rs
@@ -361,7 +361,7 @@ impl Process {
             None => return Err(SysError::new(EOVERFLOW)),
             _ => {}
         }
-        let mem = MemoryMapping::from_fd_offset(&shm, length as usize, offset as usize)
+        let mem = MemoryMapping::from_fd_offset(&shm, length as usize, offset)
             .map_err(mmap_to_sys_err)?;
         let slot = vm.add_mmio_memory(GuestAddress(start), mem, read_only, dirty_log)?;
         entry.insert(PluginObject::Memory {
diff --git a/sys_util/src/guest_memory.rs b/sys_util/src/guest_memory.rs
index 2390b92..e8f620b 100644
--- a/sys_util/src/guest_memory.rs
+++ b/sys_util/src/guest_memory.rs
@@ -84,7 +84,7 @@ impl Display for Error {
 struct MemoryRegion {
     mapping: MemoryMapping,
     guest_base: GuestAddress,
-    memfd_offset: usize,
+    memfd_offset: u64,
 }
 
 fn region_end(region: &MemoryRegion) -> GuestAddress {
@@ -175,7 +175,7 @@ impl GuestMemory {
                 memfd_offset: offset,
             });
 
-            offset += size;
+            offset += size as u64;
         }
 
         Ok(GuestMemory {
@@ -262,7 +262,7 @@ impl GuestMemory {
     ///  * memfd_offset: usize
     pub fn with_regions<F, E>(&self, mut cb: F) -> result::Result<(), E>
     where
-        F: FnMut(usize, GuestAddress, usize, usize, usize) -> result::Result<(), E>,
+        F: FnMut(usize, GuestAddress, usize, usize, u64) -> result::Result<(), E>,
     {
         for (index, region) in self.regions.iter().enumerate() {
             cb(
@@ -584,10 +584,10 @@ impl GuestMemory {
     ///                .expect("failed to get offset");
     /// assert_eq!(offset, 0x3500);
     /// ```
-    pub fn offset_from_base(&self, guest_addr: GuestAddress) -> Result<usize> {
+    pub fn offset_from_base(&self, guest_addr: GuestAddress) -> Result<u64> {
         for region in self.regions.iter() {
             if guest_addr >= region.guest_base && guest_addr < region_end(region) {
-                return Ok(region.memfd_offset + guest_addr.offset_from(region.guest_base) as usize);
+                return Ok(region.memfd_offset + guest_addr.offset_from(region.guest_base) as u64);
             }
         }
         Err(Error::InvalidGuestAddress(guest_addr))
diff --git a/sys_util/src/mmap.rs b/sys_util/src/mmap.rs
index d8ba6b6..006b0a8 100644
--- a/sys_util/src/mmap.rs
+++ b/sys_util/src/mmap.rs
@@ -164,7 +164,7 @@ impl MemoryMapping {
         MemoryMapping::from_fd_offset(fd, size, 0)
     }
 
-    pub fn from_fd_offset(fd: &dyn AsRawFd, size: usize, offset: usize) -> Result<MemoryMapping> {
+    pub fn from_fd_offset(fd: &dyn AsRawFd, size: usize, offset: u64) -> Result<MemoryMapping> {
         MemoryMapping::from_fd_offset_protection(fd, size, offset, Protection::read_write())
     }
 
@@ -177,7 +177,7 @@ impl MemoryMapping {
     pub fn from_fd_offset_populate(
         fd: &dyn AsRawFd,
         size: usize,
-        offset: usize,
+        offset: u64,
     ) -> Result<MemoryMapping> {
         MemoryMapping::from_fd_offset_flags(
             fd,
@@ -199,7 +199,7 @@ impl MemoryMapping {
     fn from_fd_offset_flags(
         fd: &dyn AsRawFd,
         size: usize,
-        offset: usize,
+        offset: u64,
         flags: c_int,
         prot: Protection,
     ) -> Result<MemoryMapping> {
@@ -220,7 +220,7 @@ impl MemoryMapping {
     pub fn from_fd_offset_protection(
         fd: &dyn AsRawFd,
         size: usize,
-        offset: usize,
+        offset: u64,
         prot: Protection,
     ) -> Result<MemoryMapping> {
         MemoryMapping::from_fd_offset_flags(fd, size, offset, libc::MAP_SHARED, prot)
@@ -261,7 +261,7 @@ impl MemoryMapping {
         addr: *mut u8,
         fd: &dyn AsRawFd,
         size: usize,
-        offset: usize,
+        offset: u64,
         prot: Protection,
     ) -> Result<MemoryMapping> {
         MemoryMapping::try_mmap(
@@ -280,7 +280,7 @@ impl MemoryMapping {
         size: usize,
         prot: c_int,
         flags: c_int,
-        fd: Option<(&dyn AsRawFd, usize)>,
+        fd: Option<(&dyn AsRawFd, u64)>,
     ) -> Result<MemoryMapping> {
         let mut flags = flags;
         // If addr is provided, set the FIXED flag, and validate addr alignment
@@ -297,7 +297,7 @@ impl MemoryMapping {
         // If fd is provided, validate fd offset is within bounds
         let (fd, offset) = match fd {
             Some((fd, offset)) => {
-                if offset > libc::off_t::max_value() as usize {
+                if offset > libc::off_t::max_value() as u64 {
                     return Err(Error::InvalidOffset);
                 }
                 (fd.as_raw_fd(), offset as libc::off_t)
@@ -714,7 +714,7 @@ impl MemoryMappingArena {
         offset: usize,
         size: usize,
         fd: &dyn AsRawFd,
-        fd_offset: usize,
+        fd_offset: u64,
     ) -> Result<()> {
         self.add_fd_offset_protection(offset, size, fd, fd_offset, Protection::read_write())
     }
@@ -734,7 +734,7 @@ impl MemoryMappingArena {
         offset: usize,
         size: usize,
         fd: &dyn AsRawFd,
-        fd_offset: usize,
+        fd_offset: u64,
         prot: Protection,
     ) -> Result<()> {
         self.try_add(offset, size, prot, Some((fd, fd_offset)))
@@ -747,7 +747,7 @@ impl MemoryMappingArena {
         offset: usize,
         size: usize,
         prot: Protection,
-        fd: Option<(&dyn AsRawFd, usize)>,
+        fd: Option<(&dyn AsRawFd, u64)>,
     ) -> Result<()> {
         self.validate_range(offset, size)?;
 
@@ -953,7 +953,7 @@ mod tests {
     #[test]
     fn from_fd_offset_invalid() {
         let fd = unsafe { std::fs::File::from_raw_fd(-1) };
-        let res = MemoryMapping::from_fd_offset(&fd, 4096, (libc::off_t::max_value() as usize) + 1)
+        let res = MemoryMapping::from_fd_offset(&fd, 4096, (libc::off_t::max_value() as u64) + 1)
             .unwrap_err();
         match res {
             Error::InvalidOffset => {}
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index cb836be..4d48ff4 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -235,7 +235,7 @@ pub enum VmMemoryRequest {
     RegisterMmapMemory {
         fd: MaybeOwnedFd,
         size: usize,
-        offset: usize,
+        offset: u64,
         gpa: u64,
     },
 }
@@ -303,7 +303,7 @@ impl VmMemoryRequest {
                 offset,
                 gpa,
             } => {
-                let mmap = match MemoryMapping::from_fd_offset(fd, size, offset) {
+                let mmap = match MemoryMapping::from_fd_offset(fd, size, offset as u64) {
                     Ok(v) => v,
                     Err(_e) => return VmMemoryResponse::Err(SysError::new(EINVAL)),
                 };
-- 
cgit 1.4.1


From e1980a9c360b04705a16434bdaf1a56161dafb56 Mon Sep 17 00:00:00 2001
From: Daniel Verkamp <dverkamp@chromium.org>
Date: Fri, 7 Feb 2020 11:00:55 -0800
Subject: devices: pmem: implement flush using msync()

Previously, writable pmem devices implemented the flush command using
fsync(); however, this does not guarantee synchronization of memory
mappings via mmap() to the file on disk.  What we actually need is
msync() on the pmem file mapping, but we don't have access to that
mapping in the pmem child process, and it isn't trivial to pass it along
since it is owned by the Vm object once it has been added as a
mmap_arena.

In order to call msync() on the mapping, add a new VmControl socket so
that the pmem device can request that the main process issues an msync()
on the MemoryMappingArena identified by its slot number.

BUG=chromium:1007535
TEST=mount filesystem on /dev/pmem0 and sync; verify msync in strace

Change-Id: Id0484757c422cf81d454fd54012a12dbcc1baaf6
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2044365
Reviewed-by: Stephen Barber <smbarber@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Daniel Verkamp <dverkamp@chromium.org>
---
 devices/src/virtio/pmem.rs | 60 ++++++++++++++++++++++++++++++++---------
 src/linux.rs               | 67 +++++++++++++++++++++++++++++++++++++++-------
 vm_control/src/lib.rs      | 47 ++++++++++++++++++++++++++++++++
 3 files changed, 152 insertions(+), 22 deletions(-)

(limited to 'vm_control/src')

diff --git a/devices/src/virtio/pmem.rs b/devices/src/virtio/pmem.rs
index 931b037..499e110 100644
--- a/devices/src/virtio/pmem.rs
+++ b/devices/src/virtio/pmem.rs
@@ -13,6 +13,10 @@ use sys_util::{error, EventFd, GuestAddress, GuestMemory, PollContext, PollToken
 
 use data_model::{DataInit, Le32, Le64};
 
+use msg_socket::{MsgReceiver, MsgSender};
+
+use vm_control::{VmMsyncRequest, VmMsyncRequestSocket, VmMsyncResponse};
+
 use super::{
     copy_config, DescriptorChain, DescriptorError, Interrupt, Queue, Reader, VirtioDevice, Writer,
     TYPE_PMEM, VIRTIO_F_VERSION_1,
@@ -83,19 +87,38 @@ struct Worker {
     interrupt: Interrupt,
     queue: Queue,
     memory: GuestMemory,
-    disk_image: File,
+    pmem_device_socket: VmMsyncRequestSocket,
+    mapping_arena_slot: u32,
 }
 
 impl Worker {
     fn execute_request(&self, request: virtio_pmem_req) -> u32 {
         match request.type_.to_native() {
-            VIRTIO_PMEM_REQ_TYPE_FLUSH => match self.disk_image.sync_all() {
-                Ok(()) => VIRTIO_PMEM_RESP_TYPE_OK,
-                Err(e) => {
-                    error!("failed flushing disk image: {}", e);
-                    VIRTIO_PMEM_RESP_TYPE_EIO
+            VIRTIO_PMEM_REQ_TYPE_FLUSH => {
+                let request = VmMsyncRequest::MsyncArena {
+                    slot: self.mapping_arena_slot,
+                    offset: 0, // The pmem backing file is always at offset 0 in the arena.
+                };
+
+                if let Err(e) = self.pmem_device_socket.send(&request) {
+                    error!("failed to send request: {}", e);
+                    return VIRTIO_PMEM_RESP_TYPE_EIO;
+                }
+
+                match self.pmem_device_socket.recv() {
+                    Ok(response) => match response {
+                        VmMsyncResponse::Ok => VIRTIO_PMEM_RESP_TYPE_OK,
+                        VmMsyncResponse::Err(e) => {
+                            error!("failed flushing disk image: {}", e);
+                            VIRTIO_PMEM_RESP_TYPE_EIO
+                        }
+                    },
+                    Err(e) => {
+                        error!("failed to receive data: {}", e);
+                        VIRTIO_PMEM_RESP_TYPE_EIO
+                    }
                 }
-            },
+            }
             _ => {
                 error!("unknown request type: {}", request.type_.to_native());
                 VIRTIO_PMEM_RESP_TYPE_EIO
@@ -199,21 +222,27 @@ pub struct Pmem {
     worker_thread: Option<thread::JoinHandle<()>>,
     disk_image: Option<File>,
     mapping_address: GuestAddress,
+    mapping_arena_slot: u32,
     mapping_size: u64,
+    pmem_device_socket: Option<VmMsyncRequestSocket>,
 }
 
 impl Pmem {
     pub fn new(
         disk_image: File,
         mapping_address: GuestAddress,
+        mapping_arena_slot: u32,
         mapping_size: u64,
+        pmem_device_socket: Option<VmMsyncRequestSocket>,
     ) -> SysResult<Pmem> {
         Ok(Pmem {
             kill_event: None,
             worker_thread: None,
             disk_image: Some(disk_image),
             mapping_address,
+            mapping_arena_slot,
             mapping_size,
+            pmem_device_socket,
         })
     }
 }
@@ -233,11 +262,15 @@ impl Drop for Pmem {
 
 impl VirtioDevice for Pmem {
     fn keep_fds(&self) -> Vec<RawFd> {
+        let mut keep_fds = Vec::new();
         if let Some(disk_image) = &self.disk_image {
-            vec![disk_image.as_raw_fd()]
-        } else {
-            vec![]
+            keep_fds.push(disk_image.as_raw_fd());
         }
+
+        if let Some(ref pmem_device_socket) = self.pmem_device_socket {
+            keep_fds.push(pmem_device_socket.as_raw_fd());
+        }
+        keep_fds
     }
 
     fn device_type(&self) -> u32 {
@@ -274,7 +307,9 @@ impl VirtioDevice for Pmem {
         let queue = queues.remove(0);
         let queue_event = queue_events.remove(0);
 
-        if let Some(disk_image) = self.disk_image.take() {
+        let mapping_arena_slot = self.mapping_arena_slot;
+
+        if let Some(pmem_device_socket) = self.pmem_device_socket.take() {
             let (self_kill_event, kill_event) =
                 match EventFd::new().and_then(|e| Ok((e.try_clone()?, e))) {
                     Ok(v) => v,
@@ -291,8 +326,9 @@ impl VirtioDevice for Pmem {
                     let mut worker = Worker {
                         interrupt,
                         memory,
-                        disk_image,
                         queue,
+                        pmem_device_socket,
+                        mapping_arena_slot,
                     };
                     worker.run(queue_event, kill_event);
                 });
diff --git a/src/linux.rs b/src/linux.rs
index 5bc2dcc..9dbdb5c 100644
--- a/src/linux.rs
+++ b/src/linux.rs
@@ -56,7 +56,8 @@ use vm_control::{
     BalloonControlResult, DiskControlCommand, DiskControlRequestSocket, DiskControlResponseSocket,
     DiskControlResult, UsbControlSocket, VmControlResponseSocket, VmIrqRequest, VmIrqResponse,
     VmIrqResponseSocket, VmMemoryControlRequestSocket, VmMemoryControlResponseSocket,
-    VmMemoryRequest, VmMemoryResponse, VmRunMode,
+    VmMemoryRequest, VmMemoryResponse, VmMsyncRequest, VmMsyncRequestSocket, VmMsyncResponse,
+    VmMsyncResponseSocket, VmRunMode,
 };
 
 use crate::{Config, DiskOption, Executable, SharedDir, SharedDirKind, TouchDeviceOption};
@@ -254,6 +255,7 @@ enum TaggedControlSocket {
     Vm(VmControlResponseSocket),
     VmMemory(VmMemoryControlResponseSocket),
     VmIrq(VmIrqResponseSocket),
+    VmMsync(VmMsyncResponseSocket),
 }
 
 impl AsRef<UnixSeqpacket> for TaggedControlSocket {
@@ -263,6 +265,7 @@ impl AsRef<UnixSeqpacket> for TaggedControlSocket {
             Vm(ref socket) => socket.as_ref(),
             VmMemory(ref socket) => socket.as_ref(),
             VmIrq(ref socket) => socket.as_ref(),
+            VmMsync(ref socket) => socket.as_ref(),
         }
     }
 }
@@ -874,6 +877,7 @@ fn create_pmem_device(
     resources: &mut SystemAllocator,
     disk: &DiskOption,
     index: usize,
+    pmem_device_socket: VmMsyncRequestSocket,
 ) -> DeviceResult {
     let fd = OpenOptions::new()
         .read(true)
@@ -935,16 +939,23 @@ fn create_pmem_device(
         )
         .map_err(Error::AllocatePmemDeviceAddress)?;
 
-    vm.add_mmap_arena(
+    let slot = vm
+        .add_mmap_arena(
+            GuestAddress(mapping_address),
+            arena,
+            /* read_only = */ disk.read_only,
+            /* log_dirty_pages = */ false,
+        )
+        .map_err(Error::AddPmemDeviceMemory)?;
+
+    let dev = virtio::Pmem::new(
+        fd,
         GuestAddress(mapping_address),
-        arena,
-        /* read_only = */ disk.read_only,
-        /* log_dirty_pages = */ false,
+        slot,
+        arena_size,
+        Some(pmem_device_socket),
     )
-    .map_err(Error::AddPmemDeviceMemory)?;
-
-    let dev = virtio::Pmem::new(fd, GuestAddress(mapping_address), arena_size)
-        .map_err(Error::PmemDeviceNew)?;
+    .map_err(Error::PmemDeviceNew)?;
 
     Ok(VirtioDeviceStub {
         dev: Box::new(dev) as Box<dyn VirtioDevice>,
@@ -964,6 +975,7 @@ fn create_virtio_devices(
     gpu_device_socket: VmMemoryControlRequestSocket,
     balloon_device_socket: BalloonControlResponseSocket,
     disk_device_sockets: &mut Vec<DiskControlResponseSocket>,
+    pmem_device_sockets: &mut Vec<VmMsyncRequestSocket>,
 ) -> DeviceResult<Vec<VirtioDeviceStub>> {
     let mut devs = Vec::new();
 
@@ -973,7 +985,15 @@ fn create_virtio_devices(
     }
 
     for (index, pmem_disk) in cfg.pmem_devices.iter().enumerate() {
-        devs.push(create_pmem_device(cfg, vm, resources, pmem_disk, index)?);
+        let pmem_device_socket = pmem_device_sockets.remove(0);
+        devs.push(create_pmem_device(
+            cfg,
+            vm,
+            resources,
+            pmem_disk,
+            index,
+            pmem_device_socket,
+        )?);
     }
 
     devs.push(create_rng_device(cfg)?);
@@ -1124,6 +1144,7 @@ fn create_devices(
     gpu_device_socket: VmMemoryControlRequestSocket,
     balloon_device_socket: BalloonControlResponseSocket,
     disk_device_sockets: &mut Vec<DiskControlResponseSocket>,
+    pmem_device_sockets: &mut Vec<VmMsyncRequestSocket>,
     usb_provider: HostBackendDeviceProvider,
 ) -> DeviceResult<Vec<(Box<dyn PciDevice>, Option<Minijail>)>> {
     let stubs = create_virtio_devices(
@@ -1136,6 +1157,7 @@ fn create_devices(
         gpu_device_socket,
         balloon_device_socket,
         disk_device_sockets,
+        pmem_device_sockets,
     )?;
 
     let mut pci_devices = Vec::new();
@@ -1606,6 +1628,15 @@ pub fn run_config(cfg: Config) -> Result<()> {
         disk_device_sockets.push(disk_device_socket);
     }
 
+    let mut pmem_device_sockets = Vec::new();
+    let pmem_count = cfg.pmem_devices.len();
+    for _ in 0..pmem_count {
+        let (pmem_host_socket, pmem_device_socket) =
+            msg_socket::pair::<VmMsyncResponse, VmMsyncRequest>().map_err(Error::CreateSocket)?;
+        pmem_device_sockets.push(pmem_device_socket);
+        control_sockets.push(TaggedControlSocket::VmMsync(pmem_host_socket));
+    }
+
     let (gpu_host_socket, gpu_device_socket) =
         msg_socket::pair::<VmMemoryResponse, VmMemoryRequest>().map_err(Error::CreateSocket)?;
     control_sockets.push(TaggedControlSocket::VmMemory(gpu_host_socket));
@@ -1633,6 +1664,7 @@ pub fn run_config(cfg: Config) -> Result<()> {
                 gpu_device_socket,
                 balloon_device_socket,
                 &mut disk_device_sockets,
+                &mut pmem_device_sockets,
                 usb_provider,
             )
         },
@@ -2031,6 +2063,21 @@ fn run_control(
                                     }
                                 }
                             },
+                            TaggedControlSocket::VmMsync(socket) => match socket.recv() {
+                                Ok(request) => {
+                                    let response = request.execute(&mut linux.vm);
+                                    if let Err(e) = socket.send(&response) {
+                                        error!("failed to send VmMsyncResponse: {}", e);
+                                    }
+                                }
+                                Err(e) => {
+                                    if let MsgError::BadRecvSize { actual: 0, .. } = e {
+                                        vm_control_indices_to_remove.push(index);
+                                    } else {
+                                        error!("failed to recv VmMsyncRequest: {}", e);
+                                    }
+                                }
+                            },
                         }
                     }
                 }
diff --git a/vm_control/src/lib.rs b/vm_control/src/lib.rs
index 4d48ff4..4035c1c 100644
--- a/vm_control/src/lib.rs
+++ b/vm_control/src/lib.rs
@@ -408,6 +408,50 @@ pub enum VmIrqResponse {
     Err(SysError),
 }
 
+#[derive(MsgOnSocket, Debug)]
+pub enum VmMsyncRequest {
+    /// Flush the content of a memory mapping to its backing file.
+    /// `slot` selects the arena (as returned by `Vm::add_mmap_arena`).
+    /// `offset` is the offset of the mapping to sync within the arena.
+    MsyncArena { slot: u32, offset: usize },
+}
+
+#[derive(MsgOnSocket, Debug)]
+pub enum VmMsyncResponse {
+    Ok,
+    Err(SysError),
+}
+
+impl VmMsyncRequest {
+    /// Executes this request on the given Vm.
+    ///
+    /// # Arguments
+    /// * `vm` - The `Vm` to perform the request on.
+    ///
+    /// This does not return a result, instead encapsulating the success or failure in a
+    /// `VmMsyncResponse` with the intended purpose of sending the response back over the socket
+    /// that received this `VmMsyncResponse`.
+    pub fn execute(&self, vm: &mut Vm) -> VmMsyncResponse {
+        use self::VmMsyncRequest::*;
+        match *self {
+            MsyncArena { slot, offset } => {
+                if let Some(arena) = vm.get_mmap_arena(slot) {
+                    match arena.msync(offset) {
+                        Ok(true) => VmMsyncResponse::Ok,
+                        Ok(false) => VmMsyncResponse::Err(SysError::new(EINVAL)),
+                        Err(e) => match e {
+                            MmapError::SystemCallFailed(errno) => VmMsyncResponse::Err(errno),
+                            _ => VmMsyncResponse::Err(SysError::new(EINVAL)),
+                        },
+                    }
+                } else {
+                    VmMsyncResponse::Err(SysError::new(EINVAL))
+                }
+            }
+        }
+    }
+}
+
 pub type BalloonControlRequestSocket = MsgSocket<BalloonControlCommand, BalloonControlResult>;
 pub type BalloonControlResponseSocket = MsgSocket<BalloonControlResult, BalloonControlCommand>;
 
@@ -422,6 +466,9 @@ pub type VmMemoryControlResponseSocket = MsgSocket<VmMemoryResponse, VmMemoryReq
 pub type VmIrqRequestSocket = MsgSocket<VmIrqRequest, VmIrqResponse>;
 pub type VmIrqResponseSocket = MsgSocket<VmIrqResponse, VmIrqRequest>;
 
+pub type VmMsyncRequestSocket = MsgSocket<VmMsyncRequest, VmMsyncResponse>;
+pub type VmMsyncResponseSocket = MsgSocket<VmMsyncResponse, VmMsyncRequest>;
+
 pub type VmControlRequestSocket = MsgSocket<VmRequest, VmResponse>;
 pub type VmControlResponseSocket = MsgSocket<VmResponse, VmRequest>;
 
-- 
cgit 1.4.1