// Copyright (C) 2018 Minio Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package sio implements the DARE format. It provides an API for secure
// en/decrypting IO operations using io.Reader and io.Writer.
package sio // import "github.com/minio/sio"

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"runtime"

	"golang.org/x/crypto/chacha20"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/sys/cpu"
)

const (
	// Version20 specifies version 2.0
	Version20 byte = 0x20
	// Version10 specifies version 1.0
	Version10 byte = 0x10
)

const (
	// AES_GCM specifies the cipher suite AES-GCM with 128 or 256 bit keys.
	AES_GCM byte = iota
	// CHACHA20_POLY1305 specifies the cipher suite ChaCha20Poly1305 with 256 bit keys.
	CHACHA20_POLY1305
)

// supportsAES indicates whether the CPU provides hardware support for AES-GCM.
// AES-GCM should only be selected as default cipher if there's hardware support.
var supportsAES = detectAESSupport()

func detectAESSupport() bool {
	// x86/x86_64: Check for AES-NI and PCLMULQDQ instructions
	if runtime.GOARCH == "amd64" || runtime.GOARCH == "386" {
		return cpu.X86.HasAES && cpu.X86.HasPCLMULQDQ
	}
	// ARM64: Check for AES and PMULL crypto extensions
	if runtime.GOARCH == "arm64" {
		return cpu.ARM64.HasAES && cpu.ARM64.HasPMULL
	}
	// s390x: Has built-in hardware acceleration
	if runtime.GOARCH == "s390x" {
		return true
	}
	// For other architectures, default to ChaCha20-Poly1305
	return false
}

const (
	headerSize     = 16
	maxPayloadSize = 1 << 16
	tagSize        = 16
	maxPackageSize = headerSize + maxPayloadSize + tagSize

	maxDecryptedSize = 1 << 48
	maxEncryptedSize = maxDecryptedSize + ((headerSize + tagSize) * 1 << 32)
)

var newAES = func(key []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

var newChaCha20 = func(key []byte) (cipher.AEAD, error) {
	// ChaCha20 requires 256 bit keys. Therefore, we expand the
	// given 128 bit key into a 256 bit key. However, we don't
	// use a simple algebraic key expansion to avoid any strutural
	// relation/pattern in the resulting 256 bit key. Even though
	// no effective attack against ChaCha20 is known that exploits
	// related keys, we don't want to make any additional security
	// assumption about ChaCha20.
	// Therefore, we use the HChaCha20 KDF to derive a uniformly
	// random 256 bit key from a 128 bit key as following:
	//
	//   Input:  K_i as 128 bit key
	//   Output: K_o as 256 bit key
	//
	//   k   = K_i | 0...0      # Expand K_i by 16 zero bytes to make it 256 bit long
	//   n   = 0...0 | 1        # Set a 128 bit nonce value n to 1
	//   K_o = HChaCha20(k, n)  # Derive K_o as output of the HCHaCha20 KDF
	//
	// This construction ensures that the K_o is 256 bit long and is not strutural related
	// to K_i since HChaCha20 is a secure PRF. As long as an attacker has no effective attack
	// distinguishing HChaCha20 from randomness, the attacker's best attack is a key space
	// search trying to find the right 128 bit key among 2^128 possible candidates - which
	// is not feasible.
	if len(key) == 128/8 {
		var (
			k     [32]byte
			nonce [16]byte
			err   error
		)
		copy(k[:], key)
		binary.BigEndian.PutUint32(nonce[12:], 1)

		if key, err = chacha20.HChaCha20(k[:], nonce[:]); err != nil {
			return nil, err
		}
	}
	return chacha20poly1305.New(key)
}

var supportedCiphers = [...]func([]byte) (cipher.AEAD, error){
	AES_GCM:           newAES,
	CHACHA20_POLY1305: newChaCha20,
}

var (
	errUnsupportedVersion = Error{"sio: unsupported version"}
	errUnsupportedCipher  = Error{"sio: unsupported cipher suite"}
	errInvalidPayloadSize = Error{"sio: invalid payload size"}
	errTagMismatch        = Error{"sio: authentication failed"}
	errUnexpectedSize     = Error{"sio: size is too large for DARE"}

	// Version 1.0 specific
	errPackageOutOfOrder = Error{"sio: sequence number mismatch"}

	// Version 2.0 specific
	errNonceMismatch   = Error{"sio: header nonce mismatch"}
	errUnexpectedEOF   = Error{"sio: unexpected EOF"}
	errUnexpectedData  = Error{"sio: unexpected data after final package"}
	errSealAfterFinal  = Error{"sio: cannot seal any package after final one"}
	errWriteAfterClose = Error{"sio: write to stream after close"}
)

// Error is the error returned by an io.Reader or io.Writer
// if the encrypted data cannot be decrypted because it is
// malformed or not authentic.
type Error struct{ msg string }

func (e Error) Error() string { return e.msg }

// Config contains the format configuration. The only field
// which must always be set manually is the secret key.
type Config struct {
	// The minimal supported version of the format. If
	// not set the default value - Version10 - is used.
	MinVersion byte

	// The highest supported version of the format. If
	// not set the default value - Version20 - is used.
	MaxVersion byte

	// A list of supported cipher suites. If not set the
	// default value is used.
	CipherSuites []byte

	// The secret encryption key. It must be 32 bytes long.
	Key []byte

	// The first expected sequence number. It should only
	// be set manually when decrypting a range within a
	// stream.
	SequenceNumber uint32

	// The RNG used to generate random values. If not set
	// the default value (crypto/rand.Reader) is used.
	Rand io.Reader

	// Nonce will override the nonce if set non-nil.
	// V2 will use all 12 bytes, V1 first 8 bytes.
	Nonce *[12]byte

	// The size of the encrypted payload in bytes. The
	// default value is 64KB. It should be used to restrict
	// the size of encrypted packages. The payload size
	// must be between 1 and 64 KB.
	//
	// This field is specific for version 1.0 and is
	// deprecated.
	PayloadSize int
}

// EncryptedSize computes the size of an encrypted data stream
// from the plaintext size. It is the inverse of DecryptedSize().
//
// EncryptedSize returns an error if the provided size is to large.
func EncryptedSize(size uint64) (uint64, error) {
	if size > maxDecryptedSize {
		return 0, errUnexpectedSize
	}

	encSize := (size / maxPayloadSize) * maxPackageSize
	if mod := size % maxPayloadSize; mod > 0 {
		encSize += mod + (headerSize + tagSize)
	}
	return encSize, nil
}

// DecryptedSize computes the size of a decrypted data stream
// from the encrypted stream size. It is the inverse of EncryptedSize().
//
// DecryptedSize returns an error if the provided size is to large
// or if the provided size is an invalid encrypted stream size.
func DecryptedSize(size uint64) (uint64, error) {
	if size > maxEncryptedSize {
		return 0, errUnexpectedSize
	}
	decSize := (size / maxPackageSize) * maxPayloadSize
	if mod := size % maxPackageSize; mod > 0 {
		if mod <= headerSize+tagSize {
			return 0, errors.New("sio: size is not valid") // last package is not valid
		}
		decSize += mod - (headerSize + tagSize)
	}
	return decSize, nil
}

// Encrypt reads from src until it encounters an io.EOF and encrypts all received
// data. The encrypted data is written to dst. It returns the number of bytes
// encrypted and the first error encountered while encrypting, if any.
//
// Encrypt returns the number of bytes written to dst.
func Encrypt(dst io.Writer, src io.Reader, config Config) (n int64, err error) {
	encReader, err := EncryptReader(src, config)
	if err != nil {
		return 0, fmt.Errorf("sio: failed to create encryption reader: %w", err)
	}
	n, err = io.CopyBuffer(dst, encReader, make([]byte, headerSize+maxPayloadSize+tagSize))
	if err != nil {
		return n, fmt.Errorf("sio: encryption failed: %w", err)
	}
	return n, nil
}

// Decrypt reads from src until it encounters an io.EOF and decrypts all received
// data. The decrypted data is written to dst. It returns the number of bytes
// decrypted and the first error encountered while decrypting, if any.
//
// Decrypt returns the number of bytes written to dst. Decrypt only writes data to
// dst if the data was decrypted successfully. It returns an error of type sio.Error
// if decryption fails.
func Decrypt(dst io.Writer, src io.Reader, config Config) (n int64, err error) {
	decReader, err := DecryptReader(src, config)
	if err != nil {
		return 0, fmt.Errorf("sio: failed to create decryption reader: %w", err)
	}
	n, err = io.CopyBuffer(dst, decReader, make([]byte, maxPayloadSize))
	if err != nil {
		return n, fmt.Errorf("sio: decryption failed: %w", err)
	}
	return n, nil
}

// DecryptBuffer decrypts all received data in src.
// The decrypted data is appended to dst.
// If the number of output bytes is unknown,
// making a dst with capacity of len(src) is reasonable.
//
// DecryptBuffer only returns data to if the data was decrypted successfully.
// It returns an error of type sio.Error if decryption fails.
func DecryptBuffer(dst, src []byte, config Config) (output []byte, err error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MinVersion == Version10 && config.MaxVersion == Version10 {
		buf := bytes.NewBuffer(dst)
		if _, err := Decrypt(buf, bytes.NewBuffer(src), config); err != nil {
			return nil, err
		}
		return buf.Bytes(), nil
	}
	if config.MinVersion == Version20 && config.MaxVersion == Version20 {
		return decryptBufferV20(dst, src, &config)
	}
	return decryptBuffer(dst, src, &config)
}

// EncryptReader wraps the given src and returns an io.Reader which encrypts
// all received data. EncryptReader returns an error if the provided encryption
// configuration is invalid.
func EncryptReader(src io.Reader, config Config) (io.Reader, error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MaxVersion == Version20 {
		return encryptReaderV20(src, &config)
	}
	return encryptReaderV10(src, &config)
}

// DecryptReader wraps the given src and returns an io.Reader which decrypts
// all received data. DecryptReader returns an error if the provided decryption
// configuration is invalid. The returned io.Reader returns an error of
// type sio.Error if the decryption fails.
func DecryptReader(src io.Reader, config Config) (io.Reader, error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MinVersion == Version10 && config.MaxVersion == Version10 {
		return decryptReaderV10(src, &config)
	}
	if config.MinVersion == Version20 && config.MaxVersion == Version20 {
		return decryptReaderV20(src, &config)
	}
	return decryptReader(src, &config), nil
}

// DecryptReaderAt wraps the given src and returns an io.ReaderAt which decrypts
// all received data. DecryptReaderAt returns an error if the provided decryption
// configuration is invalid. The returned io.ReaderAt returns an error of
// type sio.Error if the decryption fails.
func DecryptReaderAt(src io.ReaderAt, config Config) (io.ReaderAt, error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MinVersion == Version10 && config.MaxVersion == Version10 {
		return decryptReaderAtV10(src, &config)
	}
	if config.MinVersion == Version20 && config.MaxVersion == Version20 {
		return decryptReaderAtV20(src, &config)
	}
	return decryptReaderAt(src, &config), nil
}

// EncryptWriter wraps the given dst and returns an io.WriteCloser which
// encrypts all data written to it. EncryptWriter returns an error if the
// provided decryption configuration is invalid.
//
// The returned io.WriteCloser must be closed successfully to finalize the
// encryption process.
func EncryptWriter(dst io.Writer, config Config) (io.WriteCloser, error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MaxVersion == Version20 {
		return encryptWriterV20(dst, &config)
	}
	return encryptWriterV10(dst, &config)
}

// DecryptWriter wraps the given dst and returns an io.WriteCloser which
// decrypts all data written to it. DecryptWriter returns an error if the
// provided decryption configuration is invalid.
//
// The returned io.WriteCloser must be closed successfully to finalize the
// decryption process. The returned io.WriteCloser returns an error of
// type sio.Error if the decryption fails.
func DecryptWriter(dst io.Writer, config Config) (io.WriteCloser, error) {
	if err := setConfigDefaults(&config); err != nil {
		return nil, err
	}
	if config.MinVersion == Version10 && config.MaxVersion == Version10 {
		return decryptWriterV10(dst, &config)
	}
	if config.MinVersion == Version20 && config.MaxVersion == Version20 {
		return decryptWriterV20(dst, &config)
	}
	return decryptWriter(dst, &config), nil
}

func defaultCipherSuites() []byte {
	if supportsAES {
		return []byte{AES_GCM, CHACHA20_POLY1305}
	}
	return []byte{CHACHA20_POLY1305, AES_GCM}
}

func setConfigDefaults(config *Config) error {
	if config.MinVersion > Version20 {
		return errors.New("sio: unknown minimum version")
	}
	if config.MaxVersion > Version20 {
		return errors.New("sio: unknown maximum version")
	}
	if len(config.Key) != 128/8 && len(config.Key) != 256/8 {
		return errors.New("sio: invalid key size")
	}
	if len(config.CipherSuites) > 2 {
		return errors.New("sio: too many cipher suites")
	}
	for _, c := range config.CipherSuites {
		if int(c) >= len(supportedCiphers) {
			return errors.New("sio: unknown cipher suite")
		}
	}
	if config.PayloadSize > maxPayloadSize {
		return errors.New("sio: payload size is too large")
	}

	if config.MinVersion < Version10 {
		config.MinVersion = Version10
	}
	if config.MaxVersion < Version10 {
		config.MaxVersion = Version20
	}
	if config.MinVersion > config.MaxVersion {
		return errors.New("sio: minimum version cannot be larger than maximum version")
	}
	if len(config.CipherSuites) == 0 {
		config.CipherSuites = defaultCipherSuites()
	}
	if config.Rand == nil {
		config.Rand = rand.Reader
	}
	if config.PayloadSize == 0 {
		config.PayloadSize = maxPayloadSize
	}
	return nil
}
