Skip to content

Commit b78367c

Browse files
declanvkDeclan Kelly
andauthored
Record TCP connection local socket address in metadata (#3286)
## Motivation and Context I want to use this field to uniquely identify TCP connection based on their `local_addr` + `remote_addr`. See awslabs/aws-sdk-rust#990 for additional motivation for this change. ## Description - Add a new optional `local_addr` field in the `ConnectionMetadata` struct. - Transfer the `local_addr` `SocketAddress` from the `hyper::HttpInfo` to the `ConnectionMetadata` field. - Add to the `trace-serialize` example program so that it will print out the capture connection values. ## Testing `cargo test` in `rust-runtime/aws-smithy-runtime-api` and `aws-smithy-runtime`. Also ran: ``` thedeck@c889f3b04fb0 examples % cargo run --example trace-serialize Finished dev [unoptimized + debuginfo] target(s) in 0.13s Running `/Users/thedeck/repos/github/declanvk/smithy-rs/target/debug/examples/trace-serialize` 2023-12-06T00:13:15.605555Z INFO lazy_load_identity: aws_smithy_runtime::client::identity::cache::lazy: identity cache miss occurred; added new identity (took Ok(296µs)) 2023-12-06T00:13:15.608344Z INFO trace_serialize: Response received: response=Response { status: StatusCode(200), headers: Headers { headers: {"content-type": HeaderValue { _private: "application/json" }, "content-length": HeaderValue { _private: "17" }, "date": HeaderValue { _private: "Wed, 06 Dec 2023 00:13:15 GMT" }} }, body: SdkBody { inner: BoxBody, retryable: false }, extensions: Extensions } 2023-12-06T00:13:15.608388Z INFO trace_serialize: Captured connection info remote_addr=Some(127.0.0.1:13734) local_addr=Some(127.0.0.1:50199) 2023-12-06T00:13:15.608511Z INFO trace_serialize: Response received POKEMON_SERVICE_URL=http://localhost:13734 response=GetServerStatisticsOutput { calls_count: 0 } ``` You can see the log line with "Captured connection info" contains the `remote_addr` and the `local_addr` fields. ## Checklist - [x] I have updated `CHANGELOG.next.toml` if I made changes to the AWS SDK, generated SDK code, or SDK runtime crates ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ Co-authored-by: Declan Kelly <thedeck@amazon.com>
1 parent 8df5ac8 commit b78367c

File tree

5 files changed

+243
-9
lines changed

5 files changed

+243
-9
lines changed

CHANGELOG.next.toml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,3 +73,9 @@ message = "Fix documentation and examples on HyperConnector and HyperClientBuild
7373
references = ["smithy-rs#3282"]
7474
meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client" }
7575
author = "jdisanti"
76+
77+
[[smithy-rs]]
78+
message = "Expose local socket address from ConnectionMetadata."
79+
references = ["aws-sdk-rust#990"]
80+
meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "client" }
81+
author = "declanvk"

examples/pokemon-service-client-usage/examples/trace-serialize.rs

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
5+
use aws_smithy_runtime::client::http::connection_poisoning::CaptureSmithyConnection;
56
/// This example demonstrates how an interceptor can be written to trace what is being
67
/// serialized / deserialized on the wire.
78
///
@@ -59,7 +60,7 @@ impl Intercept for WireFormatInterceptor {
5960
&self,
6061
context: &BeforeDeserializationInterceptorContextRef<'_>,
6162
_runtime_components: &RuntimeComponents,
62-
_cfg: &mut ConfigBag,
63+
cfg: &mut ConfigBag,
6364
) -> Result<(), BoxError> {
6465
// Get the response type from the context.
6566
let response = context.response();
@@ -70,6 +71,18 @@ impl Intercept for WireFormatInterceptor {
7071
tracing::error!(?response);
7172
}
7273

74+
// Print the connection information
75+
let captured_connection = cfg.load::<CaptureSmithyConnection>().cloned();
76+
if let Some(captured_connection) = captured_connection.and_then(|conn| conn.get()) {
77+
tracing::info!(
78+
remote_addr = ?captured_connection.remote_addr(),
79+
local_addr = ?captured_connection.local_addr(),
80+
"Captured connection info"
81+
);
82+
} else {
83+
tracing::warn!("Connection info is missing!");
84+
}
85+
7386
Ok(())
7487
}
7588
}

rust-runtime/aws-smithy-runtime-api/src/client/connection.rs

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use std::sync::Arc;
1414
pub struct ConnectionMetadata {
1515
is_proxied: bool,
1616
remote_addr: Option<SocketAddr>,
17+
local_addr: Option<SocketAddr>,
1718
poison_fn: Arc<dyn Fn() + Send + Sync>,
1819
}
1920

@@ -25,6 +26,10 @@ impl ConnectionMetadata {
2526
}
2627

2728
/// Create a new [`ConnectionMetadata`].
29+
#[deprecated(
30+
since = "1.1.0",
31+
note = "`ConnectionMetadata::new` is deprecated in favour of `ConnectionMetadata::builder`."
32+
)]
2833
pub fn new(
2934
is_proxied: bool,
3035
remote_addr: Option<SocketAddr>,
@@ -33,21 +38,220 @@ impl ConnectionMetadata {
3338
Self {
3439
is_proxied,
3540
remote_addr,
41+
// need to use builder to set this field
42+
local_addr: None,
3643
poison_fn: Arc::new(poison),
3744
}
3845
}
3946

47+
/// Builder for this connection metadata
48+
pub fn builder() -> ConnectionMetadataBuilder {
49+
ConnectionMetadataBuilder::new()
50+
}
51+
4052
/// Get the remote address for this connection, if one is set.
4153
pub fn remote_addr(&self) -> Option<SocketAddr> {
4254
self.remote_addr
4355
}
56+
57+
/// Get the local address for this connection, if one is set.
58+
pub fn local_addr(&self) -> Option<SocketAddr> {
59+
self.local_addr
60+
}
4461
}
4562

4663
impl Debug for ConnectionMetadata {
4764
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
4865
f.debug_struct("SmithyConnection")
4966
.field("is_proxied", &self.is_proxied)
5067
.field("remote_addr", &self.remote_addr)
68+
.field("local_addr", &self.local_addr)
69+
.finish()
70+
}
71+
}
72+
73+
/// Builder type that is used to construct a [`ConnectionMetadata`] value.
74+
#[derive(Default)]
75+
pub struct ConnectionMetadataBuilder {
76+
is_proxied: Option<bool>,
77+
remote_addr: Option<SocketAddr>,
78+
local_addr: Option<SocketAddr>,
79+
poison_fn: Option<Arc<dyn Fn() + Send + Sync>>,
80+
}
81+
82+
impl Debug for ConnectionMetadataBuilder {
83+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
84+
f.debug_struct("ConnectionMetadataBuilder")
85+
.field("is_proxied", &self.is_proxied)
86+
.field("remote_addr", &self.remote_addr)
87+
.field("local_addr", &self.local_addr)
5188
.finish()
5289
}
5390
}
91+
92+
impl ConnectionMetadataBuilder {
93+
/// Creates a new builder.
94+
pub fn new() -> Self {
95+
Self::default()
96+
}
97+
98+
/// Set whether or not the associated connection is to an HTTP proxy.
99+
pub fn proxied(mut self, proxied: bool) -> Self {
100+
self.set_proxied(Some(proxied));
101+
self
102+
}
103+
104+
/// Set whether or not the associated connection is to an HTTP proxy.
105+
pub fn set_proxied(&mut self, proxied: Option<bool>) -> &mut Self {
106+
self.is_proxied = proxied;
107+
self
108+
}
109+
110+
/// Set the remote address of the connection used.
111+
pub fn remote_addr(mut self, remote_addr: SocketAddr) -> Self {
112+
self.set_remote_addr(Some(remote_addr));
113+
self
114+
}
115+
116+
/// Set the remote address of the connection used.
117+
pub fn set_remote_addr(&mut self, remote_addr: Option<SocketAddr>) -> &mut Self {
118+
self.remote_addr = remote_addr;
119+
self
120+
}
121+
122+
/// Set the local address of the connection used.
123+
pub fn local_addr(mut self, local_addr: SocketAddr) -> Self {
124+
self.set_local_addr(Some(local_addr));
125+
self
126+
}
127+
128+
/// Set the local address of the connection used.
129+
pub fn set_local_addr(&mut self, local_addr: Option<SocketAddr>) -> &mut Self {
130+
self.local_addr = local_addr;
131+
self
132+
}
133+
134+
/// Set a closure which will poison the associated connection.
135+
///
136+
/// A poisoned connection will not be reused for subsequent requests by the pool
137+
pub fn poison_fn(mut self, poison_fn: impl Fn() + Send + Sync + 'static) -> Self {
138+
self.set_poison_fn(Some(poison_fn));
139+
self
140+
}
141+
142+
/// Set a closure which will poison the associated connection.
143+
///
144+
/// A poisoned connection will not be reused for subsequent requests by the pool
145+
pub fn set_poison_fn(
146+
&mut self,
147+
poison_fn: Option<impl Fn() + Send + Sync + 'static>,
148+
) -> &mut Self {
149+
self.poison_fn =
150+
poison_fn.map(|poison_fn| Arc::new(poison_fn) as Arc<dyn Fn() + Send + Sync>);
151+
self
152+
}
153+
154+
/// Build a [`ConnectionMetadata`] value.
155+
///
156+
/// # Panics
157+
///
158+
/// If either the `is_proxied` or `poison_fn` has not been set, then this method will panic
159+
pub fn build(self) -> ConnectionMetadata {
160+
ConnectionMetadata {
161+
is_proxied: self
162+
.is_proxied
163+
.expect("is_proxied should be set for ConnectionMetadata"),
164+
remote_addr: self.remote_addr,
165+
local_addr: self.local_addr,
166+
poison_fn: self
167+
.poison_fn
168+
.expect("poison_fn should be set for ConnectionMetadata"),
169+
}
170+
}
171+
}
172+
173+
#[cfg(test)]
174+
mod tests {
175+
use std::{
176+
net::{IpAddr, Ipv6Addr},
177+
sync::Mutex,
178+
};
179+
180+
use super::*;
181+
182+
const TEST_SOCKET_ADDR: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 100);
183+
184+
#[test]
185+
#[should_panic]
186+
fn builder_panic_missing_proxied() {
187+
ConnectionMetadataBuilder::new()
188+
.poison_fn(|| {})
189+
.local_addr(TEST_SOCKET_ADDR)
190+
.remote_addr(TEST_SOCKET_ADDR)
191+
.build();
192+
}
193+
194+
#[test]
195+
#[should_panic]
196+
fn builder_panic_missing_poison_fn() {
197+
ConnectionMetadataBuilder::new()
198+
.proxied(true)
199+
.local_addr(TEST_SOCKET_ADDR)
200+
.remote_addr(TEST_SOCKET_ADDR)
201+
.build();
202+
}
203+
204+
#[test]
205+
fn builder_all_fields_successful() {
206+
let mutable_flag = Arc::new(Mutex::new(false));
207+
208+
let connection_metadata = ConnectionMetadataBuilder::new()
209+
.proxied(true)
210+
.local_addr(TEST_SOCKET_ADDR)
211+
.remote_addr(TEST_SOCKET_ADDR)
212+
.poison_fn({
213+
let mutable_flag = Arc::clone(&mutable_flag);
214+
move || {
215+
let mut guard = mutable_flag.lock().unwrap();
216+
*guard = !*guard;
217+
}
218+
})
219+
.build();
220+
221+
assert_eq!(connection_metadata.is_proxied, true);
222+
assert_eq!(connection_metadata.remote_addr(), Some(TEST_SOCKET_ADDR));
223+
assert_eq!(connection_metadata.local_addr(), Some(TEST_SOCKET_ADDR));
224+
assert_eq!(*mutable_flag.lock().unwrap(), false);
225+
connection_metadata.poison();
226+
assert_eq!(*mutable_flag.lock().unwrap(), true);
227+
}
228+
229+
#[test]
230+
fn builder_optional_fields_translate() {
231+
let metadata1 = ConnectionMetadataBuilder::new()
232+
.proxied(true)
233+
.poison_fn(|| {})
234+
.build();
235+
236+
assert_eq!(metadata1.local_addr(), None);
237+
assert_eq!(metadata1.remote_addr(), None);
238+
239+
let metadata2 = ConnectionMetadataBuilder::new()
240+
.proxied(true)
241+
.poison_fn(|| {})
242+
.local_addr(TEST_SOCKET_ADDR)
243+
.build();
244+
245+
assert_eq!(metadata2.local_addr(), Some(TEST_SOCKET_ADDR));
246+
assert_eq!(metadata2.remote_addr(), None);
247+
248+
let metadata3 = ConnectionMetadataBuilder::new()
249+
.proxied(true)
250+
.poison_fn(|| {})
251+
.remote_addr(TEST_SOCKET_ADDR)
252+
.build();
253+
254+
assert_eq!(metadata3.local_addr(), None);
255+
assert_eq!(metadata3.remote_addr(), Some(TEST_SOCKET_ADDR));
256+
}
257+
}

rust-runtime/aws-smithy-runtime/src/client/http/connection_poisoning.rs

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,6 @@ impl Intercept for ConnectionPoisoningInterceptor {
100100
type LoaderFn = dyn Fn() -> Option<ConnectionMetadata> + Send + Sync;
101101

102102
/// State for a middleware that will monitor and manage connections.
103-
#[allow(missing_debug_implementations)]
104103
#[derive(Clone, Default)]
105104
pub struct CaptureSmithyConnection {
106105
loader: Arc<Mutex<Option<Box<LoaderFn>>>>,
@@ -154,7 +153,14 @@ mod test {
154153
let retriever = CaptureSmithyConnection::new();
155154
let retriever_clone = retriever.clone();
156155
assert!(retriever.get().is_none());
157-
retriever.set_connection_retriever(|| Some(ConnectionMetadata::new(true, None, || {})));
156+
retriever.set_connection_retriever(|| {
157+
Some(
158+
ConnectionMetadata::builder()
159+
.proxied(true)
160+
.poison_fn(|| {})
161+
.build(),
162+
)
163+
});
158164

159165
assert!(retriever.get().is_some());
160166
assert!(retriever_clone.get().is_some());

rust-runtime/aws-smithy-runtime/src/client/http/hyper_014.rs

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -287,14 +287,19 @@ fn extract_smithy_connection(capture_conn: &CaptureConnection) -> Option<Connect
287287
let mut extensions = Extensions::new();
288288
conn.get_extras(&mut extensions);
289289
let http_info = extensions.get::<HttpInfo>();
290-
let smithy_connection = ConnectionMetadata::new(
291-
conn.is_proxied(),
292-
http_info.map(|info| info.remote_addr()),
293-
move || match capture_conn.connection_metadata().as_ref() {
290+
let mut builder = ConnectionMetadata::builder()
291+
.proxied(conn.is_proxied())
292+
.poison_fn(move || match capture_conn.connection_metadata().as_ref() {
294293
Some(conn) => conn.poison(),
295294
None => tracing::trace!("no connection existed to poison"),
296-
},
297-
);
295+
});
296+
297+
builder
298+
.set_local_addr(http_info.map(|info| info.local_addr()))
299+
.set_remote_addr(http_info.map(|info| info.remote_addr()));
300+
301+
let smithy_connection = builder.build();
302+
298303
Some(smithy_connection)
299304
} else {
300305
None

0 commit comments

Comments
 (0)