github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/pkg/rpctype/rpc.go (about)

     1  // Copyright 2017 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 rpctype
     5  
     6  import (
     7  	"compress/flate"
     8  	"fmt"
     9  	"io"
    10  	"net"
    11  	"net/rpc"
    12  	"time"
    13  
    14  	"github.com/google/syzkaller/pkg/log"
    15  )
    16  
    17  type RPCServer struct {
    18  	ln net.Listener
    19  	s  *rpc.Server
    20  }
    21  
    22  func NewRPCServer(addr, name string, receiver interface{}) (*RPCServer, error) {
    23  	ln, err := net.Listen("tcp", addr)
    24  	if err != nil {
    25  		return nil, fmt.Errorf("failed to listen on %v: %w", addr, err)
    26  	}
    27  	s := rpc.NewServer()
    28  	if err := s.RegisterName(name, receiver); err != nil {
    29  		return nil, err
    30  	}
    31  	serv := &RPCServer{
    32  		ln: ln,
    33  		s:  s,
    34  	}
    35  	return serv, nil
    36  }
    37  
    38  func (serv *RPCServer) Serve() {
    39  	for {
    40  		conn, err := serv.ln.Accept()
    41  		if err != nil {
    42  			log.Logf(0, "failed to accept an rpc connection: %v", err)
    43  			continue
    44  		}
    45  		setupKeepAlive(conn, time.Minute)
    46  		go serv.s.ServeConn(newFlateConn(conn))
    47  	}
    48  }
    49  
    50  func (serv *RPCServer) Addr() net.Addr {
    51  	return serv.ln.Addr()
    52  }
    53  
    54  type RPCClient struct {
    55  	conn net.Conn
    56  	c    *rpc.Client
    57  }
    58  
    59  func NewRPCClient(addr string) (*RPCClient, error) {
    60  	conn, err := net.DialTimeout("tcp", addr, 3*time.Minute)
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  	setupKeepAlive(conn, time.Minute)
    65  	cli := &RPCClient{
    66  		conn: conn,
    67  		c:    rpc.NewClient(newFlateConn(conn)),
    68  	}
    69  	return cli, nil
    70  }
    71  
    72  func (cli *RPCClient) Call(method string, args, reply interface{}) error {
    73  	// Note: SetDeadline is not implemented on fuchsia, so don't fail on error.
    74  	cli.conn.SetDeadline(time.Now().Add(10 * time.Minute))
    75  	defer cli.conn.SetDeadline(time.Time{})
    76  	return cli.c.Call(method, args, reply)
    77  }
    78  
    79  func (cli *RPCClient) Close() {
    80  	cli.c.Close()
    81  }
    82  
    83  func setupKeepAlive(conn net.Conn, keepAlive time.Duration) {
    84  	conn.(*net.TCPConn).SetKeepAlive(true)
    85  	conn.(*net.TCPConn).SetKeepAlivePeriod(keepAlive)
    86  }
    87  
    88  // flateConn wraps net.Conn in flate.Reader/Writer for compressed traffic.
    89  type flateConn struct {
    90  	r io.ReadCloser
    91  	w *flate.Writer
    92  	c io.Closer
    93  }
    94  
    95  func newFlateConn(conn io.ReadWriteCloser) io.ReadWriteCloser {
    96  	w, err := flate.NewWriter(conn, 9)
    97  	if err != nil {
    98  		panic(err)
    99  	}
   100  	return &flateConn{
   101  		r: flate.NewReader(conn),
   102  		w: w,
   103  		c: conn,
   104  	}
   105  }
   106  
   107  func (fc *flateConn) Read(data []byte) (int, error) {
   108  	return fc.r.Read(data)
   109  }
   110  
   111  func (fc *flateConn) Write(data []byte) (int, error) {
   112  	n, err := fc.w.Write(data)
   113  	if err != nil {
   114  		return n, err
   115  	}
   116  	if err := fc.w.Flush(); err != nil {
   117  		return n, err
   118  	}
   119  	return n, nil
   120  }
   121  
   122  func (fc *flateConn) Close() error {
   123  	var err0 error
   124  	if err := fc.r.Close(); err != nil {
   125  		err0 = err
   126  	}
   127  	if err := fc.w.Close(); err != nil {
   128  		err0 = err
   129  	}
   130  	if err := fc.c.Close(); err != nil {
   131  		err0 = err
   132  	}
   133  	return err0
   134  }