Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix: add explicit test for tonic
Browse files Browse the repository at this point in the history
0xAlcibiades committed Dec 22, 2024
1 parent 0997bde commit 08fa2e9
Showing 3 changed files with 146 additions and 3 deletions.
4 changes: 3 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -44,10 +44,12 @@ reqwest = { version = "0.12.9", features = ["rustls-tls", "http2"] }
ring = "0.17.8"
tokio = { version = "1.42", features = ["rt-multi-thread", "net", "test-util", "time", "signal"] }
tokio-util = { version = "0.7.13", features = ["compat"] }
tonic = "0.12.3"
tonic = { version = "0.12.3", features = ["tls"] }
tonic-health = "0.12.3"
tonic-reflection = "0.12.3"
tracing-subscriber = "0.3.19"
lazy_static = "1.5.0"
tonic-types = "0.12.3"

[[bench]]
name = "hello_world_tower_hyper_tls_tcp"
16 changes: 14 additions & 2 deletions examples/tonic.rs
Original file line number Diff line number Diff line change
@@ -9,6 +9,7 @@ use tokio::net::TcpListener;
use tokio_stream::wrappers::TcpListenerStream;
use tonic::server::NamedService;
use tonic::transport::Server;
use tower::ServiceExt;
use tracing::{info, Level};

use postel::{load_certs, load_private_key, serve_http_with_shutdown};
@@ -23,11 +24,19 @@ impl NamedService for GreeterService {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
// Initialize logging
tracing_subscriber::fmt().with_max_level(Level::INFO).init();
tracing_subscriber::fmt()
.with_max_level(Level::DEBUG)
.init();

// Configure server address
let addr = SocketAddr::from(([127, 0, 0, 1], 8443));

let reflection_server = tonic_reflection::server::Builder::configure()
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
.register_encoded_file_descriptor_set(tonic_types::pb::FILE_DESCRIPTOR_SET)
.register_encoded_file_descriptor_set(tonic_reflection::pb::v1alpha::FILE_DESCRIPTOR_SET)
.build_v1alpha()?;

// Set up gRPC health service
let (mut health_reporter, health_service) = tonic_health::server::health_reporter();
health_reporter.set_serving::<GreeterService>().await;
@@ -65,8 +74,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {

let svc = Server::builder()
.add_service(health_service)
.add_service(reflection_server)
.into_service()
.into_axum_router()
.into_service()
.into_axum_router();
.boxed_clone();

let hyper_svc = TowerToHyperService::new(svc);

129 changes: 129 additions & 0 deletions tests/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use hyper_util::rt::TokioExecutor;
use hyper_util::server::conn::auto::Builder as HttpConnectionBuilder;
use hyper_util::service::TowerToHyperService;
use rustls::ServerConfig;
use tokio::net::TcpListener;
use tokio_stream::wrappers::TcpListenerStream;
use tonic::server::NamedService;
use tonic::transport::Server;
use tonic::transport::{Certificate, Channel, ClientTlsConfig};
use tonic_health::pb::health_client::HealthClient;
use tower::ServiceExt;

use postel::{load_certs, load_private_key, serve_http_with_shutdown};

pub struct TestService;

impl NamedService for TestService {
const NAME: &'static str = "test.service";
}

async fn setup_test_server(
) -> Result<(SocketAddr, tokio::sync::oneshot::Sender<()>), Box<dyn std::error::Error + Send + Sync>>
{
rustls::crypto::aws_lc_rs::default_provider()
.install_default()
.expect("Failed to install rustls crypto provider");
// Setup server with random port
let addr = SocketAddr::from(([127, 0, 0, 1], 0));
let listener = TcpListener::bind(addr).await?;
let server_addr = listener.local_addr()?;
let incoming = TcpListenerStream::new(listener);

// Setup TLS
let certs = load_certs("examples/sample.pem")?;
let key = load_private_key("examples/sample.rsa")?;

let mut config = ServerConfig::builder()
.with_no_client_auth()
.with_single_cert(certs, key)?;
config.alpn_protocols = vec![b"h2".to_vec()];
let tls_config = Arc::new(config);

// Setup gRPC services
let reflection_server = tonic_reflection::server::Builder::configure()
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
.build_v1alpha()?;

let (mut health_reporter, health_service) = tonic_health::server::health_reporter();
health_reporter.set_serving::<TestService>().await;

let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();

// Build and start server
let builder = HttpConnectionBuilder::new(TokioExecutor::new());

let svc = Server::builder()
.add_service(health_service)
.add_service(reflection_server)
.into_service()
.into_axum_router()
.into_service()
.boxed_clone();

let hyper_svc = TowerToHyperService::new(svc);

tokio::spawn(async move {
serve_http_with_shutdown(
hyper_svc,
incoming,
builder,
Some(tls_config),
Some(async {
shutdown_rx.await.ok();
}),
)
.await
.expect("Server failed unexpectedly");
});

// Give server a moment to start
tokio::time::sleep(Duration::from_millis(100)).await;

Ok((server_addr, shutdown_tx))
}

async fn setup_test_client(
addr: SocketAddr,
) -> Result<HealthClient<Channel>, Box<dyn std::error::Error + Send + Sync>> {
let pem = std::fs::read_to_string("examples/sample.pem")?;
let ca = Certificate::from_pem(pem);

let tls_config = ClientTlsConfig::new()
.ca_certificate(ca)
.domain_name("localhost");

let channel = Channel::builder(format!("https://{}", addr).parse()?)
.tls_config(tls_config)?
.connect()
.await?;

Ok(HealthClient::new(channel))
}

#[tokio::test]
async fn test_grpc_health_check() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let (addr, shutdown_tx) = setup_test_server().await?;
let mut client = setup_test_client(addr).await?;

// Test service check
let response = client
.check(tonic_health::pb::HealthCheckRequest {
service: "test.service".to_string(),
})
.await?;

assert_eq!(
response.into_inner().status(),
tonic_health::pb::health_check_response::ServingStatus::Serving
);

// Clean shutdown
shutdown_tx.send(()).unwrap();

Ok(())
}

0 comments on commit 08fa2e9

Please sign in to comment.