66
77use crate :: error:: { ErrorCode , WcError , WcResult } ;
88use crate :: scheduler:: ResourceEnvelope ;
9- use crate :: types:: PeerIdStr ;
9+ use crate :: types:: { AttestationQuote , PeerIdStr } ;
10+ use crate :: verification:: attestation:: { self , MeasurementRegistry } ;
1011use serde:: { Deserialize , Serialize } ;
1112
1213/// Information about a node registered with the broker.
@@ -20,6 +21,10 @@ pub struct NodeInfo {
2021 pub capacity : ResourceEnvelope ,
2122 /// Trust tier (1 = basic, 2 = attested, 3 = TEE).
2223 pub trust_tier : u8 ,
24+ /// Whether this node's attestation has been verified.
25+ pub attestation_verified : bool ,
26+ /// When attestation was last verified (microseconds since epoch).
27+ pub attestation_verified_at : Option < u64 > ,
2328}
2429
2530/// Minimum resource requirements for task placement.
@@ -46,6 +51,8 @@ pub struct Broker {
4651 pub node_roster : Vec < NodeInfo > ,
4752 /// Standby pool — nodes registered but currently unavailable (draining, etc.).
4853 pub standby_pool : Vec < NodeInfo > ,
54+ /// Frozen hosts — removed from scheduling due to incident response.
55+ pub frozen_hosts : Vec < PeerIdStr > ,
4956}
5057
5158impl Broker {
@@ -56,6 +63,7 @@ impl Broker {
5663 region_code : region_code. into ( ) ,
5764 node_roster : Vec :: new ( ) ,
5865 standby_pool : Vec :: new ( ) ,
66+ frozen_hosts : Vec :: new ( ) ,
5967 }
6068 }
6169
@@ -109,6 +117,12 @@ impl Broker {
109117 . map ( |node| node. peer_id . clone ( ) )
110118 . collect ( ) ;
111119
120+ // Exclude frozen hosts (incident response)
121+ let eligible: Vec < PeerIdStr > = eligible
122+ . into_iter ( )
123+ . filter ( |p| !self . frozen_hosts . contains ( p) )
124+ . collect ( ) ;
125+
112126 if eligible. is_empty ( ) {
113127 return Err ( WcError :: new (
114128 ErrorCode :: NoEligibleNodes ,
@@ -118,6 +132,61 @@ impl Broker {
118132
119133 Ok ( eligible)
120134 }
135+
136+ /// Register a node with attestation verification (T041).
137+ ///
138+ /// Per FR-S010/FR-S011: verifies the node's attestation quote against
139+ /// the measurement registry before admitting it to the active roster.
140+ /// Nodes with invalid attestation are rejected. Nodes with no attestation
141+ /// (empty quote) are classified as T0 (trust_tier = 0).
142+ pub fn register_node_with_attestation (
143+ & mut self ,
144+ mut node_info : NodeInfo ,
145+ quote : & AttestationQuote ,
146+ registry : & MeasurementRegistry ,
147+ ) -> WcResult < ( ) > {
148+ // Verify attestation
149+ let verified = attestation:: verify_attestation_with_registry ( quote, registry)
150+ . unwrap_or ( false ) ;
151+
152+ if !verified {
153+ // If quote is non-empty but invalid, reject entirely
154+ if !quote. quote_bytes . is_empty ( ) {
155+ return Err ( WcError :: new (
156+ ErrorCode :: AttestationFailed ,
157+ format ! (
158+ "Node {} presented invalid attestation — rejected (not downgraded to T0)" ,
159+ node_info. peer_id
160+ ) ,
161+ ) ) ;
162+ }
163+ // Empty quote → classify as T0
164+ node_info. trust_tier = 0 ;
165+ node_info. attestation_verified = false ;
166+ } else {
167+ node_info. attestation_verified = true ;
168+ node_info. attestation_verified_at = Some ( crate :: types:: Timestamp :: now ( ) . 0 ) ;
169+ }
170+
171+ self . register_node ( node_info)
172+ }
173+
174+ /// Freeze a host — remove from scheduling pool (incident response).
175+ pub fn freeze_host ( & mut self , peer_id : & PeerIdStr ) {
176+ if !self . frozen_hosts . contains ( peer_id) {
177+ self . frozen_hosts . push ( peer_id. clone ( ) ) ;
178+ }
179+ }
180+
181+ /// Unfreeze a host — restore to scheduling pool.
182+ pub fn unfreeze_host ( & mut self , peer_id : & PeerIdStr ) {
183+ self . frozen_hosts . retain ( |p| p != peer_id) ;
184+ }
185+
186+ /// Check if a host is frozen.
187+ pub fn is_host_frozen ( & self , peer_id : & PeerIdStr ) -> bool {
188+ self . frozen_hosts . contains ( peer_id)
189+ }
121190}
122191
123192#[ cfg( test) ]
@@ -142,6 +211,8 @@ mod tests {
142211 region_code : "us-east-1" . to_string ( ) ,
143212 capacity : test_envelope ( cpu, ram) ,
144213 trust_tier : 1 ,
214+ attestation_verified : false ,
215+ attestation_verified_at : None ,
145216 }
146217 }
147218
@@ -227,4 +298,84 @@ mod tests {
227298 let err = broker. match_task ( & reqs) . unwrap_err ( ) ;
228299 assert_eq ! ( err. code( ) , Some ( ErrorCode :: NoEligibleNodes ) ) ;
229300 }
301+
302+ #[ test]
303+ fn frozen_host_excluded_from_matching ( ) {
304+ let mut broker = Broker :: new ( "broker-001" , "us-east-1" ) ;
305+ broker. register_node ( test_node ( "peer-frozen" , 8000 , 16 * 1024 * 1024 * 1024 ) ) . unwrap ( ) ;
306+ broker. register_node ( test_node ( "peer-active" , 8000 , 16 * 1024 * 1024 * 1024 ) ) . unwrap ( ) ;
307+
308+ broker. freeze_host ( & "peer-frozen" . to_string ( ) ) ;
309+
310+ let reqs = TaskRequirements {
311+ min_cpu_millicores : 1000 ,
312+ min_ram_bytes : 1 ,
313+ min_scratch_bytes : 1 ,
314+ min_trust_tier : 1 ,
315+ } ;
316+ let matched = broker. match_task ( & reqs) . unwrap ( ) ;
317+ assert_eq ! ( matched. len( ) , 1 ) ;
318+ assert_eq ! ( matched[ 0 ] , "peer-active" ) ;
319+ }
320+
321+ #[ test]
322+ fn unfreeze_host_restores_matching ( ) {
323+ let mut broker = Broker :: new ( "broker-001" , "us-east-1" ) ;
324+ broker. register_node ( test_node ( "peer-1" , 8000 , 16 * 1024 * 1024 * 1024 ) ) . unwrap ( ) ;
325+ broker. freeze_host ( & "peer-1" . to_string ( ) ) ;
326+ assert ! ( broker. is_host_frozen( & "peer-1" . to_string( ) ) ) ;
327+
328+ broker. unfreeze_host ( & "peer-1" . to_string ( ) ) ;
329+ assert ! ( !broker. is_host_frozen( & "peer-1" . to_string( ) ) ) ;
330+
331+ let reqs = TaskRequirements {
332+ min_cpu_millicores : 1000 ,
333+ min_ram_bytes : 1 ,
334+ min_scratch_bytes : 1 ,
335+ min_trust_tier : 1 ,
336+ } ;
337+ assert_eq ! ( broker. match_task( & reqs) . unwrap( ) . len( ) , 1 ) ;
338+ }
339+
340+ #[ test]
341+ fn attestation_with_empty_quote_classifies_t0 ( ) {
342+ use crate :: types:: { AttestationQuote , AttestationType } ;
343+ use crate :: verification:: attestation:: MeasurementRegistry ;
344+
345+ let mut broker = Broker :: new ( "broker-001" , "us-east-1" ) ;
346+ let registry = MeasurementRegistry :: new ( ) ;
347+ let mut node = test_node ( "peer-noattest" , 8000 , 16 * 1024 * 1024 * 1024 ) ;
348+ node. trust_tier = 2 ; // claims T2
349+
350+ let empty_quote = AttestationQuote {
351+ quote_type : AttestationType :: Tpm2 ,
352+ quote_bytes : Vec :: new ( ) ,
353+ platform_info : "test" . into ( ) ,
354+ } ;
355+
356+ broker. register_node_with_attestation ( node, & empty_quote, & registry) . unwrap ( ) ;
357+ // Should have been downgraded to T0
358+ assert_eq ! ( broker. node_roster[ 0 ] . trust_tier, 0 ) ;
359+ assert ! ( !broker. node_roster[ 0 ] . attestation_verified) ;
360+ }
361+
362+ #[ test]
363+ fn attestation_with_invalid_quote_rejected ( ) {
364+ use crate :: types:: { AttestationQuote , AttestationType } ;
365+ use crate :: verification:: attestation:: MeasurementRegistry ;
366+
367+ let mut broker = Broker :: new ( "broker-001" , "us-east-1" ) ;
368+ let registry = MeasurementRegistry :: new ( ) ;
369+ let node = test_node ( "peer-bad" , 8000 , 16 * 1024 * 1024 * 1024 ) ;
370+
371+ // Non-empty but garbage quote
372+ let bad_quote = AttestationQuote {
373+ quote_type : AttestationType :: Tpm2 ,
374+ quote_bytes : vec ! [ 0xFF , 0xFE , 0xFD , 0xFC , 0x00 ] ,
375+ platform_info : "test" . into ( ) ,
376+ } ;
377+
378+ let result = broker. register_node_with_attestation ( node, & bad_quote, & registry) ;
379+ assert ! ( result. is_err( ) , "Invalid attestation should reject the node" ) ;
380+ }
230381}
0 commit comments