[Wasm] Use a lazy restore frame when returning from tail calls
Component: JSC | 5223ee5
JSTests/wasm/stress/tail-call-cross-instance-gc.js
각 WebAssembly instance는 memory base, memory bounds, instance pointer 같은 instance 고유 데이터를 전용 register에 고정해 둡니다. Instance 경계를 넘는 tail call, 예를 들어 다른 module로 향하는 return_call_indirect가 발생하면, callee의 instance가 이 register들을 덮어쓸 수 있습니다. 이 chain이 결국 원래의 non-tail caller로 돌아온다면, 그 전에 register들이 복원되어야 합니다.
기존 방식은 이 문제를 컴파일 시점에 해결했습니다. computeTransitiveTailCalls와 callCanClobberInstance라는 transitive analysis를 통해 instance를 transitive하게 clobber할 수 있는 함수를 표시하고, 그 caller들이 모든 call 이후에 register를 즉시(eagerly) 복원하도록 했습니다. 다만 이 방식은 보수적이었고, inlining을 가로막는 문제가 있었습니다.
이번 commit은 이를 runtime 메커니즘으로 대체합니다. Instance 경계를 넘는 첫 tail call이 발생하는 시점에, caller의 argument 영역 바로 위에 32바이트짜리 restore frame이 lazy하게 삽입됩니다. 이 frame은 caller의 instance pointer와 원래 return address를 저장하며, 그 아래의 전체 frame은 32바이트만큼 아래로 밀립니다. 전용 thunk인 _wasm_restore_frame_return이 return 시점에 instance와 memory register를 다시 로드합니다. 이 패치는 여기에 더해 ARM64E PAC re-signing과 JIT cage gate thunk를 추가했으며, 기존 analysis와 그로 인한 inlining 제약도 함께 제거했습니다.
Before (compile-time analysis):
Caller (inst A) Callee (inst B) Return
│ │ │
├─[marked:clobbering]────────►│ return ───────────────►│
│ restore regs eagerly (inlining of B blocked)
│ after every call
After (runtime restore frame, lazy insertion):
Caller frame (inst A)
│ arg area
│ [32-byte restore frame inserted above args, first cross-inst call only]
│ slot 0: original CallerFrameAndPC
│ slot 1: saved instance A
│ slot 2: RestoreFrameCallee
│ callee-saves ← entire frame shifted down 32 bytes
│
├─ return_call ──► Callee (inst B) ──► ... ──► return
│ │
│ new cfr → _wasm_restore_frame_return
│ reload inst A regs
│ jump → original retPC
│
Reuse check: if retPC == thunk addr, skip insertion (no accumulation)
Significance
Instance 경계를 넘는 tail call의 ABI correctness가 static analysis에서 새로운 runtime frame 타입으로 옮겨갔습니다. 이제는 GC, stack unwinder, sampling profiler, 그리고 세 가지 Wasm backend 모두가 이 frame을 일관되게 처리해야 합니다. 그 대가로, instance 경계를 넘는 tail call chain에 관여하는 Wasm 함수들에서 보수적인 inlining barrier가 사라졌습니다. 스펙이 요구하는 "no accumulation" 속성은 이제 compile-time 속성이 아니라, 현재 ReturnPC와 thunk 주소를 비교하는 runtime reuse check에 의존하게 되었습니다.
Audit directions
점검할 만한 지점은 크게 여섯 곳입니다.
Stack frame shifting: restore frame은 caller의 전체 frame(callee-save 포함)을 물리적으로 32바이트 아래로 밀어내는 방식으로 삽입됩니다. emitRestoreInstanceFrameIfNeeded의 copy loop는 stack에 위치한 source가 있을 때는 do-while을 사용하고, 필요 없을 때는 frameSize 파라미터를 통해 이를 건너뜁니다. Copy 개수나 방향에서 off-by-one이 발생하면, 즉각적인 crash 없이 callee-save register나 인접 데이터가 조용히 손상될 수 있습니다.
ARM64E PAC re-signing: 원래의 return PC는 caller frame의 PAC signing context에서 restore frame의 context로 다시 서명되어야 하며, !Options::allowNonSPTagging()은 별도의 special case로 명시적으로 처리됩니다. 어떤 context나 key를 사용하는지에서 실수가 생기거나 non-SP-tagging 경로에서 edge case가 누락되면, 위조 가능한 return address나 PAC bypass로 이어질 가능성이 있습니다.
Reuse check correctness: no-accumulation invariant는 전적으로 ReturnPC == thunk-address check에 의존하며, 이 check는 BBQ, OMG, IPInt 사이에서, 그리고 assembly thunk와 JIT cage gate thunk 사이에서 서로 일치해야 합니다. 만약 정상적인 return address가 thunk 주소와 우연히 충돌하거나, gate thunk의 주소가 check가 비교하는 대상과 다르다면, restore frame이 생략되어 잘못된 instance가 복원되거나, 반대로 중복 삽입되어 스펙을 위반하는 stack 증가가 발생할 수 있습니다.
GC and unwinder correctness: 이것은 새로운 frame 타입입니다. GC는 이 frame의 Callee(RestoreInstanceCallee singleton)와 CodeBlock(wasmInstance pointer) slot을 scan해야 하며, stack unwinder와 sampling profiler도 이 frame을 인식하고 순회할 수 있어야 합니다. 이 중 어느 하나라도 실패하면, 그것이 GC 압력 상황에서 간헐적으로만 발생하더라도 use-after-free나 잘못된 profiler symbolication으로 이어질 수 있습니다. 전용 stress test인 tail-call-cross-instance-gc.js가 존재한다는 사실은, 작성자들이 이 지점을 가장 위험한 부분으로 판단했음을 시사합니다.
IPInt vs. BBQ/OMG ABI consistency: IPInt(WasmIPIntGenerator를 통해)와 BBQ/OMG JIT 양쪽 모두 restore frame을 구현하고 있으며, frame layout, slot offset, 내장된 thunk 주소 중 어느 하나라도 차이가 있으면 reuse check가 무효화되거나 thunk가 잘못된 slot에서 reload하게 됩니다.
마지막으로 JIT cage gate thunk는 JIT로 컴파일된 코드가 일관된 진입점을 사용하도록 등록된 것으로, JIT-key로 서명된 return PC를 올바르게 untag해야 합니다. _wasm_restore_frame_return의 로직과 조금이라도 어긋나면, 이는 security boundary의 불일치이며 untagging 과정을 혼란시키는 데 악용될 가능성이 있습니다. 앞으로도 유의해야 할 일반적인 패턴은 다음과 같습니다. Runtime에 삽입되는 frame 타입은 모든 backend, GC, unwinder가 동시에 동일하게 이해하고 있어야 하며, 이 네 곳 각각이 독립적으로 점검해야 할 지점이 됩니다.