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

     1  package wsutil
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/ronaksoft/rony/pools"
     7  
     8  	"github.com/gobwas/ws"
     9  )
    10  
    11  // CipherReader implements io.Reader that applies xor-cipher to the bytes read
    12  // from source.
    13  // It could help to unmask WebSocket frame payload on the fly.
    14  type CipherReader struct {
    15  	r    io.Reader
    16  	mask [4]byte
    17  	pos  int
    18  }
    19  
    20  // NewCipherReader creates xor-cipher reader from r with given mask.
    21  func NewCipherReader(r io.Reader, mask [4]byte) *CipherReader {
    22  	return &CipherReader{r, mask, 0}
    23  }
    24  
    25  // Reset resets CipherReader to read from r with given mask.
    26  func (c *CipherReader) Reset(r io.Reader, mask [4]byte) {
    27  	c.r = r
    28  	c.mask = mask
    29  	c.pos = 0
    30  }
    31  
    32  // Read implements io.Reader interface. It applies mask given during
    33  // initialization to every read byte.
    34  func (c *CipherReader) Read(p []byte) (n int, err error) {
    35  	n, err = c.r.Read(p)
    36  	ws.Cipher(p[:n], c.mask, c.pos)
    37  	c.pos += n
    38  
    39  	return
    40  }
    41  
    42  // CipherWriter implements io.Writer that applies xor-cipher to the bytes
    43  // written to the destination writer. It does not modify the original bytes.
    44  type CipherWriter struct {
    45  	w    io.Writer
    46  	mask [4]byte
    47  	pos  int
    48  }
    49  
    50  // NewCipherWriter creates xor-cipher writer to w with given mask.
    51  func NewCipherWriter(w io.Writer, mask [4]byte) *CipherWriter {
    52  	return &CipherWriter{w, mask, 0}
    53  }
    54  
    55  // Reset reset CipherWriter to write to w with given mask.
    56  func (c *CipherWriter) Reset(w io.Writer, mask [4]byte) {
    57  	c.w = w
    58  	c.mask = mask
    59  	c.pos = 0
    60  }
    61  
    62  // Write implements io.Writer interface. It applies masking during
    63  // initialization to every sent byte. It does not modify original slice.
    64  func (c *CipherWriter) Write(p []byte) (n int, err error) {
    65  	buf := pools.Buffer.GetLen(len(p))
    66  	buf.CopyFrom(p)
    67  	ws.Cipher(*buf.Bytes(), c.mask, c.pos)
    68  	n, err = c.w.Write(*buf.Bytes())
    69  	c.pos += n
    70  	pools.Buffer.Put(buf)
    71  
    72  	return
    73  }