github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/bufio/conn.go (about)

     1  package bufio
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  
     7  	"github.com/sagernet/sing/common"
     8  	"github.com/sagernet/sing/common/buf"
     9  	M "github.com/sagernet/sing/common/metadata"
    10  	N "github.com/sagernet/sing/common/network"
    11  )
    12  
    13  func NewPacketConn(conn net.PacketConn) N.NetPacketConn {
    14  	if udpConn, isUDPConn := conn.(*net.UDPConn); isUDPConn {
    15  		return &ExtendedUDPConn{udpConn}
    16  	} else if packetConn, isPacketConn := conn.(N.NetPacketConn); isPacketConn && !forceSTDIO {
    17  		return packetConn
    18  	} else {
    19  		return &ExtendedPacketConn{conn}
    20  	}
    21  }
    22  
    23  type ExtendedUDPConn struct {
    24  	*net.UDPConn
    25  }
    26  
    27  func (w *ExtendedUDPConn) ReadPacket(buffer *buf.Buffer) (M.Socksaddr, error) {
    28  	n, addr, err := w.ReadFromUDPAddrPort(buffer.FreeBytes())
    29  	if err != nil {
    30  		return M.Socksaddr{}, err
    31  	}
    32  	buffer.Truncate(n)
    33  	return M.SocksaddrFromNetIP(addr).Unwrap(), nil
    34  }
    35  
    36  func (w *ExtendedUDPConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
    37  	defer buffer.Release()
    38  	if destination.IsFqdn() {
    39  		udpAddr, err := net.ResolveUDPAddr("udp", destination.String())
    40  		if err != nil {
    41  			return err
    42  		}
    43  		return common.Error(w.UDPConn.WriteTo(buffer.Bytes(), udpAddr))
    44  	}
    45  	return common.Error(w.UDPConn.WriteToUDP(buffer.Bytes(), destination.UDPAddr()))
    46  }
    47  
    48  func (w *ExtendedUDPConn) Upstream() any {
    49  	return w.UDPConn
    50  }
    51  
    52  type ExtendedPacketConn struct {
    53  	net.PacketConn
    54  }
    55  
    56  func (w *ExtendedPacketConn) ReadPacket(buffer *buf.Buffer) (M.Socksaddr, error) {
    57  	_, addr, err := buffer.ReadPacketFrom(w)
    58  	if err != nil {
    59  		return M.Socksaddr{}, err
    60  	}
    61  	return M.SocksaddrFromNet(addr).Unwrap(), err
    62  }
    63  
    64  func (w *ExtendedPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
    65  	defer buffer.Release()
    66  	return common.Error(w.WriteTo(buffer.Bytes(), destination.UDPAddr()))
    67  }
    68  
    69  func (w *ExtendedPacketConn) Upstream() any {
    70  	return w.PacketConn
    71  }
    72  
    73  type ExtendedReaderWrapper struct {
    74  	io.Reader
    75  }
    76  
    77  func (r *ExtendedReaderWrapper) ReadBuffer(buffer *buf.Buffer) error {
    78  	n, err := r.Read(buffer.FreeBytes())
    79  	buffer.Truncate(n)
    80  	if n > 0 && err == io.EOF {
    81  		return nil
    82  	}
    83  	return err
    84  }
    85  
    86  func (r *ExtendedReaderWrapper) WriteTo(w io.Writer) (n int64, err error) {
    87  	return Copy(w, r.Reader)
    88  }
    89  
    90  func (r *ExtendedReaderWrapper) Upstream() any {
    91  	return r.Reader
    92  }
    93  
    94  func (r *ExtendedReaderWrapper) ReaderReplaceable() bool {
    95  	return true
    96  }
    97  
    98  func NewExtendedReader(reader io.Reader) N.ExtendedReader {
    99  	if forceSTDIO {
   100  		if r, ok := reader.(*ExtendedReaderWrapper); ok {
   101  			return r
   102  		}
   103  	} else {
   104  		if r, ok := reader.(N.ExtendedReader); ok {
   105  			return r
   106  		}
   107  	}
   108  	return &ExtendedReaderWrapper{reader}
   109  }
   110  
   111  type ExtendedWriterWrapper struct {
   112  	io.Writer
   113  }
   114  
   115  func (w *ExtendedWriterWrapper) WriteBuffer(buffer *buf.Buffer) error {
   116  	defer buffer.Release()
   117  	return common.Error(w.Write(buffer.Bytes()))
   118  }
   119  
   120  func (w *ExtendedWriterWrapper) ReadFrom(r io.Reader) (n int64, err error) {
   121  	return Copy(w.Writer, r)
   122  }
   123  
   124  func (w *ExtendedWriterWrapper) Upstream() any {
   125  	return w.Writer
   126  }
   127  
   128  func (w *ExtendedWriterWrapper) WriterReplaceable() bool {
   129  	return true
   130  }
   131  
   132  func NewExtendedWriter(writer io.Writer) N.ExtendedWriter {
   133  	if forceSTDIO {
   134  		if w, ok := writer.(*ExtendedWriterWrapper); ok {
   135  			return w
   136  		}
   137  	} else {
   138  		if w, ok := writer.(N.ExtendedWriter); ok {
   139  			return w
   140  		}
   141  	}
   142  	return &ExtendedWriterWrapper{writer}
   143  }
   144  
   145  type ExtendedConnWrapper struct {
   146  	net.Conn
   147  	reader N.ExtendedReader
   148  	writer N.ExtendedWriter
   149  }
   150  
   151  func (w *ExtendedConnWrapper) ReadBuffer(buffer *buf.Buffer) error {
   152  	return w.reader.ReadBuffer(buffer)
   153  }
   154  
   155  func (w *ExtendedConnWrapper) WriteBuffer(buffer *buf.Buffer) error {
   156  	return w.writer.WriteBuffer(buffer)
   157  }
   158  
   159  func (w *ExtendedConnWrapper) ReadFrom(r io.Reader) (n int64, err error) {
   160  	return Copy(w.writer, r)
   161  }
   162  
   163  func (r *ExtendedConnWrapper) WriteTo(w io.Writer) (n int64, err error) {
   164  	return Copy(w, r.reader)
   165  }
   166  
   167  func (w *ExtendedConnWrapper) UpstreamReader() any {
   168  	return w.reader
   169  }
   170  
   171  func (w *ExtendedConnWrapper) ReaderReplaceable() bool {
   172  	return true
   173  }
   174  
   175  func (w *ExtendedConnWrapper) UpstreamWriter() any {
   176  	return w.writer
   177  }
   178  
   179  func (w *ExtendedConnWrapper) WriterReplaceable() bool {
   180  	return true
   181  }
   182  
   183  func (w *ExtendedConnWrapper) Upstream() any {
   184  	return w.Conn
   185  }
   186  
   187  func NewExtendedConn(conn net.Conn) N.ExtendedConn {
   188  	if c, ok := conn.(N.ExtendedConn); ok {
   189  		return c
   190  	}
   191  	return &ExtendedConnWrapper{
   192  		Conn:   conn,
   193  		reader: NewExtendedReader(conn),
   194  		writer: NewExtendedWriter(conn),
   195  	}
   196  }