github.com/EagleQL/Xray-core@v1.4.3/common/crypto/io.go (about)

     1  package crypto
     2  
     3  import (
     4  	"crypto/cipher"
     5  	"io"
     6  
     7  	"github.com/xtls/xray-core/common/buf"
     8  )
     9  
    10  type CryptionReader struct {
    11  	stream cipher.Stream
    12  	reader io.Reader
    13  }
    14  
    15  func NewCryptionReader(stream cipher.Stream, reader io.Reader) *CryptionReader {
    16  	return &CryptionReader{
    17  		stream: stream,
    18  		reader: reader,
    19  	}
    20  }
    21  
    22  func (r *CryptionReader) Read(data []byte) (int, error) {
    23  	nBytes, err := r.reader.Read(data)
    24  	if nBytes > 0 {
    25  		r.stream.XORKeyStream(data[:nBytes], data[:nBytes])
    26  	}
    27  	return nBytes, err
    28  }
    29  
    30  var (
    31  	_ buf.Writer = (*CryptionWriter)(nil)
    32  )
    33  
    34  type CryptionWriter struct {
    35  	stream    cipher.Stream
    36  	writer    io.Writer
    37  	bufWriter buf.Writer
    38  }
    39  
    40  // NewCryptionWriter creates a new CryptionWriter.
    41  func NewCryptionWriter(stream cipher.Stream, writer io.Writer) *CryptionWriter {
    42  	return &CryptionWriter{
    43  		stream:    stream,
    44  		writer:    writer,
    45  		bufWriter: buf.NewWriter(writer),
    46  	}
    47  }
    48  
    49  // Write implements io.Writer.Write().
    50  func (w *CryptionWriter) Write(data []byte) (int, error) {
    51  	w.stream.XORKeyStream(data, data)
    52  
    53  	if err := buf.WriteAllBytes(w.writer, data); err != nil {
    54  		return 0, err
    55  	}
    56  	return len(data), nil
    57  }
    58  
    59  // WriteMultiBuffer implements buf.Writer.
    60  func (w *CryptionWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
    61  	for _, b := range mb {
    62  		w.stream.XORKeyStream(b.Bytes(), b.Bytes())
    63  	}
    64  
    65  	return w.bufWriter.WriteMultiBuffer(mb)
    66  }