dilithium.go

// Copyright (C) 2020, ISARA Corporation
//
// 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
//
// <a href="http://www.apache.org/licenses/LICENSE-2.0">http://www.apache.org/licenses/LICENSE-2.0</a>
//
// 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 iqrcrypto

// #include <stdio.h>
// #include <stdlib.h>
// #include <errno.h>
// #include "include/iqr_context.h"
// #include "include/iqr_rng.h"
// #include "include/iqr_dilithium.h"
import "C"

import (
	"bytes"
	"crypto"
	"crypto/x509/pkix"
	"encoding/asn1"
	"errors"
	"io"
	"unsafe"
)

var (
	// OidDilithiumSignatureScheme Dilithium-Signature-Scheme
	OidDilithiumSignatureScheme = asn1.ObjectIdentifier{0, 4, 0, 127, 0, 15, 1, 1, 9, 0}
)

var (
	OidDilithium_III_SHAKE_r2 = asn1.RawValue{
		Class:      0,
		Tag:        6,
		IsCompound: false,
		Bytes:      []byte{4, 0, 127, 0, 15, 6, 9, 1},
		FullBytes:  []byte{6, 8, 4, 0, 127, 0, 15, 6, 9, 1},
	}

	OidDilithium_IV_SHAKE_r2 = asn1.RawValue{
		Class:      0,
		Tag:        6,
		IsCompound: false,
		Bytes:      []byte{4, 0, 127, 0, 15, 6, 9, 2},
		FullBytes:  []byte{6, 8, 4, 0, 127, 0, 15, 6, 9, 2},
	}
)

type pkcs8 struct {
	Version    int
	Algo       pkix.AlgorithmIdentifier
	PrivateKey []byte
	// optional attributes omitted.
}

// IqrDILITHIUM128 128 bit quantum security (138 bit classical security) variant.
var IqrDILITHIUM128 = &C.IQR_DILITHIUM_128

// IqrDILITHIUM160 160 bit quantum security (176 bit classical security) variant.
var IqrDILITHIUM160 = &C.IQR_DILITHIUM_160

// IqrDilithiumVariant the Dilithium variant, one of IqrDILITHIUM128 and IqrDILITHIUM160
type IqrDilithiumVariant = C.iqr_DilithiumVariant

// IqrDilithiumParams the Dilithium domain parameter.
type IqrDilithiumParams = C.iqr_DilithiumParams

// IqrDilithiumPrivateKey the Dilithium private key.
type IqrDilithiumPrivateKey = C.iqr_DilithiumPrivateKey

// IqrDilithiumPublicKey the Dilithium public key.
type IqrDilithiumPublicKey = C.iqr_DilithiumPublicKey

// IqrDilithiumCreateParams creates Dilithium's domain parameters.
func IqrDilithiumCreateParams(ctx *IqrContext, variant *IqrDilithiumVariant, params **IqrDilithiumParams) error {
	ret := C.iqr_DilithiumCreateParams(ctx, variant, params)
	return IqrError(ret)
}

// IqrDilithiumDestroyParams destroys Dilithium's domain parameters.
func IqrDilithiumDestroyParams(params **IqrDilithiumParams) error {
	ret := C.iqr_DilithiumDestroyParams(params)
	return IqrError(ret)
}

// IqrDilithiumCreateKeyPair Golang wrapper for iqr_DilithiumCreateKeyPair
func IqrDilithiumCreateKeyPair(params *IqrDilithiumParams, rng *IqrRNG, pub **IqrDilithiumPublicKey, priv **IqrDilithiumPrivateKey) error {
	ret := C.iqr_DilithiumCreateKeyPair(params, rng, pub, priv)
	return IqrError(ret)
}

// IqrDilithiumGetPrivateKeySize Golang wrapper for iqr_DilithiumGetPrivateKeySize
func IqrDilithiumGetPrivateKeySize(params *IqrDilithiumParams, privateKeySize *int64) error {

	var privKeySize sizeT
	ret := C.iqr_DilithiumGetPrivateKeySize(params, &privKeySize)
	*privateKeySize = int64(privKeySize)
	return IqrError(ret)
}

// IqrDilithiumGetPublicKeySize gets the Dilithium public key size.
func IqrDilithiumGetPublicKeySize(params *IqrDilithiumParams, publicKeySize *int64) error {

	var pubKeySize sizeT
	ret := C.iqr_DilithiumGetPublicKeySize(params, &pubKeySize)
	*publicKeySize = int64(pubKeySize)
	return IqrError(ret)
}

// IqrDilithiumGetSignatureSize gets the Dilithium signature size.
func IqrDilithiumGetSignatureSize(params *IqrDilithiumParams, sigSize *int64) error {

	var signatureSize sizeT
	ret := C.iqr_DilithiumGetSignatureSize(params, &signatureSize)
	*sigSize = int64(signatureSize)
	return IqrError(ret)
}

// IqrDilithiumExportPublicKey exports the Dilithium private key's data into a buffer.
func IqrDilithiumExportPublicKey(pubKey *IqrDilithiumPublicKey, buf []byte, size int64) error {
	ret := C.iqr_DilithiumExportPublicKey(pubKey, (*C.uint8_t)(unsafe.Pointer(&buf[0])), sizeT(size))
	return IqrError(ret)
}

func equalOID(a, b asn1.ObjectIdentifier) bool {
	if len(a) != len(b) {
		return false
	}
	for i, v := range a {
		if v != b[i] {
			return false
		}
	}
	return true
}

// IqrDilithiumImportPublicKey imports a Dilithium private key object from a buffer.
// This function only accepts data that has been generated by
// iqr_DilithiumExportPrivateKey(). There is currently no standard for saving
// Dilithium private keys.
func IqrDilithiumImportPublicKey(params *IqrDilithiumParams, buf []byte, size int64, publicKey **IqrDilithiumPublicKey) error {
	ret := C.iqr_DilithiumImportPublicKey(params, (*C.uint8_t)(unsafe.Pointer(&buf[0])), sizeT(size), publicKey)
	return IqrError(ret)
}

// IqrDilithiumImportPublicKeyFromASN1 imports public key from a asn.1 encoded public key buffer.
func IqrDilithiumImportPublicKeyFromASN1(ctx *IqrContext, der []byte, size int64, publicKey **IqrDilithiumPublicKey,
	variant **IqrDilithiumVariant, params **IqrDilithiumParams) error {

	if len(der) != int(size) {
		return errors.New("Size not matched")
	}

	pubkey := pkixPublicKey{}
	_, err := asn1.Unmarshal(der, &pubkey)
	if err != nil {
		return err
	}

	algo := pubkey.Algo.Algorithm
	//Check algo is dilithium
	if !equalOID(algo, OidDilithiumSignatureScheme) {
		return errors.New("Unexpected algorithm")
	}
	parameters := pubkey.Algo.Parameters

	// Find out the variant from the parameters.
	if bytes.Equal(OidDilithium_III_SHAKE_r2.Bytes, parameters.Bytes) {
		*variant = IqrDILITHIUM128
	} else if bytes.Equal(OidDilithium_IV_SHAKE_r2.Bytes, parameters.Bytes) {
		*variant = IqrDILITHIUM160
	} else {
		//TODO: compare other variants.
		return errors.New("Unexpected variant")
	}

	err = IqrDilithiumCreateParams(ctx, *variant, params)
	if err != nil {
		return err
	}

	var publicKeySize int64
	err = IqrDilithiumGetPublicKeySize(*params, &publicKeySize)
	if err != nil {
		return err
	}
	*publicKey = nil
	buf := pubkey.BitString.Bytes
	bufLen := len(buf)
	if int(publicKeySize) > bufLen {
		return errors.New("Incorrect public key size")
	}

	var keyData []byte
	_, err = asn1.Unmarshal(buf, &keyData)
	err = IqrDilithiumImportPublicKey(*params, keyData, publicKeySize, publicKey)
	return err
}

// IqrDilithiumExportPrivateKey exports a Dilithium private key.
func IqrDilithiumExportPrivateKey(privKey *IqrDilithiumPrivateKey, buf []byte, size int64) error {
	ret := C.iqr_DilithiumExportPrivateKey(privKey, (*C.uint8_t)(unsafe.Pointer(&buf[0])), sizeT(size))
	return IqrError(ret)
}

// IqrDilithiumExportPrivateKeyPKCS8 exports a Dilithium private key into PKCS8 format.
func IqrDilithiumExportPrivateKeyPKCS8(key *DilithiumPrivateKey) (der []byte, err error) {
	params := key.Params
	variant := key.Variant
	var keySizePriv int64
	err = IqrDilithiumGetPrivateKeySize(params, &keySizePriv)
	if err != nil {
		return nil, err
	}

	bufPriv := make([]byte, int(keySizePriv))
	err = IqrDilithiumExportPrivateKey(key.PrivKey, bufPriv, keySizePriv)
	if err != nil {
		return nil, err
	}

	var keySizePub int64
	err = IqrDilithiumGetPublicKeySize(params, &keySizePub)
	if err != nil {
		return nil, err
	}

	bufPub := make([]byte, int(keySizePub))
	err = IqrDilithiumExportPublicKey(key.PubKey, bufPub, keySizePub)
	if err != nil {
		return nil, err
	}

	var algoParam asn1.RawValue
	if variant == IqrDILITHIUM128 {
		algoParam = OidDilithium_III_SHAKE_r2
	} else if variant == IqrDILITHIUM160 {
		algoParam = OidDilithium_IV_SHAKE_r2
	}

	type internalDilithium struct {
		Version int
		KeyPri  []byte
		KeyPub  []byte
	}
	privKeyData := internalDilithium{
		Version: 0,
		KeyPri:  bufPriv,
		KeyPub:  bufPub,
	}
	bufDer, err := asn1.Marshal(privKeyData)
	if err != nil {
		return nil, err
	}

	privPkcs8 := pkcs8{
		Version: 0,
		Algo: pkix.AlgorithmIdentifier{
			Algorithm:  OidDilithiumSignatureScheme,
			Parameters: algoParam,
		},
		PrivateKey: bufDer,
	}
	der, err = asn1.Marshal(privPkcs8)
	if err != nil {
		return nil, err
	}
	return der, nil
}

// IqrDilithiumImportPrivateKey imports a Dilithium private key object from a buffer.
// This function only accepts data that has been generated by
// IqrDilithiumExportPrivateKey(). There is currently no standard for saving
// Dilithium private keys.
func IqrDilithiumImportPrivateKey(params *IqrDilithiumParams, buf []byte, size int64, privateKey **IqrDilithiumPrivateKey) error {
	ret := C.iqr_DilithiumImportPrivateKey(params, (*C.uint8_t)(unsafe.Pointer(&buf[0])), sizeT(size), privateKey)
	return IqrError(ret)
}

// IqrDilithiumImportPrivateKeyFromPKCS8
func IqrDilithiumImportPrivateKeyFromPKCS8(ctx *IqrContext, der []byte, size int64, privateKey **IqrDilithiumPrivateKey,
	variant **IqrDilithiumVariant, params **IqrDilithiumParams) error {

	if len(der) != int(size) {
		return errors.New("Size not matched")
	}

	privKey := pkcs8{}
	_, err := asn1.Unmarshal(der, &privKey)
	if err != nil {
		return err
	}

	algo := privKey.Algo.Algorithm
	//Check algo is dilithium
	if !equalOID(algo, OidDilithiumSignatureScheme) {
		return errors.New("Unexpected algorithm")
	}
	parameters := privKey.Algo.Parameters

	// Find out the variant from the parameters.
	if bytes.Equal(OidDilithium_III_SHAKE_r2.Bytes, parameters.Bytes) {
		*variant = IqrDILITHIUM128
	} else if bytes.Equal(OidDilithium_IV_SHAKE_r2.Bytes, parameters.Bytes) {
		*variant = IqrDILITHIUM160
	} else {
		//TODO: compare with other variants.
		return errors.New("Unexpected variant")
	}

	err = IqrDilithiumCreateParams(ctx, *variant, params)
	if err != nil {
		return err
	}

	var privateKeySize int64
	IqrDilithiumGetPrivateKeySize(*params, &privateKeySize)
	if int(privateKeySize) > len(privKey.PrivateKey) {
		return errors.New("Incorrect private key size")
	}

	*privateKey = nil
	buf := privKey.PrivateKey
	bufSize := privateKeySize

	type internalDilithium struct {
		Version int
		KeyData []byte
	}
	var privKeyData internalDilithium
	_, err = asn1.Unmarshal(buf, &privKeyData)
	if err != nil {
		return err
	}
	err = IqrDilithiumImportPrivateKey(*params, privKeyData.KeyData, bufSize, privateKey)
	return err
}

// IqrDilithiumSign signs a message using a Dilithium private key.
// sigSize must be exact the same value returned by iqrDilithiumGetSingatureSize
func IqrDilithiumSign(privKey *IqrDilithiumPrivateKey, message []byte, messageSize int64, sig []byte, sigSize int64) error {
	ret := C.iqr_DilithiumSign(privKey, (*C.uint8_t)(unsafe.Pointer(&message[0])), sizeT(messageSize),
		(*C.uint8_t)(unsafe.Pointer(&sig[0])), sizeT(sigSize))
	return IqrError(ret)
}

// IqrDilithiumVerify verifies the signature of a message using a Dilithium public key.
// sigSize must be exact the same value returned by iqrDilithiumGetSingatureSize
func IqrDilithiumVerify(pubKey *IqrDilithiumPublicKey, message []byte, messageSize int64, sig []byte, sigSize int64) error {
	ret := C.iqr_DilithiumVerify(pubKey, (*C.uint8_t)(unsafe.Pointer(&message[0])), sizeT(messageSize),
		(*C.uint8_t)(unsafe.Pointer(&sig[0])), sizeT(sigSize))
	return IqrError(ret)
}

// IqrDilithiumDestroyPublicKey destroys a Dilithium public key.
func IqrDilithiumDestroyPublicKey(publicKey **IqrDilithiumPublicKey) error {
	ret := C.iqr_DilithiumDestroyPublicKey(publicKey)
	return IqrError(ret)
}

// IqrDilithiumDestroyPrivateKey destroys a Dilithium private key.
func IqrDilithiumDestroyPrivateKey(privateKey **IqrDilithiumPrivateKey) error {
	ret := C.iqr_DilithiumDestroyPrivateKey(privateKey)
	return IqrError(ret)
}

// A DilithiumPrivateKey represents a Dilithium key
// It also implements the crypto.Signer interface.
type DilithiumPrivateKey struct {
	Variant *IqrDilithiumVariant
	Params  *IqrDilithiumParams
	PubKey  *IqrDilithiumPublicKey
	PrivKey *IqrDilithiumPrivateKey
}

// QSKeyType returns wether is the key is Quantum-Safe.
func (priv *DilithiumPrivateKey) QSKeyType() string {
	return "dilithium"
}

// Public returns the public key corresponding to priv.
func (priv *DilithiumPrivateKey) Public() crypto.PublicKey {
	return priv.PubKey
}

// Sign signs digest with priv, reading randomness from rand.
func (priv *DilithiumPrivateKey) Sign(rand io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
	var qsSigSize int64
	err := IqrDilithiumGetSignatureSize(priv.Params, &qsSigSize)
	if err != nil {
		return nil, err
	}

	var qsSig = make([]byte, int(qsSigSize))
	err = IqrDilithiumSign(priv.PrivKey, digest, int64(len(digest)), qsSig, qsSigSize)
	if err != nil {
		return nil, err
	}
	return qsSig, nil
}

// Destroy destroy the private and its public key and release the memory
func (priv *DilithiumPrivateKey) Destroy() error {

	if priv.Params != nil {
		IqrDilithiumDestroyParams(&priv.Params)
	}

	if priv.PubKey != nil {
		IqrDilithiumDestroyPublicKey(&priv.PubKey)
	}

	if priv.PubKey != nil {
		IqrDilithiumDestroyPrivateKey(&priv.PrivKey)
	}

	return nil
}

// GenerateDilithiumPrivateKey generates a new dilithium private key using the vaiant specified
// Call IqrDilithiumDestroyPrivateKey if the returned private key no longer used.
// variant: one of IqrDILITHIUM128 and IqrDILITHIUM160
func GenerateDilithiumPrivateKey(variant *IqrDilithiumVariant, rand io.Reader) (*DilithiumPrivateKey, error) {
	var ctx *IqrContext
	var rng *IqrRNG

	err := IqrCreateContext(&ctx)
	if err != nil {
		return nil, err
	}
	defer IqrDestroyContext(&ctx)

	err = IqrInitRNG(&ctx, &rng, rand)
	if err != nil {
		return nil, err
	}
	defer IqrRNGDestroy(&rng)

	var params *IqrDilithiumParams
	var priv *IqrDilithiumPrivateKey
	var pub *IqrDilithiumPublicKey

	err = IqrDilithiumCreateParams(ctx, variant, &params)
	if err != nil {
		return nil, err
	}

	err = IqrDilithiumCreateKeyPair(params, rng, &pub, &priv)
	if err != nil {
		return nil, err
	}

	privateKey := new(DilithiumPrivateKey)
	privateKey.Params = params
	privateKey.PubKey = pub
	privateKey.PrivKey = priv
	privateKey.Variant = variant
	return privateKey, nil
}