bcprov/src/main/java/org/bouncycastle/crypto/agreement/SM2KeyExchange.java

package org.bouncycastle.crypto.agreement;

import java.math.BigInteger;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithID;
import org.bouncycastle.crypto.params.SM2KeyExchangePrivateParameters;
import org.bouncycastle.crypto.params.SM2KeyExchangePublicParameters;
import org.bouncycastle.math.ec.ECFieldElement;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.BigIntegers;

/**
 * SM2 Key Exchange protocol - based on https://tools.ietf.org/html/draft-shen-sm2-ecdsa-02
 */
public class SM2KeyExchange
{
    private final Digest digest;

    private byte[] userID;
    private ECPrivateKeyParameters staticKey;
    private ECPoint staticPubPoint;
    private ECPoint ephemeralPubPoint;
    private ECDomainParameters ecParams;
    private int curveLength;
    private int w;
    private ECPrivateKeyParameters ephemeralKey;
    private boolean initiator;

    public SM2KeyExchange()
    {
        this(new SM3Digest());
    }

    public SM2KeyExchange(Digest digest)
    {
        this.digest = digest;
    }

    public void init(
        CipherParameters privParam)
    {
        SM2KeyExchangePrivateParameters baseParam;

        if (privParam instanceof ParametersWithID)
        {
            baseParam = (SM2KeyExchangePrivateParameters)((ParametersWithID)privParam).getParameters();
            userID = ((ParametersWithID)privParam).getID();
        }
        else
        {
            baseParam = (SM2KeyExchangePrivateParameters)privParam;
            userID = new byte[0];
        }

        initiator = baseParam.isInitiator();
        staticKey = baseParam.getStaticPrivateKey();
        ephemeralKey = baseParam.getEphemeralPrivateKey();
        ecParams = staticKey.getParameters();
        staticPubPoint = baseParam.getStaticPublicPoint();
        ephemeralPubPoint = baseParam.getEphemeralPublicPoint();

        curveLength = (ecParams.getCurve().getFieldSize() + 7) / 8;
        w = ecParams.getCurve().getFieldSize() / 2 - 1;
    }

    public int getFieldSize()
    {
        return (staticKey.getParameters().getCurve().getFieldSize() + 7) / 8;
    }

    public byte[] calculateKey(int kLen, CipherParameters pubParam)
    {
        SM2KeyExchangePublicParameters otherPub;
        byte[] otherUserID;

        if (pubParam instanceof ParametersWithID)
        {
            otherPub = (SM2KeyExchangePublicParameters)((ParametersWithID)pubParam).getParameters();
            otherUserID = ((ParametersWithID)pubParam).getID();
        }
        else
        {
            otherPub = (SM2KeyExchangePublicParameters)pubParam;
            otherUserID = new byte[0];
        }

        byte[] za = getZ(digest, userID, staticPubPoint);
        byte[] zb = getZ(digest, otherUserID, otherPub.getStaticPublicKey().getQ());

        ECPoint U = calculateU(otherPub);

        byte[] rv;
        if (initiator)
        {
            rv = kdf(U, za, zb, kLen);
        }
        else
        {
            rv = kdf(U, zb, za, kLen);
        }

        return rv;
    }

    public byte[][] calculateKeyWithConfirmation(int kLen, byte[] confirmationTag, CipherParameters pubParam)
    {
        SM2KeyExchangePublicParameters otherPub;
        byte[] otherUserID;

        if (pubParam instanceof ParametersWithID)
        {
            otherPub = (SM2KeyExchangePublicParameters)((ParametersWithID)pubParam).getParameters();
            otherUserID = ((ParametersWithID)pubParam).getID();
        }
        else
        {
            otherPub = (SM2KeyExchangePublicParameters)pubParam;
            otherUserID = new byte[0];
        }

        if (initiator && confirmationTag == null)
        {
            throw new IllegalArgumentException("if initiating, confirmationTag must be set");
        }
        
        byte[] za = getZ(digest, userID, staticPubPoint);
        byte[] zb = getZ(digest, otherUserID, otherPub.getStaticPublicKey().getQ());

        ECPoint U = calculateU(otherPub);

        byte[] rv;
        if (initiator)
        {
            rv = kdf(U, za, zb, kLen);

            byte[] inner = calculateInnerHash(digest, U, za, zb, ephemeralPubPoint, otherPub.getEphemeralPublicKey().getQ());

            byte[] s1 = S1(digest, U, inner);

            if (!Arrays.constantTimeAreEqual(s1, confirmationTag))
            {
                throw new IllegalStateException("confirmation tag mismatch");
            }
            
            return new byte[][] { rv, S2(digest, U, inner)};
        }
        else
        {
            rv = kdf(U, zb, za, kLen);

            byte[] inner = calculateInnerHash(digest, U, zb, za, otherPub.getEphemeralPublicKey().getQ(), ephemeralPubPoint);

            return new byte[][] { rv, S1(digest, U, inner), S2(digest, U, inner) };
        }
    }

    private ECPoint calculateU(SM2KeyExchangePublicParameters otherPub)
    {
        BigInteger x1 = reduce(ephemeralPubPoint.getAffineXCoord().toBigInteger());

        BigInteger tA = staticKey.getD().add(x1.multiply(ephemeralKey.getD())).mod(ecParams.getN());

        BigInteger x2 = reduce(otherPub.getEphemeralPublicKey().getQ().getAffineXCoord().toBigInteger());

        ECPoint B0 = otherPub.getEphemeralPublicKey().getQ().multiply(x2).normalize();

        ECPoint B1 = otherPub.getStaticPublicKey().getQ().add(B0).normalize();

        return B1.multiply(ecParams.getH().multiply(tA)).normalize();
    }

    private byte[] kdf(ECPoint u, byte[] za, byte[] zb, int klen)
    {
         int ct = 1;
         int v = digest.getDigestSize() * 8;

         byte[] buf = new byte[digest.getDigestSize()];
         byte[] rv = new byte[(klen + 7) / 8];
         int off = 0;

         for (int i = 1; i <= ((klen + v - 1) / v); i++)
         {
             addFieldElement(digest, u.getAffineXCoord());
             addFieldElement(digest, u.getAffineYCoord());
             digest.update(za, 0, za.length);
             digest.update(zb, 0, zb.length);
             digest.update((byte)(ct >> 24));
             digest.update((byte)(ct >> 16));
             digest.update((byte)(ct >> 8));
             digest.update((byte)ct);

             digest.doFinal(buf, 0);

             if (off + buf.length < rv.length)
             {
                 System.arraycopy(buf, 0, rv, off, buf.length);
             }
             else
             {
                 System.arraycopy(buf, 0, rv, off, rv.length - off);
             }

             off += buf.length;
             ct++;
         }

         return rv;
    }

    //x1~=2^w+(x1 AND (2^w-1))
    private BigInteger reduce(BigInteger x)
    {
        return x.and(BigInteger.valueOf(1).shiftLeft(w).subtract(BigInteger.valueOf(1))).setBit(w);
    }

    private byte[] S1(Digest digest, ECPoint u, byte[] inner)
    {
        byte[] rv = new byte[digest.getDigestSize()];

        digest.update((byte)0x02);
        addFieldElement(digest, u.getAffineYCoord());
        digest.update(inner, 0, inner.length);

        digest.doFinal(rv, 0);

        return rv;
    }

    private byte[] calculateInnerHash(Digest digest, ECPoint u, byte[] za, byte[] zb, ECPoint p1, ECPoint p2)
    {
        addFieldElement(digest, u.getAffineXCoord());
        digest.update(za, 0, za.length);
        digest.update(zb, 0, zb.length);
        addFieldElement(digest, p1.getAffineXCoord());
        addFieldElement(digest, p1.getAffineYCoord());
        addFieldElement(digest, p2.getAffineXCoord());
        addFieldElement(digest, p2.getAffineYCoord());

        byte[] rv = new byte[digest.getDigestSize()];

        digest.doFinal(rv, 0);
        return rv;
    }

    private byte[] S2(Digest digest, ECPoint u, byte[] inner)
    {
        byte[] rv = new byte[digest.getDigestSize()];

        digest.update((byte)0x03);
        addFieldElement(digest, u.getAffineYCoord());
        digest.update(inner, 0, inner.length);

        digest.doFinal(rv, 0);

        return rv;
    }

    private byte[] getZ(Digest digest, byte[] userID, ECPoint pubPoint)
    {
        addUserID(digest, userID);

        addFieldElement(digest, ecParams.getCurve().getA());
        addFieldElement(digest, ecParams.getCurve().getB());
        addFieldElement(digest, ecParams.getG().getAffineXCoord());
        addFieldElement(digest, ecParams.getG().getAffineYCoord());
        addFieldElement(digest, pubPoint.getAffineXCoord());
        addFieldElement(digest, pubPoint.getAffineYCoord());

        byte[] rv = new byte[digest.getDigestSize()];

        digest.doFinal(rv, 0);

        return rv;
    }

    private void addUserID(Digest digest, byte[] userID)
    {
        int len = userID.length * 8;

        digest.update((byte)(len >> 8 & 0xFF));
        digest.update((byte)(len & 0xFF));
        digest.update(userID, 0, userID.length);
    }

    private void addFieldElement(Digest digest, ECFieldElement v)
    {
        byte[] p = BigIntegers.asUnsignedByteArray(curveLength, v.toBigInteger());

        digest.update(p, 0, p.length);
    }
}