Fix gRPC stream status and timeout validation

Author: benthecarman

Cargo.lock

@@ -12,3 +12,11 @@ ldk-server-protos = { path = "../ldk-server-protos" }
 reqwest = { version = "0.11.13", default-features = false, features = ["rustls-tls"] }
 prost = { version = "0.11.6", default-features = false, features = ["std", "prost-derive"] }
 bitcoin_hashes = "0.14"
+http-body = "0.4"
+hyper = { version = "0.14", default-features = false, features = ["client", "http2", "runtime", "tcp"] }
+hyper-rustls = { version = "0.24", default-features = false, features = ["http2", "tls12", "tokio-runtime"] }
+rustls = "0.21"
+rustls-pemfile = "1"
+
+[dev-dependencies]
+tokio = { version = "1", default-features = false, features = ["macros", "rt"] }

e2e-tests/Cargo.lock

@@ -7,10 +7,14 @@
 // You may not use this file except in accordance with one or both of these
 // licenses.
 
+use std::io::Cursor;
 use std::time::{SystemTime, UNIX_EPOCH};
 
 use bitcoin_hashes::hmac::{Hmac, HmacEngine};
 use bitcoin_hashes::{sha256, Hash, HashEngine};
+use hyper::body::HttpBody as _;
+use hyper::{Body as HyperBody, Client as HyperClient, Request as HyperRequest, Version};
+use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder};
 use ldk_server_protos::api::SubscribeEventsRequest;
 use ldk_server_protos::api::{
 	Bolt11ClaimForHashRequest, Bolt11ClaimForHashResponse, Bolt11FailForHashRequest,
@@ -50,7 +54,9 @@ use ldk_server_protos::endpoints::{
 };
 use ldk_server_protos::events::EventEnvelope;
 use prost::Message;
-use reqwest::{Certificate, Client};
+use reqwest::{header::HeaderMap, Certificate, Client};
+use rustls::{ClientConfig, RootCertStore};
+use rustls_pemfile::certs;
 
 use crate::error::LdkServerError;
 use crate::error::LdkServerErrorCode::{
@@ -60,6 +66,8 @@ use crate::error::LdkServerErrorCode::{
 /// gRPC path prefix for the LightningNode service.
 const GRPC_SERVICE_PREFIX: &str = "/api.LightningNode/";
 
+type StreamingClient = HyperClient<HttpsConnector<hyper::client::HttpConnector>, HyperBody>;
+
 /// Client to access a hosted instance of LDK Server via gRPC.
 ///
 /// The client requires the server's TLS certificate to be provided for verification.
@@ -69,6 +77,7 @@ const GRPC_SERVICE_PREFIX: &str = "/api.LightningNode/";
 pub struct LdkServerClient {
 	base_url: String,
 	client: Client,
+	streaming_client: StreamingClient,
 	api_key: String,
 }
 
@@ -82,13 +91,14 @@ impl LdkServerClient {
 	pub fn new(base_url: String, api_key: String, server_cert_pem: &[u8]) -> Result<Self, String> {
 		let cert = Certificate::from_pem(server_cert_pem)
 			.map_err(|e| format!("Failed to parse server certificate: {e}"))?;
+		let streaming_client = build_streaming_client(server_cert_pem)?;
 
 		let client = Client::builder()
 			.add_root_certificate(cert)
 			.build()
 			.map_err(|e| format!("Failed to build HTTP client: {e}"))?;
 
-		Ok(Self { base_url, client, api_key })
+		Ok(Self { base_url, client, streaming_client, api_key })
 	}
 
 	/// Computes the HMAC-SHA256 authentication header value.
@@ -411,36 +421,32 @@ impl LdkServerClient {
 		let auth_header = self.compute_auth_header();
 
 		let response = self
-			.client
-			.post(&url)
-			.header("content-type", "application/grpc+proto")
-			.header("te", "trailers")
-			.header("x-auth", auth_header)
-			.body(grpc_body)
-			.send()
+			.streaming_client
+			.request(
+				HyperRequest::post(&url)
+					.version(Version::HTTP_2)
+					.header("content-type", "application/grpc+proto")
+					.header("te", "trailers")
+					.header("x-auth", auth_header)
+					.body(HyperBody::from(grpc_body))
+					.map_err(|e| {
+						LdkServerError::new(
+							InternalError,
+							format!("Failed to build gRPC request: {e}"),
+						)
+					})?,
+			)
 			.await
 			.map_err(|e| {
 				LdkServerError::new(InternalError, format!("gRPC request failed: {}", e))
 			})?;
 
-		// Check for Trailers-Only error
-		if let Some(status_val) = response.headers().get("grpc-status") {
-			if let Ok(status_str) = status_val.to_str() {
-				if let Ok(code) = status_str.parse::<u32>() {
-					if code != 0 {
-						let message = response
-							.headers()
-							.get("grpc-message")
-							.and_then(|v| v.to_str().ok())
-							.map(percent_decode)
-							.unwrap_or_default();
-						return Err(grpc_code_to_error(code, message));
-					}
-				}
-			}
+		let (parts, body) = response.into_parts();
+		if let Some(error) = grpc_error_from_headers(&parts.headers) {
+			return Err(error);
 		}
 
-		Ok(EventStream { response, buf: Vec::new() })
+		Ok(EventStream { body, buf: Vec::new(), trailers_checked: false })
 	}
 
 	/// Send a unary gRPC request and decode the response.
@@ -473,20 +479,8 @@ impl LdkServerClient {
 		// Check for Trailers-Only error responses (grpc-status in response headers).
 		// In gRPC, when there is no response body (error case), the server sends
 		// grpc-status as part of the initial HEADERS frame, readable as a regular header.
-		if let Some(status_val) = response.headers().get("grpc-status") {
-			if let Ok(status_str) = status_val.to_str() {
-				if let Ok(code) = status_str.parse::<u32>() {
-					if code != 0 {
-						let message = response
-							.headers()
-							.get("grpc-message")
-							.and_then(|v| v.to_str().ok())
-							.map(percent_decode)
-							.unwrap_or_default();
-						return Err(grpc_code_to_error(code, message));
-					}
-				}
-			}
+		if let Some(error) = grpc_error_from_headers(response.headers()) {
+			return Err(error);
 		}
 
 		// Read the response body
@@ -515,14 +509,42 @@ impl LdkServerClient {
 
 /// Map a gRPC status code to an LdkServerError.
 fn grpc_code_to_error(code: u32, message: String) -> LdkServerError {
-	let error_code = match code {
-		3 => InvalidRequestError,  // INVALID_ARGUMENT
-		16 => AuthError,           // UNAUTHENTICATED
-		9 => LightningError,       // FAILED_PRECONDITION
-		13 => InternalServerError, // INTERNAL
-		_ => InternalError,
-	};
-	LdkServerError::new(error_code, message)
+	match code {
+		3 => LdkServerError::new(InvalidRequestError, message), // INVALID_ARGUMENT
+		9 => LdkServerError::new(LightningError, message),      // FAILED_PRECONDITION
+		13 => LdkServerError::new(InternalServerError, message), // INTERNAL
+		14 => LdkServerError::new(
+			InternalError,
+			if message.is_empty() {
+				"gRPC stream became unavailable".to_string()
+			} else {
+				format!("gRPC stream became unavailable: {message}")
+			},
+		),
+		16 => LdkServerError::new(AuthError, message), // UNAUTHENTICATED
+		_ => LdkServerError::new(
+			InternalError,
+			if message.is_empty() {
+				format!("gRPC status {code}")
+			} else {
+				format!("gRPC status {code}: {message}")
+			},
+		),
+	}
+}
+
+fn grpc_error_from_headers(headers: &HeaderMap) -> Option<LdkServerError> {
+	let code = headers.get("grpc-status")?.to_str().ok()?.parse::<u32>().ok()?;
+	if code == 0 {
+		return None;
+	}
+
+	let message = headers
+		.get("grpc-message")
+		.and_then(|v| v.to_str().ok())
+		.map(percent_decode)
+		.unwrap_or_default();
+	Some(grpc_code_to_error(code, message))
 }
 
 /// Minimal percent-decoding for grpc-message values.
@@ -556,8 +578,9 @@ fn hex_val(b: u8) -> Option<u8> {
 ///
 /// Call [`next_event`](EventStream::next_event) to receive the next event from the server.
 pub struct EventStream {
-	response: reqwest::Response,
+	body: hyper::Body,
 	buf: Vec<u8>,
+	trailers_checked: bool,
 }
 
 impl EventStream {
@@ -582,23 +605,126 @@ impl EventStream {
 			}
 
 			// Need more data — read the next chunk from the response body
-			match self.response.chunk().await {
-				Ok(Some(chunk)) => self.buf.extend_from_slice(&chunk),
-				Ok(None) => return None, // stream ended
-				Err(e) => {
+			match self.body.data().await {
+				Some(Ok(chunk)) => self.buf.extend_from_slice(&chunk),
+				Some(Err(e)) => {
 					return Some(Err(LdkServerError::new(
 						InternalError,
 						format!("Failed to read event stream: {}", e),
 					)));
 				},
+				None => {
+					if self.trailers_checked {
+						return None;
+					}
+					self.trailers_checked = true;
+					return self.finish_stream().await;
+				},
 			}
 		}
 	}
+
+	async fn finish_stream(&mut self) -> Option<Result<EventEnvelope, LdkServerError>> {
+		match self.body.trailers().await {
+			Ok(Some(trailers)) => {
+				if let Some(error) = grpc_error_from_headers(&trailers) {
+					return Some(Err(error));
+				}
+			},
+			Ok(None) => {},
+			Err(e) => {
+				return Some(Err(LdkServerError::new(
+					InternalError,
+					format!("Failed to read event stream trailers: {}", e),
+				)));
+			},
+		}
+
+		if self.buf.is_empty() {
+			None
+		} else {
+			Some(Err(LdkServerError::new(
+				InternalError,
+				"Event stream ended with an incomplete gRPC frame",
+			)))
+		}
+	}
+}
+
+fn build_streaming_client(server_cert_pem: &[u8]) -> Result<StreamingClient, String> {
+	let mut pem_reader = Cursor::new(server_cert_pem);
+	let certs =
+		certs(&mut pem_reader).map_err(|e| format!("Failed to parse server certificate: {e}"))?;
+	if certs.is_empty() {
+		return Err("Failed to parse server certificate: no certificates found in PEM".to_string());
+	}
+
+	let mut roots = RootCertStore::empty();
+	let (added, _ignored) = roots.add_parsable_certificates(&certs);
+	if added == 0 {
+		return Err("Failed to build streaming client: certificate was not accepted".to_string());
+	}
+
+	let tls_config = ClientConfig::builder()
+		.with_safe_defaults()
+		.with_root_certificates(roots)
+		.with_no_client_auth();
+	let connector = HttpsConnectorBuilder::new()
+		.with_tls_config(tls_config)
+		.https_only()
+		.enable_http2()
+		.build();
+
+	Ok(HyperClient::builder().http2_only(true).build(connector))
 }
 
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use hyper::Body;
+	use reqwest::header::HeaderValue;
+
+	#[test]
+	fn test_grpc_error_from_headers_ignores_ok_status() {
+		let mut headers = HeaderMap::new();
+		headers.insert("grpc-status", HeaderValue::from_static("0"));
+		assert!(grpc_error_from_headers(&headers).is_none());
+	}
+
+	#[test]
+	fn test_grpc_error_from_headers_decodes_message() {
+		let mut headers = HeaderMap::new();
+		headers.insert("grpc-status", HeaderValue::from_static("3"));
+		headers.insert("grpc-message", HeaderValue::from_static("bad%20request"));
+
+		let err = grpc_error_from_headers(&headers).unwrap();
+		assert_eq!(err.error_code, InvalidRequestError);
+		assert_eq!(err.message, "bad request");
+	}
+
+	#[test]
+	fn test_grpc_code_to_error_marks_unavailable_streams() {
+		let err = grpc_code_to_error(14, "server shutting down".to_string());
+		assert_eq!(err.error_code, InternalError);
+		assert_eq!(err.message, "gRPC stream became unavailable: server shutting down");
+	}
+
+	#[tokio::test]
+	async fn test_event_stream_surfaces_terminal_grpc_status() {
+		let (mut sender, body) = Body::channel();
+		let mut trailers = HeaderMap::new();
+		trailers.insert("grpc-status", HeaderValue::from_static("14"));
+		trailers.insert("grpc-message", HeaderValue::from_static("server%20restarting"));
+		sender.send_trailers(trailers).await.unwrap();
+		drop(sender);
+
+		let mut stream = EventStream { body, buf: Vec::new(), trailers_checked: false };
+
+		let result = stream.next_event().await.unwrap().unwrap_err();
+		assert_eq!(result.error_code, InternalError);
+		assert_eq!(result.message, "gRPC stream became unavailable: server restarting");
+		assert!(stream.next_event().await.is_none());
+	}
 
 	#[test]
 	fn test_hex_val() {