Skip to content

Commit bc68fc9

Browse files
nolagnolag
committed
Use subscriptions to determine where things run
1 parent dc7ba03 commit bc68fc9

9 files changed

Lines changed: 174 additions & 68 deletions

File tree

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ update-capabilities:
6969

7070
# Override if you wish to test against a branch.
7171
# Alternatively, you can override the directory in chainlink-common to point to this repository
72-
COMMON_VERSION ?= 660e58a08c23698a5115b5c95ae6f64f97957932
72+
COMMON_VERSION ?= b8745afcfc871d3c3e311b37bf085fe07d057153
7373
MODULE := github.com/smartcontractkit/chainlink-common
7474

7575
# Override on Windows if you aren't using a Unix-like shell:

‎cre/execution_handler.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,18 @@ func Handler[C any, M proto.Message, T any, O any](trigger Trigger[M, T], callba
3030
// HandlerInTee creates a coupling of a Trigger and a callback function to be used in TEE (Trusted Execution Environment) mode.
3131
// The coupling ensures that when the Trigger is invoked, the callback function is called with a TeeRuntime.
3232
func HandlerInTee[C any, M proto.Message, T any, O any, A AcceptedTees](trigger Trigger[M, T], callback func(config C, runtime TeeRuntime, payload T) (O, error), tees A) ExecutionHandler[C, Runtime] {
33-
requirements := &sdk.Requirements{Tee: &sdk.Tee{}}
33+
wrapped := func(config C, runtime Runtime, t T) (O, error) {
34+
helper, ok := runtime.(interface{ Tee() TeeRuntime })
35+
if !ok {
36+
panic("Runner did not provide an extractable TEERuntime. If you wrapped the runtime, wrap the method Tee() TeeRuntime instead.")
37+
}
38+
39+
return callback(config, helper.Tee(), t)
40+
}
41+
return handler(trigger, wrapped, teeRequirements(tees))
42+
}
43+
44+
func teeRequirements[A AcceptedTees](tees A) *sdk.Requirements {
3445
reqs := &sdk.Requirements{Tee: &sdk.Tee{}}
3546
switch typedTees := any(tees).(type) {
3647
case []TeeAndRegions:
@@ -42,17 +53,7 @@ func HandlerInTee[C any, M proto.Message, T any, O any, A AcceptedTees](trigger
4253
case AnyTee:
4354
reqs.Tee.Type = &sdk.Tee_Any{Any: &emptypb.Empty{}}
4455
}
45-
46-
wrapped := func(config C, runtime Runtime, t T) (O, error) {
47-
// hack to allow it to pass us a teeRuntime
48-
helper, ok := runtime.(interface{ Tee() TeeRuntime })
49-
if !ok {
50-
panic("Runner did not provide an extractable TEERuntime. If you wrapped the runtime, wrap the method Tee() TeeRuntime instead.")
51-
}
52-
53-
return callback(config, helper.Tee(), t)
54-
}
55-
return handler(trigger, wrapped, requirements)
56+
return reqs
5657
}
5758

5859
func handler[R, C any, M proto.Message, T any, O any](trigger Trigger[M, T], callback func(config C, runtime R, payload T) (O, error), requirements *sdk.Requirements) ExecutionHandler[C, R] {
@@ -77,7 +78,7 @@ func handler[R, C any, M proto.Message, T any, O any](trigger Trigger[M, T], cal
7778
return eh
7879
}
7980

80-
return executionHandlerWithRequirementsImpl[C, R]{ExecutionHandler: eh, requirements: requirements}
81+
return &executionHandlerWithRequirementsImpl[C, R]{ExecutionHandler: eh, requirements: requirements}
8182
}
8283

8384
type executionHandlerImpl[C, R any, M proto.Message, T any] struct {

‎cre/runner.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,14 @@ import (
66
"github.com/smartcontractkit/chainlink-protos/cre/go/sdk"
77
)
88

9+
type InitFn[C any] = func(config C, logger *slog.Logger, secretsProvider SecretsProvider) (Workflow[C], error)
10+
911
// Runner is the entry point to running a CRE workflow.
1012
type Runner[C any] interface {
1113
// Run creates the workflow and starts it.
1214
// Upon registration of a workflow, a run is used to register to `Trigger`s.
1315
// Upon receiving a trigger, the appropriate handler's callback is invoked.
14-
Run(initFn func(config C, logger *slog.Logger, secretsProvider SecretsProvider) (Workflow[C], error))
16+
Run(initFn InitFn[C])
1517
}
1618

1719
type AnyTee struct{}

‎cre/wasm/runner.go‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@ type runnerInternals interface {
2222
switchModes(mode int32)
2323
now(response unsafe.Pointer) int32
2424
exit()
25-
requirements(data unsafe.Pointer, dataLen int32)
2625
}
2726

2827
func newRunner[C Config](parse func(configBytes []byte) (C, error), runnerInternals runnerInternals, runtimeInternals runtimeInternals) cre.Runner[C] {
@@ -42,7 +41,7 @@ func newRunner[C Config](parse func(configBytes []byte) (C, error), runnerIntern
4241
&runner[C, cre.Runtime]{
4342
sp: drt,
4443
runtime: drt,
45-
switchRunner: switchRuntimeWrapper{Runtime: drt},
44+
switchRuntime: &switchRuntimeWrapper{Runtime: drt},
4645
runnerInternals: runnerInternals,
4746
setRuntime: func(maxResponseSize uint64) {
4847
drt.MaxResponseSize = maxResponseSize
@@ -54,13 +53,13 @@ func newRunner[C Config](parse func(configBytes []byte) (C, error), runnerIntern
5453

5554
type runner[C, T any] struct {
5655
runnerInternals
57-
trigger *sdk.Trigger
58-
id string
59-
runtime T
60-
switchRunner T
61-
setRuntime func(maxResponseSize uint64)
62-
config C
63-
sp cre.SecretsProvider
56+
trigger *sdk.Trigger
57+
id string
58+
runtime T
59+
switchRuntime T
60+
setRuntime func(maxResponseSize uint64)
61+
config C
62+
sp cre.SecretsProvider
6463
}
6564

6665
var _ baseRunner[any, cre.Runtime] = (*runner[any, cre.Runtime])(nil)
@@ -77,12 +76,8 @@ func (r *runner[C, T]) run(wfs []cre.ExecutionHandler[C, T]) {
7776
runtime := r.runtime
7877
for idx, handler := range wfs {
7978
if uint64(idx) == r.trigger.Id {
80-
if reqsProvider, ok := handler.(interface{ Requirements() *sdk.Requirements }); ok {
81-
reqs := reqsProvider.Requirements()
82-
requirementsBuffered, _ := proto.Marshal(reqs)
83-
marshalledPtr, marshalledLen, _ := bufferToPointerLen(requirementsBuffered)
84-
r.runnerInternals.requirements(marshalledPtr, marshalledLen)
85-
runtime = r.switchRunner
79+
if _, ok := handler.(cre.ExecutionHandlerWithRequirements[C, T]); ok {
80+
runtime = r.switchRuntime
8681
}
8782

8883
response, err := handler.Callback()(r.config, runtime, r.trigger.Payload)
@@ -122,11 +117,15 @@ func (s *subscriber[C, T]) secretsProvider() cre.SecretsProvider {
122117
func (s *subscriber[C, T]) run(wfs []cre.ExecutionHandler[C, T]) {
123118
subscriptions := make([]*sdk.TriggerSubscription, len(wfs))
124119
for i, handler := range wfs {
125-
subscriptions[i] = &sdk.TriggerSubscription{
120+
sub := &sdk.TriggerSubscription{
126121
Id: handler.CapabilityID(),
127122
Payload: handler.TriggerCfg(),
128123
Method: handler.Method(),
129124
}
125+
if reqsProvider, ok := handler.(cre.ExecutionHandlerWithRequirements[C, T]); ok {
126+
sub.Requirements = reqsProvider.Requirements()
127+
}
128+
subscriptions[i] = sub
130129
}
131130
triggerSubscription := &sdk.TriggerSubscriptionRequest{Subscriptions: subscriptions}
132131

@@ -225,3 +224,4 @@ type switchRuntimeWrapper struct {
225224
func (s *switchRuntimeWrapper) Tee() cre.TeeRuntime {
226225
return sdkimpl.NewTeeRuntime(s.Runtime)
227226
}
227+

‎cre/wasm/runner_test.go‎

Lines changed: 117 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -154,33 +154,131 @@ func TestRunner_Run(t *testing.T) {
154154
})
155155
}
156156

157-
func TestNewTeeRunner(t *testing.T) {
158-
t.Run("Specified list", func(t *testing.T) {
157+
func TestHandlerInTee(t *testing.T) {
158+
t.Run("specified list sets requirements on subscription", func(t *testing.T) {
159159
acceptedTees := []cre.TeeAndRegions{{Type: cre.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}
160160

161-
teeRunner := newTeeRunner(
162-
acceptedTees, func(b []byte) (string, error) { return string(b), nil }, testRunnerInternals(t, subscribeRequest), testRuntimeInternals(t))
161+
internals := testRunnerInternals(t, subscribeRequest)
162+
dr := newRunner(func(b []byte) (string, error) { return string(b), nil }, internals, testRuntimeInternals(t))
163163

164-
requirements := teeRunner.(teeRunnerWrapper[string]).baseRunner.(*subscriber[string, cre.TeeRuntime]).runnerInternals.(*runnerInternalsTestHook).requirementsSent
165-
actual := &sdk.Requirements{}
166-
require.NoError(t, proto.Unmarshal(requirements, actual))
167-
expected := &sdk.Requirements{
168-
Tee: &sdk.Tee{Type: &sdk.Tee_TypeSelection{TypeSelection: &sdk.TeeTypeSelection{Types: []*sdk.TeeTypeAndRegions{{Type: sdk.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}}}},
164+
dr.Run(func(string, *slog.Logger, cre.SecretsProvider) (cre.Workflow[string], error) {
165+
return cre.Workflow[string]{
166+
cre.HandlerInTee(
167+
basictrigger.Trigger(testworkflow.TestWorkflowTriggerConfig()),
168+
func(_ string, _ cre.TeeRuntime, _ *basictrigger.Outputs) (string, error) {
169+
return "tee-result", nil
170+
},
171+
acceptedTees,
172+
),
173+
}, nil
174+
})
175+
176+
actual := &sdk.ExecutionResult{}
177+
require.NoError(t, proto.Unmarshal(internals.sentResponse, actual))
178+
switch result := actual.Result.(type) {
179+
case *sdk.ExecutionResult_TriggerSubscriptions:
180+
subs := result.TriggerSubscriptions.Subscriptions
181+
require.Len(t, subs, 1)
182+
expected := &sdk.Requirements{
183+
Tee: &sdk.Tee{Type: &sdk.Tee_TypeSelection{TypeSelection: &sdk.TeeTypeSelection{Types: []*sdk.TeeTypeAndRegions{{Type: sdk.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}}}},
184+
}
185+
assert.True(t, proto.Equal(expected, subs[0].Requirements))
186+
default:
187+
assert.Fail(t, "unexpected result type", result)
188+
}
189+
})
190+
191+
t.Run("any tee sets requirements on subscription", func(t *testing.T) {
192+
internals := testRunnerInternals(t, subscribeRequest)
193+
dr := newRunner(func(b []byte) (string, error) { return string(b), nil }, internals, testRuntimeInternals(t))
194+
195+
dr.Run(func(string, *slog.Logger, cre.SecretsProvider) (cre.Workflow[string], error) {
196+
return cre.Workflow[string]{
197+
cre.HandlerInTee(
198+
basictrigger.Trigger(testworkflow.TestWorkflowTriggerConfig()),
199+
func(_ string, _ cre.TeeRuntime, _ *basictrigger.Outputs) (string, error) {
200+
return "tee-result", nil
201+
},
202+
cre.AnyTee{},
203+
),
204+
}, nil
205+
})
206+
207+
actual := &sdk.ExecutionResult{}
208+
require.NoError(t, proto.Unmarshal(internals.sentResponse, actual))
209+
switch result := actual.Result.(type) {
210+
case *sdk.ExecutionResult_TriggerSubscriptions:
211+
subs := result.TriggerSubscriptions.Subscriptions
212+
require.Len(t, subs, 1)
213+
expected := &sdk.Requirements{
214+
Tee: &sdk.Tee{Type: &sdk.Tee_Any{Any: &emptypb.Empty{}}},
215+
}
216+
assert.True(t, proto.Equal(expected, subs[0].Requirements))
217+
default:
218+
assert.Fail(t, "unexpected result type", result)
219+
}
220+
})
221+
222+
t.Run("regular handler has no requirements on subscription", func(t *testing.T) {
223+
internals := testRunnerInternals(t, subscribeRequest)
224+
dr := newRunner(func(b []byte) (string, error) { return string(b), nil }, internals, testRuntimeInternals(t))
225+
226+
dr.Run(func(string, *slog.Logger, cre.SecretsProvider) (cre.Workflow[string], error) {
227+
return cre.Workflow[string]{
228+
cre.Handler(
229+
basictrigger.Trigger(testworkflow.TestWorkflowTriggerConfig()),
230+
func(_ string, _ cre.Runtime, _ *basictrigger.Outputs) (string, error) {
231+
return "no-tee", nil
232+
},
233+
),
234+
}, nil
235+
})
236+
237+
actual := &sdk.ExecutionResult{}
238+
require.NoError(t, proto.Unmarshal(internals.sentResponse, actual))
239+
switch result := actual.Result.(type) {
240+
case *sdk.ExecutionResult_TriggerSubscriptions:
241+
subs := result.TriggerSubscriptions.Subscriptions
242+
require.Len(t, subs, 1)
243+
assert.Nil(t, subs[0].Requirements)
244+
default:
245+
assert.Fail(t, "unexpected result type", result)
169246
}
170-
assert.True(t, proto.Equal(expected, actual))
171247
})
172248

173-
t.Run("any tee", func(t *testing.T) {
174-
teeRunner := newTeeRunner(
175-
cre.AnyTee{}, func(b []byte) (string, error) { return string(b), nil }, testRunnerInternals(t, subscribeRequest), testRuntimeInternals(t))
249+
t.Run("tee handler callback receives TeeRuntime", func(t *testing.T) {
250+
acceptedTees := []cre.TeeAndRegions{{Type: cre.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}
176251

177-
requirements := teeRunner.(teeRunnerWrapper[string]).baseRunner.(*subscriber[string, cre.TeeRuntime]).runnerInternals.(*runnerInternalsTestHook).requirementsSent
178-
actual := &sdk.Requirements{}
179-
require.NoError(t, proto.Unmarshal(requirements, actual))
180-
expected := &sdk.Requirements{
181-
Tee: &sdk.Tee{Type: &sdk.Tee_Any{Any: &emptypb.Empty{}}},
252+
triggerReq := &sdk.ExecuteRequest{
253+
Config: anyConfig,
254+
MaxResponseSize: anyMaxResponseSize,
255+
Request: &sdk.ExecuteRequest_Trigger{
256+
Trigger: &sdk.Trigger{
257+
Id: 0,
258+
Payload: mustAny(testworkflow.TestWorkflowTrigger()),
259+
},
260+
},
182261
}
183-
assert.True(t, proto.Equal(expected, actual))
262+
263+
internals := testRunnerInternals(t, triggerReq)
264+
dr := newRunner(func(b []byte) (string, error) { return string(b), nil }, internals, testRuntimeInternals(t))
265+
266+
callbackInvoked := false
267+
dr.Run(func(string, *slog.Logger, cre.SecretsProvider) (cre.Workflow[string], error) {
268+
return cre.Workflow[string]{
269+
cre.HandlerInTee(
270+
basictrigger.Trigger(testworkflow.TestWorkflowTriggerConfig()),
271+
func(_ string, rt cre.TeeRuntime, _ *basictrigger.Outputs) (string, error) {
272+
callbackInvoked = true
273+
assert.NotNil(t, rt, "TeeRuntime should not be nil")
274+
return "done", nil
275+
},
276+
acceptedTees,
277+
),
278+
}, nil
279+
})
280+
281+
assert.True(t, callbackInvoked, "tee callback should have been invoked")
184282
})
185283
}
186284

‎cre/wasm/runner_test_hooks.go‎

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,12 @@ import (
88
)
99

1010
type runnerInternalsTestHook struct {
11-
testTb testing.TB
12-
execId string
13-
arguments []string
14-
sentResponse []byte
15-
modeSwitched bool
16-
mode int32
17-
requirementsSent []byte
11+
testTb testing.TB
12+
execId string
13+
arguments []string
14+
sentResponse []byte
15+
modeSwitched bool
16+
mode int32
1817
}
1918

2019
func (r *runnerInternalsTestHook) args() []string {
@@ -41,8 +40,4 @@ func (r *runnerInternalsTestHook) exit() {
4140
// Unlike the WASM, tests continue to execute
4241
}
4342

44-
func (r *runnerInternalsTestHook) requirements(data unsafe.Pointer, dataLen int32) {
45-
r.requirementsSent = unsafe.Slice((*byte)(data), dataLen)
46-
}
47-
4843
var _ runnerInternals = (*runnerInternalsTestHook)(nil)

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ go 1.25.3
55
require (
66
github.com/ethereum/go-ethereum v1.17.2
77
github.com/shopspring/decimal v1.4.0
8-
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260421194300-2c8da85a337a
8+
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260427170224-3b3204904066
99
github.com/stretchr/testify v1.11.1
1010
google.golang.org/protobuf v1.36.11
1111
)

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0t
2727
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
2828
github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k=
2929
github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME=
30-
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260421194300-2c8da85a337a h1:2mwWuRputcmFMzehSUlk95q9NQp9cspupb6FZxgCh7w=
31-
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260421194300-2c8da85a337a/go.mod h1:Jqt53s27Tr0jDl8mdBXg1xhu6F8Fci8JOuq43tgHOM8=
30+
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260427170224-3b3204904066 h1:XgmfrVnD6Z2yf6f+4qcGZlqvdJlffRippMmvqE8Yl3c=
31+
github.com/smartcontractkit/chainlink-protos/cre/go v0.0.0-20260427170224-3b3204904066/go.mod h1:Jqt53s27Tr0jDl8mdBXg1xhu6F8Fci8JOuq43tgHOM8=
3232
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
3333
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
3434
golang.org/x/sys v0.40.0 h1:DBZZqJ2Rkml6QMQsZywtnjnnGvHza6BTfYFWY9kjEWQ=

‎standard_tests/tee_runtime/main_wasip1.go‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,27 @@ import (
88
"github.com/smartcontractkit/cre-sdk-go/internal_testing/capabilities/basictrigger"
99
)
1010

11+
var teeRequirements = []cre.TeeAndRegions{{Type: cre.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}
12+
1113
func subscribe(_ []byte, _ *slog.Logger, _ cre.SecretsProvider) (cre.Workflow[[]byte], error) {
12-
teeRequiements := []cre.TeeAndRegions{{Type: cre.TeeType_TEE_TYPE_AWS_NITRO, Regions: []string{"us-west-2"}}}
1314
return cre.Workflow[[]byte]{
1415
cre.HandlerInTee(
1516
basictrigger.Trigger(&basictrigger.Config{Name: "first-trigger", Number: 100}),
16-
trigger,
17-
teeRequiements),
17+
teeTrigger,
18+
teeRequirements,
19+
),
20+
cre.Handler(
21+
basictrigger.Trigger(&basictrigger.Config{Name: "second-trigger", Number: 200}),
22+
regularTrigger,
23+
),
1824
}, nil
1925
}
2026

21-
func trigger(config []byte, runtime cre.TeeRuntime, payload *basictrigger.Outputs) (int32, error) {
27+
func teeTrigger(_ []byte, _ cre.TeeRuntime, _ *basictrigger.Outputs) (int32, error) {
28+
return 0, nil
29+
}
30+
31+
func regularTrigger(_ []byte, _ cre.Runtime, _ *basictrigger.Outputs) (int32, error) {
2232
return 0, nil
2333
}
2434

0 commit comments

Comments
 (0)