@@ -18,6 +18,7 @@ use pyth_sdk::{
18
18
} ;
19
19
use solana_program:: clock:: Clock ;
20
20
use solana_program:: pubkey:: Pubkey ;
21
+ use std:: cmp:: min;
21
22
use std:: mem:: size_of;
22
23
23
24
pub use pyth_sdk:: {
@@ -192,7 +193,10 @@ pub struct ProductAccount {
192
193
impl ProductAccount {
193
194
pub fn iter ( & self ) -> AttributeIter {
194
195
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
+ ) ] ,
196
200
}
197
201
}
198
202
}
@@ -574,21 +578,21 @@ impl<'a> Iterator for AttributeIter<'a> {
574
578
if self . attrs . is_empty ( ) {
575
579
return None ;
576
580
}
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) ? ;
579
583
self . attrs = data;
580
584
Some ( ( key, val) )
581
585
}
582
586
}
583
587
584
- fn get_attr_str ( buf : & [ u8 ] ) -> ( & str , & [ u8 ] ) {
588
+ fn get_attr_str ( buf : & [ u8 ] ) -> Option < ( & str , & [ u8 ] ) > {
585
589
if buf. is_empty ( ) {
586
- return ( "" , & [ ] ) ;
590
+ return Some ( ( "" , & [ ] ) ) ;
587
591
}
588
592
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) )
592
596
}
593
597
594
598
#[ cfg( test) ]
@@ -601,6 +605,11 @@ mod test {
601
605
use solana_program:: clock:: Clock ;
602
606
use solana_program:: pubkey:: Pubkey ;
603
607
608
+ use crate :: state:: {
609
+ PROD_ACCT_SIZE ,
610
+ PROD_HDR_SIZE ,
611
+ } ;
612
+
604
613
use super :: {
605
614
PriceInfo ,
606
615
PriceStatus ,
@@ -965,4 +974,70 @@ mod test {
965
974
assert_eq ! ( old_b, new_b) ;
966
975
}
967
976
}
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
+ }
968
1043
}
0 commit comments