diff --git a/protocol/types.go b/protocol/types.go index 57bacb9..2b24500 100644 --- a/protocol/types.go +++ b/protocol/types.go @@ -279,6 +279,27 @@ type ProviderCapabilityReport struct { Reason string `json:"reason,omitempty"` } +type HardwareAttestation struct { + Class string `json:"class"` + Provider string `json:"provider"` + VerifierID string `json:"verifier_id"` + KeyID string `json:"key_id"` + Verified bool `json:"verified"` + ExpiresAt time.Time `json:"expires_at,omitzero"` +} + +type HardwareSecurityCapabilities struct { + TEE []string `json:"tee,omitempty"` + HardwareClasses []string `json:"hardware_classes,omitempty"` + HardwareAttestations []HardwareAttestation `json:"hardware_attestations,omitempty"` +} + +type HardwarePlacementCapabilities struct { + GPUCount int `json:"gpu_count,omitempty"` + Security HardwareSecurityCapabilities `json:"security,omitzero"` + Now time.Time `json:"now,omitzero"` +} + type PlacementRequirementCapabilities struct { CapabilityTags []string `json:"capability_tags,omitempty"` ExecutorProviders []string `json:"executor_providers,omitempty"` @@ -332,6 +353,55 @@ func ValidatePlacementRequirementsAgainstCapabilities(req PlacementRequirements, return errors.Join(errs...) } +func RequiredHardwareClass(req PlacementRequirements) string { + if req.HardwareClass != "" { + return req.HardwareClass + } + switch req.ExecutionSecurityTier { + case ExecutionConfidentialCPU: + return string(ExecutionConfidentialCPU) + case ExecutionConfidentialGPU: + return string(ExecutionConfidentialGPU) + default: + return "" + } +} + +func PlacementRequiresVerifiedHardwareAttestation(req PlacementRequirements) bool { + return req.ExecutionSecurityTier == ExecutionConfidentialCPU || req.ExecutionSecurityTier == ExecutionConfidentialGPU +} + +func HardwareCapabilitiesSatisfyPlacementRequirements(req PlacementRequirements, caps HardwarePlacementCapabilities) bool { + class := RequiredHardwareClass(req) + if class == "" { + return true + } + if class == string(ExecutionConfidentialGPU) && caps.GPUCount == 0 { + return false + } + requireVerified := PlacementRequiresVerifiedHardwareAttestation(req) || + class == string(ExecutionConfidentialCPU) || + class == string(ExecutionConfidentialGPU) + if requireVerified { + now := caps.Now + if now.IsZero() { + now = time.Now().UTC() + } + for _, attestation := range caps.Security.HardwareAttestations { + if attestation.Class == class && + attestation.Provider != "" && + attestation.VerifierID != "" && + attestation.KeyID != "" && + attestation.Verified && + (attestation.ExpiresAt.IsZero() || attestation.ExpiresAt.After(now)) { + return true + } + } + return false + } + return contains(caps.Security.HardwareClasses, class) || contains(caps.Security.TEE, class) +} + func ExecutorCapabilitiesHaveSupportedMatch(req PlacementRequirements, caps PlacementRequirementCapabilities) bool { for _, executor := range caps.Executors { if ExecutorMatchesPlacementRequirements(executor, req) { diff --git a/protocol/types_test.go b/protocol/types_test.go index 72746dc..ce56a49 100644 --- a/protocol/types_test.go +++ b/protocol/types_test.go @@ -486,6 +486,132 @@ func TestValidatePlacementRequirementsAgainstCapabilities(t *testing.T) { } } +func TestRequiredHardwareClass(t *testing.T) { + for name, tc := range map[string]struct { + req protocol.PlacementRequirements + want string + }{ + "explicit": { + req: protocol.PlacementRequirements{HardwareClass: "sev-snp"}, + want: "sev-snp", + }, + "confidential cpu": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialCPU}, + want: "confidential-cpu", + }, + "confidential gpu": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialGPU}, + want: "confidential-gpu", + }, + "none": { + req: protocol.PlacementRequirements{}, + }, + } { + t.Run(name, func(t *testing.T) { + if got := protocol.RequiredHardwareClass(tc.req); got != tc.want { + t.Fatalf("RequiredHardwareClass() = %q, want %q", got, tc.want) + } + }) + } +} + +func TestHardwareCapabilitiesSatisfyPlacementRequirements(t *testing.T) { + now := time.Date(2026, 5, 26, 5, 0, 0, 0, time.UTC) + verifiedCPU := protocol.HardwareAttestation{ + Class: "confidential-cpu", + Provider: "sev-snp", + VerifierID: "verifier-1", + KeyID: "key-1", + Verified: true, + ExpiresAt: now.Add(time.Hour), + } + verifiedGPU := protocol.HardwareAttestation{ + Class: "confidential-gpu", + Provider: "nvidia-cc", + VerifierID: "verifier-1", + KeyID: "key-1", + Verified: true, + ExpiresAt: now.Add(time.Hour), + } + + for name, tc := range map[string]struct { + req protocol.PlacementRequirements + caps protocol.HardwarePlacementCapabilities + want bool + }{ + "confidential cpu requires verified attestation": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialCPU}, + caps: protocol.HardwarePlacementCapabilities{ + Now: now, + Security: protocol.HardwareSecurityCapabilities{ + HardwareAttestations: []protocol.HardwareAttestation{verifiedCPU}, + }, + }, + want: true, + }, + "confidential cpu rejects expired attestation": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialCPU}, + caps: protocol.HardwarePlacementCapabilities{ + Now: now, + Security: protocol.HardwareSecurityCapabilities{HardwareAttestations: []protocol.HardwareAttestation{{ + Class: "confidential-cpu", + Provider: "sev-snp", + VerifierID: "verifier-1", + KeyID: "key-1", + Verified: true, + ExpiresAt: now.Add(-time.Second), + }}}, + }, + }, + "confidential gpu requires gpu and verified attestation": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialGPU}, + caps: protocol.HardwarePlacementCapabilities{ + GPUCount: 1, + Now: now, + Security: protocol.HardwareSecurityCapabilities{ + HardwareAttestations: []protocol.HardwareAttestation{verifiedGPU}, + }, + }, + want: true, + }, + "confidential gpu rejects missing gpu even with attestation": { + req: protocol.PlacementRequirements{ExecutionSecurityTier: protocol.ExecutionConfidentialGPU}, + caps: protocol.HardwarePlacementCapabilities{Security: protocol.HardwareSecurityCapabilities{ + HardwareAttestations: []protocol.HardwareAttestation{verifiedGPU}, + }}, + }, + "custom class accepts hardware_classes": { + req: protocol.PlacementRequirements{HardwareClass: "arm-neoverse"}, + caps: protocol.HardwarePlacementCapabilities{Security: protocol.HardwareSecurityCapabilities{ + HardwareClasses: []string{"arm-neoverse"}, + }}, + want: true, + }, + "custom class accepts tee": { + req: protocol.PlacementRequirements{HardwareClass: "sev-snp"}, + caps: protocol.HardwarePlacementCapabilities{Security: protocol.HardwareSecurityCapabilities{ + TEE: []string{"sev-snp"}, + }}, + want: true, + }, + "custom class with confidential tier requires verified attestation": { + req: protocol.PlacementRequirements{ + ExecutionSecurityTier: protocol.ExecutionConfidentialCPU, + HardwareClass: "sev-snp", + }, + caps: protocol.HardwarePlacementCapabilities{Security: protocol.HardwareSecurityCapabilities{ + HardwareClasses: []string{"sev-snp"}, + }}, + }, + } { + t.Run(name, func(t *testing.T) { + if got := protocol.HardwareCapabilitiesSatisfyPlacementRequirements(tc.req, tc.caps); got != tc.want { + t.Fatalf("HardwareCapabilitiesSatisfyPlacementRequirements() = %v, want %v", got, tc.want) + } + }) + } +} + func TestResourceLimitsRejectNegativeValues(t *testing.T) { limits := protocol.ResourceLimits{ CPUPercent: -1,