github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/grpc/grpc.go (about)

     1  package grpc
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"net"
     7  	"sync"
     8  	"time"
     9  
    10  	"google.golang.org/protobuf/types/known/wrapperspb"
    11  )
    12  
    13  //go:generate protoc --go-grpc_out=. --go-grpc_opt=paths=source_relative chunk.proto
    14  
    15  type stream_conn interface {
    16  	Send(*wrapperspb.BytesValue) error
    17  	Recv() (*wrapperspb.BytesValue, error)
    18  }
    19  
    20  var _ net.Conn = (*conn)(nil)
    21  
    22  type conn struct {
    23  	raw stream_conn
    24  
    25  	buf  *bytes.Reader
    26  	rmux sync.Mutex
    27  
    28  	raddr net.Addr
    29  	laddr net.Addr
    30  
    31  	mu     sync.Mutex
    32  	closed bool
    33  	close  context.CancelFunc
    34  
    35  	deadline *time.Timer
    36  }
    37  
    38  func (c *conn) Read(b []byte) (int, error) {
    39  	c.rmux.Lock()
    40  	defer c.rmux.Unlock()
    41  
    42  	if c.buf != nil && c.buf.Len() > 0 {
    43  		return c.buf.Read(b)
    44  	}
    45  
    46  	data, err := c.raw.Recv()
    47  	if err != nil {
    48  		return 0, err
    49  	}
    50  
    51  	if c.buf == nil {
    52  		c.buf = bytes.NewReader(data.Value)
    53  	} else {
    54  		c.buf.Reset(data.Value)
    55  	}
    56  
    57  	return c.buf.Read(b)
    58  }
    59  
    60  func (c *conn) Write(b []byte) (int, error) {
    61  	if err := c.raw.Send(&wrapperspb.BytesValue{Value: b}); err != nil {
    62  		return 0, err
    63  	}
    64  
    65  	return len(b), nil
    66  }
    67  
    68  func (c *conn) Close() error {
    69  	c.mu.Lock()
    70  	defer c.mu.Unlock()
    71  
    72  	if c.closed {
    73  		return nil
    74  	}
    75  
    76  	c.close()
    77  	c.closed = true
    78  	return nil
    79  }
    80  
    81  func (c *conn) LocalAddr() net.Addr  { return c.laddr }
    82  func (c *conn) RemoteAddr() net.Addr { return c.raddr }
    83  
    84  func (c *conn) SetDeadline(t time.Time) error {
    85  	if c.deadline == nil {
    86  		if !t.IsZero() {
    87  			c.deadline = time.AfterFunc(time.Until(t), func() { c.Close() })
    88  		}
    89  		return nil
    90  	}
    91  
    92  	if t.IsZero() {
    93  		c.deadline.Stop()
    94  	} else {
    95  		c.deadline.Reset(time.Until(t))
    96  	}
    97  
    98  	return nil
    99  }
   100  
   101  func (c *conn) SetReadDeadline(t time.Time) error  { return c.SetDeadline(t) }
   102  func (c *conn) SetWriteDeadline(t time.Time) error { return c.SetDeadline(t) }