From dd829afebd76123a2abe0b3cd7acc595ad7d9085 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Mon, 12 Nov 2018 15:39:55 +0000
Subject: [PATCH] firmware: fix another TOCTTOU race in sync/async RPC code.

---
 artiq/firmware/ksupport/lib.rs                | 13 +++++++++++++
 artiq/firmware/libproto_artiq/kernel_proto.rs |  1 +
 artiq/firmware/runtime/session.rs             | 10 ++++++++--
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/artiq/firmware/ksupport/lib.rs b/artiq/firmware/ksupport/lib.rs
index 5b3c2bdfd..74ab757ef 100644
--- a/artiq/firmware/ksupport/lib.rs
+++ b/artiq/firmware/ksupport/lib.rs
@@ -514,6 +514,19 @@ pub unsafe fn main() {
         attribute_writeback(typeinfo as *const ());
     }
 
+    // Make sure all async RPCs are processed before exiting.
+    // Otherwise, if the comms and kernel CPU run in the following sequence:
+    //
+    //    comms                     kernel
+    //    -----------------------   -----------------------
+    //    check for async RPC
+    //                              post async RPC
+    //                              post RunFinished
+    //    check for mailbox
+    //
+    // the async RPC would be missed.
+    send(&RpcFlush);
+
     send(&RunFinished);
 
     loop {}
diff --git a/artiq/firmware/libproto_artiq/kernel_proto.rs b/artiq/firmware/libproto_artiq/kernel_proto.rs
index cefc27ee5..f7fef6eb0 100644
--- a/artiq/firmware/libproto_artiq/kernel_proto.rs
+++ b/artiq/firmware/libproto_artiq/kernel_proto.rs
@@ -64,6 +64,7 @@ pub enum Message<'a> {
     },
     RpcRecvRequest(*mut ()),
     RpcRecvReply(Result<usize, Exception<'a>>),
+    RpcFlush,
 
     CacheGetRequest { key: &'a str },
     CacheGetReply   { value: &'static [i32] },
diff --git a/artiq/firmware/runtime/session.rs b/artiq/firmware/runtime/session.rs
index cab4ef25a..c0498ef4a 100644
--- a/artiq/firmware/runtime/session.rs
+++ b/artiq/firmware/runtime/session.rs
@@ -410,7 +410,13 @@ fn process_kern_message(io: &Io, aux_mutex: &Mutex,
                         kern_acknowledge()
                     }
                 }
-            }
+            },
+            &kern::RpcFlush => {
+                // See ksupport/lib.rs for the reason this request exists.
+                // We do not need to do anything here because of how the main loop is
+                // structured.
+                kern_acknowledge()
+            },
 
             &kern::CacheGetRequest { key } => {
                 let value = session.congress.cache.get(key);
@@ -581,7 +587,7 @@ fn respawn<F>(io: &Io, handle: &mut Option<ThreadHandle>, f: F)
     *handle = Some(io.spawn(16384, f))
 }
 
-pub fn thread(io: Io, aux_mutex: &Mutex, 
+pub fn thread(io: Io, aux_mutex: &Mutex,
         routing_table: &Urc<RefCell<drtio_routing::RoutingTable>>,
         up_destinations: &Urc<RefCell<[bool; drtio_routing::DEST_COUNT]>>) {
     let listener = TcpListener::new(&io, 65535);