1 /*
  2  * Copyright (c) 2019, 2021, Oracle and/or its affiliates. All rights reserved.
  3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  4  *
  5  * This code is free software; you can redistribute it and/or modify it
  6  * under the terms of the GNU General Public License version 2 only, as
  7  * published by the Free Software Foundation.
  8  *
  9  * This code is distributed in the hope that it will be useful, but WITHOUT
 10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 12  * version 2 for more details (a copy is included in the LICENSE file that
 13  * accompanied this code).
 14  *
 15  * You should have received a copy of the GNU General Public License version
 16  * 2 along with this work; if not, write to the Free Software Foundation,
 17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 18  *
 19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 20  * or visit www.oracle.com if you need additional information or have any
 21  * questions.
 22  */
 23   
 24 package org.openjdk.bench.jdk.incubator.vector.crypto;
 25 
 26 import org.openjdk.jmh.annotations.*;
 27 import jdk.incubator.vector.*;
 28 
 29 import java.nio.ByteOrder;
 30 import java.util.Arrays;
 31 
 32 @State(Scope.Thread)
 33 @BenchmarkMode(Mode.Throughput)
 34 @Fork(value = 1, jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
 35 @Warmup(iterations = 3, time = 3)
 36 @Measurement(iterations = 8, time = 2)
 37 public class ChaChaBench {
 38 
 39     @Param({"16384", "65536"})
 40     private int dataSize;
 41     
 42     private ChaChaVector cc20_S128 = makeCC20(VectorShape.S_128_BIT);
 43     private ChaChaVector cc20_S256 = makeCC20(VectorShape.S_256_BIT);
 44     private ChaChaVector cc20_S512 = makeCC20(VectorShape.S_512_BIT);
 45  
 46     private byte[] in;
 47     private byte[] out;
 48     
 49     private byte[] key = new byte[32];
 50     private byte[] nonce = new byte[12];
 51     private long counter = 0;
 52 
 53     private static ChaChaVector makeCC20(VectorShape shape) {
 54         ChaChaVector cc20 = new ChaChaVector(shape);
 55         runKAT(cc20);
 56         return cc20;
 57     }
 58 
 59     @Setup
 60     public void setup() {
 61         
 62         in = new byte[dataSize];
 63         out = new byte[dataSize];
 64     }
 65 
 66     @Benchmark
 67     public void encrypt128() {
 68         cc20_S128.chacha20(key, nonce, counter, in, out);
 69     }
 70 
 71     @Benchmark
 72     public void encrypt256() {
 73         cc20_S256.chacha20(key, nonce, counter, in, out);
 74     }
 75 
 76     @Benchmark
 77     public void encrypt512() {
 78         cc20_S512.chacha20(key, nonce, counter, in, out);
 79     }
 80 
 81     private static class ChaChaVector {
 82 
 83         private static final int[] STATE_CONSTANTS =
 84             new int[]{0x61707865, 0x3320646e, 0x79622d32, 0x6b206574};
 85 
 86         private final VectorSpecies<Integer> intSpecies;
 87         private final int numBlocks;
 88 
 89         private final VectorShuffle<Integer> rot1;
 90         private final VectorShuffle<Integer> rot2;
 91         private final VectorShuffle<Integer> rot3;
 92 
 93         private final IntVector counterAdd;
 94 
 95         private final VectorShuffle<Integer> shuf0;
 96         private final VectorShuffle<Integer> shuf1;
 97         private final VectorShuffle<Integer> shuf2;
 98         private final VectorShuffle<Integer> shuf3;
 99 
100         private final VectorMask<Integer> mask0;
101         private final VectorMask<Integer> mask1;
102         private final VectorMask<Integer> mask2;
103         private final VectorMask<Integer> mask3;
104 
105         private final int[] state;
106 
107         public ChaChaVector(VectorShape shape) {
108             this.intSpecies = VectorSpecies.of(int.class, shape);
109             this.numBlocks = intSpecies.length() / 4;
110 
111             this.rot1 = makeRotate(1);
112             this.rot2 = makeRotate(2);
113             this.rot3 = makeRotate(3);
114 
115             this.counterAdd = makeCounterAdd();
116 
117             this.shuf0 = makeRearrangeShuffle(0);
118             this.shuf1 = makeRearrangeShuffle(1);
119             this.shuf2 = makeRearrangeShuffle(2);
120             this.shuf3 = makeRearrangeShuffle(3);
121 
122             this.mask0 = makeRearrangeMask(0);
123             this.mask1 = makeRearrangeMask(1);
124             this.mask2 = makeRearrangeMask(2);
125             this.mask3 = makeRearrangeMask(3);
126 
127             this.state = new int[numBlocks * 16];
128         }
129 
130         private VectorShuffle<Integer>  makeRotate(int amount) {
131             int[] shuffleArr = new int[intSpecies.length()];
132 
133             for (int i = 0; i < intSpecies.length(); i ++) {
134                 int offset = (i / 4) * 4;
135                 shuffleArr[i] = offset + ((i + amount) % 4);
136             }
137 
138             return VectorShuffle.fromValues(intSpecies, shuffleArr);
139         }
140 
141         private IntVector makeCounterAdd() {
142             int[] addArr = new int[intSpecies.length()];
143             for(int i = 0; i < numBlocks; i++) {
144                 addArr[4 * i] = numBlocks;
145             }
146             return IntVector.fromArray(intSpecies, addArr, 0);
147         }
148 
149         private VectorShuffle<Integer>  makeRearrangeShuffle(int order) {
150             int[] shuffleArr = new int[intSpecies.length()];
151             int start = order * 4;
152             for (int i = 0; i < shuffleArr.length; i++) {
153                 shuffleArr[i] = (i % 4) + start;
154             }
155             return VectorShuffle.fromArray(intSpecies, shuffleArr, 0);
156         }
157 
158         private VectorMask<Integer> makeRearrangeMask(int order) {
159             boolean[] maskArr = new boolean[intSpecies.length()];
160             int start = order * 4;
161             if (start < maskArr.length) {
162                 for (int i = 0; i < 4; i++) {
163                     maskArr[i + start] = true;
164                 }
165             }
166 
167             return VectorMask.fromValues(intSpecies, maskArr);
168         }
169 
170         public void makeState(byte[] key, byte[] nonce, long counter,
171             int[] out) {
172 
173             // first field is constants
174             for (int i = 0; i < 4; i++) {
175                 for (int j = 0; j < numBlocks; j++) {
176                     out[4*j + i] = STATE_CONSTANTS[i];
177                 }
178             }
179 
180             // second field is first part of key
181             int fieldStart = 4 * numBlocks;
182             for (int i = 0; i < 4; i++) {
183                 int keyInt = 0;
184                 for (int j = 0; j < 4; j++) {
185                     keyInt += (0xFF & key[4 * i + j]) << 8 * j;
186                 }
187                 for (int j = 0; j < numBlocks; j++) {
188                     out[fieldStart + j*4 + i] = keyInt;
189                 }
190             }
191 
192             // third field is second part of key
193             fieldStart = 8 * numBlocks;
194             for (int i = 0; i < 4; i++) {
195                 int keyInt = 0;
196                 for (int j = 0; j < 4; j++) {
197                     keyInt += (0xFF & key[4 * (i + 4) + j]) << 8 * j;
198                 }
199 
200                 for (int j = 0; j < numBlocks; j++) {
201                     out[fieldStart + j*4 + i] = keyInt;
202                 }
203             }
204 
205             // fourth field is counter and nonce
206             fieldStart = 12 * numBlocks;
207             for (int j = 0; j < numBlocks; j++) {
208                 out[fieldStart + j*4] = (int) (counter + j);
209             }
210 
211             for (int i = 0; i < 3; i++) {
212                 int nonceInt = 0;
213                 for (int j = 0; j < 4; j++) {
214                     nonceInt += (0xFF & nonce[4 * i + j]) << 8 * j;
215                 }
216 
217                 for (int j = 0; j < numBlocks; j++) {
218                     out[fieldStart + j*4 + 1 + i] = nonceInt;
219                 }
220             }
221         }
222 
223         public void chacha20(byte[] key, byte[] nonce, long counter,
224             byte[] in, byte[] out) {
225 
226             makeState(key, nonce, counter, state);
227 
228             int len = intSpecies.length();
229 
230             IntVector sa = IntVector.fromArray(intSpecies, state, 0);
231             IntVector sb = IntVector.fromArray(intSpecies, state, len);
232             IntVector sc = IntVector.fromArray(intSpecies, state, 2 * len);
233             IntVector sd = IntVector.fromArray(intSpecies, state, 3 * len);
234 
235             int stateLenBytes = state.length * 4;
236             int numStates = (in.length + stateLenBytes - 1) / stateLenBytes;
237             for (int j = 0; j < numStates; j++){
238 
239                 IntVector a = sa;
240                 IntVector b = sb;
241                 IntVector c = sc;
242                 IntVector d = sd;
243 
244                 for (int i = 0; i < 10; i++) {
245                     // first round
246                     a = a.add(b);
247                     d = d.lanewise(VectorOperators.XOR, a);
248                     d = d.lanewise(VectorOperators.ROL, 16);
249 
250                     c = c.add(d);
251                     b = b.lanewise(VectorOperators.XOR, c);
252                     b = b.lanewise(VectorOperators.ROL,12);
253 
254                     a = a.add(b);
255                     d = d.lanewise(VectorOperators.XOR, a);
256                     d = d.lanewise(VectorOperators.ROL,8);
257 
258                     c = c.add(d);
259                     b = b.lanewise(VectorOperators.XOR, c);
260                     b = b.lanewise(VectorOperators.ROL,7);
261 
262                     // makeRotate
263                     b = b.rearrange(rot1);
264                     c = c.rearrange(rot2);
265                     d = d.rearrange(rot3);
266 
267                     // second round
268                     a = a.add(b);
269                     d = d.lanewise(VectorOperators.XOR, a);
270                     d = d.lanewise(VectorOperators.ROL,16);
271 
272                     c = c.add(d);
273                     b = b.lanewise(VectorOperators.XOR, c);
274                     b = b.lanewise(VectorOperators.ROL,12);
275 
276                     a = a.add(b);
277                     d = d.lanewise(VectorOperators.XOR, a);
278                     d = d.lanewise(VectorOperators.ROL,8);
279 
280                     c = c.add(d);
281                     b = b.lanewise(VectorOperators.XOR, c);
282                     b = b.lanewise(VectorOperators.ROL,7);
283 
284                     // makeRotate
285                     b = b.rearrange(rot3);
286                     c = c.rearrange(rot2);
287                     d = d.rearrange(rot1);
288                 }
289 
290                 a = a.add(sa);
291                 b = b.add(sb);
292                 c = c.add(sc);
293                 d = d.add(sd);
294 
295                 // rearrange the vectors
296                 if (intSpecies.length() == 4) {
297                     // no rearrange needed
298                 } else if (intSpecies.length() == 8) {
299                     IntVector a_r =
300                             a.rearrange(shuf0).blend(b.rearrange(shuf0), mask1);
301                     IntVector b_r =
302                             c.rearrange(shuf0).blend(d.rearrange(shuf0), mask1);
303                     IntVector c_r =
304                             a.rearrange(shuf1).blend(b.rearrange(shuf1), mask1);
305                     IntVector d_r =
306                             c.rearrange(shuf1).blend(d.rearrange(shuf1), mask1);
307 
308                     a = a_r;
309                     b = b_r;
310                     c = c_r;
311                     d = d_r;
312                 } else if (intSpecies.length() == 16) {
313                     IntVector a_r = a;
314                     a_r = a_r.blend(b.rearrange(shuf0), mask1);
315                     a_r = a_r.blend(c.rearrange(shuf0), mask2);
316                     a_r = a_r.blend(d.rearrange(shuf0), mask3);
317 
318                     IntVector b_r = b;
319                     b_r = b_r.blend(a.rearrange(shuf1), mask0);
320                     b_r = b_r.blend(c.rearrange(shuf1), mask2);
321                     b_r = b_r.blend(d.rearrange(shuf1), mask3);
322 
323                     IntVector c_r = c;
324                     c_r = c_r.blend(a.rearrange(shuf2), mask0);
325                     c_r = c_r.blend(b.rearrange(shuf2), mask1);
326                     c_r = c_r.blend(d.rearrange(shuf2), mask3);
327 
328                     IntVector d_r = d;
329                     d_r = d_r.blend(a.rearrange(shuf3), mask0);
330                     d_r = d_r.blend(b.rearrange(shuf3), mask1);
331                     d_r = d_r.blend(c.rearrange(shuf3), mask2);
332 
333                     a = a_r;
334                     b = b_r;
335                     c = c_r;
336                     d = d_r;
337                 } else {
338                     throw new RuntimeException("not supported");
339                 }
340 
341                 // xor keystream with input
342                 int inOff = stateLenBytes * j;
343                 IntVector ina = IntVector.fromByteArray(intSpecies, in, inOff, ByteOrder.LITTLE_ENDIAN);
344                 IntVector inb = IntVector.fromByteArray(intSpecies, in, inOff + 4 * len, ByteOrder.LITTLE_ENDIAN);
345                 IntVector inc = IntVector.fromByteArray(intSpecies, in, inOff + 8 * len, ByteOrder.LITTLE_ENDIAN);
346                 IntVector ind = IntVector.fromByteArray(intSpecies, in, inOff + 12 * len, ByteOrder.LITTLE_ENDIAN);
347 
348                 ina.lanewise(VectorOperators.XOR, a).intoByteArray(out, inOff, ByteOrder.LITTLE_ENDIAN);
349                 inb.lanewise(VectorOperators.XOR, b).intoByteArray(out, inOff + 4 * len, ByteOrder.LITTLE_ENDIAN);
350                 inc.lanewise(VectorOperators.XOR, c).intoByteArray(out, inOff + 8 * len, ByteOrder.LITTLE_ENDIAN);
351                 ind.lanewise(VectorOperators.XOR, d).intoByteArray(out, inOff + 12 * len, ByteOrder.LITTLE_ENDIAN);
352 
353                 // increment counter
354                 sd = sd.add(counterAdd);
355             }
356         }
357 
358         public int implBlockSize() {
359             return numBlocks * 64;
360         }
361     }
362 
363     private static byte[] hexStringToByteArray(String str) {
364         byte[] result = new byte[str.length() / 2];
365         for (int i = 0; i < result.length; i++) {
366             result[i] = (byte) Character.digit(str.charAt(2 * i), 16);
367             result[i] <<= 4;
368             result[i] += Character.digit(str.charAt(2 * i + 1), 16);
369         }
370         return result;
371     }
372 
373     private static void runKAT(ChaChaVector cc20, String keyStr,
374         String nonceStr, long counter, String inStr, String outStr) {
375 
376         byte[] key = hexStringToByteArray(keyStr);
377         byte[] nonce = hexStringToByteArray(nonceStr);
378         byte[] in = hexStringToByteArray(inStr);
379         byte[] expOut = hexStringToByteArray(outStr);
380 
381         // implementation only works at multiples of some size
382         int blockSize = cc20.implBlockSize();
383 
384         int length = blockSize * ((in.length + blockSize - 1) / blockSize);
385         in = Arrays.copyOf(in, length);
386         byte[] out = new byte[length];
387 
388         cc20.chacha20(key, nonce, counter, in, out);
389 
390         byte[] actOut = new byte[expOut.length];
391         System.arraycopy(out, 0, actOut, 0, expOut.length);
392 
393         if (!Arrays.equals(out, 0, expOut.length, expOut, 0, expOut.length)) {
394             throw new RuntimeException("Incorrect result");
395         }
396     }
397 
398     /*
399      * ChaCha20 Known Answer Tests to ensure that the implementation is correct.
400      */
401     private static void runKAT(ChaChaVector cc20) {
402         runKAT(cc20,
403         "0000000000000000000000000000000000000000000000000000000000000001",
404         "000000000000000000000002",
405         1,
406         "416e79207375626d697373696f6e20746f20746865204945544620696e74656e" +
407         "6465642062792074686520436f6e7472696275746f7220666f72207075626c69" +
408         "636174696f6e20617320616c6c206f722070617274206f6620616e2049455446" +
409         "20496e7465726e65742d4472616674206f722052464320616e6420616e792073" +
410         "746174656d656e74206d6164652077697468696e2074686520636f6e74657874" +
411         "206f6620616e204945544620616374697669747920697320636f6e7369646572" +
412         "656420616e20224945544620436f6e747269627574696f6e222e205375636820" +
413         "73746174656d656e747320696e636c756465206f72616c2073746174656d656e" +
414         "747320696e20494554462073657373696f6e732c2061732077656c6c20617320" +
415         "7772697474656e20616e6420656c656374726f6e696320636f6d6d756e696361" +
416         "74696f6e73206d61646520617420616e792074696d65206f7220706c6163652c" +
417         "207768696368206172652061646472657373656420746f",
418         "a3fbf07df3fa2fde4f376ca23e82737041605d9f4f4f57bd8cff2c1d4b7955ec" +
419         "2a97948bd3722915c8f3d337f7d370050e9e96d647b7c39f56e031ca5eb6250d" +
420         "4042e02785ececfa4b4bb5e8ead0440e20b6e8db09d881a7c6132f420e527950" +
421         "42bdfa7773d8a9051447b3291ce1411c680465552aa6c405b7764d5e87bea85a" +
422         "d00f8449ed8f72d0d662ab052691ca66424bc86d2df80ea41f43abf937d3259d" +
423         "c4b2d0dfb48a6c9139ddd7f76966e928e635553ba76c5c879d7b35d49eb2e62b" +
424         "0871cdac638939e25e8a1e0ef9d5280fa8ca328b351c3c765989cbcf3daa8b6c" +
425         "cc3aaf9f3979c92b3720fc88dc95ed84a1be059c6499b9fda236e7e818b04b0b" +
426         "c39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f" +
427         "5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e6" +
428         "98ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab" +
429         "7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221"
430         );
431     }
432 }