This is an automated email from the ASF dual-hosted git repository.

mssun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-teaclave.git


The following commit(s) were added to refs/heads/master by this push:
     new 131aa73  feat(management): add disable function API (#620)
131aa73 is described below

commit 131aa7361691973d4c85c5045c9310d2463efbe8
Author: Qinkun Bao <[email protected]>
AuthorDate: Wed Feb 9 16:39:57 2022 -0800

    feat(management): add disable function API (#620)
---
 cmake/scripts/test.sh                              |  1 +
 examples/python/test_disable_function.py           | 98 ++++++++++++++++++++++
 sdk/python/teaclave.py                             | 19 +++++
 services/frontend/enclave/src/service.rs           | 39 ++++++---
 services/management/enclave/src/service.rs         | 79 +++++++++++++++--
 .../src/proto/teaclave_frontend_service.proto      |  8 ++
 .../src/proto/teaclave_management_service.proto    |  1 +
 services/proto/src/teaclave_frontend_service.rs    | 50 +++++++++++
 services/proto/src/teaclave_management_service.rs  |  2 +
 tests/functional/enclave/src/management_service.rs | 23 +++++
 10 files changed, 299 insertions(+), 21 deletions(-)

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 @@ run_examples() {
   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 @@ class DeleteFunctionRequest(Request):
         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 @@ class FrontendService(TeaclaveService):
         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_common::UserCredential;
 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 @@ enum Endpoints {
     UpdateFunction,
     ListFunctions,
     DeleteFunction,
+    DisableFunction,
     CreateTask,
     GetTask,
     AssignData,
@@ -127,9 +129,10 @@ fn authorize(claims: &UserAuthClaims, request: Endpoints) 
-> bool {
     }
 
     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 @@ impl TeaclaveFrontend for TeaclaveFrontendService {
         )
     }
 
+    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 std::sync::{Arc, SgxMutex as Mutex};
 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 @@ impl TeaclaveManagement for TeaclaveManagementService {
         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 DeleteFunctionRequest {
 
 message DeleteFunctionResponse { }
 
+message DisableFunctionRequest {
+  string function_id = 1;
+}
+
+message DisableFunctionResponse { }
+
+
 message ListFunctionsRequest {
   string user_id = 1;
 }
@@ -254,6 +261,7 @@ service TeaclaveFrontend {
   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 @@ service TeaclaveManagement {
   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 @@ impl DeleteFunctionRequest {
 #[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 From<DeleteFunctionResponse> for 
proto::DeleteFunctionResponse {
     }
 }
 
+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 From<DeleteFunctionRequest> for 
proto::DeleteFunctionRequest {
     }
 }
 
+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 UpdateFunctionRequestBuilder =
 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 @@ fn test_delete_function() {
 }
 
 #[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);

---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to