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

     1  // Copyright 2022 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  	"math/rand"
     8  	"sort"
     9  	"testing"
    10  
    11  	"github.com/google/syzkaller/pkg/image"
    12  	"github.com/google/syzkaller/pkg/testutil"
    13  )
    14  
    15  func TestGenericHeatmap(t *testing.T) {
    16  	t.Parallel()
    17  	// A test case is some data with the regions the heatmap is permitted to choose.
    18  	testData := []struct {
    19  		data    []byte
    20  		regions []region
    21  	}{
    22  		{
    23  			// Normal usage test.
    24  			[]byte(
    25  				"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    26  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    27  					"4eHh4eHh4eHh4eHh4Q5GTbHh4eHh4eHh4eHh4eHhcOHh4eHh4eHh4eHh4eHh4eHh4eHh4eEfNuHh4XPh" +
    28  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHjd+GRzcLh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    29  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4dpiSwpoReHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    30  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4bGfM+Hh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    31  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    32  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    33  					"4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    34  					"mpNKOZnS4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh4eHh" +
    35  					"4eHh4eHh4eHh4eHh4eHh4Q=="),
    36  			[]region{{128, 384}, {512, 576}},
    37  		},
    38  		{
    39  			// Test all constant bytes, i.e. falling back to uniform selection.
    40  			[]byte(
    41  				"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" +
    42  					"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" +
    43  					"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" +
    44  					"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" +
    45  					"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" +
    46  					"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"),
    47  			[]region{{0, 324}}, // Anywhere in the data.
    48  		},
    49  	}
    50  
    51  	const tries = 10
    52  	iters := testutil.IterCount() / tries
    53  
    54  	r := rand.New(testutil.RandSource(t))
    55  	for _, test := range testData {
    56  		data, err := image.DecodeB64(test.data)
    57  		if err != nil {
    58  			t.Fatalf("bad decode: %v", err)
    59  		}
    60  		for i := 0; i < iters; i++ {
    61  			hm := MakeGenericHeatmap(data, r).(*GenericHeatmap)
    62  
    63  			for j := 0; j < tries; j++ {
    64  				index := hm.ChooseLocation()
    65  				if !checkIndex(index, len(data), test.regions) {
    66  					hm.debugPrint(t, data, test.regions)
    67  					t.Fatalf("selected index %d does not fall in a region", index)
    68  				}
    69  			}
    70  		}
    71  	}
    72  }
    73  
    74  // Check an index is within some regions.
    75  func checkIndex(index, maxIndex int, regions []region) bool {
    76  	if index < 0 || index >= maxIndex {
    77  		return false
    78  	}
    79  
    80  	for _, region := range regions {
    81  		if region.start <= index && index < region.end {
    82  			return true
    83  		}
    84  	}
    85  	return false
    86  }
    87  
    88  type region struct {
    89  	start int
    90  	end   int
    91  }
    92  
    93  func (hm *GenericHeatmap) debugPrint(t *testing.T, data []byte, regions []region) {
    94  	// Print data.
    95  	t.Logf("data: len = %d", len(data))
    96  	for j := 0; j < len(data); j += granularity {
    97  		end := j + granularity
    98  		if end > len(data) {
    99  			end = len(data)
   100  		}
   101  		t.Logf("%8d: %x", j*granularity, data[j:end])
   102  	}
   103  	t.Log("\n")
   104  
   105  	// Print selected regions in data.
   106  	sort.Slice(regions, func(i, j int) bool {
   107  		return regions[i].start < regions[j].start
   108  	})
   109  	for j, region := range regions {
   110  		t.Logf("region  %4d: %8v - %8v", j, region.start, region.end)
   111  	}
   112  	t.Log("\n")
   113  
   114  	// Print heatmap.
   115  	t.Logf("generic heatmap (total segment length %d)", hm.length)
   116  	for j, seg := range hm.segments {
   117  		t.Logf("segment %4d: %8v - %8v", j, seg.offset, seg.offset+seg.length)
   118  	}
   119  	t.Log("\n\n\n")
   120  }