Skip to content

Commit 5104c6d

Browse files
Fix transport type to schema mapping and add interval field
- Map transport types to URI schemes (http/https->grpc, http1/https1->http1, dns->dns) - Error if URI already contains query parameters - Add optional interval field to TransportConfig - Build proper URIs with schemes and query parameters Co-authored-by: Hulto <hulto@users.noreply.github.com>
1 parent 5575298 commit 5104c6d

File tree

1 file changed

+31
-9
lines changed

1 file changed

+31
-9
lines changed

implants/lib/pb/build.rs

Lines changed: 31 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ struct TransportConfig {
1111
#[serde(rename = "type")]
1212
transport_type: String,
1313
extra: String,
14+
#[serde(default)]
15+
interval: Option<u64>,
1416
}
1517

1618
#[derive(Debug, Deserialize)]
@@ -65,8 +67,8 @@ fn parse_yaml_config() -> Result<bool, Box<dyn std::error::Error>> {
6567
for transport in &config.transports {
6668
// Validate transport type
6769
let transport_type_lower = transport.transport_type.to_lowercase();
68-
if !["grpc", "http1", "dns"].contains(&transport_type_lower.as_str()) {
69-
return Err(format!("Invalid transport type '{}'. Must be one of: GRPC, http1, DNS", transport.transport_type).into());
70+
if !["grpc", "http1", "dns", "http", "https", "https1"].contains(&transport_type_lower.as_str()) {
71+
return Err(format!("Invalid transport type '{}'. Must be one of: GRPC, http, https, http1, https1, DNS", transport.transport_type).into());
7072
}
7173

7274
// Validate that extra is valid JSON
@@ -75,19 +77,39 @@ fn parse_yaml_config() -> Result<bool, Box<dyn std::error::Error>> {
7577
.map_err(|e| format!("Invalid JSON in 'extra' field for transport '{}': {}", transport.uri, e))?;
7678
}
7779

78-
// Build DSN part with query parameters
79-
let mut dsn_part = transport.uri.clone();
80+
// Map transport type to URI scheme
81+
let scheme = match transport_type_lower.as_str() {
82+
"http" | "https" => "grpc",
83+
"http1" | "https1" => "http1",
84+
"dns" => "dns",
85+
"grpc" => "grpc",
86+
_ => return Err(format!("Unhandled transport type '{}'", transport.transport_type).into()),
87+
};
88+
89+
// Parse URI and check if it already has query parameters
90+
let uri = &transport.uri;
91+
if uri.contains('?') {
92+
return Err(format!("URI '{}' already contains query parameters. Query parameters should not be in the URI field.", uri).into());
93+
}
94+
95+
// Build DSN with scheme and query parameters
96+
let mut dsn_part = format!("{}://{}", scheme, uri);
8097

81-
// Check if URI already has query parameters
82-
let separator = if dsn_part.contains('?') { "&" } else { "?" };
98+
// Start query parameters
99+
dsn_part.push('?');
83100

84-
// Add transport type as query parameter
85-
dsn_part.push_str(&format!("{}transport={}", separator, transport_type_lower));
101+
// Add interval as query parameter if present
102+
if let Some(interval) = transport.interval {
103+
dsn_part.push_str(&format!("interval={}", interval));
104+
}
86105

87106
// Add extra as query parameter if not empty
88107
if !transport.extra.is_empty() {
108+
if transport.interval.is_some() {
109+
dsn_part.push('&');
110+
}
89111
let encoded_extra = urlencoding::encode(&transport.extra);
90-
dsn_part.push_str(&format!("&extra={}", encoded_extra));
112+
dsn_part.push_str(&format!("extra={}", encoded_extra));
91113
}
92114

93115
dsn_parts.push(dsn_part);

0 commit comments

Comments
 (0)