v2ray-core/io/encryption.go

package io

import (
  "crypto/cipher"
  "io"
)

// CryptionReader is a general purpose reader that applies
// block cipher on top of a regular reader.
type CryptionReader struct {
  mode cipher.BlockMode
  reader io.Reader
}

func NewCryptionReader(mode cipher.BlockMode, reader io.Reader) *CryptionReader {
  this := new(CryptionReader)
  this.mode = mode
  this.reader = reader
  return this
}

// Read reads blocks from underlying reader, the length of blocks must be
// a multiply of BlockSize()
func (reader CryptionReader) Read(blocks []byte) (int, error) {
  nBytes, err := reader.reader.Read(blocks)
  if err != nil && err != io.EOF {
    return nBytes, err
  }
  if nBytes < len(blocks) {
    for i, _ := range blocks[nBytes:] {
      blocks[i] = 0
    }
  }
  reader.mode.CryptBlocks(blocks, blocks)
  return nBytes, err
}

func (reader CryptionReader) BlockSize() int {
  return reader.mode.BlockSize()
}

// Cryption writer is a general purpose of byte stream writer that applies
// block cipher on top of a regular writer.
type CryptionWriter struct {
  mode cipher.BlockMode
  writer io.Writer
}

func NewCryptionWriter(mode cipher.BlockMode, writer io.Writer) *CryptionWriter {
  this := new(CryptionWriter)
  this.mode = mode
  this.writer = writer
  return this
}

// Write writes the give blocks to underlying writer. The length of the blocks
// must be a multiply of BlockSize()
func (writer CryptionWriter) Write(blocks []byte) (int, error) {
  writer.mode.CryptBlocks(blocks, blocks)
  return writer.writer.Write(blocks)
}

func (writer CryptionWriter) BlockSize() int {
  return writer.mode.BlockSize()
}