v2ray-core/io/vmess/decryptionreader.go

package vmess

import (
  "bytes"
	"crypto/aes"
	"crypto/cipher"
	"fmt"
	"io"
)

const (
	blockSize = 16
)

type DecryptionReader struct {
	cipher             cipher.Block
	reader             io.Reader
	buffer             *bytes.Buffer
}

func NewDecryptionReader(reader io.Reader, key []byte) (*DecryptionReader, error) {
	decryptionReader := new(DecryptionReader)
	cipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	decryptionReader.cipher = cipher
	decryptionReader.reader = reader
  decryptionReader.buffer = bytes.NewBuffer(make([]byte, 0, 2 * blockSize))
	return decryptionReader, nil
}

func (reader *DecryptionReader) readBlock() error {
  buffer := make([]byte, blockSize)
  nBytes, err := reader.reader.Read(buffer)
	if err != nil {
		return err
	}
	if nBytes < blockSize {
		return fmt.Errorf("Expected to read %d bytes, but got %d bytes", blockSize, nBytes)
	}
	reader.cipher.Decrypt(buffer, buffer)
  reader.buffer.Write(buffer)
	return nil
}

func (reader *DecryptionReader) Read(p []byte) (int, error) {
  if reader.buffer.Len() == 0 {
    err := reader.readBlock()
    if err != nil {
      return 0, err
    }
  }
	nBytes, err := reader.buffer.Read(p)
  if err != nil {
    return nBytes, err
  }
  if nBytes < len(p) {
    err = reader.readBlock()
    if err != nil {
      return nBytes, err
    }
    moreBytes, err := reader.buffer.Read(p[nBytes:])
    if err != nil {
      return nBytes, err
    }
    nBytes += moreBytes
    if nBytes != len(p) {
      return nBytes, fmt.Errorf("Unable to read %d bytes", len(p))
    }
  }
  return nBytes, err
}