1/*
2 * Copyright (c) 2004, 2007, 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/*
25 * @test
26 * @bug 4979126
27 * @summary Ensure update()/doFinal() matches javadoc description
28 * when ShortBufferException is thrown.
29 * @author Valerie Peng
30 */
31import java.util.Arrays;
32import java.security.*;
33import java.security.spec.*;
34import java.math.BigInteger;
35import javax.crypto.*;
36import javax.crypto.spec.*;
37import java.security.Provider;
38
39public class TestShortBuffer {
40    private static final String ALGO = "AES";
41    private static final String[] MODES = {
42        "ECB", "CBC", "PCBC", "CFB16", "OFB8"
43    };
44    private static final SecretKey KEY =
45        new SecretKeySpec(new byte[16], ALGO);
46    private static byte[] SHORTBUFFER = new byte[1];
47
48    private static final byte[] PLAINTEXT  = new byte[30];
49    static {
50        PLAINTEXT[0] = (byte)0x15;
51    };
52    private Cipher ci = null;
53    private byte[] in = null;
54    private byte[] expected = null;
55    private byte[] out = null;
56    private int outOffset = 0;
57
58    private TestShortBuffer(Cipher ci) {
59        this.ci = ci;
60    }
61
62    private void init(byte[] in, byte[] expected) {
63        this.in = (byte[]) in.clone();
64        this.expected = (byte[]) expected.clone();
65        this.out = new byte[expected.length];
66        this.outOffset = 0;
67   }
68
69    private static void runTest() throws Exception {
70        // Initialization
71        for (int i = 0; i < MODES.length; i++) {
72            System.out.println("===== TESTING MODE " + MODES[i] + " =====");
73            Cipher ci = Cipher.getInstance(ALGO+"/"+MODES[i]+"/PKCS5Padding",
74                                           "SunJCE");
75            TestShortBuffer test = null;
76            int stored = 0;
77            AlgorithmParameters params = null;
78            byte[] cipherText = null;
79            byte[] shortBuffer = new byte[8];
80            for (int k = 0; k < 2; k++) {
81                byte[] expected = null;
82                switch (k) {
83                case 0: // Encryption
84                    System.out.println("Testing with Cipher.ENCRYPT_MODE");
85                    ci.init(Cipher.ENCRYPT_MODE, KEY);
86                    cipherText = ci.doFinal(PLAINTEXT);
87                    test = new TestShortBuffer(ci);
88                    test.init(PLAINTEXT, cipherText);
89                    params = ci.getParameters();
90                    break;
91                case 1: // Decryption
92                    System.out.println("Testing with Cipher.DECRYPT_MODE");
93                    ci.init(Cipher.DECRYPT_MODE, KEY, params);
94                    test = new TestShortBuffer(ci);
95                    test.init(cipherText, PLAINTEXT);
96                    break;
97                }
98                int offset = 2 + i*5;
99                test.testUpdate();
100                test.testUpdateWithUpdate(offset);
101                test.testDoFinal();
102                test.testDoFinalWithUpdate(offset);
103            }
104        }
105    }
106
107    private void checkOutput() throws Exception {
108        if (!Arrays.equals(out, expected)) {
109            System.out.println("got: " + new BigInteger(out));
110            System.out.println("expect: " + new BigInteger(expected));
111            throw new Exception("Generated different outputs");
112        }
113    }
114    private void testUpdate() throws Exception {
115        outOffset = 0;
116        int stored = 0;
117        try {
118            stored = ci.update(in, 0, in.length, SHORTBUFFER);
119            throw new Exception("Should throw ShortBufferException!");
120        } catch (ShortBufferException sbe) {
121            System.out.println("Expected SBE thrown....");
122            // retry with a large-enough buffer according to javadoc
123            stored = ci.update(in, 0, in.length, out);
124            stored = ci.doFinal(out, outOffset += stored);
125            if (out.length != (stored + outOffset)) {
126                throw new Exception("Wrong number of output bytes");
127            }
128        }
129        checkOutput();
130    }
131    private void testUpdateWithUpdate(int offset) throws Exception {
132        outOffset = 0;
133        int stored = 0;
134        byte[] out1 = ci.update(in, 0, offset);
135        if (out1 != null) {
136            System.arraycopy(out1, 0, out, 0, out1.length);
137            outOffset += out1.length;
138        }
139        try {
140            stored = ci.update(in, offset, in.length-offset, SHORTBUFFER);
141            throw new Exception("Should throw ShortBufferException!");
142        } catch (ShortBufferException sbe) {
143            System.out.println("Expected SBE thrown....");
144            // retry with a large-enough buffer according to javadoc
145            stored = ci.update(in, offset, in.length-offset,
146                               out, outOffset);
147            stored = ci.doFinal(out, outOffset+=stored);
148            if (out.length != (stored + outOffset)) {
149                throw new Exception("Wrong number of output bytes");
150            }
151        }
152        checkOutput();
153    }
154    private void testDoFinal() throws Exception {
155        outOffset = 0;
156        int stored = 0;
157        try {
158            stored = ci.doFinal(in, 0, in.length, SHORTBUFFER);
159            throw new Exception("Should throw ShortBufferException!");
160        } catch (ShortBufferException sbe) {
161            System.out.println("Expected SBE thrown....");
162            // retry with a large-enough buffer according to javadoc
163            stored = ci.doFinal(in, 0, in.length, out, 0);
164            if (out.length != stored) {
165                throw new Exception("Wrong number of output bytes");
166            }
167        }
168        checkOutput();
169    }
170
171    private void testDoFinalWithUpdate(int offset) throws Exception {
172        outOffset = 0;
173        int stored = 0;
174        byte[] out1 = ci.update(in, 0, offset);
175        if (out1 != null) {
176            System.arraycopy(out1, 0, out, 0, out1.length);
177            outOffset += out1.length;
178        }
179        try {
180            stored = ci.doFinal(in, offset, in.length-offset, SHORTBUFFER);
181            throw new Exception("Should throw ShortBufferException!");
182        } catch (ShortBufferException sbe) {
183            System.out.println("Expected SBE thrown....");
184            // retry with a large-enough buffer according to javadoc
185            stored = ci.doFinal(in, offset, in.length-offset,
186                                out, outOffset);
187            if (out.length != (stored + outOffset)) {
188                throw new Exception("Wrong number of output bytes");
189            }
190        }
191        checkOutput();
192    }
193    public static void main(String[] argv) throws Exception {
194        runTest();
195        System.out.println("Test Passed");
196    }
197}
198