feat(management): add disable function API (#620)
diff --git a/cmake/scripts/test.sh b/cmake/scripts/test.sh
index 3b18f33..e4e5740 100755
--- a/cmake/scripts/test.sh
+++ b/cmake/scripts/test.sh
@@ -250,6 +250,7 @@
python3 wasm_c_simple_add.py
python3 wasm_rust_psi.py
python3 wasm_tvm_mnist.py
+ python3 test_disable_function.py
popd
pushd ${TEACLAVE_PROJECT_ROOT}/examples/c
diff --git a/examples/python/test_disable_function.py b/examples/python/test_disable_function.py
new file mode 100644
index 0000000..5c51e6d
--- /dev/null
+++ b/examples/python/test_disable_function.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import sys
+
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
+
+
+class UserData:
+
+ def __init__(self, user_id, password):
+ self.user_id = user_id
+ self.password = password
+
+
+USER_DATA_0 = UserData("user0", "password")
+
+
+class ConfigClient:
+
+ def __init__(self, user_id, user_password):
+ self.user_id = user_id
+ self.user_password = user_password
+ with connect_authentication_service() as client:
+ print(f"[+] {self.user_id} login")
+ token = client.user_login(self.user_id, self.user_password)
+ self.client = connect_frontend_service()
+ metadata = {"id": self.user_id, "token": token}
+ self.client.metadata = metadata
+
+ def register_function(self, func_name):
+ client = self.client
+
+ print(f"[+] {self.user_id} registering function")
+
+ function_id = client.register_function(name=func_name,
+ description=func_name,
+ executor_type="builtin",
+ arguments=["num_user"],
+ inputs=[],
+ outputs=[])
+
+ return function_id
+
+ def list_function(self):
+ client = self.client
+
+ print(f"[+] {self.user_id} list function")
+ functions = client.list_functions(self.user_id)
+ return functions
+
+ def disable_function(self, function_id):
+ client = self.client
+
+ print(f"[+] {self.user_id} disable function")
+ client.disable_function(function_id)
+
+
+def main():
+ platform_admin = PlatformAdmin("admin", "teaclave")
+ try:
+ platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
+ except Exception:
+ pass
+
+ config_client = ConfigClient(USER_DATA_0.user_id, USER_DATA_0.password)
+ function_id1 = config_client.register_function("func_test1")
+ function_id2 = config_client.register_function("func_test1")
+ functions = config_client.list_function()
+ func_nums_before = len(functions['registered_functions'])
+ print(f"{func_nums_before} functions registered")
+ print(f"Disable {function_id2}")
+ config_client.disable_function(function_id2)
+ functions = config_client.list_function()
+ func_nums_after = len(functions['registered_functions'])
+ print(f"{func_nums_after} functions registered")
+ assert (func_nums_before == func_nums_after + 1)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/sdk/python/teaclave.py b/sdk/python/teaclave.py
index 88080f3..b7a59b5 100644
--- a/sdk/python/teaclave.py
+++ b/sdk/python/teaclave.py
@@ -579,6 +579,14 @@
self.function_id = function_id
+class DisableFunctionRequest(Request):
+
+ def __init__(self, metadata: Metadata, function_id: str):
+ self.request = "disable_function"
+ self.metadata = metadata
+ self.function_id = function_id
+
+
class GetFunctionRequest(Request):
def __init__(self, metadata: Metadata, function_id: str):
@@ -796,6 +804,17 @@
else:
raise TeaclaveException("Failed to delete function")
+ def disable_function(self, function_id: str):
+ self.check_metadata()
+ self.check_channel()
+ request = DisableFunctionRequest(self.metadata, function_id)
+ _write_message(self.channel, request)
+ response = _read_message(self.channel)
+ if response["result"] == "ok":
+ return response["content"]
+ else:
+ raise TeaclaveException("Failed to disable function")
+
def register_input_file(self, url: str, schema: str, key: List[int],
iv: List[int], cmac: List[int]):
self.check_metadata()
diff --git a/services/frontend/enclave/src/service.rs b/services/frontend/enclave/src/service.rs
index bd0d3bc..924f606 100644
--- a/services/frontend/enclave/src/service.rs
+++ b/services/frontend/enclave/src/service.rs
@@ -28,15 +28,16 @@
use teaclave_proto::teaclave_frontend_service::{
ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse,
CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, CreateTaskResponse,
- DeleteFunctionRequest, DeleteFunctionResponse, 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,
+ 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,
};
use teaclave_proto::teaclave_management_service::TeaclaveManagementClient;
use teaclave_rpc::endpoint::Endpoint;
@@ -108,6 +109,7 @@
UpdateFunction,
ListFunctions,
DeleteFunction,
+ DisableFunction,
CreateTask,
GetTask,
AssignData,
@@ -127,9 +129,10 @@
}
match request {
- Endpoints::RegisterFunction | Endpoints::UpdateFunction | Endpoints::DeleteFunction => {
- role.is_function_owner()
- }
+ Endpoints::RegisterFunction
+ | Endpoints::UpdateFunction
+ | Endpoints::DeleteFunction
+ | Endpoints::DisableFunction => role.is_function_owner(),
Endpoints::RegisterInputFile
| Endpoints::RegisterOutputFile
| Endpoints::UpdateInputFile
@@ -338,6 +341,18 @@
)
}
+ fn disable_function(
+ &self,
+ request: Request<DisableFunctionRequest>,
+ ) -> TeaclaveServiceResponseResult<DisableFunctionResponse> {
+ authentication_and_forward_to_management!(
+ self,
+ request,
+ disable_function,
+ Endpoints::DisableFunction
+ )
+ }
+
fn list_functions(
&self,
request: Request<ListFunctionsRequest>,
diff --git a/services/management/enclave/src/service.rs b/services/management/enclave/src/service.rs
index 55e730d..fd7adbf 100644
--- a/services/management/enclave/src/service.rs
+++ b/services/management/enclave/src/service.rs
@@ -24,15 +24,16 @@
use teaclave_proto::teaclave_frontend_service::{
ApproveTaskRequest, ApproveTaskResponse, AssignDataRequest, AssignDataResponse,
CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, CreateTaskResponse,
- DeleteFunctionRequest, DeleteFunctionResponse, 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,
+ 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,
};
use teaclave_proto::teaclave_management_service::TeaclaveManagement;
use teaclave_proto::teaclave_storage_service::{
@@ -394,6 +395,66 @@
Ok(response)
}
+ // access control: function.owner == user_id
+ // disable function
+ // 1. `List functions` does not show this function
+ // 2. `Create new task` with the function id fails
+ fn disable_function(
+ &self,
+ request: Request<DisableFunctionRequest>,
+ ) -> TeaclaveServiceResponseResult<DisableFunctionResponse> {
+ let user_id = self.get_request_user_id(request.metadata())?;
+ let role = self.get_request_role(request.metadata())?;
+
+ let mut function: Function = self
+ .read_from_db(&request.message.function_id)
+ .map_err(|_| TeaclaveManagementServiceError::PermissionDenied)?;
+
+ if role != UserRole::PlatformAdmin {
+ ensure!(
+ function.owner == user_id,
+ TeaclaveManagementServiceError::PermissionDenied
+ );
+ }
+ let func_id = function.external_id().to_string();
+
+ // Updated function owner
+ let mut u = User::default();
+ u.id = function.owner.clone();
+ let external_id = u.external_id();
+ let user: Result<User> = self.read_from_db(&external_id);
+ if let Ok(mut us) = user {
+ us.allowed_functions.retain(|f| !f.eq(&func_id));
+ us.registered_functions.retain(|f| !f.eq(&func_id));
+ self.write_to_db(&us)
+ .map_err(|_| TeaclaveManagementServiceError::StorageError)?;
+ } else {
+ log::warn!("Invalid user id from functions");
+ }
+
+ // Update allowed function list for users
+ for user_id in &function.user_allowlist {
+ let mut u = User::default();
+ u.id = user_id.into();
+ let external_id = u.external_id();
+ let user: Result<User> = self.read_from_db(&external_id);
+ if let Ok(mut us) = user {
+ us.allowed_functions.retain(|f| !f.eq(&func_id));
+ us.registered_functions.retain(|f| !f.eq(&func_id));
+ self.write_to_db(&us)
+ .map_err(|_| TeaclaveManagementServiceError::StorageError)?;
+ } else {
+ log::warn!("Invalid user id from functions");
+ }
+ }
+
+ function.user_allowlist.clear();
+ self.write_to_db(&function)
+ .map_err(|_| TeaclaveManagementServiceError::StorageError)?;
+ let response = DisableFunctionResponse {};
+ Ok(response)
+ }
+
// access contro: user_id = request.user_id
fn list_functions(
&self,
diff --git a/services/proto/src/proto/teaclave_frontend_service.proto b/services/proto/src/proto/teaclave_frontend_service.proto
index 71d3628..31d8fd0 100644
--- a/services/proto/src/proto/teaclave_frontend_service.proto
+++ b/services/proto/src/proto/teaclave_frontend_service.proto
@@ -168,6 +168,13 @@
message DeleteFunctionResponse { }
+message DisableFunctionRequest {
+ string function_id = 1;
+}
+
+message DisableFunctionResponse { }
+
+
message ListFunctionsRequest {
string user_id = 1;
}
@@ -254,6 +261,7 @@
rpc UpdateFunction (UpdateFunctionRequest) returns (UpdateFunctionResponse);
rpc ListFunctions (ListFunctionsRequest) returns (ListFunctionsResponse);
rpc DeleteFunction (DeleteFunctionRequest) returns (DeleteFunctionResponse);
+ rpc DisableFunction (DisableFunctionRequest) returns (DisableFunctionResponse);
rpc CreateTask (CreateTaskRequest) returns (CreateTaskResponse);
rpc GetTask (GetTaskRequest) returns (GetTaskResponse);
rpc AssignData (AssignDataRequest) returns (AssignDataResponse);
diff --git a/services/proto/src/proto/teaclave_management_service.proto b/services/proto/src/proto/teaclave_management_service.proto
index 82c50a4..d79bd42 100644
--- a/services/proto/src/proto/teaclave_management_service.proto
+++ b/services/proto/src/proto/teaclave_management_service.proto
@@ -37,6 +37,7 @@
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 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);
rpc CreateTask (teaclave_frontend_service_proto.CreateTaskRequest) returns (teaclave_frontend_service_proto.CreateTaskResponse);
rpc GetTask (teaclave_frontend_service_proto.GetTaskRequest) returns (teaclave_frontend_service_proto.GetTaskResponse);
diff --git a/services/proto/src/teaclave_frontend_service.rs b/services/proto/src/teaclave_frontend_service.rs
index 26d3198..b353ad0 100644
--- a/services/proto/src/teaclave_frontend_service.rs
+++ b/services/proto/src/teaclave_frontend_service.rs
@@ -552,6 +552,23 @@
#[derive(Debug)]
pub struct DeleteFunctionResponse {}
+#[into_request(TeaclaveManagementRequest::DisableFunction)]
+#[into_request(TeaclaveFrontendRequest::DisableFunction)]
+#[derive(Debug)]
+pub struct DisableFunctionRequest {
+ pub function_id: ExternalID,
+}
+
+impl DisableFunctionRequest {
+ pub fn new(function_id: ExternalID) -> Self {
+ Self { function_id }
+ }
+}
+
+#[into_request(TeaclaveManagementResponse::DisableFunction)]
+#[derive(Debug)]
+pub struct DisableFunctionResponse {}
+
#[into_request(TeaclaveManagementRequest::CreateTask)]
#[into_request(TeaclaveFrontendRequest::CreateTask)]
#[derive(Default)]
@@ -1282,6 +1299,20 @@
}
}
+impl std::convert::TryFrom<proto::DisableFunctionResponse> for DisableFunctionResponse {
+ type Error = Error;
+
+ fn try_from(_proto: proto::DisableFunctionResponse) -> Result<Self> {
+ Ok(DisableFunctionResponse {})
+ }
+}
+
+impl From<DisableFunctionResponse> for proto::DisableFunctionResponse {
+ fn from(_response: DisableFunctionResponse) -> Self {
+ Self {}
+ }
+}
+
impl std::convert::TryFrom<proto::ListFunctionsRequest> for ListFunctionsRequest {
type Error = Error;
@@ -1342,6 +1373,25 @@
}
}
+impl std::convert::TryFrom<proto::DisableFunctionRequest> for DisableFunctionRequest {
+ type Error = Error;
+
+ fn try_from(proto: proto::DisableFunctionRequest) -> Result<Self> {
+ let function_id = proto.function_id.try_into()?;
+ let ret = Self { function_id };
+
+ Ok(ret)
+ }
+}
+
+impl From<DisableFunctionRequest> for proto::DisableFunctionRequest {
+ fn from(request: DisableFunctionRequest) -> Self {
+ Self {
+ function_id: request.function_id.to_string(),
+ }
+ }
+}
+
impl std::convert::TryFrom<proto::GetFunctionResponse> for GetFunctionResponse {
type Error = Error;
diff --git a/services/proto/src/teaclave_management_service.rs b/services/proto/src/teaclave_management_service.rs
index c868d59..6f8cbcb 100644
--- a/services/proto/src/teaclave_management_service.rs
+++ b/services/proto/src/teaclave_management_service.rs
@@ -52,6 +52,8 @@
pub type UpdateFunctionResponse = crate::teaclave_frontend_service::UpdateFunctionResponse;
pub type DeleteFunctionRequest = crate::teaclave_frontend_service::DeleteFunctionRequest;
pub type DeleteFunctionResponse = crate::teaclave_frontend_service::DeleteFunctionResponse;
+pub type DisableFunctionRequest = crate::teaclave_frontend_service::DisableFunctionRequest;
+pub type DisableFunctionResponse = crate::teaclave_frontend_service::DisableFunctionResponse;
pub type GetFunctionRequest = crate::teaclave_frontend_service::GetFunctionRequest;
pub type GetFunctionResponse = crate::teaclave_frontend_service::GetFunctionResponse;
pub type ListFunctionsRequest = crate::teaclave_frontend_service::ListFunctionsRequest;
diff --git a/tests/functional/enclave/src/management_service.rs b/tests/functional/enclave/src/management_service.rs
index 765192b..ad016a0 100644
--- a/tests/functional/enclave/src/management_service.rs
+++ b/tests/functional/enclave/src/management_service.rs
@@ -189,6 +189,29 @@
}
#[test_case]
+fn test_disable_function() {
+ let function_input = FunctionInput::new("input", "input_desc", false);
+ let function_output = FunctionOutput::new("output", "output_desc", false);
+ let request = RegisterFunctionRequestBuilder::new()
+ .name("mock_function")
+ .executor_type(ExecutorType::Python)
+ .payload(b"def entrypoint:\n\treturn".to_vec())
+ .public(true)
+ .arguments(vec!["arg"])
+ .inputs(vec![function_input])
+ .outputs(vec![function_output])
+ .build();
+
+ let mut client = authorized_client("mock_user");
+ let response = client.register_function(request);
+ let function_id = response.unwrap().function_id;
+
+ let request = DisableFunctionRequest::new(function_id);
+ let response = client.disable_function(request);
+ assert!(response.is_ok());
+}
+
+#[test_case]
fn test_update_function() {
let function_input = FunctionInput::new("input", "input_desc", false);
let function_output = FunctionOutput::new("output", "output_desc", false);
