Repository: project-oak/oak (rust)

Issue:
unused closure in handle visit usage
A [new version](https://github.com/project-oak/oak/pull/1597/files) of the Rust compiler now makes this an error:

```
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ error: unused closure that must be used
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║    --> /opt/my-project/oak_io/src/handle.rs:129:5
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     |
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 129 | /     msg.visit(|handle: &mut Handle| {
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 130 | |         handles.push(*handle);
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 131 | |         *handle = oak_abi::INVALID_HANDLE;
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 132 | |     });
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     | |_______^
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     |
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     = note: `-D unused-must-use` implied by `-D warnings`
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     = note: closures are lazy and do nothing unless called
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ error: unused closure that must be used
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║    --> /opt/my-project/oak_io/src/handle.rs:171:5
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     |
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 171 | /     msg.visit(|handle| {
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 172 | |         if let Some(to_inject) = handles.next() {
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 173 | |             *handle = *to_inject;
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 174 | |         } else {
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 175 | |             result = Err(OakError::ProtobufDecodeError(None));
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 176 | |         }
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ 177 | |     });
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     | |_______^
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     |
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║     = note: closures are lazy and do nothing unless called
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ error: aborting due to 2 previous errors
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║
tests ❯ cargo clippy ❯ ./oak_runtime/Cargo.toml ║ error: could not compile `oak_io`
```

I guess [the type of `visit` is a bit awkward](https://github.com/project-oak/oak/pull/1108#discussion_r440454021) in order to support recursion, perhaps it should be revisited.

Pull request:
Change signature of `HandleVisit` impl function to a fold.
# Checklist

- [x] Pull request affects core Oak functionality (e.g. runtime, SDK, ABI)
  - [x] I have written tests that cover the code changes.
  - [ ] I have checked that these tests are run by
        [Cloudbuild](/cloudbuild.yaml)
  - [x] I have updated [documentation](/docs/) accordingly.
  - [ ] I have raised an [issue](https://github.com/project-oak/oak/issues) to
        cover any TODOs and/or unfinished work.
- [ ] Pull request includes prototype/experimental work that is under
      construction.

Fixes #1599.

Patch:
diff --git a/docs/programming-oak.md b/docs/programming-oak.md
index f84df726a2..fd465eac5c 100644
--- a/docs/programming-oak.md
+++ b/docs/programming-oak.md
@@ -64,7 +64,7 @@ oak::entrypoint!(grpc_oak_main<ConfigMap> => |_receiver| {
     };
     let grpc_channel =
         oak::grpc::server::init("[::]:8080").expect("could not create gRPC server pseudo-Node");
-    oak::run_command_loop(node, grpc_channel.iter());
+    oak::run_command_loop(MachineLearningDispatcher::new(node), grpc_channel.iter());
 });
 ```
 <!-- prettier-ignore-end -->
@@ -579,7 +579,7 @@ async fn test_say_hello() {
     let runtime = oak_tests::run_single_module_default(MODULE_WASM_FILE_NAME)
         .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = HelloWorldClient::with_interceptor(channel, interceptor);
 
     let req = HelloRequest {
diff --git a/examples/Cargo.lock b/examples/Cargo.lock
index 220491ee72..783caaa4dd 100644
--- a/examples/Cargo.lock
+++ b/examples/Cargo.lock
@@ -912,6 +912,7 @@ dependencies = [
  "oak",
  "oak_abi",
  "oak_io",
+ "oak_utils",
  "prost",
  "rand 0.7.3",
  "rand_distr",
@@ -1157,6 +1158,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/examples/machine_learning/module/rust/Cargo.toml b/examples/machine_learning/module/rust/Cargo.toml
index 889d38047d..ba33c2a6fc 100644
--- a/examples/machine_learning/module/rust/Cargo.toml
+++ b/examples/machine_learning/module/rust/Cargo.toml
@@ -17,3 +17,6 @@ oak_abi = "=0.1.0"
 oak_io = "=0.1.0"
 log = "*"
 prost = "*"
+
+[build-dependencies]
+oak_utils = "*"
diff --git a/examples/machine_learning/module/rust/build.rs b/examples/machine_learning/module/rust/build.rs
new file mode 100644
index 0000000000..9b4ee9a02e
--- /dev/null
+++ b/examples/machine_learning/module/rust/build.rs
@@ -0,0 +1,22 @@
+//
+// Copyright 2019 The Project Oak Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+fn main() {
+    oak_utils::compile_protos(
+        &["../../proto/machine_learning.proto"],
+        &["../../proto", "../../../../third_party"],
+    );
+}
diff --git a/examples/machine_learning/module/rust/src/lib.rs b/examples/machine_learning/module/rust/src/lib.rs
index 1d719facaa..cac8ceb3f7 100644
--- a/examples/machine_learning/module/rust/src/lib.rs
+++ b/examples/machine_learning/module/rust/src/lib.rs
@@ -17,12 +17,20 @@
 // This model was inspired by
 // https://github.com/AtheMathmo/rusty-machine/blob/master/examples/naive_bayes_dogs.rs .
 
-use log::{error, info, warn};
+pub mod proto {
+    include!(concat!(
+        env!("OUT_DIR"),
+        "/oak.examples.machine_learning.rs"
+    ));
+}
+
+use log::{info, warn};
 use oak::{
     grpc,
     io::{Receiver, ReceiverExt},
 };
 use oak_abi::proto::oak::application::ConfigMap;
+use proto::{MachineLearning, MachineLearningDispatcher, MlData, MlLearn, MlPredict, MlResponse};
 use rand::prelude::*;
 use rand_distr::{Distribution, Normal, Standard};
 use rusty_machine::{
@@ -158,7 +166,7 @@ oak::entrypoint!(oak_main<grpc::Invocation> => |receiver: Receiver<grpc::Invocat
         config: None,
         model: NaiveBayes::new(),
     };
-    oak::run_command_loop(node, receiver.iter());
+    oak::run_command_loop(MachineLearningDispatcher::new(node), receiver.iter());
 });
 
 oak::entrypoint!(grpc_oak_main<ConfigMap> => |_receiver| {
@@ -171,7 +179,7 @@ oak::entrypoint!(grpc_oak_main<ConfigMap> => |_receiver| {
     };
     let grpc_channel =
         oak::grpc::server::init("[::]:8080").expect("could not create gRPC server pseudo-Node");
-    oak::run_command_loop(node, grpc_channel.iter());
+    oak::run_command_loop(MachineLearningDispatcher::new(node), grpc_channel.iter());
 });
 
 struct Node {
@@ -181,107 +189,102 @@ struct Node {
     model: NaiveBayes<naive_bayes::Gaussian>,
 }
 
-impl oak::grpc::ServerNode for Node {
-    fn invoke(&mut self, method: &str, _req: &[u8], writer: grpc::ChannelResponseWriter) {
-        match method {
-            "/oak.examples.machine_learning.MachineLearning/Data" => {
-                info!("Data");
-                //(self.training_set_size, self.test_set_size) = (1000, 1000);
-                // Generate all of our train and test data
-                let (training_matrix, target_matrix, test_matrix, test_animals) =
-                    generate_animal_data(self.training_set_size, self.test_set_size);
-                self.config = Some(Config {
-                    training_matrix,
-                    target_matrix,
-                    test_matrix,
-                    test_animals,
-                });
+impl MachineLearning for Node {
+    fn data(&mut self, _req: MlData) -> grpc::Result<MlResponse> {
+        info!("Data");
+        //(self.training_set_size, self.test_set_size) = (1000, 1000);
+        // Generate all of our train and test data
+        let (training_matrix, target_matrix, test_matrix, test_animals) =
+            generate_animal_data(self.training_set_size, self.test_set_size);
+        self.config = Some(Config {
+            training_matrix,
+            target_matrix,
+            test_matrix,
+            test_animals,
+        });
+        Ok(MlResponse::default())
+    }
+    fn learn(&mut self, _req: MlLearn) -> grpc::Result<MlResponse> {
+        info!("Training model");
+        //self.model = NaiveBayes::<naive_bayes::Gaussian>::new();
+        match self.config {
+            Some(ref c) => self
+                .model
+                .train(&c.training_matrix, &c.target_matrix)
+                .expect("failed to train model of dogs"),
+            None => {
+                warn!("config not set");
             }
-            "/oak.examples.machine_learning.MachineLearning/Learn" => {
-                info!("Training model");
-                //self.model = NaiveBayes::<naive_bayes::Gaussian>::new();
-                match self.config {
-                    Some(ref c) => self
-                        .model
-                        .train(&c.training_matrix, &c.target_matrix)
-                        .expect("failed to train model of dogs"),
-                    None => {
-                        warn!("config not set");
-                    }
-                }
+        }
+        Ok(MlResponse::default())
+    }
+    fn predict(&mut self, _req: MlPredict) -> grpc::Result<MlResponse> {
+        info!("Predicting");
+        let mut predictions = None;
+        match self.config {
+            Some(ref c) => {
+                predictions = Some(
+                    self.model
+                        .predict(&c.test_matrix)
+                        .expect("failed to predict dogs!?"),
+                )
             }
-            "/oak.examples.machine_learning.MachineLearning/Predict" => {
-                info!("Predicting");
-                let mut predictions = None;
-                match self.config {
-                    Some(ref c) => {
-                        predictions = Some(
-                            self.model
-                                .predict(&c.test_matrix)
-                                .expect("failed to predict dogs!?"),
-                        )
-                    }
-                    None => {
-                        warn!("config not set");
-                    }
-                }
-                // Score how well we did.
-                let mut hits = 0;
-                let unprinted_total = self.test_set_size.saturating_sub(10) as usize;
-                match self.config {
-                    Some(ref c) => {
-                        if let Some(ref p) = predictions {
-                            for (animal, prediction) in c
-                                .test_animals
-                                .iter()
-                                .zip(p.iter_rows())
-                                .take(unprinted_total)
-                            {
-                                evaluate_prediction(&mut hits, animal, prediction);
-                            }
-                        }
-                    }
-                    None => {
-                        warn!("test_animals not set");
-                        return;
+            None => {
+                warn!("config not set");
+            }
+        }
+        // Score how well we did.
+        let mut hits = 0;
+        let unprinted_total = self.test_set_size.saturating_sub(10) as usize;
+        match self.config {
+            Some(ref c) => {
+                if let Some(ref p) = predictions {
+                    for (animal, prediction) in c
+                        .test_animals
+                        .iter()
+                        .zip(p.iter_rows())
+                        .take(unprinted_total)
+                    {
+                        evaluate_prediction(&mut hits, animal, prediction);
                     }
                 }
+            }
+            None => {
+                warn!("test_animals not set");
+                return Err(grpc::build_status(
+                    grpc::Code::InvalidArgument,
+                    "test_animals not set",
+                ));
+            }
+        }
 
-                if unprinted_total > 0 {
-                    info!("...");
-                }
+        if unprinted_total > 0 {
+            info!("...");
+        }
 
-                if let Some(ref c) = self.config {
-                    if let Some(ref p) = predictions {
-                        for (animal, prediction) in c
-                            .test_animals
-                            .iter()
-                            .zip(p.iter_rows())
-                            .skip(unprinted_total)
-                        {
-                            let (actual_type, accurate) =
-                                evaluate_prediction(&mut hits, animal, prediction);
-                            info!(
-                                "Predicted: {:?}; Actual: {:?}; Accurate? {:?}",
-                                animal.type_, actual_type, accurate
-                            );
-                        }
-                    }
+        if let Some(ref c) = self.config {
+            if let Some(ref p) = predictions {
+                for (animal, prediction) in c
+                    .test_animals
+                    .iter()
+                    .zip(p.iter_rows())
+                    .skip(unprinted_total)
+                {
+                    let (actual_type, accurate) =
+                        evaluate_prediction(&mut hits, animal, prediction);
+                    info!(
+                        "Predicted: {:?}; Actual: {:?}; Accurate? {:?}",
+                        animal.type_, actual_type, accurate
+                    );
                 }
-                info!(
-                    "Accuracy: {}/{} = {:.1}%",
-                    hits,
-                    self.test_set_size,
-                    (f64::from(hits)) / (f64::from(self.test_set_size as u32)) * 100.
-                );
-            }
-            _ => {
-                error!("unknown method name: {}", method);
-                panic!("unknown method name");
             }
         }
-        writer
-            .write_empty(grpc::WriteMode::Close)
-            .expect("Failed to write empty response");
+        info!(
+            "Accuracy: {}/{} = {:.1}%",
+            hits,
+            self.test_set_size,
+            (f64::from(hits)) / (f64::from(self.test_set_size as u32)) * 100.
+        );
+        Ok(MlResponse::default())
     }
 }
diff --git a/experimental/Cargo.lock b/experimental/Cargo.lock
index a43c2d467c..2420b3a9c9 100644
--- a/experimental/Cargo.lock
+++ b/experimental/Cargo.lock
@@ -835,6 +835,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/experimental/benchmark/src/application/oak.rs b/experimental/benchmark/src/application/oak.rs
index 93095d1300..d38f9ccfc0 100644
--- a/experimental/benchmark/src/application/oak.rs
+++ b/experimental/benchmark/src/application/oak.rs
@@ -60,7 +60,7 @@ impl OakApplication {
         let runtime = oak_runtime::configure_and_run(config)
             .expect("Couldn't configure runtime with test wasm");
 
-        let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+        let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
         let client = TrustedDatabaseClient::with_interceptor(channel, interceptor);
 
         OakApplication { runtime, client }
diff --git a/oak_derive/Cargo.lock b/oak_derive/Cargo.lock
index bcac1bdd32..6470d568b9 100644
--- a/oak_derive/Cargo.lock
+++ b/oak_derive/Cargo.lock
@@ -138,6 +138,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/oak_derive/src/lib.rs b/oak_derive/src/lib.rs
index 9d1d2d01cc..66b94a0498 100644
--- a/oak_derive/src/lib.rs
+++ b/oak_derive/src/lib.rs
@@ -52,11 +52,11 @@ fn struct_impls(name: &Ident, data: &syn::DataStruct) -> TokenStream {
             .collect(),
         Fields::Unit => Vec::new(),
     };
-    let body = accessors_visit(&accessors);
+    let body = accessors_fold(&accessors);
 
     quote! {
         impl ::oak_io::handle::HandleVisit for #name {
-            fn visit<F: FnMut(&mut ::oak_io::Handle)>(&mut self, visitor: F) -> F {
+            fn fold<B>(&mut self, init: B, f: fn(B, &mut ::oak_io::Handle) -> B) -> B {
                 #body
             }
         }
@@ -67,7 +67,7 @@ fn enum_impls(name: &Ident, data: &syn::DataEnum) -> TokenStream {
     let variants: Vec<TokenStream> = data.variants.iter().map(variant_impl).collect();
     quote! {
         impl ::oak_io::handle::HandleVisit for #name {
-            fn visit<F: FnMut(&mut ::oak_io::Handle)>(&mut self, visitor: F) -> F {
+            fn fold<B>(&mut self, init: B, f: fn(B, &mut ::oak_io::Handle) -> B) -> B {
                 match self {
                     #(
                         #name::#variants,
@@ -83,7 +83,7 @@ fn variant_impl(variant: &syn::Variant) -> TokenStream {
     match &variant.fields {
         Fields::Named(fields) => {
             let fields: Vec<Ident> = fields.named.iter().flat_map(|f| f.ident.clone()).collect();
-            let body = accessors_visit(&fields);
+            let body = accessors_fold(&fields);
             quote! {
                 #variant_ident { #( #fields ),* } => {
                     #body
@@ -98,7 +98,7 @@ fn variant_impl(variant: &syn::Variant) -> TokenStream {
                 .enumerate()
                 .map(|(i, _)| format_ident!("_{}", i))
                 .collect();
-            let body = accessors_visit(&accessors);
+            let body = accessors_fold(&accessors);
             quote! {
                 #variant_ident( #( #accessors ),* ) => {
                     #body
@@ -106,17 +106,17 @@ fn variant_impl(variant: &syn::Variant) -> TokenStream {
             }
         }
         Fields::Unit => quote! {
-            #variant_ident => visitor
+            #variant_ident => init
         },
     }
 }
 
-fn accessors_visit<T: ToTokens>(accessors: &[T]) -> TokenStream {
+fn accessors_fold<T: ToTokens>(accessors: &[T]) -> TokenStream {
     quote! {
-        let mut _v = visitor;
+        let accumulator = init;
         #(
-            _v  = #accessors.visit(_v);
+            let accumulator = #accessors.fold(accumulator, f);
         )*
-        _v
+        accumulator
     }
 }
diff --git a/oak_io/Cargo.lock b/oak_io/Cargo.lock
index 9fcf522214..6fc8fde95a 100644
--- a/oak_io/Cargo.lock
+++ b/oak_io/Cargo.lock
@@ -144,6 +144,7 @@ dependencies = [
  "assert_matches",
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/oak_io/Cargo.toml b/oak_io/Cargo.toml
index 7e4cec5ece..26c20b87c2 100644
--- a/oak_io/Cargo.toml
+++ b/oak_io/Cargo.toml
@@ -14,6 +14,7 @@ oak_derive = { path = "../oak_derive" }
 prost = "*"
 prost-types = "*"
 either = "*"
+log = "*"
 
 [dev-dependencies]
 assert_matches = "*"
diff --git a/oak_io/src/handle.rs b/oak_io/src/handle.rs
index a41ae7a576..efdf2a53de 100644
--- a/oak_io/src/handle.rs
+++ b/oak_io/src/handle.rs
@@ -16,6 +16,7 @@
 include!(concat!(env!("OUT_DIR"), "/oak.handle.rs"));
 
 use crate::OakError;
+use log::error;
 use oak_abi::Handle;
 
 /// Wrapper for a handle to the read half of a channel.
@@ -82,9 +83,8 @@ impl From<Handle> for WriteHandle {
 /// }
 ///
 /// impl HandleVisit for Thing {
-///     fn visit<F: FnMut(&mut Handle)>(&mut self, mut visitor: F) -> F {
-///         visitor(&mut self.handle);
-///         visitor
+///     fn fold<B>(&mut self, init: B, mut f: fn(B, &mut Handle) -> B) -> B {
+///         f(init, &mut self.handle)
 ///     }
 /// }
 /// ```
@@ -92,7 +92,7 @@ pub trait HandleVisit {
     /// Invokes the provided closure on every handle contained in `self`.
     ///
     /// The mutable reference allows modifying the handles.
-    fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F;
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B;
 }
 
 /// Return all handles in `T`.
@@ -112,9 +112,8 @@ pub trait HandleVisit {
 ///     handle: Handle,
 /// }
 /// # impl HandleVisit for Thing {
-/// #   fn visit<F: FnMut(&mut Handle)>(&mut self, mut visitor: F) -> F {
-/// #     visitor(&mut self.handle);
-/// #     visitor
+/// #   fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
+/// #       f(init, &mut self.handle)
 /// #   }
 /// # }
 ///
@@ -125,13 +124,11 @@ pub trait HandleVisit {
 /// assert_eq!(handles, vec![42]);
 /// ```
 pub fn extract_handles<T: HandleVisit>(msg: &mut T) -> Vec<Handle> {
-    let mut handles = Vec::new();
-    // TODO(#1599): Remove the `let _` when the underlying issue is fixed.
-    let _ = msg.visit(|handle: &mut Handle| {
+    msg.fold(Vec::new(), |mut handles, handle| {
         handles.push(*handle);
         *handle = oak_abi::INVALID_HANDLE;
-    });
-    handles
+        handles
+    })
 }
 
 /// Inject handles into a message.
@@ -153,9 +150,8 @@ pub fn extract_handles<T: HandleVisit>(msg: &mut T) -> Vec<Handle> {
 ///     handle: Handle,
 /// }
 /// # impl HandleVisit for Thing {
-/// #   fn visit<F: FnMut(&mut Handle)>(&mut self, mut visitor: F) -> F {
-/// #     visitor(&mut self.handle);
-/// #     visitor
+/// #   fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
+/// #       f(init, &mut self.handle)
 /// #   }
 /// # }
 ///
@@ -167,20 +163,24 @@ pub fn extract_handles<T: HandleVisit>(msg: &mut T) -> Vec<Handle> {
 /// assert_eq!(thing, Thing { handle: 42 });
 /// ```
 pub fn inject_handles<T: HandleVisit>(msg: &mut T, handles: &[Handle]) -> Result<(), OakError> {
-    let mut handles = handles.iter();
-    let mut result = Ok(());
-    // TODO(#1599): Remove the `let _` when the underlying issue is fixed.
-    let _ = msg.visit(|handle| {
-        if let Some(to_inject) = handles.next() {
-            *handle = *to_inject;
+    msg.fold(Ok(handles.iter()), |handles, handle| {
+        let mut handles = handles?;
+        let to_inject = handles.next().ok_or_else(|| {
+            error!("Not enough handles provided to populate message");
+            OakError::ProtobufDecodeError(None)
+        })?;
+        *handle = *to_inject;
+        Ok(handles)
+    })
+    // Check that there are no remaining handles
+    .and_then(|mut remaining_handles| {
+        if remaining_handles.next().is_some() {
+            error!("Too many handles provided for message",);
+            Err(OakError::ProtobufDecodeError(None))
         } else {
-            result = Err(OakError::ProtobufDecodeError(None));
+            Ok(())
         }
-    });
-    if handles.next().is_some() {
-        result = Err(OakError::ProtobufDecodeError(None));
-    }
-    result
+    })
 }
 
 // A default implementation of the HandleVisit trait that does nothing
@@ -189,8 +189,8 @@ macro_rules! handle_visit_blanket_impl {
     ($($t:ty),+) => {
         $(
             impl HandleVisit for $t {
-                fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F {
-                    visitor
+                fn fold<B>(&mut self, init: B, _: fn(B, &mut Handle) -> B) -> B {
+                    init
                 }
             }
         )+
@@ -220,27 +220,26 @@ pub use oak_derive::HandleVisit;
 
 // Optional fields
 impl<T: HandleVisit> HandleVisit for Option<T> {
-    fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F {
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
         if let Some(inner) = self {
-            inner.visit(visitor)
+            inner.fold(init, f)
         } else {
-            visitor
+            init
         }
     }
 }
 
 // For repeated fields.
 impl<T: HandleVisit> HandleVisit for Vec<T> {
-    fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F {
-        self.iter_mut()
-            .fold(visitor, |visitor, item| item.visit(visitor))
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
+        self.iter_mut().fold(init, |init, item| item.fold(init, f))
     }
 }
 
 // For recursive messages.
 impl<T: HandleVisit> HandleVisit for Box<T> {
-    fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F {
-        self.as_mut().visit(visitor)
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
+        self.as_mut().fold(init, f)
     }
 }
 
@@ -252,13 +251,13 @@ impl<T: HandleVisit> HandleVisit for Box<T> {
 impl<K: Ord + core::hash::Hash, V: HandleVisit, S> HandleVisit
     for std::collections::HashMap<K, V, S>
 {
-    fn visit<F: FnMut(&mut Handle)>(&mut self, visitor: F) -> F {
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
         let mut entries: Vec<(&K, &mut V)> = self.iter_mut().collect();
         // Can be unstable because keys are guaranteed to be unique.
         entries.sort_unstable_by_key(|&(k, _)| k);
         entries
             .into_iter()
             .map(|(_, v)| v)
-            .fold(visitor, |visitor, value| value.visit(visitor))
+            .fold(init, |init, value| value.fold(init, f))
     }
 }
diff --git a/oak_io/src/receiver.rs b/oak_io/src/receiver.rs
index 34f543865a..70c712b5ba 100644
--- a/oak_io/src/receiver.rs
+++ b/oak_io/src/receiver.rs
@@ -59,9 +59,8 @@ impl<T: Decodable> From<ReadHandle> for Receiver<T> {
 }
 
 impl<T: Decodable> crate::handle::HandleVisit for Receiver<T> {
-    fn visit<F: FnMut(&mut crate::Handle)>(&mut self, mut visitor: F) -> F {
-        visitor(&mut self.handle.handle);
-        visitor
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut crate::Handle) -> B) -> B {
+        f(init, &mut self.handle.handle)
     }
 }
 
diff --git a/oak_io/src/sender.rs b/oak_io/src/sender.rs
index f24f571a46..fd5be34332 100644
--- a/oak_io/src/sender.rs
+++ b/oak_io/src/sender.rs
@@ -42,9 +42,8 @@ impl<T: Encodable> Sender<T> {
 }
 
 impl<T: Encodable> crate::handle::HandleVisit for Sender<T> {
-    fn visit<F: FnMut(&mut crate::Handle)>(&mut self, mut visitor: F) -> F {
-        visitor(&mut self.handle.handle);
-        visitor
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut crate::Handle) -> B) -> B {
+        f(init, &mut self.handle.handle)
     }
 }
 
diff --git a/oak_loader/Cargo.lock b/oak_loader/Cargo.lock
index 2b6598ea6b..e4d07f6e6f 100644
--- a/oak_loader/Cargo.lock
+++ b/oak_loader/Cargo.lock
@@ -718,6 +718,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
@@ -1240,9 +1241,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.58"
+version = "1.0.59"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a230ea9107ca2220eea9d46de97eddcb04cd00e92d13dda78e478dd33fa82bd4"
+checksum = "dcac07dbffa1c65e7f816ab9eba78eb142c6d44410f4eeba1e26e4f5dfa56b95"
 dependencies = [
  "itoa",
  "ryu",
diff --git a/oak_runtime/Cargo.lock b/oak_runtime/Cargo.lock
index 6fc077d234..3a371d522c 100644
--- a/oak_runtime/Cargo.lock
+++ b/oak_runtime/Cargo.lock
@@ -730,6 +730,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/oak_services/Cargo.lock b/oak_services/Cargo.lock
index 65e7726c9b..0f9dc97f64 100644
--- a/oak_services/Cargo.lock
+++ b/oak_services/Cargo.lock
@@ -166,6 +166,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",
diff --git a/sdk/Cargo.lock b/sdk/Cargo.lock
index 86706d966c..1833a9c5bf 100644
--- a/sdk/Cargo.lock
+++ b/sdk/Cargo.lock
@@ -828,6 +828,7 @@ version = "0.1.0"
 dependencies = [
  "either",
  "hex",
+ "log",
  "oak_abi",
  "oak_derive",
  "oak_utils",


Test patch:
diff --git a/examples/abitest/module_0/rust/tests/integration_test.rs b/examples/abitest/module_0/rust/tests/integration_test.rs
index 56e366f096..116e2d23cb 100644
--- a/examples/abitest/module_0/rust/tests/integration_test.rs
+++ b/examples/abitest/module_0/rust/tests/integration_test.rs
@@ -58,7 +58,7 @@ async fn setup() -> (
     let runtime =
         oak_runtime::configure_and_run(config).expect("unable to configure runtime with test wasm");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let client = OakAbiTestServiceClient::with_interceptor(channel, interceptor);
 
     (runtime, client)
diff --git a/examples/aggregator/module/rust/tests/integration_test.rs b/examples/aggregator/module/rust/tests/integration_test.rs
index e37fc8511c..a454c6f80a 100644
--- a/examples/aggregator/module/rust/tests/integration_test.rs
+++ b/examples/aggregator/module/rust/tests/integration_test.rs
@@ -55,7 +55,7 @@ async fn test_aggregator() {
     )
     .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = AggregatorClient::with_interceptor(channel, interceptor);
 
     for i in 0..SAMPLE_THRESHOLD as u32 {
diff --git a/examples/chat/module/rust/tests/integration_test.rs b/examples/chat/module/rust/tests/integration_test.rs
index b44d29f68a..ece65d12b6 100644
--- a/examples/chat/module/rust/tests/integration_test.rs
+++ b/examples/chat/module/rust/tests/integration_test.rs
@@ -149,7 +149,7 @@ impl<'a> Chatter<'a> {
             user_handle,
             base64::encode(&room_key_pair.pkcs8_public_key())
         );
-        let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+        let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
         // TODO(#1357): Use key pair to authenticate this client and label requests.
         let client = ChatClient::with_interceptor(channel, interceptor);
         Chatter {
diff --git a/examples/hello_world/module/rust/tests/integration_test.rs b/examples/hello_world/module/rust/tests/integration_test.rs
index 77a2c9c16e..be0281fbc9 100644
--- a/examples/hello_world/module/rust/tests/integration_test.rs
+++ b/examples/hello_world/module/rust/tests/integration_test.rs
@@ -28,7 +28,7 @@ async fn test_say_hello() {
     let runtime = oak_tests::run_single_module_default(MODULE_WASM_FILE_NAME)
         .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = HelloWorldClient::with_interceptor(channel, interceptor);
 
     let req = HelloRequest {
diff --git a/examples/private_set_intersection/module/rust/tests/integration_test.rs b/examples/private_set_intersection/module/rust/tests/integration_test.rs
index d40d5ff459..77a612d5fa 100644
--- a/examples/private_set_intersection/module/rust/tests/integration_test.rs
+++ b/examples/private_set_intersection/module/rust/tests/integration_test.rs
@@ -31,7 +31,7 @@ async fn test_set_intersection() {
     let runtime = oak_tests::run_single_module_default(MODULE_WASM_FILE_NAME)
         .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = PrivateSetIntersectionClient::with_interceptor(channel, interceptor);
 
     let req = SubmitSetRequest {
diff --git a/examples/running_average/module/rust/tests/integration_test.rs b/examples/running_average/module/rust/tests/integration_test.rs
index 7227ff954d..6c7b1255a0 100644
--- a/examples/running_average/module/rust/tests/integration_test.rs
+++ b/examples/running_average/module/rust/tests/integration_test.rs
@@ -34,7 +34,7 @@ async fn test_running_average() {
     let runtime = oak_tests::run_single_module_default(MODULE_WASM_FILE_NAME)
         .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = RunningAverageClient::with_interceptor(channel, interceptor);
 
     submit_sample(&mut client, 100).await;
diff --git a/examples/translator/module/rust/tests/integration_test.rs b/examples/translator/module/rust/tests/integration_test.rs
index a32cf3199e..29e0dd2e05 100644
--- a/examples/translator/module/rust/tests/integration_test.rs
+++ b/examples/translator/module/rust/tests/integration_test.rs
@@ -27,7 +27,7 @@ async fn test_translate() {
     let runtime = oak_tests::run_single_module(MODULE_WASM_FILE_NAME, "grpc_oak_main")
         .expect("Unable to configure runtime with test wasm!");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = TranslatorClient::with_interceptor(channel, interceptor);
 
     let req = TranslateRequest {
diff --git a/examples/trusted_database/module/rust/tests/integration_test.rs b/examples/trusted_database/module/rust/tests/integration_test.rs
index f5b178ec48..adc97faff8 100644
--- a/examples/trusted_database/module/rust/tests/integration_test.rs
+++ b/examples/trusted_database/module/rust/tests/integration_test.rs
@@ -130,7 +130,7 @@ async fn test_trusted_database() {
     let runtime =
         oak_runtime::configure_and_run(config).expect("Couldn't configure runtime with test wasm");
 
-    let (channel, interceptor) = oak_tests::channel_and_interceptor().await;
+    let (channel, interceptor) = oak_tests::public_channel_and_interceptor().await;
     let mut client = TrustedDatabaseClient::with_interceptor(channel, interceptor);
 
     // Test nearest point of interest.
diff --git a/oak_derive/tests/visit.rs b/oak_derive/tests/visit.rs
index 716c7489bf..1380334ed7 100644
--- a/oak_derive/tests/visit.rs
+++ b/oak_derive/tests/visit.rs
@@ -20,9 +20,8 @@ use oak_io::{handle::HandleVisit, Handle};
 struct Visited(Handle);
 
 impl HandleVisit for Visited {
-    fn visit<F: FnMut(&mut Handle)>(&mut self, mut f: F) -> F {
-        f(&mut self.0);
-        f
+    fn fold<B>(&mut self, init: B, f: fn(B, &mut Handle) -> B) -> B {
+        f(init, &mut self.0)
     }
 }
 
@@ -127,14 +126,9 @@ mod enums {
     }
 }
 
-// Asserts that `t` is visited exactly `count` times when calling `HandleVisit::visit`.
-fn assert_visit<T: HandleVisit>(mut t: T, count: usize) {
-    let mut counter = 0;
+// Asserts that `t` is visited exactly `expected` times when calling [`HandleVisit::fold`].
+fn assert_visit<T: HandleVisit>(mut t: T, expected: usize) {
+    let count = t.fold(0, |count, _| count + 1);
 
-    // TODO(#1599): Remove the `let _` when the underlying issue is fixed.
-    let _ = t.visit(|_| {
-        counter += 1;
-    });
-
-    assert_eq!(counter, count);
+    assert_eq!(count, expected);
 }
diff --git a/sdk/rust/oak_tests/src/lib.rs b/sdk/rust/oak_tests/src/lib.rs
index d45899fca0..f3d452c979 100644
--- a/sdk/rust/oak_tests/src/lib.rs
+++ b/sdk/rust/oak_tests/src/lib.rs
@@ -197,9 +197,9 @@ fn create_http_config() -> Option<oak_runtime::HttpConfiguration> {
     Some(oak_runtime::HttpConfiguration { tls_config })
 }
 
-/// Build a channel and interceptor suitable for building a client that connects
-/// to a Runtime under test.
-pub async fn channel_and_interceptor() -> (Channel, impl Into<tonic::Interceptor>) {
+/// Build a labeled channel and interceptor suitable for building a client that
+/// connects to a Runtime under test.
+pub async fn channel_and_interceptor(label: &Label) -> (Channel, impl Into<tonic::Interceptor>) {
     // Build a channel that connects to the Runtime under test.
     let uri = RUNTIME_URI.parse().expect("Error parsing URI");
     let tls_config = ClientTlsConfig::new()
@@ -235,17 +235,23 @@ pub async fn channel_and_interceptor() -> (Channel, impl Into<tonic::Interceptor
 
     // Build an interceptor that will attach a public-untrusted Oak label to
     // every gRPC request.
-    let mut label = Vec::new();
-    Label::public_untrusted()
-        .encode(&mut label)
+    let mut encoded_label = Vec::new();
+    label
+        .encode(&mut encoded_label)
         .expect("Error encoding label");
     let interceptor = move |mut request: Request<()>| {
         request.metadata_mut().insert_bin(
             oak_abi::OAK_LABEL_GRPC_METADATA_KEY,
-            MetadataValue::from_bytes(label.as_ref()),
+            MetadataValue::from_bytes(encoded_label.as_ref()),
         );
         Ok(request)
     };
 
     (channel, interceptor)
 }
+
+/// Build a public-untrusted channel and interceptor suitable for building a
+/// client that connects to a Runtime under test.
+pub async fn public_channel_and_interceptor() -> (Channel, impl Into<tonic::Interceptor>) {
+    channel_and_interceptor(&Label::public_untrusted()).await
+}