github.com/eagleql/xray-core@v1.4.4/common/net/cnc/connection.go (about)

     1  package cnc
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/eagleql/xray-core/common"
     8  	"github.com/eagleql/xray-core/common/buf"
     9  	"github.com/eagleql/xray-core/common/net"
    10  	"github.com/eagleql/xray-core/common/signal/done"
    11  )
    12  
    13  type ConnectionOption func(*connection)
    14  
    15  func ConnectionLocalAddr(a net.Addr) ConnectionOption {
    16  	return func(c *connection) {
    17  		c.local = a
    18  	}
    19  }
    20  
    21  func ConnectionRemoteAddr(a net.Addr) ConnectionOption {
    22  	return func(c *connection) {
    23  		c.remote = a
    24  	}
    25  }
    26  
    27  func ConnectionInput(writer io.Writer) ConnectionOption {
    28  	return func(c *connection) {
    29  		c.writer = buf.NewWriter(writer)
    30  	}
    31  }
    32  
    33  func ConnectionInputMulti(writer buf.Writer) ConnectionOption {
    34  	return func(c *connection) {
    35  		c.writer = writer
    36  	}
    37  }
    38  
    39  func ConnectionOutput(reader io.Reader) ConnectionOption {
    40  	return func(c *connection) {
    41  		c.reader = &buf.BufferedReader{Reader: buf.NewReader(reader)}
    42  	}
    43  }
    44  
    45  func ConnectionOutputMulti(reader buf.Reader) ConnectionOption {
    46  	return func(c *connection) {
    47  		c.reader = &buf.BufferedReader{Reader: reader}
    48  	}
    49  }
    50  
    51  func ConnectionOutputMultiUDP(reader buf.Reader) ConnectionOption {
    52  	return func(c *connection) {
    53  		c.reader = &buf.BufferedReader{
    54  			Reader:  reader,
    55  			Spliter: buf.SplitFirstBytes,
    56  		}
    57  	}
    58  }
    59  
    60  func ConnectionOnClose(n io.Closer) ConnectionOption {
    61  	return func(c *connection) {
    62  		c.onClose = n
    63  	}
    64  }
    65  
    66  func NewConnection(opts ...ConnectionOption) net.Conn {
    67  	c := &connection{
    68  		done: done.New(),
    69  		local: &net.TCPAddr{
    70  			IP:   []byte{0, 0, 0, 0},
    71  			Port: 0,
    72  		},
    73  		remote: &net.TCPAddr{
    74  			IP:   []byte{0, 0, 0, 0},
    75  			Port: 0,
    76  		},
    77  	}
    78  
    79  	for _, opt := range opts {
    80  		opt(c)
    81  	}
    82  
    83  	return c
    84  }
    85  
    86  type connection struct {
    87  	reader  *buf.BufferedReader
    88  	writer  buf.Writer
    89  	done    *done.Instance
    90  	onClose io.Closer
    91  	local   net.Addr
    92  	remote  net.Addr
    93  }
    94  
    95  func (c *connection) Read(b []byte) (int, error) {
    96  	return c.reader.Read(b)
    97  }
    98  
    99  // ReadMultiBuffer implements buf.Reader.
   100  func (c *connection) ReadMultiBuffer() (buf.MultiBuffer, error) {
   101  	return c.reader.ReadMultiBuffer()
   102  }
   103  
   104  // Write implements net.Conn.Write().
   105  func (c *connection) Write(b []byte) (int, error) {
   106  	if c.done.Done() {
   107  		return 0, io.ErrClosedPipe
   108  	}
   109  
   110  	l := len(b)
   111  	mb := make(buf.MultiBuffer, 0, l/buf.Size+1)
   112  	mb = buf.MergeBytes(mb, b)
   113  	return l, c.writer.WriteMultiBuffer(mb)
   114  }
   115  
   116  func (c *connection) WriteMultiBuffer(mb buf.MultiBuffer) error {
   117  	if c.done.Done() {
   118  		buf.ReleaseMulti(mb)
   119  		return io.ErrClosedPipe
   120  	}
   121  
   122  	return c.writer.WriteMultiBuffer(mb)
   123  }
   124  
   125  // Close implements net.Conn.Close().
   126  func (c *connection) Close() error {
   127  	common.Must(c.done.Close())
   128  	common.Interrupt(c.reader)
   129  	common.Close(c.writer)
   130  	if c.onClose != nil {
   131  		return c.onClose.Close()
   132  	}
   133  
   134  	return nil
   135  }
   136  
   137  // LocalAddr implements net.Conn.LocalAddr().
   138  func (c *connection) LocalAddr() net.Addr {
   139  	return c.local
   140  }
   141  
   142  // RemoteAddr implements net.Conn.RemoteAddr().
   143  func (c *connection) RemoteAddr() net.Addr {
   144  	return c.remote
   145  }
   146  
   147  // SetDeadline implements net.Conn.SetDeadline().
   148  func (c *connection) SetDeadline(t time.Time) error {
   149  	return nil
   150  }
   151  
   152  // SetReadDeadline implements net.Conn.SetReadDeadline().
   153  func (c *connection) SetReadDeadline(t time.Time) error {
   154  	return nil
   155  }
   156  
   157  // SetWriteDeadline implements net.Conn.SetWriteDeadline().
   158  func (c *connection) SetWriteDeadline(t time.Time) error {
   159  	return nil
   160  }