This is an automated email from the ASF dual-hosted git repository. hsun pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/incubator-teaclave.git
commit 37d8895e284780b7745fbec46a8126cbd4c93b43 Author: sunhe05 <[email protected]> AuthorDate: Tue Mar 14 09:44:12 2023 +0000 Add usage statistics and quota for function --- examples/c/builtin_echo.c | 3 +- examples/c/builtin_ordered_set_intersect.c | 3 +- examples/rust/builtin_echo/src/main.rs | 1 + .../rust/builtin_ordered_set_intersect/src/main.rs | 1 + sdk/c/teaclave_client_sdk.h | 26 +++++ sdk/python/teaclave.py | 35 +++++- sdk/rust/src/bindings.rs | 5 + sdk/rust/src/lib.rs | 74 ++++++++++++- services/frontend/enclave/src/service.rs | 34 ++++-- services/management/enclave/src/error.rs | 2 + services/management/enclave/src/lib.rs | 1 + services/management/enclave/src/service.rs | 118 ++++++++++++++++++--- .../src/proto/teaclave_frontend_service.proto | 12 +++ .../src/proto/teaclave_management_service.proto | 1 + services/proto/src/teaclave_frontend_service.rs | 83 ++++++++++++++- services/proto/src/teaclave_management_service.rs | 4 + tests/functional/enclave/src/management_service.rs | 1 + types/src/function.rs | 29 +++++ 18 files changed, 399 insertions(+), 34 deletions(-) diff --git a/examples/c/builtin_echo.c b/examples/c/builtin_echo.c index 9a838356..0fed2594 100644 --- a/examples/c/builtin_echo.c +++ b/examples/c/builtin_echo.c @@ -37,7 +37,8 @@ const char *register_function_request_serialized = QUOTE({ "arguments" : [{"key": "message", "default_value": "", "allow_overwrite": true}], "inputs" : [], "outputs" : [], - "user_allowlist": [] + "user_allowlist": [], + "usage_quota": -1 }); const char *create_task_request_serialized = QUOTE({ diff --git a/examples/c/builtin_ordered_set_intersect.c b/examples/c/builtin_ordered_set_intersect.c index 2b278202..70037e13 100644 --- a/examples/c/builtin_ordered_set_intersect.c +++ b/examples/c/builtin_ordered_set_intersect.c @@ -69,7 +69,8 @@ const char *register_function_request_serialized = QUOTE( {"name": "output_result1", "description": "Output data.", "optional": false}, {"name": "output_result2", "description": "Output data.", "optional": false} ], - "user_allowlist": ["user0", "user1"] + "user_allowlist": ["user0", "user1"], + "usage_quota": -1 }); const char *create_task_request_serialized = QUOTE( diff --git a/examples/rust/builtin_echo/src/main.rs b/examples/rust/builtin_echo/src/main.rs index 378b43ba..f595274d 100644 --- a/examples/rust/builtin_echo/src/main.rs +++ b/examples/rust/builtin_echo/src/main.rs @@ -69,6 +69,7 @@ fn echo(message: &str) -> Result<Vec<u8>> { )]), None, None, + None, )?; println!( diff --git a/examples/rust/builtin_ordered_set_intersect/src/main.rs b/examples/rust/builtin_ordered_set_intersect/src/main.rs index d39127d4..75ead1f6 100644 --- a/examples/rust/builtin_ordered_set_intersect/src/main.rs +++ b/examples/rust/builtin_ordered_set_intersect/src/main.rs @@ -130,6 +130,7 @@ impl Client { teaclave_client_sdk::FunctionOutput::new("output_result1", "Output data.", false), teaclave_client_sdk::FunctionOutput::new("output_result2", "Output data.", false), ]), + None, )?; self.client.get_function(&function_id)?; let function_arguments = hashmap!("order" => "ascending", "save_log" => "true"); // Order can be ascending or desending diff --git a/sdk/c/teaclave_client_sdk.h b/sdk/c/teaclave_client_sdk.h index affaa623..95519a4a 100644 --- a/sdk/c/teaclave_client_sdk.h +++ b/sdk/c/teaclave_client_sdk.h @@ -307,6 +307,32 @@ int teaclave_get_function_serialized(struct FrontendClient *client, char *serialized_response, size_t *serialized_response_len); +/** + * Send JSON serialized request to the service with the `client` and + * get the serialized response. + * + * # Arguments + * + * * `client`: service client. + * * `serialized_request`; JSON serialized request + * * `serialized_response`: buffer to store the JSON serialized response. + * * `serialized_response_len`: length of the allocated + * `serialized_response`, will be set as the length of + * `serialized_response` when return successfully. + * + * # Return + * + * The function returns 0 for success. On error, the function returns 1. + * + * # Safety + * + * Inconsistent length of allocated buffer may caused overflow. + */ +int teaclave_get_function_usage_stats_serialized(struct FrontendClient *client, + const char *serialized_request, + char *serialized_response, + size_t *serialized_response_len); + /** * Send JSON serialized request to the service with the `client` and * get the serialized response. diff --git a/sdk/python/teaclave.py b/sdk/python/teaclave.py index c4a4873c..1dad6159 100644 --- a/sdk/python/teaclave.py +++ b/sdk/python/teaclave.py @@ -568,7 +568,7 @@ class RegisterFunctionRequest(Request): executor_type: str, public: bool, payload: List[int], arguments: List[FunctionArgument], inputs: List[FunctionInput], outputs: List[FunctionOutput], - user_allowlist: List[str]): + user_allowlist: List[str], usage_quota: int): self.request = "register_function" self.metadata = metadata self.name = name @@ -580,6 +580,7 @@ class RegisterFunctionRequest(Request): self.inputs = inputs self.outputs = outputs self.user_allowlist = user_allowlist + self.usage_quota = usage_quota class UpdateFunctionRequest(Request): @@ -588,7 +589,7 @@ class UpdateFunctionRequest(Request): description: str, executor_type: str, public: bool, payload: List[int], arguments: List[FunctionArgument], inputs: List[FunctionInput], outputs: List[FunctionOutput], - user_allowlist: List[str]): + user_allowlist: List[str], usage_quota: int): self.request = "update_function" self.metadata = metadata self.function_id = function_id @@ -601,6 +602,7 @@ class UpdateFunctionRequest(Request): self.inputs = inputs self.outputs = outputs self.user_allowlist = user_allowlist + self.usage_quota = usage_quota class ListFunctionsRequest(Request): @@ -635,6 +637,14 @@ class GetFunctionRequest(Request): self.function_id = function_id +class GetFunctionUsageStatsRequest(Request): + + def __init__(self, metadata: Metadata, function_id: str): + self.request = "get_function_usage_stats" + self.metadata = metadata + self.function_id = function_id + + class RegisterInputFileRequest(Request): def __init__(self, metadata: Metadata, url: str, cmac: List[int], @@ -763,13 +773,14 @@ class FrontendService(TeaclaveService): inputs: List[FunctionInput] = [], outputs: List[FunctionOutput] = [], user_allowlist: List[str] = [], + usage_quota: int = -1, ): self.check_metadata() self.check_channel() request = RegisterFunctionRequest(self.metadata, name, description, executor_type, public, payload, arguments, inputs, outputs, - user_allowlist) + user_allowlist, usage_quota) _write_message(self.channel, request) response = _read_message(self.channel) if response["result"] == "ok": @@ -792,13 +803,14 @@ class FrontendService(TeaclaveService): inputs: List[FunctionInput] = [], outputs: List[FunctionOutput] = [], user_allowlist: List[str] = [], + usage_quota: int = -1, ): self.check_metadata() self.check_channel() request = UpdateFunctionRequest(self.metadata, function_id, name, description, executor_type, public, payload, arguments, inputs, outputs, - user_allowlist) + user_allowlist, usage_quota) _write_message(self.channel, request) response = _read_message(self.channel) if response["result"] == "ok": @@ -834,6 +846,21 @@ class FrontendService(TeaclaveService): reason = response["request_error"] raise TeaclaveException(f"Failed to get function ({reason})") + def get_function_usage_stats(self, user_id: str, function_id: str): + self.check_metadata() + self.check_channel() + request = GetFunctionUsageStatsRequest(self.metadata, function_id) + _write_message(self.channel, request) + response = _read_message(self.channel) + if response["result"] == "ok": + return response["content"] + else: + reason = "unknown" + if "request_error" in response: + reason = response["request_error"] + raise TeaclaveException( + f"Failed to get function usage statistics ({reason})") + def delete_function(self, function_id: str): self.check_metadata() self.check_channel() diff --git a/sdk/rust/src/bindings.rs b/sdk/rust/src/bindings.rs index 5b388ed3..f48b8c71 100644 --- a/sdk/rust/src/bindings.rs +++ b/sdk/rust/src/bindings.rs @@ -460,6 +460,11 @@ generate_function_serialized!( teaclave_get_function_serialized, get_function_serialized ); +generate_function_serialized!( + FrontendClient, + teaclave_get_function_usage_stats_serialized, + get_function_usage_stats_serialized +); generate_function_serialized!( FrontendClient, teaclave_register_input_file_serialized, diff --git a/sdk/rust/src/lib.rs b/sdk/rust/src/lib.rs index 65d80fed..a212624e 100644 --- a/sdk/rust/src/lib.rs +++ b/sdk/rust/src/lib.rs @@ -35,13 +35,15 @@ pub use teaclave_proto::teaclave_frontend_service::GetFunctionResponse as Functi pub use teaclave_proto::teaclave_frontend_service::{ ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse, CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, CreateTaskResponse, - GetFunctionRequest, GetFunctionResponse, GetTaskRequest, GetTaskResponse, InvokeTaskRequest, + GetFunctionRequest, GetFunctionResponse, GetFunctionUsageStatsRequest, + GetFunctionUsageStatsResponse, GetTaskRequest, GetTaskResponse, InvokeTaskRequest, InvokeTaskResponse, RegisterFunctionRequest, RegisterFunctionRequestBuilder, RegisterFunctionResponse, RegisterInputFileRequest, RegisterInputFileResponse, RegisterOutputFileRequest, RegisterOutputFileResponse, }; pub use teaclave_types::{ - EnclaveInfo, Executor, FileCrypto, FunctionArgument, FunctionInput, FunctionOutput, TaskResult, + EnclaveInfo, Executor, FileCrypto, FunctionArgument, FunctionInput, FunctionOutput, + FunctionUsage, TaskResult, }; pub mod bindings; @@ -214,6 +216,7 @@ impl FrontendClient { arguments: Option<Vec<FunctionArgument>>, inputs: Option<Vec<FunctionInput>>, outputs: Option<Vec<FunctionOutput>>, + usage_quota: Option<i32>, ) -> Result<String> { let executor_type = executor_type.try_into()?; let mut builder = RegisterFunctionRequestBuilder::new() @@ -233,6 +236,10 @@ impl FrontendClient { if let Some(outputs) = outputs { builder = builder.outputs(outputs); } + if let Some(usage_quota) = usage_quota { + let usage_quota = (usage_quota >= 0).then_some(usage_quota); + builder = builder.usage_quota(usage_quota) + } let request = builder.build(); let response = self.register_function_with_request(request)?; @@ -266,6 +273,37 @@ impl FrontendClient { Ok(response) } + pub fn get_function_usage_stats_serialized( + &mut self, + serialized_request: &str, + ) -> Result<String> { + let request: frontend_proto::GetFunctionUsageStatsRequest = + serde_json::from_str(serialized_request)?; + let response: frontend_proto::GetFunctionUsageStatsResponse = self + .get_function_usage_stats_with_request(request.try_into()?)? + .into(); + let serialized_response = serde_json::to_string(&response)?; + + Ok(serialized_response) + } + + pub fn get_function_usage_stats(&mut self, function_id: &str) -> Result<i32> { + let function_id = function_id.try_into()?; + let request = GetFunctionUsageStatsRequest::new(function_id); + let response = self.get_function_usage_stats_with_request(request)?; + + Ok(response.current_usage) + } + + pub fn get_function_usage_stats_with_request( + &mut self, + request: GetFunctionUsageStatsRequest, + ) -> Result<GetFunctionUsageStatsResponse> { + let response = self.api_client.get_function_usage_stats(request)?; + + Ok(response) + } + pub fn register_input_file_with_request( &mut self, request: RegisterInputFileRequest, @@ -595,6 +633,7 @@ mod tests { Some(vec![FunctionArgument::new("message", "", true)]), None, None, + Some(2), ) .unwrap(); let _ = client.get_function(&function_id).unwrap(); @@ -602,7 +641,7 @@ mod tests { let task_id = client .create_task( &function_id, - Some(function_arguments), + Some(function_arguments.clone()), "builtin", None, None, @@ -611,8 +650,37 @@ mod tests { let _ = client.invoke_task(&task_id).unwrap(); let (result, log) = client.get_task_result(&task_id).unwrap(); + let usage_number = client.get_function_usage_stats(&function_id).unwrap(); assert_eq!(result, b"Hello, Teaclave!"); assert!(log.is_empty()); + assert_eq!(1, usage_number); + + let task_id = client + .create_task( + &function_id, + Some(function_arguments.clone()), + "builtin", + None, + None, + ) + .unwrap(); + let _ = client.invoke_task(&task_id).unwrap(); + let (result, log) = client.get_task_result(&task_id).unwrap(); + let usage_number = client.get_function_usage_stats(&function_id).unwrap(); + assert_eq!(result, b"Hello, Teaclave!"); + assert!(log.is_empty()); + assert_eq!(2, usage_number); + + let task_id = client + .create_task( + &function_id, + Some(function_arguments), + "builtin", + None, + None, + ) + .unwrap(); + assert!(client.invoke_task(&task_id).is_err()); } #[test] diff --git a/services/frontend/enclave/src/service.rs b/services/frontend/enclave/src/service.rs index b59bdd5e..bc072769 100644 --- a/services/frontend/enclave/src/service.rs +++ b/services/frontend/enclave/src/service.rs @@ -31,15 +31,15 @@ use teaclave_proto::teaclave_frontend_service::{ ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse, CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, CreateTaskResponse, DeleteFunctionRequest, DeleteFunctionResponse, DisableFunctionRequest, DisableFunctionResponse, - GetFunctionRequest, GetFunctionResponse, GetInputFileRequest, GetInputFileResponse, - GetOutputFileRequest, GetOutputFileResponse, GetTaskRequest, GetTaskResponse, - InvokeTaskRequest, InvokeTaskResponse, ListFunctionsRequest, ListFunctionsResponse, - RegisterFunctionRequest, RegisterFunctionResponse, RegisterFusionOutputRequest, - RegisterFusionOutputResponse, RegisterInputFileRequest, RegisterInputFileResponse, - RegisterInputFromOutputRequest, RegisterInputFromOutputResponse, RegisterOutputFileRequest, - RegisterOutputFileResponse, TeaclaveFrontend, UpdateFunctionRequest, UpdateFunctionResponse, - UpdateInputFileRequest, UpdateInputFileResponse, UpdateOutputFileRequest, - UpdateOutputFileResponse, + GetFunctionRequest, GetFunctionResponse, GetFunctionUsageStatsRequest, + GetFunctionUsageStatsResponse, GetInputFileRequest, GetInputFileResponse, GetOutputFileRequest, + GetOutputFileResponse, GetTaskRequest, GetTaskResponse, InvokeTaskRequest, InvokeTaskResponse, + ListFunctionsRequest, ListFunctionsResponse, RegisterFunctionRequest, RegisterFunctionResponse, + RegisterFusionOutputRequest, RegisterFusionOutputResponse, RegisterInputFileRequest, + RegisterInputFileResponse, RegisterInputFromOutputRequest, RegisterInputFromOutputResponse, + RegisterOutputFileRequest, RegisterOutputFileResponse, TeaclaveFrontend, UpdateFunctionRequest, + UpdateFunctionResponse, UpdateInputFileRequest, UpdateInputFileResponse, + UpdateOutputFileRequest, UpdateOutputFileResponse, }; use teaclave_proto::teaclave_management_service::TeaclaveManagementClient; use teaclave_rpc::endpoint::Endpoint; @@ -108,6 +108,7 @@ enum Endpoints { GetInputFile, RegisterFunction, GetFunction, + GetFunctionUsageStats, UpdateFunction, ListFunctions, DeleteFunction, @@ -149,7 +150,7 @@ fn authorize(claims: &UserAuthClaims, request: Endpoints) -> bool { | Endpoints::ApproveTask | Endpoints::InvokeTask | Endpoints::CancelTask => role.is_data_owner(), - Endpoints::GetFunction | Endpoints::ListFunctions => { + Endpoints::GetFunction | Endpoints::ListFunctions | Endpoints::GetFunctionUsageStats => { role.is_function_owner() || role.is_data_owner() } } @@ -271,6 +272,7 @@ impl TeaclaveFrontend for TeaclaveFrontendService { Endpoints::RegisterInputFromOutput ) } + fn get_output_file( &self, request: Request<GetOutputFileRequest>, @@ -331,6 +333,18 @@ impl TeaclaveFrontend for TeaclaveFrontendService { ) } + fn get_function_usage_stats( + &self, + request: Request<GetFunctionUsageStatsRequest>, + ) -> TeaclaveServiceResponseResult<GetFunctionUsageStatsResponse> { + authentication_and_forward_to_management!( + self, + request, + get_function_usage_stats, + Endpoints::GetFunctionUsageStats + ) + } + fn delete_function( &self, request: Request<DeleteFunctionRequest>, diff --git a/services/management/enclave/src/error.rs b/services/management/enclave/src/error.rs index 1df0b992..ed330144 100644 --- a/services/management/enclave/src/error.rs +++ b/services/management/enclave/src/error.rs @@ -46,6 +46,8 @@ pub(crate) enum ManagementServiceError { TaskInvokeError, #[error("failed to cancel task, reason: {0}")] TaskCancelError(String), + #[error("function quota has been used up")] + FunctionQuotaError, } impl From<ManagementServiceError> for TeaclaveServiceResponseError { diff --git a/services/management/enclave/src/lib.rs b/services/management/enclave/src/lib.rs index c98f698a..66a4ec0b 100644 --- a/services/management/enclave/src/lib.rs +++ b/services/management/enclave/src/lib.rs @@ -140,6 +140,7 @@ pub mod tests { service::tests::handle_input_file, service::tests::handle_output_file, service::tests::handle_function, + service::tests::check_function_quota, service::tests::handle_task, service::tests::handle_staged_task, ) diff --git a/services/management/enclave/src/service.rs b/services/management/enclave/src/service.rs index 1557de35..2cb4b4f8 100644 --- a/services/management/enclave/src/service.rs +++ b/services/management/enclave/src/service.rs @@ -23,15 +23,15 @@ use teaclave_proto::teaclave_frontend_service::{ ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse, CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, CreateTaskResponse, DeleteFunctionRequest, DeleteFunctionResponse, DisableFunctionRequest, DisableFunctionResponse, - GetFunctionRequest, GetFunctionResponse, GetInputFileRequest, GetInputFileResponse, - GetOutputFileRequest, GetOutputFileResponse, GetTaskRequest, GetTaskResponse, - InvokeTaskRequest, InvokeTaskResponse, ListFunctionsRequest, ListFunctionsResponse, - RegisterFunctionRequest, RegisterFunctionResponse, RegisterFusionOutputRequest, - RegisterFusionOutputResponse, RegisterInputFileRequest, RegisterInputFileResponse, - RegisterInputFromOutputRequest, RegisterInputFromOutputResponse, RegisterOutputFileRequest, - RegisterOutputFileResponse, UpdateFunctionRequest, UpdateFunctionResponse, - UpdateInputFileRequest, UpdateInputFileResponse, UpdateOutputFileRequest, - UpdateOutputFileResponse, + GetFunctionRequest, GetFunctionResponse, GetFunctionUsageStatsRequest, + GetFunctionUsageStatsResponse, GetInputFileRequest, GetInputFileResponse, GetOutputFileRequest, + GetOutputFileResponse, GetTaskRequest, GetTaskResponse, InvokeTaskRequest, InvokeTaskResponse, + ListFunctionsRequest, ListFunctionsResponse, RegisterFunctionRequest, RegisterFunctionResponse, + RegisterFusionOutputRequest, RegisterFusionOutputResponse, RegisterInputFileRequest, + RegisterInputFileResponse, RegisterInputFromOutputRequest, RegisterInputFromOutputResponse, + RegisterOutputFileRequest, RegisterOutputFileResponse, UpdateFunctionRequest, + UpdateFunctionResponse, UpdateInputFileRequest, UpdateInputFileResponse, + UpdateOutputFileRequest, UpdateOutputFileResponse, }; use teaclave_proto::teaclave_management_service::TeaclaveManagement; use teaclave_proto::teaclave_storage_service::{ @@ -297,6 +297,12 @@ impl TeaclaveManagement for TeaclaveManagementService { } } + let usage = FunctionUsage { + function_id: function.id, + ..Default::default() + }; + self.write_to_db(&usage)?; + let response = RegisterFunctionResponse::new(function.external_id()); Ok(response) } @@ -364,6 +370,42 @@ impl TeaclaveManagement for TeaclaveManagementService { } } + // access control: + // function.public || request.role == PlatformAdmin || requested user_id in the user_allowlist + fn get_function_usage_stats( + &self, + request: Request<GetFunctionUsageStatsRequest>, + ) -> TeaclaveServiceResponseResult<GetFunctionUsageStatsResponse> { + let user_id = get_request_user_id(&request)?; + let role = get_request_role(&request)?; + let request = request.message; + let function: Function = self + .read_from_db(&request.function_id) + .map_err(|_| ManagementServiceError::InvalidFunctionId)?; + + ensure!( + function.public + || role == UserRole::PlatformAdmin + || function.user_allowlist.contains(&user_id.to_string()), + ManagementServiceError::PermissionDenied + ); + + let usage = FunctionUsage { + function_id: function.id, + ..Default::default() + }; + let external_id = usage.external_id(); + let function_usage = self + .read_from_db::<FunctionUsage>(&external_id) + .map_err(|_| ManagementServiceError::InvalidFunctionId)?; + let function_quota = function.usage_quota.unwrap_or(-1); + let response = GetFunctionUsageStatsResponse { + function_quota, + current_usage: function_usage.use_numbers, + }; + Ok(response) + } + // access control: function.owner == user_id fn delete_function( &self, @@ -692,6 +734,22 @@ impl TeaclaveManagement for TeaclaveManagementService { log::debug!("InvokeTask: get function: {:?}", function); + let usage = FunctionUsage { + function_id: function.id, + ..Default::default() + }; + let external_id = usage.external_id(); + let mut function_usage = self + .read_from_db::<FunctionUsage>(&external_id) + .map_err(|_| ManagementServiceError::InvalidFunctionId)?; + let function_current_use_numbers = function_usage.use_numbers; + + if let Some(quota) = function.usage_quota { + if quota <= function_current_use_numbers { + return Err(ManagementServiceError::FunctionQuotaError.into()); + } + } + let mut task: Task<Stage> = ts.try_into().map_err(|e| { log::warn!("Stage state error: {:?}", e); ManagementServiceError::TaskInvokeError @@ -707,6 +765,8 @@ impl TeaclaveManagement for TeaclaveManagementService { let ts: TaskState = task.into(); self.write_to_db(&ts)?; + function_usage.use_numbers = function_current_use_numbers + 1; + self.write_to_db(&function_usage)?; Ok(InvokeTaskResponse) } @@ -894,8 +954,9 @@ impl TeaclaveManagementService { let function_arg1 = FunctionArgument::new("arg1", "", true); let function_arg2 = FunctionArgument::new("arg2", "", true); + let function_id = Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap(); let function = FunctionBuilder::new() - .id(Uuid::parse_str("00000000-0000-0000-0000-000000000001").unwrap()) + .id(function_id) .name("mock-func-1") .description("mock-desc") .payload(b"mock-payload".to_vec()) @@ -906,12 +967,20 @@ impl TeaclaveManagementService { .owner("teaclave".to_string()) .build(); + let function_usage = FunctionUsage { + function_id, + use_numbers: 0, + }; + self.write_to_db(&function)?; + self.write_to_db(&function_usage)?; let function_output = FunctionOutput::new("output", "output_desc", false); let function_arg1 = FunctionArgument::new("arg1", "", true); + let function_id = Uuid::parse_str("00000000-0000-0000-0000-000000000002").unwrap(); + let function = FunctionBuilder::new() - .id(Uuid::parse_str("00000000-0000-0000-0000-000000000002").unwrap()) + .id(function_id) .name("mock-func-2") .description("mock-desc") .payload(b"mock-payload".to_vec()) @@ -921,10 +990,17 @@ impl TeaclaveManagementService { .owner("teaclave".to_string()) .build(); + let function_usage = FunctionUsage { + function_id, + use_numbers: 0, + }; + self.write_to_db(&function)?; + self.write_to_db(&function_usage)?; + let function_id = Uuid::parse_str("00000000-0000-0000-0000-000000000003").unwrap(); let function = FunctionBuilder::new() - .id(Uuid::parse_str("00000000-0000-0000-0000-000000000003").unwrap()) + .id(function_id) .name("mock-func-3") .description("Private mock function") .payload(b"mock-payload".to_vec()) @@ -934,7 +1010,14 @@ impl TeaclaveManagementService { .user_allowlist(vec!["mock_user".to_string(), "mock_user1".to_string()]) .build(); + let function_usage = FunctionUsage { + function_id, + use_numbers: 0, + }; + self.write_to_db(&function)?; + self.write_to_db(&function_usage)?; + Ok(()) } } @@ -1016,6 +1099,17 @@ pub mod tests { debug!("function: {:?}", deserialized_function); } + pub fn check_function_quota() { + let function = FunctionBuilder::new().build(); + assert_eq!(function.usage_quota, None); + + let function = FunctionBuilder::new().usage_quota(Some(-5)).build(); + assert_eq!(function.usage_quota, None); + + let function = FunctionBuilder::new().usage_quota(Some(5)).build(); + assert_eq!(function.usage_quota, Some(5)); + } + pub fn handle_task() { let function_arg = FunctionArgument::new("arg", "", true); let function = FunctionBuilder::new() diff --git a/services/proto/src/proto/teaclave_frontend_service.proto b/services/proto/src/proto/teaclave_frontend_service.proto index afd0dbb4..07c23987 100644 --- a/services/proto/src/proto/teaclave_frontend_service.proto +++ b/services/proto/src/proto/teaclave_frontend_service.proto @@ -128,6 +128,7 @@ message RegisterFunctionRequest { repeated FunctionInput inputs = 10; repeated FunctionOutput outputs = 11; repeated string user_allowlist = 12; + int32 usage_quota = 13; } message RegisterFunctionResponse { @@ -145,6 +146,7 @@ message UpdateFunctionRequest { repeated FunctionInput inputs = 10; repeated FunctionOutput outputs = 11; repeated string user_allowlist = 12; + int32 usage_quota = 13; } message UpdateFunctionResponse { @@ -168,6 +170,15 @@ message GetFunctionResponse { repeated string user_allowlist = 12; } +message GetFunctionUsageStatsRequest { + string function_id = 1; +} + +message GetFunctionUsageStatsResponse { + int32 function_quota = 1; + int32 current_usage = 2; +} + message DeleteFunctionRequest { string function_id = 1; } @@ -264,6 +275,7 @@ service TeaclaveFrontend { rpc GetInputFile (GetInputFileRequest) returns (GetInputFileResponse); rpc RegisterFunction (RegisterFunctionRequest) returns (RegisterFunctionResponse); rpc GetFunction (GetFunctionRequest) returns (GetFunctionResponse); + rpc GetFunctionUsageStats (GetFunctionUsageStatsRequest) returns (GetFunctionUsageStatsResponse); rpc UpdateFunction (UpdateFunctionRequest) returns (UpdateFunctionResponse); rpc ListFunctions (ListFunctionsRequest) returns (ListFunctionsResponse); rpc DeleteFunction (DeleteFunctionRequest) returns (DeleteFunctionResponse); diff --git a/services/proto/src/proto/teaclave_management_service.proto b/services/proto/src/proto/teaclave_management_service.proto index d79bd42e..7d9639ab 100644 --- a/services/proto/src/proto/teaclave_management_service.proto +++ b/services/proto/src/proto/teaclave_management_service.proto @@ -36,6 +36,7 @@ service TeaclaveManagement { rpc RegisterFunction (teaclave_frontend_service_proto.RegisterFunctionRequest) returns (teaclave_frontend_service_proto.RegisterFunctionResponse); rpc UpdateFunction (teaclave_frontend_service_proto.UpdateFunctionRequest) returns (teaclave_frontend_service_proto.UpdateFunctionResponse); rpc GetFunction (teaclave_frontend_service_proto.GetFunctionRequest) returns (teaclave_frontend_service_proto.GetFunctionResponse); + rpc GetFunctionUsageStats (teaclave_frontend_service_proto.GetFunctionUsageStatsRequest) returns (teaclave_frontend_service_proto.GetFunctionUsageStatsResponse); rpc DeleteFunction (teaclave_frontend_service_proto.DeleteFunctionRequest) returns (teaclave_frontend_service_proto.DeleteFunctionResponse); rpc DisableFunction (teaclave_frontend_service_proto.DisableFunctionRequest) returns (teaclave_frontend_service_proto.DisableFunctionResponse); rpc ListFunctions (teaclave_frontend_service_proto.ListFunctionsRequest) returns (teaclave_frontend_service_proto.ListFunctionsResponse); diff --git a/services/proto/src/teaclave_frontend_service.rs b/services/proto/src/teaclave_frontend_service.rs index 2e3c9150..ce3c26cc 100644 --- a/services/proto/src/teaclave_frontend_service.rs +++ b/services/proto/src/teaclave_frontend_service.rs @@ -288,6 +288,7 @@ pub struct RegisterFunctionRequest { pub inputs: Vec<FunctionInput>, pub outputs: Vec<FunctionOutput>, pub user_allowlist: Vec<String>, + pub usage_quota: Option<i32>, } #[derive(Default)] @@ -351,6 +352,11 @@ impl RegisterFunctionRequestBuilder { self } + pub fn usage_quota(mut self, usage_quota: Option<i32>) -> Self { + self.request.usage_quota = usage_quota; + self + } + pub fn build(self) -> RegisterFunctionRequest { self.request } @@ -369,6 +375,7 @@ impl From<RegisterFunctionRequest> for FunctionBuilder { .inputs(request.inputs) .outputs(request.outputs) .user_allowlist(request.user_allowlist) + .usage_quota(request.usage_quota) } } @@ -398,6 +405,7 @@ pub struct UpdateFunctionRequest { pub inputs: Vec<FunctionInput>, pub outputs: Vec<FunctionOutput>, pub user_allowlist: Vec<String>, + pub usage_quota: Option<i32>, } #[derive(Default)] @@ -466,6 +474,11 @@ impl UpdateFunctionRequestBuilder { self } + pub fn usage_quota(mut self, usage_quota: Option<i32>) -> Self { + self.request.usage_quota = usage_quota; + self + } + pub fn build(self) -> UpdateFunctionRequest { self.request } @@ -485,6 +498,7 @@ impl From<UpdateFunctionRequest> for FunctionBuilder { .inputs(request.inputs) .outputs(request.outputs) .user_allowlist(request.user_allowlist) + .usage_quota(request.usage_quota) } } @@ -528,6 +542,26 @@ pub struct GetFunctionResponse { pub user_allowlist: Vec<String>, } +#[into_request(TeaclaveManagementRequest::GetFunctionUsageStats)] +#[into_request(TeaclaveFrontendRequest::GetFunctionUsageStats)] +#[derive(Debug)] +pub struct GetFunctionUsageStatsRequest { + pub function_id: ExternalID, +} + +impl GetFunctionUsageStatsRequest { + pub fn new(function_id: ExternalID) -> Self { + Self { function_id } + } +} + +#[into_request(TeaclaveManagementResponse::GetFunctionUsageStats)] +#[derive(Debug)] +pub struct GetFunctionUsageStatsResponse { + pub function_quota: i32, + pub current_usage: i32, +} + #[into_request(TeaclaveManagementRequest::DeleteFunction)] #[into_request(TeaclaveFrontendRequest::DeleteFunction)] #[derive(Debug)] @@ -1123,6 +1157,8 @@ impl std::convert::TryFrom<proto::RegisterFunctionRequest> for RegisterFunctionR .map(FunctionArgument::try_from) .collect(); + let usage_quota = (proto.usage_quota >= 0).then_some(proto.usage_quota); + let ret = Self { name: proto.name, description: proto.description, @@ -1133,6 +1169,7 @@ impl std::convert::TryFrom<proto::RegisterFunctionRequest> for RegisterFunctionR inputs: inputs?, outputs: outputs?, user_allowlist: proto.user_allowlist, + usage_quota, }; Ok(ret) } @@ -1166,6 +1203,7 @@ impl From<RegisterFunctionRequest> for proto::RegisterFunctionRequest { inputs, outputs, user_allowlist: request.user_allowlist, + usage_quota: request.usage_quota.unwrap_or(-1), } } } @@ -1175,9 +1213,7 @@ impl std::convert::TryFrom<proto::RegisterFunctionResponse> for RegisterFunction fn try_from(proto: proto::RegisterFunctionResponse) -> Result<Self> { let function_id = proto.function_id.try_into()?; - let ret = Self { function_id }; - - Ok(ret) + Ok(Self { function_id }) } } @@ -1211,6 +1247,8 @@ impl std::convert::TryFrom<proto::UpdateFunctionRequest> for UpdateFunctionReque .map(FunctionArgument::try_from) .collect(); + let usage_quota = (proto.usage_quota >= 0).then_some(proto.usage_quota); + let ret = Self { function_id, name: proto.name, @@ -1222,6 +1260,7 @@ impl std::convert::TryFrom<proto::UpdateFunctionRequest> for UpdateFunctionReque inputs: inputs?, outputs: outputs?, user_allowlist: proto.user_allowlist, + usage_quota, }; Ok(ret) } @@ -1256,6 +1295,7 @@ impl From<UpdateFunctionRequest> for proto::UpdateFunctionRequest { inputs, outputs, user_allowlist: request.user_allowlist, + usage_quota: request.usage_quota.unwrap_or(-1), } } } @@ -1298,6 +1338,43 @@ impl From<GetFunctionRequest> for proto::GetFunctionRequest { } } +impl From<GetFunctionUsageStatsRequest> for proto::GetFunctionUsageStatsRequest { + fn from(request: GetFunctionUsageStatsRequest) -> Self { + Self { + function_id: request.function_id.to_string(), + } + } +} + +impl std::convert::TryFrom<proto::GetFunctionUsageStatsRequest> for GetFunctionUsageStatsRequest { + type Error = Error; + + fn try_from(proto: proto::GetFunctionUsageStatsRequest) -> Result<Self> { + let function_id: ExternalID = proto.function_id.try_into()?; + Ok(Self { function_id }) + } +} + +impl From<GetFunctionUsageStatsResponse> for proto::GetFunctionUsageStatsResponse { + fn from(request: GetFunctionUsageStatsResponse) -> Self { + Self { + function_quota: request.function_quota, + current_usage: request.current_usage, + } + } +} + +impl std::convert::TryFrom<proto::GetFunctionUsageStatsResponse> for GetFunctionUsageStatsResponse { + type Error = Error; + + fn try_from(proto: proto::GetFunctionUsageStatsResponse) -> Result<Self> { + Ok(Self { + function_quota: proto.function_quota, + current_usage: proto.current_usage, + }) + } +} + impl std::convert::TryFrom<proto::DeleteFunctionResponse> for DeleteFunctionResponse { type Error = Error; diff --git a/services/proto/src/teaclave_management_service.rs b/services/proto/src/teaclave_management_service.rs index 6f8cbcb0..e9612eaf 100644 --- a/services/proto/src/teaclave_management_service.rs +++ b/services/proto/src/teaclave_management_service.rs @@ -50,6 +50,10 @@ pub type UpdateFunctionRequest = crate::teaclave_frontend_service::UpdateFunctio pub type UpdateFunctionRequestBuilder = crate::teaclave_frontend_service::UpdateFunctionRequestBuilder; pub type UpdateFunctionResponse = crate::teaclave_frontend_service::UpdateFunctionResponse; +pub type GetFunctionUsageStatsRequest = + crate::teaclave_frontend_service::GetFunctionUsageStatsRequest; +pub type GetFunctionUsageStatsResponse = + crate::teaclave_frontend_service::GetFunctionUsageStatsResponse; pub type DeleteFunctionRequest = crate::teaclave_frontend_service::DeleteFunctionRequest; pub type DeleteFunctionResponse = crate::teaclave_frontend_service::DeleteFunctionResponse; pub type DisableFunctionRequest = crate::teaclave_frontend_service::DisableFunctionRequest; diff --git a/tests/functional/enclave/src/management_service.rs b/tests/functional/enclave/src/management_service.rs index 95249edd..29b164cb 100644 --- a/tests/functional/enclave/src/management_service.rs +++ b/tests/functional/enclave/src/management_service.rs @@ -135,6 +135,7 @@ fn test_register_function() { .arguments(vec![FunctionArgument::new("arg", "", true)]) .inputs(vec![function_input]) .outputs(vec![function_output]) + .usage_quota(Some(0)) .build(); let mut client = authorized_client("mock_user"); diff --git a/types/src/function.rs b/types/src/function.rs index 0180d8e8..13e0f331 100644 --- a/types/src/function.rs +++ b/types/src/function.rs @@ -87,6 +87,7 @@ pub struct Function { pub outputs: Vec<FunctionOutput>, pub owner: UserID, pub user_allowlist: Vec<String>, + pub usage_quota: Option<i32>, } #[derive(Default)] @@ -156,6 +157,16 @@ impl FunctionBuilder { self } + pub fn usage_quota(mut self, usage_quota: Option<i32>) -> Self { + let usage_quota = match usage_quota { + Some(quota) if quota < 0 => None, + _ => usage_quota, + }; + + self.function.usage_quota = usage_quota; + self + } + pub fn build(self) -> Function { self.function } @@ -191,3 +202,21 @@ impl FunctionArgument { } } } + +const FUNCION_USAGE_PREFIX: &str = "usage"; + +#[derive(Default, Debug, Deserialize, Serialize)] +pub struct FunctionUsage { + pub function_id: Uuid, + pub use_numbers: i32, +} + +impl Storable for FunctionUsage { + fn key_prefix() -> &'static str { + FUNCION_USAGE_PREFIX + } + + fn uuid(&self) -> Uuid { + self.function_id + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
