Skip to content

Commit 5ac8353

Browse files
committed
One file for hmac
1 parent 806143d commit 5ac8353

File tree

4 files changed

+129
-103
lines changed

4 files changed

+129
-103
lines changed

lib/executor/src/execution/hmac.rs

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
use std::collections::BTreeMap;
2+
3+
use bytes::{BufMut, Bytes};
4+
use hive_router_config::hmac_signature::{BooleanOrExpression, HMACSignatureConfig};
5+
use hmac::{Hmac, Mac};
6+
use sha2::Sha256;
7+
use vrl::{compiler::Program as VrlProgram, core::Value as VrlValue};
8+
9+
use crate::{
10+
execution::client_request_details::ClientRequestDetails,
11+
executors::{error::SubgraphExecutorError, http::FIRST_EXTENSION_STR},
12+
utils::{
13+
consts::{CLOSE_BRACE, COLON, COMMA, QUOTE},
14+
expression::{compile_expression, execute_expression_with_value},
15+
},
16+
};
17+
18+
#[derive(Debug)]
19+
pub enum BooleanOrProgram {
20+
Boolean(bool),
21+
Program(Box<VrlProgram>),
22+
}
23+
24+
pub fn compile_hmac_config(
25+
config: &HMACSignatureConfig,
26+
) -> Result<BooleanOrProgram, SubgraphExecutorError> {
27+
match &config.enabled {
28+
BooleanOrExpression::Boolean(b) => Ok(BooleanOrProgram::Boolean(*b)),
29+
BooleanOrExpression::Expression { expression } => {
30+
let program = compile_expression(expression, None)
31+
.map_err(SubgraphExecutorError::HMACExpressionBuild)?;
32+
Ok(BooleanOrProgram::Program(Box::new(program)))
33+
}
34+
}
35+
}
36+
type HmacSha256 = Hmac<Sha256>;
37+
38+
pub fn sign_hmac(
39+
hmac_program: &BooleanOrProgram,
40+
hmac_config: &HMACSignatureConfig,
41+
subgraph_name: &str,
42+
client_request: &ClientRequestDetails,
43+
first_extension: &mut bool,
44+
body: &mut Vec<u8>,
45+
) -> Result<(), SubgraphExecutorError> {
46+
let should_sign_hmac = match &hmac_program {
47+
BooleanOrProgram::Boolean(b) => *b,
48+
BooleanOrProgram::Program(expr) => {
49+
// .subgraph
50+
let subgraph_value = VrlValue::Object(BTreeMap::from([(
51+
"name".into(),
52+
VrlValue::Bytes(Bytes::from(subgraph_name.to_owned())),
53+
)]));
54+
// .request
55+
let request_value: VrlValue = client_request.into();
56+
let target_value = VrlValue::Object(BTreeMap::from([
57+
("subgraph".into(), subgraph_value),
58+
("request".into(), request_value),
59+
]));
60+
let result = execute_expression_with_value(expr, target_value);
61+
match result {
62+
Ok(VrlValue::Boolean(b)) => b,
63+
Ok(_) => {
64+
return Err(SubgraphExecutorError::HMACSignatureError(
65+
"HMAC signature expression did not evaluate to a boolean".to_string(),
66+
));
67+
}
68+
Err(e) => {
69+
return Err(SubgraphExecutorError::HMACSignatureError(format!(
70+
"HMAC signature expression evaluation error: {}",
71+
e
72+
)));
73+
}
74+
}
75+
}
76+
};
77+
78+
if should_sign_hmac {
79+
if hmac_config.secret.is_empty() {
80+
return Err(SubgraphExecutorError::HMACSignatureError(
81+
"HMAC signature secret is empty".to_string(),
82+
));
83+
}
84+
let mut mac = HmacSha256::new_from_slice(hmac_config.secret.as_bytes()).map_err(|e| {
85+
SubgraphExecutorError::HMACSignatureError(format!(
86+
"Failed to create HMAC instance: {}",
87+
e
88+
))
89+
})?;
90+
let mut body_without_extensions = body.clone();
91+
body_without_extensions.put(CLOSE_BRACE);
92+
mac.update(&body_without_extensions);
93+
let result = mac.finalize();
94+
let result_bytes = result.into_bytes();
95+
if *first_extension {
96+
body.put(FIRST_EXTENSION_STR);
97+
*first_extension = false;
98+
} else {
99+
body.put(COMMA);
100+
}
101+
body.put(QUOTE);
102+
body.put(hmac_config.extension_name.as_bytes());
103+
body.put(QUOTE);
104+
body.put(COLON);
105+
let hmac_hex = hex::encode(result_bytes);
106+
body.put(QUOTE);
107+
body.put(hmac_hex.as_bytes());
108+
body.put(QUOTE);
109+
}
110+
Ok(())
111+
}

lib/executor/src/execution/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
pub mod client_request_details;
22
pub mod error;
3+
pub mod hmac;
34
pub mod jwt_forward;
45
pub mod plan;
56
pub mod rewrites;

lib/executor/src/executors/http.rs

Lines changed: 10 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,24 @@
1-
use std::collections::BTreeMap;
21
use std::sync::Arc;
32

3+
use crate::execution::hmac::{sign_hmac, BooleanOrProgram};
44
use crate::executors::common::HttpExecutionResponse;
55
use crate::executors::dedupe::{request_fingerprint, ABuildHasher, SharedResponse};
6-
use crate::utils::expression::execute_expression_with_value;
76
use dashmap::DashMap;
87
use hive_router_config::HiveRouterConfig;
98
use tokio::sync::OnceCell;
109

1110
use async_trait::async_trait;
1211

1312
use bytes::{BufMut, Bytes, BytesMut};
14-
use hmac::{Hmac, Mac};
1513
use http::HeaderMap;
1614
use http::HeaderValue;
1715
use http_body_util::BodyExt;
1816
use http_body_util::Full;
1917
use hyper::Version;
2018
use hyper_tls::HttpsConnector;
2119
use hyper_util::client::legacy::{connect::HttpConnector, Client};
22-
use sha2::Sha256;
2320
use tokio::sync::Semaphore;
2421
use tracing::debug;
25-
use vrl::compiler::Program as VrlProgram;
2622

2723
use crate::executors::common::HttpExecutionRequest;
2824
use crate::executors::error::SubgraphExecutorError;
@@ -32,7 +28,6 @@ use crate::utils::consts::COLON;
3228
use crate::utils::consts::COMMA;
3329
use crate::utils::consts::QUOTE;
3430
use crate::{executors::common::SubgraphExecutor, json_writer::write_and_escape_string};
35-
use vrl::core::Value as VrlValue;
3631

3732
#[derive(Debug)]
3833
pub struct HTTPSubgraphExecutor {
@@ -48,18 +43,10 @@ pub struct HTTPSubgraphExecutor {
4843

4944
const FIRST_VARIABLE_STR: &[u8] = b",\"variables\":{";
5045
const FIRST_QUOTE_STR: &[u8] = b"{\"query\":";
51-
const FIRST_EXTENSION_STR: &[u8] = b",\"extensions\":{";
46+
pub const FIRST_EXTENSION_STR: &[u8] = b",\"extensions\":{";
5247

5348
pub type HttpClient = Client<HttpsConnector<HttpConnector>, Full<Bytes>>;
5449

55-
type HmacSha256 = Hmac<Sha256>;
56-
57-
#[derive(Debug)]
58-
pub enum BooleanOrProgram {
59-
Boolean(bool),
60-
Program(Box<VrlProgram>),
61-
}
62-
6350
impl HTTPSubgraphExecutor {
6451
pub fn new(
6552
subgraph_name: String,
@@ -92,80 +79,6 @@ impl HTTPSubgraphExecutor {
9279
}
9380
}
9481

95-
fn sign_hmac(
96-
&self,
97-
execution_request: &HttpExecutionRequest,
98-
body: &mut Vec<u8>,
99-
first_extension: &mut bool,
100-
) -> Result<(), SubgraphExecutorError> {
101-
let should_sign_hmac = match &self.should_sign_hmac.as_ref() {
102-
BooleanOrProgram::Boolean(b) => *b,
103-
BooleanOrProgram::Program(expr) => {
104-
// .subgraph
105-
let subgraph_value = VrlValue::Object(BTreeMap::from([(
106-
"name".into(),
107-
VrlValue::Bytes(Bytes::from(self.subgraph_name.to_owned())),
108-
)]));
109-
// .request
110-
let request_value: VrlValue = execution_request.client_request.into();
111-
let target_value = VrlValue::Object(BTreeMap::from([
112-
("subgraph".into(), subgraph_value),
113-
("request".into(), request_value),
114-
]));
115-
let result = execute_expression_with_value(expr, target_value);
116-
match result {
117-
Ok(VrlValue::Boolean(b)) => b,
118-
Ok(_) => {
119-
return Err(SubgraphExecutorError::HMACSignatureError(
120-
"HMAC signature expression did not evaluate to a boolean".to_string(),
121-
));
122-
}
123-
Err(e) => {
124-
return Err(SubgraphExecutorError::HMACSignatureError(format!(
125-
"HMAC signature expression evaluation error: {}",
126-
e
127-
)));
128-
}
129-
}
130-
}
131-
};
132-
133-
if should_sign_hmac {
134-
if self.config.hmac_signature.secret.is_empty() {
135-
return Err(SubgraphExecutorError::HMACSignatureError(
136-
"HMAC signature secret is empty".to_string(),
137-
));
138-
}
139-
let mut mac = HmacSha256::new_from_slice(self.config.hmac_signature.secret.as_bytes())
140-
.map_err(|e| {
141-
SubgraphExecutorError::HMACSignatureError(format!(
142-
"Failed to create HMAC instance: {}",
143-
e
144-
))
145-
})?;
146-
let mut body_without_extensions = body.clone();
147-
body_without_extensions.put(CLOSE_BRACE);
148-
mac.update(&body_without_extensions);
149-
let result = mac.finalize();
150-
let result_bytes = result.into_bytes();
151-
if *first_extension {
152-
body.put(FIRST_EXTENSION_STR);
153-
*first_extension = false;
154-
} else {
155-
body.put(COMMA);
156-
}
157-
body.put(QUOTE);
158-
body.put(self.config.hmac_signature.extension_name.as_bytes());
159-
body.put(QUOTE);
160-
body.put(COLON);
161-
let hmac_hex = hex::encode(result_bytes);
162-
body.put(QUOTE);
163-
body.put(hmac_hex.as_bytes());
164-
body.put(QUOTE);
165-
}
166-
Ok(())
167-
}
168-
16982
fn build_request_body<'exec, 'req>(
17083
&self,
17184
execution_request: &HttpExecutionRequest<'exec, 'req>,
@@ -213,7 +126,14 @@ impl HTTPSubgraphExecutor {
213126
let mut first_extension = true;
214127

215128
if !self.config.hmac_signature.is_disabled() {
216-
self.sign_hmac(execution_request, &mut body, &mut first_extension)?;
129+
sign_hmac(
130+
&self.should_sign_hmac,
131+
&self.config.hmac_signature,
132+
&self.subgraph_name,
133+
execution_request.client_request,
134+
&mut first_extension,
135+
&mut body,
136+
)?;
217137
}
218138

219139
if let Some(extensions) = &execution_request.extensions {

lib/executor/src/executors/map.rs

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,7 @@ use std::{
55

66
use bytes::{BufMut, Bytes, BytesMut};
77
use dashmap::DashMap;
8-
use hive_router_config::{
9-
hmac_signature::BooleanOrExpression, override_subgraph_urls::UrlOrExpression, HiveRouterConfig,
10-
};
8+
use hive_router_config::{override_subgraph_urls::UrlOrExpression, HiveRouterConfig};
119
use http::Uri;
1210
use hyper_tls::HttpsConnector;
1311
use hyper_util::{
@@ -19,14 +17,17 @@ use tracing::error;
1917
use vrl::{compiler::Program as VrlProgram, core::Value as VrlValue};
2018

2119
use crate::{
22-
execution::client_request_details::ClientRequestDetails,
20+
execution::{
21+
client_request_details::ClientRequestDetails,
22+
hmac::{compile_hmac_config, BooleanOrProgram},
23+
},
2324
executors::{
2425
common::{
2526
HttpExecutionRequest, HttpExecutionResponse, SubgraphExecutor, SubgraphExecutorBoxedArc,
2627
},
2728
dedupe::{ABuildHasher, SharedResponse},
2829
error::SubgraphExecutorError,
29-
http::{BooleanOrProgram, HTTPSubgraphExecutor, HttpClient},
30+
http::{HTTPSubgraphExecutor, HttpClient},
3031
},
3132
response::graphql_error::GraphQLError,
3233
utils::expression::{compile_expression, execute_expression_with_value},
@@ -65,14 +66,7 @@ impl SubgraphExecutorMap {
6566

6667
let max_connections_per_host = config.traffic_shaping.max_connections_per_host;
6768

68-
let should_sign_hmac = match &config.hmac_signature.enabled {
69-
BooleanOrExpression::Boolean(b) => BooleanOrProgram::Boolean(*b),
70-
BooleanOrExpression::Expression { expression } => {
71-
let program = compile_expression(expression, None)
72-
.map_err(SubgraphExecutorError::HMACExpressionBuild)?;
73-
BooleanOrProgram::Program(Box::new(program))
74-
}
75-
};
69+
let should_sign_hmac = compile_hmac_config(&config.hmac_signature)?;
7670

7771
Ok(SubgraphExecutorMap {
7872
executors_by_subgraph: Default::default(),

0 commit comments

Comments
 (0)