github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/internal/gateway/tcp/conn_wrap.go (about)

     1  package tcpGateway
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"net"
     7  	"os"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/mailru/easygo/netpoll"
    12  )
    13  
    14  /*
    15     Creation Time: 2020 - Aug - 10
    16     Created by:  (ehsan)
    17     Maintainers:
    18        1.  Ehsan N. Moosa (E2)
    19     Auditor: Ehsan N. Moosa (E2)
    20     Copyright Ronak Software Group 2020
    21  */
    22  
    23  var (
    24  	wrapConnPool sync.Pool
    25  )
    26  
    27  // filer describes an object that has ability to return os.File.
    28  type filer interface {
    29  	// File returns a copy of object's file descriptor.
    30  	File() (*os.File, error)
    31  }
    32  
    33  type wrapConn struct {
    34  	c   net.Conn
    35  	r   io.Reader
    36  	buf *bytes.Buffer
    37  }
    38  
    39  func (wc *wrapConn) Write(b []byte) (n int, err error) {
    40  	return wc.c.Write(b)
    41  }
    42  
    43  func (wc *wrapConn) Close() error {
    44  	err := wc.c.Close()
    45  	releaseWrapConn(wc)
    46  
    47  	return err
    48  }
    49  
    50  func (wc *wrapConn) LocalAddr() net.Addr {
    51  	return wc.c.LocalAddr()
    52  }
    53  
    54  func (wc *wrapConn) RemoteAddr() net.Addr {
    55  	return wc.c.RemoteAddr()
    56  }
    57  
    58  func (wc *wrapConn) SetDeadline(t time.Time) error {
    59  	return wc.c.SetDeadline(t)
    60  }
    61  
    62  func (wc *wrapConn) SetReadDeadline(t time.Time) error {
    63  	return wc.c.SetReadDeadline(t)
    64  }
    65  
    66  func (wc *wrapConn) SetWriteDeadline(t time.Time) error {
    67  	return wc.c.SetWriteDeadline(t)
    68  }
    69  
    70  func newWrapConn(c net.Conn) *wrapConn {
    71  	wc := &wrapConn{
    72  		c:   c,
    73  		buf: bytes.NewBuffer(make([]byte, 0, 128)),
    74  	}
    75  	wc.r = io.TeeReader(wc.c, wc.buf)
    76  
    77  	return wc
    78  }
    79  
    80  func acquireWrapConn(c net.Conn) *wrapConn {
    81  	wc, ok := wrapConnPool.Get().(*wrapConn)
    82  	if !ok {
    83  		return newWrapConn(c)
    84  	}
    85  	wc.c = c
    86  	wc.r = io.TeeReader(wc.c, wc.buf)
    87  
    88  	return wc
    89  }
    90  
    91  func releaseWrapConn(wc *wrapConn) {
    92  	wc.buf.Reset()
    93  	wrapConnPool.Put(wc)
    94  }
    95  
    96  func (wc *wrapConn) UnsafeConn() net.Conn {
    97  	return wc.c
    98  }
    99  
   100  func (wc *wrapConn) Read(p []byte) (int, error) {
   101  	return wc.r.Read(p)
   102  }
   103  
   104  func (wc *wrapConn) File() (*os.File, error) {
   105  	x, ok := wc.c.(filer)
   106  	if !ok {
   107  		return nil, netpoll.ErrNotFiler
   108  	}
   109  
   110  	return x.File()
   111  }
   112  
   113  func (wc *wrapConn) ReadyForUpgrade() {
   114  	wc.r = io.MultiReader(wc.buf, wc.c)
   115  }