【SGX】Ocall Proxy处调用Ocall出现异常的解决办法

问题描述

有的Ocall Proxy（或者称之为Wrapper/Bridge）样子如下

sgx_status_t SGX_CDECL ocall_print_string(const char* str)
{
	...
	__tmp = sgx_ocalloc(ocalloc_size);
	...
	status = sgx_ocall(8, ms);
	...
	sgx_ocfree();
	...
	return status;
}

如果我们在sgx_ocalloc后又试图调用了包含sgx_ocalloc和sgx_ocfree的Ocall Proxy，就会出现异常。

此时，调用顺序变成了sgx_ocalloc$\to$sgx_ocalloc$\to$sgx_ocfree$\to$sgx_ocfree。

每一次sgx_ocalloc都会在Untrusted Thread Stack上调整（减）Untrusted SP寄存器来容纳OCALL想要传出的Marshaling Parameter，这个调整是依据每次Ocall的需求逐渐调整的。

而sgx_ocfree恢复Untrusted SP寄存器时，是从thread_data_t.last_sp中直接找到第一次sgx_ocalloc前USP值，因此无论sgx_ocalloc有几次，一次sgx_ocfree就会将USP恢复到若干次sgx_ocalloc前。

按照正常情况，Ocall Proxy中是不允许再调用Ocall的，这是SGX及其Edger8r工具做出的假设，也挺合理的。

但是当我有特殊的需求，比如在Ocall Proxy中插桩Ocall打印一些东西或其它功能，那么就会出现异常。

解决方法

在子Ocall的sgx_ocalloc前获取当前USP值

size_t get_untrust_sp()
{
    thread_data_t *thread_data = get_thread_data();
    ssa_gpr_t *ssa_gpr = reinterpret_cast<ssa_gpr_t *>(thread_data->first_ssa_gpr);
    return ssa_gpr->REG(sp_u);
}

在子Ocall的sgx_ocfree后（此时USP已经恢复到任意次sgx_ocalloc前）设置USP为之前保存的值

void set_untrust_sp(size_t addr)
{
    thread_data_t *thread_data = get_thread_data();
    ssa_gpr_t *ssa_gpr = reinterpret_cast<ssa_gpr_t *>(thread_data->first_ssa_gpr);
    ssa_gpr->REG(sp_u) = addr;
}

LLVM Pass对此插桩的例子：

void adjustUntrustedSPRegisterAtOcallAllocAndFree(Function &F)
{
    // initialize
    Module *M = F.getParent();
    IRBuilder<> IRB(M->getContext());
    FunctionCallee SetUSP = M->getOrInsertFunction("set_untrust_sp", IRB.getVoidTy(), IRB.getInt64Ty());
    FunctionCallee GetUSP = M->getOrInsertFunction("get_untrust_sp", IRB.getInt64Ty());

    // get interesting callinst
    SmallVector<CallInst *, 8> OcallocVec, OcfreeVec;
    for (auto &BB : F)
    {
        for (auto &I : BB)
        {
            if (CallInst *CI = dyn_cast<CallInst>(&I))
            {
                Function *callee = CI->getCalledFunction();
                if (callee != nullptr)
                {
                    StringRef callee_name = callee->getName();
                    assert(callee_name != "");
                    if (callee_name == "sgx_ocalloc")
                    {
                        OcallocVec.push_back(CI);
                    }
                    else if (callee_name == "sgx_ocfree")
                    {
                        OcfreeVec.push_back(CI);
                    }
                }
            }
        }
    }

    // instrument
    IRB.SetInsertPoint(&F.getEntryBlock().getInstList().front());
    Value *usp = IRB.CreateAlloca(IRB.getInt64Ty());
    for (auto CI : OcallocVec)
    {
        IRB.SetInsertPoint(CI);
        IRB.CreateStore(IRB.CreateCall(GetUSP), usp);
    }
    for (auto CI : OcfreeVec)
    {
        IRB.SetInsertPoint(CI->getNextNode());
        IRB.CreateCall(SetUSP, IRB.CreateLoad(usp));
    }
}

由于这些特殊的操作的声明只要SGX Internal内部出现，还需要定义一些Internal Structural

#pragma once
#include <stdint.h>
#include <stddef.h>

typedef size_t sys_word_t;

typedef struct _thread_data_t
{
    sys_word_t self_addr;
    sys_word_t last_sp;          /* set by urts, relative to TCS */
    sys_word_t stack_base_addr;  /* set by urts, relative to TCS */
    sys_word_t stack_limit_addr; /* set by urts, relative to TCS */
    sys_word_t first_ssa_gpr;    /* set by urts, relative to TCS */
    sys_word_t stack_guard;      /* GCC expects start_guard at 0x14 on x86 and 0x28 on x64 */

    sys_word_t flags;
    sys_word_t xsave_size; /* in bytes (se_ptrace.c needs to know its offset).*/
    sys_word_t last_error; /* init to be 0. Used by trts. */
    struct _thread_data_t *m_next;
    sys_word_t tls_addr;  /* points to TLS pages */
    sys_word_t tls_array; /* points to TD.tls_addr relative to TCS */
    intptr_t exception_flag;
    sys_word_t cxx_thread_info[6];
    sys_word_t stack_commit_addr;
} thread_data_t;

#define REG(name) r##name
#define REGISTER(name) uint64_t REG(name)

/****************************************************************************
 * Definitions for SSA
 ****************************************************************************/
typedef struct _exit_info_t
{
    uint32_t vector : 8;    /* Exception number of exceptions reported inside enclave */
    uint32_t exit_type : 3; /* 3: Hardware exceptions, 6: Software exceptions */
    uint32_t reserved : 20;
    uint32_t valid : 1; /* 0: unsupported exceptions, 1: Supported exceptions */
} exit_info_t;

typedef struct _ssa_gpr_t
{
    REGISTER(ax);          /* (0) */
    REGISTER(cx);          /* (8) */
    REGISTER(dx);          /* (16) */
    REGISTER(bx);          /* (24) */
    REGISTER(sp);          /* (32) */
    REGISTER(bp);          /* (40) */
    REGISTER(si);          /* (48) */
    REGISTER(di);          /* (56) */
    uint64_t r8;           /* (64) */
    uint64_t r9;           /* (72) */
    uint64_t r10;          /* (80) */
    uint64_t r11;          /* (88) */
    uint64_t r12;          /* (96) */
    uint64_t r13;          /* (104) */
    uint64_t r14;          /* (112) */
    uint64_t r15;          /* (120) */
    REGISTER(flags);       /* (128) */
    REGISTER(ip);          /* (136) */
    REGISTER(sp_u);        /* (144) untrusted stack pointer. saved by EENTER */
    REGISTER(bp_u);        /* (152) untrusted frame pointer. saved by EENTER */
    exit_info_t exit_info; /* (160) contain information for exits */
    uint32_t reserved;     /* (164) padding to multiple of 8 bytes */
    uint64_t fs;           /* (168) FS register */
    uint64_t gs;           /* (176) GS register */
} ssa_gpr_t;

#if defined(__cplusplus)
extern "C"
{
#endif
    thread_data_t *get_thread_data(void);
#if defined(__cplusplus)
}
#endif