4
4
TODO: Leverage Hypothesis!
5
5
"""
6
6
from abc import ABC , abstractmethod
7
+ from enum import Enum , auto
7
8
import random
8
9
import string
9
- from typing import List , Optional , Sequence , cast
10
+ from typing import List , Optional , Sequence , Type , cast
10
11
11
12
from algosdk import abi , encoding
12
13
14
+
13
15
from graviton .models import PyTypes
14
16
15
17
16
18
class ABIStrategy (ABC ):
19
+ """
20
+ TODO: when incorporating hypothesis strategies, we'll need a more holistic
21
+ approach that looks at relationships amongst various args.
22
+ Current approach only looks at each argument as a completely independent entity.
23
+ """
24
+
25
+ @abstractmethod
26
+ def __init__ (self , abi_instance : abi .ABIType , dynamic_length : Optional [int ] = None ):
27
+ pass
28
+
17
29
@abstractmethod
18
30
def get (self ) -> PyTypes :
19
31
pass
@@ -167,6 +179,15 @@ def address_logic(x):
167
179
168
180
169
181
class RandomABIStrategyHalfSized (RandomABIStrategy ):
182
+ """
183
+ This strategy only generates data that is half the size that _ought_ to be possible.
184
+ This is useful in the case that operations involving the generated arguments
185
+ could overflow due to multiplication.
186
+
187
+ Since this only makes sense for `abi.UintType`, it degenerates to the standard
188
+ `RandomABIStrategy` for other types.
189
+ """
190
+
170
191
def __init__ (
171
192
self ,
172
193
abi_instance : abi .ABIType ,
@@ -183,3 +204,229 @@ def get(self) -> PyTypes:
183
204
return cast (int , full_random ) % (
184
205
1 << (cast (abi .UintType , self .abi_type ).bit_size // 2 )
185
206
)
207
+
208
+
209
+ class ABIArgsMod (Enum ):
210
+ # insert a random byte into selector:
211
+ selector_byte_insert = auto ()
212
+ # delete a byte at a random position from the selector:
213
+ selector_byte_delete = auto ()
214
+ # replace a random byte in the selector:
215
+ selector_byte_replace = auto ()
216
+ # delete a random argument:
217
+ parameter_delete = auto ()
218
+ # insert a random argument:
219
+ parameter_append = auto ()
220
+
221
+
222
+ class CallStrategy (ABC ):
223
+ def __init__ (
224
+ self ,
225
+ argument_strategy : Type [ABIStrategy ] = RandomABIStrategy ,
226
+ * ,
227
+ num_dryruns : int = 1 ,
228
+ ):
229
+ self .argument_strategy : Type [ABIStrategy ] = argument_strategy
230
+ self .num_dryruns : int = num_dryruns
231
+
232
+ def generate_value (self , gen_type : abi .ABIType ) -> PyTypes :
233
+ return cast (Type [ABIStrategy ], self .argument_strategy )(gen_type ).get ()
234
+
235
+ @abstractmethod
236
+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
237
+ pass
238
+
239
+
240
+ class ABICallStrategy (CallStrategy ):
241
+ """
242
+ TODO: refactor to comport with ABIStrategy + Hypothesis
243
+ TODO: make this generic on the strategy type
244
+ """
245
+
246
+ append_args_type : abi .ABIType = abi .ByteType ()
247
+
248
+ def __init__ (
249
+ self ,
250
+ contract : str ,
251
+ argument_strategy : Type [ABIStrategy ] = RandomABIStrategy ,
252
+ * ,
253
+ num_dryruns : int = 1 ,
254
+ handle_selector : bool = True ,
255
+ abi_args_mod : Optional [ABIArgsMod ] = None ,
256
+ ):
257
+ """
258
+ contract - ABI Contract JSON
259
+
260
+ argument_strategy (default=RandomABIStrategy) - ABI strategy for generating arguments
261
+
262
+ num_dry_runs (default=1) - the number of dry runs to run
263
+ (generates different inputs each time)
264
+
265
+ handle_selector (default=True) - usually we'll want to let
266
+ `ABIContractExecutor.run_sequence()`
267
+ handle adding the method selector so this param.
268
+ But if set False: when providing `inputs`
269
+ ensure that the 0'th argument for method calls is the selector.
270
+ And when set True: when NOT providing `inputs`, the selector arg
271
+ at index 0 will be added automatically.
272
+
273
+ abi_args_mod (optional) - when desiring to mutate the args, provide an ABIArgsMod value
274
+ """
275
+ super ().__init__ (argument_strategy , num_dryruns = num_dryruns )
276
+ self .contract : abi .Contract = abi .Contract .from_json (contract )
277
+ self .handle_selector = handle_selector
278
+ self .abi_args_mod = abi_args_mod
279
+
280
+ def abi_method (self , method : Optional [str ]) -> abi .Method :
281
+ assert method , "cannot get abi.Method for bare app call"
282
+
283
+ return self .contract .get_method_by_name (method )
284
+
285
+ def method_signature (self , method : Optional [str ]) -> Optional [str ]:
286
+ """Returns None, for a bare app call (method=None signals this)"""
287
+ if method is None :
288
+ return None
289
+
290
+ return self .abi_method (method ).get_signature ()
291
+
292
+ def method_selector (self , method : Optional [str ]) -> bytes :
293
+ assert method , "cannot get method_selector for bare app call"
294
+
295
+ return self .abi_method (method ).get_selector ()
296
+
297
+ def argument_types (self , method : Optional [str ]) -> List [abi .ABIType ]:
298
+ """
299
+ Argument types (excluding selector)
300
+ """
301
+ if method is None :
302
+ return []
303
+
304
+ return [cast (abi .ABIType , arg .type ) for arg in self .abi_method (method ).args ]
305
+
306
+ def num_args (self , method : Optional [str ]) -> int :
307
+ return len (self .argument_types (method ))
308
+
309
+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
310
+ """
311
+ Generates inputs appropriate for bare app calls and method calls
312
+ according to available argument_strategy.
313
+ """
314
+ assert (
315
+ self .argument_strategy
316
+ ), "cannot generate inputs without an argument_strategy"
317
+
318
+ mutating = self .abi_args_mod is not None
319
+
320
+ if not (method or mutating ):
321
+ # bare calls receive no arguments (unless mutating)
322
+ return [tuple () for _ in range (self .num_dryruns )]
323
+
324
+ arg_types = self .argument_types (method )
325
+
326
+ prefix : List [bytes ] = []
327
+ if self .handle_selector and method :
328
+ prefix = [self .method_selector (method )]
329
+
330
+ modify_selector = False
331
+ if (action := self .abi_args_mod ) in (
332
+ ABIArgsMod .selector_byte_delete ,
333
+ ABIArgsMod .selector_byte_insert ,
334
+ ABIArgsMod .selector_byte_replace ,
335
+ ):
336
+ assert (
337
+ prefix
338
+ ), f"{ self .abi_args_mod = } which means we need to modify the selector, but we don't have one available to modify"
339
+ modify_selector = True
340
+
341
+ def selector_mod (prefix ):
342
+ """
343
+ modifies the selector by mutating a random byte (when modify_selector == True)
344
+ ^^^
345
+ """
346
+ assert isinstance (prefix , list ) and len (prefix ) <= 1
347
+ if not (prefix and modify_selector ):
348
+ return prefix
349
+
350
+ selector = prefix [0 ]
351
+ idx = random .randint (0 , 4 )
352
+ x , y = selector [:idx ], selector [idx :]
353
+ if action == ABIArgsMod .selector_byte_insert :
354
+ selector = x + random .randbytes (1 ) + y
355
+ elif action == ABIArgsMod .selector_byte_delete :
356
+ selector = (x [:- 1 ] + y ) if x else y [:- 1 ]
357
+ else :
358
+ assert (
359
+ action == ABIArgsMod .selector_byte_replace
360
+ ), f"expected action={ ABIArgsMod .selector_byte_replace } but got [{ action } ]"
361
+ idx = random .randint (0 , 3 )
362
+ selector = (
363
+ selector [:idx ]
364
+ + bytes ([(selector [idx ] + 1 ) % 256 ])
365
+ + selector [idx + 1 :]
366
+ )
367
+ return [selector ]
368
+
369
+ def args_mod (args ):
370
+ """
371
+ modifies the args by appending or deleting a random value (for appropriate `action`)
372
+ ^^^
373
+ """
374
+ if action not in (ABIArgsMod .parameter_append , ABIArgsMod .parameter_delete ):
375
+ return args
376
+
377
+ if action == ABIArgsMod .parameter_delete :
378
+ return args if not args else tuple (args [:- 1 ])
379
+
380
+ assert action == ABIArgsMod .parameter_append
381
+ return args + (self .generate_value (self .append_args_type ),)
382
+
383
+ def gen_args ():
384
+ # TODO: when incorporating hypothesis strategies, we'll need a more holistic
385
+ # approach that looks at relationships amongst various args
386
+ args = tuple (
387
+ selector_mod (prefix )
388
+ + [self .generate_value (atype ) for atype in arg_types ]
389
+ )
390
+ return args_mod (args )
391
+
392
+ return [gen_args () for _ in range (self .num_dryruns )]
393
+
394
+
395
+ class RandomArgLengthCallStrategy (CallStrategy ):
396
+ """
397
+ Generate a random number or arguments using the single
398
+ argument_strategy provided.
399
+ """
400
+
401
+ def __init__ (
402
+ self ,
403
+ argument_strategy : Type [ABIStrategy ],
404
+ max_args : int ,
405
+ * ,
406
+ num_dryruns : int = 1 ,
407
+ min_args : int = 0 ,
408
+ type_for_args : abi .ABIType = abi .ABIType .from_string ("byte[8]" ),
409
+ ):
410
+ super ().__init__ (argument_strategy , num_dryruns = num_dryruns )
411
+ self .max_args : int = max_args
412
+ self .min_args : int = min_args
413
+ self .type_for_args : abi .ABIType = type_for_args
414
+
415
+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
416
+ assert (
417
+ method is None
418
+ ), f"actual non-None method={ method } not supported for RandomArgLengthCallStrategy"
419
+ assert (
420
+ self .argument_strategy
421
+ ), "cannot generate inputs without an argument_strategy"
422
+
423
+ def gen_args ():
424
+ num_args = random .randint (self .min_args , self .max_args )
425
+ abi_args = [
426
+ self .generate_value (self .type_for_args ) for _ in range (num_args )
427
+ ]
428
+ # because cannot provide a method signature to include the arg types,
429
+ # we need to convert back to raw bytes
430
+ return tuple (bytes (arg ) for arg in abi_args )
431
+
432
+ return [gen_args () for _ in range (self .num_dryruns )]
0 commit comments