178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xupackage com.android.server.wifi; 278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport static org.mockito.Matchers.anyInt; 478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport static org.mockito.Matchers.anyString; 578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport static org.mockito.Mockito.mock; 678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport static org.mockito.Mockito.when; 778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport android.security.KeyStore; 978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport android.util.SparseArray; 1078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 1178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport org.mockito.Matchers; 1278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport org.mockito.invocation.InvocationOnMock; 1378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport org.mockito.stubbing.Answer; 1478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 1578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport java.util.Arrays; 1678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuimport java.util.HashMap; 1778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 1878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xuclass MockKeyStore { 1978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 2078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public static class KeyBlob { 2178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public byte[] blob; 2278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public int flag; 2378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 2478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public void update(byte[] blob, int flag) { 2578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu this.blob = Arrays.copyOf(blob, blob.length); 2678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu this.flag = flag; 2778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 2878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 2978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private SparseArray<HashMap<String, KeyBlob>> mStore; 3078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 3178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public MockKeyStore() { 3278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu mStore = new SparseArray<HashMap<String, KeyBlob>>(); 3378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 3478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 3578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public KeyStore createMock() { 3678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu KeyStore mock = mock(KeyStore.class); 3778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.state()).thenReturn(KeyStore.State.UNLOCKED); 3878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 3978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.put(anyString(), Matchers.any(byte[].class), anyInt(), anyInt())) 4078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu .thenAnswer(new Answer<Boolean>() { 4178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 4278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 4378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public Boolean answer(InvocationOnMock invocation) throws Throwable { 4478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu Object[] args = invocation.getArguments(); 4578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return put((String) args[0], (byte[]) args[1], (Integer) args[2], 4678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu (Integer) args[3]); 4778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 4878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu }); 4978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 5078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.importKey(anyString(), Matchers.any(byte[].class), anyInt(), anyInt())) 5178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu .thenAnswer(new Answer<Boolean>() { 5278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 5378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 5478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public Boolean answer(InvocationOnMock invocation) throws Throwable { 5578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu Object[] args = invocation.getArguments(); 5678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return importKey((String) args[0], (byte[]) args[1], (Integer) args[2], 5778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu (Integer) args[3]); 5878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 5978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu }); 6078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 6178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.delete(anyString(), anyInt())).thenAnswer(new Answer<Boolean>() { 6278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 6378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 6478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public Boolean answer(InvocationOnMock invocation) throws Throwable { 6578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu Object[] args = invocation.getArguments(); 6678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return delete((String) args[0], (Integer) args[1]); 6778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 6878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu }); 6978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 7078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.contains(anyString(), anyInt())).thenAnswer(new Answer<Boolean>() { 7178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 7278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 7378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public Boolean answer(InvocationOnMock invocation) throws Throwable { 7478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu Object[] args = invocation.getArguments(); 7578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return contains((String) args[0], (Integer) args[1]); 7678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 7778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu }); 7878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 7978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu when(mock.duplicate(anyString(), anyInt(), anyString(), anyInt())) 8078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu .thenAnswer(new Answer<Boolean>() { 8178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 8278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public Boolean answer(InvocationOnMock invocation) throws Throwable { 8378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu Object[] args = invocation.getArguments(); 8478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return duplicate((String) args[0], (Integer) args[1], (String) args[2], 8578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu (Integer) args[3]); 8678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 8778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu }); 8878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return mock; 8978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 9078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 9178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private KeyBlob access(int uid, String key, boolean createIfNotExist) { 9278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (mStore.get(uid) == null) { 9378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu mStore.put(uid, new HashMap<String, KeyBlob>()); 9478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 9578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu HashMap<String, KeyBlob> map = mStore.get(uid); 9678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (map.containsKey(key)) { 9778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return map.get(key); 9878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } else { 9978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (createIfNotExist) { 10078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu KeyBlob blob = new KeyBlob(); 10178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu map.put(key, blob); 10278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return blob; 10378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } else { 10478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return null; 10578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 10678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 10778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 10878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 10978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public KeyBlob getKeyBlob(int uid, String key) { 11078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return access(uid, key, false); 11178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 11278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 11378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private boolean put(String key, byte[] value, int uid, int flags) { 11478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu access(uid, key, true).update(value, flags); 11578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return true; 11678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 11778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 11878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private boolean importKey(String keyName, byte[] key, int uid, int flags) { 11978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return put(keyName, key, uid, flags); 12078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 12178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 12278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private boolean delete(String key, int uid) { 12378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (mStore.get(uid) != null) { 12478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu mStore.get(uid).remove(key); 12578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 12678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return true; 12778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 12878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 12978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private boolean contains(String key, int uid) { 13078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return access(uid, key, false) != null; 13178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 13278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 13378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu private boolean duplicate(String srcKey, int srcUid, String destKey, int destUid) { 13478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu for (int i = 0; i < mStore.size(); i++) { 13578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu int key = mStore.keyAt(i); 13678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu // Cannot copy to itself 13778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (srcKey.equals(destKey) && key == destUid) { 13878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu continue; 13978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 14078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (srcUid == -1 || srcUid == key) { 14178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu HashMap<String, KeyBlob> map = mStore.get(key); 14278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu if (map.containsKey(srcKey)) { 14378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu KeyBlob blob = map.get(srcKey); 14478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu access(destUid, destKey, true).update(blob.blob, blob.flag); 14578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu break; 14678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 14778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 14878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 14978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return true; 15078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 15178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu 15278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu @Override 15378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu public String toString() { 15478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu StringBuilder sb = new StringBuilder(); 15578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu sb.append("KeyStore {"); 15678463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu for (int i = 0; i < mStore.size(); i++) { 15778463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu int uid = mStore.keyAt(i); 15878463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu for (String keyName : mStore.get(uid).keySet()) { 15978463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu sb.append(String.format("%d:%s, ", uid, keyName)); 16078463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 16178463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 16278463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu sb.append("}"); 16378463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu return sb.toString(); 16478463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu } 16578463be32c01259841f4c0b670b742c3c2040cf3Rubin Xu}