github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/net/connection.go (about)

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