github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/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{upstream: writer, rawConn: rawConn}, true
    37  		}
    38  	case syscall.RawConn:
    39  		return &SyscallVectorisedWriter{upstream: writer, rawConn: 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{upstream: writer, rawConn: rawConn}, true
    52  		}
    53  	case syscall.RawConn:
    54  		return &SyscallVectorisedPacketWriter{upstream: writer, rawConn: 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.NewSize(bufferLen)
    78  		defer buffer.Release()
    79  		bufferBytes = buffer.FreeBytes()
    80  	}
    81  	buf.CopyMulti(bufferBytes, buffers)
    82  	return common.Error(w.upstream.Write(bufferBytes))
    83  }
    84  
    85  func (w *BufferedVectorisedWriter) Upstream() any {
    86  	return w.upstream
    87  }
    88  
    89  var _ N.VectorisedWriter = (*NetVectorisedWriterWrapper)(nil)
    90  
    91  type NetVectorisedWriterWrapper struct {
    92  	upstream io.Writer
    93  }
    94  
    95  func (w *NetVectorisedWriterWrapper) WriteVectorised(buffers []*buf.Buffer) error {
    96  	defer buf.ReleaseMulti(buffers)
    97  	netBuffers := net.Buffers(buf.ToSliceMulti(buffers))
    98  	return common.Error(netBuffers.WriteTo(w.upstream))
    99  }
   100  
   101  func (w *NetVectorisedWriterWrapper) Upstream() any {
   102  	return w.upstream
   103  }
   104  
   105  func (w *NetVectorisedWriterWrapper) WriterReplaceable() bool {
   106  	return true
   107  }
   108  
   109  var _ N.VectorisedWriter = (*SyscallVectorisedWriter)(nil)
   110  
   111  type SyscallVectorisedWriter struct {
   112  	upstream any
   113  	rawConn  syscall.RawConn
   114  	syscallVectorisedWriterFields
   115  }
   116  
   117  func (w *SyscallVectorisedWriter) Upstream() any {
   118  	return w.upstream
   119  }
   120  
   121  func (w *SyscallVectorisedWriter) WriterReplaceable() bool {
   122  	return true
   123  }
   124  
   125  var _ N.VectorisedPacketWriter = (*SyscallVectorisedPacketWriter)(nil)
   126  
   127  type SyscallVectorisedPacketWriter struct {
   128  	upstream any
   129  	rawConn  syscall.RawConn
   130  	syscallVectorisedWriterFields
   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  }