github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/common/net/connection.go (about)

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