mercredi 8 septembre 2021

A segmentation fault problem may caused by stack while I try to code my simple coroutine

I try to write a simple coroutine on x86_64 platform, but I met with problems as I make a context switch. The code is organized as follows, and finally it will output "succeed to save _current's environment\nSegmentation fault" on my terminal.

//coroutine.hpp
#ifndef _COROUTINE_HPP
#define _COROUTINE_HPP

#include <iostream>
#include <cstdint>
#include <map>
#include <vector>

typedef std::uint64_t addr_t;
typedef std::uint64_t reg_t;

using std::map;
using std::string;
using std::cout;
using std::endl;

class CoroutineEnvFactory;
class CoroutineManager;
struct CoroutineEnv {
    reg_t rip;
    reg_t rbx;
    reg_t rsi;
    reg_t rdi;
    reg_t rsp;
    reg_t rbp;
    char* stackTop;
    reg_t stackSize;
    friend class CoroutineEnvFactory;
    friend class CoroutineManager;

protected:
    CoroutineEnv() {}
    ~CoroutineEnv() {
        delete[] (stackTop-stackSize);
        stackTop = nullptr;
    }
};

class CoroutineEnvFactory {
public:
    CoroutineEnvFactory() {}
    ~CoroutineEnvFactory() {}
    CoroutineEnv* emptyCoroutineEnv() {
        return new CoroutineEnv;
    }
    template <typename func_t>
    CoroutineEnv* funcCoroutineEnv(func_t rip, reg_t stackSize=4096) {
        CoroutineEnv* env = new CoroutineEnv;
        // initialize the stack of coroutine
        env->stackSize = stackSize < 4096 ? 4096 : stackSize;
        env->stackTop = new char[env->stackSize] + env->stackSize;
        // initialize register environment
        env->rbx = env->rsi = env->rdi = 0;
        env->rbp = env->rsp = (reg_t)(env->stackTop);
        // initialize entry point
        env->rip = (reg_t)rip;
        return env;
    }
};

class CoroutineManager {
public:
    static constexpr const char* MAIN_COROUTINE = "MAIN";
private:
    CoroutineEnvFactory _factory;
    map<string, CoroutineEnv*> _list;
    string _current;
private:
    CoroutineManager() {
        _list[MAIN_COROUTINE] = _factory.emptyCoroutineEnv();
        _current = string(MAIN_COROUTINE);
    }

    void remove(CoroutineEnv* env) {
        env->~CoroutineEnv();
        map<string, CoroutineEnv*>::iterator it;
        for (it=_list.begin(); it!=_list.end(); ++it) {
            if (it->second == env) {
                _list.erase(it);
                break;
            }
        }
    }

    static void endPoint() {
        reg_t rbp;
        asm volatile ("movq %%rbp, %0": "=r"(rbp) :: "memory");
        rbp += sizeof(addr_t);
        CoroutineEnv* coroutine = reinterpret_cast<CoroutineEnv*>(*(addr_t*)rbp);
        rbp += sizeof(addr_t);
        CoroutineManager* manager = reinterpret_cast<CoroutineManager*>(*(addr_t*)rbp);

        manager->remove(coroutine);
        manager->switch_to(MAIN_COROUTINE);
    }

public:
    // singleton
    static CoroutineManager& getManager() {
        static CoroutineManager manager;
        return manager;
    }

    void manage(string name, CoroutineEnv* coroutine) {
        coroutine->rsp -= 3*sizeof(addr_t);
        coroutine->rbp = coroutine->rsp;
        _list[name] = coroutine;
        addr_t *rsp = (addr_t*)coroutine->rsp;
        // push ending function
        *rsp = (addr_t)endPoint;
        // push it's parameters
        *(rsp+1) = (addr_t)coroutine;
        *(rsp+2) = (addr_t)this;
    }

    void switch_to(string next) {
        if (_list.end() == _list.find(next)) return;
        if (next == _current) return;
        CoroutineEnv* target = _list[next];
        CoroutineEnv* cur = _list[_current];

        // save _current's environment
        asm volatile (
            "movq 8(%%rbp), %0;"
            "movq %%rbx, %1;"
            "movq %%rsi, %2;"
            "movq %%rdi, %3;"
            "movq %%rbp, %4;"
            "movq %%rsp, %5;"
            : "=r"(cur->rip), "=r"(cur->rbx), "=r"(cur->rsi), "=r"(cur->rdi), "=r"(cur->rbp), "=r"(cur->rsp)
            :: "memory"
        );

        #ifdef _DEBUG
        cout << "succeed to save _current's environment" << endl;
        #endif

        // switch to next coroutine
        _current = next;
        asm volatile (
            "movq %0, %%rbx;"
            "movq %1, %%rsi;"
            "movq %2, %%rdi;"
            "movq %3, %%rbp;"
            "movq %4, %%rsp;"
            "pushq %5; ret;"
            :: "r"(target->rbx), "r"(target->rsi), "r"(target->rdi), "r"(target->rbp), "r"(target->rsp), "r"(target->rip)
        );
    }
};
#endif

//couroutine.cpp
#include <iostream>
#include <unistd.h>
#define _DEBUG
#include "coroutine.hpp"

using namespace std;

void foo() {
    int i = 3;
    while(--i) {
        cout << "Hello World!" << endl;
        sleep(1);
    }
}

int main() {
    CoroutineEnv* env = CoroutineEnvFactory().funcCoroutineEnv(foo);
    CoroutineManager& manager = CoroutineManager::getManager();
    manager.manage("foo", env);
    manager.switch_to("foo");
    cout << "end main" << endl;
}

I use a struct CoroutineEnv to save necessary coroutine context, which mainly contains some general purpose registers and stack-related registers. Class CoroutineEnvFactory is just used to create CoroutineEnv object. Finally, class CoroutineManager manages all coroutines created in run-time environment, whose function switch_to provides the ability to switch the context to another assigned coroutine.

However, as I assign the %rsp(in my code it's "movq %4, %%rsp;") of target coroutine, the next instruction "pushq %5;"(Here, %5 is equal to the rip of target coroutine) triggers a segmentation fault(I verify it by comment the instruction "pushq %5;"). I cannot tell whether the memory that %rsp points to is invalid or not, so I try to make some memory access with the address %rsp points to before that context switch, but there's nothing happened. I would appreciate it if you can give some advice, thanks a lot!

Aucun commentaire:

Enregistrer un commentaire