From 64334951139332024af6744cd51f21333d687662 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 22:18:28 -0700 Subject: [PATCH 1/8] feat(ateletpb): add signed object access to external checkpoint config --- internal/proto/ateletpb/atelet.pb.go | 357 +++++++++++++++++++++------ internal/proto/ateletpb/atelet.proto | 56 +++++ 2 files changed, 337 insertions(+), 76 deletions(-) diff --git a/internal/proto/ateletpb/atelet.pb.go b/internal/proto/ateletpb/atelet.pb.go index 90aac619f3..bc73820758 100644 --- a/internal/proto/ateletpb/atelet.pb.go +++ b/internal/proto/ateletpb/atelet.pb.go @@ -1994,6 +1994,122 @@ func (x *ExternalCheckpointConfiguration) GetSnapshotUri() string { return "" } +// SignedObjectAccess is a short-lived, prefix-scoped capability that lets atelet +// read or write a single snapshot's objects with no cloud credential on the +// node. atelet dispatches on the mechanics carried here, never on the cloud: +// the control-plane broker fills in whichever fields its backend needs. +type SignedObjectAccess struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Base URL of the account/bucket+container, e.g. + // "https://acct.blob.core.windows.net/snapshots" or "http://s3-host/bucket". + PrefixUrl string `protobuf:"bytes,1,opt,name=prefix_url,json=prefixUrl,proto3" json:"prefix_url,omitempty"` + // Read mechanics (populated on restore). atelet fetches an object by either + // looking it up in read_object_urls (S3: per-object presigned GET), or + // appending "?" to prefix_url (Azure: one container SAS). + ReadToken string `protobuf:"bytes,2,opt,name=read_token,json=readToken,proto3" json:"read_token,omitempty"` + ReadObjectUrls map[string]string `protobuf:"bytes,3,rep,name=read_object_urls,json=readObjectUrls,proto3" json:"read_object_urls,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + // Write mechanics (populated on checkpoint). write_method is "PUT" or "POST". + // + // PUT (Azure): atelet PUTs "/?" with + // write_headers (block blob needs x-ms-blob-type: BlockBlob). + // POST (S3): atelet submits a multipart form to post_url with post_fields, + // setting the "key" field per object; one signed starts-with policy + // covers every file under the prefix. + WriteMethod string `protobuf:"bytes,4,opt,name=write_method,json=writeMethod,proto3" json:"write_method,omitempty"` + WriteToken string `protobuf:"bytes,5,opt,name=write_token,json=writeToken,proto3" json:"write_token,omitempty"` + WriteHeaders map[string]string `protobuf:"bytes,6,rep,name=write_headers,json=writeHeaders,proto3" json:"write_headers,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + PostUrl string `protobuf:"bytes,7,opt,name=post_url,json=postUrl,proto3" json:"post_url,omitempty"` + PostFields map[string]string `protobuf:"bytes,8,rep,name=post_fields,json=postFields,proto3" json:"post_fields,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SignedObjectAccess) Reset() { + *x = SignedObjectAccess{} + mi := &file_atelet_proto_msgTypes[30] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SignedObjectAccess) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SignedObjectAccess) ProtoMessage() {} + +func (x *SignedObjectAccess) ProtoReflect() protoreflect.Message { + mi := &file_atelet_proto_msgTypes[30] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SignedObjectAccess.ProtoReflect.Descriptor instead. +func (*SignedObjectAccess) Descriptor() ([]byte, []int) { + return file_atelet_proto_rawDescGZIP(), []int{30} +} + +func (x *SignedObjectAccess) GetPrefixUrl() string { + if x != nil { + return x.PrefixUrl + } + return "" +} + +func (x *SignedObjectAccess) GetReadToken() string { + if x != nil { + return x.ReadToken + } + return "" +} + +func (x *SignedObjectAccess) GetReadObjectUrls() map[string]string { + if x != nil { + return x.ReadObjectUrls + } + return nil +} + +func (x *SignedObjectAccess) GetWriteMethod() string { + if x != nil { + return x.WriteMethod + } + return "" +} + +func (x *SignedObjectAccess) GetWriteToken() string { + if x != nil { + return x.WriteToken + } + return "" +} + +func (x *SignedObjectAccess) GetWriteHeaders() map[string]string { + if x != nil { + return x.WriteHeaders + } + return nil +} + +func (x *SignedObjectAccess) GetPostUrl() string { + if x != nil { + return x.PostUrl + } + return "" +} + +func (x *SignedObjectAccess) GetPostFields() map[string]string { + if x != nil { + return x.PostFields + } + return nil +} + type CheckpointRequest struct { state protoimpl.MessageState `protogen:"open.v1"` TargetAteomUid string `protobuf:"bytes,1,opt,name=target_ateom_uid,json=targetAteomUid,proto3" json:"target_ateom_uid,omitempty"` @@ -2015,14 +2131,19 @@ type CheckpointRequest struct { // *CheckpointRequest_ExternalConfig Config isCheckpointRequest_Config `protobuf_oneof:"config"` // What should be included in the checkpoint. - Scope SnapshotScope `protobuf:"varint,11,opt,name=scope,proto3,enum=atelet.SnapshotScope" json:"scope,omitempty"` + Scope SnapshotScope `protobuf:"varint,11,opt,name=scope,proto3,enum=atelet.SnapshotScope" json:"scope,omitempty"` + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to write that snapshot's objects over plain HTTP, + // with no cloud storage SDK or long-lived credential of its own. Empty means + // the node uses its built-in storage client. + SignedAccess map[string]*SignedObjectAccess `protobuf:"bytes,12,rep,name=signed_access,json=signedAccess,proto3" json:"signed_access,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *CheckpointRequest) Reset() { *x = CheckpointRequest{} - mi := &file_atelet_proto_msgTypes[30] + mi := &file_atelet_proto_msgTypes[31] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2034,7 +2155,7 @@ func (x *CheckpointRequest) String() string { func (*CheckpointRequest) ProtoMessage() {} func (x *CheckpointRequest) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[30] + mi := &file_atelet_proto_msgTypes[31] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2047,7 +2168,7 @@ func (x *CheckpointRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use CheckpointRequest.ProtoReflect.Descriptor instead. func (*CheckpointRequest) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{30} + return file_atelet_proto_rawDescGZIP(), []int{31} } func (x *CheckpointRequest) GetTargetAteomUid() string { @@ -2138,6 +2259,13 @@ func (x *CheckpointRequest) GetScope() SnapshotScope { return SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED } +func (x *CheckpointRequest) GetSignedAccess() map[string]*SignedObjectAccess { + if x != nil { + return x.SignedAccess + } + return nil +} + type isCheckpointRequest_Config interface { isCheckpointRequest_Config() } @@ -2162,7 +2290,7 @@ type CheckpointResponse struct { func (x *CheckpointResponse) Reset() { *x = CheckpointResponse{} - mi := &file_atelet_proto_msgTypes[31] + mi := &file_atelet_proto_msgTypes[32] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2174,7 +2302,7 @@ func (x *CheckpointResponse) String() string { func (*CheckpointResponse) ProtoMessage() {} func (x *CheckpointResponse) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[31] + mi := &file_atelet_proto_msgTypes[32] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2187,7 +2315,7 @@ func (x *CheckpointResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use CheckpointResponse.ProtoReflect.Descriptor instead. func (*CheckpointResponse) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{31} + return file_atelet_proto_rawDescGZIP(), []int{32} } type UploadPausedCheckpointRequest struct { @@ -2208,14 +2336,20 @@ type UploadPausedCheckpointRequest struct { // the local snapshot's own manifest, which is authoritative. When they // differ, atelet converts where possible (micro-VM FULL capture to a DATA // upload by selecting the durable-dir tar) and rejects otherwise. - DesiredScope SnapshotScope `protobuf:"varint,8,opt,name=desired_scope,json=desiredScope,proto3,enum=atelet.SnapshotScope" json:"desired_scope,omitempty"` + DesiredScope SnapshotScope `protobuf:"varint,8,opt,name=desired_scope,json=desiredScope,proto3,enum=atelet.SnapshotScope" json:"desired_scope,omitempty"` + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to write the paused snapshot's objects to their + // destination over plain HTTP, with no cloud storage SDK or long-lived + // credential of its own. Empty means the node uses its built-in storage + // client. + SignedAccess map[string]*SignedObjectAccess `protobuf:"bytes,9,rep,name=signed_access,json=signedAccess,proto3" json:"signed_access,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *UploadPausedCheckpointRequest) Reset() { *x = UploadPausedCheckpointRequest{} - mi := &file_atelet_proto_msgTypes[32] + mi := &file_atelet_proto_msgTypes[33] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2227,7 +2361,7 @@ func (x *UploadPausedCheckpointRequest) String() string { func (*UploadPausedCheckpointRequest) ProtoMessage() {} func (x *UploadPausedCheckpointRequest) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[32] + mi := &file_atelet_proto_msgTypes[33] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2240,7 +2374,7 @@ func (x *UploadPausedCheckpointRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use UploadPausedCheckpointRequest.ProtoReflect.Descriptor instead. func (*UploadPausedCheckpointRequest) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{32} + return file_atelet_proto_rawDescGZIP(), []int{33} } func (x *UploadPausedCheckpointRequest) GetAtespace() string { @@ -2299,6 +2433,13 @@ func (x *UploadPausedCheckpointRequest) GetDesiredScope() SnapshotScope { return SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED } +func (x *UploadPausedCheckpointRequest) GetSignedAccess() map[string]*SignedObjectAccess { + if x != nil { + return x.SignedAccess + } + return nil +} + type UploadPausedCheckpointResponse struct { state protoimpl.MessageState `protogen:"open.v1"` unknownFields protoimpl.UnknownFields @@ -2307,7 +2448,7 @@ type UploadPausedCheckpointResponse struct { func (x *UploadPausedCheckpointResponse) Reset() { *x = UploadPausedCheckpointResponse{} - mi := &file_atelet_proto_msgTypes[33] + mi := &file_atelet_proto_msgTypes[34] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2319,7 +2460,7 @@ func (x *UploadPausedCheckpointResponse) String() string { func (*UploadPausedCheckpointResponse) ProtoMessage() {} func (x *UploadPausedCheckpointResponse) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[33] + mi := &file_atelet_proto_msgTypes[34] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2332,7 +2473,7 @@ func (x *UploadPausedCheckpointResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use UploadPausedCheckpointResponse.ProtoReflect.Descriptor instead. func (*UploadPausedCheckpointResponse) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{33} + return file_atelet_proto_rawDescGZIP(), []int{34} } type RestoreRequest struct { @@ -2370,15 +2511,22 @@ type RestoreRequest struct { // gVisor and micro-VM DATA-scope restores the sandbox is (re)sized to these; // for a FULL micro-VM restore the size baked into the snapshot wins. Zero // means "unset": keep the runtime default. - CpuMilli int64 `protobuf:"varint,14,opt,name=cpu_milli,json=cpuMilli,proto3" json:"cpu_milli,omitempty"` // CPU limit in millicores (1000 = one core). - MemoryBytes int64 `protobuf:"varint,15,opt,name=memory_bytes,json=memoryBytes,proto3" json:"memory_bytes,omitempty"` // Memory limit in bytes. + CpuMilli int64 `protobuf:"varint,14,opt,name=cpu_milli,json=cpuMilli,proto3" json:"cpu_milli,omitempty"` // CPU limit in millicores (1000 = one core). + MemoryBytes int64 `protobuf:"varint,15,opt,name=memory_bytes,json=memoryBytes,proto3" json:"memory_bytes,omitempty"` // Memory limit in bytes. + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to read that snapshot's objects over plain HTTP, + // with no cloud storage SDK or long-lived credential of its own. A restore + // that combines the actor's own snapshot with a golden snapshot carries one + // entry per snapshot it reads. Empty means the node uses its built-in + // storage client. + SignedAccess map[string]*SignedObjectAccess `protobuf:"bytes,16,rep,name=signed_access,json=signedAccess,proto3" json:"signed_access,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *RestoreRequest) Reset() { *x = RestoreRequest{} - mi := &file_atelet_proto_msgTypes[34] + mi := &file_atelet_proto_msgTypes[35] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2390,7 +2538,7 @@ func (x *RestoreRequest) String() string { func (*RestoreRequest) ProtoMessage() {} func (x *RestoreRequest) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[34] + mi := &file_atelet_proto_msgTypes[35] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2403,7 +2551,7 @@ func (x *RestoreRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RestoreRequest.ProtoReflect.Descriptor instead. func (*RestoreRequest) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{34} + return file_atelet_proto_rawDescGZIP(), []int{35} } func (x *RestoreRequest) GetTargetAteomUid() string { @@ -2522,6 +2670,13 @@ func (x *RestoreRequest) GetMemoryBytes() int64 { return 0 } +func (x *RestoreRequest) GetSignedAccess() map[string]*SignedObjectAccess { + if x != nil { + return x.SignedAccess + } + return nil +} + type isRestoreRequest_Config interface { isRestoreRequest_Config() } @@ -2546,7 +2701,7 @@ type RestoreResponse struct { func (x *RestoreResponse) Reset() { *x = RestoreResponse{} - mi := &file_atelet_proto_msgTypes[35] + mi := &file_atelet_proto_msgTypes[36] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2558,7 +2713,7 @@ func (x *RestoreResponse) String() string { func (*RestoreResponse) ProtoMessage() {} func (x *RestoreResponse) ProtoReflect() protoreflect.Message { - mi := &file_atelet_proto_msgTypes[35] + mi := &file_atelet_proto_msgTypes[36] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2571,7 +2726,7 @@ func (x *RestoreResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RestoreResponse.ProtoReflect.Descriptor instead. func (*RestoreResponse) Descriptor() ([]byte, []int) { - return file_atelet_proto_rawDescGZIP(), []int{35} + return file_atelet_proto_rawDescGZIP(), []int{36} } var File_atelet_proto protoreflect.FileDescriptor @@ -2703,9 +2858,31 @@ const file_atelet_proto_rawDesc = "" + "\x04port\x18\x02 \x01(\x05R\x04port\"\r\n" + "\vRunResponse\"C\n" + "\x1cLocalCheckpointConfiguration\x12#\n" + - "\rsnapshot_name\x18\x01 \x01(\tR\fsnapshotName\"D\n" + + "\rsnapshot_name\x18\x01 \x01(\tR\fsnapshotName\"Y\n" + "\x1fExternalCheckpointConfiguration\x12!\n" + - "\fsnapshot_uri\x18\x01 \x01(\tR\vsnapshotUri\"\xab\x04\n" + + "\fsnapshot_uri\x18\x01 \x01(\tR\vsnapshotUriJ\x04\b\x02\x10\x03R\rsigned_access\"\xee\x04\n" + + "\x12SignedObjectAccess\x12\x1d\n" + + "\n" + + "prefix_url\x18\x01 \x01(\tR\tprefixUrl\x12\x1d\n" + + "\n" + + "read_token\x18\x02 \x01(\tR\treadToken\x12X\n" + + "\x10read_object_urls\x18\x03 \x03(\v2..atelet.SignedObjectAccess.ReadObjectUrlsEntryR\x0ereadObjectUrls\x12!\n" + + "\fwrite_method\x18\x04 \x01(\tR\vwriteMethod\x12\x1f\n" + + "\vwrite_token\x18\x05 \x01(\tR\n" + + "writeToken\x12Q\n" + + "\rwrite_headers\x18\x06 \x03(\v2,.atelet.SignedObjectAccess.WriteHeadersEntryR\fwriteHeaders\x12\x19\n" + + "\bpost_url\x18\a \x01(\tR\apostUrl\x12K\n" + + "\vpost_fields\x18\b \x03(\v2*.atelet.SignedObjectAccess.PostFieldsEntryR\n" + + "postFields\x1aA\n" + + "\x13ReadObjectUrlsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1a?\n" + + "\x11WriteHeadersEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1a=\n" + + "\x0fPostFieldsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xda\x05\n" + "\x11CheckpointRequest\x12(\n" + "\x10target_ateom_uid\x18\x01 \x01(\tR\x0etargetAteomUid\x12\x1a\n" + "\batespace\x18\x02 \x01(\tR\batespace\x12\x1d\n" + @@ -2719,9 +2896,13 @@ const file_atelet_proto_rawDesc = "" + "\flocal_config\x18\t \x01(\v2$.atelet.LocalCheckpointConfigurationH\x00R\vlocalConfig\x12R\n" + "\x0fexternal_config\x18\n" + " \x01(\v2'.atelet.ExternalCheckpointConfigurationH\x00R\x0eexternalConfig\x12+\n" + - "\x05scope\x18\v \x01(\x0e2\x15.atelet.SnapshotScopeR\x05scopeB\b\n" + + "\x05scope\x18\v \x01(\x0e2\x15.atelet.SnapshotScopeR\x05scope\x12P\n" + + "\rsigned_access\x18\f \x03(\v2+.atelet.CheckpointRequest.SignedAccessEntryR\fsignedAccess\x1a[\n" + + "\x11SignedAccessEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x120\n" + + "\x05value\x18\x02 \x01(\v2\x1a.atelet.SignedObjectAccessR\x05value:\x028\x01B\b\n" + "\x06config\"\x14\n" + - "\x12CheckpointResponse\"\x87\x03\n" + + "\x12CheckpointResponse\"\xc2\x04\n" + "\x1dUploadPausedCheckpointRequest\x12\x1a\n" + "\batespace\x18\x01 \x01(\tR\batespace\x12\x1d\n" + "\n" + @@ -2731,8 +2912,12 @@ const file_atelet_proto_rawDesc = "" + "\x13actor_template_name\x18\x05 \x01(\tR\x11actorTemplateName\x12.\n" + "\x13local_snapshot_name\x18\x06 \x01(\tR\x11localSnapshotName\x128\n" + "\x18destination_snapshot_uri\x18\a \x01(\tR\x16destinationSnapshotUri\x12:\n" + - "\rdesired_scope\x18\b \x01(\x0e2\x15.atelet.SnapshotScopeR\fdesiredScope\" \n" + - "\x1eUploadPausedCheckpointResponse\"\xee\x05\n" + + "\rdesired_scope\x18\b \x01(\x0e2\x15.atelet.SnapshotScopeR\fdesiredScope\x12\\\n" + + "\rsigned_access\x18\t \x03(\v27.atelet.UploadPausedCheckpointRequest.SignedAccessEntryR\fsignedAccess\x1a[\n" + + "\x11SignedAccessEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x120\n" + + "\x05value\x18\x02 \x01(\v2\x1a.atelet.SignedObjectAccessR\x05value:\x028\x01\" \n" + + "\x1eUploadPausedCheckpointResponse\"\x9a\a\n" + "\x0eRestoreRequest\x12(\n" + "\x10target_ateom_uid\x18\x01 \x01(\tR\x0etargetAteomUid\x12\x1a\n" + "\batespace\x18\x02 \x01(\tR\batespace\x12\x1d\n" + @@ -2750,7 +2935,11 @@ const file_atelet_proto_rawDesc = "" + "\x13golden_snapshot_uri\x18\f \x01(\tR\x11goldenSnapshotUri\x12A\n" + "\x0eegress_gateway\x18\r \x01(\v2\x15.atelet.EgressGatewayH\x01R\regressGateway\x88\x01\x01\x12\x1b\n" + "\tcpu_milli\x18\x0e \x01(\x03R\bcpuMilli\x12!\n" + - "\fmemory_bytes\x18\x0f \x01(\x03R\vmemoryBytesB\b\n" + + "\fmemory_bytes\x18\x0f \x01(\x03R\vmemoryBytes\x12M\n" + + "\rsigned_access\x18\x10 \x03(\v2(.atelet.RestoreRequest.SignedAccessEntryR\fsignedAccess\x1a[\n" + + "\x11SignedAccessEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x120\n" + + "\x05value\x18\x02 \x01(\v2\x1a.atelet.SignedObjectAccessR\x05value:\x028\x01B\b\n" + "\x06configB\x11\n" + "\x0f_egress_gateway\"\x11\n" + "\x0fRestoreResponse*\x9a\x01\n" + @@ -2791,7 +2980,7 @@ func file_atelet_proto_rawDescGZIP() []byte { } var file_atelet_proto_enumTypes = make([]protoimpl.EnumInfo, 3) -var file_atelet_proto_msgTypes = make([]protoimpl.MessageInfo, 39) +var file_atelet_proto_msgTypes = make([]protoimpl.MessageInfo, 46) var file_atelet_proto_goTypes = []any{ (ActorMetadataField)(0), // 0: atelet.ActorMetadataField (CheckpointType)(0), // 1: atelet.CheckpointType @@ -2826,26 +3015,33 @@ var file_atelet_proto_goTypes = []any{ (*RunResponse)(nil), // 30: atelet.RunResponse (*LocalCheckpointConfiguration)(nil), // 31: atelet.LocalCheckpointConfiguration (*ExternalCheckpointConfiguration)(nil), // 32: atelet.ExternalCheckpointConfiguration - (*CheckpointRequest)(nil), // 33: atelet.CheckpointRequest - (*CheckpointResponse)(nil), // 34: atelet.CheckpointResponse - (*UploadPausedCheckpointRequest)(nil), // 35: atelet.UploadPausedCheckpointRequest - (*UploadPausedCheckpointResponse)(nil), // 36: atelet.UploadPausedCheckpointResponse - (*RestoreRequest)(nil), // 37: atelet.RestoreRequest - (*RestoreResponse)(nil), // 38: atelet.RestoreResponse - nil, // 39: atelet.ArchAssets.FilesEntry - nil, // 40: atelet.SandboxAssets.AssetsEntry - nil, // 41: atelet.ExternalVolumeSource.VolumeContextEntry + (*SignedObjectAccess)(nil), // 33: atelet.SignedObjectAccess + (*CheckpointRequest)(nil), // 34: atelet.CheckpointRequest + (*CheckpointResponse)(nil), // 35: atelet.CheckpointResponse + (*UploadPausedCheckpointRequest)(nil), // 36: atelet.UploadPausedCheckpointRequest + (*UploadPausedCheckpointResponse)(nil), // 37: atelet.UploadPausedCheckpointResponse + (*RestoreRequest)(nil), // 38: atelet.RestoreRequest + (*RestoreResponse)(nil), // 39: atelet.RestoreResponse + nil, // 40: atelet.ArchAssets.FilesEntry + nil, // 41: atelet.SandboxAssets.AssetsEntry + nil, // 42: atelet.ExternalVolumeSource.VolumeContextEntry + nil, // 43: atelet.SignedObjectAccess.ReadObjectUrlsEntry + nil, // 44: atelet.SignedObjectAccess.WriteHeadersEntry + nil, // 45: atelet.SignedObjectAccess.PostFieldsEntry + nil, // 46: atelet.CheckpointRequest.SignedAccessEntry + nil, // 47: atelet.UploadPausedCheckpointRequest.SignedAccessEntry + nil, // 48: atelet.RestoreRequest.SignedAccessEntry } var file_atelet_proto_depIdxs = []int32{ 12, // 0: atelet.TerminateRequest.spec:type_name -> atelet.WorkloadSpec 12, // 1: atelet.RunRequest.spec:type_name -> atelet.WorkloadSpec 11, // 2: atelet.RunRequest.sandbox_assets:type_name -> atelet.SandboxAssets 8, // 3: atelet.RunRequest.egress_gateway:type_name -> atelet.EgressGateway - 39, // 4: atelet.ArchAssets.files:type_name -> atelet.ArchAssets.FilesEntry - 40, // 5: atelet.SandboxAssets.assets:type_name -> atelet.SandboxAssets.AssetsEntry + 40, // 4: atelet.ArchAssets.files:type_name -> atelet.ArchAssets.FilesEntry + 41, // 5: atelet.SandboxAssets.assets:type_name -> atelet.SandboxAssets.AssetsEntry 23, // 6: atelet.WorkloadSpec.containers:type_name -> atelet.Container 21, // 7: atelet.WorkloadSpec.volumes:type_name -> atelet.Volume - 41, // 8: atelet.ExternalVolumeSource.volume_context:type_name -> atelet.ExternalVolumeSource.VolumeContextEntry + 42, // 8: atelet.ExternalVolumeSource.volume_context:type_name -> atelet.ExternalVolumeSource.VolumeContextEntry 0, // 9: atelet.ActorMetadataItem.field:type_name -> atelet.ActorMetadataField 16, // 10: atelet.ActorMetadataDataSource.items:type_name -> atelet.ActorMetadataItem 17, // 11: atelet.SystemInfoDataSource.actor_metadata:type_name -> atelet.ActorMetadataDataSource @@ -2862,37 +3058,46 @@ var file_atelet_proto_depIdxs = []int32{ 26, // 22: atelet.Container.resources:type_name -> atelet.ResourceLimits 25, // 23: atelet.SecurityContext.capabilities:type_name -> atelet.Capabilities 29, // 24: atelet.Readyz.http_get:type_name -> atelet.HTTPGetAction - 12, // 25: atelet.CheckpointRequest.spec:type_name -> atelet.WorkloadSpec - 1, // 26: atelet.CheckpointRequest.type:type_name -> atelet.CheckpointType - 31, // 27: atelet.CheckpointRequest.local_config:type_name -> atelet.LocalCheckpointConfiguration - 32, // 28: atelet.CheckpointRequest.external_config:type_name -> atelet.ExternalCheckpointConfiguration - 2, // 29: atelet.CheckpointRequest.scope:type_name -> atelet.SnapshotScope - 2, // 30: atelet.UploadPausedCheckpointRequest.desired_scope:type_name -> atelet.SnapshotScope - 12, // 31: atelet.RestoreRequest.spec:type_name -> atelet.WorkloadSpec - 1, // 32: atelet.RestoreRequest.type:type_name -> atelet.CheckpointType - 31, // 33: atelet.RestoreRequest.local_config:type_name -> atelet.LocalCheckpointConfiguration - 32, // 34: atelet.RestoreRequest.external_config:type_name -> atelet.ExternalCheckpointConfiguration - 2, // 35: atelet.RestoreRequest.scope:type_name -> atelet.SnapshotScope - 8, // 36: atelet.RestoreRequest.egress_gateway:type_name -> atelet.EgressGateway - 9, // 37: atelet.ArchAssets.FilesEntry.value:type_name -> atelet.AssetFile - 10, // 38: atelet.SandboxAssets.AssetsEntry.value:type_name -> atelet.ArchAssets - 3, // 39: atelet.CredentialBroker.MintActorCertificate:input_type -> atelet.MintActorCertificateRequest - 7, // 40: atelet.AteomHerder.Run:input_type -> atelet.RunRequest - 33, // 41: atelet.AteomHerder.Checkpoint:input_type -> atelet.CheckpointRequest - 37, // 42: atelet.AteomHerder.Restore:input_type -> atelet.RestoreRequest - 35, // 43: atelet.AteomHerder.UploadPausedCheckpoint:input_type -> atelet.UploadPausedCheckpointRequest - 5, // 44: atelet.AteomHerder.Terminate:input_type -> atelet.TerminateRequest - 4, // 45: atelet.CredentialBroker.MintActorCertificate:output_type -> atelet.MintActorCertificateResponse - 30, // 46: atelet.AteomHerder.Run:output_type -> atelet.RunResponse - 34, // 47: atelet.AteomHerder.Checkpoint:output_type -> atelet.CheckpointResponse - 38, // 48: atelet.AteomHerder.Restore:output_type -> atelet.RestoreResponse - 36, // 49: atelet.AteomHerder.UploadPausedCheckpoint:output_type -> atelet.UploadPausedCheckpointResponse - 6, // 50: atelet.AteomHerder.Terminate:output_type -> atelet.TerminateResponse - 45, // [45:51] is the sub-list for method output_type - 39, // [39:45] is the sub-list for method input_type - 39, // [39:39] is the sub-list for extension type_name - 39, // [39:39] is the sub-list for extension extendee - 0, // [0:39] is the sub-list for field type_name + 43, // 25: atelet.SignedObjectAccess.read_object_urls:type_name -> atelet.SignedObjectAccess.ReadObjectUrlsEntry + 44, // 26: atelet.SignedObjectAccess.write_headers:type_name -> atelet.SignedObjectAccess.WriteHeadersEntry + 45, // 27: atelet.SignedObjectAccess.post_fields:type_name -> atelet.SignedObjectAccess.PostFieldsEntry + 12, // 28: atelet.CheckpointRequest.spec:type_name -> atelet.WorkloadSpec + 1, // 29: atelet.CheckpointRequest.type:type_name -> atelet.CheckpointType + 31, // 30: atelet.CheckpointRequest.local_config:type_name -> atelet.LocalCheckpointConfiguration + 32, // 31: atelet.CheckpointRequest.external_config:type_name -> atelet.ExternalCheckpointConfiguration + 2, // 32: atelet.CheckpointRequest.scope:type_name -> atelet.SnapshotScope + 46, // 33: atelet.CheckpointRequest.signed_access:type_name -> atelet.CheckpointRequest.SignedAccessEntry + 2, // 34: atelet.UploadPausedCheckpointRequest.desired_scope:type_name -> atelet.SnapshotScope + 47, // 35: atelet.UploadPausedCheckpointRequest.signed_access:type_name -> atelet.UploadPausedCheckpointRequest.SignedAccessEntry + 12, // 36: atelet.RestoreRequest.spec:type_name -> atelet.WorkloadSpec + 1, // 37: atelet.RestoreRequest.type:type_name -> atelet.CheckpointType + 31, // 38: atelet.RestoreRequest.local_config:type_name -> atelet.LocalCheckpointConfiguration + 32, // 39: atelet.RestoreRequest.external_config:type_name -> atelet.ExternalCheckpointConfiguration + 2, // 40: atelet.RestoreRequest.scope:type_name -> atelet.SnapshotScope + 8, // 41: atelet.RestoreRequest.egress_gateway:type_name -> atelet.EgressGateway + 48, // 42: atelet.RestoreRequest.signed_access:type_name -> atelet.RestoreRequest.SignedAccessEntry + 9, // 43: atelet.ArchAssets.FilesEntry.value:type_name -> atelet.AssetFile + 10, // 44: atelet.SandboxAssets.AssetsEntry.value:type_name -> atelet.ArchAssets + 33, // 45: atelet.CheckpointRequest.SignedAccessEntry.value:type_name -> atelet.SignedObjectAccess + 33, // 46: atelet.UploadPausedCheckpointRequest.SignedAccessEntry.value:type_name -> atelet.SignedObjectAccess + 33, // 47: atelet.RestoreRequest.SignedAccessEntry.value:type_name -> atelet.SignedObjectAccess + 3, // 48: atelet.CredentialBroker.MintActorCertificate:input_type -> atelet.MintActorCertificateRequest + 7, // 49: atelet.AteomHerder.Run:input_type -> atelet.RunRequest + 34, // 50: atelet.AteomHerder.Checkpoint:input_type -> atelet.CheckpointRequest + 38, // 51: atelet.AteomHerder.Restore:input_type -> atelet.RestoreRequest + 36, // 52: atelet.AteomHerder.UploadPausedCheckpoint:input_type -> atelet.UploadPausedCheckpointRequest + 5, // 53: atelet.AteomHerder.Terminate:input_type -> atelet.TerminateRequest + 4, // 54: atelet.CredentialBroker.MintActorCertificate:output_type -> atelet.MintActorCertificateResponse + 30, // 55: atelet.AteomHerder.Run:output_type -> atelet.RunResponse + 35, // 56: atelet.AteomHerder.Checkpoint:output_type -> atelet.CheckpointResponse + 39, // 57: atelet.AteomHerder.Restore:output_type -> atelet.RestoreResponse + 37, // 58: atelet.AteomHerder.UploadPausedCheckpoint:output_type -> atelet.UploadPausedCheckpointResponse + 6, // 59: atelet.AteomHerder.Terminate:output_type -> atelet.TerminateResponse + 54, // [54:60] is the sub-list for method output_type + 48, // [48:54] is the sub-list for method input_type + 48, // [48:48] is the sub-list for extension type_name + 48, // [48:48] is the sub-list for extension extendee + 0, // [0:48] is the sub-list for field type_name } func init() { file_atelet_proto_init() } @@ -2911,11 +3116,11 @@ func file_atelet_proto_init() { (*Volume_SystemInfo)(nil), (*Volume_Image)(nil), } - file_atelet_proto_msgTypes[30].OneofWrappers = []any{ + file_atelet_proto_msgTypes[31].OneofWrappers = []any{ (*CheckpointRequest_LocalConfig)(nil), (*CheckpointRequest_ExternalConfig)(nil), } - file_atelet_proto_msgTypes[34].OneofWrappers = []any{ + file_atelet_proto_msgTypes[35].OneofWrappers = []any{ (*RestoreRequest_LocalConfig)(nil), (*RestoreRequest_ExternalConfig)(nil), } @@ -2925,7 +3130,7 @@ func file_atelet_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_atelet_proto_rawDesc), len(file_atelet_proto_rawDesc)), NumEnums: 3, - NumMessages: 39, + NumMessages: 46, NumExtensions: 0, NumServices: 2, }, diff --git a/internal/proto/ateletpb/atelet.proto b/internal/proto/ateletpb/atelet.proto index b189cc50ec..df891bf619 100644 --- a/internal/proto/ateletpb/atelet.proto +++ b/internal/proto/ateletpb/atelet.proto @@ -300,6 +300,41 @@ message ExternalCheckpointConfiguration { // The object storage URI of the snapshot to write. Object names are appended // to it, so it addresses the snapshot as a whole rather than any one object. string snapshot_uri = 1; + + // Field 2 previously held a single SignedObjectAccess. Signed access is now a + // per-snapshot-URI map on the request messages (RestoreRequest, + // CheckpointRequest, UploadPausedCheckpointRequest), so a restore that reads + // more than one snapshot can carry a capability for each. + reserved 2; + reserved "signed_access"; +} + +// SignedObjectAccess is a short-lived, prefix-scoped capability that lets atelet +// read or write a single snapshot's objects with no cloud credential on the +// node. atelet dispatches on the mechanics carried here, never on the cloud: +// the control-plane broker fills in whichever fields its backend needs. +message SignedObjectAccess { + // Base URL of the account/bucket+container, e.g. + // "https://acct.blob.core.windows.net/snapshots" or "http://s3-host/bucket". + string prefix_url = 1; + + // Read mechanics (populated on restore). atelet fetches an object by either + // looking it up in read_object_urls (S3: per-object presigned GET), or + // appending "?" to prefix_url (Azure: one container SAS). + string read_token = 2; + map read_object_urls = 3; + + // Write mechanics (populated on checkpoint). write_method is "PUT" or "POST". + // PUT (Azure): atelet PUTs "/?" with + // write_headers (block blob needs x-ms-blob-type: BlockBlob). + // POST (S3): atelet submits a multipart form to post_url with post_fields, + // setting the "key" field per object; one signed starts-with policy + // covers every file under the prefix. + string write_method = 4; + string write_token = 5; + map write_headers = 6; + string post_url = 7; + map post_fields = 8; } enum CheckpointType { @@ -355,6 +390,12 @@ message CheckpointRequest { // What should be included in the checkpoint. SnapshotScope scope = 11; + + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to write that snapshot's objects over plain HTTP, + // with no cloud storage SDK or long-lived credential of its own. Empty means + // the node uses its built-in storage client. + map signed_access = 12; } message CheckpointResponse { @@ -382,6 +423,13 @@ message UploadPausedCheckpointRequest { // differ, atelet converts where possible (micro-VM FULL capture to a DATA // upload by selecting the durable-dir tar) and rejects otherwise. SnapshotScope desired_scope = 8; + + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to write the paused snapshot's objects to their + // destination over plain HTTP, with no cloud storage SDK or long-lived + // credential of its own. Empty means the node uses its built-in storage + // client. + map signed_access = 9; } message UploadPausedCheckpointResponse { @@ -430,6 +478,14 @@ message RestoreRequest { // means "unset": keep the runtime default. int64 cpu_milli = 14; // CPU limit in millicores (1000 = one core). int64 memory_bytes = 15; // Memory limit in bytes. + + // signed_access maps a snapshot URI to a short-lived, prefix-scoped + // capability the node uses to read that snapshot's objects over plain HTTP, + // with no cloud storage SDK or long-lived credential of its own. A restore + // that combines the actor's own snapshot with a golden snapshot carries one + // entry per snapshot it reads. Empty means the node uses its built-in + // storage client. + map signed_access = 16; } message RestoreResponse { From 449df38825a2add3e2ab96b0c9f62943638e9a79 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 22:18:30 -0700 Subject: [PATCH 2/8] feat(ateapi): add pluggable storage broker interface and registry --- cmd/ateapi/internal/storagebroker/broker.go | 110 ++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 cmd/ateapi/internal/storagebroker/broker.go diff --git a/cmd/ateapi/internal/storagebroker/broker.go b/cmd/ateapi/internal/storagebroker/broker.go new file mode 100644 index 0000000000..a64f6ff836 --- /dev/null +++ b/cmd/ateapi/internal/storagebroker/broker.go @@ -0,0 +1,110 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package storagebroker is the vendor-neutral seam by which ate-api-server +// mints short-lived, snapshot-scoped capabilities (signed URLs) for atelet to +// read and write snapshot objects over plain HTTP. Core imports only this +// interface and a registry; a cloud-specific implementation registers itself +// from an out-of-tree build (see the azure sub-implementation, which links the +// Azure SDK). This keeps the control-plane binary free of any cloud SDK unless +// a provider is compiled in, and keeps the node (atelet) free of one entirely. +package storagebroker + +import ( + "context" + "fmt" + "sort" + "sync" + "time" +) + +// Capability is a control-plane-minted, snapshot-scoped grant. atelet forms a +// per-object request from these mechanics; it needs no cloud credential of its +// own. A read capability (from MintRead) populates the read fields; a write +// capability (from MintWrite) populates the write fields. +type Capability struct { + // PrefixURL is the account/bucket base URL, e.g. + // "https://acct.blob.core.windows.net/snapshots" or "http://s3-host/bucket". + PrefixURL string + + // Read mechanics. Either a token appended to / (Azure), or a + // per-object presigned URL map (S3). + ReadToken string + ReadObjectURLs map[string]string + + // Write mechanics. WriteMethod is "PUT" (Azure: append WriteToken, send + // WriteHeaders) or "POST" (S3: multipart form to PostURL with PostFields). + WriteMethod string + WriteToken string + WriteHeaders map[string]string + PostURL string + PostFields map[string]string +} + +// Write methods a Capability may carry. +const ( + WriteMethodPUT = "PUT" + WriteMethodPOST = "POST" +) + +// Broker mints read/write capabilities for one snapshot, identified by its URI +// prefix. Implementations authorize and sign; they never touch snapshot bytes. +type Broker interface { + MintRead(ctx context.Context, snapshotURI string, ttl time.Duration) (Capability, error) + MintWrite(ctx context.Context, snapshotURI string, ttl time.Duration) (Capability, error) +} + +// Factory builds a Broker. Cloud implementations register one under a backend +// name; New selects it. Registration keeps core free of cloud SDK imports. +type Factory func(ctx context.Context) (Broker, error) + +var ( + mu sync.RWMutex + factories = map[string]Factory{} +) + +// Register adds a broker backend. Implementations call this from an init in +// their own (out-of-tree) package. +func Register(name string, f Factory) { + mu.Lock() + defer mu.Unlock() + factories[name] = f +} + +// New constructs the broker registered under backend, or an error naming the +// backends that are compiled in. An empty backend disables signing (nil, nil), +// so callers fall back to the legacy in-atelet storage client. +func New(ctx context.Context, backend string) (Broker, error) { + if backend == "" { + return nil, nil + } + mu.RLock() + f, ok := factories[backend] + mu.RUnlock() + if !ok { + return nil, fmt.Errorf("storage broker backend %q not registered (have: %v)", backend, registered()) + } + return f(ctx) +} + +func registered() []string { + mu.RLock() + defer mu.RUnlock() + names := make([]string, 0, len(factories)) + for n := range factories { + names = append(names, n) + } + sort.Strings(names) + return names +} From dcc57e8e918df70dab413339bf7d6caab5fc8cf6 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 22:18:31 -0700 Subject: [PATCH 3/8] feat(ateapi): pre-mint per-snapshot signed URLs at resume and suspend --- .../internal/controlapi/broker_wiring.go | 87 +++++++++++ .../internal/controlapi/broker_wiring_test.go | 139 ++++++++++++++++++ .../controlapi/functionaltest/common_test.go | 2 +- cmd/ateapi/internal/controlapi/service.go | 4 +- cmd/ateapi/internal/controlapi/workflow.go | 6 + .../internal/controlapi/workflow_resume.go | 10 ++ .../internal/controlapi/workflow_suspend.go | 9 ++ .../controlapi/workflow_suspend_test.go | 2 +- .../controlapi/workflow_testutil_test.go | 2 +- cmd/ateapi/main.go | 9 +- 10 files changed, 265 insertions(+), 5 deletions(-) create mode 100644 cmd/ateapi/internal/controlapi/broker_wiring.go create mode 100644 cmd/ateapi/internal/controlapi/broker_wiring_test.go diff --git a/cmd/ateapi/internal/controlapi/broker_wiring.go b/cmd/ateapi/internal/controlapi/broker_wiring.go new file mode 100644 index 0000000000..3f42540c01 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/broker_wiring.go @@ -0,0 +1,87 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "fmt" + "time" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/storagebroker" + "github.com/agent-substrate/substrate/internal/proto/ateletpb" +) + +// snapshotCapabilityTTL bounds how long a minted snapshot read/write capability +// is valid. It must comfortably exceed a restore or checkpoint transfer, but +// stay short so a leaked URL exposes one snapshot only briefly. +const snapshotCapabilityTTL = 15 * time.Minute + +// readAccessFor mints a read capability for each snapshot URI the node will +// read, returning them keyed by URI. A restore that combines the actor's own +// snapshot with a golden snapshot passes both; the caller may pass an unset +// (empty) URI unconditionally and it is skipped. With no broker configured it +// returns nil and the node reads with its built-in storage client. +func (w *ActorWorkflow) readAccessFor(ctx context.Context, snapshotURIs ...string) (map[string]*ateletpb.SignedObjectAccess, error) { + if w.broker == nil { + return nil, nil + } + return mintAccess(ctx, w.broker.MintRead, snapshotURIs...) +} + +// writeAccessFor mints a write capability for each snapshot URI the node will +// write, returning them keyed by URI. With no broker configured it returns nil. +func (w *ActorWorkflow) writeAccessFor(ctx context.Context, snapshotURIs ...string) (map[string]*ateletpb.SignedObjectAccess, error) { + if w.broker == nil { + return nil, nil + } + return mintAccess(ctx, w.broker.MintWrite, snapshotURIs...) +} + +// mintAccess mints one capability per distinct, non-empty URI with mint and maps +// each onto the wire message keyed by URI. It returns nil (not an empty map) +// when nothing was minted, so a caller can assign the result directly. +func mintAccess(ctx context.Context, mint func(context.Context, string, time.Duration) (storagebroker.Capability, error), snapshotURIs ...string) (map[string]*ateletpb.SignedObjectAccess, error) { + out := make(map[string]*ateletpb.SignedObjectAccess, len(snapshotURIs)) + for _, uri := range snapshotURIs { + if uri == "" || out[uri] != nil { + continue + } + cap, err := mint(ctx, uri, snapshotCapabilityTTL) + if err != nil { + return nil, fmt.Errorf("minting snapshot capability for %q: %w", uri, err) + } + out[uri] = signedAccessFromCapability(cap) + } + if len(out) == 0 { + return nil, nil + } + return out, nil +} + +// signedAccessFromCapability maps a broker capability onto the wire message +// atelet consumes. A read capability populates the read fields, a write +// capability the write fields; the rest stay empty. +func signedAccessFromCapability(cap storagebroker.Capability) *ateletpb.SignedObjectAccess { + return &ateletpb.SignedObjectAccess{ + PrefixUrl: cap.PrefixURL, + ReadToken: cap.ReadToken, + ReadObjectUrls: cap.ReadObjectURLs, + WriteMethod: cap.WriteMethod, + WriteToken: cap.WriteToken, + WriteHeaders: cap.WriteHeaders, + PostUrl: cap.PostURL, + PostFields: cap.PostFields, + } +} diff --git a/cmd/ateapi/internal/controlapi/broker_wiring_test.go b/cmd/ateapi/internal/controlapi/broker_wiring_test.go new file mode 100644 index 0000000000..ee062e9b9f --- /dev/null +++ b/cmd/ateapi/internal/controlapi/broker_wiring_test.go @@ -0,0 +1,139 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/storagebroker" +) + +// fakeBroker records the URIs it is asked to sign and returns a capability whose +// tokens embed the URI, so a test can tell entries apart. +type fakeBroker struct { + err error + mintedRead []string + mintedWrite []string +} + +func (f *fakeBroker) MintRead(_ context.Context, uri string, _ time.Duration) (storagebroker.Capability, error) { + if f.err != nil { + return storagebroker.Capability{}, f.err + } + f.mintedRead = append(f.mintedRead, uri) + return storagebroker.Capability{PrefixURL: "https://read/" + uri, ReadToken: "rt-" + uri}, nil +} + +func (f *fakeBroker) MintWrite(_ context.Context, uri string, _ time.Duration) (storagebroker.Capability, error) { + if f.err != nil { + return storagebroker.Capability{}, f.err + } + f.mintedWrite = append(f.mintedWrite, uri) + return storagebroker.Capability{ + PrefixURL: "https://write/" + uri, + WriteMethod: storagebroker.WriteMethodPUT, + WriteToken: "wt-" + uri, + }, nil +} + +func TestReadAccessFor_MintsOnePerURI(t *testing.T) { + fb := &fakeBroker{} + w := &ActorWorkflow{broker: fb} + + // The empty URI (an unset golden) must be skipped, not signed. + m, err := w.readAccessFor(context.Background(), "s3://bucket/actor", "s3://bucket/golden", "") + if err != nil { + t.Fatalf("readAccessFor: %v", err) + } + if len(m) != 2 { + t.Fatalf("want 2 capabilities, got %d: %v", len(m), m) + } + if got := m["s3://bucket/actor"].GetReadToken(); got != "rt-s3://bucket/actor" { + t.Errorf("actor read token = %q, want rt-s3://bucket/actor", got) + } + if got := m["s3://bucket/golden"].GetReadToken(); got != "rt-s3://bucket/golden" { + t.Errorf("golden read token = %q, want rt-s3://bucket/golden", got) + } + if _, ok := m[""]; ok { + t.Errorf("empty URI must not be signed") + } +} + +func TestReadAccessFor_DedupsRepeatedURI(t *testing.T) { + fb := &fakeBroker{} + w := &ActorWorkflow{broker: fb} + + m, err := w.readAccessFor(context.Background(), "s3://bucket/a", "s3://bucket/a") + if err != nil { + t.Fatalf("readAccessFor: %v", err) + } + if len(m) != 1 { + t.Fatalf("want 1 capability, got %d", len(m)) + } + if len(fb.mintedRead) != 1 { + t.Errorf("broker minted %d times, want 1 (deduped)", len(fb.mintedRead)) + } +} + +func TestReadWriteAccessFor_NilBrokerReturnsNil(t *testing.T) { + w := &ActorWorkflow{} // no broker configured + + if m, err := w.readAccessFor(context.Background(), "s3://bucket/a"); err != nil || m != nil { + t.Errorf("readAccessFor with no broker = (%v, %v), want (nil, nil)", m, err) + } + if m, err := w.writeAccessFor(context.Background(), "s3://bucket/a"); err != nil || m != nil { + t.Errorf("writeAccessFor with no broker = (%v, %v), want (nil, nil)", m, err) + } +} + +func TestAccessFor_AllEmptyReturnsNil(t *testing.T) { + w := &ActorWorkflow{broker: &fakeBroker{}} + + m, err := w.readAccessFor(context.Background(), "", "") + if err != nil { + t.Fatalf("readAccessFor: %v", err) + } + if m != nil { + t.Errorf("want nil map when nothing to sign, got %v", m) + } +} + +func TestWriteAccessFor_MintsWriteCapability(t *testing.T) { + fb := &fakeBroker{} + w := &ActorWorkflow{broker: fb} + + m, err := w.writeAccessFor(context.Background(), "s3://bucket/dest") + if err != nil { + t.Fatalf("writeAccessFor: %v", err) + } + sa := m["s3://bucket/dest"] + if sa.GetWriteMethod() != storagebroker.WriteMethodPUT || sa.GetWriteToken() != "wt-s3://bucket/dest" { + t.Errorf("write capability = %+v, want PUT with wt-s3://bucket/dest", sa) + } +} + +func TestAccessFor_PropagatesMintError(t *testing.T) { + w := &ActorWorkflow{broker: &fakeBroker{err: errors.New("boom")}} + + if _, err := w.readAccessFor(context.Background(), "s3://bucket/a"); err == nil { + t.Errorf("readAccessFor: want error, got nil") + } + if _, err := w.writeAccessFor(context.Background(), "s3://bucket/a"); err == nil { + t.Errorf("writeAccessFor: want error, got nil") + } +} diff --git a/cmd/ateapi/internal/controlapi/functionaltest/common_test.go b/cmd/ateapi/internal/controlapi/functionaltest/common_test.go index 8676cd2261..0d31c55b4e 100644 --- a/cmd/ateapi/internal/controlapi/functionaltest/common_test.go +++ b/cmd/ateapi/internal/controlapi/functionaltest/common_test.go @@ -178,7 +178,7 @@ func setupTestWithVolumePlugins(t *testing.T, ns string, plugins map[string]volu mockDriverName: mockPlugin, } } - service := controlapi.NewRPCService(persistence, wc, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, scLister, dialer, instruments, "", volPlugins) + service := controlapi.NewRPCService(persistence, wc, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, scLister, dialer, instruments, "", volPlugins, nil) // 5. Start REAL gRPC Server for ATE API grpcServer := grpc.NewServer(grpc.ChainUnaryInterceptor( diff --git a/cmd/ateapi/internal/controlapi/service.go b/cmd/ateapi/internal/controlapi/service.go index d19845e05e..fe8e072887 100644 --- a/cmd/ateapi/internal/controlapi/service.go +++ b/cmd/ateapi/internal/controlapi/service.go @@ -18,6 +18,7 @@ import ( "context" "sync" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/storagebroker" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" "github.com/agent-substrate/substrate/internal/resources" @@ -70,6 +71,7 @@ func NewRPCService( instruments *Instruments, egressGatewayAddress string, volumePlugins map[string]volume.VolumePluginControlPlane, + broker storagebroker.Broker, ) *RPCService { impl := newServiceImpl(persistence, actorTemplateLister, storageClassLister) s := &RPCService{ @@ -81,7 +83,7 @@ func NewRPCService( instruments: instruments, volumePlugins: volumePlugins, } - s.actorWorkflow = NewActorWorkflow(impl, workerCache, dialer, actorTemplateLister, workerPoolLister, sandboxConfigLister, storageClassLister, instruments, egressGatewayAddress, s) + s.actorWorkflow = NewActorWorkflow(impl, workerCache, dialer, actorTemplateLister, workerPoolLister, sandboxConfigLister, storageClassLister, instruments, egressGatewayAddress, s, broker) s.workerWorkflow = NewWorkerWorkflow(impl) return s } diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index b7daedc4ea..0063426542 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/cmd/ateapi/internal/controlapi/workflow.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/agent-substrate/substrate/cmd/ateapi/internal/scheduling" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/storagebroker" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" "github.com/agent-substrate/substrate/internal/resources" @@ -79,6 +80,9 @@ type ActorWorkflow struct { instruments *Instruments egressGatewayAddress string pluginRegistry VolumePluginRegistry + // broker, when set, mints short-lived per-snapshot signed access so atelet + // reads and writes snapshots over plain HTTP with no cloud credential. + broker storagebroker.Broker } // NewActorWorkflow creates a new ActorWorkflow. instruments may be nil. @@ -93,6 +97,7 @@ func NewActorWorkflow( instruments *Instruments, egressGatewayAddress string, pluginRegistry VolumePluginRegistry, + broker storagebroker.Broker, ) *ActorWorkflow { return &ActorWorkflow{ store: store, @@ -106,6 +111,7 @@ func NewActorWorkflow( instruments: instruments, egressGatewayAddress: egressGatewayAddress, pluginRegistry: pluginRegistry, + broker: broker, } } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 662d228723..343fc740e6 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -708,6 +708,11 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou if !src.GoldenSnapshotURI.IsZero() { req.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN req.GoldenSnapshotUri = src.GoldenSnapshotURI.String() + // atelet combines the local pause snapshot with the golden snapshot + // it reads from object storage; mint a read capability for it. + if req.SignedAccess, err = w.readAccessFor(ctx, src.GoldenSnapshotURI.String()); err != nil { + return tele, maybeCrashActor(ctx, w.store, actorRef, err, "while minting golden snapshot read access", ateattr.OperationResume) + } } tele.WireSnapshotScope = ateattr.SnapshotScopeValue(req.Scope) @@ -751,6 +756,11 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou CpuMilli: cpuMilli, MemoryBytes: memBytes, } + // The node reads the actor's snapshot and, on a golden-data restore, the + // golden snapshot too; mint a read capability scoped to each. + if req.SignedAccess, err = w.readAccessFor(ctx, src.SnapshotURI.String(), src.GoldenSnapshotURI.String()); err != nil { + return tele, maybeCrashActor(ctx, w.store, actorRef, err, "while minting snapshot read access", ateattr.OperationResume) + } _, err = client.Restore(ctx, req) return tele, maybeCrashActor(ctx, w.store, actorRef, err, "while restoring durable snapshot", ateattr.OperationResume) } else { diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index eeccb4659b..afc9a95c72 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -259,6 +259,10 @@ func (w *ActorWorkflow) ensureAteletSuspended(ctx context.Context, actorRef reso Scope: actorSnapshotContentScopeToAtelet(commitSnapshotScope(actor.GetMetadata().GetAtespace(), actorTemplate)), ActorUid: actor.GetMetadata().Uid, } + // The node writes the actor's snapshot; mint a write capability for it. + if req.SignedAccess, err = w.writeAccessFor(ctx, snapshotURI.String()); err != nil { + return "", err + } wireSnapshotScope = ateattr.SnapshotScopeValue(req.Scope) _, err = client.Checkpoint(ctx, req) @@ -311,6 +315,11 @@ func (w *ActorWorkflow) ensurePausedSnapshotUploaded(ctx context.Context, actorR // from the captured scope in the snapshot's manifest where possible. DesiredScope: actorSnapshotContentScopeToAtelet(commitSnapshotScope(actor.GetMetadata().GetAtespace(), actorTemplate)), } + // The node writes the paused snapshot to its destination; mint a write + // capability for it. + if req.SignedAccess, err = w.writeAccessFor(ctx, snapshotURI.String()); err != nil { + return "", err + } wireSnapshotScope = ateattr.SnapshotScopeValue(req.DesiredScope) _, err = client.UploadPausedCheckpoint(ctx, req) diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index 72da34d48e..fae2798f8b 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -703,7 +703,7 @@ func TestSuspendActor_PausedWithoutLocalSnapshotCrashes(t *testing.T) { }); err != nil { t.Fatalf("add template to indexer: %v", err) } - w := NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil) + w := NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil, nil) seedWorkflowActor(t, ctx, st, resources.ActorRef{Atespace: "team-a", Name: "id1"}, "ns", "tmpl1", ateapipb.ActorState_ACTOR_STATE_PAUSED) diff --git a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go index 130b0441d2..52ce0ef1fe 100644 --- a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go @@ -50,7 +50,7 @@ func newTestActorWorkflow(t *testing.T, st store.Interface, tmplNamespace, tmplN }); err != nil { t.Fatalf("add template to indexer: %v", err) } - return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil) + return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil, nil, "", nil, nil) } // seedWorkflowActor stores an actor with the given state, bound to the given diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index e13c9cbc0e..bf04b58a83 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -30,6 +30,7 @@ import ( "github.com/agent-substrate/substrate/cmd/ateapi/internal/controlapi" "github.com/agent-substrate/substrate/cmd/ateapi/internal/debugapi" "github.com/agent-substrate/substrate/cmd/ateapi/internal/oidcjwt" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/storagebroker" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/atepg" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" @@ -187,7 +188,13 @@ func main() { volPlugins := make(map[string]volume.VolumePluginControlPlane) ateletDialer := controlapi.NewAteletDialer(workerPodInformer.GetIndexer(), ateletPodInformer.GetIndexer(), *ateletClientCredBundle, *podIdentityCACerts) - controlSrv := controlapi.NewRPCService(persistence, workerCache, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, storageClassLister, ateletDialer, instruments, *egressGatewayAddress, volPlugins) + // ATE_STORAGE_BACKEND selects a snapshot storage broker (e.g. "s3"); empty + // disables signing and atelet reads/writes snapshots with its own client. + storageBroker, err := storagebroker.New(ctx, os.Getenv("ATE_STORAGE_BACKEND")) + if err != nil { + serverboot.Fatal(ctx, "Failed to create storage broker", err) + } + controlSrv := controlapi.NewRPCService(persistence, workerCache, actorTemplateLister, workerPoolLister, sandboxConfigLister, csiDriverConfigLister, storageClassLister, ateletDialer, instruments, *egressGatewayAddress, volPlugins, storageBroker) // Drive stored ActorTemplates through the golden actor flow. templateReconciler := controlapi.NewActorTemplateReconciler(persistence, controlSrv, sandboxConfigLister) From dcbdf86059718797acc1b04add8e7762cb286d06 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 22:18:31 -0700 Subject: [PATCH 4/8] feat(atelet): read and write snapshots via signed URLs with no cloud SDK --- cmd/atelet/main.go | 34 ++--- cmd/atelet/main_test.go | 1 + cmd/atelet/signedobjectstore.go | 206 +++++++++++++++++++++++++++ cmd/atelet/signedobjectstore_test.go | 51 +++++++ 4 files changed, 276 insertions(+), 16 deletions(-) create mode 100644 cmd/atelet/signedobjectstore.go create mode 100644 cmd/atelet/signedobjectstore_test.go diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index c1994865b3..2071579360 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -763,7 +763,7 @@ func (s *AteomHerder) uploadExternalCheckpoint(ctx context.Context, req *ateletp if err != nil { return err } - return s.uploadSnapshot(ctx, uri, checkpointDir, rec, req.GetActorTemplateNamespace(), req.GetActorTemplateName()) + return s.uploadSnapshot(ctx, req.GetSignedAccess(), uri, checkpointDir, rec, req.GetActorTemplateNamespace(), req.GetActorTemplateName()) } // uploadSnapshot uploads rec's snapshot files from srcDir to uri (each @@ -772,7 +772,8 @@ func (s *AteomHerder) uploadExternalCheckpoint(ctx context.Context, req *ateletp // assume every file it lists is already present. A crash mid-upload thus // leaves only orphaned files, never a manifest pointing at files that never // landed; retries overwrite the deterministic object names. -func (s *AteomHerder) uploadSnapshot(ctx context.Context, uri resources.SnapshotURI, srcDir string, rec *sandboxAssetsRecord, templateNamespace, templateName string) error { +func (s *AteomHerder) uploadSnapshot(ctx context.Context, signedAccess map[string]*ateletpb.SignedObjectAccess, uri resources.SnapshotURI, srcDir string, rec *sandboxAssetsRecord, templateNamespace, templateName string) error { + store := s.storeForURI(signedAccess, uri.String()) g, gCtx := errgroup.WithContext(ctx) for _, fileName := range rec.SnapshotFiles { local := filepath.Join(srcDir, fileName) @@ -782,7 +783,7 @@ func (s *AteomHerder) uploadSnapshot(ctx context.Context, uri resources.Snapshot if err != nil { return fmt.Errorf("while addressing %s in GCS: %w", fileName, err) } - if err := ategcs.SendLocalFileToGCSWithZstd(gCtx, s.gcsClient, objectURI, local); err != nil { + if err := ategcs.SendLocalFileToGCSWithZstd(gCtx, store, objectURI, local); err != nil { return fmt.Errorf("while uploading %s to GCS: %w", fileName, err) } return nil @@ -800,7 +801,7 @@ func (s *AteomHerder) uploadSnapshot(ctx context.Context, uri resources.Snapshot if err != nil { return fmt.Errorf("while addressing snapshot manifest in GCS: %w", err) } - if err := ategcs.SendBytesToGCS(ctx, s.gcsClient, manifestURI, manifest); err != nil { + if err := ategcs.SendBytesToGCS(ctx, store, manifestURI, manifest); err != nil { return fmt.Errorf("while uploading snapshot manifest: %w", err) } return nil @@ -872,7 +873,7 @@ func (s *AteomHerder) uploadLocalCheckpointDir(ctx context.Context, req *ateletp // means the whole snapshot is committed and this retry already // succeeded. Absent on both sides, the paused actor's state is // unrecoverable. - _, fetchErr := ategcs.FetchFromGCS(ctx, s.gcsClient, manifestURI) + _, fetchErr := ategcs.FetchFromGCS(ctx, s.storeForURI(req.GetSignedAccess(), req.GetDestinationSnapshotUri()), manifestURI) if fetchErr == nil { slog.InfoContext(ctx, "Local snapshot already uploaded and pruned; nothing to do", slog.String("snapshot_uri", req.GetDestinationSnapshotUri())) return "", nil @@ -910,7 +911,7 @@ func (s *AteomHerder) uploadLocalCheckpointDir(ctx context.Context, req *ateletp } } - return rec.SandboxClass, s.uploadSnapshot(ctx, uri, localDir, rec, req.GetActorTemplateNamespace(), req.GetActorTemplateName()) + return rec.SandboxClass, s.uploadSnapshot(ctx, req.GetSignedAccess(), uri, localDir, rec, req.GetActorTemplateNamespace(), req.GetActorTemplateName()) } // narrowFullCaptureToData rewrites rec so a FULL capture uploads as a DATA @@ -1009,7 +1010,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) if err != nil { return nil, ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonInvalidObjectURL) } - manifest, err := ategcs.FetchFromGCS(ctx, s.gcsClient, manifestURI) + manifest, err := ategcs.FetchFromGCS(ctx, s.storeForURI(req.GetSignedAccess(), req.GetExternalConfig().GetSnapshotUri()), manifestURI) if err != nil { return nil, ateerrors.CrashIfReason(ctx, fmt.Errorf("while fetching snapshot manifest: %w", err), ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonFailedGetExternalObject) } @@ -1047,7 +1048,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) if err != nil { return nil, ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonInvalidObjectURL) } - manifest, err := ategcs.FetchFromGCS(ctx, s.gcsClient, manifestURI) + manifest, err := ategcs.FetchFromGCS(ctx, s.storeForURI(req.GetSignedAccess(), req.GetGoldenSnapshotUri()), manifestURI) if err != nil { return nil, ateerrors.CrashIfReason(ctx, fmt.Errorf("while fetching golden snapshot manifest: %w", err), ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonFailedGetExternalObject) } @@ -1101,10 +1102,10 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) if goldenRec == nil { return fmt.Errorf("no golden snapshot record for a %s restore", req.GetScope()) } - if err := s.downloadCombinedCheckpoint(gctx, req.GetExternalConfig().GetSnapshotUri(), req.GetGoldenSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles); err != nil { + if err := s.downloadCombinedCheckpoint(gctx, req.GetSignedAccess(), req.GetExternalConfig().GetSnapshotUri(), req.GetGoldenSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles); err != nil { return ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonFailedGetExternalObject, ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonTerminalFileSystemError) } - } else if err := s.downloadExternalCheckpoint(gctx, req.GetExternalConfig().GetSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles); err != nil { + } else if err := s.downloadExternalCheckpoint(gctx, req.GetSignedAccess(), req.GetExternalConfig().GetSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles); err != nil { return ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonFailedGetExternalObject, ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonTerminalFileSystemError) } case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: @@ -1124,7 +1125,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest) }) if combineWithGolden { gLocal.Go(func() error { - if err := s.downloadExternalCheckpoint(gLocalCtx, req.GetGoldenSnapshotUri(), checkpointDir, goldenOnlyFiles(sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles)); err != nil { + if err := s.downloadExternalCheckpoint(gLocalCtx, req.GetSignedAccess(), req.GetGoldenSnapshotUri(), checkpointDir, goldenOnlyFiles(sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles)); err != nil { return ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonFailedGetExternalObject, ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonTerminalFileSystemError) } return nil @@ -1478,22 +1479,23 @@ func goldenOnlyFiles(actorFiles, goldenFiles []string) []string { // as a single folder: every file of the actor's own snapshot (the durable-dir // data) plus the golden snapshot's files the actor's set does not shadow, so // the result looks like a Full snapshot whose durable-dir data is the actor's. -func (s *AteomHerder) downloadCombinedCheckpoint(ctx context.Context, actorURI, goldenURI, dstDir string, actorFiles, goldenFiles []string) error { +func (s *AteomHerder) downloadCombinedCheckpoint(ctx context.Context, signedAccess map[string]*ateletpb.SignedObjectAccess, actorURI, goldenURI, dstDir string, actorFiles, goldenFiles []string) error { g, gctx := errgroup.WithContext(ctx) g.Go(func() error { - return s.downloadExternalCheckpoint(gctx, actorURI, dstDir, actorFiles) + return s.downloadExternalCheckpoint(gctx, signedAccess, actorURI, dstDir, actorFiles) }) g.Go(func() error { - return s.downloadExternalCheckpoint(gctx, goldenURI, dstDir, goldenOnlyFiles(actorFiles, goldenFiles)) + return s.downloadExternalCheckpoint(gctx, signedAccess, goldenURI, dstDir, goldenOnlyFiles(actorFiles, goldenFiles)) }) return g.Wait() } -func (s *AteomHerder) downloadExternalCheckpoint(ctx context.Context, snapshotURI string, dstDir string, files []string) error { +func (s *AteomHerder) downloadExternalCheckpoint(ctx context.Context, signedAccess map[string]*ateletpb.SignedObjectAccess, snapshotURI string, dstDir string, files []string) error { uri, err := resources.ParseSnapshotURI(snapshotURI) if err != nil { return err } + store := s.storeForURI(signedAccess, snapshotURI) g, gCtx := errgroup.WithContext(ctx) for _, fileName := range files { fileName := fileName @@ -1503,7 +1505,7 @@ func (s *AteomHerder) downloadExternalCheckpoint(ctx context.Context, snapshotUR if err != nil { return fmt.Errorf("while addressing %s in GCS: %w", fileName, err) } - if err := ategcs.FetchLocalFileFromGCSWithZstd(gCtx, s.gcsClient, objectURI, local); err != nil { + if err := ategcs.FetchLocalFileFromGCSWithZstd(gCtx, store, objectURI, local); err != nil { return fmt.Errorf("while downloading %s from GCS: %w", fileName, err) } return nil diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index cdbd030156..d52db96173 100644 --- a/cmd/atelet/main_test.go +++ b/cmd/atelet/main_test.go @@ -1172,6 +1172,7 @@ func TestDownloadCombinedCheckpoint(t *testing.T) { dstDir := t.TempDir() err := s.downloadCombinedCheckpoint(context.Background(), + nil, "gs://bucket/root/snapshots/ate-demo/counter-1-snap", "gs://bucket/golden-root/snapshots/ate-golden/golden-1", dstDir, diff --git a/cmd/atelet/signedobjectstore.go b/cmd/atelet/signedobjectstore.go new file mode 100644 index 0000000000..e20a45dd81 --- /dev/null +++ b/cmd/atelet/signedobjectstore.go @@ -0,0 +1,206 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "bytes" + "context" + "fmt" + "io" + "mime/multipart" + "net/http" + "strings" + + "github.com/agent-substrate/substrate/cmd/atelet/internal/ategcs" + "github.com/agent-substrate/substrate/internal/proto/ateletpb" +) + +// signedObjectStore reads and writes snapshot objects using control-plane-minted +// signed access, over plain net/http. It links no cloud SDK and holds no cloud +// credential. It dispatches on the mechanics carried in SignedObjectAccess, not +// on the cloud: +// +// read: a per-object URL map (S3 presigned GET) or a token appended to +// / (Azure container SAS). +// write: a multipart POST to postURL with postFields (S3 presigned POST + +// starts-with) or a PUT to /? (Azure). +// +// It satisfies ategcs.ObjectStorage, so the snapshot manifest and zstd code use +// it unchanged. It deliberately omits the streaming-put fast path so PUT uploads +// take the buffered, seekable, Content-Length route Azure block blob requires. +type signedObjectStore struct { + prefixURL string + readToken string + readObjectURLs map[string]string + writeMethod string + writeToken string + writeHeaders map[string]string + postURL string + postFields map[string]string + httpClient *http.Client +} + +var _ ategcs.ObjectStorage = (*signedObjectStore)(nil) + +func newSignedObjectStore(sa *ateletpb.SignedObjectAccess) *signedObjectStore { + return &signedObjectStore{ + prefixURL: strings.TrimRight(sa.GetPrefixUrl(), "/"), + readToken: sa.GetReadToken(), + readObjectURLs: sa.GetReadObjectUrls(), + writeMethod: sa.GetWriteMethod(), + writeToken: sa.GetWriteToken(), + writeHeaders: sa.GetWriteHeaders(), + postURL: sa.GetPostUrl(), + postFields: sa.GetPostFields(), + httpClient: http.DefaultClient, + } +} + +// GetObject fetches object over HTTP. bucket is ignored: the capability already +// names the account/bucket. A per-object URL wins; otherwise the key and read +// token are appended to the prefix URL. +func (s *signedObjectStore) GetObject(ctx context.Context, bucket, object string) (io.ReadCloser, error) { + url, ok := s.readObjectURLs[object] + if !ok { + if s.readToken == "" { + return nil, fmt.Errorf("no read capability for object %q", object) + } + url = s.prefixURL + "/" + strings.TrimPrefix(object, "/") + "?" + s.readToken + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + resp, err := s.httpClient.Do(req) + if err != nil { + return nil, err + } + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + resp.Body.Close() + return nil, fmt.Errorf("GET %s: status %d: %s", object, resp.StatusCode, strings.TrimSpace(string(body))) + } + return resp.Body, nil +} + +// PutObject writes object over HTTP: POST for S3 (form policy) or PUT for Azure. +func (s *signedObjectStore) PutObject(ctx context.Context, bucket, object string, reader io.Reader) error { + if s.writeMethod == "POST" { + return s.postObject(ctx, object, reader) + } + return s.putObject(ctx, object, reader) +} + +// putObject does a plain PUT with the object key and write token appended, the +// Azure block-blob path. +func (s *signedObjectStore) putObject(ctx context.Context, object string, reader io.Reader) error { + length, err := seekableLen(reader) + if err != nil { + return fmt.Errorf("signed PUT %s: %w", object, err) + } + url := s.prefixURL + "/" + strings.TrimPrefix(object, "/") + "?" + s.writeToken + req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, reader) + if err != nil { + return err + } + req.ContentLength = length + for k, v := range s.writeHeaders { + req.Header.Set(k, v) + } + resp, err := s.httpClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + return fmt.Errorf("PUT %s: status %d: %s", object, resp.StatusCode, strings.TrimSpace(string(body))) + } + return nil +} + +// postObject submits a multipart form to the S3 POST endpoint, reusing the one +// signed starts-with policy and setting the key to this object. The signature +// covers the policy, not the body, so the same fields upload every file under +// the prefix. +func (s *signedObjectStore) postObject(ctx context.Context, object string, reader io.Reader) error { + var body bytes.Buffer + form := multipart.NewWriter(&body) + for k, v := range s.postFields { + if k == "key" { + continue + } + _ = form.WriteField(k, v) + } + _ = form.WriteField("key", object) + fw, err := form.CreateFormFile("file", object) + if err != nil { + return err + } + if _, err := io.Copy(fw, reader); err != nil { + return fmt.Errorf("buffering %s for POST: %w", object, err) + } + if err := form.Close(); err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.postURL, &body) + if err != nil { + return err + } + req.Header.Set("Content-Type", form.FormDataContentType()) + resp, err := s.httpClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode >= 300 { + msg, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + return fmt.Errorf("POST %s: status %d: %s", object, resp.StatusCode, strings.TrimSpace(string(msg))) + } + return nil +} + +// seekableLen returns the remaining length of a seekable reader without +// consuming it. +func seekableLen(reader io.Reader) (int64, error) { + seeker, ok := reader.(io.Seeker) + if !ok { + return 0, fmt.Errorf("a seekable body is required to set Content-Length") + } + cur, err := seeker.Seek(0, io.SeekCurrent) + if err != nil { + return 0, err + } + end, err := seeker.Seek(0, io.SeekEnd) + if err != nil { + return 0, err + } + if _, err := seeker.Seek(cur, io.SeekStart); err != nil { + return 0, err + } + return end - cur, nil +} + +// storeForURI returns the signed HTTP store for the snapshot at snapshotURI when +// the control plane minted a capability for it, otherwise atelet's built-in +// storage client. A restore that reads several snapshots (actor plus golden) +// carries one capability per URI, so the store is resolved per object source. +// This is the single switch that keeps the node cloud-agnostic. +func (s *AteomHerder) storeForURI(signedAccess map[string]*ateletpb.SignedObjectAccess, snapshotURI string) ategcs.ObjectStorage { + if sa := signedAccess[snapshotURI]; sa != nil && sa.GetPrefixUrl() != "" { + return newSignedObjectStore(sa) + } + return s.gcsClient +} diff --git a/cmd/atelet/signedobjectstore_test.go b/cmd/atelet/signedobjectstore_test.go new file mode 100644 index 0000000000..c0000e2d50 --- /dev/null +++ b/cmd/atelet/signedobjectstore_test.go @@ -0,0 +1,51 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "testing" + + "github.com/agent-substrate/substrate/internal/proto/ateletpb" +) + +// TestStoreForURI verifies the per-URI switch: a restore that reads several +// snapshots resolves each object source independently, so the actor's snapshot +// can use a signed store while an unsigned golden falls back to the node's +// built-in client. +func TestStoreForURI(t *testing.T) { + s := &AteomHerder{gcsClient: fakeObjectStorage{}} + signedAccess := map[string]*ateletpb.SignedObjectAccess{ + "s3://bucket/actor": {PrefixUrl: "https://account/container", ReadToken: "tok"}, + "s3://bucket/empty": {PrefixUrl: ""}, // no prefix -> not a usable capability + } + + isSigned := func(uri string, saMap map[string]*ateletpb.SignedObjectAccess) bool { + _, ok := s.storeForURI(saMap, uri).(*signedObjectStore) + return ok + } + + if !isSigned("s3://bucket/actor", signedAccess) { + t.Errorf("URI with a capability: want signed store") + } + if isSigned("s3://bucket/other", signedAccess) { + t.Errorf("URI absent from the map: want built-in fallback") + } + if isSigned("s3://bucket/empty", signedAccess) { + t.Errorf("empty-prefix capability: want built-in fallback") + } + if isSigned("s3://bucket/actor", nil) { + t.Errorf("nil capability map: want built-in fallback") + } +} From 793b402d5a981eee848e99241f1d85cf28aac931 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 22:18:32 -0700 Subject: [PATCH 5/8] feat(ateapi): add out-of-process UDS storage broker client --- cmd/ateapi/internal/storagebroker/uds.go | 109 +++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 cmd/ateapi/internal/storagebroker/uds.go diff --git a/cmd/ateapi/internal/storagebroker/uds.go b/cmd/ateapi/internal/storagebroker/uds.go new file mode 100644 index 0000000000..23319eb7c7 --- /dev/null +++ b/cmd/ateapi/internal/storagebroker/uds.go @@ -0,0 +1,109 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package storagebroker + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "os" + "time" +) + +// The "uds" backend dials an out-of-process broker over a Unix domain socket. +// The broker holds the cloud identity and links the cloud SDK; ate-api-server +// links none. This is the end-state that keeps the control-plane core free of +// every cloud SDK (the sidecar is swappable per cloud), mirroring the mTLS-UDS +// pattern atelet already uses for its actor credential broker. +func init() { + Register("uds", newUDSBroker) +} + +// mintRequest and mintReply are the tiny wire contract between ate-api-server +// and the out-of-process broker sidecar. +type mintRequest struct { + Verb string `json:"verb"` + SnapshotURI string `json:"snapshotUri"` + TTLSeconds int `json:"ttlSeconds"` +} + +type mintReply struct { + PrefixURL string `json:"prefixUrl"` + ReadToken string `json:"readToken,omitempty"` + ReadObjectURLs map[string]string `json:"readObjectUrls,omitempty"` + WriteMethod string `json:"writeMethod,omitempty"` + WriteToken string `json:"writeToken,omitempty"` + WriteHeaders map[string]string `json:"writeHeaders,omitempty"` + PostURL string `json:"postUrl,omitempty"` + PostFields map[string]string `json:"postFields,omitempty"` +} + +type udsBroker struct { + client *http.Client +} + +func newUDSBroker(ctx context.Context) (Broker, error) { + socket := os.Getenv("BROKER_UDS_PATH") + if socket == "" { + return nil, fmt.Errorf("uds storage broker needs BROKER_UDS_PATH") + } + transport := &http.Transport{ + DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, "unix", socket) + }, + } + return &udsBroker{client: &http.Client{Transport: transport}}, nil +} + +func (b *udsBroker) MintRead(ctx context.Context, snapshotURI string, ttl time.Duration) (Capability, error) { + return b.mint(ctx, "read", snapshotURI, ttl) +} + +func (b *udsBroker) MintWrite(ctx context.Context, snapshotURI string, ttl time.Duration) (Capability, error) { + return b.mint(ctx, "write", snapshotURI, ttl) +} + +func (b *udsBroker) mint(ctx context.Context, verb, snapshotURI string, ttl time.Duration) (Capability, error) { + body, err := json.Marshal(mintRequest{Verb: verb, SnapshotURI: snapshotURI, TTLSeconds: int(ttl.Seconds())}) + if err != nil { + return Capability{}, err + } + // Host is ignored (the transport dials the socket); the path selects the op. + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://broker/mint", bytes.NewReader(body)) + if err != nil { + return Capability{}, err + } + req.Header.Set("Content-Type", "application/json") + resp, err := b.client.Do(req) + if err != nil { + return Capability{}, fmt.Errorf("dialing broker over UDS: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + msg, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + return Capability{}, fmt.Errorf("broker mint %s: status %d: %s", verb, resp.StatusCode, string(msg)) + } + var reply mintReply + if err := json.NewDecoder(resp.Body).Decode(&reply); err != nil { + return Capability{}, fmt.Errorf("decoding broker reply: %w", err) + } + // mintReply mirrors Capability field for field, so a struct conversion + // copies it (struct tags are ignored in conversions). + return Capability(reply), nil +} From 498daa9107c392ae0256479b85484fae8c930258 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Wed, 26 Aug 2026 23:57:32 -0700 Subject: [PATCH 6/8] feat(s3-broker-server): add out-of-process S3 storage broker sidecar --- .../internal/storagebroker/mint_bench_test.go | 189 ++++++++++++++ cmd/s3-broker-server/main.go | 233 ++++++++++++++++++ 2 files changed, 422 insertions(+) create mode 100644 cmd/ateapi/internal/storagebroker/mint_bench_test.go create mode 100644 cmd/s3-broker-server/main.go diff --git a/cmd/ateapi/internal/storagebroker/mint_bench_test.go b/cmd/ateapi/internal/storagebroker/mint_bench_test.go new file mode 100644 index 0000000000..63b84a73bf --- /dev/null +++ b/cmd/ateapi/internal/storagebroker/mint_bench_test.go @@ -0,0 +1,189 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package storagebroker + +import ( + "bytes" + "context" + "fmt" + "net/url" + "os" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" +) + +// These benchmarks quantify the cost of running the S3 broker out-of-process +// over a UDS sidecar versus doing the same signing in-process, against the same +// S3 endpoint. The in-process path is replicated here (this is a test file, so +// the AWS SDK it links never enters the ate-api-server binary) so the comparison +// survives the removal of the in-tree broker. They are env-gated; `go test` +// skips them unless an S3 endpoint and a running s3-broker-server socket are set: +// +// S3_ENDPOINT=http://localhost:9000 S3_BUCKET=ate-snapshots \ +// AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... \ +// BROKER_UDS_PATH=/tmp/s3broker.sock \ +// go test -run '^$' -bench Mint -benchtime 300x ./cmd/ateapi/internal/storagebroker/ + +const ( + benchSnapshotURI = "s3://ate-snapshots/bench/snap" + benchObjectCount = 5 +) + +func benchEnvReady(tb testing.TB) { + if os.Getenv("S3_BUCKET") == "" || os.Getenv("BROKER_UDS_PATH") == "" { + tb.Skip("mint benchmarks need S3_BUCKET, S3_ENDPOINT, AWS creds, and BROKER_UDS_PATH (s3-broker-server running)") + } +} + +// inProcS3 replicates the in-process presign the (removed) in-tree broker did, +// so the benchmark can measure it against the same logic behind the UDS sidecar. +type inProcS3 struct { + client *s3.Client + presign *s3.PresignClient + bucket string +} + +func newInProcS3(tb testing.TB) *inProcS3 { + tb.Helper() + endpoint := firstBenchEnv("S3_ENDPOINT", "AWS_ENDPOINT_URL") + region := firstBenchEnv("AWS_REGION", "AWS_DEFAULT_REGION") + if region == "" { + region = "us-east-1" + } + cfg := aws.Config{ + Region: region, + Credentials: credentials.NewStaticCredentialsProvider(os.Getenv("AWS_ACCESS_KEY_ID"), os.Getenv("AWS_SECRET_ACCESS_KEY"), ""), + } + client := s3.NewFromConfig(cfg, func(o *s3.Options) { + if endpoint != "" { + o.BaseEndpoint = aws.String(endpoint) + o.UsePathStyle = true + } + }) + return &inProcS3{client: client, presign: s3.NewPresignClient(client), bucket: os.Getenv("S3_BUCKET")} +} + +func benchKeyPrefix(snapshotURI string) string { + u, err := url.Parse(snapshotURI) + if err != nil { + return strings.Trim(snapshotURI, "/") + } + return strings.Trim(u.Path, "/") +} + +func (b *inProcS3) mintWrite(ctx context.Context, snapshotURI string, ttl time.Duration) error { + prefix := benchKeyPrefix(snapshotURI) + "/" + _, err := b.presign.PresignPostObject(ctx, &s3.PutObjectInput{Bucket: aws.String(b.bucket), Key: aws.String(prefix)}, func(o *s3.PresignPostOptions) { + o.Conditions = []any{[]any{"starts-with", "$key", prefix}} + o.Expires = ttl + }) + return err +} + +func (b *inProcS3) mintRead(ctx context.Context, snapshotURI string, ttl time.Duration) error { + prefix := benchKeyPrefix(snapshotURI) + "/" + out, err := b.client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{Bucket: aws.String(b.bucket), Prefix: aws.String(prefix)}) + if err != nil { + return err + } + for _, obj := range out.Contents { + if _, err := b.presign.PresignGetObject(ctx, &s3.GetObjectInput{Bucket: aws.String(b.bucket), Key: obj.Key}, s3.WithPresignExpires(ttl)); err != nil { + return err + } + } + return nil +} + +func firstBenchEnv(keys ...string) string { + for _, k := range keys { + if v := os.Getenv(k); v != "" { + return v + } + } + return "" +} + +// benchSeed writes benchObjectCount tiny objects under the bench prefix so the +// read mint has a realistic object list to enumerate and presign. +func benchSeed(tb testing.TB) { + tb.Helper() + b := newInProcS3(tb) + for i := 0; i < benchObjectCount; i++ { + key := fmt.Sprintf("bench/snap/file-%d.img.zstd", i) + if _, err := b.client.PutObject(context.Background(), &s3.PutObjectInput{Bucket: aws.String(b.bucket), Key: aws.String(key), Body: bytes.NewReader([]byte("x"))}); err != nil { + tb.Fatalf("seeding %s: %v", key, err) + } + } +} + +func benchUDS(tb testing.TB) Broker { + tb.Helper() + u, err := newUDSBroker(context.Background()) + if err != nil { + tb.Fatalf("uds broker: %v", err) + } + return u +} + +func BenchmarkMintWriteInProc(b *testing.B) { + benchEnvReady(b) + it := newInProcS3(b) + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := it.mintWrite(context.Background(), benchSnapshotURI, 15*time.Minute); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkMintWriteUDS(b *testing.B) { + benchEnvReady(b) + u := benchUDS(b) + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := u.MintWrite(context.Background(), benchSnapshotURI, 15*time.Minute); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkMintReadInProc(b *testing.B) { + benchEnvReady(b) + benchSeed(b) + it := newInProcS3(b) + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := it.mintRead(context.Background(), benchSnapshotURI, 15*time.Minute); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkMintReadUDS(b *testing.B) { + benchEnvReady(b) + benchSeed(b) + u := benchUDS(b) + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := u.MintRead(context.Background(), benchSnapshotURI, 15*time.Minute); err != nil { + b.Fatal(err) + } + } +} diff --git a/cmd/s3-broker-server/main.go b/cmd/s3-broker-server/main.go new file mode 100644 index 0000000000..b0f0fcc600 --- /dev/null +++ b/cmd/s3-broker-server/main.go @@ -0,0 +1,233 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Command s3-broker-server is the out-of-process S3 storage broker sidecar. It +// links the AWS SDK and holds the S3 credential, and serves signed-URL mints to +// ate-api-server over a Unix domain socket, so ate-api-server's own binary links +// no cloud SDK: the control-plane core dials this socket and mints nothing +// itself. S3's SDK is already vendored for atelet's object store, so this sidecar +// ships in-repo (unlike a cloud whose SDK is not vendored, which lives in a +// separate module); the point is only to keep it out of the core process. +// +// It mints the two S3 mechanisms: +// - write: one presigned POST policy with a starts-with condition on the key, +// so a single capability covers every file atelet writes under the snapshot +// prefix without the control plane knowing the file names in advance, and +// - read: a presigned GET per object, enumerated by listing the prefix (the +// broker holds the credential; the node never does). +package main + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "net" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/s3" +) + +// mintRequest and mintReply mirror the wire contract in +// cmd/ateapi/internal/storagebroker/uds.go. +type mintRequest struct { + Verb string `json:"verb"` + SnapshotURI string `json:"snapshotUri"` + TTLSeconds int `json:"ttlSeconds"` +} + +type mintReply struct { + PrefixURL string `json:"prefixUrl"` + ReadToken string `json:"readToken,omitempty"` + ReadObjectURLs map[string]string `json:"readObjectUrls,omitempty"` + WriteMethod string `json:"writeMethod,omitempty"` + WriteToken string `json:"writeToken,omitempty"` + WriteHeaders map[string]string `json:"writeHeaders,omitempty"` + PostURL string `json:"postUrl,omitempty"` + PostFields map[string]string `json:"postFields,omitempty"` +} + +type s3Broker struct { + client *s3.Client + presign *s3.PresignClient + bucket string + prefixURL string +} + +func main() { + ctx := context.Background() + socket := getenv("BROKER_UDS_PATH", "/run/broker/broker.sock") + b, err := newS3Broker() + if err != nil { + fatal(fmt.Sprintf("building s3 broker: %v", err)) + } + + mux := http.NewServeMux() + mux.HandleFunc("/mint", func(w http.ResponseWriter, r *http.Request) { + var req mintRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + ttl := time.Duration(req.TTLSeconds) * time.Second + if ttl <= 0 { + ttl = 15 * time.Minute + } + var ( + reply mintReply + mErr error + ) + if req.Verb == "write" { + reply, mErr = b.mintWrite(r.Context(), req.SnapshotURI, ttl) + } else { + reply, mErr = b.mintRead(r.Context(), req.SnapshotURI, ttl) + } + if mErr != nil { + http.Error(w, mErr.Error(), http.StatusInternalServerError) + return + } + slog.InfoContext(r.Context(), "Minted snapshot capability", slog.String("verb", req.Verb), slog.String("snapshotUri", req.SnapshotURI)) + _ = json.NewEncoder(w).Encode(reply) + }) + + _ = os.Remove(socket) + if err := os.MkdirAll(dir(socket), 0o755); err != nil { + fatal(fmt.Sprintf("mkdir socket dir: %v", err)) + } + ln, err := net.Listen("unix", socket) + if err != nil { + fatal(fmt.Sprintf("listen on %s: %v", socket, err)) + } + slog.InfoContext(ctx, "s3-broker-server listening", slog.String("socket", socket), slog.String("bucket", b.bucket)) + if err := http.Serve(ln, mux); err != nil { + fatal(fmt.Sprintf("serve: %v", err)) + } +} + +func newS3Broker() (*s3Broker, error) { + endpoint := firstEnv("S3_ENDPOINT", "AWS_ENDPOINT_URL") + bucket := os.Getenv("S3_BUCKET") + if bucket == "" { + return nil, fmt.Errorf("S3_BUCKET is required") + } + region := firstEnv("AWS_REGION", "AWS_DEFAULT_REGION") + if region == "" { + region = "us-east-1" + } + cfg := aws.Config{ + Region: region, + Credentials: credentials.NewStaticCredentialsProvider(os.Getenv("AWS_ACCESS_KEY_ID"), os.Getenv("AWS_SECRET_ACCESS_KEY"), ""), + } + client := s3.NewFromConfig(cfg, func(o *s3.Options) { + if endpoint != "" { + o.BaseEndpoint = aws.String(endpoint) + o.UsePathStyle = true + } + }) + return &s3Broker{ + client: client, + presign: s3.NewPresignClient(client), + bucket: bucket, + prefixURL: strings.TrimRight(endpoint, "/") + "/" + bucket, + }, nil +} + +// snapshotKeyPrefix turns a snapshot URI into the object-key prefix atelet's +// file keys share, mirroring how atelet derives the object path. +func snapshotKeyPrefix(snapshotURI string) string { + u, err := url.Parse(snapshotURI) + if err != nil { + return strings.Trim(snapshotURI, "/") + } + return strings.Trim(u.Path, "/") +} + +func (b *s3Broker) mintWrite(ctx context.Context, snapshotURI string, ttl time.Duration) (mintReply, error) { + prefix := snapshotKeyPrefix(snapshotURI) + "/" + post, err := b.presign.PresignPostObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(b.bucket), + Key: aws.String(prefix), + }, func(o *s3.PresignPostOptions) { + o.Conditions = []any{[]any{"starts-with", "$key", prefix}} + o.Expires = ttl + }) + if err != nil { + return mintReply{}, fmt.Errorf("presigning S3 POST policy: %w", err) + } + return mintReply{ + PrefixURL: b.prefixURL, + WriteMethod: "POST", + PostURL: post.URL, + PostFields: post.Values, + }, nil +} + +func (b *s3Broker) mintRead(ctx context.Context, snapshotURI string, ttl time.Duration) (mintReply, error) { + prefix := snapshotKeyPrefix(snapshotURI) + "/" + out, err := b.client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{ + Bucket: aws.String(b.bucket), + Prefix: aws.String(prefix), + }) + if err != nil { + return mintReply{}, fmt.Errorf("listing snapshot prefix %q: %w", prefix, err) + } + urls := make(map[string]string, len(out.Contents)) + for _, obj := range out.Contents { + key := aws.ToString(obj.Key) + pres, err := b.presign.PresignGetObject(ctx, &s3.GetObjectInput{ + Bucket: aws.String(b.bucket), + Key: aws.String(key), + }, s3.WithPresignExpires(ttl)) + if err != nil { + return mintReply{}, fmt.Errorf("presigning GET for %q: %w", key, err) + } + urls[key] = pres.URL + } + return mintReply{PrefixURL: b.prefixURL, ReadObjectURLs: urls}, nil +} + +func firstEnv(keys ...string) string { + for _, k := range keys { + if v := os.Getenv(k); v != "" { + return v + } + } + return "" +} + +func getenv(k, def string) string { + if v := os.Getenv(k); v != "" { + return v + } + return def +} + +func dir(path string) string { + i := strings.LastIndex(path, "/") + if i <= 0 { + return "." + } + return path[:i] +} + +func fatal(msg string) { + slog.Error(msg) + os.Exit(1) +} From 380cf9adc18d0ef20b6e732eb86b200adf8fa478 Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Thu, 27 Aug 2026 22:00:39 -0700 Subject: [PATCH 7/8] refactor(ateapi): select the storage broker via a dedicated env var --- cmd/ateapi/main.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index bf04b58a83..689989b2d8 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -188,9 +188,11 @@ func main() { volPlugins := make(map[string]volume.VolumePluginControlPlane) ateletDialer := controlapi.NewAteletDialer(workerPodInformer.GetIndexer(), ateletPodInformer.GetIndexer(), *ateletClientCredBundle, *podIdentityCACerts) - // ATE_STORAGE_BACKEND selects a snapshot storage broker (e.g. "s3"); empty - // disables signing and atelet reads/writes snapshots with its own client. - storageBroker, err := storagebroker.New(ctx, os.Getenv("ATE_STORAGE_BACKEND")) + // ATE_SNAPSHOT_BROKER selects the out-of-process snapshot storage broker + // (e.g. "uds"); empty disables signing and atelet reads and writes snapshots + // with its own client. Distinct from atelet's ATE_STORAGE_BACKEND, which + // selects the node's own object-store client (gcs or s3). + storageBroker, err := storagebroker.New(ctx, os.Getenv("ATE_SNAPSHOT_BROKER")) if err != nil { serverboot.Fatal(ctx, "Failed to create storage broker", err) } From c915cc4b640846043b3b190100b8fb1053b4e72a Mon Sep 17 00:00:00 2001 From: Anish Ramasekar Date: Thu, 27 Aug 2026 22:00:41 -0700 Subject: [PATCH 8/8] test(atelet): cover the signed object store over HTTP --- cmd/atelet/signedobjectstore_test.go | 209 +++++++++++++++++++++++++++ 1 file changed, 209 insertions(+) diff --git a/cmd/atelet/signedobjectstore_test.go b/cmd/atelet/signedobjectstore_test.go index c0000e2d50..24e4bd2ce4 100644 --- a/cmd/atelet/signedobjectstore_test.go +++ b/cmd/atelet/signedobjectstore_test.go @@ -15,6 +15,12 @@ package main import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" "testing" "github.com/agent-substrate/substrate/internal/proto/ateletpb" @@ -49,3 +55,206 @@ func TestStoreForURI(t *testing.T) { t.Errorf("nil capability map: want built-in fallback") } } + +// TestSignedObjectStoreGetObjectPerObjectURL covers the S3 read path: a +// per-object presigned GET URL is fetched verbatim. +func TestSignedObjectStoreGetObjectPerObjectURL(t *testing.T) { + const want = "snapshot-bytes" + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + t.Errorf("method = %s, want GET", r.Method) + } + if r.URL.Path != "/presigned/manifest" { + t.Errorf("path = %s, want /presigned/manifest", r.URL.Path) + } + _, _ = io.WriteString(w, want) + })) + defer ts.Close() + + s := &signedObjectStore{ + readObjectURLs: map[string]string{"manifest.json": ts.URL + "/presigned/manifest"}, + httpClient: ts.Client(), + } + rc, err := s.GetObject(context.Background(), "bucket", "manifest.json") + if err != nil { + t.Fatalf("GetObject: %v", err) + } + defer rc.Close() + got, _ := io.ReadAll(rc) + if string(got) != want { + t.Errorf("body = %q, want %q", got, want) + } +} + +// TestSignedObjectStoreGetObjectPrefixToken covers the Azure read path: the key +// and read token are appended to the prefix URL. +func TestSignedObjectStoreGetObjectPrefixToken(t *testing.T) { + const want = "azure-bytes" + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/container/dir/file.bin" { + t.Errorf("path = %s, want /container/dir/file.bin", r.URL.Path) + } + if r.URL.RawQuery != "sig=abc123" { + t.Errorf("query = %s, want sig=abc123", r.URL.RawQuery) + } + _, _ = io.WriteString(w, want) + })) + defer ts.Close() + + s := &signedObjectStore{ + prefixURL: ts.URL + "/container", + readToken: "sig=abc123", + httpClient: ts.Client(), + } + rc, err := s.GetObject(context.Background(), "bucket", "dir/file.bin") + if err != nil { + t.Fatalf("GetObject: %v", err) + } + defer rc.Close() + got, _ := io.ReadAll(rc) + if string(got) != want { + t.Errorf("body = %q, want %q", got, want) + } +} + +// TestSignedObjectStoreGetObjectNoCapability verifies a read with neither a +// per-object URL nor a read token fails instead of fetching something wrong. +func TestSignedObjectStoreGetObjectNoCapability(t *testing.T) { + s := &signedObjectStore{httpClient: http.DefaultClient} + if _, err := s.GetObject(context.Background(), "bucket", "missing"); err == nil { + t.Fatal("GetObject with no read capability: want error, got nil") + } +} + +// TestSignedObjectStoreGetObjectHTTPError verifies a non-200 response surfaces +// as an error carrying the status code. +func TestSignedObjectStoreGetObjectHTTPError(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "denied", http.StatusForbidden) + })) + defer ts.Close() + + s := &signedObjectStore{ + readObjectURLs: map[string]string{"o": ts.URL + "/o"}, + httpClient: ts.Client(), + } + _, err := s.GetObject(context.Background(), "bucket", "o") + if err == nil { + t.Fatal("GetObject on 403: want error, got nil") + } + if !strings.Contains(err.Error(), "403") { + t.Errorf("error = %v, want it to mention status 403", err) + } +} + +// TestSignedObjectStorePutObjectPOST covers the S3 write path: a multipart form +// POST whose key field is set to the object, overriding the policy placeholder, +// while the other signed fields are forwarded. +func TestSignedObjectStorePutObjectPOST(t *testing.T) { + const payload = "counter-state" + var gotKey, gotPolicy, gotBody string + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method = %s, want POST", r.Method) + } + if err := r.ParseMultipartForm(1 << 20); err != nil { + t.Errorf("ParseMultipartForm: %v", err) + w.WriteHeader(http.StatusBadRequest) + return + } + gotKey = r.FormValue("key") + gotPolicy = r.FormValue("policy") + f, _, err := r.FormFile("file") + if err != nil { + t.Errorf("FormFile: %v", err) + w.WriteHeader(http.StatusBadRequest) + return + } + defer f.Close() + b, _ := io.ReadAll(f) + gotBody = string(b) + w.WriteHeader(http.StatusNoContent) + })) + defer ts.Close() + + s := &signedObjectStore{ + writeMethod: "POST", + postURL: ts.URL + "/upload", + postFields: map[string]string{"key": "ignored-placeholder", "policy": "signed-policy"}, + httpClient: ts.Client(), + } + if err := s.PutObject(context.Background(), "bucket", "snap/dir/counter.bin", strings.NewReader(payload)); err != nil { + t.Fatalf("PutObject: %v", err) + } + if gotKey != "snap/dir/counter.bin" { + t.Errorf("form key = %q, want the object key (placeholder overridden)", gotKey) + } + if gotPolicy != "signed-policy" { + t.Errorf("policy field = %q, want it forwarded from postFields", gotPolicy) + } + if gotBody != payload { + t.Errorf("uploaded body = %q, want %q", gotBody, payload) + } +} + +// TestSignedObjectStorePutObjectPUT covers the Azure write path: a PUT to the +// prefixed key with the write token, an explicit Content-Length, and the signed +// write headers. +func TestSignedObjectStorePutObjectPUT(t *testing.T) { + const payload = "block-blob-bytes" + var gotPath, gotQuery, gotHeader, gotBody string + var gotLen int64 + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotQuery = r.URL.RawQuery + gotHeader = r.Header.Get("x-ms-blob-type") + gotLen = r.ContentLength + b, _ := io.ReadAll(r.Body) + gotBody = string(b) + w.WriteHeader(http.StatusCreated) + })) + defer ts.Close() + + s := &signedObjectStore{ + prefixURL: ts.URL + "/container", + writeMethod: "PUT", + writeToken: "sv=2021&sig=xyz", + writeHeaders: map[string]string{"x-ms-blob-type": "BlockBlob"}, + httpClient: ts.Client(), + } + if err := s.PutObject(context.Background(), "bucket", "dir/file.bin", bytes.NewReader([]byte(payload))); err != nil { + t.Fatalf("PutObject: %v", err) + } + if gotPath != "/container/dir/file.bin" { + t.Errorf("path = %q, want /container/dir/file.bin", gotPath) + } + if gotQuery != "sv=2021&sig=xyz" { + t.Errorf("query = %q, want the write token", gotQuery) + } + if gotHeader != "BlockBlob" { + t.Errorf("x-ms-blob-type = %q, want BlockBlob", gotHeader) + } + if gotLen != int64(len(payload)) { + t.Errorf("Content-Length = %d, want %d", gotLen, len(payload)) + } + if gotBody != payload { + t.Errorf("body = %q, want %q", gotBody, payload) + } +} + +// TestSignedObjectStorePutObjectPUTRequiresSeekable documents that the PUT path +// needs a seekable body to set Content-Length, so a streaming (non-seekable) +// reader is rejected before any request is made. +func TestSignedObjectStorePutObjectPUTRequiresSeekable(t *testing.T) { + s := &signedObjectStore{ + prefixURL: "http://unused/container", + writeMethod: "PUT", + writeToken: "t", + httpClient: http.DefaultClient, + } + pr, _ := io.Pipe() + defer pr.Close() + if err := s.PutObject(context.Background(), "bucket", "obj", pr); err == nil { + t.Fatal("PutObject with a non-seekable body: want error, got nil") + } +}