Add PersistentKeyChainSnapshot serialization/deserialization methods.

/*

 * Copyright (C) 2017 The Android Open Source Project

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.android.server.locksettings.recoverablekeystore.storage;

import android.annotation.NonNull;

import android.annotation.Nullable;

import android.security.keystore.recovery.KeyChainProtectionParams;

import android.security.keystore.recovery.KeyChainSnapshot;

import android.security.keystore.recovery.KeyDerivationParams;

import android.security.keystore.recovery.WrappedApplicationKey;

import com.android.internal.annotations.VisibleForTesting;

import com.android.internal.util.ArrayUtils;

import java.io.ByteArrayInputStream;

import java.io.ByteArrayOutputStream;

import java.io.DataInputStream;

import java.io.DataOutputStream;

import java.io.IOException;

import java.util.ArrayList;

import java.util.List;

/**

 * This class provides helper methods serialize and deserialize {@link KeyChainSnapshot}.

 *

 * <p> It is necessary since {@link android.os.Parcelable} is not designed for persistent storage.

 *

 * <p> For every list, length is stored before the elements.

 *

 */

public class PersistentKeyChainSnapshot {

    private static final int VERSION = 1;

    private static final int NULL_LIST_LENGTH = -1;

    private DataInputStream mInput;

    private DataOutputStream mOut;

    private ByteArrayOutputStream mOutStream;

    @VisibleForTesting

    PersistentKeyChainSnapshot() {

    }

    @VisibleForTesting

    void initReader(byte[] input) {

        mInput = new DataInputStream(new ByteArrayInputStream(input));

    }

    @VisibleForTesting

    void initWriter() {

        mOutStream = new ByteArrayOutputStream();

        mOut = new DataOutputStream(mOutStream);

    }

    @VisibleForTesting

    byte[] getOutput() {

        return mOutStream.toByteArray();

    }

    /**

     * Converts {@link KeyChainSnapshot} to its binary representation.

     *

     * @param snapshot The snapshot.

     *

     * @throws IOException if serialization failed.

     */

    public static byte[] serialize(@NonNull KeyChainSnapshot snapshot) throws IOException {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        writer.writeInt(VERSION);

        writer.writeKeyChainSnapshot(snapshot);

        return writer.getOutput();

    }

    /**

     * deserializes {@link KeyChainSnapshot}.

     *

     * @input input - byte array produced by {@link serialize} method.

     * @throws IOException if parsing failed.

     */

    public static @NonNull KeyChainSnapshot deserialize(@NonNull byte[] input)

            throws IOException {

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(input);

        try {

            int version = reader.readInt();

            if (version != VERSION) {

                throw new IOException("Unsupported version " + version);

            }

            return reader.readKeyChainSnapshot();

        } catch (IOException e) {

            throw new IOException("Malformed KeyChainSnapshot", e);

        }

    }

    /**

     * Must be in sync with {@link KeyChainSnapshot.writeToParcel}

     */

    @VisibleForTesting

    void writeKeyChainSnapshot(KeyChainSnapshot snapshot) throws IOException {

        writeInt(snapshot.getSnapshotVersion());

        writeProtectionParamsList(snapshot.getKeyChainProtectionParams());

        writeBytes(snapshot.getEncryptedRecoveryKeyBlob());

        writeKeysList(snapshot.getWrappedApplicationKeys());

        writeInt(snapshot.getMaxAttempts());

        writeLong(snapshot.getCounterId());

        writeBytes(snapshot.getServerParams());

        writeBytes(snapshot.getTrustedHardwarePublicKey());

    }

    @VisibleForTesting

    KeyChainSnapshot readKeyChainSnapshot() throws IOException {

        int snapshotVersion = readInt();

        List<KeyChainProtectionParams> protectionParams = readProtectionParamsList();

        byte[] encryptedRecoveryKey = readBytes();

        List<WrappedApplicationKey> keysList = readKeysList();

        int maxAttempts = readInt();

        long conterId = readLong();

        byte[] serverParams = readBytes();

        byte[] trustedHardwarePublicKey = readBytes();

        return new KeyChainSnapshot.Builder()

                .setSnapshotVersion(snapshotVersion)

                .setKeyChainProtectionParams(protectionParams)

                .setEncryptedRecoveryKeyBlob(encryptedRecoveryKey)

                .setWrappedApplicationKeys(keysList)

                .setMaxAttempts(maxAttempts)

                .setCounterId(conterId)

                .setServerParams(serverParams)

                .setTrustedHardwarePublicKey(trustedHardwarePublicKey)

                .build();

    }

    @VisibleForTesting

    void writeProtectionParamsList(

            @NonNull List<KeyChainProtectionParams> ProtectionParamsList) throws IOException {

        writeInt(ProtectionParamsList.size());

        for (KeyChainProtectionParams protectionParams : ProtectionParamsList) {

            writeProtectionParams(protectionParams);

        }

    }

    @VisibleForTesting

    List<KeyChainProtectionParams> readProtectionParamsList() throws IOException {

        int length = readInt();

        List<KeyChainProtectionParams> result = new ArrayList<>(length);

        for (int i = 0; i < length; i++) {

            result.add(readProtectionParams());

        }

        return result;

    }

    /**

     * Must be in sync with {@link KeyChainProtectionParams.writeToParcel}

     */

    @VisibleForTesting

    void writeProtectionParams(@NonNull KeyChainProtectionParams protectionParams)

            throws IOException {

        if (!ArrayUtils.isEmpty(protectionParams.getSecret())) {

            // Extra security check.

            throw new RuntimeException("User generated secret should not be stored");

        }

        writeInt(protectionParams.getUserSecretType());

        writeInt(protectionParams.getLockScreenUiFormat());

        writeKeyDerivationParams(protectionParams.getKeyDerivationParams());

        writeBytes(protectionParams.getSecret());

    }

    @VisibleForTesting

    KeyChainProtectionParams readProtectionParams() throws IOException {

        int userSecretType = readInt();

        int lockScreenUiFormat = readInt();

        KeyDerivationParams derivationParams = readKeyDerivationParams();

        byte[] secret = readBytes();

        return new KeyChainProtectionParams.Builder()

                .setUserSecretType(userSecretType)

                .setLockScreenUiFormat(lockScreenUiFormat)

                .setKeyDerivationParams(derivationParams)

                .setSecret(secret)

                .build();

    }

    /**

     * Must be in sync with {@link KeyDerivationParams.writeToParcel}

     */

    @VisibleForTesting

    void writeKeyDerivationParams(@NonNull KeyDerivationParams Params) throws IOException {

        writeInt(Params.getAlgorithm());

        writeBytes(Params.getSalt());

    }

    @VisibleForTesting

    KeyDerivationParams readKeyDerivationParams() throws IOException {

        int algorithm = readInt();

        byte[] salt = readBytes();

        return KeyDerivationParams.createSha256Params(salt);

    }

    @VisibleForTesting

    void writeKeysList(@NonNull List<WrappedApplicationKey> applicationKeys) throws IOException {

        writeInt(applicationKeys.size());

        for (WrappedApplicationKey keyEntry : applicationKeys) {

            writeKeyEntry(keyEntry);

        }

    }

    @VisibleForTesting

    List<WrappedApplicationKey> readKeysList() throws IOException {

        int length = readInt();

        List<WrappedApplicationKey> result = new ArrayList<>(length);

        for (int i = 0; i < length; i++) {

            result.add(readKeyEntry());

        }

        return result;

    }

    /**

     * Must be in sync with {@link WrappedApplicationKey.writeToParcel}

     */

    @VisibleForTesting

    void writeKeyEntry(@NonNull WrappedApplicationKey keyEntry) throws IOException {

        mOut.writeUTF(keyEntry.getAlias());

        writeBytes(keyEntry.getEncryptedKeyMaterial());

        writeBytes(keyEntry.getAccount());

    }

    @VisibleForTesting

    WrappedApplicationKey readKeyEntry() throws IOException {

        String alias = mInput.readUTF();

        byte[] keyMaterial = readBytes();

        byte[] account = readBytes();

        return new WrappedApplicationKey.Builder()

                .setAlias(alias)

                .setEncryptedKeyMaterial(keyMaterial)

                .setAccount(account)

                .build();

    }

    @VisibleForTesting

    void writeInt(int value) throws IOException {

        mOut.writeInt(value);

    }

    @VisibleForTesting

    int readInt() throws IOException {

        return mInput.readInt();

    }

    @VisibleForTesting

    void writeLong(long value) throws IOException {

        mOut.writeLong(value);

    }

    @VisibleForTesting

    long readLong() throws IOException {

        return mInput.readLong();

    }

    @VisibleForTesting

    void writeBytes(@Nullable byte[] value) throws IOException {

        if (value == null) {

            writeInt(NULL_LIST_LENGTH);

            return;

        }

        writeInt(value.length);

        mOut.write(value, 0, value.length);

    }

    /**

     * Reads @code{byte[]} from current position. Converts {@code null} to an empty array.

     */

    @VisibleForTesting

    @NonNull byte[] readBytes() throws IOException {

        int length = readInt();

        if (length == NULL_LIST_LENGTH) {

            return new byte[]{};

        }

        byte[] result = new byte[length];

        mInput.read(result, 0, result.length);

        return result;

    }

}

/*

 * Copyright (C) 2017 The Android Open Source Project

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.android.server.locksettings.recoverablekeystore.storage;

import static com.google.common.truth.Truth.assertThat;

import static org.testng.Assert.assertThrows;

import android.security.keystore.recovery.KeyDerivationParams;

import android.security.keystore.recovery.WrappedApplicationKey;

import android.security.keystore.recovery.KeyChainSnapshot;

import android.security.keystore.recovery.KeyChainProtectionParams;

import android.support.test.filters.SmallTest;

import android.support.test.runner.AndroidJUnit4;

import org.junit.Test;

import org.junit.runner.RunWith;

import java.io.IOException;

import java.nio.charset.StandardCharsets;

import java.util.ArrayList;

@SmallTest

@RunWith(AndroidJUnit4.class)

public class PersistentKeyChainSnapshotTest {

    private static final String ALIAS = "some_key";

    private static final String ALIAS2 = "another_key";

    private static final byte[] RECOVERY_KEY_MATERIAL = "recovery_key_data"

            .getBytes(StandardCharsets.UTF_8);

    private static final byte[] KEY_MATERIAL = "app_key_data".getBytes(StandardCharsets.UTF_8);

    private static final byte[] PUBLIC_KEY = "public_key_data".getBytes(StandardCharsets.UTF_8);

    private static final byte[] ACCOUNT = "test_account".getBytes(StandardCharsets.UTF_8);

    private static final byte[] SALT = "salt".getBytes(StandardCharsets.UTF_8);

    private static final int SNAPSHOT_VERSION = 2;

    private static final int MAX_ATTEMPTS = 10;

    private static final long COUNTER_ID = 123456789L;

    private static final byte[] SERVER_PARAMS = "server_params".getBytes(StandardCharsets.UTF_8);

    private static final byte[] ZERO_BYTES = new byte[0];

    private static final byte[] ONE_BYTE = new byte[]{(byte) 11};

    private static final byte[] TWO_BYTES = new byte[]{(byte) 222,(byte) 222};

    @Test

    public void testWriteInt() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        writer.writeInt(Integer.MIN_VALUE);

        writer.writeInt(Integer.MAX_VALUE);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        assertThat(reader.readInt()).isEqualTo(Integer.MIN_VALUE);

        assertThat(reader.readInt()).isEqualTo(Integer.MAX_VALUE);

        assertThrows(

                IOException.class,

                () -> reader.readInt());

    }

    @Test

    public void testWriteLong() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        writer.writeLong(Long.MIN_VALUE);

        writer.writeLong(Long.MAX_VALUE);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        assertThat(reader.readLong()).isEqualTo(Long.MIN_VALUE);

        assertThat(reader.readLong()).isEqualTo(Long.MAX_VALUE);

        assertThrows(

                IOException.class,

                () -> reader.readLong());

    }

    @Test

    public void testWriteBytes() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        writer.writeBytes(ZERO_BYTES);

        writer.writeBytes(ONE_BYTE);

        writer.writeBytes(TWO_BYTES);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        assertThat(reader.readBytes()).isEqualTo(ZERO_BYTES);

        assertThat(reader.readBytes()).isEqualTo(ONE_BYTE);

        assertThat(reader.readBytes()).isEqualTo(TWO_BYTES);

        assertThrows(

                IOException.class,

                () -> reader.readBytes());

    }

    @Test

    public void testReadBytes_returnsNullArrayAsEmpty() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        writer.writeBytes(null);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        assertThat(reader.readBytes()).isEqualTo(new byte[]{}); // null -> empty array

    }

    @Test

    public void testWriteKeyEntry() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        WrappedApplicationKey entry = new WrappedApplicationKey.Builder()

                .setAlias(ALIAS)

                .setEncryptedKeyMaterial(KEY_MATERIAL)

                .setAccount(ACCOUNT)

                .build();

        writer.writeKeyEntry(entry);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        WrappedApplicationKey copy = reader.readKeyEntry();

        assertThat(copy.getAlias()).isEqualTo(ALIAS);

        assertThat(copy.getEncryptedKeyMaterial()).isEqualTo(KEY_MATERIAL);

        assertThat(copy.getAccount()).isEqualTo(ACCOUNT);

        assertThrows(

                IOException.class,

                () -> reader.readKeyEntry());

    }

    public void testWriteProtectionParams() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        KeyDerivationParams derivationParams = KeyDerivationParams.createSha256Params(SALT);

        KeyChainProtectionParams protectionParams =  new KeyChainProtectionParams.Builder()

                .setUserSecretType(1)

                .setLockScreenUiFormat(2)

                .setKeyDerivationParams(derivationParams)

                .build();

        writer.writeProtectionParams(protectionParams);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        KeyChainProtectionParams copy = reader.readProtectionParams();

        assertThat(copy.getUserSecretType()).isEqualTo(1);

        assertThat(copy.getLockScreenUiFormat()).isEqualTo(2);

        assertThat(copy.getKeyDerivationParams().getSalt()).isEqualTo(SALT);

        assertThrows(

                IOException.class,

                () -> reader.readProtectionParams());

    }

    public void testKeyChainSnapshot() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        KeyDerivationParams derivationParams = KeyDerivationParams.createSha256Params(SALT);

        ArrayList<KeyChainProtectionParams> protectionParamsList = new ArrayList<>();

        protectionParamsList.add(new KeyChainProtectionParams.Builder()

                .setUserSecretType(1)

                .setLockScreenUiFormat(2)

                .setKeyDerivationParams(derivationParams)

                .build());

        ArrayList<WrappedApplicationKey> appKeysList = new ArrayList<>();

        appKeysList.add(new WrappedApplicationKey.Builder()

                .setAlias(ALIAS)

                .setEncryptedKeyMaterial(KEY_MATERIAL)

                .setAccount(ACCOUNT)

                .build());

        KeyChainSnapshot snapshot =  new KeyChainSnapshot.Builder()

                .setSnapshotVersion(SNAPSHOT_VERSION)

                .setKeyChainProtectionParams(protectionParamsList)

                .setEncryptedRecoveryKeyBlob(KEY_MATERIAL)

                .setWrappedApplicationKeys(appKeysList)

                .setMaxAttempts(MAX_ATTEMPTS)

                .setCounterId(COUNTER_ID)

                .setServerParams(SERVER_PARAMS)

                .setTrustedHardwarePublicKey(PUBLIC_KEY)

                .build();

        writer.writeKeyChainSnapshot(snapshot);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        KeyChainSnapshot copy = reader.readKeyChainSnapshot();

        assertThat(copy.getSnapshotVersion()).isEqualTo(SNAPSHOT_VERSION);

        assertThat(copy.getKeyChainProtectionParams()).hasSize(2);

        assertThat(copy.getKeyChainProtectionParams().get(0).getUserSecretType()).isEqualTo(1);

        assertThat(copy.getKeyChainProtectionParams().get(1).getUserSecretType()).isEqualTo(2);

        assertThat(copy.getEncryptedRecoveryKeyBlob()).isEqualTo(RECOVERY_KEY_MATERIAL);

        assertThat(copy.getWrappedApplicationKeys()).hasSize(2);

        assertThat(copy.getWrappedApplicationKeys().get(0).getAlias()).isEqualTo(ALIAS);

        assertThat(copy.getWrappedApplicationKeys().get(1).getAlias()).isEqualTo(ALIAS2);

        assertThat(copy.getMaxAttempts()).isEqualTo(MAX_ATTEMPTS);

        assertThat(copy.getCounterId()).isEqualTo(COUNTER_ID);

        assertThat(copy.getServerParams()).isEqualTo(SERVER_PARAMS);

        assertThat(copy.getTrustedHardwarePublicKey()).isEqualTo(PUBLIC_KEY);

        assertThrows(

                IOException.class,

                () -> reader.readKeyChainSnapshot());

        verifyDeserialize(snapshot);

    }

    public void testKeyChainSnapshot_withManyKeysAndProtectionParams() throws Exception {

        PersistentKeyChainSnapshot writer = new PersistentKeyChainSnapshot();

        writer.initWriter();

        KeyDerivationParams derivationParams = KeyDerivationParams.createSha256Params(SALT);

        ArrayList<KeyChainProtectionParams> protectionParamsList = new ArrayList<>();

        protectionParamsList.add(new KeyChainProtectionParams.Builder()

                .setUserSecretType(1)

                .setLockScreenUiFormat(2)

                .setKeyDerivationParams(derivationParams)

                .build());

        protectionParamsList.add(new KeyChainProtectionParams.Builder()

                .setUserSecretType(2)

                .setLockScreenUiFormat(3)

                .setKeyDerivationParams(derivationParams)

                .build());

        ArrayList<WrappedApplicationKey> appKeysList = new ArrayList<>();

        appKeysList.add(new WrappedApplicationKey.Builder()

                .setAlias(ALIAS)

                .setEncryptedKeyMaterial(KEY_MATERIAL)

                .setAccount(ACCOUNT)

                .build());

        appKeysList.add(new WrappedApplicationKey.Builder()

                .setAlias(ALIAS2)

                .setEncryptedKeyMaterial(KEY_MATERIAL)

                .setAccount(ACCOUNT)

                .build());

        KeyChainSnapshot snapshot =  new KeyChainSnapshot.Builder()

                .setSnapshotVersion(SNAPSHOT_VERSION)

                .setKeyChainProtectionParams(protectionParamsList)

                .setEncryptedRecoveryKeyBlob(KEY_MATERIAL)

                .setWrappedApplicationKeys(appKeysList)

                .setMaxAttempts(MAX_ATTEMPTS)

                .setCounterId(COUNTER_ID)

                .setServerParams(SERVER_PARAMS)

                .setTrustedHardwarePublicKey(PUBLIC_KEY)

                .build();

        writer.writeKeyChainSnapshot(snapshot);

        byte[] result = writer.getOutput();

        PersistentKeyChainSnapshot reader = new PersistentKeyChainSnapshot();

        reader.initReader(result);

        KeyChainSnapshot copy = reader.readKeyChainSnapshot();

        assertThat(copy.getSnapshotVersion()).isEqualTo(SNAPSHOT_VERSION);

        assertThat(copy.getKeyChainProtectionParams().get(0).getUserSecretType()).isEqualTo(1);

        assertThat(copy.getEncryptedRecoveryKeyBlob()).isEqualTo(RECOVERY_KEY_MATERIAL);

        assertThat(copy.getWrappedApplicationKeys().get(0).getAlias()).isEqualTo(ALIAS);

        assertThat(copy.getMaxAttempts()).isEqualTo(MAX_ATTEMPTS);

        assertThat(copy.getCounterId()).isEqualTo(COUNTER_ID);

        assertThat(copy.getServerParams()).isEqualTo(SERVER_PARAMS);

        assertThat(copy.getTrustedHardwarePublicKey()).isEqualTo(PUBLIC_KEY);

        assertThrows(

                IOException.class,

                () -> reader.readKeyChainSnapshot());

        verifyDeserialize(snapshot);

    }

    private void verifyDeserialize(KeyChainSnapshot snapshot) throws Exception {

        byte[] serialized = PersistentKeyChainSnapshot.serialize(snapshot);

        KeyChainSnapshot copy = PersistentKeyChainSnapshot.deserialize(serialized);

        assertThat(copy.getSnapshotVersion())

                .isEqualTo(snapshot.getSnapshotVersion());

        assertThat(copy.getKeyChainProtectionParams().size())

                .isEqualTo(copy.getKeyChainProtectionParams().size());

        assertThat(copy.getEncryptedRecoveryKeyBlob())

                .isEqualTo(snapshot.getEncryptedRecoveryKeyBlob());

        assertThat(copy.getWrappedApplicationKeys().size())

                .isEqualTo(snapshot.getWrappedApplicationKeys().size());

        assertThat(copy.getMaxAttempts()).isEqualTo(snapshot.getMaxAttempts());

        assertThat(copy.getCounterId()).isEqualTo(snapshot.getCounterId());

        assertThat(copy.getServerParams()).isEqualTo(snapshot.getServerParams());

        assertThat(copy.getTrustedHardwarePublicKey())

                .isEqualTo(snapshot.getTrustedHardwarePublicKey());

    }

    public void testDeserialize_failsForNewerVersion() throws Exception {

        byte[] newVersion = new byte[]{(byte) 2, (byte) 0, (byte) 0, (byte) 0};

        assertThrows(

                IOException.class,

                () -> PersistentKeyChainSnapshot.deserialize(newVersion));

    }

    public void testDeserialize_failsForEmptyData() throws Exception {

        byte[] empty = new byte[]{};

        assertThrows(

                IOException.class,

                () -> PersistentKeyChainSnapshot.deserialize(empty));

    }

}