1from typing import Optional
2import logging
3
4from . import MemoryMixin
5from ...errors import SimMemoryLimitError, SimMemoryError, SimUnsatError
6
7l = logging.getLogger(__name__)
8
9
10class SizeNormalizationMixin(MemoryMixin):
11    """
12    Provides basic services related to normalizing sizes. After this mixin, sizes will always be a plain int.
13    Assumes that the data is a BV.
14
15    - load will throw a TypeError if no size is provided
16    - store will default to len(data)//byte_width if no size is provided
17    """
18    def load(self, addr, size=None, **kwargs):
19        if size is None:
20            raise TypeError("Must provide size to load")
21        elif type(size) is int:
22            out_size = size
23        elif getattr(size, 'op', None) == 'BVV':
24            out_size = size.args[0]
25        else:
26            raise Exception("Size must be concretely resolved by this point in the memory stack")
27
28        return super().load(addr, size=out_size, **kwargs)
29
30    def store(self, addr, data, size=None, **kwargs):
31        max_size = len(data) // self.state.arch.byte_width
32        if size is None:
33            out_size = max_size
34        elif type(size) is int:
35            out_size = size
36        elif getattr(size, 'op', None) == 'BVV':
37            out_size = size.args[0]
38        else:
39            raise Exception("Size must be concretely resolved by this point in the memory stack")
40
41        if out_size > max_size:
42            raise SimMemoryError("Not enough data for store")
43        if out_size == 0:
44            # skip zero-sized stores
45            return
46
47        super().store(addr, data, size=out_size, **kwargs)
48
49
50class SizeConcretizationMixin(MemoryMixin):
51    """
52    This mixin allows memory to process symbolic sizes. It will not touch any sizes which are not ASTs with non-BVV ops.
53    Assumes that the data is a BV.
54
55    - symbolic load sizes will be concretized as their maximum and a warning will be logged
56    - symbolic store sizes will be dispatched as several conditional stores with concrete sizes
57    """
58    def __init__(self, concretize_symbolic_write_size: bool=False, max_concretize_count: Optional[int]=256,
59                 max_symbolic_size: int=0x400000, raise_memory_limit_error: bool=False, size_limit: int=257, **kwargs):
60        super().__init__(**kwargs)
61        self._concretize_symbolic_write_size = concretize_symbolic_write_size  # in place of the state option CONCRETIZE_SYMBOLIC_WRITE_SIZES
62        self._max_concretize_count = max_concretize_count
63        self._max_symbolic_size = max_symbolic_size
64        self._raise_memory_limit_error = raise_memory_limit_error
65        self._size_limit = size_limit
66
67    def load(self, addr, size=None, **kwargs):
68        if getattr(size, 'op', 'BVV') == 'BVV':
69            return super().load(addr, size=size, **kwargs)
70
71        l.warning("Loading symbolic size via max. be careful.")
72        out_size = self.state.solver.max(size)
73        return super().load(addr, size=out_size, **kwargs)
74
75    def store(self, addr, data, size=None, condition=None, **kwargs):
76        if getattr(size, 'op', 'BVV') == 'BVV':
77            super().store(addr, data, size=size, condition=condition, **kwargs)
78            return
79
80        max_size = len(data) // self.state.arch.byte_width
81        try:
82            if self._raise_memory_limit_error:
83                conc_sizes = list(self.state.solver.eval_upto(
84                    size,
85                    self._size_limit,
86                    extra_constraints=(size <= max_size,)
87                    )
88                )
89                if len(conc_sizes) == self._size_limit:
90                    raise SimMemoryLimitError("Extremely unconstrained store size")
91            else:
92                conc_sizes = list(self.state.solver.eval_upto(
93                    size,
94                    self._max_concretize_count,
95                    extra_constraints=(size <= max_size,)
96                    )
97                )
98        except SimUnsatError:
99            # size has to be greater than max_size
100            raise SimMemoryError("Not enough data for store")
101
102        # filter out all concrete sizes that are greater than max_size
103        # Note that the VSA solver (used in static mode) cannot precisely handle extra constraints. As a result, we may
104        # get conc_sizes with values that violate the extra constraint (size <= max_size).
105        conc_sizes = [ cs for cs in conc_sizes if cs <= max_size ]
106        conc_sizes.sort()
107        if not conc_sizes:
108            raise SimMemoryError("Not enough data for store")
109        if self._max_concretize_count is not None:
110            conc_sizes = conc_sizes[:self._max_concretize_count]
111
112        if size.symbolic:
113            if any(cs > self._max_symbolic_size for cs in conc_sizes):
114                l.warning("At least one concretized size is over the limit of %d bytes. Constrain them to the limit.",
115                        self._max_symbolic_size)
116            conc_sizes = [min(cs, self._max_symbolic_size) for cs in conc_sizes]
117
118        if condition is None:
119            condition = self.state.solver.true
120        for conc_size in conc_sizes:
121            if conc_size == 0:
122                continue
123            super().store(addr, data, size=conc_size, condition=condition & (size == conc_size), **kwargs)
124
125
126