PocketPy中实现安全终止Python代码执行的方法

背景介绍

在嵌入式Python解释器PocketPy中，当我们需要执行用户提供的Python代码时，安全性是一个重要考虑因素。特别是当用户代码包含无限循环等可能耗尽系统资源的操作时，我们需要一种机制来安全地终止这些代码的执行，而不影响主程序的运行。

核心解决方案

PocketPy提供了一个名为_ceval_on_step的回调机制，允许开发者在每条字节码执行前插入自定义逻辑。通过这个机制，我们可以实现代码执行的超时控制和安全终止功能。

实现步骤

1. 启用回调功能

首先需要在用户配置中启用字节码执行回调功能：

#define PK_ENABLE_CEVAL_CALLBACK 1

2. 创建自定义VM类

创建一个继承自VM的自定义虚拟机类，添加控制标志和中断处理逻辑：

#include <atomic>
#include "pocketpy.h"

using namespace pkpy;

class SafeVM : public VM {
public:
    std::atomic<bool> should_stop{false};

    SafeVM() : VM() {
        this->_ceval_on_step = [](VM* vm, Frame* frame, Bytecode bc) {
            SafeVM* safe_vm = static_cast<SafeVM*>(vm);
            if(safe_vm->should_stop) {
                safe_vm->should_stop = false;
                safe_vm->raiseKeyboardInterrupt();
            }
        };
    }

    void raiseKeyboardInterrupt() {
        _error("KeyboardInterrupt", "");
    }
};

3. 使用自定义VM执行代码

在实际使用时，可以通过设置should_stop标志来请求终止执行：

SafeVM vm;
std::thread exec_thread([&vm](){
    try {
        vm.exec("while True: pass");
    } catch(const Exception& e) {
        std::cout << "Execution stopped: " << e.summary() << std::endl;
    }
});

// 在另一个线程中设置停止标志
std::this_thread::sleep_for(std::chrono::seconds(1));
vm.should_stop = true;
exec_thread.join();

技术原理

1. 字节码级控制：_ceval_on_step回调在每条字节码执行前被调用，提供了细粒度的执行控制。

2. 线程安全：使用std::atomic确保多线程环境下的标志访问安全。

3. 异常处理：通过抛出KeyboardInterrupt异常来终止执行，这与CPython的行为一致。

注意事项

1. 回调频率较高，应确保回调逻辑尽可能轻量。

2. 终止不是即时的，需要等待当前字节码执行完成。

3. 不同版本的PocketPy可能有不同的API，需要检查对应版本的实现。

扩展应用

这种机制不仅可以用于防止无限循环，还可以实现：

• 执行时间限制
• 内存使用监控
• 自定义调试功能
• 性能分析

通过合理利用PocketPy提供的回调机制，开发者可以构建更安全、更可控的Python执行环境。