Skip to content

Commit af8cf15

Browse files
Merge pull request #34928 from nate-chandler/concurrency/irgen/ptrauth
[Async CC] Pointer authentication.
2 parents e365fc5 + 84e7ba7 commit af8cf15

File tree

49 files changed

+117
-85
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+117
-85
lines changed

include/swift/AST/IRGenOptions.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,9 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
140140

141141
/// The parent async context stored within a child async context.
142142
PointerAuthSchema AsyncContextParent;
143+
144+
/// The function to call to resume running in the parent context.
145+
PointerAuthSchema AsyncContextResume;
143146
};
144147

145148
enum class JITDebugArtifact : unsigned {

lib/IRGen/GenCall.cpp

Lines changed: 77 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1867,12 +1867,15 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
18671867
FunctionPointer functionPointer, llvm::Value *thickContext,
18681868
std::pair<bool, bool> values, Size initialContextSize) {
18691869
assert(values.first || values.second);
1870+
assert(functionPointer.getKind() ==
1871+
FunctionPointer::KindTy::AsyncFunctionPointer);
18701872
bool emitFunction = values.first;
18711873
bool emitSize = values.second;
18721874
// TODO: This calculation should be extracted out into standalone functions
18731875
// emitted on-demand per-module to improve codesize.
18741876
switch (representation) {
18751877
case SILFunctionTypeRepresentation::Thick: {
1878+
assert(!functionPointer.useStaticContextSize());
18761879
// If the called function is thick, the size of the called function's
18771880
// async context is not statically knowable.
18781881
//
@@ -1941,16 +1944,25 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
19411944
SmallVector<std::pair<llvm::BasicBlock *, llvm::Value *>, 2> sizePhiValues;
19421945
{ // thin
19431946
IGF.Builder.emitBlock(thinBlock);
1947+
auto *ptr = functionPointer.getRawPointer();
1948+
if (auto authInfo = functionPointer.getAuthInfo()) {
1949+
ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
1950+
}
1951+
auto *afpPtr =
1952+
IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
19441953
if (emitFunction) {
1945-
auto *uncastFnPtr = functionPointer.getPointer(IGF);
1954+
llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
1955+
auto *uncastFnPtr = IGF.emitLoadOfRelativePointer(
1956+
Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
1957+
/*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
19461958
auto *fnPtr = IGF.Builder.CreateBitCast(uncastFnPtr, IGF.IGM.Int8PtrTy);
1959+
if (auto authInfo = functionPointer.getAuthInfo()) {
1960+
fnPtr = emitPointerAuthSign(IGF, fnPtr, authInfo);
1961+
}
19471962
fnPhiValues.push_back({thinBlock, fnPtr});
19481963
}
19491964
if (emitSize) {
1950-
auto *ptr = functionPointer.getRawPointer();
1951-
auto *descriptorPtr =
1952-
IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
1953-
auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
1965+
auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
19541966
auto *size =
19551967
IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
19561968
sizePhiValues.push_back({thinBlock, size});
@@ -2026,20 +2038,34 @@ std::pair<llvm::Value *, llvm::Value *> irgen::getAsyncFunctionAndSize(
20262038
case SILFunctionTypeRepresentation::WitnessMethod:
20272039
case SILFunctionTypeRepresentation::Closure:
20282040
case SILFunctionTypeRepresentation::Block: {
2041+
auto *ptr = functionPointer.getRawPointer();
2042+
if (auto authInfo = functionPointer.getAuthInfo()) {
2043+
ptr = emitPointerAuthAuth(IGF, ptr, authInfo);
2044+
}
2045+
auto *afpPtr =
2046+
IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
20292047
llvm::Value *fn = nullptr;
20302048
if (emitFunction) {
2031-
fn = functionPointer.getPointer(IGF);
2049+
if (functionPointer.useStaticContextSize()) {
2050+
fn = functionPointer.getRawPointer();
2051+
} else {
2052+
llvm::Value *addrPtr = IGF.Builder.CreateStructGEP(afpPtr, 0);
2053+
fn = IGF.emitLoadOfRelativePointer(
2054+
Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
2055+
/*expectedType*/ functionPointer.getFunctionType()->getPointerTo());
2056+
}
2057+
if (auto authInfo = functionPointer.getAuthInfo()) {
2058+
fn = emitPointerAuthSign(IGF, fn, authInfo);
2059+
}
20322060
}
20332061
llvm::Value *size = nullptr;
20342062
if (emitSize) {
20352063
if (functionPointer.useStaticContextSize()) {
20362064
size = llvm::ConstantInt::get(IGF.IGM.Int32Ty,
20372065
initialContextSize.getValue());
2038-
} else {
2039-
auto *ptr = functionPointer.getRawPointer();
2040-
auto *descriptorPtr =
2041-
IGF.Builder.CreateBitCast(ptr, IGF.IGM.AsyncFunctionPointerPtrTy);
2042-
auto *sizePtr = IGF.Builder.CreateStructGEP(descriptorPtr, 1);
2066+
} else {
2067+
assert(!functionPointer.useStaticContextSize());
2068+
auto *sizePtr = IGF.Builder.CreateStructGEP(afpPtr, 1);
20432069
size = IGF.Builder.CreateLoad(sizePtr, IGF.IGM.getPointerAlignment());
20442070
}
20452071
}
@@ -2336,7 +2362,8 @@ class AsyncCallEmission final : public CallEmission {
23362362
}
23372363
FunctionPointer getCalleeFunctionPointer() override {
23382364
return FunctionPointer(
2339-
FunctionPointer::KindTy::Function, calleeFunction, PointerAuthInfo(),
2365+
FunctionPointer::KindTy::Function, calleeFunction,
2366+
CurCallee.getFunctionPointer().getAuthInfo(),
23402367
IGF.IGM.getSignature(getCallee().getSubstFunctionType()));
23412368
}
23422369
SILType getParameterType(unsigned index) override {
@@ -2374,8 +2401,7 @@ class AsyncCallEmission final : public CallEmission {
23742401
llvm::Intrinsic::coro_async_resume, {});
23752402
auto fnVal = currentResumeFn;
23762403
// Sign the pointer.
2377-
// TODO: use a distinct schema.
2378-
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
2404+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextResume) {
23792405
Address fieldAddr =
23802406
fieldLayout.project(IGF, this->context, /*offsets*/ llvm::None);
23812407
auto authInfo = PointerAuthInfo::emit(
@@ -2454,6 +2480,9 @@ class AsyncCallEmission final : public CallEmission {
24542480
Builder.CreateBitOrPointerCast(dispatchFn, IGM.Int8PtrTy));
24552481
arguments.push_back(
24562482
Builder.CreateBitOrPointerCast(fn.getRawPointer(), IGM.Int8PtrTy));
2483+
if (auto authInfo = fn.getAuthInfo()) {
2484+
arguments.push_back(fn.getAuthInfo().getDiscriminator());
2485+
}
24572486
for (auto arg: args)
24582487
arguments.push_back(arg);
24592488
return IGF.emitSuspendAsyncCall(arguments);
@@ -3640,14 +3669,34 @@ llvm::Value *irgen::emitTaskCreate(
36403669
auto layout = getAsyncContextLayout(
36413670
IGF.IGM, taskFunctionCanSILType, taskFunctionCanSILType, subs);
36423671

3672+
CanSILFunctionType taskContinuationFunctionTy = [&]() {
3673+
ASTContext &ctx = IGF.IGM.IRGen.SIL.getASTContext();
3674+
auto extInfo =
3675+
ASTExtInfoBuilder()
3676+
.withRepresentation(FunctionTypeRepresentation::CFunctionPointer)
3677+
.build();
3678+
// FIXME: Use the appropriate signature for TaskContinuationFunction:
3679+
//
3680+
// using TaskContinuationFunction =
3681+
// SWIFT_CC(swift)
3682+
// void (AsyncTask *, ExecutorRef, AsyncContext *);
3683+
auto ty = FunctionType::get({}, ctx.TheEmptyTupleType, extInfo);
3684+
return IGF.IGM.getLoweredType(ty).castTo<SILFunctionType>();
3685+
}();
3686+
36433687
// Call the function.
36443688
llvm::CallInst *result;
36453689
llvm::Value *theSize, *theFunction;
3690+
auto taskFunctionPointer = FunctionPointer::forExplosionValue(
3691+
IGF, taskFunction, taskFunctionCanSILType);
36463692
std::tie(theFunction, theSize) =
36473693
getAsyncFunctionAndSize(IGF, SILFunctionTypeRepresentation::Thick,
3648-
FunctionPointer::forExplosionValue(
3649-
IGF, taskFunction, taskFunctionCanSILType),
3650-
localContextInfo);
3694+
taskFunctionPointer, localContextInfo);
3695+
if (auto authInfo = PointerAuthInfo::forFunctionPointer(
3696+
IGF.IGM, taskContinuationFunctionTy)) {
3697+
theFunction = emitPointerAuthResign(
3698+
IGF, theFunction, taskFunctionPointer.getAuthInfo(), authInfo);
3699+
}
36513700
theFunction = IGF.Builder.CreateBitOrPointerCast(
36523701
theFunction, IGF.IGM.TaskContinuationFunctionPtrTy);
36533702
theSize = IGF.Builder.CreateZExtOrBitCast(theSize, IGF.IGM.SizeTy);
@@ -4519,18 +4568,25 @@ llvm::Value *FunctionPointer::getPointer(IRGenFunction &IGF) const {
45194568
return Value;
45204569
case KindTy::Value::AsyncFunctionPointer: {
45214570
if (!isFunctionPointerWithoutContext) {
4571+
auto *fnPtr = Value;
4572+
if (auto authInfo = AuthInfo) {
4573+
fnPtr = emitPointerAuthAuth(IGF, fnPtr, authInfo);
4574+
}
45224575
auto *descriptorPtr =
4523-
IGF.Builder.CreateBitCast(Value, IGF.IGM.AsyncFunctionPointerPtrTy);
4576+
IGF.Builder.CreateBitCast(fnPtr, IGF.IGM.AsyncFunctionPointerPtrTy);
45244577
auto *addrPtr = IGF.Builder.CreateStructGEP(descriptorPtr, 0);
4525-
return IGF.emitLoadOfRelativePointer(
4578+
auto *result = IGF.emitLoadOfRelativePointer(
45264579
Address(addrPtr, IGF.IGM.getPointerAlignment()), /*isFar*/ false,
45274580
/*expectedType*/ getFunctionType()->getPointerTo());
4581+
if (auto authInfo = AuthInfo) {
4582+
result = emitPointerAuthSign(IGF, result, authInfo);
4583+
}
4584+
return result;
45284585
} else {
45294586
return IGF.Builder.CreateBitOrPointerCast(
45304587
Value, getFunctionType()->getPointerTo());
45314588
}
45324589
}
4533-
45344590
}
45354591
}
45364592

@@ -4577,8 +4633,7 @@ void irgen::emitAsyncReturn(IRGenFunction &IGF, AsyncContextLayout &asyncLayout,
45774633
.loadAsCopy(IGF, returnToCallerAddr, fn);
45784634
llvm::Value *fnVal = fn.claimNext();
45794635

4580-
// TODO: use distinct schema
4581-
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
4636+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextResume) {
45824637
Address fieldAddr =
45834638
returnToCallerLayout.project(IGF, contextAddr, /*offsets*/ llvm::None);
45844639
auto authInfo = PointerAuthInfo::emit(IGF, schema, fieldAddr.getAddress(),

lib/IRGen/GenFunc.cpp

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -769,7 +769,7 @@ class PartialApplicationForwarderEmission {
769769
virtual void addArgument(llvm::Value *argValue, unsigned index) = 0;
770770
virtual SILParameterInfo getParameterInfo(unsigned index) = 0;
771771
virtual llvm::Value *getContext() = 0;
772-
virtual llvm::Value *getDynamicFunctionPointer() = 0;
772+
virtual llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) = 0;
773773
virtual llvm::Value *getDynamicFunctionContext() = 0;
774774
virtual void addDynamicFunctionContext(Explosion &explosion,
775775
DynamicFunctionKind kind) = 0;
@@ -931,7 +931,9 @@ class SyncPartialApplicationForwarderEmission
931931
return substType->getParameters()[index];
932932
}
933933
llvm::Value *getContext() override { return origParams.claimNext(); }
934-
llvm::Value *getDynamicFunctionPointer() override { return args.takeLast(); }
934+
llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
935+
return args.takeLast();
936+
}
935937
llvm::Value *getDynamicFunctionContext() override { return args.takeLast(); }
936938
void addDynamicFunctionContext(Explosion &explosion,
937939
DynamicFunctionKind kind) override {
@@ -1127,15 +1129,14 @@ class AsyncPartialApplicationForwarderEmission
11271129
llvm::Value *getContext() override {
11281130
return loadValue(layout.getLocalContextLayout());
11291131
}
1130-
llvm::Value *getDynamicFunctionPointer() override {
1132+
llvm::Value *getDynamicFunctionPointer(PointerAuthInfo &authInfo) override {
11311133
assert(dynamicFunction && dynamicFunction->pointer);
11321134
auto *context = dynamicFunction->context;
11331135
if (!context) {
11341136
return dynamicFunction->pointer;
11351137
}
11361138
auto *rawFunction = subIGF.Builder.CreateBitCast(
11371139
dynamicFunction->pointer, origSig.getType()->getPointerTo());
1138-
auto authInfo = PointerAuthInfo::forFunctionPointer(IGM, origType);
11391140
auto functionPointer =
11401141
FunctionPointer(FunctionPointer::KindTy::AsyncFunctionPointer,
11411142
rawFunction, authInfo, origSig);
@@ -1719,19 +1720,19 @@ static llvm::Function *emitPartialApplicationForwarder(IRGenModule &IGM,
17191720

17201721
// Otherwise, it was the last thing we added to the layout.
17211722

1722-
// The dynamic function pointer is packed "last" into the context,
1723-
// and we pulled it out as an argument. Just pop it off.
1724-
auto fnPtr = emission->getDynamicFunctionPointer();
1725-
1726-
// It comes out of the context as an i8*. Cast to the function type.
1727-
fnPtr = subIGF.Builder.CreateBitCast(fnPtr, fnTy);
1728-
17291723
assert(lastCapturedFieldPtr);
17301724
auto authInfo = PointerAuthInfo::emit(subIGF,
17311725
IGM.getOptions().PointerAuth.PartialApplyCapture,
17321726
lastCapturedFieldPtr,
17331727
PointerAuthEntity::Special::PartialApplyCapture);
17341728

1729+
// The dynamic function pointer is packed "last" into the context,
1730+
// and we pulled it out as an argument. Just pop it off.
1731+
auto fnPtr = emission->getDynamicFunctionPointer(authInfo);
1732+
1733+
// It comes out of the context as an i8*. Cast to the function type.
1734+
fnPtr = subIGF.Builder.CreateBitCast(fnPtr, fnTy);
1735+
17351736
return FunctionPointer(FunctionPointer::KindTy::Function, fnPtr, authInfo,
17361737
origSig);
17371738
}();
@@ -2391,7 +2392,12 @@ llvm::Function *IRGenFunction::getOrCreateResumePrjFn() {
23912392
auto &Builder = IGF.Builder;
23922393
auto addr = Builder.CreateBitOrPointerCast(&(*it), IGF.IGM.Int8PtrPtrTy);
23932394
Address callerContextAddr(addr, IGF.IGM.getPointerAlignment());
2394-
auto callerContext = Builder.CreateLoad(callerContextAddr);
2395+
llvm::Value *callerContext = Builder.CreateLoad(callerContextAddr);
2396+
if (auto schema = IGF.IGM.getOptions().PointerAuth.AsyncContextParent) {
2397+
auto authInfo =
2398+
PointerAuthInfo::emit(IGF, schema, addr, PointerAuthEntity());
2399+
callerContext = emitPointerAuthAuth(IGF, callerContext, authInfo);
2400+
}
23952401
Builder.CreateRet(callerContext);
23962402
},
23972403
false /*isNoInline*/));
@@ -2412,6 +2418,10 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
24122418
ArrayRef<llvm::Type *> argTypes) {
24132419
SmallVector<llvm::Type*, 8> argTys;
24142420
argTys.push_back(IGM.Int8PtrTy); // Function pointer to be called.
2421+
auto originalAuthInfo = fnPtr.getAuthInfo();
2422+
if (fnPtr.getAuthInfo()) {
2423+
argTys.push_back(IGM.Int64Ty); // Discriminator for the function pointer.
2424+
}
24152425
for (auto ty : argTypes) {
24162426
argTys.push_back(ty);
24172427
}
@@ -2431,13 +2441,18 @@ IRGenFunction::createAsyncDispatchFn(const FunctionPointer &fnPtr,
24312441
IGM.DebugInfo->emitArtificialFunction(dispatchIGF, dispatch);
24322442
auto &Builder = dispatchIGF.Builder;
24332443
auto it = dispatchIGF.CurFn->arg_begin(), end = dispatchIGF.CurFn->arg_end();
2434-
llvm::Value *ptrArg = &*(it++);
2444+
llvm::Value *fnPtrArg = &*(it++);
2445+
llvm::Value *discriminatorArg = ((bool)originalAuthInfo) ? &*(it++) : nullptr;
24352446
SmallVector<llvm::Value *, 8> callArgs;
24362447
for (; it != end; ++it) {
24372448
callArgs.push_back(&*it);
24382449
}
2439-
ptrArg = Builder.CreateBitOrPointerCast(ptrArg, calleeFnPtrType);
2440-
auto callee = FunctionPointer(fnPtr.getKind(), ptrArg, fnPtr.getAuthInfo(),
2450+
fnPtrArg = Builder.CreateBitOrPointerCast(fnPtrArg, calleeFnPtrType);
2451+
PointerAuthInfo newAuthInfo =
2452+
((bool)originalAuthInfo)
2453+
? PointerAuthInfo(fnPtr.getAuthInfo().getKey(), discriminatorArg)
2454+
: originalAuthInfo;
2455+
auto callee = FunctionPointer(fnPtr.getKind(), fnPtrArg, newAuthInfo,
24412456
fnPtr.getSignature());
24422457
auto call = Builder.CreateCall(callee, callArgs);
24432458
call->setTailCall();

lib/IRGen/GenOpaque.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -484,7 +484,7 @@ IRGenFunction::emitValueWitnessFunctionRef(SILType type,
484484

485485
auto vwtable = emitValueWitnessTableRef(type, &metadataSlot);
486486
auto witness = emitLoadOfValueWitnessFunction(*this, vwtable, index);
487-
setScopedLocalTypeDataForLayout(type, key, witness.getPointer(*this));
487+
setScopedLocalTypeDataForLayout(type, key, witness.getRawPointer());
488488
if (auto &authInfo = witness.getAuthInfo()) {
489489
setScopedLocalTypeDataForLayout(type,
490490
LocalTypeDataKind::forValueWitnessDiscriminator(index),

lib/IRGen/GenPointerAuth.cpp

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,7 @@ llvm::Value *irgen::emitPointerAuthStrip(IRGenFunction &IGF,
7171
FunctionPointer irgen::emitPointerAuthResign(IRGenFunction &IGF,
7272
const FunctionPointer &fn,
7373
const PointerAuthInfo &newAuthInfo) {
74-
// TODO: Handle resigning AsyncFunctionPointers.
75-
assert(fn.getKind().value == FunctionPointer::KindTy::Value::Function);
76-
llvm::Value *fnPtr = emitPointerAuthResign(IGF, fn.getPointer(IGF),
74+
llvm::Value *fnPtr = emitPointerAuthResign(IGF, fn.getRawPointer(),
7775
fn.getAuthInfo(), newAuthInfo);
7876
return FunctionPointer(fn.getKind(), fnPtr, newAuthInfo, fn.getSignature());
7977
}

lib/IRGen/IRGen.cpp

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -690,8 +690,12 @@ static void setPointerAuthOptions(PointerAuthOptions &opts,
690690
SpecialPointerAuthDiscriminators::ResilientClassStubInitCallback);
691691

692692
opts.AsyncContextParent =
693-
PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
693+
PointerAuthSchema(dataKey, /*address*/ true, Discrimination::Constant,
694694
SpecialPointerAuthDiscriminators::AsyncContextParent);
695+
696+
opts.AsyncContextResume =
697+
PointerAuthSchema(codeKey, /*address*/ true, Discrimination::Constant,
698+
SpecialPointerAuthDiscriminators::AsyncContextResume);
695699
}
696700

697701
std::unique_ptr<llvm::TargetMachine>

test/Concurrency/Runtime/async_task_priority_basic.swift

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
// REQUIRES: executable_test
33
// REQUIRES: concurrency
44
// REQUIRES: libdispatch
5-
// XFAIL: CPU=arm64e
65

76
import Dispatch
87

test/Concurrency/Runtime/basic_future.swift

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
// REQUIRES: executable_test
44
// REQUIRES: concurrency
55
// REQUIRES: libdispatch
6-
// XFAIL: CPU=arm64e
76

87
import Dispatch
98

test/Concurrency/Runtime/future_fibonacci.swift

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
// REQUIRES: executable_test
44
// REQUIRES: concurrency
55
// REQUIRES: libdispatch
6-
// XFAIL: CPU=arm64e
76

87
import Dispatch
98

test/IRGen/async/partial_apply.sil

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
// RUN: %target-swift-frontend -emit-module -enable-library-evolution -emit-module-path=%t/resilient_struct.swiftmodule -module-name=resilient_struct %S/../../Inputs/resilient_struct.swift
33
// RUN: %target-swift-frontend -I %t -emit-ir %s | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
44

5-
// REQUIRES: CPU=x86_64
65
// REQUIRES: concurrency
76

87
import Builtin

0 commit comments

Comments
 (0)