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

     1  package bufio
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"time"
     7  
     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  type CachedConn struct {
    14  	net.Conn
    15  	buffer *buf.Buffer
    16  }
    17  
    18  func NewCachedConn(conn net.Conn, buffer *buf.Buffer) *CachedConn {
    19  	buffer.IncRef()
    20  	return &CachedConn{
    21  		Conn:   conn,
    22  		buffer: buffer,
    23  	}
    24  }
    25  
    26  func (c *CachedConn) ReadCached() *buf.Buffer {
    27  	buffer := c.buffer
    28  	c.buffer = nil
    29  	if buffer != nil {
    30  		buffer.DecRef()
    31  	}
    32  	return buffer
    33  }
    34  
    35  func (c *CachedConn) Read(p []byte) (n int, err error) {
    36  	if c.buffer != nil {
    37  		n, err = c.buffer.Read(p)
    38  		if err == nil {
    39  			return
    40  		}
    41  		c.buffer.DecRef()
    42  		c.buffer.Release()
    43  		c.buffer = nil
    44  	}
    45  	return c.Conn.Read(p)
    46  }
    47  
    48  func (c *CachedConn) WriteTo(w io.Writer) (n int64, err error) {
    49  	if c.buffer != nil {
    50  		wn, wErr := w.Write(c.buffer.Bytes())
    51  		if wErr != nil {
    52  			c.buffer.DecRef()
    53  			c.buffer.Release()
    54  			c.buffer = nil
    55  		}
    56  		n += int64(wn)
    57  	}
    58  	cn, err := Copy(w, c.Conn)
    59  	n += cn
    60  	return
    61  }
    62  
    63  func (c *CachedConn) SetReadDeadline(t time.Time) error {
    64  	if c.buffer != nil && !c.buffer.IsEmpty() {
    65  		return nil
    66  	}
    67  	return c.Conn.SetReadDeadline(t)
    68  }
    69  
    70  func (c *CachedConn) ReadFrom(r io.Reader) (n int64, err error) {
    71  	return Copy(c.Conn, r)
    72  }
    73  
    74  func (c *CachedConn) Upstream() any {
    75  	return c.Conn
    76  }
    77  
    78  func (c *CachedConn) ReaderReplaceable() bool {
    79  	return c.buffer == nil
    80  }
    81  
    82  func (c *CachedConn) WriterReplaceable() bool {
    83  	return true
    84  }
    85  
    86  func (c *CachedConn) Close() error {
    87  	if buffer := c.buffer; buffer != nil {
    88  		buffer.DecRef()
    89  		buffer.Release()
    90  		c.buffer = nil
    91  	}
    92  	return c.Conn.Close()
    93  }
    94  
    95  type CachedReader struct {
    96  	upstream io.Reader
    97  	buffer   *buf.Buffer
    98  }
    99  
   100  func NewCachedReader(upstream io.Reader, buffer *buf.Buffer) *CachedReader {
   101  	buffer.IncRef()
   102  	return &CachedReader{
   103  		upstream: upstream,
   104  		buffer:   buffer,
   105  	}
   106  }
   107  
   108  func (r *CachedReader) ReadCached() *buf.Buffer {
   109  	buffer := r.buffer
   110  	r.buffer = nil
   111  	if buffer != nil {
   112  		buffer.DecRef()
   113  	}
   114  	return buffer
   115  }
   116  
   117  func (r *CachedReader) Read(p []byte) (n int, err error) {
   118  	if r.buffer != nil {
   119  		n, err = r.buffer.Read(p)
   120  		if err == nil {
   121  			return
   122  		}
   123  		r.buffer.DecRef()
   124  		r.buffer.Release()
   125  		r.buffer = nil
   126  	}
   127  	return r.upstream.Read(p)
   128  }
   129  
   130  func (r *CachedReader) WriteTo(w io.Writer) (n int64, err error) {
   131  	if r.buffer != nil {
   132  		wn, wErr := w.Write(r.buffer.Bytes())
   133  		if wErr != nil {
   134  			return 0, wErr
   135  		}
   136  		n += int64(wn)
   137  	}
   138  	cn, err := Copy(w, r.upstream)
   139  	n += cn
   140  	return
   141  }
   142  
   143  func (r *CachedReader) Upstream() any {
   144  	return r.upstream
   145  }
   146  
   147  func (r *CachedReader) ReaderReplaceable() bool {
   148  	return r.buffer == nil
   149  }
   150  
   151  func (r *CachedReader) Close() error {
   152  	if buffer := r.buffer; buffer != nil {
   153  		buffer.DecRef()
   154  		buffer.Release()
   155  		r.buffer = nil
   156  	}
   157  	return nil
   158  }
   159  
   160  type CachedPacketConn struct {
   161  	N.PacketConn
   162  	buffer      *buf.Buffer
   163  	destination M.Socksaddr
   164  }
   165  
   166  func NewCachedPacketConn(conn N.PacketConn, buffer *buf.Buffer, destination M.Socksaddr) *CachedPacketConn {
   167  	buffer.IncRef()
   168  	return &CachedPacketConn{
   169  		PacketConn:  conn,
   170  		buffer:      buffer,
   171  		destination: destination,
   172  	}
   173  }
   174  
   175  func (c *CachedPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
   176  	if c.buffer != nil {
   177  		_, err = buffer.ReadOnceFrom(c.buffer)
   178  		if err != nil {
   179  			return M.Socksaddr{}, err
   180  		}
   181  		c.buffer.DecRef()
   182  		c.buffer.Release()
   183  		c.buffer = nil
   184  		return c.destination, nil
   185  	}
   186  	return c.PacketConn.ReadPacket(buffer)
   187  }
   188  
   189  func (c *CachedPacketConn) ReadCachedPacket() *N.PacketBuffer {
   190  	buffer := c.buffer
   191  	c.buffer = nil
   192  	if buffer != nil {
   193  		buffer.DecRef()
   194  	}
   195  	return &N.PacketBuffer{
   196  		Buffer:      buffer,
   197  		Destination: c.destination,
   198  	}
   199  }
   200  
   201  func (c *CachedPacketConn) Upstream() any {
   202  	return c.PacketConn
   203  }
   204  
   205  func (c *CachedPacketConn) ReaderReplaceable() bool {
   206  	return c.buffer == nil
   207  }
   208  
   209  func (c *CachedPacketConn) WriterReplaceable() bool {
   210  	return true
   211  }
   212  
   213  func (c *CachedPacketConn) Close() error {
   214  	if buffer := c.buffer; buffer != nil {
   215  		buffer.DecRef()
   216  		buffer.Release()
   217  		c.buffer = nil
   218  	}
   219  	return c.PacketConn.Close()
   220  }