Python: TLS Support for servers (#2002)

* Add `PyTlsConfig` struct

* Support TLS server

* Add TLS support to Pokemon service

* Make sure to create `tokio::net::TcpListener` in a Tokio context

* Fix doc link

* Add missing imports in tests

* Add `tls::Listener`

* Reload TLS config periodically

* Add context to `TODO`

* Return `&'static str` from `base_url()`

* Flatten `match` in `tls::Listener`

* Propogate listener errors but ignore handshake errors in `tls::Listener`

* Add tests to `tls::Listener`

* Add test to `tls::Listener` to make sure we are propogating listener errors

* Use `PathBuf` instead of plain `String`s for paths

Author: unexge

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonApplicationGenerator.kt

@@ -264,17 +264,18 @@ class PythonApplicationGenerator(
                     Ok(())
                 }
                 /// Main entrypoint: start the server on multiple workers.
-                ##[pyo3(text_signature = "(${'$'}self, address, port, backlog, workers)")]
+                ##[pyo3(text_signature = "(${'$'}self, address, port, backlog, workers, tls)")]
                 pub fn run(
                     &mut self,
                     py: #{pyo3}::Python,
                     address: Option<String>,
                     port: Option<i32>,
                     backlog: Option<i32>,
                     workers: Option<usize>,
+                    tls: Option<#{SmithyPython}::tls::PyTlsConfig>,
                 ) -> #{pyo3}::PyResult<()> {
                     use #{SmithyPython}::PyApp;
-                    self.run_server(py, address, port, backlog, workers)
+                    self.run_server(py, address, port, backlog, workers, tls)
                 }
                 /// Lambda entrypoint: start the server on Lambda.
                 ##[cfg(feature = "aws-lambda")]
@@ -287,17 +288,18 @@ class PythonApplicationGenerator(
                     self.run_lambda_handler(py)
                 }
                 /// Build the service and start a single worker.
-                ##[pyo3(text_signature = "(${'$'}self, socket, worker_number)")]
+                ##[pyo3(text_signature = "(${'$'}self, socket, worker_number, tls)")]
                 pub fn start_worker(
                     &mut self,
                     py: pyo3::Python,
                     socket: &pyo3::PyCell<#{SmithyPython}::PySocket>,
                     worker_number: isize,
+                    tls: Option<#{SmithyPython}::tls::PyTlsConfig>,
                 ) -> pyo3::PyResult<()> {
                     use #{SmithyPython}::PyApp;
                     let event_loop = self.configure_python_event_loop(py)?;
                     let service = self.build_and_configure_service(py, event_loop)?;
-                    self.start_hyper_worker(py, socket, event_loop, service, worker_number)
+                    self.start_hyper_worker(py, socket, event_loop, service, worker_number, tls)
                 }
                 """,
                 *codegenScope,

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerModuleGenerator.kt

@@ -48,6 +48,7 @@ class PythonServerModuleGenerator(
                 renderPySocketType()
                 renderPyLogging()
                 renderPyMiddlewareTypes()
+                renderPyTlsTypes()
                 renderPyApplicationType()
             }
         }
@@ -162,6 +163,22 @@ class PythonServerModuleGenerator(
         )
     }
 
+    private fun RustWriter.renderPyTlsTypes() {
+        rustTemplate(
+            """
+            let tls = #{pyo3}::types::PyModule::new(py, "tls")?;
+            tls.add_class::<#{SmithyPython}::tls::PyTlsConfig>()?;
+            pyo3::py_run!(
+                py,
+                tls,
+                "import sys; sys.modules['$libName.tls'] = tls"
+            );
+            m.add_submodule(tls)?;
+            """,
+            *codegenScope,
+        )
+    }
+
     // Render Python application type.
     private fun RustWriter.renderPyApplicationType() {
         rustTemplate(

rust-runtime/aws-smithy-http-server-python/Cargo.toml

@@ -26,6 +26,9 @@ futures = "0.3"
 http = "0.2"
 hyper = { version = "0.14.20", features = ["server", "http1", "http2", "tcp", "stream"] }
 lambda_http = { version = "0.7.1", optional = true }
+tls-listener = { version = "0.5.1", features = ["rustls", "hyper-h2"] }
+rustls-pemfile = "1.0.1"
+tokio-rustls = "0.23.4"
 num_cpus = "1.13.1"
 parking_lot = "0.12.1"
 pin-project-lite = "0.2"
@@ -47,6 +50,8 @@ futures-util = "0.3"
 tower-test = "0.4"
 tokio-test = "0.4"
 pyo3-asyncio = { version = "0.17.0", features = ["testing", "attributes", "tokio-runtime"] }
+rcgen = "0.10.0"
+hyper-rustls = { version = "0.23.1", features = ["http2"] }
 
 [[test]]
 name = "middleware_tests"

rust-runtime/aws-smithy-http-server-python/examples/pokemon-service-test/Cargo.toml

@@ -9,8 +9,13 @@ description = "Run tests against the Python server implementation"
 [dev-dependencies]
 command-group = "1.0"
 tokio = { version = "1.20.1", features = ["full"] }
+serial_test = "0.9.0"
+rustls-pemfile = "1.0.1"
+tokio-rustls = "0.23.4"
+hyper-rustls = { version = "0.23.0", features = ["http2"] }
 
 # Local paths
 aws-smithy-client  = { path = "../../../aws-smithy-client/", features = ["rustls"] }
+aws-smithy-http = { path = "../../../aws-smithy-http/" }
 aws-smithy-types  = { path = "../../../aws-smithy-types/" }
 pokemon-service-client = { path = "../pokemon-service-client/" }

rust-runtime/aws-smithy-http-server-python/examples/pokemon-service-test/tests/helpers.rs

@@ -3,9 +3,61 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+use std::fs::File;
+use std::io::BufReader;
+use std::process::Command;
+use std::time::Duration;
+
+use aws_smithy_client::{erase::DynConnector, hyper_ext::Adapter};
+use aws_smithy_http::operation::Request;
 use command_group::{CommandGroup, GroupChild};
 use pokemon_service_client::{Builder, Client, Config};
-use std::{process::Command, thread, time::Duration};
+use tokio::time;
+
+const TEST_KEY: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/testdata/localhost.key");
+const TEST_CERT: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/testdata/localhost.crt");
+
+pub type PokemonClient = Client<
+    aws_smithy_client::erase::DynConnector,
+    aws_smithy_client::erase::DynMiddleware<aws_smithy_client::erase::DynConnector>,
+>;
+
+enum PokemonServiceVariant {
+    Http,
+    Http2,
+}
+
+impl PokemonServiceVariant {
+    async fn run_process(&self) -> GroupChild {
+        let mut args = vec!["../pokemon_service.py".to_string()];
+
+        match self {
+            PokemonServiceVariant::Http => {}
+            PokemonServiceVariant::Http2 => {
+                args.push("--enable-tls".to_string());
+                args.push(format!("--tls-key-path={TEST_KEY}"));
+                args.push(format!("--tls-cert-path={TEST_CERT}"));
+            }
+        }
+
+        let process = Command::new("python3")
+            .args(args)
+            .group_spawn()
+            .expect("failed to spawn the Pokémon Service program");
+
+        // The Python interpreter takes a little to startup.
+        time::sleep(Duration::from_secs(2)).await;
+
+        process
+    }
+
+    fn base_url(&self) -> &'static str {
+        match self {
+            PokemonServiceVariant::Http => "http://localhost:13734",
+            PokemonServiceVariant::Http2 => "https://localhost:13734",
+        }
+    }
+}
 
 pub(crate) struct PokemonService {
     // We need to ensure all processes forked by the Python interpreter
@@ -16,15 +68,16 @@ pub(crate) struct PokemonService {
 
 impl PokemonService {
     #[allow(dead_code)]
-    pub(crate) fn run() -> Self {
-        let process = Command::new("python3")
-            .arg("../pokemon_service.py")
-            .group_spawn()
-            .expect("failed to spawn the Pokémon Service program");
-        // The Python interpreter takes a little to startup.
-        thread::sleep(Duration::from_secs(2));
+    pub(crate) async fn run() -> Self {
         Self {
-            child_process: process,
+            child_process: PokemonServiceVariant::Http.run_process().await,
+        }
+    }
+
+    #[allow(dead_code)]
+    pub(crate) async fn run_http2() -> Self {
+        Self {
+            child_process: PokemonServiceVariant::Http2.run_process().await,
         }
     }
 }
@@ -39,19 +92,49 @@ impl Drop for PokemonService {
 }
 
 #[allow(dead_code)]
-pub fn client() -> Client<
-    aws_smithy_client::erase::DynConnector,
-    aws_smithy_client::erase::DynMiddleware<aws_smithy_client::erase::DynConnector>,
-> {
+pub fn client() -> PokemonClient {
+    let base_url = PokemonServiceVariant::Http.base_url();
     let raw_client = Builder::new()
         .rustls_connector(Default::default())
-        .middleware_fn(|mut req| {
-            let http_req = req.http_mut();
-            let uri = format!("http://localhost:13734{}", http_req.uri().path());
-            *http_req.uri_mut() = uri.parse().unwrap();
-            req
-        })
+        .middleware_fn(rewrite_base_url(base_url))
+        .build_dyn();
+    let config = Config::builder().build();
+    Client::with_config(raw_client, config)
+}
+
+#[allow(dead_code)]
+pub fn http2_client() -> PokemonClient {
+    // Create custom cert store and add our test certificate to prevent unknown cert issues.
+    let mut reader = BufReader::new(File::open(TEST_CERT).expect("could not open certificate"));
+    let certs = rustls_pemfile::certs(&mut reader).expect("could not parse certificate");
+    let mut roots = tokio_rustls::rustls::RootCertStore::empty();
+    roots.add_parsable_certificates(&certs);
+
+    let connector = hyper_rustls::HttpsConnectorBuilder::new()
+        .with_tls_config(
+            tokio_rustls::rustls::ClientConfig::builder()
+                .with_safe_defaults()
+                .with_root_certificates(roots)
+                .with_no_client_auth(),
+        )
+        .https_only()
+        .enable_http2()
+        .build();
+
+    let base_url = PokemonServiceVariant::Http2.base_url();
+    let raw_client = Builder::new()
+        .connector(DynConnector::new(Adapter::builder().build(connector)))
+        .middleware_fn(rewrite_base_url(base_url))
         .build_dyn();
     let config = Config::builder().build();
     Client::with_config(raw_client, config)
 }
+
+fn rewrite_base_url(base_url: &'static str) -> impl Fn(Request) -> Request + Clone {
+    move |mut req| {
+        let http_req = req.http_mut();
+        let uri = format!("{base_url}{}", http_req.uri().path());
+        *http_req.uri_mut() = uri.parse().unwrap();
+        req
+    }
+}

rust-runtime/aws-smithy-http-server-python/examples/pokemon-service-test/tests/simple_integration_test.rs

@@ -7,35 +7,43 @@
 // These tests only have access to your crate's public API.
 // See: https://doc.rust-lang.org/book/ch11-03-test-organization.html#integration-tests
 
-use std::time::Duration;
-
-use crate::helpers::{client, PokemonService};
 use aws_smithy_types::error::display::DisplayErrorContext;
-use tokio::time;
+use serial_test::serial;
+
+use crate::helpers::{client, http2_client, PokemonClient, PokemonService};
 
 mod helpers;
 
 #[tokio::test]
+#[serial]
 async fn simple_integration_test() {
-    let _program = PokemonService::run();
-    // Give PokemonSérvice some time to start up.
-    time::sleep(Duration::from_millis(50)).await;
+    let _program = PokemonService::run().await;
+    simple_integration_test_with_client(client()).await;
+}
+
+#[tokio::test]
+#[serial]
+async fn simple_integration_test_http2() {
+    let _program = PokemonService::run_http2().await;
+    simple_integration_test_with_client(http2_client()).await;
+}
 
-    let service_statistics_out = client().get_server_statistics().send().await.unwrap();
+async fn simple_integration_test_with_client(client: PokemonClient) {
+    let service_statistics_out = client.get_server_statistics().send().await.unwrap();
     assert_eq!(0, service_statistics_out.calls_count.unwrap());
 
-    let pokemon_species_output = client()
+    let pokemon_species_output = client
         .get_pokemon_species()
         .name("pikachu")
         .send()
         .await
         .unwrap();
     assert_eq!("pikachu", pokemon_species_output.name().unwrap());
 
-    let service_statistics_out = client().get_server_statistics().send().await.unwrap();
+    let service_statistics_out = client.get_server_statistics().send().await.unwrap();
     assert_eq!(1, service_statistics_out.calls_count.unwrap());
 
-    let pokemon_species_error = client()
+    let pokemon_species_error = client
         .get_pokemon_species()
         .name("some_pokémon")
         .send()
@@ -49,8 +57,8 @@ async fn simple_integration_test() {
         "expected '{message}' to contain '{expected}'"
     );
 
-    let service_statistics_out = client().get_server_statistics().send().await.unwrap();
+    let service_statistics_out = client.get_server_statistics().send().await.unwrap();
     assert_eq!(2, service_statistics_out.calls_count.unwrap());
 
-    let _health_check = client().check_health().send().await.unwrap();
+    let _health_check = client.check_health().send().await.unwrap();
 }

rust-runtime/aws-smithy-http-server-python/examples/pokemon-service-test/tests/testdata/localhost.crt

@@ -0,0 +1,30 @@
+-----BEGIN CERTIFICATE-----
+MIIFGTCCAwGgAwIBAgIUN/FD3OayKwJt9hXNKo4JKxqFSK4wDQYJKoZIhvcNAQEL
+BQAwFDESMBAGA1UEAwwJbG9jYWxob3N0MB4XDTIyMDgxNzE1MjQzMFoXDTMyMDgx
+NDE1MjQzMFowFDESMBAGA1UEAwwJbG9jYWxob3N0MIICIjANBgkqhkiG9w0BAQEF
+AAOCAg8AMIICCgKCAgEAulMGcyA69ioNMT8Kz0CdP2QP5elLNnltBykoqoJwbvKS
+94+l5XA//29M4NpLphHcDxNXx3qB318bixUIPBtu66OiIsTGX8yrYPA4IO3Xt5/2
+wp2z1lNLouyW1+gPaPjKzcrjnHmqHS90CFDQqxdv9I0rIFIQ+U5hm5T9Hjr5xs36
+43l2FXAjeigoEuwtVBDt44yhEyeLSDwFJES3sH73AvpruMdxGv2KDVN4whuajWll
+RLTqpqBvVSM6JbaV/VD2simpZeolSl8yKIenM2PWPdLIHSMEBg6IaYgpSpzoyvmh
+089peAaiJfVrN53QjqDVyaN5os9ST03ZEzXQUI38lpvWGmV9Tcs5WfidLA1EbPjv
+yE1zBbZh0SrP/+EALwkoIRslI8DXvz/9U5Cq7q9U4OHjWB+yjE5/BX6o6hfrqfJ1
+Ldg2fTp/TYEudmefM8eRzx6sdYtTPZBrSpkRgvmxd+6k3QUtsAQhtBTMpvJpWsgs
+sD7Uo6G2JRag53oT/2cxG03Qy5HqySZUK1bpFW03W5FL3Pq6AkpGy1hnSxlifkHp
+si61dbjCV5uRdxRCLyH9fD3HImecet+vnuZlvsP0MAzh0vbli/dcFZ7xUoSqFWnj
+egnPohdOmF6C8kXvWBt51N4jjW+eLxPAr9H0mJtdIvEHWBNNW9iitzGz5Gw0g4sC
+AwEAAaNjMGEwHQYDVR0OBBYEFEoLkB78Z6jgPPmOyf0XnWo/LjA9MB8GA1UdIwQY
+MBaAFEoLkB78Z6jgPPmOyf0XnWo/LjA9MBQGA1UdEQQNMAuCCWxvY2FsaG9zdDAJ
+BgNVHRMEAjAAMA0GCSqGSIb3DQEBCwUAA4ICAQC17OljBEEVYefzk2mwg20AXDtL
+PUJ46hLrUM7BcNBjd8AbtrLH/pdCRCexnv7tzYbwMhDNdqHiIcXDHEMNP3gXryB2
+ckU5ms/LzfKADM2/hrDZiR03XYSL4thjFkQNVfYnk9k7LTv9pKW0b+J2OrMun7+w
+bdXcNw+igvnYiBgNJRo0IC9O5nejqLGWwBfveAJPetxjy6PvBkLqgIw2glivmTrh
+Kdoq/I2/ZcxT0GyhEVIHP9W8Hh5goNm+RbsB/hDYhK+5s2+rL1lwJrwhNBrHhG1u
+CtYmd2rD0J/mGf1cAw7t+hmwW0O7J9BVZw4YL/m4vDAsTO4zaeoAvDwsgQwPzPF1
+rmRtV+7jJHyIP/b021XIdIZU5KsXCCA3+B31mHJF1GLreG7WI+wClRsiNSbP7Zuw
+OnUOTDZc77Y4oaDKl0UL8tz1GNwX5G9U5h+FciTPKCtg1gGiqSkB/3BOON2WaVOb
+6Di9iAoH+dIjvWR/7ez7DAk/ITpGvBXS5RqaIXfB9pSJlVYsGp03ikgng1eJdXy4
+57XZnd47upHH88NTvIH9G/iOXQQCzF3MQXOqrJ/gem3ICeelvOoyNseHLvi8ZEqa
+s693CJWaQAK/jD1mhka7yQzmb/Y1I53crc2UqSxX4FqFYP8xymza4Cg/E6pPJerG
+LE/drJtbrIHTUlJB2Q==
+-----END CERTIFICATE-----