开发TEE的踩坑之开发TEE

系统：Ubuntu20.04（双系统，非虚拟机）

一、前置说明
- 1、TEE平台的选择
- 2、机器间的通信方式
- 3、程序和数据集的示例
- 4、结果文件的解密
二、服务器部署
三、客户端部署
四、工程应用

本系列为笔者开发TEE（Trusted Execution Environment，可信执行环境）的系列踩坑文，给广大开发者分享自己所谓的“经验”，希望对大家有帮助。

开发路线为：

装Ubuntu20.04的双系统（不展开，B站有很多教程）
配置SGX环境（前提是SGX机器，最好是SGX2）
配置PCCS（很蛋疼，英特尔官方github的步骤也未必行得通）
开发TEE（花的时间最多）

本篇文章分享开发TEE的踩坑
笔者是在某蚁集团的Occlum的基础上，结合ChatGPT给出的回答去解决途中遇到的问题的

一、前置说明

1、TEE平台的选择

TEE平台有众多选择，在此之前，笔者已在某语的TrustedFlow踩过无数坑了。但奈何若要使用他们的平台开发TEE，需要有可以支持SGX_2的CPU，否则一切都是徒劳，只能用仿真模式，这无异于把程序放在裸机上跑，已与TEE毫无相关；还有一个致命之处是他们例程的角色过于复杂，若要用于实际的落地工程，需要具备良好的密码学功底，并花费大量时间刨出自己所需要的模块为己所用。
基于以上缺陷，笔者最后选择了某蚁集团的Occlum开发TEE。

2、机器间的通信方式

在介绍通信方式之前，首先要明确一点：TEE沙箱是作为服务器（Server），而机构是作为客户端（Client），且二者同处一个局域网内，否则是否“可信”要打上问号了。明确角色后，且两台机器同处于一个局域网内，那么通信方式便怎么简单怎么来。
此处笔者以flask（基于Python的Web框架）进行客户端和服务器的通信（该通信方式不含双向远程认证），

3、程序和数据集的示例

对于没接触过Occlum的读者可以模仿这篇文章跑一份程序体验一下。
【注意：机器一定要预先配置好SGX环境，可参考该系列踩坑文的配置SGX环境】
笔者开发TEE的场景是隐私求交，但目前做的还只是明文求交，只不过程序确实是在TEE环境（Enclave，飞地）运行的，是可以保证“可信”。
程序语言选用C，之所以不采用Python，是因为求交所需要安装的库和依赖之间有版本冲突不兼容的问题，此处笔者已踩过坑，读者可尝试并留下评论互相学习噢。
进行求交的数据集此处采用经典的乳腺癌数据集：alice和bob.

4、结果文件的解密

若读者有接触过Occlum，相信第一次使用一定会遇到读取数据集的路径问题；若还定义了生成结果文件的路径，也肯定会遇到生成结果文件是密文且多个的问题，这些密文文件是保存在/occlum_instance/run/mount/__ROOT，如下图：
在这里插入图片描述
笔者已询问过Occlum的官方人员，但回答是使用挂载mount进行解密。笔者至今仍不明白目录挂载与结果文件的解密有何联系，亦没有渠道对这些密文结果文件（里面都是乱码）进行解密，有兴趣的读者可研究一番并留下评论互相学习噢。
基于Occlum的这种生成密文结果的机制，笔者将计就计，绕开这个问题，却依然能输出明文结果并返回给客户端。读者且保持耐心继续往下浏览。

二、服务器部署

1、起特权模式、映射/dev/sgx和mount文件夹（mount用于在外部宿主机也存储结果文件）、与宿主机的时间同步、映射50054和50055端口号（一个用于通信请求，一个用于被访问结果文件，此处以50054和50055为例）的容器。

sudo docker run -it --privileged -v /dev/sgx:/dev/sgx -v /home/yunqi/mount:/root/mount -v /etc/localtime:/etc/localtime:ro -v /etc/timezone:/etc/timezone:ro -p 50054:50054 -p 50055:50055 <你的镜像源>/occlum:latest-ubuntu20.04

【温馨提示：6月份的时候，由于某些不可抗拒的原因，国内的dockerhub已经禁止访问了，因此读者可以参考该视频，把你需要的docker镜像用Github放在自己的阿里云上】

2、在容器内安装flask库（采用阿里云镜像源更快），用于客户端（机构）和服务器（TEE沙箱）间的通信，前面已有介绍。

pip3 install Flask -i https://mirrors.aliyun.com/pypi/simple/

3、用于存放程序的文件夹

mkdir program

4、用于存放服务器的流程代码和自动化指令

mkdir server_file

5、新建用于求交的程序

cd ./program
touch psi.c
vim psi.c

运行的程序一般是放在/occlum_instance/image/bin，且由于Occlum的特性，读者不必质疑define的路径

#psi.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define FILE1_PATH "/bin/file1.csv"
#define FILE2_PATH "/bin/file2.csv"

#define MAX_LINE_LENGTH 1024

typedef struct {
  char id[100];
  char attributes[900];
} Record;

void trim_newline(char *str) {
  char *newline = strchr(str, '\n');
  if (newline) {
    *newline = '\0';
  }
}

int main() {
  FILE *File1 = fopen(FILE1_PATH, "r");
  FILE *File2 = fopen(FILE2_PATH, "r");

  if (!File1 || !File2) {
    fprintf(stderr, "Error opening file.\n");
    return 1;
  }

  char line[MAX_LINE_LENGTH];
  Record Records1[1000];
  Record Records2[1000];
  int Count1 = 0, Count2 = 0;

  // Read file1
  while (fgets(line, MAX_LINE_LENGTH, File1)) {
    trim_newline(line);
    char *token = strtok(line, ",");
    strcpy(Records1[Count1].id, token);
    token = strtok(NULL, "");
    if (token != NULL) {
      strcpy(Records1[Count1].attributes, token);
    } else {
      strcpy(Records1[Count1].attributes, "");
    }
    Count1++;
  }
  fclose(File1);

  // Read file2
  while (fgets(line, MAX_LINE_LENGTH, File2)) {
    trim_newline(line);
    char *token = strtok(line, ",");
    strcpy(Records2[Count2].id, token);
    token = strtok(NULL, "");
    if (token != NULL) {
      strcpy(Records2[Count2].attributes, token);
    } else {
      strcpy(Records2[Count2].attributes, "");
    }
    Count2++;
  }
  fclose(File2);

  // Find intersection and print results
  printf("The intersection result is as follows\n");
  for (int i = 0; i < Count1; i++) {
    for (int j = 0; j < Count2; j++) {
      if (strcmp(Records1[i].id, Records2[j].id) == 0) {
        printf("%s,%s,%s\n", Records1[i].id, Records1[i].attributes, Records2[j].attributes);
      }
    }
  }

  return 0;
}

6、新建存放自动化指令的文件

cd ../server_file
touch commands.txt
vim commands.txt

occlum-gcc -o psi psi.c
mkdir occlum_instance
cd ./occlum_instance
occlum init
cp ../file1.csv ./image/bin
cp ../file2.csv ./image/bin
cp ../psi ./image/bin
occlum build --sgx-mode HW
occlum run /bin/psi

7、新建服务器的流程代码文件

touch task_server.py
vim task_server.py

# task_server.py
from flask import Flask, jsonify, request, send_from_directory
import subprocess
import shutil
import uuid
import os
import csv
import threading
import time

app = Flask(__name__)
http_server_process = None

def start_http_server(): 
    global http_server_process
    if not http_server_process:
        # 用于被客户端访问结果文件
        http_server_process = subprocess.Popen(['python3', '-m', 'http.server', '50055'], cwd="/root/mount")
        time.sleep(2)

def stop_http_server():
    global http_server_running
    if http_server_process:
        http_server_process.terminate()
        http_server_process.wait()
        http_server_process = None

@app.route('/execute-task', methods=['POST'])
def execute_task():
    global http_server_thread, http_server_running
    try:
        # 生成唯一的任务ID，格式为uuid
        task_id = str(uuid.uuid4())

        # 获取当前时间，并创建以当前时间命名的occlum工程文件夹
        current_time = subprocess.check_output("date +'%Y-%m-%d_%H-%M-%S'", shell=True).decode('utf-8').strip()
        folder_name = f'/root/{current_time}'
        os.makedirs(folder_name, exist_ok=True)

        # 拷贝求交的C程序到occlum工程文件夹
        source_c_file = '/root/program/psi.c'
        destination_c_file = os.path.join(folder_name, 'psi.c')
        shutil.copy(source_c_file, destination_c_file)

        # 保存客户端发送来的csv数据集到occlum工程文件夹
        file1_path = os.path.join(folder_name, 'file1.csv')
        file2_path = os.path.join(folder_name, 'file2.csv')
        if 'file1' in request.files:
          file1 = request.files['file1']
          file1.save(file1_path)
        if 'file2' in request.files:
          file2 = request.files['file2']
          file2.save(file2_path)

        # 进入occlum工程文件夹
        os.chdir(folder_name)

        # 读取并执行命令文件中的occlum指令
        command_file = '/root/server_file/commands.txt'
        with open(command_file, 'r') as file:
          commands = file.readlines()

        # 使用 `/bin/bash -c` 执行整个命令字符串，并捕获输出
        command_str = '\n'.join(commands)
        result =  subprocess.run(f'/bin/bash -c "{command_str}"', shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        # 获取求交C程序的打印输出
        c_program_output = result.stdout

        # 找到C程序输出的"The intersection result is as follows"这一行的位置
        lines = c_program_output.splitlines()
        start_index = lines.index("The intersection result is as follows") + 1

        # 提取该行之后的内容
        output_lines = lines[start_index:]

        # 将提取的内容写入并保存csv结果文件
        output_csv_file_name = f'{current_time}_result_psi.csv'
        output_csv_file = os.path.join("/root/mount", output_csv_file_name)
        with open(output_csv_file, 'w', newline='') as csvfile:
            csvwriter = csv.writer(csvfile)
            for line in output_lines:
                if line.strip():
                    csvwriter.writerow(line.split(','))

        # 启动http服务器，以便客户端可以访问TEE沙箱的/root/mount的csv结果文件
        start_http_server()

        # 返回任务ID和生成的结果文件名给客户端
        return jsonify({'task_id': task_id, 'file_name': output_csv_file_name})

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/download/<filename>', methods=['GET'])
def download_file(filename):
    directory = "/root/mount"
    response = send_from_directory(directory, filename)
    stop_http_server()
    return response

if __name__ == '__main__':
  # 用于接受通信请求
  app.run(host='0.0.0.0', port=50054)

三、客户端部署

相较于服务器的部署工作，客户端就显得尤为简洁了，只需做好发送数据集和通信请求的工作就好了。
【注意：服务器和客户端的系统都是Ubuntu20.04】

mkdir client_file
<存放alice和bob的数据集在此处>
touch request_client.py
vim request_client.py

# request_client.py
import requests

# 替换为TEE沙箱的IP地址和Occlum容器用于接受通信请求的端口号
url = 'http://<TEE沙箱的IP地址>:50054/execute-task' 
files = {
    'file1': open('/home/<用户名>/client_file/alice.csv', 'rb'),
    'file2': open('/home/<用户名>/client_file/bob.csv', 'rb')
}

response = requests.post(url, files=files)

if response.status_code == 200:
    data = response.json()
    print("Task execution triggered successfully.")
    print("Task ID:", data['task_id'])
    print("Generated file name:", data['file_name'])
else:
    print("Failed to trigger task execution.")
    print("Error:", response.text)

四、工程应用

做好以上部署工作后，读者们可以开始正式感受TEE了。

1、服务器先运行

python3 task_server.py

2、客户端再运行

python3 request_client.py

3、服务器和客户端的输出分别是
在这里插入图片描述
可以看到服务器已经打开了50055的端口号，等待客户端访问。

在这里插入图片描述
可以看到客户端已经接收到服务器返回的任务ID和结果文件名.

4、此时，在客户端的浏览器可以输入<服务器的IP地址>:50055进入服务器的/root/mount文件夹，查看并下载求交的结果文件，如下图
在这里插入图片描述
5、此工程代码还有一大优势：
只要服务器一旦执行了python3 task_server.py后，理论上来说客户端是可以无限次向服务器执行求交请求的，服务器可以不再需要人工操作。