github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/shadowsocks/shadowaead/stream.go (about)

     1  package shadowaead
     2  
     3  import (
     4  	"crypto/rand"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/shadowsocks/internal"
     9  	ycrypto "github.com/Asutorufa/yuhaiin/pkg/net/proxy/yuubinsya/crypto"
    10  )
    11  
    12  // payloadSizeMax is the maximum size of payload in bytes.
    13  const payloadSizeMax = 0x3FFF // 16*1024 - 1
    14  
    15  type streamConn struct {
    16  	net.Conn
    17  	Cipher
    18  	r io.Reader
    19  	w io.Writer
    20  }
    21  
    22  func (c *streamConn) initReader() error {
    23  	salt := make([]byte, c.SaltSize())
    24  	if _, err := io.ReadFull(c.Conn, salt); err != nil {
    25  		return err
    26  	}
    27  	aead, err := c.Decrypter(salt)
    28  	if err != nil {
    29  		return err
    30  	}
    31  
    32  	if internal.CheckSalt(salt) {
    33  		return ErrRepeatedSalt
    34  	}
    35  
    36  	c.r = ycrypto.NewReader(c.Conn, make([]byte, aead.NonceSize()), aead, payloadSizeMax)
    37  	return nil
    38  }
    39  
    40  func (c *streamConn) Read(b []byte) (int, error) {
    41  	if c.r == nil {
    42  		if err := c.initReader(); err != nil {
    43  			return 0, err
    44  		}
    45  	}
    46  	return c.r.Read(b)
    47  }
    48  
    49  func (c *streamConn) initWriter() error {
    50  	salt := make([]byte, c.SaltSize())
    51  	if _, err := io.ReadFull(rand.Reader, salt); err != nil {
    52  		return err
    53  	}
    54  	aead, err := c.Encrypter(salt)
    55  	if err != nil {
    56  		return err
    57  	}
    58  	_, err = c.Conn.Write(salt)
    59  	if err != nil {
    60  		return err
    61  	}
    62  	internal.AddSalt(salt)
    63  	c.w = ycrypto.NewWriter(c.Conn, make([]byte, aead.NonceSize()), aead, payloadSizeMax)
    64  	return nil
    65  }
    66  
    67  func (c *streamConn) Write(b []byte) (int, error) {
    68  	if c.w == nil {
    69  		if err := c.initWriter(); err != nil {
    70  			return 0, err
    71  		}
    72  	}
    73  	return c.w.Write(b)
    74  }
    75  
    76  // NewConn wraps a stream-oriented net.Conn with cipher.
    77  func NewConn(c net.Conn, ciph Cipher) net.Conn { return &streamConn{Conn: c, Cipher: ciph} }