--- a/jdk/src/share/classes/sun/security/pkcs12/PKCS12KeyStore.java	Wed Jun 23 17:03:40 2010 -0700
+++ b/jdk/src/share/classes/sun/security/pkcs12/PKCS12KeyStore.java	Thu Jun 24 14:26:22 2010 +0800
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 1999, 2007, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 1999, 2010, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -180,24 +180,15 @@
         String alias;
     };
 
-    private static class KeyId {
-        byte[] keyId;
-
-        KeyId(byte[] keyId) {
+    // A certificate with its PKCS #9 attributes
+    private static class CertEntry {
+        final X509Certificate cert;
+        final byte[] keyId;
+        final String alias;
+        CertEntry(X509Certificate cert, byte[] keyId, String alias) {
+            this.cert = cert;
             this.keyId = keyId;
-        }
-        public int hashCode() {
-            int hash = 0;
-
-            for (int i = 0; i < keyId.length; i++)
-                hash += keyId[i];
-            return hash;
-        }
-        public boolean equals(Object obj) {
-            if (!(obj instanceof KeyId))
-                return false;
-            KeyId that = (KeyId)obj;
-            return (Arrays.equals(this.keyId, that.keyId));
+            this.alias = alias;
         }
     }
 
@@ -209,7 +200,9 @@
                                 new Hashtable<String, KeyEntry>();
 
     private ArrayList<KeyEntry> keyList = new ArrayList<KeyEntry>();
-    private LinkedHashMap<Object, X509Certificate> certs = new LinkedHashMap<Object, X509Certificate>();
+    private LinkedHashMap<X500Principal, X509Certificate> certsMap =
+            new LinkedHashMap<X500Principal, X509Certificate>();
+    private ArrayList<CertEntry> certEntries = new ArrayList<CertEntry>();
 
     /**
      * Returns the key associated with the given alias, using the given
@@ -472,6 +465,15 @@
         KeyEntry entry = new KeyEntry();
         entry.date = new Date();
 
+        try {
+            // set the keyId to current date
+            entry.keyId = ("Time " + (entry.date).getTime()).getBytes("UTF8");
+        } catch (UnsupportedEncodingException ex) {
+            // Won't happen
+        }
+        // set the alias
+        entry.alias = alias.toLowerCase();
+
         entry.protectedPrivKey = key.clone();
         if (chain != null) {
            entry.chain = chain.clone();
@@ -1027,10 +1029,9 @@
                 // All Certs should have a unique friendlyName.
                 // This change is made to meet NSS requirements.
                 byte[] bagAttrs = null;
-                String friendlyName = cert.getSubjectX500Principal().getName();
                 if (i == 0) {
                     // Only End-Entity Cert should have a localKeyId.
-                    bagAttrs = getBagAttributes(friendlyName, entry.keyId);
+                    bagAttrs = getBagAttributes(entry.alias, entry.keyId);
                 } else {
                     // Trusted root CA certs and Intermediate CA certs do not
                     // need to have a localKeyId, and hence localKeyId is null
@@ -1038,7 +1039,8 @@
                     // NSS pkcs12 library requires trusted CA certs in the
                     // certificate chain to have unique or null localKeyID.
                     // However, IE/OpenSSL do not impose this restriction.
-                    bagAttrs = getBagAttributes(friendlyName, null);
+                    bagAttrs = getBagAttributes(
+                            cert.getSubjectX500Principal().getName(), null);
                 }
                 if (bagAttrs != null) {
                     safeBag.write(bagAttrs);
@@ -1333,24 +1335,49 @@
             if (entry.keyId != null) {
                 ArrayList<X509Certificate> chain =
                                 new ArrayList<X509Certificate>();
-                X509Certificate cert = certs.get(new KeyId(entry.keyId));
+                X509Certificate cert = findMatchedCertificate(entry);
                 while (cert != null) {
                     chain.add(cert);
                     X500Principal issuerDN = cert.getIssuerX500Principal();
                     if (issuerDN.equals(cert.getSubjectX500Principal())) {
                         break;
                     }
-                    cert = certs.get(issuerDN);
+                    cert = certsMap.get(issuerDN);
                 }
                 /* Update existing KeyEntry in entries table */
                 if (chain.size() > 0)
                     entry.chain = chain.toArray(new Certificate[chain.size()]);
             }
         }
-        certs.clear();
+        certEntries.clear();
+        certsMap.clear();
         keyList.clear();
     }
 
+    /**
+     * Locates a matched CertEntry from certEntries, and returns its cert.
+     * @param entry the KeyEntry to match
+     * @return a certificate, null if not found
+     */
+    private X509Certificate findMatchedCertificate(KeyEntry entry) {
+        CertEntry keyIdMatch = null;
+        CertEntry aliasMatch = null;
+        for (CertEntry ce: certEntries) {
+            if (Arrays.equals(entry.keyId, ce.keyId)) {
+                keyIdMatch = ce;
+                if (entry.alias.equalsIgnoreCase(ce.alias)) {
+                    // Full match!
+                    return ce.cert;
+                }
+            } else if (entry.alias.equalsIgnoreCase(ce.alias)) {
+                aliasMatch = ce;
+            }
+        }
+        // keyId match first, for compatibility
+        if (keyIdMatch != null) return keyIdMatch.cert;
+        else if (aliasMatch != null) return aliasMatch.cert;
+        else return null;
+    }
 
     private void loadSafeContents(DerInputStream stream, char[] password)
         throws IOException, NoSuchAlgorithmException, CertificateException
@@ -1491,19 +1518,12 @@
                         keyId = "01".getBytes("UTF8");
                     }
                 }
-                if (keyId != null) {
-                    KeyId keyid = new KeyId(keyId);
-                    if (!certs.containsKey(keyid))
-                        certs.put(keyid, cert);
-                }
-                if (alias != null) {
-                    if (!certs.containsKey(alias))
-                        certs.put(alias, cert);
-                }
+                certEntries.add(new CertEntry(cert, keyId, alias));
                 X500Principal subjectDN = cert.getSubjectX500Principal();
                 if (subjectDN != null) {
-                    if (!certs.containsKey(subjectDN))
-                        certs.put(subjectDN, cert);
+                    if (!certsMap.containsKey(subjectDN)) {
+                        certsMap.put(subjectDN, cert);
+                    }
                 }
             }
         }

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/jdk/test/sun/security/pkcs12/PKCS12SameKeyId.java	Thu Jun 24 14:26:22 2010 +0800
@@ -0,0 +1,120 @@
+/*
+ * Copyright (c) 2010, Oracle and/or its affiliates. All rights reserved.
+ * 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 6958026
+ * @summary Problem with PKCS12 keystore
+ */
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.security.AlgorithmParameters;
+import java.security.KeyStore;
+import java.security.cert.Certificate;
+import java.security.cert.X509Certificate;
+import javax.crypto.Cipher;
+import javax.crypto.SecretKey;
+import javax.crypto.SecretKeyFactory;
+import javax.crypto.spec.PBEKeySpec;
+import javax.crypto.spec.PBEParameterSpec;
+import sun.security.pkcs.EncryptedPrivateKeyInfo;
+import sun.security.tools.KeyTool;
+import sun.security.util.ObjectIdentifier;
+import sun.security.x509.AlgorithmId;
+import sun.security.x509.X500Name;
+
+public class PKCS12SameKeyId {
+
+    private static final String JKSFILE = "PKCS12SameKeyId.jks";
+    private static final String P12FILE = "PKCS12SameKeyId.p12";
+    private static final char[] PASSWORD = "changeit".toCharArray();
+    private static final int SIZE = 10;
+
+    public static void main(String[] args) throws Exception {
+
+        // Prepare a JKS keystore with many entries
+        new File(JKSFILE).delete();
+        for (int i=0; i<SIZE; i++) {
+            System.err.print(".");
+            String cmd = "-keystore " + JKSFILE
+                    + " -storepass changeit -keypass changeit "
+                    + "-genkeypair -alias p" + i + " -dname CN=" + i;
+            KeyTool.main(cmd.split(" "));
+        }
+
+        // Prepare EncryptedPrivateKeyInfo parameters, copied from various
+        // places in PKCS12KeyStore.java
+        AlgorithmParameters algParams =
+                AlgorithmParameters.getInstance("PBEWithSHA1AndDESede");
+        algParams.init(new PBEParameterSpec("12345678".getBytes(), 1024));
+        AlgorithmId algid = new AlgorithmId(
+                new ObjectIdentifier("1.2.840.113549.1.12.1.3"), algParams);
+
+        PBEKeySpec keySpec = new PBEKeySpec(PASSWORD);
+        SecretKeyFactory skFac = SecretKeyFactory.getInstance("PBE");
+        SecretKey skey = skFac.generateSecret(keySpec);
+
+        Cipher cipher = Cipher.getInstance("PBEWithSHA1AndDESede");
+        cipher.init(Cipher.ENCRYPT_MODE, skey, algParams);
+
+        // Pre-calculated keys and certs and aliases
+        byte[][] keys = new byte[SIZE][];
+        Certificate[][] certChains = new Certificate[SIZE][];
+        String[] aliases = new String[SIZE];
+
+        // Reads from JKS keystore and pre-calculate
+        KeyStore ks = KeyStore.getInstance("jks");
+        ks.load(new FileInputStream(JKSFILE), PASSWORD);
+        for (int i=0; i<SIZE; i++) {
+            aliases[i] = "p" + i;
+            byte[] enckey = cipher.doFinal(
+                    ks.getKey(aliases[i], PASSWORD).getEncoded());
+            keys[i] = new EncryptedPrivateKeyInfo(algid, enckey).getEncoded();
+            certChains[i] = ks.getCertificateChain(aliases[i]);
+        }
+
+        // Write into PKCS12 keystore. Use this overloaded version of
+        // setKeyEntry() to be as fast as possible, so that they would
+        // have same localKeyId.
+        KeyStore p12 = KeyStore.getInstance("pkcs12");
+        p12.load(null, PASSWORD);
+        for (int i=0; i<SIZE; i++) {
+            p12.setKeyEntry(aliases[i], keys[i], certChains[i]);
+        }
+        p12.store(new FileOutputStream(P12FILE), PASSWORD);
+
+        // Check private keys still match certs
+        p12 = KeyStore.getInstance("pkcs12");
+        p12.load(new FileInputStream(P12FILE), PASSWORD);
+        for (int i=0; i<SIZE; i++) {
+            String a = "p" + i;
+            X509Certificate x = (X509Certificate)p12.getCertificate(a);
+            X500Name name = (X500Name)x.getSubjectDN();
+            if (!name.getCommonName().equals(""+i)) {
+                throw new Exception(a + "'s cert is " + name);
+            }
+        }
+    }
+}