github.com/sagernet/sing@v0.2.6/common/bufio/vectorised.go (about)

     1  package bufio
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"syscall"
     7  
     8  	"github.com/sagernet/sing/common"
     9  	"github.com/sagernet/sing/common/buf"
    10  	M "github.com/sagernet/sing/common/metadata"
    11  	N "github.com/sagernet/sing/common/network"
    12  )
    13  
    14  func NewVectorisedWriter(writer io.Writer) N.VectorisedWriter {
    15  	if vectorisedWriter, ok := CreateVectorisedWriter(N.UnwrapWriter(writer)); ok {
    16  		return vectorisedWriter
    17  	}
    18  	return &BufferedVectorisedWriter{upstream: writer}
    19  }
    20  
    21  func CreateVectorisedWriter(writer any) (N.VectorisedWriter, bool) {
    22  	switch w := writer.(type) {
    23  	case N.VectorisedWriter:
    24  		return w, true
    25  	case *net.TCPConn:
    26  		return &NetVectorisedWriterWrapper{w}, true
    27  	case *net.UDPConn:
    28  		return &NetVectorisedWriterWrapper{w}, true
    29  	case *net.IPConn:
    30  		return &NetVectorisedWriterWrapper{w}, true
    31  	case *net.UnixConn:
    32  		return &NetVectorisedWriterWrapper{w}, true
    33  	case syscall.Conn:
    34  		rawConn, err := w.SyscallConn()
    35  		if err == nil {
    36  			return &SyscallVectorisedWriter{writer, rawConn}, true
    37  		}
    38  	case syscall.RawConn:
    39  		return &SyscallVectorisedWriter{writer, w}, true
    40  	}
    41  	return nil, false
    42  }
    43  
    44  func CreateVectorisedPacketWriter(writer any) (N.VectorisedPacketWriter, bool) {
    45  	switch w := writer.(type) {
    46  	case N.VectorisedPacketWriter:
    47  		return w, true
    48  	case syscall.Conn:
    49  		rawConn, err := w.SyscallConn()
    50  		if err == nil {
    51  			return &SyscallVectorisedPacketWriter{writer, rawConn}, true
    52  		}
    53  	case syscall.RawConn:
    54  		return &SyscallVectorisedPacketWriter{writer, w}, true
    55  	}
    56  	return nil, false
    57  }
    58  
    59  var _ N.VectorisedWriter = (*BufferedVectorisedWriter)(nil)
    60  
    61  type BufferedVectorisedWriter struct {
    62  	upstream io.Writer
    63  }
    64  
    65  func (w *BufferedVectorisedWriter) WriteVectorised(buffers []*buf.Buffer) error {
    66  	defer buf.ReleaseMulti(buffers)
    67  	bufferLen := buf.LenMulti(buffers)
    68  	if bufferLen == 0 {
    69  		return common.Error(w.upstream.Write(nil))
    70  	} else if len(buffers) == 1 {
    71  		return common.Error(w.upstream.Write(buffers[0].Bytes()))
    72  	}
    73  	var bufferBytes []byte
    74  	if bufferLen > 65535 {
    75  		bufferBytes = make([]byte, bufferLen)
    76  	} else {
    77  		_buffer := buf.StackNewSize(bufferLen)
    78  		defer common.KeepAlive(_buffer)
    79  		buffer := common.Dup(_buffer)
    80  		defer buffer.Release()
    81  		bufferBytes = buffer.FreeBytes()
    82  	}
    83  	buf.CopyMulti(bufferBytes, buffers)
    84  	return common.Error(w.upstream.Write(bufferBytes))
    85  }
    86  
    87  func (w *BufferedVectorisedWriter) Upstream() any {
    88  	return w.upstream
    89  }
    90  
    91  var _ N.VectorisedWriter = (*NetVectorisedWriterWrapper)(nil)
    92  
    93  type NetVectorisedWriterWrapper struct {
    94  	upstream io.Writer
    95  }
    96  
    97  func (w *NetVectorisedWriterWrapper) WriteVectorised(buffers []*buf.Buffer) error {
    98  	defer buf.ReleaseMulti(buffers)
    99  	netBuffers := net.Buffers(buf.ToSliceMulti(buffers))
   100  	return common.Error(netBuffers.WriteTo(w.upstream))
   101  }
   102  
   103  func (w *NetVectorisedWriterWrapper) Upstream() any {
   104  	return w.upstream
   105  }
   106  
   107  func (w *NetVectorisedWriterWrapper) WriterReplaceable() bool {
   108  	return true
   109  }
   110  
   111  var _ N.VectorisedWriter = (*SyscallVectorisedWriter)(nil)
   112  
   113  type SyscallVectorisedWriter struct {
   114  	upstream any
   115  	rawConn  syscall.RawConn
   116  }
   117  
   118  func (w *SyscallVectorisedWriter) Upstream() any {
   119  	return w.upstream
   120  }
   121  
   122  func (w *SyscallVectorisedWriter) WriterReplaceable() bool {
   123  	return true
   124  }
   125  
   126  var _ N.VectorisedPacketWriter = (*SyscallVectorisedPacketWriter)(nil)
   127  
   128  type SyscallVectorisedPacketWriter struct {
   129  	upstream any
   130  	rawConn  syscall.RawConn
   131  }
   132  
   133  func (w *SyscallVectorisedPacketWriter) Upstream() any {
   134  	return w.upstream
   135  }
   136  
   137  var _ N.VectorisedPacketWriter = (*UnbindVectorisedPacketWriter)(nil)
   138  
   139  type UnbindVectorisedPacketWriter struct {
   140  	N.VectorisedWriter
   141  }
   142  
   143  func (w *UnbindVectorisedPacketWriter) WriteVectorisedPacket(buffers []*buf.Buffer, _ M.Socksaddr) error {
   144  	return w.WriteVectorised(buffers)
   145  }