github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/common/crypto/io.go (about)

     1  package crypto
     2  
     3  import (
     4  	"crypto/cipher"
     5  	"io"
     6  
     7  	"github.com/xmplusdev/xmcore/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 _ buf.Writer = (*CryptionWriter)(nil)
    31  
    32  type CryptionWriter struct {
    33  	stream    cipher.Stream
    34  	writer    io.Writer
    35  	bufWriter buf.Writer
    36  }
    37  
    38  // NewCryptionWriter creates a new CryptionWriter.
    39  func NewCryptionWriter(stream cipher.Stream, writer io.Writer) *CryptionWriter {
    40  	return &CryptionWriter{
    41  		stream:    stream,
    42  		writer:    writer,
    43  		bufWriter: buf.NewWriter(writer),
    44  	}
    45  }
    46  
    47  // Write implements io.Writer.Write().
    48  func (w *CryptionWriter) Write(data []byte) (int, error) {
    49  	w.stream.XORKeyStream(data, data)
    50  
    51  	if err := buf.WriteAllBytes(w.writer, data, nil); err != nil {
    52  		return 0, err
    53  	}
    54  	return len(data), nil
    55  }
    56  
    57  // WriteMultiBuffer implements buf.Writer.
    58  func (w *CryptionWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
    59  	for _, b := range mb {
    60  		w.stream.XORKeyStream(b.Bytes(), b.Bytes())
    61  	}
    62  
    63  	return w.bufWriter.WriteMultiBuffer(mb)
    64  }