Skip to content

Commit 0a412c9

Browse files
freshtonicfreshtonic
authored
Merge pull request #384 from cipherstash/fix/file-descriptor-leaks
fix(proxy): prevent file descriptor leakage
2 parents 6906b18 + fdc3bd6 commit 0a412c9

5 files changed

Lines changed: 180 additions & 44 deletions

File tree

‎packages/cipherstash-proxy-integration/src/common.rs‎

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,40 @@
11
#![allow(dead_code)]
22

3+
//! # Connection Management for Integration Tests
4+
//!
5+
//! ## Preventing File Descriptor Leaks
6+
//!
7+
//! Integration tests should reuse database connections within each test to prevent
8+
//! file descriptor exhaustion. Creating a new connection for every database operation
9+
//! causes connections to accumulate faster than the proxy's 60-second timeout can clean them up.
10+
//!
11+
//! ### Pattern: Connection Reuse
12+
//!
13+
//! **Good** - Reuse single connection per test:
14+
//! ```rust
15+
//! #[tokio::test]
16+
//! async fn my_test() {
17+
//! let client = connect_with_tls(PROXY).await;
18+
//! clear_with_client(&client).await;
19+
//! insert_with_client(sql, params, &client).await;
20+
//! query_by_with_client(sql, param, &client).await;
21+
//! // Client drops and connection closes cleanly at test end
22+
//! }
23+
//! ```
24+
//!
25+
//! **Bad** - Creates new connection per operation (4+ connections per test):
26+
//! ```rust
27+
//! #[tokio::test]
28+
//! async fn my_test() {
29+
//! clear().await; // Connection 1
30+
//! insert_jsonb().await; // Connection 2
31+
//! query_by(sql, p).await; // Connection 3
32+
//! simple_query(sql).await;// Connection 4
33+
//! }
34+
//! ```
35+
//!
36+
//! Use the `*_with_client()` variants of helper functions to reuse connections.
37+
338
use rand::{distr::Alphanumeric, Rng};
439
use rustls::{
540
client::danger::ServerCertVerifier, crypto::aws_lc_rs::default_provider,
@@ -42,8 +77,10 @@ pub fn random_string() -> String {
4277
}
4378

4479
pub async fn clear() {
45-
let client = connect_with_tls(PROXY).await;
80+
clear_with_client(&connect_with_tls(PROXY).await).await;
81+
}
4682

83+
pub async fn clear_with_client(client: &Client) {
4784
let sql = "TRUNCATE encrypted";
4885
client.simple_query(sql).await.unwrap();
4986

@@ -202,11 +239,33 @@ where
202239
query_by_params(sql, &[param]).await
203240
}
204241

242+
pub async fn query_by_with_client<T>(
243+
sql: &str,
244+
param: &(dyn ToSql + Sync),
245+
client: &Client,
246+
) -> Vec<T>
247+
where
248+
T: for<'a> tokio_postgres::types::FromSql<'a> + Send + Sync,
249+
{
250+
query_by_params_with_client(sql, &[param], client).await
251+
}
252+
205253
pub async fn query_by_params<T>(sql: &str, params: &[&(dyn ToSql + Sync)]) -> Vec<T>
206254
where
207255
T: for<'a> tokio_postgres::types::FromSql<'a> + Send + Sync,
208256
{
209257
let client = connect_with_tls(PROXY).await;
258+
query_by_params_with_client(sql, params, &client).await
259+
}
260+
261+
pub async fn query_by_params_with_client<T>(
262+
sql: &str,
263+
params: &[&(dyn ToSql + Sync)],
264+
client: &Client,
265+
) -> Vec<T>
266+
where
267+
T: for<'a> tokio_postgres::types::FromSql<'a> + Send + Sync,
268+
{
210269
let rows = client.query(sql, params).await.unwrap();
211270
rows.iter().map(|row| row.get(0)).collect::<Vec<T>>()
212271
}
@@ -236,7 +295,6 @@ where
236295
<T as std::str::FromStr>::Err: std::fmt::Debug,
237296
{
238297
let client = connect_with_tls(PROXY).await;
239-
240298
simple_query_with_client(sql, &client).await
241299
}
242300

@@ -285,10 +343,19 @@ pub async fn simple_query_with_null(sql: &str) -> Vec<Option<String>> {
285343

286344
pub async fn insert(sql: &str, params: &[&(dyn ToSql + Sync)]) {
287345
let client = connect_with_tls(PROXY).await;
346+
insert_with_client(sql, params, &client).await;
347+
}
348+
349+
pub async fn insert_with_client(sql: &str, params: &[&(dyn ToSql + Sync)], client: &Client) {
288350
client.query(sql, params).await.unwrap();
289351
}
290352

291353
pub async fn insert_jsonb() -> Value {
354+
let client = connect_with_tls(PROXY).await;
355+
insert_jsonb_with_client(&client).await
356+
}
357+
358+
pub async fn insert_jsonb_with_client(client: &Client) -> Value {
292359
let id = random_id();
293360

294361
let encrypted_jsonb = serde_json::json!({
@@ -305,7 +372,7 @@ pub async fn insert_jsonb() -> Value {
305372

306373
let sql = "INSERT INTO encrypted (id, encrypted_jsonb) VALUES ($1, $2)".to_string();
307374

308-
insert(&sql, &[&id, &encrypted_jsonb]).await;
375+
insert_with_client(&sql, &[&id, &encrypted_jsonb], client).await;
309376

310377
// Verify encryption actually occurred
311378
assert_encrypted_jsonb(id, &encrypted_jsonb).await;

‎packages/cipherstash-proxy-integration/src/select/jsonb_array_elements.rs‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,55 +1,62 @@
11
#[cfg(test)]
22
mod tests {
3-
use crate::common::{clear, insert_jsonb, query_by, simple_query, trace};
3+
use crate::common::{
4+
clear_with_client, connect_with_tls, insert_jsonb_with_client, query_by_with_client,
5+
simple_query_with_client, trace, PROXY,
6+
};
47
use crate::support::assert::assert_expected;
58
use crate::support::json_path::JsonPath;
69
use serde_json::Value;
10+
use tokio_postgres::Client;
711

8-
async fn select_jsonb(selector: &str, expected: &[Value]) {
12+
async fn select_jsonb(selector: &str, expected: &[Value], client: &Client) {
913
let selector = JsonPath::new(selector);
1014

1115
let sql =
1216
"SELECT jsonb_array_elements(jsonb_path_query(encrypted_jsonb, $1)) FROM encrypted";
13-
let actual = query_by::<Value>(sql, &selector).await;
17+
let actual = query_by_with_client::<Value>(sql, &selector, client).await;
1418

1519
assert_expected(expected, &actual);
1620

1721
let sql = format!("SELECT jsonb_array_elements(jsonb_path_query(encrypted_jsonb, '{selector}')) FROM encrypted");
18-
let actual = simple_query::<Value>(&sql).await;
22+
let actual = simple_query_with_client::<Value>(&sql, client).await;
1923

2024
assert_expected(expected, &actual);
2125
}
2226

2327
#[tokio::test]
2428
async fn select_jsonb_array_elements_with_string() {
2529
trace();
30+
let client = connect_with_tls(PROXY).await;
2631

27-
clear().await;
28-
insert_jsonb().await;
32+
clear_with_client(&client).await;
33+
insert_jsonb_with_client(&client).await;
2934

3035
let expected = vec![Value::from("hello"), Value::from("world")];
31-
select_jsonb("$.array_string[@]", &expected).await;
36+
select_jsonb("$.array_string[@]", &expected, &client).await;
3237
}
3338

3439
#[tokio::test]
3540
async fn select_jsonb_array_elements_with_numeric() {
3641
trace();
42+
let client = connect_with_tls(PROXY).await;
3743

38-
clear().await;
39-
insert_jsonb().await;
44+
clear_with_client(&client).await;
45+
insert_jsonb_with_client(&client).await;
4046

4147
let expected = vec![Value::from(42), Value::from(84)];
42-
select_jsonb("$.array_number[@]", &expected).await;
48+
select_jsonb("$.array_number[@]", &expected, &client).await;
4349
}
4450

4551
#[tokio::test]
4652
async fn select_jsonb_array_elements_with_unknown_field() {
4753
trace();
54+
let client = connect_with_tls(PROXY).await;
4855

49-
clear().await;
50-
insert_jsonb().await;
56+
clear_with_client(&client).await;
57+
insert_jsonb_with_client(&client).await;
5158

5259
let expected = vec![];
53-
select_jsonb("$.blah", &expected).await;
60+
select_jsonb("$.blah", &expected, &client).await;
5461
}
5562
}

‎packages/cipherstash-proxy-integration/src/select/jsonb_path_query.rs‎

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
#[cfg(test)]
22
mod tests {
3-
use crate::common::{clear, insert_jsonb, query_by, simple_query, trace};
3+
use crate::common::{
4+
clear_with_client, connect_with_tls, insert_jsonb_with_client, query_by_with_client,
5+
simple_query_with_client, trace, PROXY,
6+
};
47
use crate::support::assert::assert_expected;
58
use crate::support::json_path::JsonPath;
69
use serde::de::DeserializeOwned;
710
use serde_json::Value;
11+
use tokio_postgres::Client;
812

9-
async fn select_jsonb<T>(selector: &str, value: T)
13+
async fn select_jsonb<T>(selector: &str, value: T, client: &Client)
1014
where
1115
T: DeserializeOwned,
1216
serde_json::Value: From<T>,
@@ -17,86 +21,86 @@ mod tests {
1721
let expected = vec![value];
1822

1923
let sql = "SELECT jsonb_path_query(encrypted_jsonb, $1) FROM encrypted";
20-
let actual = query_by::<Value>(sql, &selector).await;
24+
let actual = query_by_with_client::<Value>(sql, &selector, client).await;
2125

2226
assert_expected(&expected, &actual);
2327

2428
let sql = format!("SELECT jsonb_path_query(encrypted_jsonb, '{selector}') FROM encrypted");
25-
let actual = simple_query::<Value>(&sql).await;
29+
let actual = simple_query_with_client::<Value>(&sql, client).await;
2630

2731
assert_expected(&expected, &actual);
2832
}
2933

3034
#[tokio::test]
3135
async fn select_jsonb_path_query_number() {
3236
trace();
37+
let client = connect_with_tls(PROXY).await;
3338

34-
clear().await;
39+
clear_with_client(&client).await;
40+
insert_jsonb_with_client(&client).await;
3541

36-
insert_jsonb().await;
37-
38-
select_jsonb("$.number", 42).await;
42+
select_jsonb("$.number", 42, &client).await;
3943
}
4044

4145
#[tokio::test]
4246
async fn select_jsonb_path_query_string() {
4347
trace();
48+
let client = connect_with_tls(PROXY).await;
4449

45-
clear().await;
46-
47-
insert_jsonb().await;
50+
clear_with_client(&client).await;
51+
insert_jsonb_with_client(&client).await;
4852

49-
select_jsonb("$.nested.string", "world".to_string()).await;
53+
select_jsonb("$.nested.string", "world".to_string(), &client).await;
5054
}
5155

5256
#[tokio::test]
5357
async fn select_jsonb_path_query_value() {
5458
trace();
59+
let client = connect_with_tls(PROXY).await;
5560

56-
clear().await;
57-
58-
insert_jsonb().await;
61+
clear_with_client(&client).await;
62+
insert_jsonb_with_client(&client).await;
5963

6064
let v = serde_json::json!({
6165
"number": 1815,
6266
"string": "world",
6367
});
6468

65-
select_jsonb("$.nested", v).await;
69+
select_jsonb("$.nested", v, &client).await;
6670
}
6771

6872
#[tokio::test]
6973
async fn select_jsonb_path_query_with_unknown() {
7074
trace();
75+
let client = connect_with_tls(PROXY).await;
7176

72-
clear().await;
73-
74-
insert_jsonb().await;
77+
clear_with_client(&client).await;
78+
insert_jsonb_with_client(&client).await;
7579

7680
let selector = JsonPath::new("$.vtha");
7781

7882
let expected = vec![];
7983

8084
let sql = "SELECT jsonb_path_query(encrypted_jsonb, $1) as selected FROM encrypted";
81-
let actual = query_by::<Value>(sql, &selector).await;
85+
let actual = query_by_with_client::<Value>(sql, &selector, &client).await;
8286

8387
assert_expected(&expected, &actual);
8488

8589
let sql = format!(
8690
"SELECT jsonb_path_query(encrypted_jsonb, '{selector}') as selected FROM encrypted"
8791
);
88-
let actual = simple_query::<Value>(&sql).await;
92+
let actual = simple_query_with_client::<Value>(&sql, &client).await;
8993

9094
assert_expected(&expected, &actual);
9195
}
9296

9397
#[tokio::test]
9498
async fn select_jsonb_path_query_with_alias() {
9599
trace();
100+
let client = connect_with_tls(PROXY).await;
96101

97-
clear().await;
98-
99-
insert_jsonb().await;
102+
clear_with_client(&client).await;
103+
insert_jsonb_with_client(&client).await;
100104

101105
let value = serde_json::json!({
102106
"number": 1815,
@@ -108,14 +112,14 @@ mod tests {
108112
let expected = vec![value];
109113

110114
let sql = "SELECT jsonb_path_query(encrypted_jsonb, $1) as selected FROM encrypted";
111-
let actual = query_by::<Value>(sql, &selector).await;
115+
let actual = query_by_with_client::<Value>(sql, &selector, &client).await;
112116

113117
assert_expected(&expected, &actual);
114118

115119
let sql = format!(
116120
"SELECT jsonb_path_query(encrypted_jsonb, '{selector}') as selected FROM encrypted"
117121
);
118-
let actual = simple_query::<Value>(&sql).await;
122+
let actual = simple_query_with_client::<Value>(&sql, &client).await;
119123

120124
assert_expected(&expected, &actual);
121125
}

0 commit comments

Comments
 (0)