github.com/eagleql/xray-core@v1.4.4/transport/internet/grpc/encoding/multiconn.go (about)

     1  package encoding
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"google.golang.org/grpc/peer"
     9  
    10  	"github.com/eagleql/xray-core/common/buf"
    11  	"github.com/eagleql/xray-core/common/net/cnc"
    12  	"github.com/eagleql/xray-core/common/signal/done"
    13  )
    14  
    15  type MultiHunkConn interface {
    16  	Context() context.Context
    17  	Send(*MultiHunk) error
    18  	Recv() (*MultiHunk, error)
    19  	SendMsg(m interface{}) error
    20  	RecvMsg(m interface{}) error
    21  }
    22  
    23  type MultiHunkReaderWriter struct {
    24  	hc     MultiHunkConn
    25  	cancel context.CancelFunc
    26  	done   *done.Instance
    27  
    28  	buf [][]byte
    29  }
    30  
    31  func NewMultiHunkReadWriter(hc MultiHunkConn, cancel context.CancelFunc) *MultiHunkReaderWriter {
    32  	return &MultiHunkReaderWriter{hc, cancel, done.New(), nil}
    33  }
    34  
    35  func NewMultiHunkConn(hc MultiHunkConn, cancel context.CancelFunc) net.Conn {
    36  	var rAddr net.Addr
    37  	pr, ok := peer.FromContext(hc.Context())
    38  	if ok {
    39  		rAddr = pr.Addr
    40  	} else {
    41  		rAddr = &net.TCPAddr{
    42  			IP:   []byte{0, 0, 0, 0},
    43  			Port: 0,
    44  		}
    45  	}
    46  
    47  	wrc := NewMultiHunkReadWriter(hc, cancel)
    48  	return cnc.NewConnection(
    49  		cnc.ConnectionInputMulti(wrc),
    50  		cnc.ConnectionOutputMulti(wrc),
    51  		cnc.ConnectionOnClose(wrc),
    52  		cnc.ConnectionRemoteAddr(rAddr),
    53  	)
    54  }
    55  
    56  func (h *MultiHunkReaderWriter) forceFetch() error {
    57  	hunk, err := h.hc.Recv()
    58  	if err != nil {
    59  		if err == io.EOF {
    60  			return err
    61  		}
    62  
    63  		return newError("failed to fetch hunk from gRPC tunnel").Base(err)
    64  	}
    65  
    66  	h.buf = hunk.Data
    67  
    68  	return nil
    69  }
    70  
    71  func (h *MultiHunkReaderWriter) ReadMultiBuffer() (buf.MultiBuffer, error) {
    72  	if h.done.Done() {
    73  		return nil, io.EOF
    74  	}
    75  
    76  	if err := h.forceFetch(); err != nil {
    77  		return nil, err
    78  	}
    79  
    80  	var mb = make(buf.MultiBuffer, 0, len(h.buf))
    81  	for _, b := range h.buf {
    82  		if len(b) == 0 {
    83  			continue
    84  		}
    85  
    86  		if cap(b) >= buf.Size {
    87  			mb = append(mb, buf.NewExisted(b))
    88  		} else {
    89  			nb := buf.New()
    90  			nb.Extend(int32(len(b)))
    91  			copy(nb.Bytes(), b)
    92  
    93  			mb = append(mb, nb)
    94  		}
    95  
    96  	}
    97  	return mb, nil
    98  }
    99  
   100  func (h *MultiHunkReaderWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
   101  	defer buf.ReleaseMulti(mb)
   102  	if h.done.Done() {
   103  		return io.ErrClosedPipe
   104  	}
   105  
   106  	hunks := make([][]byte, 0, len(mb))
   107  
   108  	for _, b := range mb {
   109  		if b.Len() > 0 {
   110  			hunks = append(hunks, b.Bytes())
   111  		}
   112  	}
   113  
   114  	err := h.hc.Send(&MultiHunk{Data: hunks})
   115  	if err != nil {
   116  		return err
   117  	}
   118  	return nil
   119  }
   120  
   121  func (h *MultiHunkReaderWriter) Close() error {
   122  	if h.cancel != nil {
   123  		h.cancel()
   124  	}
   125  	if sc, match := h.hc.(StreamCloser); match {
   126  		return sc.CloseSend()
   127  	}
   128  
   129  	return h.done.Close()
   130  }