ChaCha20.swift

//
//  CryptoSwift
//
//  Copyright (C) 2014-2022 Marcin Krzyżanowski <marcin@krzyzanowskim.com>
//  This software is provided 'as-is', without any express or implied warranty.
//
//  In no event will the authors be held liable for any damages arising from the use of this software.
//
//  Permission is granted to anyone to use this software for any purpose,including commercial applications, and to alter it and redistribute it freely, subject to the following restrictions:
//
//  - The origin of this software must not be misrepresented; you must not claim that you wrote the original software. If you use this software in a product, an acknowledgment in the product documentation is required.
//  - Altered source versions must be plainly marked as such, and must not be misrepresented as being the original software.
//  - This notice may not be removed or altered from any source or binary distribution.
//

//  https://tools.ietf.org/html/rfc7539
//

public final class ChaCha20: BlockCipher {
  public enum Error: Swift.Error {
    case invalidKeyOrInitializationVector
    case notSupported
  }

  public static let blockSize = 64 // 512 / 8
  public let keySize: Int

  fileprivate let key: Key
  fileprivate var counter: Array<UInt8>

  public init(key: Array<UInt8>, counter: Array<UInt8> = [0, 0, 0, 0], iv nonce: Array<UInt8>) throws {
    precondition(counter.count == 4 && (nonce.count == 12 || nonce.count == 8))

    if key.count != 32 {
      throw Error.invalidKeyOrInitializationVector
    }

    self.key = Key(bytes: key)
    self.keySize = self.key.count

    if nonce.count == 8 {
      self.counter = counter + [0, 0, 0, 0] + nonce
    } else if nonce.count == 12 {
      self.counter = counter + nonce
    } else {
      throw Error.invalidKeyOrInitializationVector
    }

    assert(self.counter.count == 16)
  }

  /// https://tools.ietf.org/html/rfc7539#section-2.3.
  fileprivate func core(block: inout Array<UInt8>, counter: Array<UInt8>, key: Array<UInt8>) {
    precondition(block.count == ChaCha20.blockSize)
    precondition(counter.count == 16)
    precondition(key.count == 32)

    let j0: UInt32 = 0x61707865
    let j1: UInt32 = 0x3320646e // 0x3620646e sigma/tau
    let j2: UInt32 = 0x79622d32
    let j3: UInt32 = 0x6b206574
    let j4: UInt32 = UInt32(bytes: key[0..<4]).bigEndian
    let j5: UInt32 = UInt32(bytes: key[4..<8]).bigEndian
    let j6: UInt32 = UInt32(bytes: key[8..<12]).bigEndian
    let j7: UInt32 = UInt32(bytes: key[12..<16]).bigEndian
    let j8: UInt32 = UInt32(bytes: key[16..<20]).bigEndian
    let j9: UInt32 = UInt32(bytes: key[20..<24]).bigEndian
    let j10: UInt32 = UInt32(bytes: key[24..<28]).bigEndian
    let j11: UInt32 = UInt32(bytes: key[28..<32]).bigEndian
    let j12: UInt32 = UInt32(bytes: counter[0..<4]).bigEndian
    let j13: UInt32 = UInt32(bytes: counter[4..<8]).bigEndian
    let j14: UInt32 = UInt32(bytes: counter[8..<12]).bigEndian
    let j15: UInt32 = UInt32(bytes: counter[12..<16]).bigEndian

    var (x0, x1, x2, x3, x4, x5, x6, x7) = (j0, j1, j2, j3, j4, j5, j6, j7)
    var (x8, x9, x10, x11, x12, x13, x14, x15) = (j8, j9, j10, j11, j12, j13, j14, j15)

    for _ in 0..<10 { // 20 rounds
      x0 = x0 &+ x4
      x12 ^= x0
      x12 = (x12 << 16) | (x12 >> 16)
      x8 = x8 &+ x12
      x4 ^= x8
      x4 = (x4 << 12) | (x4 >> 20)
      x0 = x0 &+ x4
      x12 ^= x0
      x12 = (x12 << 8) | (x12 >> 24)
      x8 = x8 &+ x12
      x4 ^= x8
      x4 = (x4 << 7) | (x4 >> 25)
      x1 = x1 &+ x5
      x13 ^= x1
      x13 = (x13 << 16) | (x13 >> 16)
      x9 = x9 &+ x13
      x5 ^= x9
      x5 = (x5 << 12) | (x5 >> 20)
      x1 = x1 &+ x5
      x13 ^= x1
      x13 = (x13 << 8) | (x13 >> 24)
      x9 = x9 &+ x13
      x5 ^= x9
      x5 = (x5 << 7) | (x5 >> 25)
      x2 = x2 &+ x6
      x14 ^= x2
      x14 = (x14 << 16) | (x14 >> 16)
      x10 = x10 &+ x14
      x6 ^= x10
      x6 = (x6 << 12) | (x6 >> 20)
      x2 = x2 &+ x6
      x14 ^= x2
      x14 = (x14 << 8) | (x14 >> 24)
      x10 = x10 &+ x14
      x6 ^= x10
      x6 = (x6 << 7) | (x6 >> 25)
      x3 = x3 &+ x7
      x15 ^= x3
      x15 = (x15 << 16) | (x15 >> 16)
      x11 = x11 &+ x15
      x7 ^= x11
      x7 = (x7 << 12) | (x7 >> 20)
      x3 = x3 &+ x7
      x15 ^= x3
      x15 = (x15 << 8) | (x15 >> 24)
      x11 = x11 &+ x15
      x7 ^= x11
      x7 = (x7 << 7) | (x7 >> 25)
      x0 = x0 &+ x5
      x15 ^= x0
      x15 = (x15 << 16) | (x15 >> 16)
      x10 = x10 &+ x15
      x5 ^= x10
      x5 = (x5 << 12) | (x5 >> 20)
      x0 = x0 &+ x5
      x15 ^= x0
      x15 = (x15 << 8) | (x15 >> 24)
      x10 = x10 &+ x15
      x5 ^= x10
      x5 = (x5 << 7) | (x5 >> 25)
      x1 = x1 &+ x6
      x12 ^= x1
      x12 = (x12 << 16) | (x12 >> 16)
      x11 = x11 &+ x12
      x6 ^= x11
      x6 = (x6 << 12) | (x6 >> 20)
      x1 = x1 &+ x6
      x12 ^= x1
      x12 = (x12 << 8) | (x12 >> 24)
      x11 = x11 &+ x12
      x6 ^= x11
      x6 = (x6 << 7) | (x6 >> 25)
      x2 = x2 &+ x7
      x13 ^= x2
      x13 = (x13 << 16) | (x13 >> 16)
      x8 = x8 &+ x13
      x7 ^= x8
      x7 = (x7 << 12) | (x7 >> 20)
      x2 = x2 &+ x7
      x13 ^= x2
      x13 = (x13 << 8) | (x13 >> 24)
      x8 = x8 &+ x13
      x7 ^= x8
      x7 = (x7 << 7) | (x7 >> 25)
      x3 = x3 &+ x4
      x14 ^= x3
      x14 = (x14 << 16) | (x14 >> 16)
      x9 = x9 &+ x14
      x4 ^= x9
      x4 = (x4 << 12) | (x4 >> 20)
      x3 = x3 &+ x4
      x14 ^= x3
      x14 = (x14 << 8) | (x14 >> 24)
      x9 = x9 &+ x14
      x4 ^= x9
      x4 = (x4 << 7) | (x4 >> 25)
    }

    x0 = x0 &+ j0
    x1 = x1 &+ j1
    x2 = x2 &+ j2
    x3 = x3 &+ j3
    x4 = x4 &+ j4
    x5 = x5 &+ j5
    x6 = x6 &+ j6
    x7 = x7 &+ j7
    x8 = x8 &+ j8
    x9 = x9 &+ j9
    x10 = x10 &+ j10
    x11 = x11 &+ j11
    x12 = x12 &+ j12
    x13 = x13 &+ j13
    x14 = x14 &+ j14
    x15 = x15 &+ j15

    block.replaceSubrange(0..<4, with: x0.bigEndian.bytes())
    block.replaceSubrange(4..<8, with: x1.bigEndian.bytes())
    block.replaceSubrange(8..<12, with: x2.bigEndian.bytes())
    block.replaceSubrange(12..<16, with: x3.bigEndian.bytes())
    block.replaceSubrange(16..<20, with: x4.bigEndian.bytes())
    block.replaceSubrange(20..<24, with: x5.bigEndian.bytes())
    block.replaceSubrange(24..<28, with: x6.bigEndian.bytes())
    block.replaceSubrange(28..<32, with: x7.bigEndian.bytes())
    block.replaceSubrange(32..<36, with: x8.bigEndian.bytes())
    block.replaceSubrange(36..<40, with: x9.bigEndian.bytes())
    block.replaceSubrange(40..<44, with: x10.bigEndian.bytes())
    block.replaceSubrange(44..<48, with: x11.bigEndian.bytes())
    block.replaceSubrange(48..<52, with: x12.bigEndian.bytes())
    block.replaceSubrange(52..<56, with: x13.bigEndian.bytes())
    block.replaceSubrange(56..<60, with: x14.bigEndian.bytes())
    block.replaceSubrange(60..<64, with: x15.bigEndian.bytes())
  }

  // XORKeyStream
  func process(bytes: ArraySlice<UInt8>, counter: inout Array<UInt8>, key: Array<UInt8>) -> Array<UInt8> {
    precondition(counter.count == 16)
    precondition(key.count == 32)

    var block = Array<UInt8>(repeating: 0, count: ChaCha20.blockSize)
    var bytesSlice = bytes
    var out = Array<UInt8>(reserveCapacity: bytesSlice.count)

    while bytesSlice.count >= ChaCha20.blockSize {
      self.core(block: &block, counter: counter, key: key)
      for (i, x) in block.enumerated() {
        out.append(bytesSlice[bytesSlice.startIndex + i] ^ x)
      }
      var u: UInt32 = 1
      for i in 0..<4 {
        u += UInt32(counter[i])
        counter[i] = UInt8(u & 0xff)
        u >>= 8
      }
      bytesSlice = bytesSlice[bytesSlice.startIndex + ChaCha20.blockSize..<bytesSlice.endIndex]
    }

    if !bytesSlice.isEmpty {
      self.core(block: &block, counter: counter, key: key)
      for (i, v) in bytesSlice.enumerated() {
        out.append(v ^ block[i])
      }
    }
    return out
  }
}

// MARK: Cipher

extension ChaCha20: Cipher {
  public func encrypt(_ bytes: ArraySlice<UInt8>) throws -> Array<UInt8> {
    self.process(bytes: bytes, counter: &self.counter, key: Array(self.key))
  }

  public func decrypt(_ bytes: ArraySlice<UInt8>) throws -> Array<UInt8> {
    try self.encrypt(bytes)
  }
}

// MARK: Encryptor

extension ChaCha20 {
  public struct ChaChaEncryptor: Cryptor, Updatable {
    private var accumulated = Array<UInt8>()
    private let chacha: ChaCha20

    init(chacha: ChaCha20) {
      self.chacha = chacha
    }

    public mutating func update(withBytes bytes: ArraySlice<UInt8>, isLast: Bool = false) throws -> Array<UInt8> {
      self.accumulated += bytes

      var encrypted = Array<UInt8>()
      encrypted.reserveCapacity(self.accumulated.count)
      for chunk in self.accumulated.batched(by: ChaCha20.blockSize) {
        if isLast || self.accumulated.count >= ChaCha20.blockSize {
          encrypted += try self.chacha.encrypt(chunk)
          self.accumulated.removeFirst(chunk.count) // TODO: improve performance
        }
      }
      return encrypted
    }

    public func seek(to: Int) throws {
      throw Error.notSupported
    }
  }
}

// MARK: Decryptor

extension ChaCha20 {
  public struct ChaChaDecryptor: Cryptor, Updatable {
    private var accumulated = Array<UInt8>()

    private var offset: Int = 0
    private var offsetToRemove: Int = 0
    private let chacha: ChaCha20

    init(chacha: ChaCha20) {
      self.chacha = chacha
    }

    public mutating func update(withBytes bytes: ArraySlice<UInt8>, isLast: Bool = true) throws -> Array<UInt8> {
      // prepend "offset" number of bytes at the beginning
      if self.offset > 0 {
        self.accumulated += Array<UInt8>(repeating: 0, count: self.offset) + bytes
        self.offsetToRemove = self.offset
        self.offset = 0
      } else {
        self.accumulated += bytes
      }

      var plaintext = Array<UInt8>()
      plaintext.reserveCapacity(self.accumulated.count)
      for chunk in self.accumulated.batched(by: ChaCha20.blockSize) {
        if isLast || self.accumulated.count >= ChaCha20.blockSize {
          plaintext += try self.chacha.decrypt(chunk)

          // remove "offset" from the beginning of first chunk
          if self.offsetToRemove > 0 {
            plaintext.removeFirst(self.offsetToRemove) // TODO: improve performance
            self.offsetToRemove = 0
          }

          self.accumulated.removeFirst(chunk.count)
        }
      }

      return plaintext
    }

    public func seek(to: Int) throws {
      throw Error.notSupported
    }
  }
}

// MARK: Cryptors

extension ChaCha20: Cryptors {
  //TODO: Use BlockEncryptor/BlockDecryptor

  public func makeEncryptor() -> Cryptor & Updatable {
    ChaCha20.ChaChaEncryptor(chacha: self)
  }

  public func makeDecryptor() -> Cryptor & Updatable {
    ChaCha20.ChaChaDecryptor(chacha: self)
  }
}