You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@teaclave.apache.org by ms...@apache.org on 2021/12/15 00:20:00 UTC
[incubator-teaclave] branch master updated: Refactor python client SDK (#591)
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 b3f4e8e Refactor python client SDK (#591)
b3f4e8e is described below
commit b3f4e8e2d3fcc0152395c666cc94ab2907e8c111
Author: Mingshen Sun <bo...@mssun.me>
AuthorDate: Tue Dec 14 16:19:55 2021 -0800
Refactor python client SDK (#591)
---
examples/python/builtin_echo.py | 64 +++--
examples/python/builtin_face_detection.py | 21 +-
examples/python/builtin_gbdt_train.py | 22 +-
examples/python/builtin_online_decrypt.py | 20 +-
examples/python/builtin_ordered_set_intersect.py | 28 +--
examples/python/builtin_password_check.py | 27 +--
.../python/builtin_private_join_and_compute.py | 43 ++--
examples/python/builtin_rsa_sign.py | 30 +--
examples/python/mesapy_echo.py | 21 +-
examples/python/mesapy_logistic_reg.py | 20 +-
examples/python/utils.py | 20 +-
examples/python/wasm_c_simple_add.py | 22 +-
examples/python/wasm_rust_psi.py | 26 +-
examples/python/wasm_tvm_mnist.py | 28 +--
sdk/python/teaclave.py | 270 ++++++++++++---------
15 files changed, 316 insertions(+), 346 deletions(-)
diff --git a/examples/python/builtin_echo.py b/examples/python/builtin_echo.py
index f78c822..21de3a4 100644
--- a/examples/python/builtin_echo.py
+++ b/examples/python/builtin_echo.py
@@ -19,11 +19,7 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service
class BuiltinEchoExample:
@@ -32,37 +28,33 @@ class BuiltinEchoExample:
self.user_password = user_password
def echo(self, message="Hello, Teaclave!"):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
-
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- metadata = {"id": self.user_id, "token": token}
- client.metadata = metadata
-
- print("[+] registering function")
- function_id = client.register_function(
- name="builtin-echo",
- description="Native Echo Function",
- executor_type="builtin",
- arguments=["message"])
-
- print("[+] creating task")
- task_id = client.create_task(function_id=function_id,
- function_arguments={"message": message},
- executor="builtin")
-
- print("[+] invoking task")
- client.invoke_task(task_id)
-
- print("[+] getting result")
- result = client.get_task_result(task_id)
- print("[+] done")
+ with connect_authentication_service() as client:
+ print("[+] login")
+ token = client.user_login(self.user_id, self.user_password)
+
+ with connect_frontend_service() as client:
+ metadata = {"id": self.user_id, "token": token}
+ client.metadata = metadata
+
+ print("[+] registering function")
+ function_id = client.register_function(
+ name="builtin-echo",
+ description="Native Echo Function",
+ executor_type="builtin",
+ arguments=["message"])
+
+ print("[+] creating task")
+ task_id = client.create_task(
+ function_id=function_id,
+ function_arguments={"message": message},
+ executor="builtin")
+
+ print("[+] invoking task")
+ client.invoke_task(task_id)
+
+ print("[+] getting result")
+ result = client.get_task_result(task_id)
+ print("[+] done")
return bytes(result)
diff --git a/examples/python/builtin_face_detection.py b/examples/python/builtin_face_detection.py
index 3d450f4..cee43ac 100644
--- a/examples/python/builtin_face_detection.py
+++ b/examples/python/builtin_face_detection.py
@@ -24,12 +24,7 @@ import json
from PIL import Image, ImageDraw
import requests
-from teaclave import (AuthenticationService, FrontendService, OwnerList,
- AuthenticationClient, FrontendClient, DataMap,
- FunctionInput)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service
class BuiltinFaceDetectionExample:
@@ -38,16 +33,11 @@ class BuiltinFaceDetectionExample:
self.user_password = user_password
def detect_face(self, image):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print("[+] login")
+ token = client.user_login(self.user_id, self.user_password)
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ client = connect_frontend_service()
metadata = {"id": self.user_id, "token": token}
client.metadata = metadata
@@ -82,6 +72,7 @@ class BuiltinFaceDetectionExample:
print("[+] getting result")
result = client.get_task_result(task_id)
print("[+] done")
+ client.close()
return bytes(result)
diff --git a/examples/python/builtin_gbdt_train.py b/examples/python/builtin_gbdt_train.py
index 9705446..2efc9c7 100644
--- a/examples/python/builtin_gbdt_train.py
+++ b/examples/python/builtin_gbdt_train.py
@@ -19,12 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service
class BuiltinGbdtExample:
@@ -33,16 +29,11 @@ class BuiltinGbdtExample:
self.user_password = user_password
def gbdt(self):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print("[+] login")
+ token = client.user_login(self.user_id, self.user_password)
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ client = connect_frontend_service()
metadata = {"id": self.user_id, "token": token}
client.metadata = metadata
@@ -115,6 +106,7 @@ class BuiltinGbdtExample:
print("[+] getting result")
result = client.get_task_result(task_id)
print("[+] done")
+ client.close()
return bytes(result)
diff --git a/examples/python/builtin_online_decrypt.py b/examples/python/builtin_online_decrypt.py
index 3ba4e34..5749bd9 100644
--- a/examples/python/builtin_online_decrypt.py
+++ b/examples/python/builtin_online_decrypt.py
@@ -20,11 +20,7 @@
import sys
import base64
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service
class BuiltinOnlineDecryptExample:
@@ -33,16 +29,11 @@ class BuiltinOnlineDecryptExample:
self.user_password = user_password
def decrypt(self, key, nonce, encrypted_data, algorithm):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print("[+] login")
+ token = client.user_login(self.user_id, self.user_password)
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ client = connect_frontend_service()
metadata = {"id": self.user_id, "token": token}
client.metadata = metadata
@@ -71,6 +62,7 @@ class BuiltinOnlineDecryptExample:
print("[+] getting result")
result = client.get_task_result(task_id)
print("[+] done")
+ client.close()
return bytes(result)
diff --git a/examples/python/builtin_ordered_set_intersect.py b/examples/python/builtin_ordered_set_intersect.py
index 6a314ad..24a3b9c 100644
--- a/examples/python/builtin_ordered_set_intersect.py
+++ b/examples/python/builtin_ordered_set_intersect.py
@@ -19,12 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
# In the example, user 0 creates the task and user 0, 1, upload their private data.
# Then user 0 invokes the task and user 0, 1 get the result.
@@ -76,14 +72,10 @@ class Client:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
@@ -165,8 +157,12 @@ class Client:
def main():
# Register user
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
- platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ try:
+ platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
+ platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ except Exception:
+ pass
+
user0 = Client(USER_DATA_0.user_id, USER_DATA_0.password)
user1 = Client(USER_DATA_1.user_id, USER_DATA_1.password)
diff --git a/examples/python/builtin_password_check.py b/examples/python/builtin_password_check.py
index 0b921d3..aa3ba3e 100644
--- a/examples/python/builtin_password_check.py
+++ b/examples/python/builtin_password_check.py
@@ -19,12 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
# In the example, user 0 creates the task and user 0, 1, upload their private data.
# Then user 0 invokes the task and user 0, 1 get the result.
@@ -82,14 +78,10 @@ class Client:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
@@ -155,8 +147,11 @@ class Client:
def main():
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
- platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ try:
+ platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
+ platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ except Exception:
+ pass
user0 = Client(USER_DATA_0.user_id, USER_DATA_0.password)
user1 = Client(USER_DATA_1.user_id, USER_DATA_1.password)
diff --git a/examples/python/builtin_private_join_and_compute.py b/examples/python/builtin_private_join_and_compute.py
index 37c9f09..4cf5a0e 100644
--- a/examples/python/builtin_private_join_and_compute.py
+++ b/examples/python/builtin_private_join_and_compute.py
@@ -19,12 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
# In the example, user 3 creates the task and user 0, 1, 2 upload their private data.
# Then user 3 invokes the task and user 0, 1, 2 get the result.
@@ -83,14 +79,10 @@ class ConfigClient:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
@@ -151,14 +143,10 @@ class DataClient:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
@@ -197,10 +185,13 @@ class DataClient:
def main():
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
- platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
- platform_admin.register_user(USER_DATA_2.user_id, USER_DATA_2.password)
- platform_admin.register_user(USER_DATA_3.user_id, USER_DATA_3.password)
+ try:
+ platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
+ platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ platform_admin.register_user(USER_DATA_2.user_id, USER_DATA_2.password)
+ platform_admin.register_user(USER_DATA_3.user_id, USER_DATA_3.password)
+ except Exception:
+ pass
## USER 3 creates the task
config_client = ConfigClient(USER_DATA_3.user_id, USER_DATA_3.password)
diff --git a/examples/python/builtin_rsa_sign.py b/examples/python/builtin_rsa_sign.py
index 80234fd..0724756 100644
--- a/examples/python/builtin_rsa_sign.py
+++ b/examples/python/builtin_rsa_sign.py
@@ -19,24 +19,15 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
def get_client(user_id, user_password):
- auth_client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
-
- print("[+] login")
- token = auth_client.user_login(user_id, user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print(f"[+] login")
+ token = client.user_login(user_id, user_password)
+ client = connect_frontend_service()
metadata = {"id": user_id, "token": token}
client.metadata = metadata
return client
@@ -90,8 +81,11 @@ def create_task(client, function_id, input_file_user):
def rsa_task():
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user("rsa_sign_user1", "password1")
- platform_admin.register_user("rsa_sign_user2", "password2")
+ try:
+ platform_admin.register_user("rsa_sign_user1", "password1")
+ platform_admin.register_user("rsa_sign_user2", "password2")
+ except Exception:
+ pass
client1 = get_client("rsa_sign_user1", "password1")
client2 = get_client("rsa_sign_user2", "password2")
@@ -123,6 +117,8 @@ def rsa_task():
print("[+] getting result")
result = client2.get_task_result(task_id)
print("[+] done")
+ client1.close()
+ client2.close()
return bytes(result)
diff --git a/examples/python/mesapy_echo.py b/examples/python/mesapy_echo.py
index ddc5807..f8b1750 100644
--- a/examples/python/mesapy_echo.py
+++ b/examples/python/mesapy_echo.py
@@ -19,11 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
class MesaPyEchoExample:
@@ -34,16 +31,11 @@ class MesaPyEchoExample:
def echo(self,
payload_file="mesapy_echo_payload.py",
message="Hello, Teaclave!"):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print(f"[+] {self.user_id} login")
+ token = client.user_login(self.user_id, self.user_password)
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ client = connect_frontend_service()
metadata = {"id": self.user_id, "token": token}
client.metadata = metadata
@@ -69,6 +61,7 @@ class MesaPyEchoExample:
print("[+] getting result")
result = client.get_task_result(task_id)
print("[+] done")
+ client.close()
return bytes(result)
diff --git a/examples/python/mesapy_logistic_reg.py b/examples/python/mesapy_logistic_reg.py
index 93ab5bd..116335d 100644
--- a/examples/python/mesapy_logistic_reg.py
+++ b/examples/python/mesapy_logistic_reg.py
@@ -23,12 +23,8 @@ An example about Logistic Regression in MesaPy.
import sys
import binascii
from typing import List
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient, FunctionInput,
- FunctionOutput, OwnerList, DataMap)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
from enum import Enum
@@ -90,14 +86,10 @@ class ConfigClient:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
diff --git a/examples/python/utils.py b/examples/python/utils.py
index 292480b..90c6645 100644
--- a/examples/python/utils.py
+++ b/examples/python/utils.py
@@ -16,6 +16,7 @@
# under the License.
import os
+from teaclave import (AuthenticationService, FrontendService)
HOSTNAME = 'localhost'
AUTHENTICATION_SERVICE_ADDRESS = (HOSTNAME, 7776)
@@ -38,16 +39,25 @@ else:
AS_ROOT_CA_CERT_PATH = "../../keys/" + AS_ROOT_CERT_FILENAME
ENCLAVE_INFO_PATH = "../../release/examples/enclave_info.toml"
-from teaclave import (AuthenticationService, AuthenticationClient)
-
class PlatformAdmin:
def __init__(self, user_id: str, user_password: str):
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ self.client = AuthenticationService(AUTHENTICATION_SERVICE_ADDRESS,
+ AS_ROOT_CA_CERT_PATH,
+ ENCLAVE_INFO_PATH).connect()
token = self.client.user_login(user_id, user_password)
self.client.metadata = {"id": user_id, "token": token}
def register_user(self, user_id: str, user_password: str):
self.client.user_register(user_id, user_password, "PlatformAdmin", "")
+
+
+def connect_authentication_service():
+ return AuthenticationService(AUTHENTICATION_SERVICE_ADDRESS,
+ AS_ROOT_CA_CERT_PATH,
+ ENCLAVE_INFO_PATH).connect()
+
+
+def connect_frontend_service():
+ return FrontendService(FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
+ ENCLAVE_INFO_PATH).connect()
diff --git a/examples/python/wasm_c_simple_add.py b/examples/python/wasm_c_simple_add.py
index 881bde8..b5bc0b9 100644
--- a/examples/python/wasm_c_simple_add.py
+++ b/examples/python/wasm_c_simple_add.py
@@ -19,11 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService,
- AuthenticationClient, FrontendClient)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
class WASMAddExample:
@@ -35,16 +32,10 @@ class WASMAddExample:
payload_file="wasm_c_simple_add_payload/simple_add.wasm",
adder1="3",
adder2="4"):
- client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
-
- print("[+] login")
- token = client.user_login(self.user_id, self.user_password)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ with connect_authentication_service() as client:
+ print(f"[+] {self.user_id} login")
+ token = client.user_login(self.user_id, self.user_password)
+ client = connect_frontend_service()
metadata = {"id": self.user_id, "token": token}
client.metadata = metadata
@@ -73,6 +64,7 @@ class WASMAddExample:
print("[+] getting result")
result = client.get_task_result(task_id)
print("[+] done")
+ client.close()
return bytes(result)
diff --git a/examples/python/wasm_rust_psi.py b/examples/python/wasm_rust_psi.py
index 29be27f..c1742bf 100644
--- a/examples/python/wasm_rust_psi.py
+++ b/examples/python/wasm_rust_psi.py
@@ -19,11 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService, FunctionInput,
- FunctionOutput, OwnerList)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
class UserData:
@@ -78,14 +75,10 @@ class Client:
def __init__(self, user_id, user_password):
self.user_id = user_id
self.user_password = user_password
- self.client = AuthenticationService(
- AUTHENTICATION_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
- print(f"[+] {self.user_id} login")
- token = self.client.user_login(self.user_id, self.user_password)
- self.client = FrontendService(
- FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ 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
@@ -185,8 +178,11 @@ class Client:
def main():
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
- platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ try:
+ platform_admin.register_user(USER_DATA_0.user_id, USER_DATA_0.password)
+ platform_admin.register_user(USER_DATA_1.user_id, USER_DATA_1.password)
+ except Exception:
+ pass
user0 = Client(USER_DATA_0.user_id, USER_DATA_0.password)
user1 = Client(USER_DATA_1.user_id, USER_DATA_1.password)
diff --git a/examples/python/wasm_tvm_mnist.py b/examples/python/wasm_tvm_mnist.py
index 4343f12..90d1d09 100644
--- a/examples/python/wasm_tvm_mnist.py
+++ b/examples/python/wasm_tvm_mnist.py
@@ -19,11 +19,8 @@
import sys
-from teaclave import (AuthenticationService, FrontendService, FunctionInput,
- DataMap, OwnerList)
-from utils import (AUTHENTICATION_SERVICE_ADDRESS, FRONTEND_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH, ENCLAVE_INFO_PATH, USER_ID,
- USER_PASSWORD, PlatformAdmin)
+from teaclave import FunctionInput, FunctionOutput, OwnerList, DataMap
+from utils import USER_ID, USER_PASSWORD, connect_authentication_service, connect_frontend_service, PlatformAdmin
# If you're using `docker-compose` to start the Teaclave server containers,
# please change `localhost` to `teaclave-file-service`
@@ -42,17 +39,15 @@ PAYLOAD_FILE = "wasm_tvm_mnist_payload/target/wasm32-unknown-unknown/release/mni
def main():
platform_admin = PlatformAdmin("admin", "teaclave")
- platform_admin.register_user(USER_ID, USER_PASSWORD)
-
- client = AuthenticationService(AUTHENTICATION_SERVICE_ADDRESS,
- AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
-
- print(f"[+] {USER_ID} login")
- token = client.user_login(USER_ID, USER_PASSWORD)
-
- client = FrontendService(FRONTEND_SERVICE_ADDRESS, AS_ROOT_CA_CERT_PATH,
- ENCLAVE_INFO_PATH).connect().get_client()
+ try:
+ platform_admin.register_user(USER_ID, USER_PASSWORD)
+ except Exception:
+ pass
+
+ with connect_authentication_service() as client:
+ print(f"[+] login")
+ token = client.user_login(USER_ID, USER_PASSWORD)
+ client = connect_frontend_service()
metadata = {"id": USER_ID, "token": token}
client.metadata = metadata
@@ -99,6 +94,7 @@ def main():
print(f"[+] {USER_ID} getting task result")
result = client.get_task_result(task_id)
+ client.close()
print("[+] User result: " + bytes(result).decode("utf-8"))
diff --git a/sdk/python/teaclave.py b/sdk/python/teaclave.py
index c6b5dc6..191b207 100644
--- a/sdk/python/teaclave.py
+++ b/sdk/python/teaclave.py
@@ -42,9 +42,8 @@ from OpenSSL.crypto import X509Store, X509StoreContext
from OpenSSL import crypto
__all__ = [
- 'FrontendClient', 'FrontendService', 'AuthenticationClient',
- 'AuthenticationService', 'FunctionInput', 'FunctionOutput', 'OwnerList',
- 'DataMap'
+ 'FrontendService', 'AuthenticationService', 'FunctionInput',
+ 'FunctionOutput', 'OwnerList', 'DataMap'
]
Metadata = Dict[str, str]
@@ -54,6 +53,55 @@ class Request:
pass
+class TeaclaveException(Exception):
+ pass
+
+
+class TeaclaveService:
+ channel = None
+ metadata = None
+
+ def __init__(self, name: str, address: Tuple[str, int],
+ as_root_ca_cert_path: str, enclave_info_path: str):
+ self._context = ssl._create_unverified_context()
+ self._name = name
+ self._address = address
+ self._as_root_ca_cert_path = as_root_ca_cert_path
+ self._enclave_info_path = enclave_info_path
+ self._closed = False
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ if not self._closed:
+ self.close()
+
+ def close(self):
+ self._closed = True
+ if self.channel: self.channel.close()
+
+ def check_channel(self):
+ if not self.channel: raise TeaclaveException("Channel is None")
+
+ def check_metadata(self):
+ if not self.metadata: raise TeaclaveException("Metadata is None")
+
+ def connect(self):
+ """Establish trusted connection and verify remote attestation report.
+ """
+ sock = socket.create_connection(self._address)
+ channel = self._context.wrap_socket(sock,
+ server_hostname=self._address[0])
+ cert = channel.getpeercert(binary_form=True)
+ if not cert: raise TeaclaveException("Peer cert is None")
+ _verify_report(self._as_root_ca_cert_path, self._enclave_info_path,
+ cert, self._name)
+ self.channel = channel
+
+ return self
+
+
class FunctionInput:
"""Function input for registering.
@@ -134,7 +182,7 @@ class UserLoginRequest(Request):
self.password = user_password
-class AuthenticationService:
+class AuthenticationService(TeaclaveService):
"""
Establish trusted channel with the authentication service and provide
clients to send request through RPC.
@@ -146,51 +194,10 @@ class AuthenticationService:
enclave_info_path: Path of enclave info to verify the remote service in
the attestation report.
"""
- _context = ssl._create_unverified_context()
- _channel = None
-
def __init__(self, address: Tuple[str, int], as_root_ca_cert_path: str,
enclave_info_path: str):
- self.address = address
- self.as_root_ca_cert_path = as_root_ca_cert_path
- self.enclave_info_path = enclave_info_path
-
- def connect(self):
- """Establish trusted connection and verify remote attestation report.
-
- Returns:
- AuthenticationService: The original object which can be chained
- with other methods.
- """
- sock = socket.create_connection(self.address)
- channel = self._context.wrap_socket(sock,
- server_hostname=self.address[0])
- cert = channel.getpeercert(binary_form=True)
- _verify_report(self.as_root_ca_cert_path, self.enclave_info_path, cert,
- "authentication")
-
- self._channel = channel
-
- return self
-
- def get_client(self):
- """Get a client of authentication service to send RPC requests.
-
- Returns:
- AuthenticationClient: Used for send/receive RPC requests.
- """
- return AuthenticationClient(self._channel)
-
-
-class AuthenticationClient:
- """Client to communicate with the authentication service.
-
- Args:
- channel: Trusted TLS socket (verified with remote attestation).
- """
- def __init__(self, channel: ssl.SSLSocket, metadata: Metadata = None):
- self.channel = channel
- self.metadata = metadata
+ super().__init__("authentication", address, as_root_ca_cert_path,
+ enclave_info_path)
def user_register(self, user_id: str, user_password: str, role: str,
attribute: str):
@@ -202,10 +209,16 @@ class AuthenticationClient:
role: Role of user.
attribute: Attribute related to the role.
"""
+ self.check_channel()
+ self.check_metadata()
request = UserRegisterRequest(self.metadata, user_id, user_password,
role, attribute)
_write_message(self.channel, request)
- _ = _read_message(self.channel)
+ response = _read_message(self.channel)
+ if response["result"] == "ok":
+ pass
+ else:
+ raise TeaclaveException("Failed to register user")
def user_login(self, user_id: str, user_password: str) -> str:
"""Login and get a session token.
@@ -217,56 +230,14 @@ class AuthenticationClient:
Returns:
str: User login token.
"""
+ self.check_channel()
request = UserLoginRequest(user_id, user_password)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["token"]
-
-
-class FrontendService:
- """Establish trusted channel with the frontend service and provide
- clients to send request through RPC.
-
- Args:
- address: The address of the remote services in tuple.
- as_root_ca_cert_path: Root CA certification of the attestation services
- to verify the attestation report.
- enclave_info_path: Path of enclave info to verify the remote service in
- the attestation report.
- """
- _context = ssl._create_unverified_context()
- _channel = None
-
- def __init__(self, address: Tuple[str, int], as_root_ca_cert_path: str,
- enclave_info_path: str):
- self.address = address
- self.as_root_ca_cert_path = as_root_ca_cert_path
- self.enclave_info_path = enclave_info_path
-
- def connect(self):
- """Establish trusted connection and verify remote attestation report.
-
- Returns:
- FrontendService: The original object which can be chained
- with other methods.
- """
- sock = socket.create_connection(self.address)
- channel = self._context.wrap_socket(sock,
- server_hostname=self.address[0])
- cert = channel.getpeercert(binary_form=True)
- _verify_report(self.as_root_ca_cert_path, self.enclave_info_path, cert,
- "frontend")
-
- self._channel = channel
- return self
-
- def get_client(self):
- """Get a client of frontend service to send RPC requests.
-
- Returns:
- FrontendClient: Used for send/receive RPC requests.
- """
- return FrontendClient(self._channel)
+ if response["result"] == "ok":
+ return response["content"]["token"]
+ else:
+ raise TeaclaveException("Failed to login user")
class RegisterFunctionRequest(Request):
@@ -407,10 +378,21 @@ class GetTaskRequest(Request):
self.task_id = task_id
-class FrontendClient:
- def __init__(self, channel: ssl.SSLSocket, metadata: Metadata = None):
- self.channel = channel
- self.metadata = metadata
+class FrontendService(TeaclaveService):
+ """Establish trusted channel with the frontend service and provide
+ clients to send request through RPC.
+
+ Args:
+ address: The address of the remote services in tuple.
+ as_root_ca_cert_path: Root CA certification of the attestation services
+ to verify the attestation report.
+ enclave_info_path: Path of enclave info to verify the remote service in
+ the attestation report.
+ """
+ def __init__(self, address: Tuple[str, int], as_root_ca_cert_path: str,
+ enclave_info_path: str):
+ super().__init__("frontend", address, as_root_ca_cert_path,
+ enclave_info_path)
def register_function(
self,
@@ -424,13 +406,18 @@ class FrontendClient:
outputs: List[FunctionOutput] = [],
user_allowlist: List[str] = [],
):
+ self.check_metadata()
+ self.check_channel()
request = RegisterFunctionRequest(self.metadata, name, description,
executor_type, public, payload,
arguments, inputs, outputs,
user_allowlist)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["function_id"]
+ if response["result"] == "ok":
+ return response["content"]["function_id"]
+ else:
+ raise TeaclaveException("Failed to register function")
def update_function(
self,
@@ -445,47 +432,77 @@ class FrontendClient:
outputs: List[FunctionOutput] = [],
user_allowlist: List[str] = [],
):
+ self.check_metadata()
+ self.check_channel()
request = UpdateFunctionRequest(self.metadata, function_id, name,
description, executor_type, public,
payload, arguments, inputs, outputs,
user_allowlist)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["function_id"]
+ if response["result"] == "ok":
+ return response["content"]["function_id"]
+ else:
+ raise TeaclaveException("Failed to update function")
def list_functions(self, user_id: str):
+ self.check_metadata()
+ self.check_channel()
request = ListFunctionsRequest(self.metadata, user_id)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["function_ids"]
+ if response["result"] == "ok":
+ return response["content"]["function_ids"]
+ else:
+ raise TeaclaveException("Failed to list functions")
def get_function(self, function_id: str):
+ self.check_metadata()
+ self.check_channel()
request = GetFunctionRequest(self.metadata, function_id)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]
+ if response["result"] == "ok":
+ return response["content"]
+ else:
+ raise TeaclaveException("Failed to get function")
def delete_function(self, function_id: str):
+ self.check_metadata()
+ self.check_channel()
request = DeleteFunctionRequest(self.metadata, function_id)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]
+ if response["result"] == "ok":
+ return response["content"]
+ else:
+ raise TeaclaveException("Failed to delete function")
def register_input_file(self, url: str, schema: str, key: List[int],
iv: List[int], cmac: List[int]):
+ self.check_metadata()
+ self.check_channel()
request = RegisterInputFileRequest(self.metadata, url, cmac,
CryptoInfo(schema, key, iv))
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["data_id"]
+ if response["result"] == "ok":
+ return response["content"]["data_id"]
+ else:
+ raise TeaclaveException("Failed to register input file")
def register_output_file(self, url: str, schema: str, key: List[int],
iv: List[int]):
+ self.check_metadata()
+ self.check_channel()
request = RegisterOutputFileRequest(self.metadata, url,
CryptoInfo(schema, key, iv))
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["data_id"]
+ if response["result"] == "ok":
+ return response["content"]["data_id"]
+ else:
+ raise TeaclaveException("Failed to register output file")
def create_task(self,
function_id: str,
@@ -493,39 +510,63 @@ class FrontendClient:
executor: str,
inputs_ownership: List[OwnerList] = [],
outputs_ownership: List[OwnerList] = []):
+ self.check_metadata()
+ self.check_channel()
function_arguments = json.dumps(function_arguments)
request = CreateTaskRequest(self.metadata, function_id,
function_arguments, executor,
inputs_ownership, outputs_ownership)
_write_message(self.channel, request)
response = _read_message(self.channel)
- return response["content"]["task_id"]
+ if response["result"] == "ok":
+ return response["content"]["task_id"]
+ else:
+ raise TeaclaveException("Failed to create task")
def assign_data_to_task(self, task_id: str, inputs: List[DataMap],
outputs: List[DataMap]):
+ self.check_metadata()
+ self.check_channel()
request = AssignDataRequest(self.metadata, task_id, inputs, outputs)
_write_message(self.channel, request)
- _ = _read_message(self.channel)
- return
+ response = _read_message(self.channel)
+ if response["result"] == "ok":
+ pass
+ else:
+ raise TeaclaveException("Failed to assign data to task")
def approve_task(self, task_id: str):
+ self.check_metadata()
+ self.check_channel()
request = ApproveTaskRequest(self.metadata, task_id)
_write_message(self.channel, request)
- _ = _read_message(self.channel)
- return
+ response = _read_message(self.channel)
+ if response["result"] == "ok":
+ pass
+ else:
+ raise TeaclaveException("Failed to approve task")
def invoke_task(self, task_id: str):
+ self.check_metadata()
+ self.check_channel()
request = InvokeTaskRequest(self.metadata, task_id)
_write_message(self.channel, request)
response = _read_message(self.channel)
- assert (response["result"] == "ok")
+ if response["result"] == "ok":
+ pass
+ else:
+ raise TeaclaveException("Failed to invoke task")
def get_task_result(self, task_id: str):
+ self.check_metadata()
+ self.check_channel()
request = GetTaskRequest(self.metadata, task_id)
while True:
_write_message(self.channel, request)
response = _read_message(self.channel)
+ if response["result"] != "ok":
+ raise TeaclaveException("Failed to get task result")
time.sleep(1)
if response["content"]["status"] == 10:
break
@@ -533,13 +574,18 @@ class FrontendClient:
return response["content"]["result"]["result"]["Ok"]["return_value"]
def get_output_cmac_by_tag(self, task_id: str, tag: str):
+ self.check_metadata()
+ self.check_channel()
request = GetTaskRequest(self.metadata, task_id)
while True:
_write_message(self.channel, request)
response = _read_message(self.channel)
+ if response["result"] != "ok":
+ raise TeaclaveException("Failed to get task result")
time.sleep(1)
if response["content"]["status"] == 10:
break
+
return response["content"]["result"]["result"]["Ok"]["tags_map"][tag]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@teaclave.apache.org
For additional commands, e-mail: commits-help@teaclave.apache.org