/*
 * Copyright (c) 2017, Red Hat, Inc. and/or its affiliates.
 * 
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

/*
 * @test
 * @bug 8148421
 * @compile --add-exports java.base/sun.security.util=ALL-UNNAMED --add-exports java.base/sun.security.ssl=ALL-UNNAMED --add-exports java.base/java.security=ALL-UNNAMED ExtendedMasterSecretTest.java
 * @run main/othervm ExtendedMasterSecretTest
 * @summary test for Extended Master Secret TLS extension.
 * @author Martin Balao (mbalao@redhat.com)
 * 
 */

import java.io.ByteArrayInputStream;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Base64;
import java.security.KeyStore;
import java.security.ProviderException;
import java.security.SecureRandom;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLEngine;
import javax.net.ssl.SSLEngineResult;
import javax.net.ssl.SSLParameters;
import javax.net.ssl.SSLEngineResult.Status;

import sun.security.ssl.SSLEngineImpl;
import sun.security.util.HexDumpEncoder;

import java.lang.reflect.Field;
import java.lang.reflect.Method;

import javax.net.ssl.SSLSession;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import javax.net.ssl.SSLEngineResult.HandshakeStatus;

import java.security.SecureRandomSpi;

final public class ExtendedMasterSecretTest {

    // Test data
    
    private static final byte[] expectedMasterKey = { 
            (byte)0xC3, (byte)0x77, (byte)0x02, (byte)0x3C, (byte)0xA4, (byte)0x81, (byte)0xBE,
            (byte)0x1E, (byte)0xE7, (byte)0x33, (byte)0x81, (byte)0x3A, (byte)0x7E, (byte)0xE2,
            (byte)0x14, (byte)0x21, (byte)0x4C, (byte)0xFF, (byte)0x45, (byte)0xE8, (byte)0xAD,
            (byte)0x27, (byte)0xCD, (byte)0x8B, (byte)0xE2, (byte)0x13, (byte)0x1F, (byte)0xA9,
            (byte)0x67, (byte)0x06, (byte)0xEF, (byte)0x39, (byte)0xA1, (byte)0x61, (byte)0xF4,
            (byte)0x56, (byte)0xE0, (byte)0xDA, (byte)0x7C, (byte)0x2E, (byte)0xE6, (byte)0xD2,
            (byte)0xD0, (byte)0xF5, (byte)0xC6, (byte)0x15, (byte)0xB5, (byte)0x61 
            };
    
    private static final String ksString =
            "/u3+7QAAAAIAAAACAAAAAQAJcm9vdF9jYV8yAAABXULz2U0AAADIMIHFMA4GCisGAQQBKgIRAQEF" +
            "AASBstjSsr3gdM2IyxLq/vU/RtyX6wLCgeZRyX4xqGMdYoq/A17SmEwY/s+hDf+rTJsDs6tUszlo" +
            "4AnHeh8Q6jJWAUQ4di+w2MXGYYxPG0/xrWRzdiH/SULr+L8DVswjZWxPtMt1jLs6gwycdkXTnBBf" +
            "iOAISqZjZysggipSBd7ZhWX+wEcUXxjYSKKvItipm23BfWBiWiInCRODR0VGsdwKe9sQoh8ru0CS" +
            "xo4KCwl1ijSA67sAAAABAAVYLjUwOQAAAgowggIGMIIBrAIJAKuUAhXzPEFoMAoGCCqGSM49BAMC" +
            "MIGKMQswCQYDVQQGEwJBUjEVMBMGA1UECAwMQnVlbm9zIEFpcmVzMRUwEwYDVQQHDAxCdWVub3Mg" +
            "QWlyZXMxIjAgBgNVBAoMGVJlZCBIYXQgZGUgQXJnZW50aW5hIFMuQS4xEjAQBgNVBAsMCUphdmEg" +
            "VGVhbTEVMBMGA1UEAwwMUmVkIEhhdCBJbmMuMB4XDTE3MDcxNDIxMDg0NVoXDTM3MDcwOTIxMDg0" +
            "NVowgYoxCzAJBgNVBAYTAkFSMRUwEwYDVQQIDAxCdWVub3MgQWlyZXMxFTATBgNVBAcMDEJ1ZW5v" +
            "cyBBaXJlczEiMCAGA1UECgwZUmVkIEhhdCBkZSBBcmdlbnRpbmEgUy5BLjESMBAGA1UECwwJSmF2" +
            "YSBUZWFtMRUwEwYDVQQDDAxSZWQgSGF0IEluYy4wWTATBgcqhkjOPQIBBggqhkjOPQMBBwNCAARn" +
            "Q7W9EO+wXaTh7YuirlgxRLzrlKhuIvrvhr84eC+xr1qEpXrfu9rPJTw763gMzBpf1sCjmsh6rNGL" +
            "/bBg45DDMAoGCCqGSM49BAMCA0gAMEUCIBkBUEvOgZpU1JMeJL6PUJCQgPtpYpOF7FEeAQ0Gwzbd" +
            "AiEAzKqP1B5ThWERdvtAzzH+jB1kAIyzyI1HLA9IUtHFwH8AAAABAAlyb290X2NhXzEAAAFdQvOt" +
            "8wAABQEwggT9MA4GCisGAQQBKgIRAQEFAASCBOmTVlw8D+80VRcPEXD0F3Nw1CMKETr8oLtQ7TVC" +
            "jNsiGjvRaR3GOiMFQ4wlKnQdytyc553jfaDYM3VaLqLquAUeTspVdzbIMr8V/zQTfIZeuyl4LqhA" +
            "u+ohxOr+cFFj+V1bikm2bpHdlDcXiYc3cLWgjRmBOuWKrlDoBig7eWnSw4S6hqK2sGGGNeLtYbl0" +
            "Af1q/IvzuOMrwHHV81fFiCmwpfc2pi4wbkxIVjZt0fRSZ3lQc+Mtc/zmJOl6GSGIzMFwtX4IPCgw" +
            "B5b4DOLcCXdms653a7XWHO0dv5WvUE3JATLrijgEd/WVgGDpFU+/AbUAXmj56UDd3e0K5KFVuJ04" +
            "tuol1H04UN0nsX3uUwvNpFdWK/UvvO3kAaFcfal8+g8/LZzA0FVDC3WE8rLGUMMP/KRJgcLstJLc" +
            "HDHEAOYmznpMmw1bxL3PDqOHX7C2PRHrF7y+/0+wiJGWAz0rg4gzpeZ0y1U57goxC/iGwp67AeKw" +
            "jG1ih6tIWa9nWsIJqgM33Hw1yEc59WbqGakiO85JgGyDkub+jFB0JxUyMyLjMz/ICGRcLR4h03vo" +
            "qGiZ/yEI4y+9SZktIUWGq4pKO8WqXIbhhmVindxxFXpgkytifpH1hNWNol50e35bXEflLp6iSfov" +
            "pZMVkfAQMUzUeVVrlIqtfmB4gSAQQC8ykE5dt8r/5c55JsxIZbydR8a8U9b3tPkGjKJef2o+qlpW" +
            "ms8iSYa7sgsks4IAHH/bKIAzb1JERu55F0q7beA4qPyMkasAh5chBF+EaESbs/3rVZsOPz7ayqbo" +
            "1vpyuiHLFChargcrdDwTt3RQ9KvcWSVuFuy6Q/wQjlOC2Wfi3kPqK3tzgPemoiTKj9nlzL6qFihF" +
            "cp69wbMex6pJPBK/jAhBIosxYxcohEt919Tn7IVZq2ihPabYJwlzTEF5TdWrM+kceB6ZD0XBgbJh" +
            "UU/0z/apb6djIDh68F2UD0BT/fq6Yka3RjC47+2n7h3puSurqy/lzPM09ZXF9Dba0xan+dnjZpsT" +
            "d5pZc/AI+wJs1/QffdIhBwfMZ7/3UcSxTs2oI4pRKPHPqGx7ESb80C2zHfAMflAp/bb0ghU3WOxk" +
            "wgA06ZJNvTmz3QmHTUhbiuW2SKPq7vd/1e17ynivsvAi4/7d4M46DzFgZ4hCVcMTyBbYjXe49uKJ" +
            "EMgCwZDdtPaoG1xq0DcNtZXLalXN3jRaalN3FTnWPqKknhnNixyi3FyTLWr3T8tvpRY3NCCFXkTF" +
            "xPyeAji0eYbmxMJ7pbA8OSVg2C3PVcMF5kMWxESyhjvWCmhti6vEUgcsG435OfivMFMB4+bdaE2Y" +
            "69kzqUMkPfPhA0lFPyquFYxbhGZHpKlm3ZoLCgqs1lOuKL4oXcMLOU0WF0CjhuRMbN8T4BP/LxT3" +
            "BoDlCn20MHdUmWo0QtWuLCAeYI0sw2zA4CAUyPhvTw1CqXDRruHEwSgwF+FMoUEAPHtFSG1gj5FU" +
            "mqGy3bG1PrSEz6sYQf/PSImywg/4oCyKLjhifpadvxtkmnp6g7yEqaii4t6p6yGfooDlErpAs3fL" +
            "mz6wn5eYFtcgwXjI74yyx/RxhW4uuqPFWN36r+MatCvg+TcPhDRh3EsW6WbzJ75WP9fmeP/jbuAm" +
            "yHr5tEY2xT/HF2CqLdIDO0jGcv5utgiCgAnNekwdgZEAAAABAAVYLjUwOQAAA+kwggPlMIICzaAD" +
            "AgECAgkA80Wa7UwYKzQwDQYJKoZIhvcNAQELBQAwgYgxCzAJBgNVBAYTAkFSMRUwEwYDVQQIDAxC" +
            "dWVub3MgQWlyZXMxFTATBgNVBAcMDEJ1ZW5vcyBBaXJlczEiMCAGA1UECgwZUmVkIEhhdCBkZSBB" +
            "cmdlbnRpbmEgUy5BLjESMBAGA1UECwwJSmF2YSBUZWFtMRMwEQYDVQQDDApyZWRoYXQuY29tMB4X" +
            "DTE3MDcxNDIxMDc1NFoXDTM3MDcwOTIxMDc1NFowgYgxCzAJBgNVBAYTAkFSMRUwEwYDVQQIDAxC" +
            "dWVub3MgQWlyZXMxFTATBgNVBAcMDEJ1ZW5vcyBBaXJlczEiMCAGA1UECgwZUmVkIEhhdCBkZSBB" +
            "cmdlbnRpbmEgUy5BLjESMBAGA1UECwwJSmF2YSBUZWFtMRMwEQYDVQQDDApyZWRoYXQuY29tMIIB" +
            "IjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAuZqfckEubZuMrB+5YIgqkb5+cF8IX5QR4iOz" +
            "VvlwOsXOlsMT67u3ktf5kSQgvMDuDpfQIK4Ws0UYNwmwRgqSJLXGqRzm2CKfkJCjrKxwJ5+aK1gC" +
            "KrdBqLRFfb74N8jE2dpjo7UiuHIT5WPT2M49XEGAjhhHaGbdxU0kCIhx8hQWJSzkPzp8wWIgqlTX" +
            "iBbVTA2OIJGwoE7QCde5qAuwdT6+CfwFryiRpp5OQ7nPmASztNAneg5jgLXS1NY9+CP4eHwJhm6x" +
            "y2GWRGrWWJ/diLunqe/pHSZmqP309e1lqV0jTKbopwM70f0sunPz82YV4ye4cSWDSvzGDDf1s5dY" +
            "LQIDAQABo1AwTjAdBgNVHQ4EFgQUwap3zhfgVfnrfpGpBHCRya1mTsUwHwYDVR0jBBgwFoAUwap3" +
            "zhfgVfnrfpGpBHCRya1mTsUwDAYDVR0TBAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAJBpAWoOo" +
            "XFEKU/lVQJTAq/XpwWSMbgv5arVb6UsBvMdMHKEQuf9+zwcUdUDKTMf9yvz1COSQjn5H4Oqx4fLt" +
            "cSsjMKhMnwUmuHbmPDjMp6cEGqCSAMTK2dIogxbKm13TTOeo3i9xwOeRgqeqCjUF2hy1uOGFY5rW" +
            "N/9pC0OuLkRVFQugZmAhf8jpJgQ3yu5WFsIfElZVGxXn5Ai1nlEsBMbGkYtskmlH/8BcVA4qNmM/" +
            "IZd94TqZrZ5FMz4KLDSBE2YnfhgBg4sjBhLe6yZ1wkASSLuAhphLg5PTC1j25Xlgjm0Xab1mAhlx" +
            "yPxAKvvoJXrKG20s/4yjUh41Z6OfniEoz9VpCAtv1koaVgNCyty55eGc";
    private static KeyStore ks;
    private static KeyStore ts;
    private static char[] passphrase;

    private static SSLEngine clientSSLEngine;
    private static SSLEngine serverSSLEngine;

    // Test methods

    public static void main(String args[]) throws Exception {
        
        initialize();
        
        // SSLEngine code based on RedhandshakeFinished.java
        
        boolean dataDone = false;
        
        ByteBuffer clientOut = null;       
        ByteBuffer clientIn = null;       
        ByteBuffer serverOut = null;       
        ByteBuffer serverIn = null;
        ByteBuffer cTOs;
        ByteBuffer sTOc; 

        SSLSession session = clientSSLEngine.getSession();
        int appBufferMax = session.getApplicationBufferSize();
        int netBufferMax = session.getPacketBufferSize();

        clientIn = ByteBuffer.allocate(appBufferMax + 50);
        serverIn = ByteBuffer.allocate(appBufferMax + 50);

        cTOs = ByteBuffer.allocateDirect(netBufferMax);
        sTOc = ByteBuffer.allocateDirect(netBufferMax);

        clientOut = ByteBuffer.wrap("Hi Server, I'm Client".getBytes());
        serverOut = ByteBuffer.wrap("Hello Client, I'm Server".getBytes());

        SSLEngineResult clientResult;
        SSLEngineResult serverResult;

        while (!dataDone) {
            clientResult = clientSSLEngine.wrap(clientOut, cTOs);
            runDelegatedTasks(clientResult, clientSSLEngine);
            serverResult = serverSSLEngine.wrap(serverOut, sTOc);
            runDelegatedTasks(serverResult, serverSSLEngine);
            cTOs.flip();
            sTOc.flip();
            
            System.out.println("Client -> Network");
            printTlsNetworkPacket("", cTOs);
            System.out.println("");
            System.out.println("Server -> Network");
            printTlsNetworkPacket("", sTOc);
            
            clientResult = clientSSLEngine.unwrap(sTOc, clientIn);
            runDelegatedTasks(clientResult, clientSSLEngine);
            serverResult = serverSSLEngine.unwrap(cTOs, serverIn);
            runDelegatedTasks(serverResult, serverSSLEngine);
            
            cTOs.compact();
            sTOc.compact();
            
            if (!dataDone && (clientOut.limit() == serverIn.position()) &&
                    (serverOut.limit() == clientIn.position())) {

                /*
                 * A sanity check to ensure we got what was sent.
                 */
                checkTransfer(serverOut, clientIn);
                checkTransfer(clientOut, serverIn);

                dataDone = true;
            } 
        }    
        
        Class<?> SSLSessionImplClass = Class.forName("sun.security.ssl.SSLSessionImpl");
        Method getMasterSecretMethod = SSLSessionImplClass.getDeclaredMethod("getMasterSecret");
        getMasterSecretMethod.setAccessible(true);
        
        byte[] clientMasterKey = ((SecretKey)getMasterSecretMethod.invoke(clientSSLEngine.getSession())).getEncoded();
        byte[] serverMasterKey = ((SecretKey)getMasterSecretMethod.invoke(serverSSLEngine.getSession())).getEncoded();
        
        System.out.println("Client Extended Master Key");
        showBinaryPacket(clientMasterKey);
        System.out.println(".......................");
        
        System.out.println("Server Extended Master Key");
        showBinaryPacket(serverMasterKey);
        System.out.println(".......................");
        
        if (!Arrays.equals(clientMasterKey, expectedMasterKey) || 
                !Arrays.equals(serverMasterKey, expectedMasterKey)) {
            throw new Exception("Master key does not match expected master key.");
        }
        
        System.out.println("Test passed - OK");
    }
    
    static void printTlsNetworkPacket(String prefix, ByteBuffer bb) {
        HexDumpEncoder dump = new HexDumpEncoder();
        System.out.println(prefix);
        try {
            dump.encodeBuffer(bb.slice(), System.out);
        } catch (Exception e) {
            // ignore
        }
        System.out.flush();
    }
    
    private static void checkTransfer(ByteBuffer a, ByteBuffer b)
            throws Exception {
        a.flip();
        b.flip();

        if (!a.equals(b)) {
            throw new Exception("Data didn't transfer cleanly");
        } else {
        }

        a.position(a.limit());
        b.position(b.limit());
        a.limit(a.capacity());
        b.limit(b.capacity());
    }
    
    private static void runDelegatedTasks(SSLEngineResult result,
            SSLEngine engine) throws Exception {

        if (result.getHandshakeStatus() == HandshakeStatus.NEED_TASK) {
            Runnable runnable;
            while ((runnable = engine.getDelegatedTask()) != null) {
                runnable.run();
            }
            HandshakeStatus hsStatus = engine.getHandshakeStatus();
            if (hsStatus == HandshakeStatus.NEED_TASK) {
                throw new Exception(
                    "handshake shouldn't need additional tasks");
            }
        }
    }

    private static void initialize() throws Exception {
        ks = KeyStore.getInstance("JKS");
        ts = KeyStore.getInstance("JKS");
        passphrase = "123456".toCharArray();
        ks.load(new ByteArrayInputStream(Base64.getDecoder().decode(ksString)), passphrase);
        ts.load(new ByteArrayInputStream(Base64.getDecoder().decode(ksString)), passphrase);

        clientSSLEngine = createSSLEngine(true);
        serverSSLEngine = createSSLEngine(false);

        String[] preferredSuites = new String[] { "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256" };

        serverSSLEngine.setEnabledCipherSuites(preferredSuites);
        clientSSLEngine.setEnabledCipherSuites(preferredSuites);
    }
    
    static private SSLEngine createSSLEngine(boolean client) throws Exception {
        SSLEngine ssle;
        KeyManagerFactory kmf = KeyManagerFactory.getInstance("PKIX");
        kmf.init(ks, passphrase);

        TrustManagerFactory tmf = TrustManagerFactory.getInstance("PKIX");
        tmf.init(ts);

        SSLContext sslCtx = SSLContext.getInstance("TLSv1.2");

        SecureRandom random = new SecureRandom();
        
        // Patch to remove randomness so we can assert obtained values against
        // pre-computed / expected values.
        Class<?> randomClass = random.getClass();
        Field secureRandomSpiField = randomClass.getDeclaredField("secureRandomSpi");
        secureRandomSpiField.setAccessible(true);
        secureRandomSpiField.set(random, new NoRandom());
        
        sslCtx.init(kmf.getKeyManagers(), tmf.getTrustManagers(), random);

        ssle = sslCtx.createSSLEngine("localhost", 443);
        ssle.setUseClientMode(client);
        SSLParameters sslParameters = ssle.getSSLParameters();
        ssle.setSSLParameters(sslParameters);

        return ssle;
    }

    private static void showBinaryPacket(byte[] binaryPacket) throws Exception {
        int columnCounter = 0;
        for (byte binaryPacketByte : binaryPacket) {
            System.out.print(String.format("%02X ", binaryPacketByte));
            if ((columnCounter + 1) % 8 == 0) {
                System.out.print("\n");
            }
            columnCounter++;
        }
        System.out.print("\n");
    }
    
    private static class NoRandom extends SecureRandomSpi {
        @Override
        protected void engineSetSeed(byte[] seed) {
        }

        @Override
        protected void engineNextBytes(byte[] bytes) {
            for (int i = 0; i < bytes.length; i++) {
                bytes[i] = (byte)0;
            }
        }

        @Override
        protected byte[] engineGenerateSeed(int numBytes) {
            return new byte[numBytes];
        }
    }
}