Skip to content

Commit 891e470

Browse files
mathetakemathetake
authored
compiler(arm64): fixes overflow in huge executable relocations (#2181)
Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
1 parent f31be30 commit 891e470

18 files changed

+302
-57
lines changed

‎internal/engine/wazevo/backend/backend_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2375,7 +2375,8 @@ L2 (SSA Block: blk2):
23752375
fmt.Println(be.Format())
23762376
}
23772377

2378-
be.Finalize(context.Background())
2378+
err = be.Finalize(context.Background())
2379+
require.NoError(t, err)
23792380
if verbose {
23802381
fmt.Println("============ finalization result ============")
23812382
fmt.Println(be.Format())

‎internal/engine/wazevo/backend/compiler.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ type Compiler interface {
5252

5353
// Finalize performs the finalization of the compilation, including machine code emission.
5454
// This must be called after RegAlloc.
55-
Finalize(ctx context.Context)
55+
Finalize(ctx context.Context) error
5656

5757
// Buf returns the buffer of the encoded machine code. This is only used for testing purpose.
5858
Buf() []byte
@@ -175,7 +175,9 @@ func (c *compiler) Compile(ctx context.Context) ([]byte, []RelocationInfo, error
175175
if wazevoapi.DeterministicCompilationVerifierEnabled {
176176
wazevoapi.VerifyOrSetDeterministicCompilationContextValue(ctx, "After Register Allocation", c.Format())
177177
}
178-
c.Finalize(ctx)
178+
if err := c.Finalize(ctx); err != nil {
179+
return nil, nil, err
180+
}
179181
if wazevoapi.PrintFinalizedMachineCode && wazevoapi.PrintEnabledIndex(ctx) {
180182
fmt.Printf("[[[after finalize for %s]]]%s\n", wazevoapi.GetCurrentFunctionName(ctx), c.Format())
181183
}
@@ -191,9 +193,9 @@ func (c *compiler) RegAlloc() {
191193
}
192194

193195
// Finalize implements Compiler.Finalize.
194-
func (c *compiler) Finalize(ctx context.Context) {
196+
func (c *compiler) Finalize(ctx context.Context) error {
195197
c.mach.PostRegAlloc()
196-
c.mach.Encode(ctx)
198+
return c.mach.Encode(ctx)
197199
}
198200

199201
// setCurrentGroupID sets the current instruction group ID.

‎internal/engine/wazevo/backend/isa/amd64/abi_go_call_test.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -370,10 +370,9 @@ L3:
370370
t.Run(tc.name, func(t *testing.T) {
371371
_, _, m := newSetupWithMockContext()
372372
m.CompileGoFunctionTrampoline(tc.exitCode, tc.sig, tc.needModuleContextPtr)
373-
374373
require.Equal(t, tc.exp, m.Format())
375-
376-
m.Encode(context.Background())
374+
err := m.Encode(context.Background())
375+
require.NoError(t, err)
377376
})
378377
}
379378
}
@@ -521,7 +520,8 @@ L2:
521520
_, _, m := newSetupWithMockContext()
522521
m.ectx.RootInstr = m.allocateNop()
523522
m.insertStackBoundsCheck(tc.requiredStackSize, m.ectx.RootInstr)
524-
m.Encode(context.Background())
523+
err := m.Encode(context.Background())
524+
require.NoError(t, err)
525525
require.Equal(t, tc.exp, m.Format())
526526
})
527527
}

‎internal/engine/wazevo/backend/isa/amd64/machine.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1963,7 +1963,7 @@ func (m *machine) encodeWithoutSSA(root *instruction) {
19631963
}
19641964

19651965
// Encode implements backend.Machine Encode.
1966-
func (m *machine) Encode(ctx context.Context) {
1966+
func (m *machine) Encode(ctx context.Context) (err error) {
19671967
ectx := m.ectx
19681968
bufPtr := m.c.BufPtr()
19691969

@@ -2055,10 +2055,11 @@ func (m *machine) Encode(ctx context.Context) {
20552055
panic("BUG")
20562056
}
20572057
}
2058+
return
20582059
}
20592060

20602061
// ResolveRelocations implements backend.Machine.
2061-
func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []backend.RelocationInfo) {
2062+
func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []backend.RelocationInfo, _ []int) {
20622063
for _, r := range relocations {
20632064
offset := r.Offset
20642065
calleeFnOffset := refToBinaryOffset[r.FuncRef]
@@ -2072,6 +2073,9 @@ func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, bina
20722073
}
20732074
}
20742075

2076+
// CallTrampolineIslandInfo implements backend.Machine CallTrampolineIslandInfo.
2077+
func (m *machine) CallTrampolineIslandInfo(_ int) (_, _ int, _ error) { return }
2078+
20752079
func (m *machine) lowerIcmpToFlag(xd, yd *backend.SSAValueDefinition, _64 bool) {
20762080
x := m.getOperand_Reg(xd)
20772081
y := m.getOperand_Mem_Imm32_Reg(yd)

‎internal/engine/wazevo/backend/isa/amd64/machine_pro_epi_logue_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ func TestMachine_setupPrologue(t *testing.T) {
7777

7878
m.setupPrologue()
7979
require.Equal(t, root, m.ectx.RootInstr)
80-
m.Encode(context.Background())
80+
err := m.Encode(context.Background())
81+
require.NoError(t, err)
8182
require.Equal(t, tc.exp, m.Format())
8283
})
8384
}
@@ -151,7 +152,8 @@ func TestMachine_postRegAlloc(t *testing.T) {
151152
m.postRegAlloc()
152153

153154
require.Equal(t, root, m.ectx.RootInstr)
154-
m.Encode(context.Background())
155+
err := m.Encode(context.Background())
156+
require.NoError(t, err)
155157
require.Equal(t, tc.exp, m.Format())
156158
})
157159
}

‎internal/engine/wazevo/backend/isa/amd64/util_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,11 @@ func (m *mockCompiler) Buf() []byte { return m.buf }
6767
func (m *mockCompiler) TypeOf(v regalloc.VReg) (ret ssa.Type) {
6868
return m.typeOf[v.ID()]
6969
}
70-
func (m *mockCompiler) Finalize(context.Context) {}
71-
func (m *mockCompiler) RegAlloc() {}
72-
func (m *mockCompiler) Lower() {}
73-
func (m *mockCompiler) Format() string { return "" }
74-
func (m *mockCompiler) Init() {}
70+
func (m *mockCompiler) Finalize(context.Context) (err error) { return }
71+
func (m *mockCompiler) RegAlloc() {}
72+
func (m *mockCompiler) Lower() {}
73+
func (m *mockCompiler) Format() string { return "" }
74+
func (m *mockCompiler) Init() {}
7575

7676
func newMockCompilationContext() *mockCompiler { //nolint
7777
return &mockCompiler{

‎internal/engine/wazevo/backend/isa/arm64/abi_entry_preamble_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -535,7 +535,8 @@ func TestMachine_goEntryPreamblePassArg(t *testing.T) {
535535
m.executableContext.RootInstr = cur
536536
m.goEntryPreamblePassArg(cur, paramSlicePtr, &tc.arg, tc.argSlotBeginOffsetFromSP)
537537
require.Equal(t, tc.exp, m.Format())
538-
m.Encode(context.Background())
538+
err := m.Encode(context.Background())
539+
require.NoError(t, err)
539540
})
540541
}
541542
}
@@ -688,7 +689,8 @@ func TestMachine_goEntryPreamblePassResult(t *testing.T) {
688689
m.executableContext.RootInstr = cur
689690
m.goEntryPreamblePassResult(cur, paramSlicePtr, &tc.arg, tc.retStart)
690691
require.Equal(t, tc.exp, m.Format())
691-
m.Encode(context.Background())
692+
err := m.Encode(context.Background())
693+
require.NoError(t, err)
692694
})
693695
}
694696
}

‎internal/engine/wazevo/backend/isa/arm64/abi_go_call_test.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,8 @@ func TestMachine_CompileGoFunctionTrampoline(t *testing.T) {
451451

452452
require.Equal(t, tc.exp, m.Format())
453453

454-
m.Encode(context.Background())
454+
err := m.Encode(context.Background())
455+
require.NoError(t, err)
455456
})
456457
}
457458
}
@@ -519,7 +520,8 @@ func Test_goFunctionCallLoadStackArg(t *testing.T) {
519520
m.executableContext.RootInstr = nop
520521

521522
require.Equal(t, tc.exp, m.Format())
522-
m.Encode(context.Background())
523+
err := m.Encode(context.Background())
524+
require.NoError(t, err)
523525
})
524526
})
525527
}
@@ -585,7 +587,8 @@ func Test_goFunctionCallStoreStackResult(t *testing.T) {
585587
m.executableContext.RootInstr = nop
586588

587589
require.Equal(t, tc.exp, m.Format())
588-
m.Encode(context.Background())
590+
err := m.Encode(context.Background())
591+
require.NoError(t, err)
589592
})
590593
})
591594
}

‎internal/engine/wazevo/backend/isa/arm64/instr_encoding.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,13 @@ import (
1010
)
1111

1212
// Encode implements backend.Machine Encode.
13-
func (m *machine) Encode(ctx context.Context) {
13+
func (m *machine) Encode(ctx context.Context) error {
1414
m.resolveRelativeAddresses(ctx)
1515
m.encode(m.executableContext.RootInstr)
16+
if l := len(m.compiler.Buf()); l > maxFunctionExecutableSize {
17+
return fmt.Errorf("function size exceeds the limit: %d > %d", l, maxFunctionExecutableSize)
18+
}
19+
return nil
1620
}
1721

1822
func (m *machine) encode(root *instruction) {

‎internal/engine/wazevo/backend/isa/arm64/machine.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -388,11 +388,11 @@ func (m *machine) resolveRelativeAddresses(ctx context.Context) {
388388
}
389389

390390
const (
391-
maxSignedInt26 int64 = 1<<25 - 1
392-
minSignedInt26 int64 = -(1 << 25)
391+
maxSignedInt26 = 1<<25 - 1
392+
minSignedInt26 = -(1 << 25)
393393

394-
maxSignedInt19 int64 = 1<<19 - 1
395-
minSignedInt19 int64 = -(1 << 19)
394+
maxSignedInt19 = 1<<19 - 1
395+
minSignedInt19 = -(1 << 19)
396396
)
397397

398398
func (m *machine) insertConditionalJumpTrampoline(cbr *instruction, currentBlk *labelPosition, nextLabel label) {

‎internal/engine/wazevo/backend/isa/arm64/machine_pro_epi_logue_test.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,8 @@ func TestMachine_setupPrologue(t *testing.T) {
110110

111111
m.setupPrologue()
112112
require.Equal(t, root, m.executableContext.RootInstr)
113-
m.Encode(context.Background())
113+
err := m.Encode(context.Background())
114+
require.NoError(t, err)
114115
require.Equal(t, tc.exp, m.Format())
115116
})
116117
}
@@ -223,7 +224,8 @@ func TestMachine_postRegAlloc(t *testing.T) {
223224
m.postRegAlloc()
224225

225226
require.Equal(t, root, m.executableContext.RootInstr)
226-
m.Encode(context.Background())
227+
err := m.Encode(context.Background())
228+
require.NoError(t, err)
227229
require.Equal(t, tc.exp, m.Format())
228230
})
229231
}
@@ -268,7 +270,8 @@ func TestMachine_insertStackBoundsCheck(t *testing.T) {
268270
m.executableContext.RootInstr = m.allocateInstr()
269271
m.executableContext.RootInstr.asNop0()
270272
m.insertStackBoundsCheck(tc.requiredStackSize, m.executableContext.RootInstr)
271-
m.Encode(context.Background())
273+
err := m.Encode(context.Background())
274+
require.NoError(t, err)
272275
require.Equal(t, tc.exp, m.Format())
273276
})
274277
}

‎internal/engine/wazevo/backend/isa/arm64/machine_relocation.go

Lines changed: 99 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,118 @@
11
package arm64
22

33
import (
4+
"encoding/binary"
45
"fmt"
6+
"math"
7+
"sort"
58

69
"github.com/tetratelabs/wazero/internal/engine/wazevo/backend"
710
"github.com/tetratelabs/wazero/internal/engine/wazevo/ssa"
811
)
912

13+
const (
14+
// trampolineCallSize is the size of the trampoline instruction sequence for each function in an island.
15+
trampolineCallSize = 4*4 + 4 // Four instructions + 32-bit immediate.
16+
17+
// Unconditional branch offset is encoded as divided by 4 in imm26.
18+
// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/BL--Branch-with-Link-?lang=en
19+
20+
maxUnconditionalBranchOffset = maxSignedInt26 * 4
21+
minUnconditionalBranchOffset = minSignedInt26 * 4
22+
23+
// trampolineIslandInterval is the range of the trampoline island.
24+
// Half of the range is used for the trampoline island, and the other half is used for the function.
25+
trampolineIslandInterval = maxUnconditionalBranchOffset / 2
26+
27+
// maxNumFunctions explicitly specifies the maximum number of functions that can be allowed in a single executable.
28+
maxNumFunctions = trampolineIslandInterval >> 6
29+
30+
// maxFunctionExecutableSize is the maximum size of a function that can exist in a trampoline island.
31+
// Conservatively set to 1/4 of the trampoline island interval.
32+
maxFunctionExecutableSize = trampolineIslandInterval >> 2
33+
)
34+
35+
// CallTrampolineIslandInfo implements backend.Machine CallTrampolineIslandInfo.
36+
func (m *machine) CallTrampolineIslandInfo(numFunctions int) (interval, size int, err error) {
37+
if numFunctions > maxNumFunctions {
38+
return 0, 0, fmt.Errorf("too many functions: %d > %d", numFunctions, maxNumFunctions)
39+
}
40+
return trampolineIslandInterval, trampolineCallSize * numFunctions, nil
41+
}
42+
1043
// ResolveRelocations implements backend.Machine ResolveRelocations.
11-
//
12-
// TODO: unit test!
13-
func (m *machine) ResolveRelocations(refToBinaryOffset map[ssa.FuncRef]int, binary []byte, relocations []backend.RelocationInfo) {
44+
func (m *machine) ResolveRelocations(
45+
refToBinaryOffset map[ssa.FuncRef]int,
46+
executable []byte,
47+
relocations []backend.RelocationInfo,
48+
callTrampolineIslandOffsets []int,
49+
) {
50+
for _, islandOffset := range callTrampolineIslandOffsets {
51+
encodeCallTrampolineIsland(refToBinaryOffset, islandOffset, executable)
52+
}
53+
1454
for _, r := range relocations {
1555
instrOffset := r.Offset
1656
calleeFnOffset := refToBinaryOffset[r.FuncRef]
17-
brInstr := binary[instrOffset : instrOffset+4]
1857
diff := int64(calleeFnOffset) - (instrOffset)
1958
// Check if the diff is within the range of the branch instruction.
20-
if diff < -(1<<25)*4 || diff > ((1<<25)-1)*4 {
21-
panic(fmt.Sprintf("TODO: too large binary where branch target is out of the supported range +/-128MB: %#x", diff))
59+
if diff < minUnconditionalBranchOffset || diff > maxUnconditionalBranchOffset {
60+
// Find the near trampoline island from callTrampolineIslandOffsets.
61+
islandOffset := searchTrampolineIsland(callTrampolineIslandOffsets, int(instrOffset))
62+
islandTargetOffset := islandOffset + trampolineCallSize*int(r.FuncRef)
63+
diff = int64(islandTargetOffset) - (instrOffset)
64+
if diff < minUnconditionalBranchOffset || diff > maxUnconditionalBranchOffset {
65+
panic("BUG in trampoline placement")
66+
}
2267
}
23-
// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/BL--Branch-with-Link-
24-
imm26 := diff / 4
25-
brInstr[0] = byte(imm26)
26-
brInstr[1] = byte(imm26 >> 8)
27-
brInstr[2] = byte(imm26 >> 16)
28-
if diff < 0 {
29-
brInstr[3] = (byte(imm26 >> 24 & 0b000000_01)) | 0b100101_10 // Set sign bit.
30-
} else {
31-
brInstr[3] = (byte(imm26 >> 24 & 0b000000_01)) | 0b100101_00 // No sign bit.
68+
binary.LittleEndian.PutUint32(executable[instrOffset:instrOffset+4], encodeUnconditionalBranch(true, diff))
69+
}
70+
}
71+
72+
// encodeCallTrampolineIsland encodes a trampoline island for the given functions.
73+
// Each island consists of a trampoline instruction sequence for each function.
74+
// Each trampoline instruction sequence consists of 4 instructions + 32-bit immediate.
75+
func encodeCallTrampolineIsland(refToBinaryOffset map[ssa.FuncRef]int, islandOffset int, executable []byte) {
76+
for i := 0; i < len(refToBinaryOffset); i++ {
77+
trampolineOffset := islandOffset + trampolineCallSize*i
78+
79+
fnOffset := refToBinaryOffset[ssa.FuncRef(i)]
80+
diff := fnOffset - (trampolineOffset + 16)
81+
if diff > math.MaxInt32 || diff < math.MinInt32 {
82+
// This case even amd64 can't handle. 4GB is too big.
83+
panic("too big binary")
3284
}
85+
86+
// The tmpReg, tmpReg2 is safe to overwrite (in fact any caller-saved register is safe to use).
87+
tmpReg, tmpReg2 := regNumberInEncoding[tmpRegVReg.RealReg()], regNumberInEncoding[x11]
88+
89+
// adr tmpReg, PC+16: load the address of #diff into tmpReg.
90+
binary.LittleEndian.PutUint32(executable[trampolineOffset:], encodeAdr(tmpReg, 16))
91+
// ldrsw tmpReg2, [tmpReg]: Load #diff into tmpReg2.
92+
binary.LittleEndian.PutUint32(executable[trampolineOffset+4:],
93+
encodeLoadOrStore(sLoad32, tmpReg2, addressMode{kind: addressModeKindRegUnsignedImm12, rn: tmpRegVReg}))
94+
// add tmpReg, tmpReg2, tmpReg: add #diff to the address of #diff, getting the absolute address of the function.
95+
binary.LittleEndian.PutUint32(executable[trampolineOffset+8:],
96+
encodeAluRRR(aluOpAdd, tmpReg, tmpReg, tmpReg2, true, false))
97+
// br tmpReg: branch to the function without overwriting the link register.
98+
binary.LittleEndian.PutUint32(executable[trampolineOffset+12:], encodeUnconditionalBranchReg(tmpReg, false))
99+
// #diff
100+
binary.LittleEndian.PutUint32(executable[trampolineOffset+16:], uint32(diff))
101+
}
102+
}
103+
104+
// searchTrampolineIsland finds the nearest trampoline island from callTrampolineIslandOffsets.
105+
// Note that even if the offset is in the middle of two islands, it returns the latter one.
106+
// That is ok because the island is always placed in the middle of the range.
107+
//
108+
// precondition: callTrampolineIslandOffsets is sorted in ascending order.
109+
func searchTrampolineIsland(callTrampolineIslandOffsets []int, offset int) int {
110+
l := len(callTrampolineIslandOffsets)
111+
n := sort.Search(l, func(i int) bool {
112+
return callTrampolineIslandOffsets[i] >= offset
113+
})
114+
if n == l {
115+
n = l - 1
33116
}
117+
return callTrampolineIslandOffsets[n]
34118
}

0 commit comments

Comments
 (0)