Assert gRPC calls

.github/workflows/rust.yml

@@ -139,7 +139,7 @@ jobs:
       with:
         repository: proxy-wasm/proxy-wasm-rust-sdk
         path: proxy-wasm-rust-sdk
-        ref: v0.2.1
+        ref: v0.2.2
 
     - name: Update Rust
       run: |
@@ -179,6 +179,7 @@ jobs:
         cd proxy-wasm-rust-sdk/examples/hello_world && cargo build --target wasm32-unknown-unknown --release && cd ../../..
         cd proxy-wasm-rust-sdk/examples/http_auth_random && cargo build --target wasm32-unknown-unknown --release && cd ../../..
         cd proxy-wasm-rust-sdk/examples/http_headers && cargo build --target wasm32-unknown-unknown --release && cd ../../..
+        cd proxy-wasm-rust-sdk/examples/grpc_auth_random && cargo build --target wasm32-unknown-unknown --release && cd ../../..
 
     - name: Test (hello_world)
       run: target/release/examples/hello_world proxy-wasm-rust-sdk/examples/hello_world/target/wasm32-unknown-unknown/release/proxy_wasm_example_hello_world.wasm
@@ -189,6 +190,9 @@ jobs:
     - name: Test (http_headers)
       run: target/release/examples/http_headers proxy-wasm-rust-sdk/examples/http_headers/target/wasm32-unknown-unknown/release/proxy_wasm_example_http_headers.wasm -a
 
+    - name: Test (grpc_auth_random)
+      run: target/release/examples/grpc_auth_random proxy-wasm-rust-sdk/examples/grpc_auth_random/target/wasm32-unknown-unknown/release/proxy_wasm_example_grpc_auth_random.wasm
+
   outdated:
     runs-on: ubuntu-latest
 

examples/grpc_auth_random.rs

@@ -0,0 +1,69 @@
+// Copyright 2020 Google LLC
+//
+// 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.
+
+use anyhow::Result;
+use proxy_wasm_test_framework::tester;
+use proxy_wasm_test_framework::types::*;
+use structopt::StructOpt;
+
+fn main() -> Result<()> {
+    let args = tester::MockSettings::from_args();
+    let mut module = tester::mock(args)?;
+
+    module.call_start().execute_and_expect(ReturnType::None)?;
+
+    let root_context = 1;
+    module
+        .call_proxy_on_context_create(root_context, 0)
+        .execute_and_expect(ReturnType::None)?;
+
+    let http_context = 2;
+    module
+        .call_proxy_on_context_create(http_context, root_context)
+        .execute_and_expect(ReturnType::None)?;
+
+    let token_id = 42;
+    module
+        .call_proxy_on_request_headers(http_context, 0, false)
+        .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("content-type"))
+        .returning(Some("application/grpc"))
+        .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path"))
+        .returning(Some("/someService/someService.someMethod"))
+        .expect_grpc_call(
+            Some("grpcbin"),
+            Some("grpcbin.GRPCBin"),
+            Some("RandomError"),
+            Some(&[0, 0, 0, 0]),
+            Some(&[]),
+            Some(1000), // 1 sec as millis
+        )
+        .returning(Some(token_id))
+        .execute_and_expect(ReturnType::Action(Action::Pause))?;
+
+    module
+        .call_proxy_on_grpc_receive(http_context, token_id as i32, 0 as i32)
+        .expect_log(Some(LogLevel::Info), Some("Access granted."))
+        .execute_and_expect(ReturnType::None)?;
+
+    module
+        .call_proxy_on_response_headers(http_context, 0, false)
+        .expect_replace_header_map_value(
+            Some(MapType::HttpResponseHeaders),
+            Some("Powered-By"),
+            Some("proxy-wasm"),
+        )
+        .execute_and_expect(ReturnType::Action(Action::Continue))?;
+
+    return Ok(());
+}

examples/http_auth_random.rs

@@ -55,7 +55,7 @@ fn main() -> Result<()> {
     http_auth_random
         .call_proxy_on_http_call_response(http_context, 0, 0, buffer_data.len() as i32, 0)
         .expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody))
-        .returning(Some(buffer_data))
+        .returning(Some(buffer_data.as_bytes()))
         .expect_send_local_response(
             Some(403),
             Some("Access forbidden.\n"),

src/expect_interface.rs

@@ -51,7 +51,7 @@ impl<'a> ExpectGetBufferBytes<'a> {
         }
     }
 
-    pub fn returning(&mut self, buffer_data: Option<&str>) -> &mut Tester {
+    pub fn returning(&mut self, buffer_data: Option<&[u8]>) -> &mut Tester {
         self.tester
             .get_expect_handle()
             .staged
@@ -150,3 +150,48 @@ impl<'a> ExpectHttpCall<'a> {
         self.tester
     }
 }
+
+pub struct ExpectGrpcCall<'a> {
+    tester: &'a mut Tester,
+    service: Option<&'a str>,
+    service_name: Option<&'a str>,
+    method_name: Option<&'a str>,
+    initial_metadata: Option<&'a [u8]>,
+    request: Option<&'a [u8]>,
+    timeout: Option<u64>,
+}
+
+impl<'a> ExpectGrpcCall<'a> {
+    pub fn expecting(
+        tester: &'a mut Tester,
+        service: Option<&'a str>,
+        service_name: Option<&'a str>,
+        method_name: Option<&'a str>,
+        initial_metadata: Option<&'a [u8]>,
+        request: Option<&'a [u8]>,
+        timeout: Option<u64>,
+    ) -> Self {
+        Self {
+            tester,
+            service,
+            service_name,
+            method_name,
+            initial_metadata,
+            request,
+            timeout,
+        }
+    }
+
+    pub fn returning(&mut self, token_id: Option<u32>) -> &mut Tester {
+        self.tester.get_expect_handle().staged.set_expect_grpc_call(
+            self.service,
+            self.service_name,
+            self.method_name,
+            self.initial_metadata,
+            self.request,
+            self.timeout,
+            token_id,
+        );
+        self.tester
+    }
+}

src/expectations.rs

@@ -41,19 +41,20 @@ impl ExpectHandle {
         self.staged = Expect::new(allow_unexpected);
     }
 
-    pub fn assert_stage(&self) {
+    pub fn assert_stage(&self) -> Option<String> {
         if self.staged.expect_count > 0 {
-            panic!(
+            return Some(format!(
                 "Error: failed to consume all expectations - total remaining: {}",
                 self.staged.expect_count
-            );
+            ));
         } else if self.staged.expect_count < 0 {
-            panic!(
+            return Some(format!(
                 "Error: expectations failed to account for all host calls by {} \n\
             if this is intended, please use --allow-unexpected (-a) mode",
                 -1 * self.staged.expect_count
-            );
+            ));
         }
+        None
     }
 
     pub fn print_staged(&self) {
@@ -86,6 +87,15 @@ pub struct Expect {
         Option<Duration>,
         Option<u32>,
     )>,
+    grpc_call: Vec<(
+        Option<String>,
+        Option<String>,
+        Option<String>,
+        Option<Bytes>,
+        Option<Bytes>,
+        Option<Duration>,
+        Option<u32>,
+    )>,
 }
 
 impl Expect {
@@ -106,6 +116,7 @@ impl Expect {
             add_header_map_value: vec![],
             send_local_response: vec![],
             http_call: vec![],
+            grpc_call: vec![],
         }
     }
 
@@ -190,13 +201,11 @@ impl Expect {
     pub fn set_expect_get_buffer_bytes(
         &mut self,
         buffer_type: Option<i32>,
-        buffer_data: Option<&str>,
+        buffer_data: Option<&[u8]>,
     ) {
         self.expect_count += 1;
-        self.get_buffer_bytes.push((
-            buffer_type,
-            buffer_data.map(|data| data.as_bytes().to_vec()),
-        ));
+        self.get_buffer_bytes
+            .push((buffer_type, buffer_data.map(|data| data.to_vec())));
     }
 
     pub fn get_expect_get_buffer_bytes(&mut self, buffer_type: i32) -> Option<Bytes> {
@@ -571,4 +580,73 @@ impl Expect {
             }
         }
     }
+
+    pub fn set_expect_grpc_call(
+        &mut self,
+        service: Option<&str>,
+        service_name: Option<&str>,
+        method_name: Option<&str>,
+        initial_metadata: Option<&[u8]>,
+        request: Option<&[u8]>,
+        timeout: Option<u64>,
+        token_id: Option<u32>,
+    ) {
+        self.expect_count += 1;
+        self.grpc_call.push((
+            service.map(ToString::to_string),
+            service_name.map(ToString::to_string),
+            method_name.map(ToString::to_string),
+            initial_metadata.map(|s| s.to_vec()),
+            request.map(|s| s.to_vec()),
+            timeout.map(Duration::from_millis),
+            token_id,
+        ));
+    }
+
+    pub fn get_expect_grpc_call(
+        &mut self,
+        service: String,
+        service_name: String,
+        method: String,
+        initial_metadata: &[u8],
+        request: &[u8],
+        timeout: i32,
+    ) -> Option<u32> {
+        match self.grpc_call.len() {
+            0 => {
+                if !self.allow_unexpected {
+                    self.expect_count -= 1;
+                }
+                set_status(ExpectStatus::Unexpected);
+                None
+            }
+            _ => {
+                self.expect_count -= 1;
+                let (
+                    expected_service,
+                    expected_service_name,
+                    expected_method,
+                    expected_initial_metadata,
+                    expected_request,
+                    expected_duration,
+                    result,
+                ) = self.grpc_call.remove(0);
+
+                let expected = expected_service.map(|e| e == service).unwrap_or(true)
+                    && expected_service_name
+                        .map(|e| e == service_name)
+                        .unwrap_or(true)
+                    && expected_method.map(|e| e == method).unwrap_or(true)
+                    && expected_initial_metadata
+                        .map(|e| e == initial_metadata)
+                        .unwrap_or(true)
+                    && expected_request.map(|e| e == request).unwrap_or(true)
+                    && expected_duration
+                        .map(|e| e.as_millis() as i32 == timeout)
+                        .unwrap_or(true);
+                set_expect_status(expected);
+                return result;
+            }
+        }
+    }
 }

src/host_settings.rs

@@ -294,5 +294,9 @@ pub fn default_buffer_bytes() -> HashMap<i32, Bytes> {
         BufferType::HttpCallResponseBody as i32,
         "default_call_response_body".as_bytes().to_vec(),
     );
+    default_bytes.insert(
+        BufferType::GrpcReceiveBuffer as i32,
+        "default_grpc_receive_buffer".as_bytes().to_vec(),
+    );
     default_bytes
 }