github.com/google/syzkaller@v0.0.0-20240517125934-c0f1611a36d6/pkg/flatrpc/conn_test.go (about)

     1  // Copyright 2024 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 flatrpc
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  func TestConn(t *testing.T) {
    13  	connectReq := &ConnectRequest{
    14  		Name:        "foo",
    15  		Arch:        "arch",
    16  		GitRevision: "rev1",
    17  		SyzRevision: "rev2",
    18  	}
    19  	connectReply := &ConnectReply{
    20  		LeakFrames: []string{"foo", "bar"},
    21  		RaceFrames: []string{"bar", "baz"},
    22  		Features:   FeatureCoverage | FeatureLeak,
    23  		Files:      []string{"file1"},
    24  		Globs:      []string{"glob1"},
    25  	}
    26  	executorMsg := &ExecutorMessage{
    27  		Msg: &ExecutorMessages{
    28  			Type: ExecutorMessagesRawExecuting,
    29  			Value: &ExecutingMessage{
    30  				Id:     1,
    31  				ProcId: 2,
    32  				Try:    3,
    33  			},
    34  		},
    35  	}
    36  
    37  	done := make(chan bool)
    38  	defer func() {
    39  		<-done
    40  	}()
    41  	serv, err := ListenAndServe(":0", func(c *Conn) {
    42  		defer close(done)
    43  		connectReqGot, err := Recv[ConnectRequestRaw](c)
    44  		if err != nil {
    45  			t.Fatal(err)
    46  		}
    47  		assert.Equal(t, connectReq, connectReqGot.UnPack())
    48  
    49  		if err := Send(c, connectReply); err != nil {
    50  			t.Fatal(err)
    51  		}
    52  
    53  		for i := 0; i < 10; i++ {
    54  			got, err := Recv[ExecutorMessageRaw](c)
    55  			if err != nil {
    56  				t.Fatal(err)
    57  			}
    58  			assert.Equal(t, executorMsg, got.UnPack())
    59  		}
    60  	})
    61  	if err != nil {
    62  		t.Fatal(err)
    63  	}
    64  	defer serv.Close()
    65  
    66  	c, err := Dial(serv.Addr.String(), 1)
    67  	if err != nil {
    68  		t.Fatal(err)
    69  	}
    70  	defer c.Close()
    71  
    72  	if err := Send(c, connectReq); err != nil {
    73  		t.Fatal(err)
    74  	}
    75  
    76  	connectReplyGot, err := Recv[ConnectReplyRaw](c)
    77  	if err != nil {
    78  		t.Fatal(err)
    79  	}
    80  	assert.Equal(t, connectReply, connectReplyGot.UnPack())
    81  
    82  	for i := 0; i < 10; i++ {
    83  		if err := Send(c, executorMsg); err != nil {
    84  			t.Fatal(err)
    85  		}
    86  	}
    87  }
    88  
    89  func BenchmarkConn(b *testing.B) {
    90  	connectReq := &ConnectRequest{
    91  		Name:        "foo",
    92  		Arch:        "arch",
    93  		GitRevision: "rev1",
    94  		SyzRevision: "rev2",
    95  	}
    96  	connectReply := &ConnectReply{
    97  		LeakFrames: []string{"foo", "bar"},
    98  		RaceFrames: []string{"bar", "baz"},
    99  		Features:   FeatureCoverage | FeatureLeak,
   100  		Files:      []string{"file1"},
   101  		Globs:      []string{"glob1"},
   102  	}
   103  
   104  	done := make(chan bool)
   105  	defer func() {
   106  		<-done
   107  	}()
   108  	serv, err := ListenAndServe(":0", func(c *Conn) {
   109  		defer close(done)
   110  		for i := 0; i < b.N; i++ {
   111  			_, err := Recv[ConnectRequestRaw](c)
   112  			if err != nil {
   113  				b.Fatal(err)
   114  			}
   115  			if err := Send(c, connectReply); err != nil {
   116  				b.Fatal(err)
   117  			}
   118  		}
   119  	})
   120  	if err != nil {
   121  		b.Fatal(err)
   122  	}
   123  	defer serv.Close()
   124  
   125  	c, err := Dial(serv.Addr.String(), 1)
   126  	if err != nil {
   127  		b.Fatal(err)
   128  	}
   129  	defer c.Close()
   130  
   131  	b.ReportAllocs()
   132  	b.ResetTimer()
   133  	for i := 0; i < b.N; i++ {
   134  		if err := Send(c, connectReq); err != nil {
   135  			b.Fatal(err)
   136  		}
   137  		_, err := Recv[ConnectReplyRaw](c)
   138  		if err != nil {
   139  			b.Fatal(err)
   140  		}
   141  	}
   142  }