Skip to content

Commit c9b7d21

Browse files
oskarwirgaoskarwirga
authored
[CFI/MergeFunctions] Modify MergeFunctions to propagate type information (#68628)
When MergeFuncs creates a thunk, it does not modify the function in place, but creates a new one altogether. If type metadata is not properly forwarded to this new function, LowerTypeTests will be unable to put this thunk into the dispatch table. The fix here is to just forward the type metadata to the newly created functions.
1 parent 10079a2 commit c9b7d21

File tree

2 files changed

+223
-0
lines changed

2 files changed

+223
-0
lines changed

‎llvm/lib/Transforms/IPO/MergeFunctions.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -662,6 +662,13 @@ static bool canCreateThunkFor(Function *F) {
662662
return true;
663663
}
664664

665+
/// Copy metadata from one function to another.
666+
static void copyMetadataIfPresent(Function *From, Function *To, StringRef Key) {
667+
if (MDNode *MD = From->getMetadata(Key)) {
668+
To->setMetadata(Key, MD);
669+
}
670+
}
671+
665672
// Replace G with a simple tail call to bitcast(F). Also (unless
666673
// MergeFunctionsPDI holds) replace direct uses of G with bitcast(F),
667674
// delete G. Under MergeFunctionsPDI, we use G itself for creating
@@ -740,6 +747,9 @@ void MergeFunctions::writeThunk(Function *F, Function *G) {
740747
} else {
741748
NewG->copyAttributesFrom(G);
742749
NewG->takeName(G);
750+
// Ensure CFI type metadata is propagated to the new function.
751+
copyMetadataIfPresent(G, NewG, "type");
752+
copyMetadataIfPresent(G, NewG, "kcfi_type");
743753
removeUsers(G);
744754
G->replaceAllUsesWith(NewG);
745755
G->eraseFromParent();
@@ -815,6 +825,9 @@ void MergeFunctions::mergeTwoFunctions(Function *F, Function *G) {
815825
F->getAddressSpace(), "", F->getParent());
816826
NewF->copyAttributesFrom(F);
817827
NewF->takeName(F);
828+
// Ensure CFI type metadata is propagated to the new function.
829+
copyMetadataIfPresent(F, NewF, "type");
830+
copyMetadataIfPresent(F, NewF, "kcfi_type");
818831
removeUsers(F);
819832
F->replaceAllUsesWith(NewF);
820833

‎llvm/test/Transforms/MergeFunc/cfi-thunk-merging.ll

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --include-generated-funcs --version 2
2+
;; Check the cases involving internal CFI instrumented functions where we do not expect functions to be merged.
3+
; RUN: opt -S -passes=mergefunc < %s | FileCheck %s
4+
; RUN: opt -S -passes=mergefunc,lowertypetests < %s | FileCheck --check-prefix=LOWERTYPETESTS %s
5+
6+
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
7+
target triple = "x86_64-unknown-linux"
8+
9+
@0 = private unnamed_addr constant { i16, i16, [12 x i8] } { i16 -1, i16 0, [12 x i8] c"'int (int)'\00" }
10+
11+
; Function Attrs: noinline nounwind optnone
12+
define dso_local i32 @f(i32 noundef %arg) #0 !type !3 !type !4 {
13+
entry:
14+
%arg.addr = alloca i32, align 4
15+
%a = alloca i32, align 4
16+
%b = alloca i32, align 4
17+
store i32 %arg, ptr %arg.addr, align 4
18+
store i32 0, ptr %b, align 4
19+
%0 = load i32, ptr %arg.addr, align 4
20+
%cmp = icmp sgt i32 %0, 0
21+
br i1 %cmp, label %if.then, label %if.end
22+
23+
if.then: ; preds = %entry
24+
store i32 1, ptr %a, align 4
25+
br label %if.end
26+
27+
if.end: ; preds = %if.then, %entry
28+
%1 = load i32, ptr %a, align 4
29+
%2 = load i32, ptr %b, align 4
30+
%add = add nsw i32 %1, %2
31+
ret i32 %add
32+
}
33+
34+
; Function Attrs: noinline nounwind optnone
35+
define dso_local i32 @f_thunk(i32 noundef %arg) #0 !type !3 !type !4 {
36+
entry:
37+
%arg.addr = alloca i32, align 4
38+
%a = alloca i32, align 4
39+
%b = alloca i32, align 4
40+
store i32 %arg, ptr %arg.addr, align 4
41+
store i32 0, ptr %b, align 4
42+
%0 = load i32, ptr %arg.addr, align 4
43+
%cmp = icmp sgt i32 %0, 0
44+
br i1 %cmp, label %if.then, label %if.end
45+
46+
if.then: ; preds = %entry
47+
store i32 1, ptr %a, align 4
48+
br label %if.end
49+
50+
if.end: ; preds = %if.then, %entry
51+
%1 = load i32, ptr %a, align 4
52+
%2 = load i32, ptr %b, align 4
53+
%add = add nsw i32 %1, %2
54+
ret i32 %add
55+
}
56+
57+
; Function Attrs: noinline nounwind optnone
58+
define dso_local i32 @g(i32 noundef %b) #0 !type !3 !type !4 {
59+
entry:
60+
%b.addr = alloca i32, align 4
61+
%fp = alloca ptr, align 8
62+
store i32 %b, ptr %b.addr, align 4
63+
%0 = load i32, ptr %b.addr, align 4
64+
%tobool = icmp ne i32 %0, 0
65+
%1 = zext i1 %tobool to i64
66+
%cond = select i1 %tobool, ptr @f, ptr @f_thunk
67+
store ptr %cond, ptr %fp, align 8
68+
%2 = load ptr, ptr %fp, align 8
69+
%3 = call i1 @llvm.type.test(ptr %2, metadata !"_ZTSFiiE"), !nosanitize !5
70+
br i1 %3, label %cont, label %trap, !nosanitize !5
71+
72+
trap: ; preds = %entry
73+
call void @llvm.ubsantrap(i8 2) #3, !nosanitize !5
74+
unreachable, !nosanitize !5
75+
76+
cont: ; preds = %entry
77+
%4 = load i32, ptr %b.addr, align 4
78+
%call = call i32 %2(i32 noundef %4)
79+
ret i32 %call
80+
}
81+
82+
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
83+
declare i1 @llvm.type.test(ptr, metadata) #1
84+
85+
; Function Attrs: cold noreturn nounwind
86+
declare void @llvm.ubsantrap(i8 immarg) #2
87+
88+
attributes #0 = { noinline nounwind optnone "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-features"="+cx8,+mmx,+sse,+sse2,+x87" }
89+
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
90+
attributes #2 = { cold noreturn nounwind }
91+
attributes #3 = { noreturn nounwind }
92+
93+
!llvm.module.flags = !{!0, !1}
94+
95+
!0 = !{i32 1, !"wchar_size", i32 4}
96+
!1 = !{i32 4, !"CFI Canonical Jump Tables", i32 0}
97+
!3 = !{i64 0, !"_ZTSFiiE"}
98+
!4 = !{i64 0, !"_ZTSFiiE.generalized"}
99+
!5 = !{}
100+
; CHECK-LABEL: define dso_local i32 @f
101+
; CHECK-SAME: (i32 noundef [[ARG:%.*]]) #[[ATTR0:[0-9]+]] !type !2 !type !3 {
102+
; CHECK-NEXT: entry:
103+
; CHECK-NEXT: [[ARG_ADDR:%.*]] = alloca i32, align 4
104+
; CHECK-NEXT: [[A:%.*]] = alloca i32, align 4
105+
; CHECK-NEXT: [[B:%.*]] = alloca i32, align 4
106+
; CHECK-NEXT: store i32 [[ARG]], ptr [[ARG_ADDR]], align 4
107+
; CHECK-NEXT: store i32 0, ptr [[B]], align 4
108+
; CHECK-NEXT: [[TMP0:%.*]] = load i32, ptr [[ARG_ADDR]], align 4
109+
; CHECK-NEXT: [[CMP:%.*]] = icmp sgt i32 [[TMP0]], 0
110+
; CHECK-NEXT: br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
111+
; CHECK: if.then:
112+
; CHECK-NEXT: store i32 1, ptr [[A]], align 4
113+
; CHECK-NEXT: br label [[IF_END]]
114+
; CHECK: if.end:
115+
; CHECK-NEXT: [[TMP1:%.*]] = load i32, ptr [[A]], align 4
116+
; CHECK-NEXT: [[TMP2:%.*]] = load i32, ptr [[B]], align 4
117+
; CHECK-NEXT: [[ADD:%.*]] = add nsw i32 [[TMP1]], [[TMP2]]
118+
; CHECK-NEXT: ret i32 [[ADD]]
119+
;
120+
;
121+
; CHECK-LABEL: define dso_local i32 @g
122+
; CHECK-SAME: (i32 noundef [[B:%.*]]) #[[ATTR0]] !type !2 !type !3 {
123+
; CHECK-NEXT: entry:
124+
; CHECK-NEXT: [[B_ADDR:%.*]] = alloca i32, align 4
125+
; CHECK-NEXT: [[FP:%.*]] = alloca ptr, align 8
126+
; CHECK-NEXT: store i32 [[B]], ptr [[B_ADDR]], align 4
127+
; CHECK-NEXT: [[TMP0:%.*]] = load i32, ptr [[B_ADDR]], align 4
128+
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp ne i32 [[TMP0]], 0
129+
; CHECK-NEXT: [[TMP1:%.*]] = zext i1 [[TOBOOL]] to i64
130+
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], ptr @f, ptr @f_thunk
131+
; CHECK-NEXT: store ptr [[COND]], ptr [[FP]], align 8
132+
; CHECK-NEXT: [[TMP2:%.*]] = load ptr, ptr [[FP]], align 8
133+
; CHECK-NEXT: [[TMP3:%.*]] = call i1 @llvm.type.test(ptr [[TMP2]], metadata !"_ZTSFiiE"), !nosanitize !4
134+
; CHECK-NEXT: br i1 [[TMP3]], label [[CONT:%.*]], label [[TRAP:%.*]], !nosanitize !4
135+
; CHECK: trap:
136+
; CHECK-NEXT: call void @llvm.ubsantrap(i8 2) #[[ATTR3:[0-9]+]], !nosanitize !4
137+
; CHECK-NEXT: unreachable, !nosanitize !4
138+
; CHECK: cont:
139+
; CHECK-NEXT: [[TMP4:%.*]] = load i32, ptr [[B_ADDR]], align 4
140+
; CHECK-NEXT: [[CALL:%.*]] = call i32 [[TMP2]](i32 noundef [[TMP4]])
141+
; CHECK-NEXT: ret i32 [[CALL]]
142+
;
143+
;
144+
; CHECK-LABEL: define dso_local i32 @f_thunk
145+
; CHECK-SAME: (i32 noundef [[TMP0:%.*]]) #[[ATTR0]] !type !2 {
146+
; CHECK-NEXT: [[TMP2:%.*]] = tail call i32 @f(i32 noundef [[TMP0]]) #[[ATTR0]]
147+
; CHECK-NEXT: ret i32 [[TMP2]]
148+
;
149+
;
150+
; LOWERTYPETESTS-LABEL: define dso_local i32 @f
151+
; LOWERTYPETESTS-SAME: (i32 noundef [[ARG:%.*]]) #[[ATTR0:[0-9]+]] !type !2 !type !3 {
152+
; LOWERTYPETESTS-NEXT: entry:
153+
; LOWERTYPETESTS-NEXT: [[ARG_ADDR:%.*]] = alloca i32, align 4
154+
; LOWERTYPETESTS-NEXT: [[A:%.*]] = alloca i32, align 4
155+
; LOWERTYPETESTS-NEXT: [[B:%.*]] = alloca i32, align 4
156+
; LOWERTYPETESTS-NEXT: store i32 [[ARG]], ptr [[ARG_ADDR]], align 4
157+
; LOWERTYPETESTS-NEXT: store i32 0, ptr [[B]], align 4
158+
; LOWERTYPETESTS-NEXT: [[TMP0:%.*]] = load i32, ptr [[ARG_ADDR]], align 4
159+
; LOWERTYPETESTS-NEXT: [[CMP:%.*]] = icmp sgt i32 [[TMP0]], 0
160+
; LOWERTYPETESTS-NEXT: br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
161+
; LOWERTYPETESTS: if.then:
162+
; LOWERTYPETESTS-NEXT: store i32 1, ptr [[A]], align 4
163+
; LOWERTYPETESTS-NEXT: br label [[IF_END]]
164+
; LOWERTYPETESTS: if.end:
165+
; LOWERTYPETESTS-NEXT: [[TMP1:%.*]] = load i32, ptr [[A]], align 4
166+
; LOWERTYPETESTS-NEXT: [[TMP2:%.*]] = load i32, ptr [[B]], align 4
167+
; LOWERTYPETESTS-NEXT: [[ADD:%.*]] = add nsw i32 [[TMP1]], [[TMP2]]
168+
; LOWERTYPETESTS-NEXT: ret i32 [[ADD]]
169+
;
170+
;
171+
; LOWERTYPETESTS-LABEL: define dso_local i32 @g
172+
; LOWERTYPETESTS-SAME: (i32 noundef [[B:%.*]]) #[[ATTR0]] !type !2 !type !3 {
173+
; LOWERTYPETESTS-NEXT: entry:
174+
; LOWERTYPETESTS-NEXT: [[B_ADDR:%.*]] = alloca i32, align 4
175+
; LOWERTYPETESTS-NEXT: [[FP:%.*]] = alloca ptr, align 8
176+
; LOWERTYPETESTS-NEXT: store i32 [[B]], ptr [[B_ADDR]], align 4
177+
; LOWERTYPETESTS-NEXT: [[TMP0:%.*]] = load i32, ptr [[B_ADDR]], align 4
178+
; LOWERTYPETESTS-NEXT: [[TOBOOL:%.*]] = icmp ne i32 [[TMP0]], 0
179+
; LOWERTYPETESTS-NEXT: [[TMP1:%.*]] = zext i1 [[TOBOOL]] to i64
180+
; LOWERTYPETESTS-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], ptr @.cfi.jumptable, ptr getelementptr inbounds ([2 x [8 x i8]], ptr @.cfi.jumptable, i64 0, i64 1)
181+
; LOWERTYPETESTS-NEXT: store ptr [[COND]], ptr [[FP]], align 8
182+
; LOWERTYPETESTS-NEXT: [[TMP2:%.*]] = load ptr, ptr [[FP]], align 8
183+
; LOWERTYPETESTS-NEXT: [[TMP3:%.*]] = ptrtoint ptr [[TMP2]] to i64
184+
; LOWERTYPETESTS-NEXT: [[TMP4:%.*]] = sub i64 [[TMP3]], ptrtoint (ptr @.cfi.jumptable to i64)
185+
; LOWERTYPETESTS-NEXT: [[TMP5:%.*]] = lshr i64 [[TMP4]], 3
186+
; LOWERTYPETESTS-NEXT: [[TMP6:%.*]] = shl i64 [[TMP4]], 61
187+
; LOWERTYPETESTS-NEXT: [[TMP7:%.*]] = or i64 [[TMP5]], [[TMP6]]
188+
; LOWERTYPETESTS-NEXT: [[TMP8:%.*]] = icmp ule i64 [[TMP7]], 1
189+
; LOWERTYPETESTS-NEXT: br i1 [[TMP8]], label [[CONT:%.*]], label [[TRAP:%.*]], !nosanitize !4
190+
; LOWERTYPETESTS: trap:
191+
; LOWERTYPETESTS-NEXT: call void @llvm.ubsantrap(i8 2) #[[ATTR4:[0-9]+]], !nosanitize !4
192+
; LOWERTYPETESTS-NEXT: unreachable, !nosanitize !4
193+
; LOWERTYPETESTS: cont:
194+
; LOWERTYPETESTS-NEXT: [[TMP9:%.*]] = load i32, ptr [[B_ADDR]], align 4
195+
; LOWERTYPETESTS-NEXT: [[CALL:%.*]] = call i32 [[TMP2]](i32 noundef [[TMP9]])
196+
; LOWERTYPETESTS-NEXT: ret i32 [[CALL]]
197+
;
198+
;
199+
; LOWERTYPETESTS-LABEL: define dso_local i32 @f_thunk
200+
; LOWERTYPETESTS-SAME: (i32 noundef [[TMP0:%.*]]) #[[ATTR0]] !type !2 {
201+
; LOWERTYPETESTS-NEXT: [[TMP2:%.*]] = tail call i32 @f(i32 noundef [[TMP0]]) #[[ATTR0]]
202+
; LOWERTYPETESTS-NEXT: ret i32 [[TMP2]]
203+
;
204+
;
205+
; LOWERTYPETESTS-LABEL: define private void @.cfi.jumptable
206+
; LOWERTYPETESTS-SAME: () #[[ATTR3:[0-9]+]] align 8 {
207+
; LOWERTYPETESTS-NEXT: entry:
208+
; LOWERTYPETESTS-NEXT: call void asm sideeffect "jmp ${0:c}@plt\0Aint3\0Aint3\0Aint3\0Ajmp ${1:c}@plt\0Aint3\0Aint3\0Aint3\0A", "s,s"(ptr @f, ptr @f_thunk)
209+
; LOWERTYPETESTS-NEXT: unreachable
210+
;

0 commit comments

Comments
 (0)