1 /* 2 * Copyright (C) 2020 Square, Inc. and others. 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 package okio 17 18 import java.io.IOException 19 import javax.crypto.Cipher 20 21 class CipherSource( 22 private val source: BufferedSource, 23 val cipher: Cipher, 24 ) : Source { 25 private val blockSize = cipher.blockSize 26 private val buffer = Buffer() 27 private var final = false 28 private var closed = false 29 30 init { 31 // Require block cipher <lambda>null32 require(blockSize > 0) { "Block cipher required $cipher" } 33 } 34 35 @Throws(IOException::class) readnull36 override fun read(sink: Buffer, byteCount: Long): Long { 37 require(byteCount >= 0L) { "byteCount < 0: $byteCount" } 38 check(!closed) { "closed" } 39 if (byteCount == 0L) return 0L 40 41 refill() 42 43 return buffer.read(sink, byteCount) 44 } 45 refillnull46 private fun refill() { 47 while (buffer.size == 0L && !final) { 48 if (source.exhausted()) { 49 final = true 50 doFinal() 51 break 52 } else { 53 update() 54 } 55 } 56 } 57 updatenull58 private fun update() { 59 val head = source.buffer.head!! 60 var size = head.limit - head.pos 61 62 // Shorten input until output is guaranteed to fit within a segment 63 var outputSize = cipher.getOutputSize(size) 64 while (outputSize > Segment.SIZE) { 65 if (size <= blockSize) { 66 // Bug: For AES-GCM on Android `update` method never outputs any data 67 // As a consequence, `getOutputSize` just keeps increasing indefinitely after each update 68 // When that happens, the fallback is to break the streaming implementation and just cipher the rest of the data all at once 69 final = true 70 buffer.write(cipher.doFinal(source.readByteArray())) 71 return 72 } 73 size -= blockSize 74 outputSize = cipher.getOutputSize(size) 75 } 76 val s = buffer.writableSegment(outputSize) 77 78 val ciphered = 79 cipher.update(head.data, head.pos, size, s.data, s.pos) 80 81 source.skip(size.toLong()) 82 83 s.limit += ciphered 84 buffer.size += ciphered 85 86 // We allocated a tail segment, but didn't end up needing it. Recycle! 87 if (s.pos == s.limit) { 88 buffer.head = s.pop() 89 SegmentPool.recycle(s) 90 } 91 } 92 doFinalnull93 private fun doFinal() { 94 val outputSize = cipher.getOutputSize(0) 95 if (outputSize == 0) return 96 97 // For block cipher, output size cannot exceed block size in doFinal. 98 val s = buffer.writableSegment(outputSize) 99 100 val ciphered = cipher.doFinal(s.data, s.pos) 101 102 s.limit += ciphered 103 buffer.size += ciphered 104 105 // We allocated a tail segment, but didn't end up needing it. Recycle! 106 if (s.pos == s.limit) { 107 buffer.head = s.pop() 108 SegmentPool.recycle(s) 109 } 110 } 111 timeoutnull112 override fun timeout() = source.timeout() 113 114 @Throws(IOException::class) 115 override fun close() { 116 closed = true 117 source.close() 118 } 119 } 120