Skip to content

Commit 626b680

Browse files
didierofriviadidierofrivia
committed
[refactor, feature] Adding a new GrpcResponse type
* Services also building the response Message Signed-off-by: dd di cesare <[email protected]>
1 parent f9b2873 commit 626b680

File tree

4 files changed

+192
-13
lines changed

4 files changed

+192
-13
lines changed

‎src/envoy/mod.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ pub use {
3737
AttributeContext_Request,
3838
},
3939
base::Metadata,
40-
external_auth::CheckRequest,
40+
external_auth::{CheckRequest, DeniedHttpResponse, OkHttpResponse},
4141
ratelimit::{RateLimitDescriptor, RateLimitDescriptor_Entry},
4242
rls::{RateLimitRequest, RateLimitResponse, RateLimitResponse_Code},
4343
};

‎src/service/auth.rs

+33-2
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ use crate::envoy::{
33
Address, AttributeContext, AttributeContext_HttpRequest, AttributeContext_Peer,
44
AttributeContext_Request, CheckRequest, Metadata, SocketAddress,
55
};
6+
use crate::service::grpc_message::{GrpcMessageResponse, GrpcMessageResult};
67
use chrono::{DateTime, FixedOffset, Timelike};
78
use protobuf::well_known_types::Timestamp;
9+
use protobuf::Message;
810
use proxy_wasm::hostcalls;
9-
use proxy_wasm::types::MapType;
11+
use proxy_wasm::types::{Bytes, MapType};
1012
use std::collections::HashMap;
1113

1214
pub const AUTH_SERVICE_NAME: &str = "envoy.service.auth.v3.Authorization";
@@ -16,10 +18,39 @@ pub struct AuthService;
1618

1719
#[allow(dead_code)]
1820
impl AuthService {
19-
pub fn message(ce_host: String) -> CheckRequest {
21+
pub fn request_message(ce_host: String) -> CheckRequest {
2022
AuthService::build_check_req(ce_host)
2123
}
2224

25+
pub fn response_message(
26+
res_body_bytes: &Bytes,
27+
status_code: u32,
28+
) -> GrpcMessageResult<GrpcMessageResponse> {
29+
if status_code % 2 == 0 {
30+
AuthService::response_message_ok(res_body_bytes)
31+
} else {
32+
AuthService::response_message_denied(res_body_bytes)
33+
}
34+
}
35+
36+
fn response_message_ok(
37+
res_body_bytes: &Bytes,
38+
) -> GrpcMessageResult<GrpcMessageResponse> {
39+
match Message::parse_from_bytes(res_body_bytes) {
40+
Ok(res) => Ok(GrpcMessageResponse::AuthOk(res)),
41+
Err(e) => Err(e),
42+
}
43+
}
44+
45+
fn response_message_denied(
46+
res_body_bytes: &Bytes,
47+
) -> GrpcMessageResult<GrpcMessageResponse> {
48+
match Message::parse_from_bytes(res_body_bytes) {
49+
Ok(res) => Ok(GrpcMessageResponse::AuthDenied(res)),
50+
Err(e) => Err(e),
51+
}
52+
}
53+
2354
fn build_check_req(ce_host: String) -> CheckRequest {
2455
let mut auth_req = CheckRequest::default();
2556
let mut attr = AttributeContext::default();

‎src/service/grpc_message.rs

+146-7
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,16 @@
11
use crate::configuration::ExtensionType;
2-
use crate::envoy::{CheckRequest, RateLimitDescriptor, RateLimitRequest};
2+
use crate::envoy::{
3+
CheckRequest, DeniedHttpResponse, OkHttpResponse, RateLimitDescriptor, RateLimitRequest,
4+
RateLimitResponse,
5+
};
36
use crate::service::auth::AuthService;
47
use crate::service::rate_limit::RateLimitService;
58
use protobuf::reflect::MessageDescriptor;
69
use protobuf::{
7-
Clear, CodedInputStream, CodedOutputStream, Message, ProtobufResult, UnknownFields,
10+
Clear, CodedInputStream, CodedOutputStream, Message, ProtobufError, ProtobufResult,
11+
UnknownFields,
812
};
13+
use proxy_wasm::types::Bytes;
914
use std::any::Any;
1015

1116
#[derive(Clone, Debug)]
@@ -126,11 +131,145 @@ impl GrpcMessageRequest {
126131
descriptors: protobuf::RepeatedField<RateLimitDescriptor>,
127132
) -> Self {
128133
match extension_type {
129-
ExtensionType::RateLimit => GrpcMessageRequest::RateLimit(RateLimitService::message(
130-
domain.clone(),
131-
descriptors,
132-
)),
133-
ExtensionType::Auth => GrpcMessageRequest::Auth(AuthService::message(domain.clone())),
134+
ExtensionType::RateLimit => GrpcMessageRequest::RateLimit(
135+
RateLimitService::request_message(domain.clone(), descriptors),
136+
),
137+
ExtensionType::Auth => {
138+
GrpcMessageRequest::Auth(AuthService::request_message(domain.clone()))
139+
}
134140
}
135141
}
136142
}
143+
144+
#[derive(Clone, Debug)]
145+
pub enum GrpcMessageResponse {
146+
AuthOk(OkHttpResponse),
147+
AuthDenied(DeniedHttpResponse),
148+
RateLimit(RateLimitResponse),
149+
}
150+
151+
impl Default for GrpcMessageResponse {
152+
fn default() -> Self {
153+
GrpcMessageResponse::RateLimit(RateLimitResponse::new())
154+
}
155+
}
156+
157+
impl Clear for GrpcMessageResponse {
158+
fn clear(&mut self) {
159+
todo!()
160+
}
161+
}
162+
163+
impl Message for GrpcMessageResponse {
164+
fn descriptor(&self) -> &'static MessageDescriptor {
165+
match self {
166+
GrpcMessageResponse::AuthOk(res) => res.descriptor(),
167+
GrpcMessageResponse::AuthDenied(res) => res.descriptor(),
168+
GrpcMessageResponse::RateLimit(res) => res.descriptor(),
169+
}
170+
}
171+
172+
fn is_initialized(&self) -> bool {
173+
match self {
174+
GrpcMessageResponse::AuthOk(res) => res.is_initialized(),
175+
GrpcMessageResponse::AuthDenied(res) => res.is_initialized(),
176+
GrpcMessageResponse::RateLimit(res) => res.is_initialized(),
177+
}
178+
}
179+
180+
fn merge_from(&mut self, is: &mut CodedInputStream) -> ProtobufResult<()> {
181+
match self {
182+
GrpcMessageResponse::AuthOk(res) => res.merge_from(is),
183+
GrpcMessageResponse::AuthDenied(res) => res.merge_from(is),
184+
GrpcMessageResponse::RateLimit(res) => res.merge_from(is),
185+
}
186+
}
187+
188+
fn write_to_with_cached_sizes(&self, os: &mut CodedOutputStream) -> ProtobufResult<()> {
189+
match self {
190+
GrpcMessageResponse::AuthOk(res) => res.write_to_with_cached_sizes(os),
191+
GrpcMessageResponse::AuthDenied(res) => res.write_to_with_cached_sizes(os),
192+
GrpcMessageResponse::RateLimit(res) => res.write_to_with_cached_sizes(os),
193+
}
194+
}
195+
196+
fn write_to_bytes(&self) -> ProtobufResult<Vec<u8>> {
197+
match self {
198+
GrpcMessageResponse::AuthOk(res) => res.write_to_bytes(),
199+
GrpcMessageResponse::AuthDenied(res) => res.write_to_bytes(),
200+
GrpcMessageResponse::RateLimit(res) => res.write_to_bytes(),
201+
}
202+
}
203+
204+
fn compute_size(&self) -> u32 {
205+
match self {
206+
GrpcMessageResponse::AuthOk(res) => res.compute_size(),
207+
GrpcMessageResponse::AuthDenied(res) => res.compute_size(),
208+
GrpcMessageResponse::RateLimit(res) => res.compute_size(),
209+
}
210+
}
211+
212+
fn get_cached_size(&self) -> u32 {
213+
match self {
214+
GrpcMessageResponse::AuthOk(res) => res.get_cached_size(),
215+
GrpcMessageResponse::AuthDenied(res) => res.get_cached_size(),
216+
GrpcMessageResponse::RateLimit(res) => res.get_cached_size(),
217+
}
218+
}
219+
220+
fn get_unknown_fields(&self) -> &UnknownFields {
221+
match self {
222+
GrpcMessageResponse::AuthOk(res) => res.get_unknown_fields(),
223+
GrpcMessageResponse::AuthDenied(res) => res.get_unknown_fields(),
224+
GrpcMessageResponse::RateLimit(res) => res.get_unknown_fields(),
225+
}
226+
}
227+
228+
fn mut_unknown_fields(&mut self) -> &mut UnknownFields {
229+
match self {
230+
GrpcMessageResponse::AuthOk(res) => res.mut_unknown_fields(),
231+
GrpcMessageResponse::AuthDenied(res) => res.mut_unknown_fields(),
232+
GrpcMessageResponse::RateLimit(res) => res.mut_unknown_fields(),
233+
}
234+
}
235+
236+
fn as_any(&self) -> &dyn Any {
237+
match self {
238+
GrpcMessageResponse::AuthOk(res) => res.as_any(),
239+
GrpcMessageResponse::AuthDenied(res) => res.as_any(),
240+
GrpcMessageResponse::RateLimit(res) => res.as_any(),
241+
}
242+
}
243+
244+
fn new() -> Self
245+
where
246+
Self: Sized,
247+
{
248+
// Returning default value
249+
GrpcMessageResponse::default()
250+
}
251+
252+
fn default_instance() -> &'static Self
253+
where
254+
Self: Sized,
255+
{
256+
#[allow(non_upper_case_globals)]
257+
static instance: ::protobuf::rt::LazyV2<GrpcMessageResponse> = ::protobuf::rt::LazyV2::INIT;
258+
instance.get(|| GrpcMessageResponse::RateLimit(RateLimitResponse::new()))
259+
}
260+
}
261+
262+
impl GrpcMessageResponse {
263+
pub fn new(
264+
extension_type: ExtensionType,
265+
res_body_bytes: &Bytes,
266+
status_code: u32,
267+
) -> GrpcMessageResult<GrpcMessageResponse> {
268+
match extension_type {
269+
ExtensionType::RateLimit => RateLimitService::response_message(res_body_bytes),
270+
ExtensionType::Auth => AuthService::response_message(res_body_bytes, status_code),
271+
}
272+
}
273+
}
274+
275+
pub type GrpcMessageResult<T> = Result<T, ProtobufError>;

‎src/service/rate_limit.rs

+12-3
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
use crate::envoy::{RateLimitDescriptor, RateLimitRequest};
2-
use protobuf::RepeatedField;
2+
use crate::service::grpc_message::{GrpcMessageResponse, GrpcMessageResult};
3+
use protobuf::{Message, RepeatedField};
4+
use proxy_wasm::types::Bytes;
35

46
pub const RATELIMIT_SERVICE_NAME: &str = "envoy.service.ratelimit.v3.RateLimitService";
57
pub const RATELIMIT_METHOD_NAME: &str = "ShouldRateLimit";
68

79
pub struct RateLimitService;
810

911
impl RateLimitService {
10-
pub fn message(
12+
pub fn request_message(
1113
domain: String,
1214
descriptors: RepeatedField<RateLimitDescriptor>,
1315
) -> RateLimitRequest {
@@ -19,6 +21,13 @@ impl RateLimitService {
1921
cached_size: Default::default(),
2022
}
2123
}
24+
25+
pub fn response_message(res_body_bytes: &Bytes) -> GrpcMessageResult<GrpcMessageResponse> {
26+
match Message::parse_from_bytes(res_body_bytes) {
27+
Ok(res) => Ok(GrpcMessageResponse::RateLimit(res)),
28+
Err(e) => Err(e),
29+
}
30+
}
2231
}
2332

2433
#[cfg(test)]
@@ -37,7 +46,7 @@ mod tests {
3746
field.set_entries(RepeatedField::from_vec(vec![entry]));
3847
let descriptors = RepeatedField::from_vec(vec![field]);
3948

40-
RateLimitService::message(domain.to_string(), descriptors.clone())
49+
RateLimitService::request_message(domain.to_string(), descriptors.clone())
4150
}
4251
#[test]
4352
fn builds_correct_message() {

0 commit comments

Comments
 (0)