/*
 * Copyright (C) 2015 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 android.security.keystore;

import android.annotation.CallSuper;
import android.annotation.NonNull;
import android.os.IBinder;
import android.security.KeyStore;
import android.security.KeyStoreException;
import android.security.keymaster.KeymasterArguments;
import android.security.keymaster.KeymasterDefs;
import android.security.keymaster.OperationResult;

import libcore.util.EmptyArray;

import java.nio.ByteBuffer;
import java.security.InvalidKeyException;
import java.security.InvalidParameterException;
import java.security.PrivateKey;
import java.security.ProviderException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.security.SignatureSpi;

Base class for SignatureSpi implementations of Android KeyStore backed ciphers. 
@hide
/**
 * Base class for {@link SignatureSpi} implementations of Android KeyStore backed ciphers.
 *
 * @hide
 */
abstract class AndroidKeyStoreSignatureSpiBase extends SignatureSpi
        implements KeyStoreCryptoOperation {
    private final KeyStore mKeyStore;

    // Fields below are populated by SignatureSpi.engineInitSign/engineInitVerify and KeyStore.begin
    // and should be preserved after SignatureSpi.engineSign/engineVerify finishes.
    private boolean mSigning;
    private AndroidKeyStoreKey mKey;

    
Token referencing this operation inside keystore service. It is initialized by engineInitSign/engineInitVerify and is invalidated when engineSign/engineVerify succeeds and on some error conditions in between. /**
     * Token referencing this operation inside keystore service. It is initialized by
     * {@code engineInitSign}/{@code engineInitVerify} and is invalidated when
     * {@code engineSign}/{@code engineVerify} succeeds and on some error conditions in between.
     */
    private IBinder mOperationToken;
    private long mOperationHandle;
    private KeyStoreCryptoOperationStreamer mMessageStreamer;

    
Encountered exception which could not be immediately thrown because it was encountered inside a method that does not throw checked exception. This exception will be thrown from engineSign or engineVerify. Once such an exception is encountered, engineUpdate starts ignoring input data. /**
     * Encountered exception which could not be immediately thrown because it was encountered inside
     * a method that does not throw checked exception. This exception will be thrown from
     * {@code engineSign} or {@code engineVerify}. Once such an exception is encountered,
     * {@code engineUpdate} starts ignoring input data.
     */
    private Exception mCachedException;

    AndroidKeyStoreSignatureSpiBase() {
        mKeyStore = KeyStore.getInstance();
    }

    @Override
    protected final void engineInitSign(PrivateKey key) throws InvalidKeyException {
        engineInitSign(key, null);
    }

    @Override
    protected final void engineInitSign(PrivateKey privateKey, SecureRandom random)
            throws InvalidKeyException {
        resetAll();

        boolean success = false;
        try {
            if (privateKey == null) {
                throw new InvalidKeyException("Unsupported key: null");
            }
            AndroidKeyStoreKey keystoreKey;
            if (privateKey instanceof AndroidKeyStorePrivateKey) {
                keystoreKey = (AndroidKeyStoreKey) privateKey;
            } else {
                throw new InvalidKeyException("Unsupported private key type: " + privateKey);
            }
            mSigning = true;
            initKey(keystoreKey);
            appRandom = random;
            ensureKeystoreOperationInitialized();
            success = true;
        } finally {
            if (!success) {
                resetAll();
            }
        }
    }

    @Override
    protected final void engineInitVerify(PublicKey publicKey) throws InvalidKeyException {
        resetAll();

        boolean success = false;
        try {
            if (publicKey == null) {
                throw new InvalidKeyException("Unsupported key: null");
            }
            AndroidKeyStoreKey keystoreKey;
            if (publicKey instanceof AndroidKeyStorePublicKey) {
                keystoreKey = (AndroidKeyStorePublicKey) publicKey;
            } else {
                throw new InvalidKeyException("Unsupported public key type: " + publicKey);
            }
            mSigning = false;
            initKey(keystoreKey);
            appRandom = null;
            ensureKeystoreOperationInitialized();
            success = true;
        } finally {
            if (!success) {
                resetAll();
            }
        }
    }

    
Configures this signature instance to use the provided key.
Throws:
InvalidKeyException – if the key is not suitable./**
     * Configures this signature instance to use the provided key.
     *
     * @throws InvalidKeyException if the {@code key} is not suitable.
     */
    @CallSuper
    protected void initKey(AndroidKeyStoreKey key) throws InvalidKeyException {
        mKey = key;
    }

    
Resets this cipher to its pristine pre-init state. This must be equivalent to obtaining a new
cipher instance.
Subclasses storing additional state should override this method, reset the additional
state, and then chain to superclass.
/**
     * Resets this cipher to its pristine pre-init state. This must be equivalent to obtaining a new
     * cipher instance.
     *
     * <p>Subclasses storing additional state should override this method, reset the additional
     * state, and then chain to superclass.
     */
    @CallSuper
    protected void resetAll() {
        IBinder operationToken = mOperationToken;
        if (operationToken != null) {
            mOperationToken = null;
            mKeyStore.abort(operationToken);
        }
        mSigning = false;
        mKey = null;
        appRandom = null;
        mOperationToken = null;
        mOperationHandle = 0;
        mMessageStreamer = null;
        mCachedException = null;
    }

    
Resets this cipher while preserving the initialized state. This must be equivalent to rolling back the cipher's state to just after the most recent engineInit completed successfully. 
Subclasses storing additional post-init state should override this method, reset the
additional state, and then chain to superclass.
/**
     * Resets this cipher while preserving the initialized state. This must be equivalent to
     * rolling back the cipher's state to just after the most recent {@code engineInit} completed
     * successfully.
     *
     * <p>Subclasses storing additional post-init state should override this method, reset the
     * additional state, and then chain to superclass.
     */
    @CallSuper
    protected void resetWhilePreservingInitState() {
        IBinder operationToken = mOperationToken;
        if (operationToken != null) {
            mOperationToken = null;
            mKeyStore.abort(operationToken);
        }
        mOperationHandle = 0;
        mMessageStreamer = null;
        mCachedException = null;
    }

    private void ensureKeystoreOperationInitialized() throws InvalidKeyException {
        if (mMessageStreamer != null) {
            return;
        }
        if (mCachedException != null) {
            return;
        }
        if (mKey == null) {
            throw new IllegalStateException("Not initialized");
        }

        KeymasterArguments keymasterInputArgs = new KeymasterArguments();
        addAlgorithmSpecificParametersToBegin(keymasterInputArgs);

        OperationResult opResult = mKeyStore.begin(
                mKey.getAlias(),
                mSigning ? KeymasterDefs.KM_PURPOSE_SIGN : KeymasterDefs.KM_PURPOSE_VERIFY,
                true, // permit aborting this operation if keystore runs out of resources
                keymasterInputArgs,
                null, // no additional entropy for begin -- only finish might need some
                mKey.getUid());
        if (opResult == null) {
            throw new KeyStoreConnectException();
        }

        // Store operation token and handle regardless of the error code returned by KeyStore to
        // ensure that the operation gets aborted immediately if the code below throws an exception.
        mOperationToken = opResult.token;
        mOperationHandle = opResult.operationHandle;

        // If necessary, throw an exception due to KeyStore operation having failed.
        InvalidKeyException e = KeyStoreCryptoOperationUtils.getInvalidKeyExceptionForInit(
                mKeyStore, mKey, opResult.resultCode);
        if (e != null) {
            throw e;
        }

        if (mOperationToken == null) {
            throw new ProviderException("Keystore returned null operation token");
        }
        if (mOperationHandle == 0) {
            throw new ProviderException("Keystore returned invalid operation handle");
        }

        mMessageStreamer = createMainDataStreamer(mKeyStore, opResult.token);
    }

    
Creates a streamer which sends the message to be signed/verified into the provided KeyStore
This implementation returns a working streamer.
/**
     * Creates a streamer which sends the message to be signed/verified into the provided KeyStore
     *
     * <p>This implementation returns a working streamer.
     */
    @NonNull
    protected KeyStoreCryptoOperationStreamer createMainDataStreamer(
            KeyStore keyStore, IBinder operationToken) {
        return new KeyStoreCryptoOperationChunkedStreamer(
                new KeyStoreCryptoOperationChunkedStreamer.MainDataStream(
                        keyStore, operationToken));
    }

    @Override
    public final long getOperationHandle() {
        return mOperationHandle;
    }

    @Override
    protected final void engineUpdate(byte[] b, int off, int len) throws SignatureException {
        if (mCachedException != null) {
            throw new SignatureException(mCachedException);
        }

        try {
            ensureKeystoreOperationInitialized();
        } catch (InvalidKeyException e) {
            throw new SignatureException(e);
        }

        if (len == 0) {
            return;
        }

        byte[] output;
        try {
            output = mMessageStreamer.update(b, off, len);
        } catch (KeyStoreException e) {
            throw new SignatureException(e);
        }

        if (output.length != 0) {
            throw new ProviderException(
                    "Update operation unexpectedly produced output: " + output.length + " bytes");
        }
    }

    @Override
    protected final void engineUpdate(byte b) throws SignatureException {
        engineUpdate(new byte[] {b}, 0, 1);
    }

    @Override
    protected final void engineUpdate(ByteBuffer input) {
        byte[] b;
        int off;
        int len = input.remaining();
        if (input.hasArray()) {
            b = input.array();
            off = input.arrayOffset() + input.position();
            input.position(input.limit());
        } else {
            b = new byte[len];
            off = 0;
            input.get(b);
        }

        try {
            engineUpdate(b, off, len);
        } catch (SignatureException e) {
            mCachedException = e;
        }
    }

    @Override
    protected final int engineSign(byte[] out, int outOffset, int outLen)
            throws SignatureException {
        return super.engineSign(out, outOffset, outLen);
    }

    @Override
    protected final byte[] engineSign() throws SignatureException {
        if (mCachedException != null) {
            throw new SignatureException(mCachedException);
        }

        byte[] signature;
        try {
            ensureKeystoreOperationInitialized();

            byte[] additionalEntropy =
                    KeyStoreCryptoOperationUtils.getRandomBytesToMixIntoKeystoreRng(
                            appRandom, getAdditionalEntropyAmountForSign());
            signature = mMessageStreamer.doFinal(
                    EmptyArray.BYTE, 0, 0,
                    null, // no signature provided -- it'll be generated by this invocation
                    additionalEntropy);
        } catch (InvalidKeyException | KeyStoreException e) {
            throw new SignatureException(e);
        }

        resetWhilePreservingInitState();
        return signature;
    }

    @Override
    protected final boolean engineVerify(byte[] signature) throws SignatureException {
        if (mCachedException != null) {
            throw new SignatureException(mCachedException);
        }

        try {
            ensureKeystoreOperationInitialized();
        } catch (InvalidKeyException e) {
            throw new SignatureException(e);
        }

        boolean verified;
        try {
            byte[] output = mMessageStreamer.doFinal(
                    EmptyArray.BYTE, 0, 0,
                    signature,
                    null // no additional entropy needed -- verification is deterministic
                    );
            if (output.length != 0) {
                throw new ProviderException(
                        "Signature verification unexpected produced output: " + output.length
                        + " bytes");
            }
            verified = true;
        } catch (KeyStoreException e) {
            switch (e.getErrorCode()) {
                case KeymasterDefs.KM_ERROR_VERIFICATION_FAILED:
                    verified = false;
                    break;
                default:
                    throw new SignatureException(e);
            }
        }

        resetWhilePreservingInitState();
        return verified;
    }

    @Override
    protected final boolean engineVerify(byte[] sigBytes, int offset, int len)
            throws SignatureException {
        return engineVerify(ArrayUtils.subarray(sigBytes, offset, len));
    }

    @Deprecated
    @Override
    protected final Object engineGetParameter(String param) throws InvalidParameterException {
        throw new InvalidParameterException();
    }

    @Deprecated
    @Override
    protected final void engineSetParameter(String param, Object value)
            throws InvalidParameterException {
        throw new InvalidParameterException();
    }

    protected final KeyStore getKeyStore() {
        return mKeyStore;
    }

    
Returns true if this signature is initialized for signing, false if this signature is initialized for verification. /**
     * Returns {@code true} if this signature is initialized for signing, {@code false} if this
     * signature is initialized for verification.
     */
    protected final boolean isSigning() {
        return mSigning;
    }

    // The methods below need to be implemented by subclasses.

    
Returns the amount of additional entropy (in bytes) to be provided to the KeyStore's finish operation when generating a signature. 
This value should match (or exceed) the amount of Shannon entropy of the produced signature assuming the key and the message are known. For example, for ECDSA signature this should be the size of R, whereas for the RSA signature with PKCS#1 padding this should be 0. /**
     * Returns the amount of additional entropy (in bytes) to be provided to the KeyStore's
     * {@code finish} operation when generating a signature.
     *
     * <p>This value should match (or exceed) the amount of Shannon entropy of the produced
     * signature assuming the key and the message are known. For example, for ECDSA signature this
     * should be the size of {@code R}, whereas for the RSA signature with PKCS#1 padding this
     * should be {@code 0}.
     */
    protected abstract int getAdditionalEntropyAmountForSign();

    
Invoked to add algorithm-specific parameters for the KeyStore's begin operation. 
Params:
keymasterArgs – keystore/keymaster arguments to be populated with algorithm-specific
       parameters./**
     * Invoked to add algorithm-specific parameters for the KeyStore's {@code begin} operation.
     *
     * @param keymasterArgs keystore/keymaster arguments to be populated with algorithm-specific
     *        parameters.
     */
    protected abstract void addAlgorithmSpecificParametersToBegin(
            @NonNull KeymasterArguments keymasterArgs);
}