1 /*
2  * Copyright (C) 2016 The Android Open Source Project
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 com.google.android.exoplayer2.upstream.crypto;
17 
18 import com.google.android.exoplayer2.util.Assertions;
19 import com.google.android.exoplayer2.util.Util;
20 import java.nio.ByteBuffer;
21 import java.security.InvalidAlgorithmParameterException;
22 import java.security.InvalidKeyException;
23 import java.security.NoSuchAlgorithmException;
24 import javax.crypto.Cipher;
25 import javax.crypto.NoSuchPaddingException;
26 import javax.crypto.ShortBufferException;
27 import javax.crypto.spec.IvParameterSpec;
28 import javax.crypto.spec.SecretKeySpec;
29 
30 /**
31  * A flushing variant of a AES/CTR/NoPadding {@link Cipher}.
32  *
33  * Unlike a regular {@link Cipher}, the update methods of this class are guaranteed to process all
34  * of the bytes input (and hence output the same number of bytes).
35  */
36 public final class AesFlushingCipher {
37 
38   private final Cipher cipher;
39   private final int blockSize;
40   private final byte[] zerosBlock;
41   private final byte[] flushedBlock;
42 
43   private int pendingXorBytes;
44 
AesFlushingCipher(int mode, byte[] secretKey, long nonce, long offset)45   public AesFlushingCipher(int mode, byte[] secretKey, long nonce, long offset) {
46     try {
47       cipher = Cipher.getInstance("AES/CTR/NoPadding");
48       blockSize = cipher.getBlockSize();
49       zerosBlock = new byte[blockSize];
50       flushedBlock = new byte[blockSize];
51       long counter = offset / blockSize;
52       int startPadding = (int) (offset % blockSize);
53       cipher.init(
54           mode,
55           new SecretKeySpec(secretKey, Util.splitAtFirst(cipher.getAlgorithm(), "/")[0]),
56           new IvParameterSpec(getInitializationVector(nonce, counter)));
57       if (startPadding != 0) {
58         updateInPlace(new byte[startPadding], 0, startPadding);
59       }
60     } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException
61         | InvalidAlgorithmParameterException e) {
62       // Should never happen.
63       throw new RuntimeException(e);
64     }
65   }
66 
updateInPlace(byte[] data, int offset, int length)67   public void updateInPlace(byte[] data, int offset, int length) {
68     update(data, offset, length, data, offset);
69   }
70 
update(byte[] in, int inOffset, int length, byte[] out, int outOffset)71   public void update(byte[] in, int inOffset, int length, byte[] out, int outOffset) {
72     // If we previously flushed the cipher by inputting zeros up to a block boundary, then we need
73     // to manually transform the data that actually ended the block. See the comment below for more
74     // details.
75     while (pendingXorBytes > 0) {
76       out[outOffset] = (byte) (in[inOffset] ^ flushedBlock[blockSize - pendingXorBytes]);
77       outOffset++;
78       inOffset++;
79       pendingXorBytes--;
80       length--;
81       if (length == 0) {
82         return;
83       }
84     }
85 
86     // Do the bulk of the update.
87     int written = nonFlushingUpdate(in, inOffset, length, out, outOffset);
88     if (length == written) {
89       return;
90     }
91 
92     // We need to finish the block to flush out the remaining bytes. We do so by inputting zeros,
93     // so that the corresponding bytes output by the cipher are those that would have been XORed
94     // against the real end-of-block data to transform it. We store these bytes so that we can
95     // perform the transformation manually in the case of a subsequent call to this method with
96     // the real data.
97     int bytesToFlush = length - written;
98     Assertions.checkState(bytesToFlush < blockSize);
99     outOffset += written;
100     pendingXorBytes = blockSize - bytesToFlush;
101     written = nonFlushingUpdate(zerosBlock, 0, pendingXorBytes, flushedBlock, 0);
102     Assertions.checkState(written == blockSize);
103     // The first part of xorBytes contains the flushed data, which we copy out. The remainder
104     // contains the bytes that will be needed for manual transformation in a subsequent call.
105     for (int i = 0; i < bytesToFlush; i++) {
106       out[outOffset++] = flushedBlock[i];
107     }
108   }
109 
110   private int nonFlushingUpdate(byte[] in, int inOffset, int length, byte[] out, int outOffset) {
111     try {
112       return cipher.update(in, inOffset, length, out, outOffset);
113     } catch (ShortBufferException e) {
114       // Should never happen.
115       throw new RuntimeException(e);
116     }
117   }
118 
119   private byte[] getInitializationVector(long nonce, long counter) {
120     return ByteBuffer.allocate(16).putLong(nonce).putLong(counter).array();
121   }
122 
123 }
124