1 use futures::stream; 2 use futures_util::FutureExt; 3 use tokio::sync::oneshot; 4 use tonic::transport::Server; 5 use tonic::Request; 6 use tonic_reflection::server::Builder; 7 8 use pb::server_reflection_client::ServerReflectionClient; 9 use pb::server_reflection_request::MessageRequest; 10 use pb::server_reflection_response::MessageResponse; 11 use pb::ServerReflectionRequest; 12 use pb::ServiceResponse; 13 use std::net::SocketAddr; 14 use tokio_stream::wrappers::TcpListenerStream; 15 use tokio_stream::StreamExt; 16 17 mod pb { 18 #![allow(unreachable_pub)] 19 use prost::Message; 20 21 tonic::include_proto!("grpc.reflection.v1alpha"); 22 23 pub(crate) const REFLECTION_SERVICE_DESCRIPTOR: &'static [u8] = 24 tonic::include_file_descriptor_set!("reflection_v1alpha1"); 25 26 pub(crate) fn get_encoded_reflection_service_fd() -> Vec<u8> { 27 let mut expected = Vec::new(); 28 &prost_types::FileDescriptorSet::decode(REFLECTION_SERVICE_DESCRIPTOR) 29 .expect("decode reflection service file descriptor set") 30 .file[0] 31 .encode(&mut expected) 32 .expect("encode reflection service file descriptor"); 33 expected 34 } 35 } 36 37 #[tokio::test] 38 async fn test_list_services() { 39 let response = make_test_reflection_request(ServerReflectionRequest { 40 host: "".to_string(), 41 message_request: Some(MessageRequest::ListServices(String::new())), 42 }) 43 .await; 44 45 if let MessageResponse::ListServicesResponse(services) = response { 46 assert_eq!( 47 services.service, 48 vec![ServiceResponse { 49 name: String::from("grpc.reflection.v1alpha.ServerReflection") 50 }] 51 ); 52 } else { 53 panic!("Expected a ListServicesResponse variant"); 54 } 55 } 56 57 #[tokio::test] 58 async fn test_file_by_filename() { 59 let response = make_test_reflection_request(ServerReflectionRequest { 60 host: "".to_string(), 61 message_request: Some(MessageRequest::FileByFilename(String::from( 62 "reflection.proto", 63 ))), 64 }) 65 .await; 66 67 if let MessageResponse::FileDescriptorResponse(descriptor) = response { 68 let file_descriptor_proto = descriptor 69 .file_descriptor_proto 70 .first() 71 .expect("descriptor"); 72 assert_eq!( 73 file_descriptor_proto.as_ref(), 74 pb::get_encoded_reflection_service_fd() 75 ); 76 } else { 77 panic!("Expected a FileDescriptorResponse variant"); 78 } 79 } 80 81 #[tokio::test] 82 async fn test_file_containing_symbol() { 83 let response = make_test_reflection_request(ServerReflectionRequest { 84 host: "".to_string(), 85 message_request: Some(MessageRequest::FileContainingSymbol(String::from( 86 "grpc.reflection.v1alpha.ServerReflection", 87 ))), 88 }) 89 .await; 90 91 if let MessageResponse::FileDescriptorResponse(descriptor) = response { 92 let file_descriptor_proto = descriptor 93 .file_descriptor_proto 94 .first() 95 .expect("descriptor"); 96 assert_eq!( 97 file_descriptor_proto.as_ref(), 98 pb::get_encoded_reflection_service_fd() 99 ); 100 } else { 101 panic!("Expected a FileDescriptorResponse variant"); 102 } 103 } 104 105 async fn make_test_reflection_request(request: ServerReflectionRequest) -> MessageResponse { 106 // Run a test server 107 let (shutdown_tx, shutdown_rx) = oneshot::channel(); 108 109 let addr: SocketAddr = "127.0.0.1:0".parse().expect("SocketAddr parse"); 110 let listener = tokio::net::TcpListener::bind(addr).await.expect("bind"); 111 let local_addr = listener.local_addr().expect("local address"); 112 let local_addr = format!("http://{}", local_addr.to_string()); 113 let jh = tokio::spawn(async move { 114 let service = Builder::configure() 115 .register_encoded_file_descriptor_set(pb::REFLECTION_SERVICE_DESCRIPTOR) 116 .build() 117 .unwrap(); 118 119 Server::builder() 120 .add_service(service) 121 .serve_with_incoming_shutdown(TcpListenerStream::new(listener), shutdown_rx.map(drop)) 122 .await 123 .unwrap(); 124 }); 125 126 // Give the test server a few ms to become available 127 tokio::time::sleep(std::time::Duration::from_millis(100)).await; 128 129 // Construct client and send request, extract response 130 let mut client = ServerReflectionClient::connect(local_addr) 131 .await 132 .expect("connect"); 133 134 let request = Request::new(stream::iter(vec![request])); 135 let mut inbound = client 136 .server_reflection_info(request) 137 .await 138 .expect("request") 139 .into_inner(); 140 141 let response = inbound 142 .next() 143 .await 144 .expect("steamed response") 145 .expect("successful response") 146 .message_response 147 .expect("some MessageResponse"); 148 149 // We only expect one response per request 150 assert!(inbound.next().await.is_none()); 151 152 // Shut down test server 153 shutdown_tx.send(()).expect("send shutdown"); 154 jh.await.expect("server shutdown"); 155 156 response 157 } 158