github.com/v2fly/v2ray-core/v4@v4.45.2/common/net/connection.go (about)

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