CipherSink & CipherSource (#711)

* CipherSink * CipherSource * Simplify CipherSource * Simplify by requiring block cipher * Cipher tests with single byte updates Failing for cipher source * Fix wrong segment pop * Cipher tests with empty data Sink test is failing * Fix cipher test & simplify condition in close * Extract cipher sink doFinal method * Use random seeds * Add failing Cipher test for RSA * Revert "Add failing Cipher test for RSA" This reverts commit 52b8662. There's no point in trying to support it, and it would be very difficult since it's not a block cipher. * Parameterize cipher tests Test all block ciphers which are required to be implemented on JVM, as per the documentation of the `Cipher` class. * Minor optimization to Cipher stream doFinal If `getOutputSize` returns zero, exit early. * Standard file structure License, JVM file name, inline extension. * Javadoc * Extensions on `Cipher` for source/sink Classes are now private * Fix cipher source update size Noticed issue when simplifying test by using `ForwardingSource`, which changed segment usage. * Missing word in comment * Test padding * Remove useless visibility restriction Since classes are now private, there's no point in marking constructor as internal * Fix lint issue * Cryptography documentation * Fix lint check * Apply pull request feedback Co-authored-by: Martin Devillers <[email protected]>
square · Oct 4, 2020 · da3112e · da3112e
1 parent 6344cc1
commit da3112e
Show file tree

Hide file tree

Showing 6 changed files with 708 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -914,6 +914,53 @@ As with hashing, you can generate an HMAC from a `ByteString`, `Buffer`, `Hashin
 `HashingSink`. Note that Okio doesn’t implement HMAC for MD5. Okio uses Java’s
 `java.security.MessageDigest` for cryptographic hashes and `javax.crypto.Mac` for HMAC.
 
+### Encryption/Decryption
+
+Use `Okio.sink(Cipher, Sink)` or `Okio.source(Cipher, Source)` to
+encrypt or decrypt a stream using a block cipher.
+
+Callers are responsible for the initialization of the encryption or
+decryption cipher with the chosen algorithm, the key, and
+algorithm-specific additional parameters like the initialization vector.
+The following example shows a typical usage with AES encryption, in
+which `key` and `iv` parameters should both be 16 bytes long.
+
+```java
+void encryptAes(ByteString bytes, File file, byte[] key, byte[] iv) throws GeneralSecurityException, IOException {
+    Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
+    cipher.init(Cipher.ENCRYPT_MODE, new SecretKeySpec(key, "AES"), new IvParameterSpec(iv));
+    try (BufferedSink sink = Okio.buffer(Okio.sink(cipher, Okio.buffer(Okio.sink(file))))) {
+        sink.write(bytes);
+    }
+}
+
+ByteString decryptAesToByteString(File file, byte[] key, byte[] iv) throws GeneralSecurityException, IOException {
+    Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
+    cipher.init(Cipher.DECRYPT_MODE, new SecretKeySpec(key, "AES"), new IvParameterSpec(iv));
+    try (BufferedSource source = Okio.buffer(Okio.source(cipher, Okio.buffer(Okio.source(file))))) {
+        return source.readByteString();
+    }
+}
+```
+
+In Kotlin, these encryption/decryption methods are provided as
+extensions on the `Cipher` class.
+
+```kotlin
+fun encryptAes(bytes: ByteString, file: File, key: ByteArray, iv: ByteArray) {
+  val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding").apply {
+    init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
+  }
+  cipher.sink(file.sink().buffer()).buffer().use { it.write(bytes) }
+}
+
+fun decryptAesToByteString(file: File, key: ByteArray, iv: ByteArray): ByteString {
+  val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding").apply {
+    init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
+  }
+  return cipher.source(file.source().buffer()).buffer().use(BufferedSource::readByteString)
+}
+```
 
 Releases
 --------

diff --git a/okio/src/jvmMain/kotlin/okio/CipherSink.kt b/okio/src/jvmMain/kotlin/okio/CipherSink.kt
@@ -0,0 +1,139 @@
+/*
+ * Copyright (C) 2020 Square, 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.
+ */
+
+@file:JvmMultifileClass
+@file:JvmName("Okio")
+
+package okio
+
+import java.io.IOException
+import javax.crypto.Cipher
+
+private class CipherSink(
+  private val sink: BufferedSink,
+  private val cipher: Cipher
+) : Sink {
+
+  private val blockSize = cipher.blockSize
+  private var closed = false
+
+  init {
+    // Require block cipher, and check for unsupported (too large) block size (should never happen with standard algorithms)
+    require(blockSize > 0) { "Block cipher required $cipher" }
+    require(blockSize <= Segment.SIZE) { "Cipher block size $blockSize too large $cipher" }
+  }
+
+  @Throws(IOException::class)
+  override fun write(source: Buffer, byteCount: Long) {
+    checkOffsetAndCount(source.size, 0, byteCount)
+    check(!closed) { "closed" }
+
+    var remaining = byteCount
+    while (remaining > 0) {
+      val size = update(source, remaining)
+      remaining -= size
+    }
+  }
+
+  private fun update(source: Buffer, remaining: Long): Int {
+    val head = source.head!!
+    val size = minOf(remaining, head.limit - head.pos).toInt()
+    val buffer = sink.buffer
+
+    // For block cipher, output size cannot exceed input size in update
+    val s = buffer.writableSegment(size)
+
+    val ciphered = cipher.update(head.data, head.pos, size, s.data, s.limit)
+
+    s.limit += ciphered
+    buffer.size += ciphered
+
+    if (s.pos == s.limit) {
+      // We allocated a tail segment, but didn't end up needing it. Recycle!
+      buffer.head = s.pop()
+      SegmentPool.recycle(s)
+    }
+
+    // Mark those bytes as read.
+    source.size -= size
+    head.pos += size
+
+    if (head.pos == head.limit) {
+      source.head = head.pop()
+      SegmentPool.recycle(head)
+    }
+    return size
+  }
+
+  override fun flush() =
+    sink.flush()
+
+  override fun timeout() =
+    sink.timeout()
+
+  @Throws(IOException::class)
+  override fun close() {
+    if (closed) return
+    closed = true
+
+    var thrown = doFinal()
+
+    try {
+      sink.close()
+    } catch (e: Throwable) {
+      if (thrown == null) thrown = e
+    }
+
+    if (thrown != null) throw thrown
+  }
+
+  private fun doFinal(): Throwable? {
+    val outputSize = cipher.getOutputSize(0)
+    if (outputSize == 0) return null
+
+    var thrown: Throwable? = null
+    val buffer = sink.buffer
+
+    // For block cipher, output size cannot exceed block size in doFinal
+    val s = buffer.writableSegment(outputSize)
+
+    try {
+      val ciphered = cipher.doFinal(s.data, s.limit)
+
+      s.limit += ciphered
+      buffer.size += ciphered
+    } catch (e: Throwable) {
+      thrown = e
+    }
+
+    if (s.pos == s.limit) {
+      buffer.head = s.pop()
+      SegmentPool.recycle(s)
+    }
+
+    return thrown
+  }
+}
+
+/**
+ * Returns a [Sink] that processes data using this [Cipher] while writing to
+ * [sink].
+ *
+ * @throws IllegalArgumentException
+ *  If this isn't a block cipher.
+ */
+fun Cipher.sink(sink: BufferedSink): Sink =
+  CipherSink(sink, this)
diff --git a/okio/src/jvmMain/kotlin/okio/CipherSource.kt b/okio/src/jvmMain/kotlin/okio/CipherSource.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2020 Square, 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.
+ */
+
+@file:JvmMultifileClass
+@file:JvmName("Okio")
+
+package okio
+
+import java.io.IOException
+import javax.crypto.Cipher
+
+private class CipherSource(
+  private val source: BufferedSource,
+  private val cipher: Cipher
+) : Source {
+
+  private val blockSize = cipher.blockSize
+  private val buffer = Buffer()
+  private var final = false
+  private var closed = false
+
+  init {
+    // Require block cipher, and check for unsupported (too large) block size (should never happen with standard algorithms)
+    require(blockSize > 0) { "Block cipher required $cipher" }
+    require(blockSize <= Segment.SIZE) { "Cipher block size $blockSize too large $cipher" }
+  }
+
+  @Throws(IOException::class)
+  override fun read(sink: Buffer, byteCount: Long): Long {
+    require(byteCount >= 0) { "byteCount < 0: $byteCount" }
+    check(!closed) { "closed" }
+    if (byteCount == 0L) return 0
+    if (final) return buffer.read(sink, byteCount)
+
+    refill()
+
+    return buffer.read(sink, byteCount)
+  }
+
+  private fun refill() {
+    while (buffer.size == 0L) {
+      if (source.exhausted()) {
+        final = true
+        doFinal()
+        break
+      } else {
+        update()
+      }
+    }
+  }
+
+  private fun update() {
+    val head = source.buffer.head!!
+    val size = head.limit - head.pos
+
+    // For block cipher, output size cannot exceed input size in update
+    val s = buffer.writableSegment(size)
+
+    val ciphered =
+      cipher.update(head.data, head.pos, size, s.data, s.pos)
+
+    source.skip(size.toLong())
+
+    s.limit += ciphered
+    buffer.size += ciphered
+
+    if (s.pos == s.limit) {
+      // We allocated a tail segment, but didn't end up needing it. Recycle!
+      buffer.head = s.pop()
+      SegmentPool.recycle(s)
+    }
+  }
+
+  private fun doFinal() {
+    val outputSize = cipher.getOutputSize(0)
+    if (outputSize == 0) return
+
+    // For block cipher, output size cannot exceed block size in doFinal
+    val s = buffer.writableSegment(outputSize)
+
+    val ciphered = cipher.doFinal(s.data, s.pos)
+
+    s.limit += ciphered
+    buffer.size += ciphered
+
+    if (s.pos == s.limit) {
+      // We allocated a tail segment, but didn't end up needing it. Recycle!
+      buffer.head = s.pop()
+      SegmentPool.recycle(s)
+    }
+  }
+
+  override fun timeout() =
+    source.timeout()
+
+  @Throws(IOException::class)
+  override fun close() {
+    closed = true
+    source.close()
+  }
+}
+
+/**
+ * Returns a [Source] that processes data using this [Cipher] while reading
+ * from [source].
+ *
+ * @throws IllegalArgumentException
+ *  If this isn't a block cipher.
+ */
+fun Cipher.source(source: BufferedSource): Source =
+  CipherSource(source, this)
diff --git a/okio/src/jvmTest/kotlin/okio/CipherFactory.kt b/okio/src/jvmTest/kotlin/okio/CipherFactory.kt
@@ -0,0 +1,69 @@
+package okio
+
+import javax.crypto.Cipher
+import javax.crypto.spec.IvParameterSpec
+import javax.crypto.spec.SecretKeySpec
+import kotlin.random.Random
+
+class CipherFactory(
+  private val transformation: String,
+  private val init: Cipher.(mode: Int) -> Unit
+) {
+  val blockSize
+    get() = cipher.blockSize
+
+  val encrypt: Cipher
+    get() = create(Cipher.ENCRYPT_MODE)
+
+  val decrypt: Cipher
+    get() = create(Cipher.DECRYPT_MODE)
+
+  private val cipher: Cipher
+    get() = Cipher.getInstance(transformation)
+
+  private fun create(mode: Int): Cipher =
+    cipher.apply { init(mode) }
+}
+
+data class CipherAlgorithm(
+  val transformation: String,
+  val padding: Boolean,
+  val keyLength: Int,
+  val ivLength: Int? = null
+) {
+  fun createCipherFactory(random: Random): CipherFactory {
+    val key = random.nextBytes(keyLength)
+    val secretKeySpec = SecretKeySpec(key, transformation.substringBefore('/'))
+    return if (ivLength == null) {
+      CipherFactory(transformation) { mode ->
+        init(mode, secretKeySpec)
+      }
+    } else {
+      val iv = random.nextBytes(ivLength)
+      val ivParameterSpec = IvParameterSpec(iv)
+      CipherFactory(transformation) { mode ->
+        init(mode, secretKeySpec, ivParameterSpec)
+      }
+    }
+  }
+
+  override fun toString(): String =
+    transformation
+
+  companion object {
+    fun getBlockCipherAlgorithms() = listOf(
+      CipherAlgorithm("AES/CBC/NoPadding", false, 16, 16),
+      CipherAlgorithm("AES/CBC/PKCS5Padding", true, 16, 16),
+      CipherAlgorithm("AES/ECB/NoPadding", false, 16),
+      CipherAlgorithm("AES/ECB/PKCS5Padding", true, 16),
+      CipherAlgorithm("DES/CBC/NoPadding", false, 8, 8),
+      CipherAlgorithm("DES/CBC/PKCS5Padding", true, 8, 8),
+      CipherAlgorithm("DES/ECB/NoPadding", false, 8),
+      CipherAlgorithm("DES/ECB/PKCS5Padding", true, 8),
+      CipherAlgorithm("DESede/CBC/NoPadding", false, 24, 8),
+      CipherAlgorithm("DESede/CBC/PKCS5Padding", true, 24, 8),
+      CipherAlgorithm("DESede/ECB/NoPadding", false, 24),
+      CipherAlgorithm("DESede/ECB/PKCS5Padding", true, 24)
+    )
+  }
+}