1 #![allow(unused_imports)] 2 3 use self::util::*; 4 use crate::util::mock_io_channel; 5 use futures::{Stream, StreamExt}; 6 use std::convert::TryFrom; 7 use std::{ 8 pin::Pin, 9 sync::{ 10 atomic::{AtomicUsize, Ordering::SeqCst}, 11 Arc, 12 }, 13 }; 14 use tokio::net::TcpListener; 15 use tonic::{ 16 transport::{Channel, Endpoint, Server, Uri}, 17 Request, Response, Status, Streaming, 18 }; 19 use tower::{layer::layer_fn, service_fn, Service, ServiceBuilder}; 20 use tower_http::{map_request_body::MapRequestBodyLayer, map_response_body::MapResponseBodyLayer}; 21 22 mod bidirectional_stream; 23 mod client_stream; 24 mod compressing_request; 25 mod compressing_response; 26 mod server_stream; 27 mod util; 28 29 tonic::include_proto!("test"); 30 31 #[derive(Debug, Default)] 32 struct Svc { 33 disable_compressing_on_response: bool, 34 } 35 36 const UNCOMPRESSED_MIN_BODY_SIZE: usize = 1024; 37 38 impl Svc { 39 fn prepare_response<B>(&self, mut res: Response<B>) -> Response<B> { 40 if self.disable_compressing_on_response { 41 res.disable_compression(); 42 } 43 44 res 45 } 46 } 47 48 #[tonic::async_trait] 49 impl test_server::Test for Svc { 50 async fn compress_output_unary(&self, _req: Request<()>) -> Result<Response<SomeData>, Status> { 51 let data = [0_u8; UNCOMPRESSED_MIN_BODY_SIZE]; 52 53 Ok(self.prepare_response(Response::new(SomeData { 54 data: data.to_vec(), 55 }))) 56 } 57 58 async fn compress_input_unary(&self, req: Request<SomeData>) -> Result<Response<()>, Status> { 59 assert_eq!(req.into_inner().data.len(), UNCOMPRESSED_MIN_BODY_SIZE); 60 Ok(Response::new(())) 61 } 62 63 type CompressOutputServerStreamStream = 64 Pin<Box<dyn Stream<Item = Result<SomeData, Status>> + Send + 'static>>; 65 66 async fn compress_output_server_stream( 67 &self, 68 _req: Request<()>, 69 ) -> Result<Response<Self::CompressOutputServerStreamStream>, Status> { 70 let data = [0_u8; UNCOMPRESSED_MIN_BODY_SIZE].to_vec(); 71 let stream = futures::stream::repeat(SomeData { data }) 72 .take(2) 73 .map(Ok::<_, Status>); 74 Ok(self.prepare_response(Response::new(Box::pin(stream)))) 75 } 76 77 async fn compress_input_client_stream( 78 &self, 79 req: Request<Streaming<SomeData>>, 80 ) -> Result<Response<()>, Status> { 81 let mut stream = req.into_inner(); 82 while let Some(item) = stream.next().await { 83 item.unwrap(); 84 } 85 Ok(self.prepare_response(Response::new(()))) 86 } 87 88 async fn compress_output_client_stream( 89 &self, 90 req: Request<Streaming<SomeData>>, 91 ) -> Result<Response<SomeData>, Status> { 92 let mut stream = req.into_inner(); 93 while let Some(item) = stream.next().await { 94 item.unwrap(); 95 } 96 97 let data = [0_u8; UNCOMPRESSED_MIN_BODY_SIZE]; 98 99 Ok(self.prepare_response(Response::new(SomeData { 100 data: data.to_vec(), 101 }))) 102 } 103 104 type CompressInputOutputBidirectionalStreamStream = 105 Pin<Box<dyn Stream<Item = Result<SomeData, Status>> + Send + 'static>>; 106 107 async fn compress_input_output_bidirectional_stream( 108 &self, 109 req: Request<Streaming<SomeData>>, 110 ) -> Result<Response<Self::CompressInputOutputBidirectionalStreamStream>, Status> { 111 let mut stream = req.into_inner(); 112 while let Some(item) = stream.next().await { 113 item.unwrap(); 114 } 115 116 let data = [0_u8; UNCOMPRESSED_MIN_BODY_SIZE].to_vec(); 117 let stream = futures::stream::repeat(SomeData { data }) 118 .take(2) 119 .map(Ok::<_, Status>); 120 Ok(self.prepare_response(Response::new(Box::pin(stream)))) 121 } 122 } 123