github.com/google/syzkaller@v0.0.0-20240517125934-c0f1611a36d6/prog/alloc.go (about)

     1  // Copyright 2018 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package prog
     5  
     6  import (
     7  	"fmt"
     8  )
     9  
    10  // memAlloc keeps track of allocated objects in a program
    11  // and decides where to allocate new objects.
    12  // It has 2 main methods: noteAlloc which is called for existing allocations
    13  // in a program as we analyze it; and alloc which decides where to allocate
    14  // a new object.
    15  // The implementation is based on a 2-level bitmap where each bit represents
    16  // 64 bytes (memAllocGranule) of program memory.
    17  type memAlloc struct {
    18  	size uint64
    19  	mem  [memAllocL1Size]*[memAllocL0Size]uint64
    20  	buf  [memAllocL0Size]uint64
    21  }
    22  
    23  const (
    24  	memAllocGranule = 64 // 1 bit per that many bytes (all allocations are rounded to this size)
    25  	memAllocMaxMem  = 16 << 20
    26  	memAllocL0Size  = 64
    27  	bitsPerUint64   = 8 * 8
    28  	memAllocL0Mem   = memAllocL0Size * memAllocGranule * bitsPerUint64
    29  	memAllocL1Size  = memAllocMaxMem / memAllocL0Mem
    30  )
    31  
    32  func newMemAlloc(totalMemSize uint64) *memAlloc {
    33  	if totalMemSize > memAllocMaxMem {
    34  		panic(fmt.Sprintf("newMemAlloc: too much mem %v (max: %v)", totalMemSize, memAllocMaxMem))
    35  	}
    36  	if totalMemSize%memAllocL0Mem != 0 {
    37  		panic(fmt.Sprintf("newMemAlloc: unaligned size %v (align: %v)", totalMemSize, memAllocL0Mem))
    38  	}
    39  	ma := &memAlloc{
    40  		size: totalMemSize / memAllocGranule,
    41  	}
    42  	ma.mem[0] = &ma.buf
    43  	return ma
    44  }
    45  
    46  func (ma *memAlloc) noteAlloc(addr0, size0 uint64) {
    47  	addr := addr0 / memAllocGranule
    48  	size := (addr0+size0+memAllocGranule-1)/memAllocGranule - addr
    49  	for i := uint64(0); i < size; i++ {
    50  		ma.set(addr + i)
    51  	}
    52  }
    53  
    54  // alloc returns the next free address of size0 with respect to the given alignment.
    55  func (ma *memAlloc) alloc(r *randGen, size0, alignment0 uint64) uint64 {
    56  	if size0 == 0 {
    57  		size0 = 1
    58  	}
    59  	if alignment0 == 0 {
    60  		alignment0 = 1
    61  	}
    62  	size := (size0 + memAllocGranule - 1) / memAllocGranule
    63  	alignment := (alignment0 + memAllocGranule - 1) / memAllocGranule
    64  	end := ma.size - size
    65  	for start := uint64(0); start <= end; start += alignment {
    66  		empty := true
    67  		for i := uint64(0); i < size; i++ {
    68  			if ma.get(start + i) {
    69  				empty = false
    70  				break
    71  			}
    72  		}
    73  		if empty {
    74  			start0 := start * memAllocGranule
    75  			ma.noteAlloc(start0, size0)
    76  			return start0
    77  		}
    78  	}
    79  	ma.bankruptcy()
    80  	return ma.alloc(r, size0, alignment0)
    81  }
    82  
    83  func (ma *memAlloc) bankruptcy() {
    84  	for i1 := uint64(0); i1 < ma.size/(memAllocL0Size*bitsPerUint64); i1++ {
    85  		if ma.mem[i1] == nil {
    86  			continue
    87  		}
    88  		for i0 := range ma.mem[i1] {
    89  			ma.mem[i1][i0] = 0
    90  		}
    91  	}
    92  }
    93  
    94  func (ma *memAlloc) pos(idx uint64) (i1, i0, bit uint64) {
    95  	i1 = idx / (memAllocL0Size * bitsPerUint64)
    96  	r1 := idx % (memAllocL0Size * bitsPerUint64)
    97  	i0 = r1 / bitsPerUint64
    98  	bit = 1 << (r1 % bitsPerUint64)
    99  	return
   100  }
   101  
   102  func (ma *memAlloc) set(idx uint64) {
   103  	i1, i0, bit := ma.pos(idx)
   104  	if ma.mem[i1] == nil {
   105  		ma.mem[i1] = new([memAllocL0Size]uint64)
   106  	}
   107  	ma.mem[i1][i0] |= bit
   108  }
   109  
   110  func (ma *memAlloc) get(idx uint64) bool {
   111  	i1, i0, bit := ma.pos(idx)
   112  	if ma.mem[i1] == nil {
   113  		return false
   114  	}
   115  	return ma.mem[i1][i0]&bit != 0
   116  }
   117  
   118  type vmaAlloc struct {
   119  	numPages uint64
   120  	used     []uint64
   121  	m        map[uint64]struct{}
   122  }
   123  
   124  func newVmaAlloc(totalPages uint64) *vmaAlloc {
   125  	return &vmaAlloc{
   126  		numPages: totalPages,
   127  		m:        make(map[uint64]struct{}),
   128  	}
   129  }
   130  
   131  func (va *vmaAlloc) noteAlloc(page, size uint64) {
   132  	for i := page; i < page+size; i++ {
   133  		if _, ok := va.m[i]; ok {
   134  			continue
   135  		}
   136  		va.m[i] = struct{}{}
   137  		va.used = append(va.used, i)
   138  	}
   139  }
   140  
   141  func (va *vmaAlloc) alloc(r *randGen, size uint64) uint64 {
   142  	if size > va.numPages {
   143  		panic(fmt.Sprintf("vmaAlloc: bad size=%v numPages=%v", size, va.numPages))
   144  	}
   145  	var page uint64
   146  	if len(va.used) == 0 || r.oneOf(5) {
   147  		page = r.rand(4)
   148  		if !r.oneOf(100) {
   149  			page = va.numPages - page - size
   150  		}
   151  	} else {
   152  		page = va.used[r.rand(len(va.used))]
   153  		if size > 1 && r.bin() {
   154  			off := r.rand(int(size))
   155  			if off > page {
   156  				off = page
   157  			}
   158  			page -= off
   159  		}
   160  		if page+size > va.numPages {
   161  			page = va.numPages - size
   162  		}
   163  	}
   164  	if page >= va.numPages || size > va.numPages || page+size > va.numPages {
   165  		panic(fmt.Sprintf("vmaAlloc: bad page=%v size=%v numPages=%v", page, size, va.numPages))
   166  	}
   167  	va.noteAlloc(page, size)
   168  	return page
   169  }