You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@teaclave.apache.org by hs...@apache.org on 2023/03/17 13:27:12 UTC
[incubator-teaclave] 04/04: Add usage statistics and quota for function
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 <su...@baidu.com>
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: commits-unsubscribe@teaclave.apache.org
For additional commands, e-mail: commits-help@teaclave.apache.org