sbpf_vm/
memory.rs

1use {
2    crate::errors::{SbpfVmError, SbpfVmResult},
3    serde::{Deserialize, Serialize},
4};
5
6/// Memory region
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum MemoryRegion {
9    Input,
10    Rodata,
11    Stack,
12    Heap,
13}
14
15/// Memory layout
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Memory {
18    pub rodata: Vec<u8>,
19    pub stack: Vec<u8>,
20    pub heap: Vec<u8>,
21    pub input: Vec<u8>,
22    pub heap_ptr: usize,
23}
24
25impl Memory {
26    // Virtual address memory map
27    pub const RODATA_START: u64 = 0x0; // Read-only data (rodata)
28    pub const STACK_START: u64 = 0x200000000; // Stack data
29    pub const HEAP_START: u64 = 0x300000000; // Heap data
30    pub const INPUT_START: u64 = 0x400000000; // Program input parameters
31
32    pub const DEFAULT_HEAP_SIZE: usize = 32768; // 32KB
33    pub const STACK_FRAME_SIZE: u64 = 4096; // 4KB
34
35    pub fn new(input: Vec<u8>, rodata: Vec<u8>, stack_size: usize, heap_size: usize) -> Self {
36        Self {
37            input,
38            rodata,
39            stack: vec![0u8; stack_size],
40            heap: vec![0u8; heap_size],
41            heap_ptr: 0,
42        }
43    }
44
45    pub fn initial_frame_pointer(&self) -> u64 {
46        Self::STACK_START + Self::STACK_FRAME_SIZE
47    }
48
49    pub fn stack_size(max_call_depth: usize) -> usize {
50        Self::STACK_FRAME_SIZE as usize * max_call_depth
51    }
52
53    // Translate virtual address to region and offset
54    fn translate(&self, addr: u64) -> SbpfVmResult<(MemoryRegion, usize)> {
55        if addr >= Self::INPUT_START {
56            let offset = (addr - Self::INPUT_START) as usize;
57            if offset < self.input.len() {
58                Ok((MemoryRegion::Input, offset))
59            } else {
60                Err(SbpfVmError::MemoryOutOfBounds(addr, 0))
61            }
62        } else if addr >= Self::HEAP_START {
63            let offset = (addr - Self::HEAP_START) as usize;
64            if offset < self.heap.len() {
65                Ok((MemoryRegion::Heap, offset))
66            } else {
67                Err(SbpfVmError::MemoryOutOfBounds(addr, 0))
68            }
69        } else if addr >= Self::STACK_START {
70            let offset = (addr - Self::STACK_START) as usize;
71            if offset < self.stack.len() {
72                Ok((MemoryRegion::Stack, offset))
73            } else {
74                Err(SbpfVmError::MemoryOutOfBounds(addr, 0))
75            }
76        } else {
77            let offset = addr as usize;
78            if offset < self.rodata.len() {
79                Ok((MemoryRegion::Rodata, offset))
80            } else {
81                Err(SbpfVmError::MemoryOutOfBounds(addr, 0))
82            }
83        }
84    }
85
86    fn get_slice(&self, region: MemoryRegion, offset: usize, len: usize) -> SbpfVmResult<&[u8]> {
87        let data = match region {
88            MemoryRegion::Input => &self.input,
89            MemoryRegion::Rodata => &self.rodata,
90            MemoryRegion::Stack => &self.stack,
91            MemoryRegion::Heap => &self.heap,
92        };
93
94        if offset + len > data.len() {
95            return Err(SbpfVmError::MemoryOutOfBounds(offset as u64, len));
96        }
97
98        Ok(&data[offset..offset + len])
99    }
100
101    fn get_slice_mut(
102        &mut self,
103        region: MemoryRegion,
104        offset: usize,
105        len: usize,
106    ) -> SbpfVmResult<&mut [u8]> {
107        // Rodata region is read-only
108        if region == MemoryRegion::Rodata {
109            return Err(SbpfVmError::InvalidMemoryAccess(
110                Self::RODATA_START + offset as u64,
111            ));
112        }
113
114        let data = match region {
115            MemoryRegion::Input => &mut self.input,
116            MemoryRegion::Stack => &mut self.stack,
117            MemoryRegion::Heap => &mut self.heap,
118            MemoryRegion::Rodata => unreachable!(),
119        };
120
121        if offset + len > data.len() {
122            return Err(SbpfVmError::MemoryOutOfBounds(offset as u64, len));
123        }
124
125        Ok(&mut data[offset..offset + len])
126    }
127
128    pub fn read_u8(&self, addr: u64) -> SbpfVmResult<u8> {
129        let (region, offset) = self.translate(addr)?;
130        let slice = self.get_slice(region, offset, 1)?;
131        Ok(slice[0])
132    }
133
134    pub fn read_u16(&self, addr: u64) -> SbpfVmResult<u16> {
135        let (region, offset) = self.translate(addr)?;
136        let slice = self.get_slice(region, offset, 2)?;
137        Ok(u16::from_le_bytes([slice[0], slice[1]]))
138    }
139
140    pub fn read_u32(&self, addr: u64) -> SbpfVmResult<u32> {
141        let (region, offset) = self.translate(addr)?;
142        let slice = self.get_slice(region, offset, 4)?;
143        Ok(u32::from_le_bytes([slice[0], slice[1], slice[2], slice[3]]))
144    }
145
146    pub fn read_u64(&self, addr: u64) -> SbpfVmResult<u64> {
147        let (region, offset) = self.translate(addr)?;
148        let slice = self.get_slice(region, offset, 8)?;
149        Ok(u64::from_le_bytes([
150            slice[0], slice[1], slice[2], slice[3], slice[4], slice[5], slice[6], slice[7],
151        ]))
152    }
153
154    pub fn read_bytes(&self, addr: u64, len: usize) -> SbpfVmResult<&[u8]> {
155        let (region, offset) = self.translate(addr)?;
156        self.get_slice(region, offset, len)
157    }
158
159    pub fn write_u8(&mut self, addr: u64, value: u8) -> SbpfVmResult<()> {
160        let (region, offset) = self.translate(addr)?;
161        let slice = self.get_slice_mut(region, offset, 1)?;
162        slice[0] = value;
163        Ok(())
164    }
165
166    pub fn write_u16(&mut self, addr: u64, value: u16) -> SbpfVmResult<()> {
167        let (region, offset) = self.translate(addr)?;
168        let slice = self.get_slice_mut(region, offset, 2)?;
169        slice.copy_from_slice(&value.to_le_bytes());
170        Ok(())
171    }
172
173    pub fn write_u32(&mut self, addr: u64, value: u32) -> SbpfVmResult<()> {
174        let (region, offset) = self.translate(addr)?;
175        let slice = self.get_slice_mut(region, offset, 4)?;
176        slice.copy_from_slice(&value.to_le_bytes());
177        Ok(())
178    }
179
180    pub fn write_u64(&mut self, addr: u64, value: u64) -> SbpfVmResult<()> {
181        let (region, offset) = self.translate(addr)?;
182        let slice = self.get_slice_mut(region, offset, 8)?;
183        slice.copy_from_slice(&value.to_le_bytes());
184        Ok(())
185    }
186
187    pub fn write_i64(&mut self, addr: u64, value: i64) -> SbpfVmResult<()> {
188        let (region, offset) = self.translate(addr)?;
189        let slice = self.get_slice_mut(region, offset, 8)?;
190        slice.copy_from_slice(&value.to_le_bytes());
191        Ok(())
192    }
193
194    pub fn write_bytes(&mut self, addr: u64, bytes: &[u8]) -> SbpfVmResult<()> {
195        let (region, offset) = self.translate(addr)?;
196        let slice = self.get_slice_mut(region, offset, bytes.len())?;
197        slice.copy_from_slice(bytes);
198        Ok(())
199    }
200
201    pub fn alloc(&mut self, size: usize) -> SbpfVmResult<u64> {
202        if self.heap_ptr + size > self.heap.len() {
203            return Err(SbpfVmError::MemoryOutOfBounds(
204                Self::HEAP_START + self.heap_ptr as u64,
205                size,
206            ));
207        }
208        let addr = Self::HEAP_START + self.heap_ptr as u64;
209        self.heap_ptr += size;
210        Ok(addr)
211    }
212
213    pub fn reset_heap(&mut self) {
214        self.heap_ptr = 0;
215        self.heap.fill(0);
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    #[test]
224    fn test_memory_regions() {
225        let input = vec![1, 2, 3, 4];
226        let rodata = vec![5, 6, 7, 8];
227        let memory = Memory::new(input, rodata, 1024, 1024);
228
229        // Test input and rodata region
230        assert_eq!(memory.read_u8(Memory::INPUT_START).unwrap(), 1);
231        assert_eq!(memory.read_u8(Memory::INPUT_START + 3).unwrap(), 4);
232
233        assert_eq!(memory.read_u8(Memory::RODATA_START).unwrap(), 5);
234        assert_eq!(memory.read_u8(Memory::RODATA_START + 3).unwrap(), 8);
235    }
236
237    #[test]
238    fn test_read_write() {
239        let mut memory = Memory::new(
240            vec![0; 16],
241            vec![0; 16],
242            Memory::STACK_FRAME_SIZE as usize,
243            1024,
244        );
245
246        let fp = memory.initial_frame_pointer();
247
248        // Write and read u8
249        let addr = fp - 1;
250        memory.write_u8(addr, 0x5).unwrap();
251        assert_eq!(memory.read_u8(addr).unwrap(), 0x5);
252
253        // Write and read u16
254        let addr = fp - 2;
255        memory.write_u16(addr, 0xabcd).unwrap();
256        assert_eq!(memory.read_u16(addr).unwrap(), 0xabcd);
257
258        // Write and read u32
259        let addr = fp - 4;
260        memory.write_u32(addr, 0xabcd1234).unwrap();
261        assert_eq!(memory.read_u32(addr).unwrap(), 0xabcd1234);
262
263        // Write and read u64
264        let addr = fp - 8;
265        memory.write_u64(addr, 0x123456789abcdef0).unwrap();
266        assert_eq!(memory.read_u64(addr).unwrap(), 0x123456789abcdef0);
267    }
268
269    #[test]
270    fn test_heap_allocation() {
271        let mut memory = Memory::new(vec![], vec![], 1024, 1024);
272
273        let addr1 = memory.alloc(64).unwrap();
274        assert_eq!(addr1, Memory::HEAP_START);
275
276        let addr2 = memory.alloc(128).unwrap();
277        assert_eq!(addr2, Memory::HEAP_START + 64);
278
279        memory.write_u64(addr1, 0x12345678).unwrap();
280        assert_eq!(memory.read_u64(addr1).unwrap(), 0x12345678);
281    }
282
283    #[test]
284    fn test_rodata_readonly() {
285        let mut memory = Memory::new(vec![], vec![1, 2, 3, 4], 1024, 1024);
286
287        // should fail to write to read-only region
288        let result = memory.write_u8(Memory::RODATA_START, 12);
289        assert!(result.is_err());
290    }
291}
sbpf_vm/memory.rs

sbpf_vm/
memory.rs
