mz_aws_util/
lib.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use std::net::IpAddr;
11
12use aws_config::{BehaviorVersion, ConfigLoader};
13use aws_smithy_http_client::tls::{self, rustls_provider::CryptoMode};
14use aws_smithy_runtime_api::client::dns::{DnsFuture, ResolveDns, ResolveDnsError};
15use aws_smithy_runtime_api::client::http::{HttpClient, SharedHttpClient};
16
17use crate::http_pool::{PooledHttpClient, SwitchedHttpClient};
18
19pub use crate::http_pool::HttpPoolMetrics;
20
21mod http_pool;
22#[cfg(feature = "s3")]
23pub mod s3;
24#[cfg(feature = "s3")]
25pub mod s3_uploader;
26
27/// Creates an AWS SDK configuration loader with the defaults for the latest
28/// behavior version plus some Materialize-specific overrides.
29pub fn defaults() -> ConfigLoader {
30    // Use the SDK's latest behavior version. We already pin the crate versions,
31    // and CI puts version upgrades through rigorous testing, so we're happy to
32    // take the latest behavior version. We can adjust this in the future as
33    // necessary, if the AWS SDK ships a new behavior version that causes
34    // trouble.
35    let behavior_version = BehaviorVersion::latest();
36
37    // This is the only method allowed to call `aws_config::defaults`.
38    #[allow(clippy::disallowed_methods)]
39    let loader = aws_config::defaults(behavior_version);
40
41    // Install our custom HTTP client.
42    let loader = loader.http_client(http_client());
43
44    loader
45}
46
47/// Returns an HTTP client for use with the AWS SDK that is appropriately
48/// configured for Materialize.
49pub fn http_client() -> impl HttpClient {
50    aws_smithy_http_client::Builder::new()
51        .tls_provider(tls_provider())
52        .build_https()
53}
54
55/// Returns an HTTP client for use with the AWS SDK, like [`http_client`], that
56/// opens more connections to a host as requests to it pile up.
57///
58/// Use it for servers that speak HTTP/2, like GCS. Over HTTP/2, [`http_client`]
59/// sends every request to a host over one connection, which caps throughput.
60///
61/// A request opens a new connection once every connection has
62/// `max_in_flight()` requests in flight, read per request. 0 counts as 1.
63pub fn pooled_http_client(
64    max_in_flight: impl Fn() -> usize + Send + Sync + 'static,
65    metrics: HttpPoolMetrics,
66) -> impl HttpClient {
67    PooledHttpClient::new(
68        || SharedHttpClient::new(http_client()),
69        max_in_flight,
70        metrics,
71    )
72}
73
74/// Returns an HTTP client for use with the AWS SDK that sends each request
75/// through [`pooled_http_client`] if `use_pool` returns true when the request
76/// is sent, and through [`http_client`] otherwise.
77pub fn switched_http_client(
78    use_pool: impl Fn() -> bool + Send + Sync + 'static,
79    pool_max_in_flight: impl Fn() -> usize + Send + Sync + 'static,
80    pool_metrics: HttpPoolMetrics,
81) -> impl HttpClient {
82    SwitchedHttpClient::new(
83        SharedHttpClient::new(http_client()),
84        SharedHttpClient::new(pooled_http_client(pool_max_in_flight, pool_metrics)),
85        use_pool,
86    )
87}
88
89/// Returns an AWS SDK HTTP client whose DNS resolver delegates to
90/// [`mz_ore::netio::resolve_address`].
91///
92/// Only the IP resolution step is overridden. The SDK still uses the original
93/// hostname for SNI and TLS certificate validation, so HTTPS endpoints work
94/// unchanged.
95pub fn http_client_with_resolver(enforce_external_addresses: bool) -> SharedHttpClient {
96    aws_smithy_http_client::Builder::new()
97        .tls_provider(tls_provider())
98        .build_with_resolver(MzAwsResolver {
99            enforce_external_addresses,
100        })
101}
102
103/// rustls on the aws-lc-rs provider, trusting the system's root certificates.
104fn tls_provider() -> tls::Provider {
105    // TODO(SEC-218): the rest of the workspace is still migrating from OpenSSL
106    // to rustls on aws-lc-rs.
107    tls::Provider::Rustls(CryptoMode::AwsLc)
108}
109
110/// A [`ResolveDns`] implementation that delegates to
111/// [`mz_ore::netio::resolve_address`], used by [`http_client_with_resolver`].
112#[derive(Clone, Debug)]
113struct MzAwsResolver {
114    enforce_external_addresses: bool,
115}
116
117impl MzAwsResolver {
118    async fn resolve(&self, host: &str) -> Result<Vec<IpAddr>, mz_ore::netio::DnsResolutionError> {
119        let ips = mz_ore::netio::resolve_address(host, self.enforce_external_addresses).await?;
120        Ok(ips.into_iter().collect())
121    }
122}
123
124impl ResolveDns for MzAwsResolver {
125    fn resolve_dns<'a>(&'a self, name: &'a str) -> DnsFuture<'a> {
126        DnsFuture::new(async move { self.resolve(name).await.map_err(ResolveDnsError::new) })
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use mz_ore::netio::DnsResolutionError;
133
134    use super::*;
135
136    #[mz_ore::test(tokio::test)]
137    #[cfg_attr(miri, ignore)]
138    async fn resolver_rejects_loopback_when_enforced() {
139        let resolver = MzAwsResolver {
140            enforce_external_addresses: true,
141        };
142        let err = resolver
143            .resolve("127.0.0.1")
144            .await
145            .expect_err("must reject loopback");
146        assert!(
147            matches!(err, DnsResolutionError::PrivateAddress),
148            "got {err:?}"
149        );
150        let err = resolver
151            .resolve_dns("127.0.0.1")
152            .await
153            .expect_err("must reject loopback");
154        let source = std::error::Error::source(&err).expect("wraps the resolution error");
155        assert!(
156            matches!(
157                source.downcast_ref::<DnsResolutionError>(),
158                Some(DnsResolutionError::PrivateAddress)
159            ),
160            "got {source:?}"
161        );
162    }
163
164    #[mz_ore::test(tokio::test)]
165    #[cfg_attr(miri, ignore)]
166    async fn resolver_allows_loopback_when_not_enforced() {
167        let resolver = MzAwsResolver {
168            enforce_external_addresses: false,
169        };
170        let addrs = resolver
171            .resolve_dns("127.0.0.1")
172            .await
173            .expect("loopback should resolve when enforcement is off");
174        assert!(addrs.contains(&IpAddr::from([127, 0, 0, 1])));
175    }
176
177    #[mz_ore::test(tokio::test)]
178    #[cfg_attr(miri, ignore)]
179    async fn resolver_allows_public_when_enforced() {
180        let resolver = MzAwsResolver {
181            enforce_external_addresses: true,
182        };
183        let addrs = resolver
184            .resolve_dns("8.8.8.8")
185            .await
186            .expect("public IP should resolve");
187        assert!(addrs.contains(&IpAddr::from([8, 8, 8, 8])));
188    }
189}
mz_aws_util/lib.rs

mz_aws_util/
lib.rs
