68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
import t, c
|
|
from stdint import *
|
|
from w32.win32process import TlsAlloc, TlsGetValue, TlsSetValue
|
|
import stdlib
|
|
import stdio
|
|
import string
|
|
|
|
|
|
# 最大提供者栈深度（编译器会检查是否够用，不够再往上提）
|
|
_MAX_PROVIDER_DEPTH: t.CDefine = 32
|
|
|
|
|
|
@t.Object
|
|
class _ProviderStack:
|
|
"""提供者栈，存储字段名和指针对。"""
|
|
names: t.CArray[VOIDPTR, 32]
|
|
ptrs: t.CArray[VOIDPTR, 32]
|
|
top: UINT64
|
|
|
|
|
|
# TLS 索引（全局，所有线程共享同一个索引）
|
|
_tls_index: ULONG = 0
|
|
|
|
|
|
@t.TLS
|
|
def _init_tls():
|
|
"""初始化 TLS 索引（一次性执行，后续调用跳过）。"""
|
|
global _tls_index
|
|
_tls_index = TlsAlloc()
|
|
|
|
|
|
def _get_stack() -> _ProviderStack | t.CPtr:
|
|
"""获取当前线程的提供者栈，不存在则分配。"""
|
|
_init_tls()
|
|
stack: _ProviderStack | t.CPtr = TlsGetValue(_tls_index)
|
|
if stack == None:
|
|
stack = stdlib.malloc(_ProviderStack.__sizeof__())
|
|
string.memset(stack, 0, _ProviderStack.__sizeof__())
|
|
TlsSetValue(_tls_index, stack)
|
|
return stack
|
|
|
|
|
|
def _push_provider(field_name: str, ptr: t.CPtr):
|
|
"""压入提供者 (field_name, ptr) 到栈中。"""
|
|
stack: _ProviderStack | t.CPtr = _get_stack()
|
|
if stack.top >= _MAX_PROVIDER_DEPTH:
|
|
stdio.printf("[with] ERROR: provider stack overflow (depth=%lu, max=%d), field '%s' not injected\n", stack.top, _MAX_PROVIDER_DEPTH, field_name)
|
|
return
|
|
stack.names[stack.top] = field_name
|
|
stack.ptrs[stack.top] = ptr
|
|
stack.top = stack.top + 1
|
|
|
|
|
|
def _pop_provider():
|
|
"""弹出栈顶提供者。"""
|
|
stack: _ProviderStack | t.CPtr = _get_stack()
|
|
if stack.top > 0:
|
|
stack.top = stack.top - 1
|
|
|
|
|
|
def _find_provider(field_name: str) -> t.CPtr:
|
|
"""从栈顶向下查找匹配 field_name 的提供者，返回对应的指针。找不到返回 NULL。"""
|
|
stack: _ProviderStack | t.CPtr = _get_stack()
|
|
for i in range(stack.top - 1, -1, -1):
|
|
if stack.names[i] == field_name:
|
|
return stack.ptrs[i]
|
|
return None
|