feat(sgx): support multi-measurements attestation of tensorflow compo…

…nents (#1092) * feat(sgx): support tensorflow remote attestation in list * fix(sgx): improve EPC cost * fix(sgx): fix memory leak of json_engine * fix(sgx):misspell Ps->ps
bytedance · Apr 7, 2024 · ab1a6bb · ab1a6bb
1 parent c3ea95d
commit ab1a6bb
Show file tree

Hide file tree

Showing 14 changed files with 470 additions and 246 deletions.
diff --git a/deploy/scripts/sgx/enclave_env.sh b/deploy/scripts/sgx/enclave_env.sh
@@ -15,11 +15,46 @@
 # limitations under the License.
 
 EXEC_DIR=/app/exec_dir
+SGX_CONFIG_PATH="$GRPC_PATH/examples/dynamic_config.json"
+TEMPLATE_PATH="/gramine/CI-Examples/generate-token/python.manifest.template"
+
+# 更新sgx的认证策略
+update_sgx_dynamic_config() {
+  local mr_enclave="$1"
+  local mr_signer="$2"
+  local isv_prod_id="$3"
+  local isv_svn="$4"
+  local json_data=$(cat "$SGX_CONFIG_PATH")
+
+  # 创建json体
+  local new_mrs=$(jq -n --arg mr_enclave "$mr_enclave" \
+                         --arg mr_signer "$mr_signer" \
+                         --arg isv_prod_id "$isv_prod_id" \
+                         --arg isv_svn "$isv_svn" \
+                         '{
+                             mr_enclave: $mr_enclave,
+                             mr_signer: $mr_signer,
+                             isv_prod_id: $isv_prod_id,
+                             isv_svn: $isv_svn
+                         }')
+
+  # 检查sgx_mrs数组中是否存在相同的条目
+  local exists=$(echo "$json_data" | jq --argjson check_mrs "$new_mrs" \
+                    '.sgx_mrs[] | select(.mr_enclave == $check_mrs.mr_enclave and .mr_signer == $check_mrs.mr_signer and .isv_prod_id == $check_mrs.isv_prod_id and .isv_svn == $check_mrs.isv_svn) | . != null')
+
+  # 不重复添加
+  if [[ -z "$exists" ]]; then
+    json_data=$(echo "$json_data" | jq --argjson new_mrs "$new_mrs" '.sgx_mrs += [$new_mrs]')
+    echo "$json_data" > "$SGX_CONFIG_PATH"
+  fi
+}
 
+# 从sig中获取度量值hex
 function get_env() {
     gramine-sgx-get-token -s python.sig -o /dev/null | grep $1 | awk -F ":" '{print $2}' | xargs
 }
 
+# 设置自定义环境
 function make_custom_env() {
     cd $EXEC_DIR
 
@@ -44,10 +79,6 @@ function make_custom_env() {
     export JAVA_HOME=/opt/tiger/jdk/openjdk-1.8.0_265
     export LD_LIBRARY_PATH=${HADOOP_HOME}/lib/native:${JAVA_HOME}/jre/lib/amd64/server:${LD_LIBRARY_PATH}
     export CLASSPATH=.:$CLASSPATH:$JAVA_HOME/lib/dt.jar:$JAVA_HOME/lib/tools.jar:$($HADOOP_HOME/bin/hadoop classpath --glob)
-    export MR_ENCLAVE=`get_env mr_enclave`
-    export MR_SIGNER=`get_env mr_signer`
-    export ISV_PROD_ID=`get_env isv_prod_id`
-    export ISV_SVN=`get_env isv_svn`
     export RA_TLS_ALLOW_OUTDATED_TCB_INSECURE=1
 
     if [ -z "$PEER_MR_SIGNER" ]; then
@@ -58,19 +89,16 @@ function make_custom_env() {
         export PEER_MR_ENCLAVE=`get_env mr_enclave`
     fi
 
-    # network proxy
-    unset http_proxy https_proxy
-    # need meituan's
-    jq --arg mr_enclave "$PEER_MR_ENCLAVE" --arg mr_signer "$PEER_MR_SIGNER" \
-        '.sgx_mrs[0].mr_enclave = $mr_enclave | .sgx_mrs[0].mr_signer = $mr_signer' \
-        $GRPC_PATH/examples/dynamic_config.json > $EXEC_DIR/dynamic_config.json
-
+    update_sgx_dynamic_config $PEER_MR_ENCLAVE $PEER_MR_SIGNER 0 0
+    cp $SGX_CONFIG_PATH $EXEC_DIR
     cd -
 }
 
+# 生成enclave和token
 function generate_token() {
     cd /gramine/CI-Examples/generate-token/
     ./generate.sh
+    update_sgx_dynamic_config `get_env mr_enclave` `get_env mr_signer` 0 0
     mkdir -p $EXEC_DIR
     cp /app/sgx/gramine/CI-Examples/tensorflow_io.py $EXEC_DIR
     cp python.sig $EXEC_DIR
@@ -80,15 +108,66 @@ function generate_token() {
     cd -
 }
 
+# 根据enclave_size调整enclave
+function build_enclave(){
+    local enclave_size="$1"
+    local need_clean="$2"
+    sed -i "/sgx.enclave_size/ s/\"[^\"]*\"/\"$enclave_size\"/" "$TEMPLATE_PATH"
+    if [ $? -eq 0 ]; then
+        echo "Enclave size changed to $enclave_size in $TEMPLATE_PATH"
+    else
+        echo "Failed to change enclave size in $TEMPLATE_PATH"
+    fi
+    generate_token
+    if [ -n "$need_clean" ]; then
+        rm -rf $EXEC_DIR
+    fi
+}
+
+function build_enclave_all(){
+    local enclave_size="8G"
+    if [ -n "$1" ] && [ $1 == "ps" ]; then
+        # build worker/master
+        if [ -n "$GRAMINE_ENCLAVE_SIZE" ]; then
+            enclave_size=$GRAMINE_ENCLAVE_SIZE
+        fi
+        build_enclave $enclave_size 1
+
+        # build ps
+        if [ -n "$COSTOM_PS_SIZE" ]; then
+            enclave_size=$COSTOM_PS_SIZE
+        else
+            enclave_size="16G"
+        fi
+        build_enclave $enclave_size
+    else
+        # build ps
+        if [ -n "$COSTOM_PS_SIZE" ]; then
+            enclave_size=$COSTOM_PS_SIZE
+        else
+            enclave_size="16G"
+        fi
+        build_enclave $enclave_size 1
+
+        # build worker/master
+        if [ -n "$GRAMINE_ENCLAVE_SIZE" ]; then
+            enclave_size=$GRAMINE_ENCLAVE_SIZE
+        else
+            enclave_size="8G"
+        fi
+        build_enclave $enclave_size
+    fi
+}
+
 if [ -n "$PCCS_IP" ]; then
         sed -i "s|PCCS_URL=https://[^ ]*|PCCS_URL=https://pccs_url:8081/sgx/certification/v3/|" /etc/sgx_default_qcnl.conf
         echo >> /etc/hosts
         echo "$PCCS_IP   pccs_url" | tee -a /etc/hosts
 elif [ -n "$PCCS_URL" ]; then
         sed -i "s|PCCS_URL=[^ ]*|PCCS_URL=$PCCS_URL|" /etc/sgx_default_qcnl.conf
 fi
+sed -i 's/USE_SECURE_CERT=TRUE/USE_SECURE_CERT=FALSE/' /etc/sgx_default_qcnl.conf
 
-TEMPLATE_PATH="/gramine/CI-Examples/generate-token/python.manifest.template"
 if [ -n "$GRAMINE_LOG_LEVEL" ]; then
         sed -i "/loader.log_level/ s/\"[^\"]*\"/\"$GRAMINE_LOG_LEVEL\"/" "$TEMPLATE_PATH"
         if [ $? -eq 0 ]; then
@@ -98,15 +177,6 @@ if [ -n "$GRAMINE_LOG_LEVEL" ]; then
         fi
 fi
 
-if [ -n "$GRAMINE_ENCLAVE_SIZE" ]; then
-    sed -i "/sgx.enclave_size/ s/\"[^\"]*\"/\"$GRAMINE_ENCLAVE_SIZE\"/" "$TEMPLATE_PATH"
-    if [ $? -eq 0 ]; then
-        echo "Enclave size changed to $GRAMINE_ENCLAVE_SIZE in $TEMPLATE_PATH"
-    else
-        echo "Failed to change enclave size in $TEMPLATE_PATH"
-    fi
-fi
-
 if [ -n "$GRAMINE_THREAD_NUM" ]; then
     sed -i "s/sgx.thread_num = [0-9]\+/sgx.thread_num = $GRAMINE_THREAD_NUM/" "$TEMPLATE_PATH"
     if [ $? -eq 0 ]; then
@@ -125,7 +195,5 @@ if [ -n "$GRAMINE_STACK_SIZE" ]; then
     fi
 fi
 
-sed -i 's/USE_SECURE_CERT=TRUE/USE_SECURE_CERT=FALSE/' /etc/sgx_default_qcnl.conf
 mkdir -p /data
-
-generate_token
+build_enclave_all $1
diff --git a/deploy/scripts/sgx/run_trainer_master_sgx.sh b/deploy/scripts/sgx/run_trainer_master_sgx.sh
@@ -92,7 +92,7 @@ fi
 
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/follower/
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/leader/
-source /app/deploy/scripts/sgx/enclave_env.sh
+source /app/deploy/scripts/sgx/enclave_env.sh master
 
 unset HTTPS_PROXY https_proxy http_proxy ftp_proxy
 

diff --git a/deploy/scripts/sgx/run_trainer_ps_sgx.sh b/deploy/scripts/sgx/run_trainer_ps_sgx.sh
@@ -35,7 +35,7 @@ fi
 
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/leader
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/follower
-source /app/deploy/scripts/sgx/enclave_env.sh
+source /app/deploy/scripts/sgx/enclave_env.sh ps
 
 make_custom_env 4
 source /root/start_aesm_service.sh

diff --git a/deploy/scripts/sgx/run_trainer_worker_sgx.sh b/deploy/scripts/sgx/run_trainer_worker_sgx.sh
@@ -46,7 +46,7 @@ fi
 
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/follower/
 cp /app/sgx/gramine/CI-Examples/tensorflow_io.py /gramine/leader/
-source /app/deploy/scripts/sgx/enclave_env.sh
+source /app/deploy/scripts/sgx/enclave_env.sh worker
 
 unset HTTPS_PROXY https_proxy http_proxy ftp_proxy
 

diff --git a/fedlearner/channel/client_interceptor.py b/fedlearner/channel/client_interceptor.py
@@ -364,6 +364,8 @@ def _grpc_error_need_recover(e):
             #if 400 <= httpstatus < 500:
             #    return True
             return True
+    if e.code() == grpc.StatusCode.UNAUTHENTICATED:
+        return False
     return True # recover in any case
     #return False
 

diff --git a/sgx/gramine/CI-Examples/generate-token/python.manifest.template b/sgx/gramine/CI-Examples/generate-token/python.manifest.template
@@ -9,15 +9,22 @@ loader.pal_internal_mem_size = "200M"
 
 loader.insecure__use_cmdline_argv = true
 loader.insecure__use_host_env = true
+loader.insecure_disable_aslr = true
+sgx.allow_file_creation = true
+sgx.static_address = 1
 
 loader.env.LD_LIBRARY_PATH = "/opt/meituan/hadoop/lib/native:/opt/tiger/jdk/openjdk-1.8.0_265/jre/lib/amd64/server:/opt/meituan/hadoop/lib/native:/opt/meituan/hadoop/lib/native/nfs:/usr/lib/jvm/java-8-openjdk-amd64/jre/lib/amd64/server:/usr/local/lib/x86_64-linux-gnu:/opt/tiger/yarn_deploy/hadoop_current/lib/native:/opt/tiger/jdk/openjdk-1.8.0_265/jre/lib/amd64/server:/opt/tiger/jdk/openjdk-1.8.0_265/jre/lib/amd64/jli/:/opt/tiger/jdk/openjdk-1.8.0_265/jre/lib/amd64:{{ python.stdlib }}/lib:/lib:{{ arch_libdir }}:/usr/local/lib:/usr/local/{{ arch_libdir }}:/usr/lib:/usr/{{ arch_libdir }}"
-
 loader.env.SECRET_PROVISION_CONSTRUCTOR = "1"
 loader.env.SECRET_PROVISION_SET_PF_KEY = "1"
 
 sys.enable_sigterm_injection = true
 sys.enable_extra_runtime_domain_names_conf = true
-sys.stack.size = "1M"
+sgx.ioctl_structs.ifconf = [ {name = "ifc_len",size = 4, direction = "inout" }, {size = 4}, { ptr=[ {size= "ifc_len", direction = "in"} ] } ]
+sgx.ioctl_structs.ifreq = [ {  name ="ifreq", size=40, direction ="inout"} ]
+sgx.allowed_ioctls = [
+  { request_code = 0x8912, struct = "ifconf"},
+  { request_code = 0x8927, struct = "ifreq"},
+]
 
 fs.mounts = [
   { path = "/lib", uri = "file:{{ gramine.runtimedir() }}" },
@@ -48,7 +55,7 @@ sgx.remote_attestation = "dcap"
 sgx.preheat_enclave = false
 sgx.nonpie_binary = true
 sgx.enable_stats = false
-sgx.enclave_size = "16G"
+sgx.enclave_size = "8G"
 sgx.thread_num = 256
 sgx.rpc_thread_num = 0
 sgx.protected_files_key = "ffeeddccbbaa99887766554433221100"
@@ -60,7 +67,9 @@ sgx.trusted_files = [
   "file:/usr/{{ arch_libdir }}/",
   "file:/etc/ssl/certs/ca-certificates.crt",
   "file:/etc/default/apport",
-  "file:/etc/mime.types"
+  "file:/etc/mime.types",
+  "file:/gramine/leader/",
+  "file:/gramine/follower/"
 ]
 
 sgx.allowed_files = [
@@ -94,7 +103,4 @@ sgx.allowed_files = [
   "file:/lib/",
   "file:/bin/",
   "file:/data/",
-  "file:/gramine/leader/",
-  "file:/gramine/follower/"
-]
-
+]
diff --git a/sgx/grpc/v1.38.1/examples/dynamic_config.json b/sgx/grpc/v1.38.1/examples/dynamic_config.json
@@ -5,12 +5,6 @@
     "verify_isv_prod_id" : "on",
     "verify_isv_svn" : "on",
     "sgx_mrs": [
-        {
-            "mr_enclave" : "",
-            "mr_signer" : "",
-            "isv_prod_id" : "0",
-            "isv_svn" : "0"
-        }
     ],
     "other" : []
 }