github.com/moqsien/xraycore@v1.8.5/transport/internet/grpc/encoding/hunkconn.go (about)

     1  package encoding
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/moqsien/xraycore/common/buf"
     9  	xnet "github.com/moqsien/xraycore/common/net"
    10  	"github.com/moqsien/xraycore/common/net/cnc"
    11  	"github.com/moqsien/xraycore/common/signal/done"
    12  	"google.golang.org/grpc/metadata"
    13  	"google.golang.org/grpc/peer"
    14  )
    15  
    16  type HunkConn interface {
    17  	Context() context.Context
    18  	Send(*Hunk) error
    19  	Recv() (*Hunk, error)
    20  	SendMsg(m interface{}) error
    21  	RecvMsg(m interface{}) error
    22  }
    23  
    24  type StreamCloser interface {
    25  	CloseSend() error
    26  }
    27  
    28  type HunkReaderWriter struct {
    29  	hc     HunkConn
    30  	cancel context.CancelFunc
    31  	done   *done.Instance
    32  
    33  	buf   []byte
    34  	index int
    35  }
    36  
    37  func NewHunkReadWriter(hc HunkConn, cancel context.CancelFunc) *HunkReaderWriter {
    38  	return &HunkReaderWriter{hc, cancel, done.New(), nil, 0}
    39  }
    40  
    41  func NewHunkConn(hc HunkConn, cancel context.CancelFunc) net.Conn {
    42  	var rAddr net.Addr
    43  	pr, ok := peer.FromContext(hc.Context())
    44  	if ok {
    45  		rAddr = pr.Addr
    46  	} else {
    47  		rAddr = &net.TCPAddr{
    48  			IP:   []byte{0, 0, 0, 0},
    49  			Port: 0,
    50  		}
    51  	}
    52  
    53  	md, ok := metadata.FromIncomingContext(hc.Context())
    54  	if ok {
    55  		header := md.Get("x-real-ip")
    56  		if len(header) > 0 {
    57  			realip := xnet.ParseAddress(header[0])
    58  			if realip.Family().IsIP() {
    59  				rAddr = &net.TCPAddr{
    60  					IP:   realip.IP(),
    61  					Port: 0,
    62  				}
    63  			}
    64  		}
    65  	}
    66  	wrc := NewHunkReadWriter(hc, cancel)
    67  	return cnc.NewConnection(
    68  		cnc.ConnectionInput(wrc),
    69  		cnc.ConnectionOutput(wrc),
    70  		cnc.ConnectionOnClose(wrc),
    71  		cnc.ConnectionRemoteAddr(rAddr),
    72  	)
    73  }
    74  
    75  func (h *HunkReaderWriter) forceFetch() error {
    76  	hunk, err := h.hc.Recv()
    77  	if err != nil {
    78  		if err == io.EOF {
    79  			return err
    80  		}
    81  
    82  		return newError("failed to fetch hunk from gRPC tunnel").Base(err)
    83  	}
    84  
    85  	h.buf = hunk.Data
    86  	h.index = 0
    87  
    88  	return nil
    89  }
    90  
    91  func (h *HunkReaderWriter) Read(buf []byte) (int, error) {
    92  	if h.done.Done() {
    93  		return 0, io.EOF
    94  	}
    95  
    96  	if h.index >= len(h.buf) {
    97  		if err := h.forceFetch(); err != nil {
    98  			return 0, err
    99  		}
   100  	}
   101  	n := copy(buf, h.buf[h.index:])
   102  	h.index += n
   103  
   104  	return n, nil
   105  }
   106  
   107  func (h *HunkReaderWriter) ReadMultiBuffer() (buf.MultiBuffer, error) {
   108  	if h.done.Done() {
   109  		return nil, io.EOF
   110  	}
   111  	if h.index >= len(h.buf) {
   112  		if err := h.forceFetch(); err != nil {
   113  			return nil, err
   114  		}
   115  	}
   116  
   117  	if cap(h.buf) >= buf.Size {
   118  		b := h.buf
   119  		h.index = len(h.buf)
   120  		return buf.MultiBuffer{buf.NewExisted(b)}, nil
   121  	}
   122  
   123  	b := buf.New()
   124  	_, err := b.ReadFrom(h)
   125  	if err != nil {
   126  		return nil, err
   127  	}
   128  	return buf.MultiBuffer{b}, nil
   129  }
   130  
   131  func (h *HunkReaderWriter) Write(buf []byte) (int, error) {
   132  	if h.done.Done() {
   133  		return 0, io.ErrClosedPipe
   134  	}
   135  
   136  	err := h.hc.Send(&Hunk{Data: buf[:]})
   137  	if err != nil {
   138  		return 0, newError("failed to send data over gRPC tunnel").Base(err)
   139  	}
   140  	return len(buf), nil
   141  }
   142  
   143  func (h *HunkReaderWriter) Close() error {
   144  	if h.cancel != nil {
   145  		h.cancel()
   146  	}
   147  	if sc, match := h.hc.(StreamCloser); match {
   148  		return sc.CloseSend()
   149  	}
   150  
   151  	return h.done.Close()
   152  }