zei 0.0.10

// Copyright 2019 Stichting Organism
//
// 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.
//! Asset Proofs

use mohan::{
    ser,
    dalek::{
        ristretto::{
            RistrettoPoint, 
            CompressedRistretto
        },
        scalar::Scalar,
        constants::RISTRETTO_BASEPOINT_COMPRESSED,
        traits::Identity
    }
};

use crate::{
    ZeiError,
    core::{
        AssetId,
        OpenAssetOut,
        AssetOut
    },
    proofs::{
        bulletproofs::{
            RangeProof,
            BulletproofGens
        },
        EqualityProof,
        BatchEqualityProof
    },
    PedersenGens
};

use bacteria::Transcript;

//700 + 64log(m)s
/// Proove Me Baby One More Time
pub struct XProofs {
    /// Proove that amounts are non zero and within range
    pub range_proof: RangeProof,
    /// Proove that total amounts in-out are equal
    pub equality_proof: EqualityProof,
    /// Proove that asset types havent mixed
    pub asset_proof: BatchEqualityProof
}

impl Default for XProofs {
    fn default() -> XProofs {
        XProofs {
            range_proof: RangeProof::default(),
            equality_proof: EqualityProof::default(),
            asset_proof: BatchEqualityProof::default(),
        }
    }
}


impl XProofs {

    /// Create Proofs
    pub fn proove(
        inputs: &Vec<OpenAssetOut>, 
        outputs: &Vec<OpenAssetOut>,
    ) -> Result<XProofs, ZeiError> {

        let mut proof = XProofs::default();
        proof.proove_internal(inputs, outputs)?;
        Ok(proof)
    }
    
    /// Create Proofs
    pub fn proove_internal(
        &mut self,
        inputs: &Vec<OpenAssetOut>, 
        outputs: &Vec<OpenAssetOut>,
    ) -> Result<(), ZeiError> {

        self.internal_step_one(inputs, outputs)?;
        self.internal_step_two(inputs, outputs)?;
        self.internal_step_three(inputs, outputs)?;
        Ok(())
    }
     

    /// Equality Proof
    fn internal_step_one(
        &mut self,
        inputs: &Vec<OpenAssetOut>, 
        outputs: &Vec<OpenAssetOut>,
    ) -> Result<(), ZeiError> {
        
        // Calculate Commitments
        let amount_comms_input = inputs.into_iter().fold(
            RistrettoPoint::identity(),
            |r, i| r + i.blinded.value.decompress().unwrap()
        );

        let amount_comms_output = outputs.into_iter().fold(
            RistrettoPoint::identity(),
            |r, i| r + i.blinded.value.decompress().unwrap()
        );

        // Sum Blinds
        let amount_blinds_input = inputs.into_iter().fold(
            Scalar::zero(),
            |r, i| r + i.value_blind
        );

        let amount_blinds_output = outputs.into_iter().fold(
            Scalar::zero(),
            |r, i| r + i.value_blind
        );

        let eq_proof = EqualityProof::proove(
            &amount_comms_input, 
            &amount_blinds_input, 
            &amount_comms_output,
            &amount_blinds_output
        )?;

        self.equality_proof = eq_proof;

        Ok(())
    }

    /// Batch Equality Proof
    fn internal_step_two(
        &mut self,
        inputs: &Vec<OpenAssetOut>, 
        outputs: &Vec<OpenAssetOut>,
    ) -> Result<(), ZeiError> {
        let mut input_asset_comms = Vec::new();
        let mut input_asset_blinds = Vec::new();

        for i in inputs.iter() {
            input_asset_comms.push(i.blinded.flavor.decompress().unwrap());
            input_asset_blinds.push(i.flavor_blind);
        }

        let mut output_asset_comms = Vec::new();
        let mut output_asset_blinds = Vec::new();

        for i in outputs.iter() {
            output_asset_comms.push(i.blinded.flavor.decompress().unwrap());
            output_asset_blinds.push(i.flavor_blind);
        }
        
        let batcheq_proof = BatchEqualityProof::proove(
            &input_asset_comms,
            &input_asset_blinds,
            &output_asset_comms,
            &output_asset_blinds
        )?;

        self.asset_proof = batcheq_proof;

        Ok(())
    }

    /// Range Proofs
    fn internal_step_three(
        &mut self,
        inputs: &Vec<OpenAssetOut>, 
        outputs: &Vec<OpenAssetOut>,
    ) -> Result<(), ZeiError> {
    
        //setup generators
        let pc_gens = PedersenGens::default();
        // Generators for Bulletproofs, valid for proofs up to bitsize 64
        // and aggregation size up to 32. This puts an upper bound on outputs
        let bp_gens = BulletproofGens::new(64, 32);

        if outputs.len() > 32 {
            return Err(ZeiError::OutputOverflowError);
        }

        // create a list of values
        let mut values = outputs.into_iter().map(
            |i| i.value
        ).collect::<Vec<u64>>();
        
        let mut blindings = outputs.into_iter().map(
            |i| i.value_blind
        ).collect::<Vec<Scalar>>();

        // Verification requires a transcript with identical initial state:
        let mut transcript = Transcript::new(b"zei_asset_xproofs");

        // Hack for bulletproofs
        if !blindings.len().is_power_of_two() {
            let gap = blindings.len().next_power_of_two() - blindings.len();
            for i in 0..gap {
                values.push(0u64);
                blindings.push(Scalar::zero());
            }
        }

        //create range proof for outputs
        let (range_proof, value_commitments) = RangeProof::prove_multiple(
                &bp_gens,
                &pc_gens,
                &mut transcript,
                &values,
                &blindings,
                64, //64bit values
        ).unwrap();

        self.range_proof = range_proof;

        Ok(())
    }

    pub fn verify(
        &self,
        inputs: &Vec<AssetOut>,
        outputs: &Vec<AssetOut>
    ) -> Result<(), ZeiError> {
        
        self.verify_step_one(inputs, outputs)?;
        self.verify_step_two(inputs, outputs)?;
        self.verify_step_three(inputs, outputs)?;
        Ok(())
    }

    /// Equality Proof
    pub fn verify_step_one(
        &self,
        inputs: &Vec<AssetOut>,
        outputs: &Vec<AssetOut>
    ) -> Result<(), ZeiError> {

        // Sum Input Ammounts Commitments
        let amount_comms_input = inputs.into_iter().fold(
            RistrettoPoint::identity(),
            |r, i| r + i.value.decompress().unwrap()
        );

        // Sum Output Ammounts Commitments
        let amount_comms_output = outputs.into_iter().fold(
            RistrettoPoint::identity(),
            |r, i| r + i.value.decompress().unwrap()
        );

        self.equality_proof.verify(&amount_comms_input, &amount_comms_output)
       
    }

    /// Verify Batch Equality Proof
    pub fn verify_step_two(
        &self,
        inputs: &Vec<AssetOut>,
        outputs: &Vec<AssetOut>
    ) -> Result<(), ZeiError> {

        let mut input_asset_comms = inputs.into_iter().map(
            |i| i.flavor.decompress().unwrap() 
        ).collect::<Vec<RistrettoPoint>>();

        let mut output_asset_comms = outputs.into_iter().map(
            |i| i.flavor.decompress().unwrap() 
        ).collect::<Vec<RistrettoPoint>>();

        self.asset_proof.verify(&input_asset_comms, &output_asset_comms) 
    }


    /// Range Proofs
    fn verify_step_three(
        &self,
        inputs: &Vec<AssetOut>, 
        outputs: &Vec<AssetOut>,
    ) -> Result<(), ZeiError> {
        // Verification requires a transcript with identical initial state:
        let mut transcript = Transcript::new(b"zei_asset_xproofs");
     
        //setup generators
        let pc_gens = PedersenGens::default();
        // Generators for Bulletproofs, valid for proofs up to bitsize 64
        // and aggregation size up to 32. This puts an upper bound on outputs
        let bp_gens = BulletproofGens::new(64, 32);

        //3. Range Proof Last Most Expensive
        // create a list of values
        let mut values = outputs.into_iter().map(
            |i| i.value
        ).collect::<Vec<CompressedRistretto>>();

        // Hack for bulletproofs
        if !values.len().is_power_of_two() {
            let gap = values.len().next_power_of_two() - values.len();
            for i in 0..gap {
                values.push(CompressedRistretto::identity());
            }
        }

        self.range_proof.verify_multiple(&bp_gens, &pc_gens, &mut transcript, &values, 64)
    }


}


impl ser::Writeable for XProofs {
	fn write<W: ser::Writer>(&self, writer: &mut W) -> Result<(), ser::Error> {
        self.range_proof.write(writer)?;
		self.equality_proof.write(writer)?;
		self.asset_proof.write(writer)?;
		Ok(())
	}
}

impl ser::Readable for XProofs {
	fn read(reader: &mut dyn ser::Reader) -> Result<XProofs, ser::Error> {
		Ok(XProofs {
            range_proof: RangeProof::read(reader)?,
			equality_proof: EqualityProof::read(reader)?,
			asset_proof: BatchEqualityProof::read(reader)?,
		})
	}
}


#[cfg(test)]
mod test {
    use super::*;
    use mohan::fast_merkle_root;
    use std::iter::FromIterator;
    use mohan::{
        mohan_rand,
        hash::H256
    };
    use crate::Lockbox;

    // Helper functions to make the tests easier to read
    // Uses switch commitments
    fn satoshi(q: u64) -> OpenAssetOut {
        let mut transcript = Transcript::new(b"satoshi_land");
        let gens = PedersenGens::default();

        //randomize transcript
        let mut rng = transcript
            .build_rng()
            .finalize(&mut mohan_rand());

        //k , H(kPK || 0), H(kPK || 1)
        let blind_share = Scalar::random(&mut rng);
        let blind_share_value: H256 = fast_merkle_root(vec![H256::from(blind_share.to_bytes()), H256::zero()]);
        let blind_share_asset: H256 = fast_merkle_root(vec![H256::from(blind_share.to_bytes()), H256::from_vec(&vec![1])]);

        let flav_id = AssetId::from_inner(H256::from_vec(b"satoshi"));
        let pc_gens = PedersenGens::default();

        return OpenAssetOut {
            blinded: AssetOut {
                flavor: pc_gens.switch_commit(&flav_id.into_scalar(), &blind_share_asset.into_scalar()).compress(),
                value: pc_gens.switch_commit(&q.into(), &blind_share_value.into_scalar()).compress(),
                vessel: Lockbox::default()
            },
            value: q,
            value_blind: gens.blind_switch(&q.into(), &blind_share_value.into_scalar()),
            flavor_id: flav_id,
            flavor_blind:  gens.blind_switch(&flav_id.into_scalar(), &blind_share_asset.into_scalar()),
        };
    }

    // Helper functions to make the tests easier to read,
    fn turing(q: u64) -> OpenAssetOut {
        let mut transcript = Transcript::new(b"turing_land");

        //randomize transcript
        let mut rng = transcript
            .build_rng()
            .finalize(&mut rand::thread_rng());

        //k , H(kPK || 0), H(kPK || 1)
        let blind_share = Scalar::random(&mut rng);
        let blind_share_value: H256 = fast_merkle_root(vec![H256::from(blind_share.to_bytes()), H256::zero()]);
        let blind_share_asset: H256 = fast_merkle_root(vec![H256::from(blind_share.to_bytes()), H256::from_vec(&vec![1])]);

        let flav_id = AssetId::from_inner(H256::from_vec(b"turing"));
        let pc_gens = PedersenGens::default();

        return OpenAssetOut {
            blinded: AssetOut {
                flavor: pc_gens.commit(&flav_id.into_scalar(), &blind_share_asset.into_scalar()).compress(),
                value: pc_gens.commit(&q.into(), &blind_share_value.into_scalar()).compress(),
                vessel: Lockbox::default()
            },
            value: q,
            value_blind: blind_share_value.into_scalar(),
            flavor_id: flav_id,
            flavor_blind: blind_share_asset.into_scalar(),
        };
    }
   

    #[test]
    fn test_xproof_simple() { 
        let sat = satoshi(1);
        let tsat = satoshi(2);
        let ghost_sat = satoshi(1);

        let p = XProofs::proove(&vec![tsat.clone()], &vec![sat.clone(), ghost_sat.clone()]).unwrap();

        assert!(p.verify(&vec![tsat.blinded], &vec![sat.blinded, ghost_sat.blinded]).is_ok());
    }

    #[test]
    fn test_xproof_simple_excess_out() { 
        let sat = satoshi(11);
        let tsat = satoshi(2);
        let ghost_sat = satoshi(1);

        let p = XProofs::proove(&vec![tsat.clone()], &vec![sat.clone()]).unwrap();

        assert!(p.verify(&vec![tsat.blinded.clone()], &vec![sat.blinded.clone(), ghost_sat.blinded.clone()]).is_err());
    }

    #[test]
    fn test_xproof_simple_excess_out2() { 
        let sat = satoshi(1);
        let tsat = satoshi(21);
        let ghost_sat = satoshi(1);

        let p = XProofs::proove(&vec![tsat.clone()], &vec![sat.clone()]).unwrap();

        assert!(p.verify(&vec![tsat.blinded.clone()], &vec![sat.blinded.clone(), ghost_sat.blinded.clone()]).is_err());
    }

    #[test]
    fn test_xproof_simple_asset_mismatch() { 
        let sat = satoshi(1);
        let tsat = satoshi(2);
        let ghost_sat = turing(2);

        let p = XProofs::proove(&vec![tsat.clone()], &vec![sat.clone(), ghost_sat.clone()]).unwrap();

        assert!(p.verify(&vec![tsat.blinded], &vec![sat.blinded, ghost_sat.blinded]).is_err());
    }

    #[test]
    fn test_xproof_simple_asset_value_mismatch() { 
        let sat = satoshi(10);
        let tsat = satoshi(2);
        let ghost_sat = turing(2);

        let p = XProofs::proove(&vec![tsat.clone()], &vec![sat.clone(), ghost_sat.clone()]).unwrap();

        assert!(p.verify(&vec![tsat.blinded], &vec![sat.blinded, ghost_sat.blinded]).is_err());
    }

    #[test]
    fn test_xproof_12in_4out() { 
        let mut inputs = Vec::new();
        let mut outputs = Vec::new();

        for i in 0..12 {
            inputs.push(satoshi(1));
        }

        for i in 0..4 {
            outputs.push(satoshi(3));
        }

        let p = XProofs::proove(&inputs, &outputs).unwrap();

        let in_blinds = inputs.into_iter().map(|i| i.blinded).collect::<Vec<AssetOut>>();
        let out_blinds = outputs.into_iter().map(|i| i.blinded).collect::<Vec<AssetOut>>();

        assert!(p.verify(&in_blinds, &out_blinds).is_ok());
    }


}