github.com/xxf098/lite-proxy@v0.15.1-0.20230422081941-12c69f323218/transport/ssr/protocol/stream.go (about)

     1  package protocol
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  
     7  	"github.com/xxf098/lite-proxy/common/pool"
     8  )
     9  
    10  type Conn struct {
    11  	net.Conn
    12  	Protocol
    13  	decoded      bytes.Buffer
    14  	underDecoded bytes.Buffer
    15  }
    16  
    17  func (c *Conn) Read(b []byte) (int, error) {
    18  	if c.decoded.Len() > 0 {
    19  		return c.decoded.Read(b)
    20  	}
    21  
    22  	buf := pool.Get(pool.RelayBufferSize)
    23  	defer pool.Put(buf)
    24  	n, err := c.Conn.Read(buf)
    25  	if err != nil {
    26  		return 0, err
    27  	}
    28  	c.underDecoded.Write(buf[:n])
    29  	err = c.Decode(&c.decoded, &c.underDecoded)
    30  	if err != nil {
    31  		return 0, err
    32  	}
    33  	n, _ = c.decoded.Read(b)
    34  	return n, nil
    35  }
    36  
    37  func (c *Conn) Write(b []byte) (int, error) {
    38  	bLength := len(b)
    39  	buf := pool.GetBuffer()
    40  	defer pool.PutBuffer(buf)
    41  	err := c.Encode(buf, b)
    42  	if err != nil {
    43  		return 0, err
    44  	}
    45  	_, err = c.Conn.Write(buf.Bytes())
    46  	if err != nil {
    47  		return 0, err
    48  	}
    49  	return bLength, nil
    50  }