Skip to content

Commit d8a75a6

Browse files
committed
Better regions
1 parent fbfac1c commit d8a75a6

8 files changed

Lines changed: 455 additions & 54 deletions

File tree

cre/execution_handler.go

Lines changed: 8 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -28,31 +28,18 @@ func Handler[C any, M proto.Message, T any, O any](trigger Trigger[M, T], callba
2828

2929
// HandlerInTee creates a coupling of a Trigger and a callback function to be used in TEE (Trusted Execution Environment) mode.
3030
// The coupling ensures that when the Trigger is invoked, the callback function is called with a TeeRuntime.
31-
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] {
32-
wrapped := func(config C, runtime Runtime, t T) (O, error) {
31+
func HandlerInTee[C any, M proto.Message, T any, O any](trigger Trigger[M, T], callback func(config C, runtime TeeRuntime, payload T) (O, error), tees TeeConstraint) ExecutionHandler[C, Runtime] {
32+
return handler(trigger, wrapTeeCallback(callback), tees.toRequirements())
33+
}
34+
35+
func wrapTeeCallback[C any, T any, O any](callback func(config C, runtime TeeRuntime, payload T) (O, error)) func(config C, runtime Runtime, payload T) (O, error) {
36+
return func(config C, runtime Runtime, t T) (O, error) {
3337
helper, ok := runtime.(interface{ Tee() TeeRuntime })
3438
if !ok {
3539
panic("Runner did not provide an extractable TEERuntime. If you wrapped the runtime, wrap the method Tee() TeeRuntime instead.")
3640
}
37-
3841
return callback(config, helper.Tee(), t)
3942
}
40-
return handler(trigger, wrapped, teeRequirements(tees))
41-
}
42-
43-
func teeRequirements[A AcceptedTees](tees A) *sdk.Requirements {
44-
reqs := &sdk.Requirements{Tee: &sdk.Tee{}}
45-
switch typedTees := any(tees).(type) {
46-
case []TeeAndRegions:
47-
typesRegions := make([]*sdk.TeeTypeAndRegions, len(typedTees))
48-
for i, tee := range typedTees {
49-
typesRegions[i] = &sdk.TeeTypeAndRegions{Type: tee.Type, Regions: tee.Regions}
50-
}
51-
reqs.Tee.Item = &sdk.Tee_TeeTypesAndRegions{TeeTypesAndRegions: &sdk.TeeTypesAndRegions{TeeTypeAndRegions: typesRegions}}
52-
case AnyTee:
53-
reqs.Tee.Item = &sdk.Tee_AnyRegions{AnyRegions: &sdk.Regions{Regions: typedTees.Regions}}
54-
}
55-
return reqs
5643
}
5744

5845
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] {
@@ -100,6 +87,8 @@ type executionHandlerWithRequirementsImpl[C, R any] struct {
10087
requirements *sdk.Requirements
10188
}
10289

90+
var _ ExecutionHandlerWithRequirements[any, any] = (*executionHandlerWithRequirementsImpl[any, any])(nil)
91+
10392
func (h *executionHandlerWithRequirementsImpl[C, R]) Requirements() *sdk.Requirements {
10493
return h.requirements
10594
}

cre/runner.go

Lines changed: 1 addition & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,6 @@
11
package cre
22

3-
import (
4-
"log/slog"
5-
6-
"github.com/smartcontractkit/chainlink-protos/cre/go/sdk"
7-
)
3+
import "log/slog"
84

95
type InitFn[C any] = func(config C, logger *slog.Logger, secretsProvider SecretsProvider) (Workflow[C], error)
106

@@ -15,26 +11,3 @@ type Runner[C any] interface {
1511
// Upon receiving a trigger, the appropriate handler's callback is invoked.
1612
Run(initFn InitFn[C])
1713
}
18-
19-
type Type = sdk.TeeType
20-
21-
type TeeAndRegions struct {
22-
Type
23-
24-
// Regions limits what regions the TEE can run in.
25-
// If empty or nil, there is no region limitation.
26-
Regions []string
27-
}
28-
29-
// AnyTee specifies that any TEE type is acceptable, with optional region constraints.
30-
type AnyTee struct {
31-
// Regions limits what regions the TEE can run in.
32-
// If empty or nil, there is no region limitation.
33-
Regions []string
34-
}
35-
36-
const TeeType_TEE_TYPE_AWS_NITRO = sdk.TeeType_TEE_TYPE_AWS_NITRO
37-
38-
type AcceptedTees interface {
39-
[]TeeAndRegions | AnyTee
40-
}

cre/tee_constraints.go

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
package cre
2+
3+
import (
4+
"bytes"
5+
"encoding/json"
6+
"fmt"
7+
8+
"github.com/smartcontractkit/chainlink-protos/cre/go/sdk"
9+
)
10+
11+
// TeeConstraint describes the set of (TEE, region) pairs a handler will accept.
12+
// Sealed via the unexported toRequirements method: only AnyTee, AnyTeeInRegions,
13+
// and OneOfTees (declared in this package) can implement it.
14+
type TeeConstraint interface {
15+
toRequirements() *sdk.Requirements
16+
}
17+
18+
// Tees is the config-friendly wrapper around a TeeConstraint. Its UnmarshalJSON
19+
// inspects the JSON shape and produces the matching variant; null and unknown
20+
// shapes are rejected to fail closed. Use *Tees in config structs when the
21+
// field is optional.
22+
type Tees struct {
23+
TeeConstraint
24+
}
25+
26+
// AnyTee accepts any TEE in any region. JSON form: {}.
27+
type AnyTee struct{}
28+
29+
func (AnyTee) toRequirements() *sdk.Requirements {
30+
return &sdk.Requirements{Tee: &sdk.Tee{Item: &sdk.Tee_AnyRegions{AnyRegions: &sdk.Regions{}}}}
31+
}
32+
33+
// AnyTeeInRegions accepts any TEE provided it runs in one of the listed regions.
34+
// JSON form: {"regions": [...]}.
35+
type AnyTeeInRegions struct {
36+
Regions []Region `json:"regions"`
37+
}
38+
39+
func (a AnyTeeInRegions) toRequirements() *sdk.Requirements {
40+
regs := make([]string, len(a.Regions))
41+
for i, r := range a.Regions {
42+
regs[i] = string(r)
43+
}
44+
return &sdk.Requirements{Tee: &sdk.Tee{Item: &sdk.Tee_AnyRegions{AnyRegions: &sdk.Regions{Regions: regs}}}}
45+
}
46+
47+
// OneOfTees accepts any of the listed per-TEE bindings. Each binding may carry
48+
// its own region set. JSON form: [{"tee": "...", "regions": [...]}, ...].
49+
type OneOfTees []TeeBinding
50+
51+
func (o OneOfTees) toRequirements() *sdk.Requirements {
52+
out := make([]*sdk.TeeTypeAndRegions, len(o))
53+
for i, b := range o {
54+
out[i] = b.teeTypeAndRegions()
55+
}
56+
return &sdk.Requirements{Tee: &sdk.Tee{Item: &sdk.Tee_TeeTypesAndRegions{TeeTypesAndRegions: &sdk.TeeTypesAndRegions{TeeTypeAndRegions: out}}}}
57+
}
58+
59+
// TeeBinding is the sealed interface for a single entry inside OneOfTees. Each
60+
// implementing type owns the enum of regions that TEE supports (e.g. Nitro <->
61+
// NitroRegion) so a region for the wrong TEE is a compile-time error. Sealed
62+
// via the unexported teeTypeAndRegions method.
63+
type TeeBinding interface {
64+
teeTypeAndRegions() *sdk.TeeTypeAndRegions
65+
}
66+
67+
// Nitro is the AWS Nitro TEE. JSON tag: "nitro".
68+
type Nitro struct {
69+
Regions []NitroRegion `json:"regions,omitempty"`
70+
}
71+
72+
func (n Nitro) teeTypeAndRegions() *sdk.TeeTypeAndRegions {
73+
var regs []string
74+
if n.Regions != nil {
75+
regs = make([]string, len(n.Regions))
76+
for i, r := range n.Regions {
77+
regs[i] = string(r)
78+
}
79+
}
80+
return &sdk.TeeTypeAndRegions{Type: sdk.TeeType_TEE_TYPE_AWS_NITRO, Regions: regs}
81+
}
82+
83+
// Region is the global region enum used by AnyTeeInRegions, where no specific
84+
// TEE is pinned. Add a new region by adding a const plus a case in UnmarshalText.
85+
type Region string
86+
87+
const (
88+
AwsUsWest2 Region = "us-west-2"
89+
)
90+
91+
func (r *Region) UnmarshalText(b []byte) error {
92+
switch v := Region(b); v {
93+
case AwsUsWest2:
94+
*r = v
95+
return nil
96+
default:
97+
return fmt.Errorf("unknown region %q", b)
98+
}
99+
}
100+
101+
// NitroRegion enumerates the regions AWS Nitro is supported in. Distinct from
102+
// Region so that a non-Nitro region cannot appear inside a Nitro binding.
103+
type NitroRegion string
104+
105+
const (
106+
NitroUsWest2 = NitroRegion(AwsUsWest2)
107+
)
108+
109+
func (r *NitroRegion) UnmarshalText(b []byte) error {
110+
switch v := NitroRegion(b); v {
111+
case NitroUsWest2:
112+
*r = v
113+
return nil
114+
default:
115+
return fmt.Errorf("aws nitro does not support region %q", b)
116+
}
117+
}
118+
119+
// teeBindingFactories maps the JSON "tee" discriminator to a fresh binding to
120+
// unmarshal into. Add a new TEE by registering it here.
121+
var teeBindingFactories = map[string]func() teeBindingUnmarshaler{
122+
"nitro": func() teeBindingUnmarshaler { return &Nitro{} },
123+
}
124+
125+
// teeBindingUnmarshaler is the concrete pointer type that both implements
126+
// TeeBinding and can be json.Unmarshal'd into. asValue returns the value form
127+
// of the binding, which is what gets stored in OneOfTees.
128+
type teeBindingUnmarshaler interface {
129+
TeeBinding
130+
json.Unmarshaler
131+
asValue() TeeBinding
132+
}
133+
134+
func (n *Nitro) asValue() TeeBinding { return *n }
135+
136+
func (n *Nitro) UnmarshalJSON(b []byte) error {
137+
type raw struct {
138+
Tee string `json:"tee"`
139+
Regions json.RawMessage `json:"regions"`
140+
}
141+
dec := json.NewDecoder(bytes.NewReader(b))
142+
dec.DisallowUnknownFields()
143+
var r raw
144+
if err := dec.Decode(&r); err != nil {
145+
return err
146+
}
147+
if len(r.Regions) == 0 {
148+
n.Regions = nil
149+
return nil
150+
}
151+
var regions []NitroRegion
152+
if err := json.Unmarshal(r.Regions, &regions); err != nil {
153+
return err
154+
}
155+
if len(regions) == 0 {
156+
return fmt.Errorf(`"regions" must not be empty`)
157+
}
158+
n.Regions = regions
159+
return nil
160+
}
161+
162+
// UnmarshalJSON dispatches on the JSON shape: {} -> AnyTee, {"regions":[...]} ->
163+
// AnyTeeInRegions, [...] -> OneOfTees. null and any other shape error.
164+
func (t *Tees) UnmarshalJSON(b []byte) error {
165+
b = bytes.TrimSpace(b)
166+
if len(b) == 0 || bytes.Equal(b, []byte("null")) {
167+
return fmt.Errorf("tee constraint is required; use *Tees for optional fields")
168+
}
169+
switch b[0] {
170+
case '[':
171+
return t.decodeArray(b)
172+
case '{':
173+
return t.decodeObject(b)
174+
default:
175+
return fmt.Errorf("tee constraint must be an object or array, got: %s", b)
176+
}
177+
}
178+
179+
func (t *Tees) decodeArray(b []byte) error {
180+
var raw []json.RawMessage
181+
if err := json.Unmarshal(b, &raw); err != nil {
182+
return err
183+
}
184+
if len(raw) == 0 {
185+
return fmt.Errorf("tee constraint list must not be empty")
186+
}
187+
out := make(OneOfTees, 0, len(raw))
188+
for i, entry := range raw {
189+
bind, err := unmarshalBinding(entry)
190+
if err != nil {
191+
return fmt.Errorf("tee constraint[%d]: %w", i, err)
192+
}
193+
out = append(out, bind)
194+
}
195+
t.TeeConstraint = out
196+
return nil
197+
}
198+
199+
func (t *Tees) decodeObject(b []byte) error {
200+
var raw map[string]json.RawMessage
201+
if err := json.Unmarshal(b, &raw); err != nil {
202+
return err
203+
}
204+
if len(raw) == 0 {
205+
t.TeeConstraint = AnyTee{}
206+
return nil
207+
}
208+
regionsRaw, ok := raw["regions"]
209+
if !ok || len(raw) != 1 {
210+
return fmt.Errorf(`unrecognized tee constraint object; expected {} or {"regions":[...]}`)
211+
}
212+
var regions []Region
213+
if err := json.Unmarshal(regionsRaw, &regions); err != nil {
214+
return err
215+
}
216+
if len(regions) == 0 {
217+
return fmt.Errorf(`"regions" must not be empty`)
218+
}
219+
t.TeeConstraint = AnyTeeInRegions{Regions: regions}
220+
return nil
221+
}
222+
223+
func unmarshalBinding(b []byte) (TeeBinding, error) {
224+
var probe struct {
225+
Tee string `json:"tee"`
226+
}
227+
if err := json.Unmarshal(b, &probe); err != nil {
228+
return nil, err
229+
}
230+
if probe.Tee == "" {
231+
return nil, fmt.Errorf(`missing "tee" discriminator`)
232+
}
233+
factory, ok := teeBindingFactories[probe.Tee]
234+
if !ok {
235+
return nil, fmt.Errorf("unknown tee %q", probe.Tee)
236+
}
237+
target := factory()
238+
if err := target.UnmarshalJSON(b); err != nil {
239+
return nil, err
240+
}
241+
return target.asValue(), nil
242+
}

0 commit comments

Comments
 (0)