// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! An implementation of a verifiable oblivious pseudorandom function (VOPRF)
//!
//! Note: This implementation is in sync with
//! [draft-irtf-cfrg-voprf-19](https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-19.html),
//! but this specification is subject to change, until the final version
//! published by the IETF.
//!
//! # Overview
//!
//! A verifiable oblivious pseudorandom function is a protocol that is evaluated
//! between a client and a server. They must first agree on a finite cyclic
//! group along with a point representation.
//!
//! We will use the following choice in this example:
//!
//! ```ignore
//! type CipherSuite = voprf::Ristretto255;
//! ```
//!
//! ## Modes of Operation
//!
//! VOPRF can be used in three modes:
//! - [Base Mode](#base-mode), which corresponds to a normal OPRF evaluation
//!   with no support for the verification of the OPRF outputs
//! - [Verifiable Mode](#verifiable-mode), which corresponds to an OPRF
//!   evaluation where the outputs can be verified against a server public key
//!   (VOPRF)
//! - [Partially Oblivious Verifiable Mode](#metadata), which corresponds to a
//!   VOPRF, where a public input can be supplied to the PRF computation
//!
//! In all of these modes, the protocol begins with a client blinding, followed
//! by a server evaluation, and finishes with a client finalization and server
//! evaluation.
//!
//! ## Base Mode
//!
//! In base mode, an [OprfClient] interacts with an [OprfServer] to compute the
//! output of the OPRF.
//!
//! ### Server Setup
//!
//! The protocol begins with a setup phase, in which the server must run
//! [OprfServer::new()] to produce an instance of itself. This instance must be
//! persisted on the server and used for online client evaluations.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! use rand::rngs::OsRng;
//! use rand::RngCore;
//! use voprf::OprfServer;
//!
//! let mut server_rng = OsRng;
//! let server = OprfServer::<CipherSuite>::new(&mut server_rng);
//! ```
//!
//! ### Client Blinding
//!
//! In the first step, the client chooses an input, and runs [OprfClient::blind]
//! to produce an [OprfClientBlindResult], which consists of a [BlindedElement]
//! to be sent to the server and an [OprfClient] which must be persisted on the
//! client for the final step of the VOPRF protocol.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! use rand::rngs::OsRng;
//! use rand::RngCore;
//! use voprf::OprfClient;
//!
//! let mut client_rng = OsRng;
//! let client_blind_result = OprfClient::<CipherSuite>::blind(b"input", &mut client_rng)
//!     .expect("Unable to construct client");
//! ```
//!
//! ### Server Blind Evaluation
//!
//! In the second step, the server takes as input the message from
//! [OprfClient::blind] (a [BlindedElement]), and runs
//! [OprfServer::blind_evaluate] to produce [EvaluationElement] to be sent to
//! the client.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::OprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = OprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::OprfServer;
//! # let mut server_rng = OsRng;
//! # let server = OprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! let server_evaluate_result = server.blind_evaluate(&client_blind_result.message);
//! ```
//!
//! ### Client Finalization
//!
//! In the final step on the client side, the client takes as input the message
//! from [OprfServer::evaluate] (an [EvaluationElement]), and runs
//! [OprfClient::finalize] to produce an output for the protocol.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::OprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = OprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::OprfServer;
//! # let mut server_rng = OsRng;
//! # let server = OprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! # let message = server.blind_evaluate(&client_blind_result.message);
//! let client_finalize_result = client_blind_result
//!     .state
//!     .finalize(b"input", &message)
//!     .expect("Unable to perform client finalization");
//!
//! println!("VOPRF output: {:?}", client_finalize_result.to_vec());
//! ```
//!
//! ### Server Evaluation
//!
//! Optionally, if the server has direct access to the PRF input, then it need
//! not perform the oblivious computation and can simply run
//! [OprfServer::evaluate] to generate an output which matches the output
//! produced by an execution of the oblivious protocol on the same input and
//! key.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::OprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = OprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::OprfServer;
//! # let mut server_rng = OsRng;
//! # let server = OprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! # let message = server.blind_evaluate(&client_blind_result.message);
//! let client_finalize_result = client_blind_result
//!     .state
//!     .finalize(b"input", &message)
//!     .expect("Unable to perform client finalization");
//!
//! let server_evaluate_result = server
//!     .evaluate(b"input")
//!     .expect("Unable to perform the server evaluation");
//!
//! assert_eq!(client_finalize_result, server_evaluate_result);
//! ```
//!
//! ## Verifiable Mode
//!
//! In verifiable mode, a [VoprfClient] interacts with a [VoprfServer] to
//! compute the output of the VOPRF. In order to verify the server's
//! computation, the client checks a server-generated proof against the server's
//! public key. If the proof fails to verify, then the client does not receive
//! an output.
//!
//! In batch mode, a single proof can be used for multiple VOPRF evaluations.
//! See [the batching section](#batching) for more details on how to perform
//! batch evaluations.
//!
//! ### Server Setup
//!
//! The protocol begins with a setup phase, in which the server must run
//! [VoprfServer::new()] to produce an instance of itself. This instance must be
//! persisted on the server and used for online client evaluations.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! use rand::rngs::OsRng;
//! use rand::RngCore;
//! use voprf::VoprfServer;
//!
//! let mut server_rng = OsRng;
//! let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//!
//! // To be sent to the client
//! println!("Server public key: {:?}", server.get_public_key());
//! ```
//!
//! The public key should be sent to the client, since the client will need it
//! in the final step of the protocol in order to complete the evaluation of the
//! VOPRF.
//!
//! ### Client Blinding
//!
//! In the first step, the client chooses an input, and runs
//! [VoprfClient::blind] to produce a [VoprfClientBlindResult], which consists
//! of a [BlindedElement] to be sent to the server and a [VoprfClient] which
//! must be persisted on the client for the final step of the VOPRF protocol.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! use rand::rngs::OsRng;
//! use rand::RngCore;
//! use voprf::VoprfClient;
//!
//! let mut client_rng = OsRng;
//! let client_blind_result = VoprfClient::<CipherSuite>::blind(b"input", &mut client_rng)
//!     .expect("Unable to construct client");
//! ```
//!
//! ### Server Blind Evaluation
//!
//! In the second step, the server takes as input the message from
//! [VoprfClient::blind] (a [BlindedElement]), and runs
//! [VoprfServer::blind_evaluate] to produce a [VoprfServerEvaluateResult],
//! which consists of an [EvaluationElement] to be sent to the client along with
//! a proof.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::{VoprfServerEvaluateResult, VoprfClient};
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::VoprfServer;
//! # let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! let VoprfServerEvaluateResult { message, proof } =
//!     server.blind_evaluate(&mut server_rng, &client_blind_result.message);
//! ```
//!
//! ### Client Finalization
//!
//! In the final step, the client takes as input the message from
//! [VoprfServer::blind_evaluate] (an [EvaluationElement]), the proof, and the
//! server's public key, and runs [VoprfClient::finalize] to produce an output
//! for the protocol.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::VoprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::VoprfServer;
//! # let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! # let server_evaluate_result = server.blind_evaluate(
//! #     &mut server_rng,
//! #     &client_blind_result.message,
//! # );
//! let client_finalize_result = client_blind_result
//!     .state
//!     .finalize(
//!         b"input",
//!         &server_evaluate_result.message,
//!         &server_evaluate_result.proof,
//!         server.get_public_key(),
//!     )
//!     .expect("Unable to perform client finalization");
//!
//! println!("VOPRF output: {:?}", client_finalize_result.to_vec());
//! ```
//!
//! ### Server Evaluation
//!
//! Optionally, if the server has direct access to the PRF input, then it need
//! not perform the oblivious computation and can simply run
//! [VoprfServer::evaluate] to generate an output which matches the output
//! produced by an execution of the oblivious protocol on the same input and
//! key.
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::VoprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #     b"input",
//! #     &mut client_rng,
//! # ).expect("Unable to construct client");
//! # use voprf::VoprfServer;
//! # let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! # let server_evaluate_result = server.blind_evaluate(
//! #     &mut server_rng,
//! #     &client_blind_result.message,
//! # );
//! let client_finalize_result = client_blind_result
//!     .state
//!     .finalize(
//!         b"input",
//!         &server_evaluate_result.message,
//!         &server_evaluate_result.proof,
//!         server.get_public_key(),
//!     )
//!     .expect("Unable to perform client finalization");
//!
//! let server_evaluate_result = server
//!     .evaluate(b"input")
//!     .expect("Unable to perform the server evaluation");
//!
//! assert_eq!(client_finalize_result, server_evaluate_result);
//! ```
//!
//! # Advanced Usage
//!
//! There are two additional (and optional) extensions to the core VOPRF
//! protocol: support for batching of evaluations, and support for public
//! metadata.
//!
//! ## Batching
//!
//! It is sometimes desirable to generate only a single, constant-size proof for
//! an unbounded number of VOPRF evaluations (on arbitrary inputs).
//! [VoprfClient] and [VoprfServer] support a batch API for handling this case.
//! In the following example, we show how to use the batch API to produce a
//! single proof for 10 parallel VOPRF evaluations.
//!
//! First, the client produces 10 blindings, storing their resulting states and
//! messages:
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::VoprfClient;
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! let mut client_rng = OsRng;
//! let mut client_states = vec![];
//! let mut client_messages = vec![];
//! for _ in 0..10 {
//!     let client_blind_result = VoprfClient::<CipherSuite>::blind(b"input", &mut client_rng)
//!         .expect("Unable to construct client");
//!     client_states.push(client_blind_result.state);
//!     client_messages.push(client_blind_result.message);
//! }
//! ```
//!
//! Next, the server calls the [VoprfServer::batch_blind_evaluate_prepare] and
//! [VoprfServer::batch_blind_evaluate_finish] function on a set of client
//! messages, to produce a corresponding set of messages to be returned to the
//! client (returned in the same order), along with a single proof:
//!
//! ```
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::{VoprfServerBatchEvaluateFinishResult, VoprfClient};
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let mut client_states = vec![];
//! # let mut client_messages = vec![];
//! # for _ in 0..10 {
//! #     let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #         b"input",
//! #        &mut client_rng,
//! #     ).expect("Unable to construct client");
//! #     client_states.push(client_blind_result.state);
//! #     client_messages.push(client_blind_result.message);
//! # }
//! # use voprf::VoprfServer;
//! let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! let prepared_evaluation_elements = server.batch_blind_evaluate_prepare(client_messages.iter());
//! let prepared_elements: Vec<_> = prepared_evaluation_elements.collect();
//! let VoprfServerBatchEvaluateFinishResult { messages, proof } = server
//!     .batch_blind_evaluate_finish(&mut server_rng, client_messages.iter(), &prepared_elements)
//!     .expect("Unable to perform server batch evaluate");
//! let messages: Vec<_> = messages.collect();
//! ```
//!
//! If `alloc` is available, `VoprfServer::batch_blind_evaluate` can be called
//! to avoid having to collect output manually:
//!
//! ```
//! # #[cfg(feature = "alloc")] {
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::{VoprfServerBatchEvaluateResult, VoprfClient};
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let mut client_states = vec![];
//! # let mut client_messages = vec![];
//! # for _ in 0..10 {
//! #     let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #         b"input",
//! #        &mut client_rng,
//! #     ).expect("Unable to construct client");
//! #     client_states.push(client_blind_result.state);
//! #     client_messages.push(client_blind_result.message);
//! # }
//! # use voprf::VoprfServer;
//! let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! let VoprfServerBatchEvaluateResult { messages, proof } = server
//!     .batch_blind_evaluate(&mut server_rng, &client_messages)
//!     .expect("Unable to perform server batch evaluate");
//! # }
//! ```
//!
//! Then, the client calls [VoprfClient::batch_finalize] on the client states
//! saved from the first step, along with the messages returned by the server,
//! along with the server's proof, in order to produce a vector of outputs if
//! the proof verifies correctly.
//!
//! ```
//! # #[cfg(feature = "alloc")] {
//! # #[cfg(feature = "ristretto255")]
//! # type CipherSuite = voprf::Ristretto255;
//! # #[cfg(not(feature = "ristretto255"))]
//! # type CipherSuite = p256::NistP256;
//! # use voprf::{VoprfServerBatchEvaluateResult, VoprfClient};
//! # use rand::{rngs::OsRng, RngCore};
//! #
//! # let mut client_rng = OsRng;
//! # let mut client_states = vec![];
//! # let mut client_messages = vec![];
//! # for _ in 0..10 {
//! #     let client_blind_result = VoprfClient::<CipherSuite>::blind(
//! #         b"input",
//! #        &mut client_rng,
//! #     ).expect("Unable to construct client");
//! #     client_states.push(client_blind_result.state);
//! #     client_messages.push(client_blind_result.message);
//! # }
//! # use voprf::VoprfServer;
//! # let mut server_rng = OsRng;
//! # let server = VoprfServer::<CipherSuite>::new(&mut server_rng).unwrap();
//! # let VoprfServerBatchEvaluateResult { messages, proof } = server
//! #     .batch_blind_evaluate(&mut server_rng, &client_messages)
//! #     .expect("Unable to perform server batch evaluate");
//! let client_batch_finalize_result = VoprfClient::batch_finalize(
//!     &[b"input"; 10],
//!     &client_states,
//!     &messages,
//!     &proof,
//!     server.get_public_key(),
//! )
//! .expect("Unable to perform client batch finalization")
//! .collect::<Vec<_>>();
//!
//! println!("VOPRF batch outputs: {:?}", client_batch_finalize_result);
//! # }
//! ```
//!
//! ## Metadata
//!
//! The optional metadata parameter included in the POPRF mode allows clients
//! and servers to cryptographically bind additional data to the VOPRF output.
//! This metadata is known to both parties at the start of the protocol, and is
//! inserted under the server's blind evaluate step and the client's finalize
//! step. This metadata can be constructed with some type of higher-level domain
//! separation to avoid cross-protocol attacks or related issues.
//!
//! The API for POPRF mode is similar to VOPRF mode, except that a [PoprfServer]
//! and [PoprfClient] are used, and that each of the functions accept an
//! additional (and optional) info parameter which represents the public input.
//! See
//! <https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-11.html#name-poprf-public-input>
//! for more detailed information on how this public input should be used.
//!
//! # Features
//!
//! - The `alloc` feature requires Rust's `alloc` crate and enables batching
//!   VOPRF evaluations.
//!
//! - The `serde` feature, enabled by default, provides convenience functions
//!   for serializing and deserializing with [serde](https://serde.rs/).
//!
//! - The `danger` feature, disabled by default, exposes functions for setting
//!   and getting internal values not available in the default API. These
//!   functions are intended for use in by higher-level cryptographic protocols
//!   that need access to these raw values and are able to perform the necessary
//!   validations on them (such as being valid group elements).
//!
//! - The `ristretto255-ciphersuite` features enables using [`Ristretto255`] as
//!   a [`CipherSuite`].
//!
//! - The `ristretto255` feature enables using [`Ristretto255`] as the
//!   underlying group for the [Group] choice. To select a specific backend see
//!   the [curve25519-dalek] documentation.
//!
//! [curve25519-dalek]:
//!     (https://docs.rs/curve25519-dalek/4.0.0-pre.5/curve25519_dalek/index.html#backends)

#![no_std]
#![cfg_attr(docsrs, feature(doc_auto_cfg))]
#![cfg_attr(not(test), deny(unsafe_code))]
#![warn(
    clippy::cargo,
    clippy::missing_errors_doc,
    missing_debug_implementations,
    missing_docs
)]
#![allow(clippy::multiple_crate_versions)]

#[cfg(any(feature = "alloc", test))]
extern crate alloc;

#[cfg(feature = "std")]
extern crate std;

mod ciphersuite;
mod common;
mod error;
mod group;
mod oprf;
mod poprf;
mod serialization;
mod voprf;

#[cfg(test)]
mod tests;

// Exports

pub use crate::ciphersuite::CipherSuite;
#[cfg(feature = "danger")]
pub use crate::common::derive_key;
pub use crate::common::{
    BlindedElement, EvaluationElement, Mode, PreparedEvaluationElement, Proof,
};
pub use crate::error::{Error, InternalError, Result};
pub use crate::group::Group;
#[cfg(feature = "ristretto255")]
pub use crate::group::Ristretto255;
pub use crate::oprf::{OprfClient, OprfClientBlindResult, OprfServer};
#[cfg(feature = "alloc")]
pub use crate::poprf::PoprfServerBatchEvaluateResult;
pub use crate::poprf::{
    PoprfClient, PoprfClientBatchFinalizeResult, PoprfPreparedTweak, PoprfServer,
    PoprfServerBatchEvaluateFinishResult, PoprfServerBatchEvaluateFinishedMessages,
    PoprfServerBatchEvaluatePrepareResult, PoprfServerBatchEvaluatePreparedEvaluationElements,
};
pub use crate::serialization::{
    BlindedElementLen, EvaluationElementLen, OprfClientLen, OprfServerLen, PoprfClientLen,
    PoprfServerLen, ProofLen, VoprfClientLen, VoprfServerLen,
};
#[cfg(feature = "alloc")]
pub use crate::voprf::VoprfServerBatchEvaluateResult;
pub use crate::voprf::{
    VoprfClient, VoprfClientBatchFinalizeResult, VoprfClientBlindResult, VoprfServer,
    VoprfServerBatchEvaluateFinishResult, VoprfServerBatchEvaluateFinishedMessages,
    VoprfServerBatchEvaluatePreparedEvaluationElements, VoprfServerEvaluateResult,
};


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Contains the main OPRF API

use core::iter::{self, Map};

use derive_where::derive_where;
use digest::core_api::BlockSizeUser;
use digest::{Digest, Output, OutputSizeUser};
use generic_array::typenum::{IsLess, IsLessOrEqual, Unsigned, U256};
use generic_array::GenericArray;
use rand_core::{CryptoRng, RngCore};

use crate::common::{
    derive_key_internal, deterministic_blind_unchecked, hash_to_group, i2osp_2,
    server_evaluate_hash_input, BlindedElement, EvaluationElement, Mode, STR_FINALIZE,
};
#[cfg(feature = "serde")]
use crate::serialization::serde::Scalar;
use crate::{CipherSuite, Error, Group, Result};

///////////////
// Constants //
// ========= //
///////////////

////////////////////////////
// High-level API Structs //
// ====================== //
////////////////////////////

/// A client which engages with a [OprfServer] in base mode, meaning
/// that the OPRF outputs are not verifiable.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct OprfClient<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) blind: <CS::Group as Group>::Scalar,
}

/// A server which engages with a [OprfClient] in base mode, meaning
/// that the OPRF outputs are not verifiable.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct OprfServer<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) sk: <CS::Group as Group>::Scalar,
}

/////////////////////////
// API Implementations //
// =================== //
/////////////////////////

impl<CS: CipherSuite> OprfClient<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Computes the first step for the multiplicative blinding version of
    /// DH-OPRF.
    ///
    /// # Errors
    /// [`Error::Input`] if the `input` is empty or longer then [`u16::MAX`].
    pub fn blind<R: RngCore + CryptoRng>(
        input: &[u8],
        blinding_factor_rng: &mut R,
    ) -> Result<OprfClientBlindResult<CS>> {
        let blind = CS::Group::random_scalar(blinding_factor_rng);
        Self::deterministic_blind_unchecked_inner(input, blind)
    }

    /// Computes the first step for the multiplicative blinding version of
    /// DH-OPRF, taking a blinding factor scalar as input instead of sampling
    /// from an RNG.
    ///
    /// # Caution
    ///
    /// This should be used with caution, since it does not perform any checks
    /// on the validity of the blinding factor!
    ///
    /// # Errors
    /// [`Error::Input`] if the `input` is empty or longer then [`u16::MAX`].
    #[cfg(any(feature = "danger", test))]
    pub fn deterministic_blind_unchecked(
        input: &[u8],
        blind: <CS::Group as Group>::Scalar,
    ) -> Result<OprfClientBlindResult<CS>> {
        Self::deterministic_blind_unchecked_inner(input, blind)
    }

    /// Can only fail with [`Error::Input`].
    fn deterministic_blind_unchecked_inner(
        input: &[u8],
        blind: <CS::Group as Group>::Scalar,
    ) -> Result<OprfClientBlindResult<CS>> {
        let blinded_element = deterministic_blind_unchecked::<CS>(input, &blind, Mode::Oprf)?;
        Ok(OprfClientBlindResult {
            state: Self { blind },
            message: BlindedElement(blinded_element),
        })
    }

    /// Computes the third step for the multiplicative blinding version of
    /// DH-OPRF, in which the client unblinds the server's message.
    ///
    /// # Errors
    /// [`Error::Input`] if the `input` is empty or longer then [`u16::MAX`].
    pub fn finalize(
        &self,
        input: &[u8],
        evaluation_element: &EvaluationElement<CS>,
    ) -> Result<Output<CS::Hash>> {
        let unblinded_element = evaluation_element.0 * &CS::Group::invert_scalar(self.blind);
        let mut outputs =
            finalize_after_unblind::<CS, _, _>(iter::once((input, unblinded_element)), &[]);
        outputs.next().unwrap()
    }

    /// Only used for test functions
    #[cfg(test)]
    pub fn from_blind(blind: <CS::Group as Group>::Scalar) -> Self {
        Self { blind }
    }

    /// Exposes the blind group element
    #[cfg(feature = "danger")]
    pub fn get_blind(&self) -> <CS::Group as Group>::Scalar {
        self.blind
    }
}

impl<CS: CipherSuite> OprfServer<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Produces a new instance of a [OprfServer] using a supplied RNG
    ///
    /// # Errors
    /// [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn new<R: RngCore + CryptoRng>(rng: &mut R) -> Result<Self> {
        let mut seed = GenericArray::<_, <CS::Group as Group>::ScalarLen>::default();
        rng.fill_bytes(&mut seed);
        Self::new_from_seed(&seed, &[])
    }

    /// Produces a new instance of a [OprfServer] using a supplied set
    /// of bytes to represent the server's private key
    ///
    /// # Errors
    /// [`Error::Deserialization`] if the private key is not a valid point on
    /// the group or zero.
    pub fn new_with_key(private_key_bytes: &[u8]) -> Result<Self> {
        let sk = CS::Group::deserialize_scalar(private_key_bytes)?;
        Ok(Self { sk })
    }

    /// Produces a new instance of a [OprfServer] using a supplied set
    /// of bytes which are used as a seed to derive the server's private key.
    ///
    /// Corresponds to DeriveKeyPair() function from the VOPRF specification.
    ///
    /// # Errors
    /// - [`Error::DeriveKeyPair`] if the `input` and `seed` together are longer
    ///   then `u16::MAX - 3`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn new_from_seed(seed: &[u8], info: &[u8]) -> Result<Self> {
        let sk = derive_key_internal::<CS>(seed, info, Mode::Oprf)?;
        Ok(Self { sk })
    }

    // Only used for tests
    #[cfg(test)]
    pub fn get_private_key(&self) -> <CS::Group as Group>::Scalar {
        self.sk
    }

    /// Computes the second step for the multiplicative blinding version of
    /// DH-OPRF. This message is sent from the server (who holds the OPRF key)
    /// to the client.
    pub fn blind_evaluate(&self, blinded_element: &BlindedElement<CS>) -> EvaluationElement<CS> {
        EvaluationElement(blinded_element.0 * &self.sk)
    }

    /// Computes the output of the OPRF on the server side
    ///
    /// # Errors
    /// [`Error::Input`]  if the `input` is longer then [`u16::MAX`].
    pub fn evaluate(&self, input: &[u8]) -> Result<Output<<CS as CipherSuite>::Hash>> {
        let input_element = hash_to_group::<CS>(input, Mode::Oprf)?;
        if CS::Group::is_identity_elem(input_element).into() {
            return Err(Error::Input);
        };
        let evaluated_element = input_element * &self.sk;

        let issued_element = CS::Group::serialize_elem(evaluated_element);

        server_evaluate_hash_input::<CS>(input, None, issued_element)
    }
}

/////////////////////////
// Convenience Structs //
//==================== //
/////////////////////////

/// Contains the fields that are returned by a non-verifiable client blind
#[derive_where(Debug; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
pub struct OprfClientBlindResult<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// The state to be persisted on the client
    pub state: OprfClient<CS>,
    /// The message to send to the server
    pub message: BlindedElement<CS>,
}

/////////////////////
// Inner functions //
// =============== //
/////////////////////

type FinalizeAfterUnblindResult<'a, C, I, IE> = Map<
    IE,
    fn((I, <<C as CipherSuite>::Group as Group>::Elem)) -> Result<Output<<C as CipherSuite>::Hash>>,
>;

/// Returned values can only fail with [`Error::Input`].
fn finalize_after_unblind<
    'a,
    CS: CipherSuite,
    I: AsRef<[u8]>,
    IE: 'a + Iterator<Item = (I, <CS::Group as Group>::Elem)>,
>(
    inputs_and_unblinded_elements: IE,
    _unused: &'a [u8],
) -> FinalizeAfterUnblindResult<CS, I, IE>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    inputs_and_unblinded_elements.map(|(input, unblinded_element)| {
        let elem_len = <CS::Group as Group>::ElemLen::U16.to_be_bytes();

        // hashInput = I2OSP(len(input), 2) || input ||
        //             I2OSP(len(unblindedElement), 2) || unblindedElement ||
        //             "Finalize"
        // return Hash(hashInput)
        Ok(CS::Hash::new()
            .chain_update(i2osp_2(input.as_ref().len()).map_err(|_| Error::Input)?)
            .chain_update(input.as_ref())
            .chain_update(elem_len)
            .chain_update(CS::Group::serialize_elem(unblinded_element))
            .chain_update(STR_FINALIZE)
            .finalize())
    })
}

///////////
// Tests //
// ===== //
///////////

#[cfg(test)]
mod tests {
    use core::ptr;

    use rand::rngs::OsRng;

    use super::*;
    use crate::common::{Dst, STR_HASH_TO_GROUP};
    use crate::Group;

    fn prf<CS: CipherSuite>(
        input: &[u8],
        key: <CS::Group as Group>::Scalar,
        info: &[u8],
        mode: Mode,
    ) -> Output<CS::Hash>
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let dst = Dst::new::<CS, _, _>(STR_HASH_TO_GROUP, mode);
        let point = CS::Group::hash_to_curve::<CS::Hash>(&[input], &dst.as_dst()).unwrap();

        let res = point * &key;

        finalize_after_unblind::<CS, _, _>(iter::once((input, res)), info)
            .next()
            .unwrap()
            .unwrap()
    }

    fn base_retrieval<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let mut rng = OsRng;
        let client_blind_result = OprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server = OprfServer::<CS>::new(&mut rng).unwrap();
        let message = server.blind_evaluate(&client_blind_result.message);
        let client_finalize_result = client_blind_result.state.finalize(input, &message).unwrap();
        let res2 = prf::<CS>(input, server.get_private_key(), &[], Mode::Oprf);
        assert_eq!(client_finalize_result, res2);
    }

    fn base_inversion_unsalted<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let mut rng = OsRng;
        let mut input = [0u8; 64];
        rng.fill_bytes(&mut input);
        let client_blind_result = OprfClient::<CS>::blind(&input, &mut rng).unwrap();
        let client_finalize_result = client_blind_result
            .state
            .finalize(&input, &EvaluationElement(client_blind_result.message.0))
            .unwrap();

        let dst = Dst::new::<CS, _, _>(STR_HASH_TO_GROUP, Mode::Oprf);
        let point = CS::Group::hash_to_curve::<CS::Hash>(&[&input], &dst.as_dst()).unwrap();
        let res2 = finalize_after_unblind::<CS, _, _>(iter::once((input.as_ref(), point)), &[])
            .next()
            .unwrap()
            .unwrap();

        assert_eq!(client_finalize_result, res2);
    }

    fn server_evaluate<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let mut rng = OsRng;
        let client_blind_result = OprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server = OprfServer::<CS>::new(&mut rng).unwrap();
        let server_result = server.blind_evaluate(&client_blind_result.message);

        let client_finalize = client_blind_result
            .state
            .finalize(input, &server_result)
            .unwrap();

        // We expect the outputs from client and server to be equal given an identical
        // input
        let server_evaluate = server.evaluate(input).unwrap();
        assert_eq!(client_finalize, server_evaluate);

        // We expect the outputs from client and server to be different given different
        // inputs
        let wrong_input = b"wrong input";
        let server_evaluate = server.evaluate(wrong_input).unwrap();
        assert!(client_finalize != server_evaluate);
    }

    fn zeroize_oprf_client<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let mut rng = OsRng;
        let client_blind_result = OprfClient::<CS>::blind(input, &mut rng).unwrap();

        let mut state = client_blind_result.state;
        unsafe { ptr::drop_in_place(&mut state) };
        assert!(state.serialize().iter().all(|&x| x == 0));

        let mut message = client_blind_result.message;
        unsafe { ptr::drop_in_place(&mut message) };
        assert!(message.serialize().iter().all(|&x| x == 0));
    }

    fn zeroize_oprf_server<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let mut rng = OsRng;
        let client_blind_result = OprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server = OprfServer::<CS>::new(&mut rng).unwrap();
        let mut message = server.blind_evaluate(&client_blind_result.message);

        let mut state = server;
        unsafe { ptr::drop_in_place(&mut state) };
        assert!(state.serialize().iter().all(|&x| x == 0));

        unsafe { ptr::drop_in_place(&mut message) };
        assert!(message.serialize().iter().all(|&x| x == 0));
    }

    #[test]
    fn test_functionality() -> Result<()> {
        use p256::NistP256;
        use p384::NistP384;

        #[cfg(feature = "ristretto255")]
        {
            use crate::Ristretto255;

            base_retrieval::<Ristretto255>();
            base_inversion_unsalted::<Ristretto255>();
            server_evaluate::<Ristretto255>();

            zeroize_oprf_client::<Ristretto255>();
            zeroize_oprf_server::<Ristretto255>();
        }

        base_retrieval::<NistP256>();
        base_inversion_unsalted::<NistP256>();
        server_evaluate::<NistP256>();

        zeroize_oprf_client::<NistP256>();
        zeroize_oprf_server::<NistP256>();

        base_retrieval::<NistP384>();
        base_inversion_unsalted::<NistP384>();
        server_evaluate::<NistP384>();

        zeroize_oprf_client::<NistP384>();
        zeroize_oprf_server::<NistP384>();

        Ok(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Contains the main POPRF API

#[cfg(feature = "alloc")]
use alloc::vec::Vec;
use core::iter::{self, Map, Repeat, Zip};

use derive_where::derive_where;
use digest::core_api::BlockSizeUser;
use digest::{Digest, Output, OutputSizeUser};
use generic_array::typenum::{IsLess, IsLessOrEqual, Unsigned, U256};
use generic_array::GenericArray;
use rand_core::{CryptoRng, RngCore};

use crate::common::{
    derive_keypair, deterministic_blind_unchecked, generate_proof, hash_to_group, i2osp_2,
    server_evaluate_hash_input, verify_proof, BlindedElement, Dst, EvaluationElement, Mode,
    PreparedEvaluationElement, Proof, STR_FINALIZE, STR_HASH_TO_SCALAR, STR_INFO,
};
#[cfg(feature = "serde")]
use crate::serialization::serde::{Element, Scalar};
use crate::{CipherSuite, Error, Group, Result};

////////////////////////////
// High-level API Structs //
// ====================== //
////////////////////////////

/// A client which engages with a [PoprfServer] in verifiable mode, meaning
/// that the OPRF outputs can be checked against a server public key.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct PoprfClient<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) blind: <CS::Group as Group>::Scalar,
    #[cfg_attr(feature = "serde", serde(with = "Element::<CS::Group>"))]
    pub(crate) blinded_element: <CS::Group as Group>::Elem,
}

/// A server which engages with a [PoprfClient] in verifiable mode, meaning
/// that the OPRF outputs can be checked against a server public key.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct PoprfServer<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    pub(crate) sk: <CS::Group as Group>::Scalar,
    #[cfg_attr(feature = "serde", serde(with = "Element::<CS::Group>"))]
    pub(crate) pk: <CS::Group as Group>::Elem,
}

/////////////////////////
// API Implementations //
// =================== //
/////////////////////////

impl<CS: CipherSuite> PoprfClient<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Computes the first step for the multiplicative blinding version of
    /// DH-OPRF.
    ///
    /// # Errors
    /// [`Error::Input`] if the `input` is empty or longer than [`u16::MAX`].
    pub fn blind<R: RngCore + CryptoRng>(
        input: &[u8],
        blinding_factor_rng: &mut R,
    ) -> Result<PoprfClientBlindResult<CS>> {
        let blind = CS::Group::random_scalar(blinding_factor_rng);
        Self::deterministic_blind_unchecked_inner(input, blind)
    }

    /// Computes the first step for the multiplicative blinding version of
    /// DH-OPRF, taking a blinding factor scalar as input instead of sampling
    /// from an RNG.
    ///
    /// # Caution
    ///
    /// This should be used with caution, since it does not perform any checks
    /// on the validity of the blinding factor!
    ///
    /// # Errors
    /// [`Error::Input`] if the `input` is empty or longer than [`u16::MAX`].
    #[cfg(any(feature = "danger", test))]
    pub fn deterministic_blind_unchecked(
        input: &[u8],
        blind: <CS::Group as Group>::Scalar,
    ) -> Result<PoprfClientBlindResult<CS>> {
        Self::deterministic_blind_unchecked_inner(input, blind)
    }

    /// Can only fail with [`Error::Input`].
    fn deterministic_blind_unchecked_inner(
        input: &[u8],
        blind: <CS::Group as Group>::Scalar,
    ) -> Result<PoprfClientBlindResult<CS>> {
        let blinded_element = deterministic_blind_unchecked::<CS>(input, &blind, Mode::Poprf)?;
        Ok(PoprfClientBlindResult {
            state: Self {
                blind,
                blinded_element,
            },
            message: BlindedElement(blinded_element),
        })
    }

    /// Computes the third step for the multiplicative blinding version of
    /// DH-OPRF, in which the client unblinds the server's message.
    ///
    /// # Errors
    /// - [`Error::Info`] if the `info` is longer than `u16::MAX`.
    /// - [`Error::Input`] if the `input` is empty or longer than [`u16::MAX`].
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    /// - [`Error::ProofVerification`] if the `proof` failed to verify.
    pub fn finalize(
        &self,
        input: &[u8],
        evaluation_element: &EvaluationElement<CS>,
        proof: &Proof<CS>,
        pk: <CS::Group as Group>::Elem,
        info: Option<&[u8]>,
    ) -> Result<Output<CS::Hash>> {
        let clients = core::array::from_ref(self);
        let messages = core::array::from_ref(evaluation_element);

        let mut batch_result =
            Self::batch_finalize(iter::once(input), clients, messages, proof, pk, info)?;
        batch_result.next().unwrap()
    }

    /// Allows for batching of the finalization of multiple [PoprfClient]
    /// and [EvaluationElement] pairs
    ///
    /// # Errors
    /// - [`Error::Info`] if the `info` is longer than `u16::MAX`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    /// - [`Error::Batch`] if the number of `inputs`, `clients` and `messages`
    ///   don't match or is longer than [`u16::MAX`].
    /// - [`Error::ProofVerification`] if the `proof` failed to verify.
    ///
    /// The resulting messages can each fail individually with [`Error::Input`]
    /// if the `input` is empty or longer than [`u16::MAX`].
    pub fn batch_finalize<'a, II: 'a + Iterator<Item = &'a [u8]> + ExactSizeIterator, IC, IM>(
        inputs: II,
        clients: &'a IC,
        messages: &'a IM,
        proof: &Proof<CS>,
        pk: <CS::Group as Group>::Elem,
        info: Option<&'a [u8]>,
    ) -> Result<PoprfClientBatchFinalizeResult<'a, CS, II, IC, IM>>
    where
        CS: 'a,
        &'a IC: 'a + IntoIterator<Item = &'a PoprfClient<CS>>,
        <&'a IC as IntoIterator>::IntoIter: ExactSizeIterator,
        &'a IM: 'a + IntoIterator<Item = &'a EvaluationElement<CS>>,
        <&'a IM as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        let unblinded_elements = poprf_unblind(clients, messages, pk, proof, info)?;

        finalize_after_unblind::<'a, CS, _, _>(unblinded_elements, inputs, info)
    }

    /// Only used for test functions
    #[cfg(test)]
    pub fn get_blind(&self) -> <CS::Group as Group>::Scalar {
        self.blind
    }
}

impl<CS: CipherSuite> PoprfServer<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Produces a new instance of a [PoprfServer] using a supplied RNG
    ///
    /// # Errors
    /// [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn new<R: RngCore + CryptoRng>(rng: &mut R) -> Result<Self> {
        let mut seed = GenericArray::<_, <CS::Group as Group>::ScalarLen>::default();
        rng.fill_bytes(&mut seed);

        Self::new_from_seed(&seed, &[])
    }

    /// Produces a new instance of a [PoprfServer] using a supplied set of
    /// bytes to represent the server's private key
    ///
    /// # Errors
    /// [`Error::Deserialization`] if the private key is not a valid point on
    /// the group or zero.
    pub fn new_with_key(key: &[u8]) -> Result<Self> {
        let sk = CS::Group::deserialize_scalar(key)?;
        let pk = CS::Group::base_elem() * &sk;
        Ok(Self { sk, pk })
    }

    /// Produces a new instance of a [PoprfServer] using a supplied set of
    /// bytes which are used as a seed to derive the server's private key.
    ///
    /// Corresponds to DeriveKeyPair() function from the VOPRF specification.
    ///
    /// # Errors
    /// - [`Error::DeriveKeyPair`] if the `input` and `seed` together are longer
    ///   then `u16::MAX - 3`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn new_from_seed(seed: &[u8], info: &[u8]) -> Result<Self> {
        let (sk, pk) = derive_keypair::<CS>(seed, info, Mode::Poprf)?;
        Ok(Self { sk, pk })
    }

    // Only used for tests
    #[cfg(test)]
    pub fn get_private_key(&self) -> <CS::Group as Group>::Scalar {
        self.sk
    }

    /// Computes the second step for the multiplicative blinding version of
    /// DH-OPRF. This message is sent from the server (who holds the OPRF key)
    /// to the client.
    ///
    /// # Errors
    /// - [`Error::Info`] if the `info` is longer than `u16::MAX`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn blind_evaluate<R: RngCore + CryptoRng>(
        &self,
        rng: &mut R,
        blinded_element: &BlindedElement<CS>,
        info: Option<&[u8]>,
    ) -> Result<PoprfServerEvaluateResult<CS>> {
        let PoprfServerBatchEvaluatePrepareResult {
            mut prepared_evaluation_elements,
            prepared_tweak,
        } = self.batch_blind_evaluate_prepare(iter::once(blinded_element), info)?;

        let prepared_evaluation_element = prepared_evaluation_elements.next().unwrap();
        let prepared_evaluation_elements = core::array::from_ref(&prepared_evaluation_element);

        let PoprfServerBatchEvaluateFinishResult {
            mut messages,
            proof,
        } = Self::batch_blind_evaluate_finish(
            rng,
            iter::once(blinded_element),
            prepared_evaluation_elements,
            &prepared_tweak,
        )
        .unwrap();

        Ok(PoprfServerEvaluateResult {
            message: messages.next().unwrap(),
            proof,
        })
    }

    /// Allows for batching of the evaluation of multiple [BlindedElement]
    /// messages from a [PoprfClient]
    ///
    /// # Errors
    /// - [`Error::Info`] if the `info` is longer than `u16::MAX`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    #[cfg(feature = "alloc")]
    pub fn batch_blind_evaluate<'a, R: RngCore + CryptoRng, IE>(
        &self,
        rng: &mut R,
        blinded_elements: &'a IE,
        info: Option<&[u8]>,
    ) -> Result<PoprfServerBatchEvaluateResult<CS>>
    where
        CS: 'a,
        &'a IE: 'a + IntoIterator<Item = &'a BlindedElement<CS>>,
        <&'a IE as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        let PoprfServerBatchEvaluatePrepareResult {
            prepared_evaluation_elements,
            prepared_tweak,
        } = self.batch_blind_evaluate_prepare(blinded_elements.into_iter(), info)?;

        let prepared_evaluation_elements: Vec<_> = prepared_evaluation_elements.collect();

        // This can't fail because we know the size of the inputs.
        let PoprfServerBatchEvaluateFinishResult { messages, proof } =
            Self::batch_blind_evaluate_finish::<_, _, Vec<_>>(
                rng,
                blinded_elements.into_iter(),
                &prepared_evaluation_elements,
                &prepared_tweak,
            )
            .unwrap();

        let messages: Vec<_> = messages.collect();

        Ok(PoprfServerBatchEvaluateResult { messages, proof })
    }

    /// Alternative version of `batch_blind_evaluate` without
    /// memory allocation. Returned [`PreparedEvaluationElement`] have to
    /// be [`collect`](Iterator::collect)ed and passed into
    /// [`batch_blind_evaluate_finish`](Self::batch_blind_evaluate_finish).
    ///
    /// # Errors
    /// - [`Error::Info`] if the `info` is longer than `u16::MAX`.
    /// - [`Error::Protocol`] if the protocol fails and can't be completed.
    pub fn batch_blind_evaluate_prepare<'a, I: Iterator<Item = &'a BlindedElement<CS>>>(
        &self,
        blinded_elements: I,
        info: Option<&[u8]>,
    ) -> Result<PoprfServerBatchEvaluatePrepareResult<CS, I>>
    where
        CS: 'a,
    {
        let tweak = compute_tweak::<CS>(self.sk, info)?;

        Ok(PoprfServerBatchEvaluatePrepareResult {
            prepared_evaluation_elements: blinded_elements.zip(iter::repeat(tweak)).map(
                |(blinded_element, tweak)| {
                    PreparedEvaluationElement(EvaluationElement(
                        blinded_element.0 * &CS::Group::invert_scalar(tweak),
                    ))
                },
            ),
            prepared_tweak: PoprfPreparedTweak(tweak),
        })
    }

    /// See [`batch_blind_evaluate_prepare`](Self::batch_blind_evaluate_prepare)
    /// for more details.
    ///
    /// # Errors
    /// [`Error::Batch`] if the number of `blinded_elements` and
    /// `prepared_evaluation_elements` don't match or is longer then
    /// [`u16::MAX`]
    pub fn batch_blind_evaluate_finish<
        'a,
        'b,
        R: RngCore + CryptoRng,
        IB: Iterator<Item = &'a BlindedElement<CS>> + ExactSizeIterator,
        IE,
    >(
        rng: &mut R,
        blinded_elements: IB,
        prepared_evaluation_elements: &'b IE,
        prepared_tweak: &PoprfPreparedTweak<CS>,
    ) -> Result<PoprfServerBatchEvaluateFinishResult<'b, CS, IE>>
    where
        CS: 'a,
        &'b IE: IntoIterator<Item = &'b PreparedEvaluationElement<CS>>,
        <&'b IE as IntoIterator>::IntoIter: ExactSizeIterator,
    {
        let g = CS::Group::base_elem();
        let tweak = prepared_tweak.0;
        let tweaked_key = g * &tweak;

        let proof = generate_proof(
            rng,
            tweak,
            g,
            tweaked_key,
            prepared_evaluation_elements
                .into_iter()
                .map(|element| element.0 .0),
            blinded_elements.map(|element| element.0),
            Mode::Poprf,
        )?;

        let messages = prepared_evaluation_elements.into_iter().map(<fn(
            &PreparedEvaluationElement<CS>,
        ) -> _>::from(
            |element| EvaluationElement(element.0 .0),
        ));

        Ok(PoprfServerBatchEvaluateFinishResult { messages, proof })
    }

    /// Computes the output of the VOPRF on the server side
    ///
    /// # Errors
    /// [`Error::Input`]  if the `input` is longer then [`u16::MAX`].
    pub fn evaluate(
        &self,
        input: &[u8],
        info: Option<&[u8]>,
    ) -> Result<Output<<CS as CipherSuite>::Hash>> {
        let input_element = hash_to_group::<CS>(input, Mode::Poprf)?;
        if CS::Group::is_identity_elem(input_element).into() {
            return Err(Error::Input);
        };

        let tweak = compute_tweak::<CS>(self.sk, info)?;

        let evaluated_element = input_element * &CS::Group::invert_scalar(tweak);

        let issued_element = CS::Group::serialize_elem(evaluated_element);

        server_evaluate_hash_input::<CS>(input, info, issued_element)
    }

    /// Retrieves the server's public key
    pub fn get_public_key(&self) -> <CS::Group as Group>::Elem {
        self.pk
    }
}

impl<CS: CipherSuite> BlindedElement<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Creates a [BlindedElement] from a raw group element.
    ///
    /// # Caution
    ///
    /// This should be used with caution, since it does not perform any checks
    /// on the validity of the value itself!
    #[cfg(feature = "danger")]
    pub fn from_value_unchecked(value: <CS::Group as Group>::Elem) -> Self {
        Self(value)
    }

    /// Exposes the internal value
    #[cfg(feature = "danger")]
    pub fn value(&self) -> <CS::Group as Group>::Elem {
        self.0
    }
}

impl<CS: CipherSuite> EvaluationElement<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Creates an [EvaluationElement] from a raw group element.
    ///
    /// # Caution
    ///
    /// This should be used with caution, since it does not perform any checks
    /// on the validity of the value itself!
    #[cfg(feature = "danger")]
    pub fn from_value_unchecked(value: <CS::Group as Group>::Elem) -> Self {
        Self(value)
    }

    /// Exposes the internal value
    #[cfg(feature = "danger")]
    pub fn value(&self) -> <CS::Group as Group>::Elem {
        self.0
    }
}

/////////////////////////
// Convenience Structs //
//==================== //
/////////////////////////

/// Contains the fields that are returned by a verifiable client blind
#[derive_where(Debug; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
pub struct PoprfClientBlindResult<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// The state to be persisted on the client
    pub state: PoprfClient<CS>,
    /// The message to send to the server
    pub message: BlindedElement<CS>,
}

/// Concrete return type for [`PoprfClient::batch_finalize`].
pub type PoprfClientBatchFinalizeResult<'a, CS, II, IC, IM> =
    FinalizeAfterUnblindResult<'a, CS, PoprfUnblindResult<'a, CS, IC, IM>, II>;

/// Contains the fields that are returned by a verifiable server evaluate
#[derive_where(Debug; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
pub struct PoprfServerEvaluateResult<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// The message to send to the client
    pub message: EvaluationElement<CS>,
    /// The proof for the client to verify
    pub proof: Proof<CS>,
}

/// Contains the fields that are returned by a verifiable server batch evaluate
#[derive_where(Debug; <CS::Group as Group>::Scalar, <CS::Group as Group>::Elem)]
#[cfg(feature = "alloc")]
pub struct PoprfServerBatchEvaluateResult<CS: CipherSuite>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// The messages to send to the client
    pub messages: Vec<EvaluationElement<CS>>,
    /// The proof for the client to verify
    pub proof: Proof<CS>,
}

/// Concrete type of [`EvaluationElement`]s in
/// [`PoprfServerBatchEvaluatePrepareResult`].
pub type PoprfServerBatchEvaluatePreparedEvaluationElements<CS, I> = Map<
    Zip<I, Repeat<<<CS as CipherSuite>::Group as Group>::Scalar>>,
    fn(
        (
            &BlindedElement<CS>,
            <<CS as CipherSuite>::Group as Group>::Scalar,
        ),
    ) -> PreparedEvaluationElement<CS>,
>;

/// Prepared tweak by a partially verifiable server batch evaluate prepare.
#[derive_where(Clone, ZeroizeOnDrop)]
#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; <CS::Group as Group>::Scalar)]
#[cfg_attr(
    feature = "serde",
    derive(serde::Deserialize, serde::Serialize),
    serde(bound = "")
)]
pub struct PoprfPreparedTweak<CS: CipherSuite>(
    #[cfg_attr(feature = "serde", serde(with = "Scalar::<CS::Group>"))]
    <CS::Group as Group>::Scalar,
)
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>;

/// Contains the fields that are returned by a partially verifiable server batch
/// evaluate prepare
#[derive_where(Debug; I, <CS::Group as Group>::Scalar)]
pub struct PoprfServerBatchEvaluatePrepareResult<CS: CipherSuite, I>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Prepared [`EvaluationElement`].
    pub prepared_evaluation_elements: PoprfServerBatchEvaluatePreparedEvaluationElements<CS, I>,
    /// Prepared tweak.
    pub prepared_tweak: PoprfPreparedTweak<CS>,
}

/// Concrete type of [`EvaluationElement`]s in
/// [`PoprfServerBatchEvaluateFinishResult`].
pub type PoprfServerBatchEvaluateFinishedMessages<'a, CS, I> = Map<
    <&'a I as IntoIterator>::IntoIter,
    fn(&PreparedEvaluationElement<CS>) -> EvaluationElement<CS>,
>;

/// Contains the fields that are returned by a verifiable server batch evaluate
/// finish.
#[derive_where(Debug; <&'a I as IntoIterator>::IntoIter, <CS::Group as Group>::Scalar)]
pub struct PoprfServerBatchEvaluateFinishResult<'a, CS: 'a + CipherSuite, I>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    &'a I: IntoIterator<Item = &'a PreparedEvaluationElement<CS>>,
{
    /// The [`EvaluationElement`]s to send to the client
    pub messages: PoprfServerBatchEvaluateFinishedMessages<'a, CS, I>,
    /// The proof for the client to verify
    pub proof: Proof<CS>,
}

/////////////////////
// Inner functions //
// =============== //
/////////////////////

/// Inner function for POPRF blind. Computes the tweaked key from the server
/// public key and info.
///
/// Can only fail with [`Error::Info`] or [`Error::Protocol`]
fn compute_tweaked_key<CS: CipherSuite>(
    pk: <CS::Group as Group>::Elem,
    info: Option<&[u8]>,
) -> Result<<CS::Group as Group>::Elem>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // None for info is treated the same as empty bytes
    let info = info.unwrap_or_default();

    // framedInfo = "Info" || I2OSP(len(info), 2) || info
    // m = G.HashToScalar(framedInfo)
    // T = G.ScalarBaseMult(m)
    // tweakedKey = T + pkS
    // if tweakedKey == G.Identity():
    //   raise InvalidInputError
    let info_len = i2osp_2(info.len()).map_err(|_| Error::Info)?;
    let framed_info = [STR_INFO.as_slice(), &info_len, info];

    let dst = Dst::new::<CS, _, _>(STR_HASH_TO_SCALAR, Mode::Poprf);
    // This can't fail, the size of the `input` is known.
    let m = CS::Group::hash_to_scalar::<CS::Hash>(&framed_info, &dst.as_dst()).unwrap();

    let t = CS::Group::base_elem() * &m;
    let tweaked_key = t + &pk;

    // Check if resulting element
    match bool::from(CS::Group::is_identity_elem(tweaked_key)) {
        true => Err(Error::Protocol),
        false => Ok(tweaked_key),
    }
}

/// Inner function for POPRF evaluate. Computes the tweak from the server
/// private key and info.
///
/// Can only fail with [`Error::Info`] and [`Error::Protocol`].
fn compute_tweak<CS: CipherSuite>(
    sk: <CS::Group as Group>::Scalar,
    info: Option<&[u8]>,
) -> Result<<CS::Group as Group>::Scalar>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    // None for info is treated the same as empty bytes
    let info = info.unwrap_or_default();

    // framedInfo = "Info" || I2OSP(len(info), 2) || info
    // m = G.HashToScalar(framedInfo)
    // t = skS + m
    // if t == 0:
    //   raise InverseError
    let info_len = i2osp_2(info.len()).map_err(|_| Error::Info)?;
    let framed_info = [STR_INFO.as_slice(), &info_len, info];

    let dst = Dst::new::<CS, _, _>(STR_HASH_TO_SCALAR, Mode::Poprf);
    // This can't fail, the size of the `input` is known.
    let m = CS::Group::hash_to_scalar::<CS::Hash>(&framed_info, &dst.as_dst()).unwrap();

    let t = sk + &m;

    // Check if resulting element is equal to zero
    match bool::from(CS::Group::is_zero_scalar(t)) {
        true => Err(Error::Protocol),
        false => Ok(t),
    }
}

type PoprfUnblindResult<'a, CS, IC, IM> = Map<
    Zip<
        Map<
            <&'a IC as IntoIterator>::IntoIter,
            fn(&PoprfClient<CS>) -> <<CS as CipherSuite>::Group as Group>::Scalar,
        >,
        <&'a IM as IntoIterator>::IntoIter,
    >,
    fn(
        (
            <<CS as CipherSuite>::Group as Group>::Scalar,
            &'a EvaluationElement<CS>,
        ),
    ) -> <<CS as CipherSuite>::Group as Group>::Elem,
>;

/// Can only fail with [`Error::Info`], [`Error::Protocol`], [`Error::Batch] or
/// [`Error::ProofVerification`].
fn poprf_unblind<'a, CS: 'a + CipherSuite, IC, IM>(
    clients: &'a IC,
    messages: &'a IM,
    pk: <CS::Group as Group>::Elem,
    proof: &Proof<CS>,
    info: Option<&[u8]>,
) -> Result<PoprfUnblindResult<'a, CS, IC, IM>>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    &'a IC: 'a + IntoIterator<Item = &'a PoprfClient<CS>>,
    <&'a IC as IntoIterator>::IntoIter: ExactSizeIterator,
    &'a IM: 'a + IntoIterator<Item = &'a EvaluationElement<CS>>,
    <&'a IM as IntoIterator>::IntoIter: ExactSizeIterator,
{
    let info = info.unwrap_or_default();
    let tweaked_key = compute_tweaked_key::<CS>(pk, Some(info))?;

    let g = CS::Group::base_elem();

    let blinds = clients
        .into_iter()
        // Convert to `fn` pointer to make a return type possible.
        .map(<fn(&PoprfClient<CS>) -> _>::from(|x| x.blind));
    let evaluation_elements = messages.into_iter().map(|element| element.0);
    let blinded_elements = clients.into_iter().map(|client| client.blinded_element);

    verify_proof(
        g,
        tweaked_key,
        evaluation_elements,
        blinded_elements,
        proof,
        Mode::Poprf,
    )?;

    Ok(blinds
        .zip(messages.into_iter())
        .map(|(blind, x)| x.0 * &CS::Group::invert_scalar(blind)))
}

type FinalizeAfterUnblindResult<'a, CS, IE, II> = Map<
    Zip<Zip<IE, II>, Repeat<&'a [u8]>>,
    fn(
        ((<<CS as CipherSuite>::Group as Group>::Elem, &[u8]), &[u8]),
    ) -> Result<GenericArray<u8, <<CS as CipherSuite>::Hash as OutputSizeUser>::OutputSize>>,
>;

/// Can only fail with [`Error::Batch`] and returned values can only fail with
/// [`Error::Info`] or [`Error::Input`] individually.
fn finalize_after_unblind<
    'a,
    CS: CipherSuite,
    IE: 'a + Iterator<Item = <CS::Group as Group>::Elem> + ExactSizeIterator,
    II: 'a + Iterator<Item = &'a [u8]> + ExactSizeIterator,
>(
    unblinded_elements: IE,
    inputs: II,
    info: Option<&'a [u8]>,
) -> Result<FinalizeAfterUnblindResult<'a, CS, IE, II>>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    if unblinded_elements.len() != inputs.len() {
        return Err(Error::Batch);
    }

    let info = info.unwrap_or_default();

    Ok(unblinded_elements.zip(inputs).zip(iter::repeat(info)).map(
        |((unblinded_element, input), info)| {
            let elem_len = <CS::Group as Group>::ElemLen::U16.to_be_bytes();

            // hashInput = I2OSP(len(input), 2) || input ||
            //             I2OSP(len(info), 2) || info ||
            //             I2OSP(len(unblindedElement), 2) || unblindedElement ||
            //             "Finalize"
            // return Hash(hashInput)
            let output = CS::Hash::new()
                .chain_update(i2osp_2(input.as_ref().len()).map_err(|_| Error::Input)?)
                .chain_update(input.as_ref())
                .chain_update(i2osp_2(info.as_ref().len()).map_err(|_| Error::Info)?)
                .chain_update(info.as_ref())
                .chain_update(elem_len)
                .chain_update(CS::Group::serialize_elem(unblinded_element))
                .chain_update(STR_FINALIZE)
                .finalize();

            Ok(output)
        },
    ))
}

///////////
// Tests //
// ===== //
///////////

#[cfg(test)]
mod tests {
    use core::ops::Add;
    use core::ptr;

    use generic_array::typenum::Sum;
    use generic_array::ArrayLength;
    use rand::rngs::OsRng;

    use super::*;
    use crate::common::STR_HASH_TO_GROUP;
    use crate::Group;

    fn prf<CS: CipherSuite>(
        input: &[u8],
        key: <CS::Group as Group>::Scalar,
        info: &[u8],
        mode: Mode,
    ) -> Output<CS::Hash>
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let t = compute_tweak::<CS>(key, Some(info)).unwrap();

        let dst = Dst::new::<CS, _, _>(STR_HASH_TO_GROUP, mode);
        let point = CS::Group::hash_to_curve::<CS::Hash>(&[input], &dst.as_dst()).unwrap();

        // evaluatedElement = G.ScalarInverse(t) * blindedElement
        let res = point * &CS::Group::invert_scalar(t);

        finalize_after_unblind::<CS, _, _>(iter::once(res), iter::once(input), Some(info))
            .unwrap()
            .next()
            .unwrap()
            .unwrap()
    }

    fn verifiable_retrieval<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let info = b"info";
        let mut rng = OsRng;
        let server = PoprfServer::<CS>::new(&mut rng).unwrap();
        let client_blind_result = PoprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server_result = server
            .blind_evaluate(&mut rng, &client_blind_result.message, Some(info))
            .unwrap();
        let client_finalize_result = client_blind_result
            .state
            .finalize(
                input,
                &server_result.message,
                &server_result.proof,
                server.get_public_key(),
                Some(info),
            )
            .unwrap();
        let res2 = prf::<CS>(input, server.get_private_key(), info, Mode::Poprf);
        assert_eq!(client_finalize_result, res2);
    }

    fn verifiable_bad_public_key<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let info = b"info";
        let mut rng = OsRng;
        let server = PoprfServer::<CS>::new(&mut rng).unwrap();
        let client_blind_result = PoprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server_result = server
            .blind_evaluate(&mut rng, &client_blind_result.message, Some(info))
            .unwrap();
        let wrong_pk = {
            let dst = Dst::new::<CS, _, _>(STR_HASH_TO_GROUP, Mode::Oprf);
            // Choose a group element that is unlikely to be the right public key
            CS::Group::hash_to_curve::<CS::Hash>(&[b"msg"], &dst.as_dst()).unwrap()
        };
        let client_finalize_result = client_blind_result.state.finalize(
            input,
            &server_result.message,
            &server_result.proof,
            wrong_pk,
            Some(info),
        );
        assert!(client_finalize_result.is_err());
    }

    fn verifiable_server_evaluate<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
    {
        let input = b"input";
        let info = Some(b"info".as_slice());
        let mut rng = OsRng;
        let client_blind_result = PoprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server = PoprfServer::<CS>::new(&mut rng).unwrap();
        let server_result = server
            .blind_evaluate(&mut rng, &client_blind_result.message, info)
            .unwrap();

        let client_finalize = client_blind_result
            .state
            .finalize(
                input,
                &server_result.message,
                &server_result.proof,
                server.get_public_key(),
                info,
            )
            .unwrap();

        // We expect the outputs from client and server to be equal given an identical
        // input
        let server_evaluate = server.evaluate(input, info).unwrap();
        assert_eq!(client_finalize, server_evaluate);

        // We expect the outputs from client and server to be different given different
        // inputs
        let wrong_input = b"wrong input";
        let server_evaluate = server.evaluate(wrong_input, info).unwrap();
        assert!(client_finalize != server_evaluate);
    }

    fn zeroize_verifiable_client<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        Sum<<CS::Group as Group>::ScalarLen, <CS::Group as Group>::ElemLen>: ArrayLength<u8>,
    {
        let input = b"input";
        let mut rng = OsRng;
        let client_blind_result = PoprfClient::<CS>::blind(input, &mut rng).unwrap();

        let mut state = client_blind_result.state;
        unsafe { ptr::drop_in_place(&mut state) };
        assert!(state.serialize().iter().all(|&x| x == 0));

        let mut message = client_blind_result.message;
        unsafe { ptr::drop_in_place(&mut message) };
        assert!(message.serialize().iter().all(|&x| x == 0));
    }

    fn zeroize_verifiable_server<CS: CipherSuite>()
    where
        <CS::Hash as OutputSizeUser>::OutputSize:
            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        Sum<<CS::Group as Group>::ScalarLen, <CS::Group as Group>::ElemLen>: ArrayLength<u8>,
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ScalarLen>,
        Sum<<CS::Group as Group>::ScalarLen, <CS::Group as Group>::ScalarLen>: ArrayLength<u8>,
    {
        let input = b"input";
        let info = b"info";
        let mut rng = OsRng;
        let server = PoprfServer::<CS>::new(&mut rng).unwrap();
        let client_blind_result = PoprfClient::<CS>::blind(input, &mut rng).unwrap();
        let server_result = server
            .blind_evaluate(&mut rng, &client_blind_result.message, Some(info))
            .unwrap();

        let mut state = server;
        unsafe { ptr::drop_in_place(&mut state) };
        assert!(state.serialize().iter().all(|&x| x == 0));

        let mut message = server_result.message;
        unsafe { ptr::drop_in_place(&mut message) };
        assert!(message.serialize().iter().all(|&x| x == 0));

        let mut proof = server_result.proof;
        unsafe { ptr::drop_in_place(&mut proof) };
        assert!(proof.serialize().iter().all(|&x| x == 0));
    }

    #[test]
    fn test_functionality() -> Result<()> {
        use p256::NistP256;
        use p384::NistP384;

        #[cfg(feature = "ristretto255")]
        {
            use crate::Ristretto255;

            verifiable_retrieval::<Ristretto255>();
            verifiable_bad_public_key::<Ristretto255>();
            verifiable_server_evaluate::<Ristretto255>();

            zeroize_verifiable_client::<Ristretto255>();
            zeroize_verifiable_server::<Ristretto255>();
        }

        verifiable_retrieval::<NistP256>();
        verifiable_bad_public_key::<NistP256>();
        verifiable_server_evaluate::<NistP256>();

        zeroize_verifiable_client::<NistP256>();
        zeroize_verifiable_server::<NistP256>();

        verifiable_retrieval::<NistP384>();
        verifiable_bad_public_key::<NistP384>();
        verifiable_server_evaluate::<NistP384>();

        zeroize_verifiable_client::<NistP384>();
        zeroize_verifiable_server::<NistP384>();

        Ok(())
    }
}


// Copyright (c) Meta Platforms, Inc. and affiliates.
//
// This source code is dual-licensed under either the MIT license found in the
// LICENSE-MIT file in the root directory of this source tree or the Apache
// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
// of this source tree. You may select, at your option, one of the above-listed
// licenses.

//! Handles the serialization of each of the components used in the VOPRF
//! protocol

use core::ops::Add;

use digest::core_api::BlockSizeUser;
use digest::OutputSizeUser;
use generic_array::sequence::Concat;
use generic_array::typenum::{IsLess, IsLessOrEqual, Sum, Unsigned, U256};
use generic_array::{ArrayLength, GenericArray};

use crate::{
    BlindedElement, CipherSuite, Error, EvaluationElement, Group, OprfClient, OprfServer,
    PoprfClient, PoprfServer, Proof, Result, VoprfClient, VoprfServer,
};

//////////////////////////////////////////////////////////
// Serialization and Deserialization for High-Level API //
// ==================================================== //
//////////////////////////////////////////////////////////

/// Length of [`OprfClient`] in bytes for serialization.
pub type OprfClientLen<CS> = <<CS as CipherSuite>::Group as Group>::ScalarLen;

impl<CS: CipherSuite> OprfClient<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, OprfClientLen<CS>> {
        CS::Group::serialize_scalar(self.blind)
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let blind = deserialize_scalar::<CS::Group>(&mut input)?;

        Ok(Self { blind })
    }
}

/// Length of [`VoprfClient`] in bytes for serialization.
pub type VoprfClientLen<CS> = Sum<
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
    <<CS as CipherSuite>::Group as Group>::ElemLen,
>;

impl<CS: CipherSuite> VoprfClient<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, VoprfClientLen<CS>>
    where
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        VoprfClientLen<CS>: ArrayLength<u8>,
    {
        <CS::Group as Group>::serialize_scalar(self.blind)
            .concat(<CS::Group as Group>::serialize_elem(self.blinded_element))
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let blind = deserialize_scalar::<CS::Group>(&mut input)?;
        let blinded_element = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self {
            blind,
            blinded_element,
        })
    }
}

/// Length of [`PoprfClient`] in bytes for serialization.
pub type PoprfClientLen<CS> = Sum<
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
    <<CS as CipherSuite>::Group as Group>::ElemLen,
>;

impl<CS: CipherSuite> PoprfClient<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, PoprfClientLen<CS>>
    where
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        PoprfClientLen<CS>: ArrayLength<u8>,
    {
        <CS::Group as Group>::serialize_scalar(self.blind)
            .concat(<CS::Group as Group>::serialize_elem(self.blinded_element))
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let blind = deserialize_scalar::<CS::Group>(&mut input)?;
        let blinded_element = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self {
            blind,
            blinded_element,
        })
    }
}

/// Length of [`OprfServer`] in bytes for serialization.
pub type OprfServerLen<CS> = <<CS as CipherSuite>::Group as Group>::ScalarLen;

impl<CS: CipherSuite> OprfServer<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, OprfServerLen<CS>> {
        CS::Group::serialize_scalar(self.sk)
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let sk = deserialize_scalar::<CS::Group>(&mut input)?;

        Ok(Self { sk })
    }
}

/// Length of [`VoprfServer`] in bytes for serialization.
pub type VoprfServerLen<CS> = Sum<
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
    <<CS as CipherSuite>::Group as Group>::ElemLen,
>;

impl<CS: CipherSuite> VoprfServer<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, VoprfServerLen<CS>>
    where
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        VoprfServerLen<CS>: ArrayLength<u8>,
    {
        CS::Group::serialize_scalar(self.sk).concat(CS::Group::serialize_elem(self.pk))
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let sk = deserialize_scalar::<CS::Group>(&mut input)?;
        let pk = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self { sk, pk })
    }
}

/// Length of [`PoprfServer`] in bytes for serialization.
pub type PoprfServerLen<CS> = Sum<
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
    <<CS as CipherSuite>::Group as Group>::ElemLen,
>;

impl<CS: CipherSuite> PoprfServer<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, PoprfServerLen<CS>>
    where
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ElemLen>,
        PoprfServerLen<CS>: ArrayLength<u8>,
    {
        CS::Group::serialize_scalar(self.sk).concat(CS::Group::serialize_elem(self.pk))
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let sk = deserialize_scalar::<CS::Group>(&mut input)?;
        let pk = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self { sk, pk })
    }
}

/// Length of [`Proof`] in bytes for serialization.
pub type ProofLen<CS> = Sum<
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
    <<CS as CipherSuite>::Group as Group>::ScalarLen,
>;

impl<CS: CipherSuite> Proof<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, ProofLen<CS>>
    where
        <CS::Group as Group>::ScalarLen: Add<<CS::Group as Group>::ScalarLen>,
        ProofLen<CS>: ArrayLength<u8>,
    {
        CS::Group::serialize_scalar(self.c_scalar)
            .concat(CS::Group::serialize_scalar(self.s_scalar))
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let c_scalar = deserialize_scalar::<CS::Group>(&mut input)?;
        let s_scalar = deserialize_scalar::<CS::Group>(&mut input)?;

        Ok(Proof { c_scalar, s_scalar })
    }
}

/// Length of [`BlindedElement`] in bytes for serialization.
pub type BlindedElementLen<CS> = <<CS as CipherSuite>::Group as Group>::ElemLen;

impl<CS: CipherSuite> BlindedElement<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, BlindedElementLen<CS>> {
        CS::Group::serialize_elem(self.0)
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let value = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self(value))
    }
}

/// Length of [`EvaluationElement`] in bytes for serialization.
pub type EvaluationElementLen<CS> = <<CS as CipherSuite>::Group as Group>::ElemLen;

impl<CS: CipherSuite> EvaluationElement<CS>
where
    <CS::Hash as OutputSizeUser>::OutputSize:
        IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
{
    /// Serialization into bytes
    pub fn serialize(&self) -> GenericArray<u8, EvaluationElementLen<CS>> {
        CS::Group::serialize_elem(self.0)
    }

    /// Deserialization from bytes
    ///
    /// # Errors
    /// [`Error::Deserialization`] if failed to deserialize `input`.
    pub fn deserialize(mut input: &[u8]) -> Result<Self> {
        let value = deserialize_elem::<CS::Group>(&mut input)?;

        Ok(Self(value))
    }
}

fn deserialize_elem<G: Group>(input: &mut &[u8]) -> Result<G::Elem> {
    let input = input
        .take_ext(G::ElemLen::USIZE)
        .ok_or(Error::Deserialization)?;
    G::deserialize_elem(input)
}

fn deserialize_scalar<G: Group>(input: &mut &[u8]) -> Result<G::Scalar> {
    let input = input
        .take_ext(G::ScalarLen::USIZE)
        .ok_or(Error::Deserialization)?;
    G::deserialize_scalar(input)
}

trait SliceExt {
    fn take_ext(self: &mut &Self, take: usize) -> Option<&Self>;
}

impl<T> SliceExt for [T] {
    fn take_ext(self: &mut &Self, take: usize) -> Option<&Self> {
        if take > self.len() {
            return None;
        }

        let (front, back) = self.split_at(take);
        *self = back;
        Some(front)
    }
}

#[cfg(feature = "serde")]
pub(crate) mod serde {
    use core::marker::PhantomData;

    use generic_array::GenericArray;
    use serde::de::{Deserializer, Error};
    use serde::ser::Serializer;
    use serde::{Deserialize, Serialize};

    use crate::Group;

    pub(crate) struct Element<G: Group>(PhantomData<G>);

    impl<'de, G: Group> Element<G> {
        pub(crate) fn deserialize<D>(deserializer: D) -> Result<G::Elem, D::Error>
        where
            D: Deserializer<'de>,
        {
            GenericArray::<_, G::ElemLen>::deserialize(deserializer)
                .and_then(|bytes| G::deserialize_elem(&bytes).map_err(D::Error::custom))
        }

        pub(crate) fn serialize<S>(self_: &G::Elem, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: Serializer,
        {
            G::serialize_elem(*self_).serialize(serializer)
        }
    }

    pub(crate) struct Scalar<G: Group>(PhantomData<G>);

    impl<'de, G: Group> Scalar<G> {
        pub(crate) fn deserialize<D>(deserializer: D) -> Result<G::Scalar, D::Error>
        where
            D: Deserializer<'de>,
        {
            GenericArray::<_, G::ScalarLen>::deserialize(deserializer)
                .and_then(|bytes| G::deserialize_scalar(&bytes).map_err(D::Error::custom))
        }

        pub(crate) fn serialize<S>(self_: &G::Scalar, serializer: S) -> Result<S::Ok, S::Error>
        where
            S: Serializer,
        {
            G::serialize_scalar(*self_).serialize(serializer)
        }
    }
}

#[cfg(test)]
mod test {
    use proptest::collection::vec;
    use proptest::prelude::*;

    use crate::{
        BlindedElement, EvaluationElement, OprfClient, OprfServer, PoprfClient, PoprfServer, Proof,
        VoprfClient, VoprfServer,
    };

    macro_rules! test_deserialize {
        ($item:ident, $bytes:ident) => {
            #[cfg(feature = "ristretto255")]
            {
                let _ = $item::<crate::Ristretto255>::deserialize(&$bytes[..]);
            }

            let _ = $item::<p256::NistP256>::deserialize(&$bytes[..]);
            let _ = $item::<p384::NistP384>::deserialize(&$bytes[..]);
        };
    }

    proptest! {
        #[test]
        fn test_nocrash_oprf_client(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(OprfClient, bytes);
        }

        #[test]
        fn test_nocrash_voprf_client(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(VoprfClient, bytes);
        }

        #[test]
        fn test_nocrash_poprf_client(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(PoprfClient, bytes);
        }

        #[test]
        fn test_nocrash_oprf_server(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(OprfServer, bytes);
        }

        #[test]
        fn test_nocrash_voprf_server(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(VoprfServer, bytes);
        }

        #[test]
        fn test_nocrash_poprf_server(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(PoprfServer, bytes);
        }

        #[test]
        fn test_nocrash_blinded_element(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(BlindedElement, bytes);
        }

        #[test]
        fn test_nocrash_evaluation_element(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(EvaluationElement, bytes);
        }

        #[test]
        fn test_nocrash_proof(bytes in vec(any::<u8>(), 0..200)) {
            test_deserialize!(Proof, bytes);
        }
    }
}