1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
use eth_types::{bytecode, Bytecode, U256};
use rand::{rngs::ThreadRng, Rng};

/// Generate Sha3 opcode
pub struct Sha3CodeGen {
    /// The offset
    pub offset: usize,
    /// The size
    pub size: usize,
    data_len: usize,
    rng: ThreadRng,
}
impl Sha3CodeGen {
    /// Construct with memory less than size
    pub fn mem_lt_size(offset: usize, size: usize) -> Self {
        let mut rng = rand::thread_rng();
        let data_len = offset
            + if size.gt(&0) {
                rng.gen_range(0..size)
            } else {
                0
            };
        Self {
            offset,
            size,
            data_len,
            rng,
        }
    }
    /// Construct with memory equal to size
    pub fn mem_eq_size(offset: usize, size: usize) -> Self {
        let data_len = offset + size;
        Self {
            offset,
            size,
            data_len,
            rng: rand::thread_rng(),
        }
    }
    /// Construct with memory greater than size
    pub fn mem_gt_size(offset: usize, size: usize) -> Self {
        let mut rng = rand::thread_rng();
        let data_len = offset
            + size
            + if size.gt(&0) {
                rng.gen_range(0..size)
            } else {
                0
            };
        Self {
            offset,
            size,
            data_len,
            rng,
        }
    }
    /// Construct with empty memory
    pub fn mem_empty(offset: usize, size: usize) -> Self {
        Self {
            offset,
            size,
            data_len: 0,
            rng: rand::thread_rng(),
        }
    }
    fn rand_bytes(&mut self) -> Vec<u8> {
        (0..self.data_len)
            .map(|_| self.rng.gen())
            .collect::<Vec<u8>>()
    }
    /// Generate bytecode for SHA3 opcode after having populated sufficient
    /// memory given the offset and size arguments for SHA3.
    pub fn gen_sha3_code(&mut self) -> (Bytecode, Vec<u8>) {
        let data = self.rand_bytes();
        let mut memory = Vec::with_capacity(self.data_len);

        // add opcodes to populate memory in the current context.
        let mut code = Bytecode::default();
        for (i, mem_chunk) in data.chunks(32).enumerate() {
            let mem_value = if mem_chunk.len() < 32 {
                std::iter::repeat(0u8)
                    .take(32 - mem_chunk.len())
                    .chain(mem_chunk.to_vec())
                    .collect::<Vec<u8>>()
            } else {
                mem_chunk.to_vec()
            };
            memory.extend_from_slice(&mem_value);
            code.op_mstore(32 * i, U256::from_big_endian(&mem_value));
        }
        // append SHA3 related opcodes at the tail end.
        let code_tail = bytecode! {
            PUSH32(self.size)
            PUSH32(self.offset)
            SHA3
            STOP
        };
        code.append(&code_tail);
        (code, memory)
    }
}