@@ -13,135 +13,145 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1313See the License for the specific language governing permissions and
1414limitations under the License.
1515*/
16- use anyhow:: { bail, Result } ;
1716use alloc:: vec:: Vec ;
18- use core:: any:: type_name;
1917use core:: slice:: from_raw_parts_mut;
2018
19+ use anyhow:: { bail, Result } ;
20+
2121pub struct InputDataSection {
2222 ptr : * mut u8 ,
23- len : usize ,
23+ len : u64 ,
2424}
2525
2626impl InputDataSection {
27- pub fn new ( ptr : * mut u8 , len : usize ) -> Self {
27+ pub fn new ( ptr : * mut u8 , len : u64 ) -> Self {
2828 InputDataSection { ptr, len }
2929 }
3030
3131 pub fn try_pop_shared_input_data_into < T > ( & self ) -> Result < T >
3232 where
3333 T : for < ' a > TryFrom < & ' a [ u8 ] > ,
3434 {
35- let input_data_buffer = unsafe { from_raw_parts_mut ( self . ptr , self . len ) } ;
35+ let shared_buffer_size = self . len as usize ;
3636
37- if input_data_buffer. is_empty ( ) {
38- bail ! ( "Got a 0-size buffer in pop_shared_input_data_into" ) ;
37+ let idb = unsafe { from_raw_parts_mut ( self . ptr , shared_buffer_size) } ;
38+
39+ if idb. is_empty ( ) {
40+ bail ! ( "Got a 0-size buffer in try_pop_shared_input_data_into" ) ;
3941 }
4042
4143 // get relative offset to next free address
42- let stack_ptr_rel: usize = usize:: from_le_bytes (
43- input_data_buffer[ ..8 ]
44- . try_into ( )
45- . expect ( "Shared input buffer too small" ) ,
46- ) ;
47-
48- if stack_ptr_rel > self . len || stack_ptr_rel < 16 {
49- bail ! ( "Invalid stack pointer: {} in pop_shared_input_data_into" , stack_ptr_rel) ;
44+ let stack_ptr_rel: u64 = u64:: from_le_bytes ( match idb[ ..8 ] . try_into ( ) {
45+ Ok ( bytes) => bytes,
46+ Err ( _) => bail ! ( "shared input buffer too small" ) ,
47+ } ) ;
48+
49+ if stack_ptr_rel as usize > shared_buffer_size || stack_ptr_rel < 16 {
50+ bail ! (
51+ "Invalid stack pointer: {} in try_pop_shared_input_data_into" ,
52+ stack_ptr_rel
53+ ) ;
5054 }
5155
5256 // go back 8 bytes and read. This is the offset to the element on top of stack
53- let last_element_offset_rel = usize:: from_le_bytes (
54- input_data_buffer[ stack_ptr_rel - 8 ..stack_ptr_rel]
55- . try_into ( )
56- . expect ( "Invalid stack pointer in pop_shared_input_data_into" ) ,
57+ let last_element_offset_rel = u64:: from_le_bytes (
58+ match idb[ stack_ptr_rel as usize - 8 ..stack_ptr_rel as usize ] . try_into ( ) {
59+ Ok ( bytes) => bytes,
60+ Err ( _) => bail ! ( "Invalid stack pointer in pop_shared_input_data_into" ) ,
61+ } ,
5762 ) ;
5863
59- let buffer = & input_data_buffer [ last_element_offset_rel..] ;
64+ let buffer = & idb [ last_element_offset_rel as usize ..] ;
6065
6166 // convert the buffer to T
6267 let type_t = match T :: try_from ( buffer) {
6368 Ok ( t) => Ok ( t) ,
64- Err ( _e) => {
65- bail ! ( "Unable to convert buffer to {}" , type_name:: <T >( ) ) ;
66- }
69+ Err ( _e) => bail ! ( "failed to convert buffer to type T in pop_shared_input_data_into" ) ,
6770 } ;
6871
6972 // update the stack pointer to point to the element we just popped of since that is now free
70- input_data_buffer [ ..8 ] . copy_from_slice ( & last_element_offset_rel. to_le_bytes ( ) ) ;
73+ idb [ ..8 ] . copy_from_slice ( & last_element_offset_rel. to_le_bytes ( ) ) ;
7174
7275 // zero out popped off buffer
73- input_data_buffer [ last_element_offset_rel..stack_ptr_rel] . fill ( 0 ) ;
76+ idb [ last_element_offset_rel as usize ..stack_ptr_rel as usize ] . fill ( 0 ) ;
7477
7578 type_t
7679 }
7780}
7881
7982pub struct OutputDataSection {
8083 pub ptr : * mut u8 ,
81- pub len : usize ,
84+ pub len : u64 ,
8285}
8386
8487impl OutputDataSection {
85- pub fn new ( ptr : * mut u8 , len : usize ) -> Self {
88+ const STACK_PTR_SIZE : usize = size_of :: < u64 > ( ) ;
89+
90+ pub fn new ( ptr : * mut u8 , len : u64 ) -> Self {
8691 OutputDataSection { ptr, len }
8792 }
8893
8994 pub fn push_shared_output_data ( & self , data : Vec < u8 > ) -> Result < ( ) > {
90- let output_data_buffer = unsafe { from_raw_parts_mut ( self . ptr , self . len ) } ;
95+ let shared_buffer_size = self . len as usize ;
96+ let odb: & mut [ u8 ] = unsafe { from_raw_parts_mut ( self . ptr , shared_buffer_size) } ;
9197
92- if output_data_buffer . is_empty ( ) {
93- bail ! ( "Got a 0-size buffer in push_shared_output_data " ) ;
98+ if odb . len ( ) < Self :: STACK_PTR_SIZE {
99+ bail ! ( "shared output buffer is too small " ) ;
94100 }
95101
96102 // get offset to next free address on the stack
97- let mut stack_ptr_rel: usize = usize:: from_le_bytes (
98- output_data_buffer[ ..8 ]
99- . try_into ( )
100- . expect ( "Shared output buffer too small" ) ,
101- ) ;
102-
103+ let mut stack_ptr_rel: u64 =
104+ u64:: from_le_bytes ( match odb[ ..Self :: STACK_PTR_SIZE ] . try_into ( ) {
105+ Ok ( bytes) => bytes,
106+ Err ( _) => bail ! ( "failed to get stack pointer in shared output buffer" ) ,
107+ } ) ;
108+
109+ // if stack_ptr_rel is 0, it means this is the first time we're using the output buffer, so
110+ // we want to offset it by 8 as to not overwrite the stack_ptr location.
103111 if stack_ptr_rel == 0 {
104112 stack_ptr_rel = 8 ;
105113 }
106114
107115 // check if the stack pointer is within the bounds of the buffer.
108116 // It can be equal to the size, but never greater
109117 // It can never be less than 8. An empty buffer's stack pointer is 8
110- if stack_ptr_rel > self . len || stack_ptr_rel < 8 {
111- bail ! ( "Invalid stack pointer: {} in push_shared_output_data" , stack_ptr_rel ) ;
118+ if stack_ptr_rel as usize > shared_buffer_size {
119+ bail ! ( "invalid stack pointer in shared output buffer" ) ;
112120 }
113121
114122 // check if there is enough space in the buffer
115- let size_required = data. len ( ) + 8 ; // the data plus the pointer pointing to the data
116- let size_available = self . len - stack_ptr_rel;
123+ let size_required: usize = data. len ( ) + 8 ; // the data plus the pointer pointing to the data
124+ let size_available: usize = shared_buffer_size - stack_ptr_rel as usize ;
117125 if size_required > size_available {
118- bail ! ( "Not enough space in shared output buffer. Required: {}, Available: {}" , size_required , size_available ) ;
126+ bail ! ( "not enough space in shared output buffer" ) ;
119127 }
120128
121129 // write the actual data
122- output_data_buffer [ stack_ptr_rel..stack_ptr_rel + data. len ( ) ] . copy_from_slice ( & data) ;
130+ odb [ stack_ptr_rel as usize ..stack_ptr_rel as usize + data. len ( ) ] . copy_from_slice ( & data) ;
123131
124132 // write the offset to the newly written data, to the top of the stack
125- let bytes = stack_ptr_rel. to_le_bytes ( ) ;
126- output_data_buffer[ stack_ptr_rel + data. len ( ) ..stack_ptr_rel + data. len ( ) + 8 ]
133+ let bytes: [ u8 ; Self :: STACK_PTR_SIZE ] = stack_ptr_rel. to_le_bytes ( ) ;
134+ odb[ stack_ptr_rel as usize + data. len ( )
135+ ..stack_ptr_rel as usize + data. len ( ) + Self :: STACK_PTR_SIZE ]
127136 . copy_from_slice ( & bytes) ;
128137
129138 // update stack pointer to point to next free address
130- let new_stack_ptr_rel = stack_ptr_rel + data. len ( ) + 8 ;
131- output_data_buffer[ 0 ..8 ] . copy_from_slice ( & new_stack_ptr_rel. to_le_bytes ( ) ) ;
139+ let new_stack_ptr_rel: u64 =
140+ ( stack_ptr_rel as usize + data. len ( ) + Self :: STACK_PTR_SIZE ) as u64 ;
141+ odb[ 0 ..Self :: STACK_PTR_SIZE ] . copy_from_slice ( & new_stack_ptr_rel. to_le_bytes ( ) ) ;
132142
133143 Ok ( ( ) )
134144 }
135145}
136146
137147impl From < ( u64 , u64 ) > for InputDataSection {
138148 fn from ( ( ptr, len) : ( u64 , u64 ) ) -> Self {
139- InputDataSection :: new ( ptr as * mut u8 , len as usize )
149+ InputDataSection :: new ( ptr as * mut u8 , len)
140150 }
141151}
142152
143153impl From < ( u64 , u64 ) > for OutputDataSection {
144154 fn from ( ( ptr, len) : ( u64 , u64 ) ) -> Self {
145- OutputDataSection :: new ( ptr as * mut u8 , len as usize )
155+ OutputDataSection :: new ( ptr as * mut u8 , len)
146156 }
147157}
0 commit comments