@@ -18,6 +18,7 @@ use pyth_sdk::{
1818} ;
1919use solana_program:: clock:: Clock ;
2020use solana_program:: pubkey:: Pubkey ;
21+ use std:: cmp:: min;
2122use std:: mem:: size_of;
2223
2324pub use pyth_sdk:: {
@@ -192,7 +193,10 @@ pub struct ProductAccount {
192193impl ProductAccount {
193194 pub fn iter ( & self ) -> AttributeIter {
194195 AttributeIter {
195- attrs : & self . attr [ ..( self . size as usize - PROD_HDR_SIZE ) ] ,
196+ attrs : & self . attr [ ..min (
197+ ( self . size as usize ) . saturating_sub ( PROD_HDR_SIZE ) ,
198+ PROD_ATTR_SIZE ,
199+ ) ] ,
196200 }
197201 }
198202}
@@ -574,21 +578,21 @@ impl<'a> Iterator for AttributeIter<'a> {
574578 if self . attrs . is_empty ( ) {
575579 return None ;
576580 }
577- let ( key, data) = get_attr_str ( self . attrs ) ;
578- let ( val, data) = get_attr_str ( data) ;
581+ let ( key, data) = get_attr_str ( self . attrs ) ? ;
582+ let ( val, data) = get_attr_str ( data) ? ;
579583 self . attrs = data;
580584 Some ( ( key, val) )
581585 }
582586}
583587
584- fn get_attr_str ( buf : & [ u8 ] ) -> ( & str , & [ u8 ] ) {
588+ fn get_attr_str ( buf : & [ u8 ] ) -> Option < ( & str , & [ u8 ] ) > {
585589 if buf. is_empty ( ) {
586- return ( "" , & [ ] ) ;
590+ return Some ( ( "" , & [ ] ) ) ;
587591 }
588592 let len = buf[ 0 ] as usize ;
589- let str = std:: str:: from_utf8 ( & buf[ 1 ..len + 1 ] ) . expect ( "attr should be ascii or utf-8" ) ;
590- let remaining_buf = & buf[ len + 1 ..] ;
591- ( str, remaining_buf)
593+ let str = std:: str:: from_utf8 ( buf. get ( 1 ..len + 1 ) ? ) . ok ( ) ? ;
594+ let remaining_buf = & buf. get ( len + 1 ..) ? ;
595+ Some ( ( str, remaining_buf) )
592596}
593597
594598#[ cfg( test) ]
@@ -601,6 +605,11 @@ mod test {
601605 use solana_program:: clock:: Clock ;
602606 use solana_program:: pubkey:: Pubkey ;
603607
608+ use crate :: state:: {
609+ PROD_ACCT_SIZE ,
610+ PROD_HDR_SIZE ,
611+ } ;
612+
604613 use super :: {
605614 PriceInfo ,
606615 PriceStatus ,
@@ -965,4 +974,70 @@ mod test {
965974 assert_eq ! ( old_b, new_b) ;
966975 }
967976 }
977+
978+ #[ test]
979+ fn test_product_account_iter_works ( ) {
980+ let mut product = super :: ProductAccount {
981+ magic : 1 ,
982+ ver : 2 ,
983+ atype : super :: AccountType :: Product as u32 ,
984+ size : PROD_HDR_SIZE as u32 + 10 ,
985+ px_acc : Pubkey :: new_from_array ( [ 3 ; 32 ] ) ,
986+ attr : [ 0 ; super :: PROD_ATTR_SIZE ] ,
987+ } ;
988+
989+ // Set some attributes
990+ product. attr [ 0 ] = 3 ; // key length
991+ product. attr [ 1 ..4 ] . copy_from_slice ( b"key" ) ;
992+ product. attr [ 4 ] = 5 ; // value length
993+ product. attr [ 5 ..10 ] . copy_from_slice ( b"value" ) ;
994+
995+ let mut iter = product. iter ( ) ;
996+ assert_eq ! ( iter. next( ) , Some ( ( "key" , "value" ) ) ) ;
997+ assert_eq ! ( iter. next( ) , None ) ;
998+
999+ // Check that the iterator does not panic on size misconfiguration
1000+ product. size = PROD_HDR_SIZE as u32 - 10 ; // Invalid size
1001+ let mut iter = product. iter ( ) ;
1002+ assert_eq ! ( iter. next( ) , None ) ; // Should not panic, just return None
1003+
1004+ product. size = PROD_ACCT_SIZE as u32 + 10 ; // Reset size to a size larger than the account size
1005+ let mut iter = product. iter ( ) ;
1006+ assert_eq ! ( iter. next( ) , Some ( ( "key" , "value" ) ) ) ;
1007+ while iter. next ( ) . is_some ( ) { } // Consume the iterator
1008+
1009+ // Check that invalid len stops the iterator. This behaviour is not perfect as it
1010+ // stops reading attributes after the first invalid one but is just a safety measure.
1011+ // In this case, we set the length byte to 255 which goes beyond the size of the
1012+ // product account.
1013+ product. attr [ 10 ] = 255 ;
1014+ for i in 11 ..266 {
1015+ product. attr [ i] = b'a' ;
1016+ }
1017+ product. attr [ 266 ] = 255 ;
1018+ for i in 267 ..super :: PROD_ATTR_SIZE {
1019+ product. attr [ i] = b'b' ;
1020+ }
1021+ let mut iter = product. iter ( ) ;
1022+ assert_eq ! ( iter. next( ) , Some ( ( "key" , "value" ) ) ) ;
1023+ assert_eq ! ( iter. next( ) , None ) ; // No more attributes because it stopped reading the invalid value
1024+
1025+ // Make sure if the value size was set to a smaller value, it would work fine
1026+ product. attr [ 266 ] = 10 ;
1027+ let mut iter = product. iter ( ) ;
1028+ assert_eq ! ( iter. next( ) , Some ( ( "key" , "value" ) ) ) ;
1029+ let ( key, val) = iter. next ( ) . unwrap ( ) ;
1030+ assert_eq ! ( key. len( ) , 255 ) ;
1031+ for byte in key. as_bytes ( ) {
1032+ assert_eq ! ( byte, & b'a' ) ;
1033+ }
1034+ assert_eq ! ( val, "bbbbbbbbbb" ) ; // No more attributes because it stopped reading the invalid value
1035+
1036+ // Check that iterator stops on non-UTF8 attributes. This behaviour is not
1037+ // perfect as it stops reading attributes after the first non-UTF8 one but
1038+ // is just a safety measure.
1039+ product. attr [ 1 ..4 ] . copy_from_slice ( b"\xff \xfe \xfa " ) ;
1040+ let mut iter = product. iter ( ) ;
1041+ assert_eq ! ( iter. next( ) , None ) ; // Should not panic, just return None
1042+ }
9681043}
0 commit comments