Repository: solana-labs/solana-program-library
Pull request #6945 (2024-06-28 08:57:01)

[confidential-transfer] Add transfer with fee proof generation and extraction
#### Problem
The logic to generate proof data for transfer with fee have not been added to the repo yet.

#### Summary
Add the logic to generate the proof data. This proof is quite technical, but it essentially follows the existing logic in [token client](https://github.com/solana-labs/solana-program-library/blob/master/token/client/src/proof_generation.rs).


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 7ef8f68f9..10d804b72 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7684,6 +7684,8 @@ dependencies = [
 name = "spl-token-confidential-transfer-proof-extraction"
 version = "0.1.0"
 dependencies = [
+ "bytemuck",
+ "solana-curve25519",
  "solana-zk-sdk",
  "thiserror",
 ]
diff --git a/libraries/pod/src/primitives.rs b/libraries/pod/src/primitives.rs
index 07eb232bf..f6759cf5f 100644
--- a/libraries/pod/src/primitives.rs
+++ b/libraries/pod/src/primitives.rs
@@ -115,6 +115,18 @@ impl_int_conversion!(PodU64, u64);
 pub struct PodI64([u8; 8]);
 impl_int_conversion!(PodI64, i64);
 
+/// `u128` type that can be used in Pods
+#[cfg_attr(
+    feature = "borsh",
+    derive(BorshDeserialize, BorshSerialize, BorshSchema)
+)]
+#[cfg_attr(feature = "serde-traits", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "serde-traits", serde(from = "u128", into = "u128"))]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
+#[repr(transparent)]
+pub struct PodU128(pub [u8; 16]);
+impl_int_conversion!(PodU128, u128);
+
 #[cfg(test)]
 mod tests {
     use {super::*, crate::bytemuck::pod_from_bytes};
@@ -230,4 +242,28 @@ mod tests {
         let deserialized = serde_json::from_str::<PodI64>(&serialized).unwrap();
         assert_eq!(pod_i64, deserialized);
     }
+
+    #[test]
+    fn test_pod_u128() {
+        assert!(pod_from_bytes::<PodU128>(&[]).is_err());
+        assert_eq!(
+            1u128,
+            u128::from(
+                *pod_from_bytes::<PodU128>(&[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+                    .unwrap()
+            )
+        );
+    }
+
+    #[cfg(feature = "serde-traits")]
+    #[test]
+    fn test_pod_u128_serde() {
+        let pod_u128: PodU128 = u128::MAX.into();
+
+        let serialized = serde_json::to_string(&pod_u128).unwrap();
+        assert_eq!(&serialized, "340282366920938463463374607431768211455");
+
+        let deserialized = serde_json::from_str::<PodU128>(&serialized).unwrap();
+        assert_eq!(pod_u128, deserialized);
+    }
 }
diff --git a/patch.crates-io.sh b/patch.crates-io.sh
index 72cc4cce9..59498ec5f 100755
--- a/patch.crates-io.sh
+++ b/patch.crates-io.sh
@@ -111,6 +111,7 @@ crates_map+=("solana-udp-client udp-client")
 crates_map+=("solana-version version")
 crates_map+=("solana-zk-token-sdk zk-token-sdk")
 crates_map+=("solana-zk-sdk zk-sdk")
+crates_map+=("solana-curve25519 curves/curve25519")
 
 patch_crates=()
 for map_entry in "${crates_map[@]}"; do
diff --git a/token-upgrade/cli/src/main.rs b/token-upgrade/cli/src/main.rs
index daf5e5875..92453f257 100644
--- a/token-upgrade/cli/src/main.rs
+++ b/token-upgrade/cli/src/main.rs
@@ -1,14 +1,12 @@
 use {
-    clap::{crate_description, crate_name, crate_version, Arg, Command},
+    clap::{crate_description, crate_name, crate_version, Arg, ArgAction, Command},
     solana_clap_v3_utils::{
         input_parsers::{
             parse_url_or_moniker,
-            signer::{try_pubkey_of, SignerSourceParserBuilder},
+            signer::{SignerSource, SignerSourceParserBuilder},
         },
         input_validators::normalize_to_url_if_moniker,
-        keypair::{
-            signer_from_path, signer_from_path_with_config, DefaultSigner, SignerFromPathConfig,
-        },
+        keypair::{signer_from_path, signer_from_source, SignerFromPathConfig},
     },
     solana_client::nonblocking::rpc_client::RpcClient,
     solana_remote_wallet::remote_wallet::RemoteWalletManager,
@@ -336,7 +334,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
                     .value_parser(SignerSourceParserBuilder::default().allow_all().build())
                     .value_name("MULTISIG_SIGNER")
                     .takes_value(true)
-                    .multiple(true)
+                    .action(ArgAction::Append)
                     .min_values(0)
                     .max_values(spl_token_2022::instruction::MAX_SIGNERS)
                     .help("Member signer of a multisig account")
@@ -348,19 +346,24 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let mut wallet_manager: Option<Rc<RemoteWalletManager>> = None;
 
     let config = {
-        let cli_config = if let Some(config_file) = matches.value_of("config_file") {
+        let cli_config = if let Some(config_file) = matches.try_get_one::<String>("config_file")? {
             solana_cli_config::Config::load(config_file).unwrap_or_default()
         } else {
             solana_cli_config::Config::default()
         };
 
-        let payer = DefaultSigner::new(
-            "payer",
-            matches
-                .value_of("payer")
-                .map(|s| s.to_string())
-                .unwrap_or_else(|| cli_config.keypair_path.clone()),
-        );
+        let payer = if let Ok(Some((signer, _))) =
+            SignerSource::try_get_signer(matches, "payer", &mut wallet_manager)
+        {
+            Box::new(signer)
+        } else {
+            signer_from_path(
+                matches,
+                &cli_config.keypair_path,
+                "payer",
+                &mut wallet_manager,
+            )?
+        };
 
         let json_rpc_url = normalize_to_url_if_moniker(
             matches
@@ -370,16 +373,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
         Config {
             commitment_config: CommitmentConfig::confirmed(),
-            payer: Arc::from(
-                payer
-                    .signer_from_path(matches, &mut wallet_manager)
-                    .unwrap_or_else(|err| {
-                        eprintln!("error: {}", err);
-                        exit(1);
-                    }),
-            ),
+            payer: Arc::from(payer),
             json_rpc_url,
-            verbose: matches.is_present("verbose"),
+            verbose: matches.try_contains_id("verbose")?,
         }
     };
     solana_logger::setup_with_default("solana=info");
@@ -394,17 +390,17 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     match (command, matches) {
         ("create-escrow", arg_matches) => {
-            let original_mint = try_pubkey_of(arg_matches, "original_mint")
-                .unwrap()
-                .unwrap();
-            let new_mint = try_pubkey_of(arg_matches, "new_mint").unwrap().unwrap();
-            let account_keypair = matches.value_of("account_keypair").map(|path| {
-                signer_from_path(arg_matches, path, "account_keypair", &mut wallet_manager)
-                    .unwrap_or_else(|err| {
-                        eprintln!("error: account keypair: {}", err);
-                        exit(1);
-                    })
-            });
+            let original_mint =
+                SignerSource::try_get_pubkey(arg_matches, "original_mint", &mut wallet_manager)
+                    .unwrap()
+                    .unwrap();
+            let new_mint =
+                SignerSource::try_get_pubkey(arg_matches, "new_mint", &mut wallet_manager)
+                    .unwrap()
+                    .unwrap();
+            let account_keypair =
+                SignerSource::try_get_signer(matches, "account_keypair", &mut wallet_manager)?
+                    .map(|(signer, _)| signer);
             let response = process_create_escrow_account(
                 &rpc_client,
                 &config.payer,
@@ -422,14 +418,16 @@ async fn main() -> Result<(), Box<dyn Error>> {
         ("exchange", arg_matches) => {
             let mut bulk_signers = vec![config.payer.clone()];
             let mut multisig_pubkeys = vec![];
-            if let Some(values) = arg_matches.values_of("multisig_signer") {
-                for (i, value) in values.enumerate() {
+
+            if let Some(sources) = arg_matches.try_get_many::<SignerSource>("multisig_signer")? {
+                for (i, source) in sources.enumerate() {
                     let name = format!("{}-{}", "multisig_signer", i.saturating_add(1));
-                    let signer = signer_from_path(arg_matches, value, &name, &mut wallet_manager)
-                        .unwrap_or_else(|e| {
-                            eprintln!("error parsing multisig signer: {}", e);
-                            exit(1);
-                        });
+                    let signer =
+                        signer_from_source(arg_matches, source, &name, &mut wallet_manager)
+                            .unwrap_or_else(|e| {
+                                eprint!("error parsing multisig signer: {}", e);
+                                exit(1);
+                            });
                     let signer_pubkey = signer.pubkey();
                     let signer = Arc::from(signer);
                     if !bulk_signers.contains(&signer) {
@@ -441,35 +439,35 @@ async fn main() -> Result<(), Box<dyn Error>> {
                 }
             }
 
-            let original_mint = try_pubkey_of(arg_matches, "original_mint")
-                .unwrap()
-                .unwrap();
-            let new_mint = try_pubkey_of(arg_matches, "new_mint").unwrap().unwrap();
+            let original_mint =
+                SignerSource::try_get_pubkey(arg_matches, "original_mint", &mut wallet_manager)
+                    .unwrap()
+                    .unwrap();
+            let new_mint =
+                SignerSource::try_get_pubkey(arg_matches, "new_mint", &mut wallet_manager)
+                    .unwrap()
+                    .unwrap();
             let signer_config = SignerFromPathConfig {
                 allow_null_signer: !multisig_pubkeys.is_empty(),
             };
-            let owner = matches
-                .value_of("owner")
-                .map_or(Ok(config.payer.clone()), |path| {
-                    signer_from_path_with_config(
-                        arg_matches,
-                        path,
-                        "owner",
-                        &mut wallet_manager,
-                        &signer_config,
-                    )
-                    .map(Arc::from)
-                })
-                .unwrap_or_else(|err| {
-                    eprintln!("error: owner signer: {}", err);
-                    exit(1);
-                });
+            let owner = if let Ok(Some((signer, _))) =
+                SignerSource::try_get_signer(matches, "owner", &mut wallet_manager)
+            {
+                Arc::from(signer)
+            } else {
+                config.payer.clone()
+            };
             if !signer_config.allow_null_signer && !bulk_signers.contains(&owner) {
                 bulk_signers.push(owner.clone());
             }
-            let burn_from = try_pubkey_of(arg_matches, "burn_from").unwrap();
-            let escrow = try_pubkey_of(arg_matches, "escrow").unwrap();
-            let destination = try_pubkey_of(arg_matches, "destination").unwrap();
+            let burn_from =
+                SignerSource::try_get_pubkey(arg_matches, "burn_from", &mut wallet_manager)
+                    .unwrap();
+            let escrow =
+                SignerSource::try_get_pubkey(arg_matches, "escrow", &mut wallet_manager).unwrap();
+            let destination =
+                SignerSource::try_get_pubkey(arg_matches, "destination", &mut wallet_manager)
+                    .unwrap();
 
             let signature = process_exchange(
                 &rpc_client,
diff --git a/token/confidential-transfer/proof-extraction/Cargo.toml b/token/confidential-transfer/proof-extraction/Cargo.toml
index fa1c92b8e..e7e091622 100644
--- a/token/confidential-transfer/proof-extraction/Cargo.toml
+++ b/token/confidential-transfer/proof-extraction/Cargo.toml
@@ -8,6 +8,8 @@ license = "Apache-2.0"
 edition = "2021"
 
 [dependencies]
+bytemuck = "1.16.1"
+solana-curve25519 = "2.0.0"
 solana-zk-sdk = "2.0.0"
 thiserror = "1.0.62"
 
diff --git a/token/confidential-transfer/proof-extraction/src/encryption.rs b/token/confidential-transfer/proof-extraction/src/encryption.rs
index d266688dd..49db2b61c 100644
--- a/token/confidential-transfer/proof-extraction/src/encryption.rs
+++ b/token/confidential-transfer/proof-extraction/src/encryption.rs
@@ -1,5 +1,11 @@
-use solana_zk_sdk::encryption::pod::grouped_elgamal::PodGroupedElGamalCiphertext3Handles;
+use solana_zk_sdk::encryption::pod::grouped_elgamal::{
+    PodGroupedElGamalCiphertext2Handles, PodGroupedElGamalCiphertext3Handles,
+};
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 #[repr(C)]
 pub struct PodTransferAmountCiphertext(pub(crate) PodGroupedElGamalCiphertext3Handles);
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[repr(C)]
+pub struct PodFeeCiphertext(pub(crate) PodGroupedElGamalCiphertext2Handles);
diff --git a/token/confidential-transfer/proof-extraction/src/errors.rs b/token/confidential-transfer/proof-extraction/src/errors.rs
index 57b3f052c..4e05a11a2 100644
--- a/token/confidential-transfer/proof-extraction/src/errors.rs
+++ b/token/confidential-transfer/proof-extraction/src/errors.rs
@@ -8,4 +8,8 @@ pub enum TokenProofExtractionError {
     PedersenCommitmentMismatch,
     #[error("Range proof length mismatch")]
     RangeProofLengthMismatch,
+    #[error("Fee parameters mismatch")]
+    FeeParametersMismatch,
+    #[error("Curve arithmetic failed")]
+    CurveArithmetic,
 }
diff --git a/token/confidential-transfer/proof-extraction/src/lib.rs b/token/confidential-transfer/proof-extraction/src/lib.rs
index f7a6dcc3d..2933275b0 100644
--- a/token/confidential-transfer/proof-extraction/src/lib.rs
+++ b/token/confidential-transfer/proof-extraction/src/lib.rs
@@ -1,3 +1,4 @@
 pub mod encryption;
 pub mod errors;
 pub mod transfer;
+pub mod transfer_with_fee;
diff --git a/token/confidential-transfer/proof-extraction/src/transfer_with_fee.rs b/token/confidential-transfer/proof-extraction/src/transfer_with_fee.rs
new file mode 100644
index 000000000..88aa610c1
--- /dev/null
+++ b/token/confidential-transfer/proof-extraction/src/transfer_with_fee.rs
@@ -0,0 +1,322 @@
+use {
+    crate::{
+        encryption::{PodFeeCiphertext, PodTransferAmountCiphertext},
+        errors::TokenProofExtractionError,
+    },
+    bytemuck::bytes_of,
+    solana_curve25519::{
+        ristretto::{self, PodRistrettoPoint},
+        scalar::PodScalar,
+    },
+    solana_zk_sdk::{
+        encryption::pod::{
+            elgamal::{PodElGamalCiphertext, PodElGamalPubkey},
+            pedersen::PodPedersenCommitment,
+        },
+        zk_elgamal_proof_program::proof_data::{
+            BatchedGroupedCiphertext2HandlesValidityProofContext,
+            BatchedGroupedCiphertext3HandlesValidityProofContext, BatchedRangeProofContext,
+            CiphertextCommitmentEqualityProofContext, PercentageWithCapProofContext,
+        },
+    },
+};
+
+const MAX_FEE_BASIS_POINTS: u64 = 10_000;
+const REMAINING_BALANCE_BIT_LENGTH: u8 = 64;
+const TRANSFER_AMOUNT_LO_BIT_LENGTH: u8 = 16;
+const TRANSFER_AMOUNT_HI_BIT_LENGTH: u8 = 32;
+const DELTA_BIT_LENGTH: u8 = 48;
+const FEE_AMOUNT_LO_BIT_LENGTH: u8 = 16;
+const FEE_AMOUNT_HI_BIT_LENGTH: u8 = 32;
+
+/// The transfer public keys associated with a transfer with fee.
+pub struct TransferWithFeePubkeys {
+    /// Source ElGamal public key
+    pub source: PodElGamalPubkey,
+    /// Destination ElGamal public key
+    pub destination: PodElGamalPubkey,
+    /// Auditor ElGamal public key
+    pub auditor: PodElGamalPubkey,
+    /// Withdraw withheld authority public key
+    pub withdraw_withheld_authority: PodElGamalPubkey,
+}
+
+/// The proof context information needed to process a [Transfer] instruction
+/// with fee.
+pub struct TransferWithFeeProofContext {
+    /// Group encryption of the low 16 bits of the transfer amount
+    pub ciphertext_lo: PodTransferAmountCiphertext,
+    /// Group encryption of the high 48 bits of the transfer amount
+    pub ciphertext_hi: PodTransferAmountCiphertext,
+    /// The public encryption keys associated with the transfer: source, dest,
+    /// auditor, and withdraw withheld authority
+    pub transfer_with_fee_pubkeys: TransferWithFeePubkeys,
+    /// The final spendable ciphertext after the transfer,
+    pub new_source_ciphertext: PodElGamalCiphertext,
+    /// The transfer fee encryption of the low 16 bits of the transfer fee
+    /// amount
+    pub fee_ciphertext_lo: PodFeeCiphertext,
+    /// The transfer fee encryption of the hi 32 bits of the transfer fee amount
+    pub fee_ciphertext_hi: PodFeeCiphertext,
+}
+
+impl TransferWithFeeProofContext {
+    pub fn verify_and_extract(
+        equality_proof_context: &CiphertextCommitmentEqualityProofContext,
+        transfer_amount_ciphertext_validity_proof_context: &BatchedGroupedCiphertext3HandlesValidityProofContext,
+        fee_sigma_proof_context: &PercentageWithCapProofContext,
+        fee_ciphertext_validity_proof_context: &BatchedGroupedCiphertext2HandlesValidityProofContext,
+        range_proof_context: &BatchedRangeProofContext,
+        expected_fee_rate_basis_points: u16,
+        expected_maximum_fee: u64,
+    ) -> Result<Self, TokenProofExtractionError> {
+        // The equality proof context consists of the source ElGamal public key, the new
+        // source available balance ciphertext, and the new source available
+        // commitment. The public key and ciphertext should be returned as part
+        // of `TransferWithFeeProofContextInfo` and the commitment should be
+        // checked with range proof for consistency.
+        let CiphertextCommitmentEqualityProofContext {
+            pubkey: source_pubkey_from_equality_proof,
+            ciphertext: new_source_ciphertext,
+            commitment: new_source_commitment,
+        } = equality_proof_context;
+
+        // The transfer amount ciphertext validity proof context consists of the
+        // destination ElGamal public key, auditor ElGamal public key, and the
+        // transfer amount ciphertexts. All of these fields should be returned
+        // as part of `TransferWithFeeProofContextInfo`. In addition, the
+        // commitments pertaining to the transfer amount ciphertexts should be
+        // checked with range proof for consistency.
+        let BatchedGroupedCiphertext3HandlesValidityProofContext {
+            first_pubkey: source_pubkey_from_validity_proof,
+            second_pubkey: destination_pubkey,
+            third_pubkey: auditor_pubkey,
+            grouped_ciphertext_lo: transfer_amount_ciphertext_lo,
+            grouped_ciphertext_hi: transfer_amount_ciphertext_hi,
+        } = transfer_amount_ciphertext_validity_proof_context;
+
+        // The fee sigma proof context consists of the fee commitment, delta commitment,
+        // claimed commitment, and max fee. The fee and claimed commitment
+        // should be checked with range proof for consistency. The delta
+        // commitment should be checked whether it is properly generated with
+        // respect to the fee parameters. The max fee should be checked for
+        // consistency with the fee parameters.
+        let PercentageWithCapProofContext {
+            percentage_commitment: fee_commitment,
+            delta_commitment,
+            claimed_commitment,
+            max_value: proof_maximum_fee,
+        } = fee_sigma_proof_context;
+
+        let proof_maximum_fee: u64 = (*proof_maximum_fee).into();
+        if expected_maximum_fee != proof_maximum_fee {
+            return Err(TokenProofExtractionError::FeeParametersMismatch);
+        }
+
+        // The transfer fee ciphertext validity proof context consists of the
+        // destination ElGamal public key, withdraw withheld authority ElGamal
+        // public key, and the transfer fee ciphertexts. The rest of the fields
+        // should be return as part of `TransferWithFeeProofContextInfo`. In
+        // addition, the destination public key should be checked for
+        // consistency with the destination public key contained in the transfer amount
+        // ciphertext validity proof, and the commitments pertaining to the transfer fee
+        // amount ciphertexts should be checked with range proof for
+        // consistency.
+        let BatchedGroupedCiphertext2HandlesValidityProofContext {
+            first_pubkey: destination_pubkey_from_transfer_fee_validity_proof,
+            second_pubkey: withdraw_withheld_authority_pubkey,
+            grouped_ciphertext_lo: fee_ciphertext_lo,
+            grouped_ciphertext_hi: fee_ciphertext_hi,
+        } = fee_ciphertext_validity_proof_context;
+
+        if destination_pubkey != destination_pubkey_from_transfer_fee_validity_proof {
+            return Err(TokenProofExtractionError::ElGamalPubkeyMismatch);
+        }
+
+        // The range proof context consists of the Pedersen commitments and bit-lengths
+        // for which the range proof is proved. The commitments must consist of
+        // seven commitments pertaining to
+        // - the new source available balance (64 bits)
+        // - the low bits of the transfer amount (16 bits)
+        // - the high bits of the transfer amount (32 bits)
+        // - the delta amount for the fee (48 bits)
+        // - the complement of the delta amount for the fee (48 bits)
+        // - the low bits of the fee amount (16 bits)
+        // - the high bits of the fee amount (32 bits)
+        let BatchedRangeProofContext {
+            commitments: range_proof_commitments,
+            bit_lengths: range_proof_bit_lengths,
+        } = range_proof_context;
+
+        // check that the range proof was created for the correct set of Pedersen
+        // commitments
+        let transfer_amount_commitment_lo = transfer_amount_ciphertext_lo.extract_commitment();
+        let transfer_amount_commitment_hi = transfer_amount_ciphertext_hi.extract_commitment();
+
+        let fee_commitment_lo = fee_ciphertext_lo.extract_commitment();
+        let fee_commitment_hi = fee_ciphertext_hi.extract_commitment();
+
+        let max_fee_basis_points_scalar = u64_to_scalar(MAX_FEE_BASIS_POINTS);
+        let max_fee_basis_points_commitment =
+            ristretto::multiply_ristretto(&max_fee_basis_points_scalar, &G)
+                .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+        let claimed_complement_commitment = ristretto::subtract_ristretto(
+            &max_fee_basis_points_commitment,
+            &commitment_to_ristretto(claimed_commitment),
+        )
+        .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+
+        let expected_commitments = [
+            bytes_of(new_source_commitment),
+            bytes_of(&transfer_amount_commitment_lo),
+            bytes_of(&transfer_amount_commitment_hi),
+            bytes_of(claimed_commitment),
+            bytes_of(&claimed_complement_commitment),
+            bytes_of(&fee_commitment_lo),
+            bytes_of(&fee_commitment_hi),
+        ];
+
+        if !range_proof_commitments
+            .iter()
+            .zip(expected_commitments.into_iter())
+            .all(|(proof_commitment, expected_commitment)| {
+                bytes_of(proof_commitment) == expected_commitment
+            })
+        {
+            return Err(TokenProofExtractionError::PedersenCommitmentMismatch);
+        }
+
+        // check that the range proof was created for the correct number of bits
+        let expected_bit_lengths = [
+            REMAINING_BALANCE_BIT_LENGTH,
+            TRANSFER_AMOUNT_LO_BIT_LENGTH,
+            TRANSFER_AMOUNT_HI_BIT_LENGTH,
+            DELTA_BIT_LENGTH,
+            DELTA_BIT_LENGTH,
+            FEE_AMOUNT_LO_BIT_LENGTH,
+            FEE_AMOUNT_HI_BIT_LENGTH,
+        ]
+        .iter();
+
+        if !range_proof_bit_lengths
+            .iter()
+            .zip(expected_bit_lengths)
+            .all(|(proof_len, expected_len)| proof_len == expected_len)
+        {
+            return Err(TokenProofExtractionError::RangeProofLengthMismatch);
+        }
+
+        // check consistency between fee sigma and fee ciphertext validity proofs
+        let sigma_proof_fee_commitment_point: PodRistrettoPoint =
+            commitment_to_ristretto(fee_commitment);
+        let validity_proof_fee_point = combine_lo_hi_pedersen_points(
+            &commitment_to_ristretto(&fee_commitment_lo),
+            &commitment_to_ristretto(&fee_commitment_hi),
+        )
+        .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+
+        if source_pubkey_from_equality_proof != source_pubkey_from_validity_proof {
+            return Err(TokenProofExtractionError::ElGamalPubkeyMismatch);
+        }
+
+        if validity_proof_fee_point != sigma_proof_fee_commitment_point {
+            return Err(TokenProofExtractionError::FeeParametersMismatch);
+        }
+
+        verify_delta_commitment(
+            &transfer_amount_commitment_lo,
+            &transfer_amount_commitment_hi,
+            fee_commitment,
+            delta_commitment,
+            expected_fee_rate_basis_points,
+        )?;
+
+        // create transfer with fee proof context info and return
+        let transfer_with_fee_pubkeys = TransferWithFeePubkeys {
+            source: *source_pubkey_from_equality_proof,
+            destination: *destination_pubkey,
+            auditor: *auditor_pubkey,
+            withdraw_withheld_authority: *withdraw_withheld_authority_pubkey,
+        };
+
+        Ok(Self {
+            ciphertext_lo: PodTransferAmountCiphertext(*transfer_amount_ciphertext_lo),
+            ciphertext_hi: PodTransferAmountCiphertext(*transfer_amount_ciphertext_hi),
+            transfer_with_fee_pubkeys,
+            new_source_ciphertext: *new_source_ciphertext,
+            fee_ciphertext_lo: PodFeeCiphertext(*fee_ciphertext_lo),
+            fee_ciphertext_hi: PodFeeCiphertext(*fee_ciphertext_hi),
+        })
+    }
+}
+
+/// Ristretto generator point for curve25519
+const G: PodRistrettoPoint = PodRistrettoPoint([
+    226, 242, 174, 10, 106, 188, 78, 113, 168, 132, 169, 97, 197, 0, 81, 95, 88, 227, 11, 106, 165,
+    130, 221, 141, 182, 166, 89, 69, 224, 141, 45, 118,
+]);
+
+/// Convert a `u64` amount into a curve25519 scalar
+fn u64_to_scalar(amount: u64) -> PodScalar {
+    let mut bytes = [0u8; 32];
+    bytes[..8].copy_from_slice(&amount.to_le_bytes());
+    PodScalar(bytes)
+}
+
+/// Convert a `u16` amount into a curve25519 scalar
+fn u16_to_scalar(amount: u16) -> PodScalar {
+    let mut bytes = [0u8; 32];
+    bytes[..2].copy_from_slice(&amount.to_le_bytes());
+    PodScalar(bytes)
+}
+
+fn commitment_to_ristretto(commitment: &PodPedersenCommitment) -> PodRistrettoPoint {
+    let mut bytes = [0u8; 32];
+    bytes.copy_from_slice(bytes_of(commitment));
+    PodRistrettoPoint(bytes)
+}
+
+/// Combine lo and hi Pedersen commitment points
+fn combine_lo_hi_pedersen_points(
+    point_lo: &PodRistrettoPoint,
+    point_hi: &PodRistrettoPoint,
+) -> Option<PodRistrettoPoint> {
+    const SCALING_CONSTANT: u64 = 65536;
+    let scaling_constant_scalar = u64_to_scalar(SCALING_CONSTANT);
+    let scaled_point_hi = ristretto::multiply_ristretto(&scaling_constant_scalar, point_hi)?;
+    ristretto::add_ristretto(point_lo, &scaled_point_hi)
+}
+
+/// Compute fee delta commitment
+fn verify_delta_commitment(
+    transfer_amount_commitment_lo: &PodPedersenCommitment,
+    transfer_amount_commitment_hi: &PodPedersenCommitment,
+    fee_commitment: &PodPedersenCommitment,
+    proof_delta_commitment: &PodPedersenCommitment,
+    transfer_fee_basis_points: u16,
+) -> Result<(), TokenProofExtractionError> {
+    let transfer_amount_point = combine_lo_hi_pedersen_points(
+        &commitment_to_ristretto(transfer_amount_commitment_lo),
+        &commitment_to_ristretto(transfer_amount_commitment_hi),
+    )
+    .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+    let transfer_fee_basis_points_scalar = u16_to_scalar(transfer_fee_basis_points);
+    let scaled_transfer_amount_point =
+        ristretto::multiply_ristretto(&transfer_fee_basis_points_scalar, &transfer_amount_point)
+            .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+
+    let max_fee_basis_points_scalar = u64_to_scalar(MAX_FEE_BASIS_POINTS);
+    let fee_point: PodRistrettoPoint = commitment_to_ristretto(fee_commitment);
+    let scaled_fee_point = ristretto::multiply_ristretto(&max_fee_basis_points_scalar, &fee_point)
+        .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+
+    let expected_delta_commitment_point =
+        ristretto::subtract_ristretto(&scaled_fee_point, &scaled_transfer_amount_point)
+            .ok_or(TokenProofExtractionError::CurveArithmetic)?;
+
+    let proof_delta_commitment_point = commitment_to_ristretto(proof_delta_commitment);
+    if expected_delta_commitment_point != proof_delta_commitment_point {
+        return Err(TokenProofExtractionError::CurveArithmetic);
+    }
+    Ok(())
+}
diff --git a/token/confidential-transfer/proof-generation/src/encryption.rs b/token/confidential-transfer/proof-generation/src/encryption.rs
index 389b6627d..8583cc712 100644
--- a/token/confidential-transfer/proof-generation/src/encryption.rs
+++ b/token/confidential-transfer/proof-generation/src/encryption.rs
@@ -47,3 +47,42 @@ impl TransferAmountCiphertext {
         self.0.handles.get(2).unwrap()
     }
 }
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[repr(C)]
+#[cfg(not(target_os = "solana"))]
+pub struct FeeCiphertext(pub(crate) GroupedElGamalCiphertext<2>);
+
+#[cfg(not(target_os = "solana"))]
+impl FeeCiphertext {
+    pub fn new(
+        amount: u64,
+        destination_pubkey: &ElGamalPubkey,
+        withdraw_withheld_authority_pubkey: &ElGamalPubkey,
+    ) -> (Self, PedersenOpening) {
+        let opening = PedersenOpening::new_rand();
+        let grouped_ciphertext = GroupedElGamal::<2>::encrypt_with(
+            [destination_pubkey, withdraw_withheld_authority_pubkey],
+            amount,
+            &opening,
+        );
+
+        (Self(grouped_ciphertext), opening)
+    }
+
+    pub fn get_commitment(&self) -> &PedersenCommitment {
+        &self.0.commitment
+    }
+
+    pub fn get_destination_handle(&self) -> &DecryptHandle {
+        // `FeeEncryption` is a wrapper for `GroupedElGamalCiphertext<2>`, which holds
+        // exactly two decryption handles.
+        self.0.handles.first().unwrap()
+    }
+
+    pub fn get_withdraw_withheld_authority_handle(&self) -> &DecryptHandle {
+        // `FeeEncryption` is a wrapper for `GroupedElGamalCiphertext<2>`, which holds
+        // exactly two decryption handles.
+        self.0.handles.get(1).unwrap()
+    }
+}
diff --git a/token/confidential-transfer/proof-generation/src/errors.rs b/token/confidential-transfer/proof-generation/src/errors.rs
index d3f6a0fc0..5cb5c14e7 100644
--- a/token/confidential-transfer/proof-generation/src/errors.rs
+++ b/token/confidential-transfer/proof-generation/src/errors.rs
@@ -8,4 +8,6 @@ pub enum TokenProofGenerationError {
     NotEnoughFunds,
     #[error("illegal amount bit length")]
     IllegalAmountBitLength,
+    #[error("fee calculation failed")]
+    FeeCalculation,
 }
diff --git a/token/confidential-transfer/proof-generation/src/lib.rs b/token/confidential-transfer/proof-generation/src/lib.rs
index 398055b36..ce8b0ea1c 100644
--- a/token/confidential-transfer/proof-generation/src/lib.rs
+++ b/token/confidential-transfer/proof-generation/src/lib.rs
@@ -9,6 +9,7 @@ use {
 pub mod encryption;
 pub mod errors;
 pub mod transfer;
+pub mod transfer_with_fee;
 
 /// The low bit length of the encrypted transfer amount
 pub const TRANSFER_AMOUNT_LO_BITS: usize = 16;
diff --git a/token/confidential-transfer/proof-generation/src/transfer_with_fee.rs b/token/confidential-transfer/proof-generation/src/transfer_with_fee.rs
new file mode 100644
index 000000000..d0705e9f1
--- /dev/null
+++ b/token/confidential-transfer/proof-generation/src/transfer_with_fee.rs
@@ -0,0 +1,339 @@
+use {
+    crate::{
+        encryption::{FeeCiphertext, TransferAmountCiphertext},
+        errors::TokenProofGenerationError,
+        try_combine_lo_hi_ciphertexts, try_combine_lo_hi_commitments, try_combine_lo_hi_openings,
+        try_split_u64, TRANSFER_AMOUNT_HI_BITS, TRANSFER_AMOUNT_LO_BITS,
+    },
+    curve25519_dalek::scalar::Scalar,
+    solana_zk_sdk::{
+        encryption::{
+            auth_encryption::{AeCiphertext, AeKey},
+            elgamal::{ElGamalCiphertext, ElGamalKeypair, ElGamalPubkey},
+            grouped_elgamal::GroupedElGamal,
+            pedersen::{Pedersen, PedersenCommitment, PedersenOpening},
+        },
+        zk_elgamal_proof_program::proof_data::{
+            BatchedGroupedCiphertext2HandlesValidityProofData,
+            BatchedGroupedCiphertext3HandlesValidityProofData, BatchedRangeProofU256Data,
+            CiphertextCommitmentEqualityProofData, PercentageWithCapProofData,
+        },
+    },
+};
+
+const MAX_FEE_BASIS_POINTS: u64 = 10_000;
+const ONE_IN_BASIS_POINTS: u128 = MAX_FEE_BASIS_POINTS as u128;
+
+const FEE_AMOUNT_LO_BITS: usize = 16;
+const FEE_AMOUNT_HI_BITS: usize = 32;
+
+const REMAINING_BALANCE_BIT_LENGTH: usize = 64;
+const DELTA_BIT_LENGTH: usize = 48;
+
+#[allow(clippy::too_many_arguments)]
+pub fn transfer_with_fee_split_proof_data(
+    current_available_balance: &ElGamalCiphertext,
+    current_decryptable_available_balance: &AeCiphertext,
+    transfer_amount: u64,
+    source_elgamal_keypair: &ElGamalKeypair,
+    aes_key: &AeKey,
+    destination_elgamal_pubkey: &ElGamalPubkey,
+    auditor_elgamal_pubkey: Option<&ElGamalPubkey>,
+    withdraw_withheld_authority_elgamal_pubkey: &ElGamalPubkey,
+    fee_rate_basis_points: u16,
+    maximum_fee: u64,
+) -> Result<
+    (
+        CiphertextCommitmentEqualityProofData,
+        BatchedGroupedCiphertext3HandlesValidityProofData,
+        PercentageWithCapProofData,
+        BatchedGroupedCiphertext2HandlesValidityProofData,
+        BatchedRangeProofU256Data,
+    ),
+    TokenProofGenerationError,
+> {
+    let default_auditor_pubkey = ElGamalPubkey::default();
+    let auditor_elgamal_pubkey = auditor_elgamal_pubkey.unwrap_or(&default_auditor_pubkey);
+
+    // Split the transfer amount into the low and high bit components
+    let (transfer_amount_lo, transfer_amount_hi) =
+        try_split_u64(transfer_amount, TRANSFER_AMOUNT_LO_BITS)
+            .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+
+    // Encrypt the `lo` and `hi` transfer amounts
+    let (transfer_amount_grouped_ciphertext_lo, transfer_amount_opening_lo) =
+        TransferAmountCiphertext::new(
+            transfer_amount_lo,
+            source_elgamal_keypair.pubkey(),
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+        );
+
+    let (transfer_amount_grouped_ciphertext_hi, transfer_amount_opening_hi) =
+        TransferAmountCiphertext::new(
+            transfer_amount_hi,
+            source_elgamal_keypair.pubkey(),
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+        );
+
+    // Decrypt the current available balance at the source
+    let current_decrypted_available_balance = current_decryptable_available_balance
+        .decrypt(aes_key)
+        .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+
+    // Compute the remaining balance at the source
+    let new_decrypted_available_balance = current_decrypted_available_balance
+        .checked_sub(transfer_amount)
+        .ok_or(TokenProofGenerationError::NotEnoughFunds)?;
+
+    // Create a new Pedersen commitment for the remaining balance at the source
+    let (new_available_balance_commitment, new_source_opening) =
+        Pedersen::new(new_decrypted_available_balance);
+
+    // Compute the remaining balance at the source as ElGamal ciphertexts
+    let transfer_amount_source_ciphertext_lo = transfer_amount_grouped_ciphertext_lo
+        .0
+        .to_elgamal_ciphertext(0)
+        .unwrap();
+
+    let transfer_amount_source_ciphertext_hi = transfer_amount_grouped_ciphertext_hi
+        .0
+        .to_elgamal_ciphertext(0)
+        .unwrap();
+
+    #[allow(clippy::arithmetic_side_effects)]
+    let new_available_balance_ciphertext = current_available_balance
+        - try_combine_lo_hi_ciphertexts(
+            &transfer_amount_source_ciphertext_lo,
+            &transfer_amount_source_ciphertext_hi,
+            TRANSFER_AMOUNT_LO_BITS,
+        )
+        .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+
+    // generate equality proof data
+    let equality_proof_data = CiphertextCommitmentEqualityProofData::new(
+        source_elgamal_keypair,
+        &new_available_balance_ciphertext,
+        &new_available_balance_commitment,
+        &new_source_opening,
+        new_decrypted_available_balance,
+    )
+    .map_err(TokenProofGenerationError::from)?;
+
+    // generate ciphertext validity data
+    let transfer_amount_ciphertext_validity_proof_data =
+        BatchedGroupedCiphertext3HandlesValidityProofData::new(
+            source_elgamal_keypair.pubkey(),
+            destination_elgamal_pubkey,
+            auditor_elgamal_pubkey,
+            &transfer_amount_grouped_ciphertext_lo.0,
+            &transfer_amount_grouped_ciphertext_hi.0,
+            transfer_amount_lo,
+            transfer_amount_hi,
+            &transfer_amount_opening_lo,
+            &transfer_amount_opening_hi,
+        )
+        .map_err(TokenProofGenerationError::from)?;
+
+    // calculate fee
+    let transfer_fee_basis_points = fee_rate_basis_points;
+    let transfer_fee_maximum_fee = maximum_fee;
+    let (raw_fee_amount, delta_fee) = calculate_fee(transfer_amount, transfer_fee_basis_points)
+        .ok_or(TokenProofGenerationError::FeeCalculation)?;
+
+    // if raw fee is greater than the maximum fee, then use the maximum fee for the
+    // fee amount
+    let fee_amount = std::cmp::min(transfer_fee_maximum_fee, raw_fee_amount);
+
+    // split and encrypt fee
+    let (fee_amount_lo, fee_amount_hi) = try_split_u64(fee_amount, FEE_AMOUNT_LO_BITS)
+        .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+    let (fee_ciphertext_lo, fee_opening_lo) = FeeCiphertext::new(
+        fee_amount_lo,
+        destination_elgamal_pubkey,
+        withdraw_withheld_authority_elgamal_pubkey,
+    );
+    let (fee_ciphertext_hi, fee_opening_hi) = FeeCiphertext::new(
+        fee_amount_hi,
+        destination_elgamal_pubkey,
+        withdraw_withheld_authority_elgamal_pubkey,
+    );
+
+    // create combined commitments and openings to be used to generate proofs
+    let combined_transfer_amount_commitment = try_combine_lo_hi_commitments(
+        transfer_amount_grouped_ciphertext_lo.get_commitment(),
+        transfer_amount_grouped_ciphertext_hi.get_commitment(),
+        TRANSFER_AMOUNT_LO_BITS,
+    )
+    .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+    let combined_transfer_amount_opening = try_combine_lo_hi_openings(
+        &transfer_amount_opening_lo,
+        &transfer_amount_opening_hi,
+        TRANSFER_AMOUNT_LO_BITS,
+    )
+    .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+
+    let combined_fee_commitment = try_combine_lo_hi_commitments(
+        fee_ciphertext_lo.get_commitment(),
+        fee_ciphertext_hi.get_commitment(),
+        FEE_AMOUNT_LO_BITS,
+    )
+    .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+    let combined_fee_opening =
+        try_combine_lo_hi_openings(&fee_opening_lo, &fee_opening_hi, FEE_AMOUNT_LO_BITS)
+            .ok_or(TokenProofGenerationError::IllegalAmountBitLength)?;
+
+    // compute claimed and real delta commitment
+    let (claimed_commitment, claimed_opening) = Pedersen::new(delta_fee);
+    let (delta_commitment, delta_opening) = compute_delta_commitment_and_opening(
+        (
+            &combined_transfer_amount_commitment,
+            &combined_transfer_amount_opening,
+        ),
+        (&combined_fee_commitment, &combined_fee_opening),
+        transfer_fee_basis_points,
+    );
+
+    // generate fee sigma proof
+    let percentage_with_cap_proof_data = PercentageWithCapProofData::new(
+        &combined_fee_commitment,
+        &combined_fee_opening,
+        fee_amount,
+        &delta_commitment,
+        &delta_opening,
+        delta_fee,
+        &claimed_commitment,
+        &claimed_opening,
+        transfer_fee_maximum_fee,
+    )
+    .map_err(TokenProofGenerationError::from)?;
+
+    // encrypt the fee amount under the destination and withdraw withheld authority
+    // ElGamal public key
+    let fee_destination_withdraw_withheld_authority_ciphertext_lo = GroupedElGamal::encrypt_with(
+        [
+            destination_elgamal_pubkey,
+            withdraw_withheld_authority_elgamal_pubkey,
+        ],
+        fee_amount_lo,
+        &fee_opening_lo,
+    );
+    let fee_destination_withdraw_withheld_authority_ciphertext_hi = GroupedElGamal::encrypt_with(
+        [
+            destination_elgamal_pubkey,
+            withdraw_withheld_authority_elgamal_pubkey,
+        ],
+        fee_amount_hi,
+        &fee_opening_hi,
+    );
+
+    // generate fee ciphertext validity data
+    let fee_ciphertext_validity_proof_data =
+        BatchedGroupedCiphertext2HandlesValidityProofData::new(
+            destination_elgamal_pubkey,
+            withdraw_withheld_authority_elgamal_pubkey,
+            &fee_destination_withdraw_withheld_authority_ciphertext_lo,
+            &fee_destination_withdraw_withheld_authority_ciphertext_hi,
+            fee_amount_lo,
+            fee_amount_hi,
+            &fee_opening_lo,
+            &fee_opening_hi,
+        )
+        .map_err(TokenProofGenerationError::from)?;
+
+    // generate range proof data
+    let delta_fee_complement = MAX_FEE_BASIS_POINTS
+        .checked_sub(delta_fee)
+        .ok_or(TokenProofGenerationError::FeeCalculation)?;
+
+    let max_fee_basis_points_commitment =
+        Pedersen::with(MAX_FEE_BASIS_POINTS, &PedersenOpening::default());
+    #[allow(clippy::arithmetic_side_effects)]
+    let claimed_complement_commitment = max_fee_basis_points_commitment - claimed_commitment;
+    #[allow(clippy::arithmetic_side_effects)]
+    let claimed_complement_opening = PedersenOpening::default() - &claimed_opening;
+
+    let range_proof_data = BatchedRangeProofU256Data::new(
+        vec![
+            &new_available_balance_commitment,
+            transfer_amount_grouped_ciphertext_lo.get_commitment(),
+            transfer_amount_grouped_ciphertext_hi.get_commitment(),
+            &claimed_commitment,
+            &claimed_complement_commitment,
+            fee_ciphertext_lo.get_commitment(),
+            fee_ciphertext_hi.get_commitment(),
+        ],
+        vec![
+            new_decrypted_available_balance,
+            transfer_amount_lo,
+            transfer_amount_hi,
+            delta_fee,
+            delta_fee_complement,
+            fee_amount_lo,
+            fee_amount_hi,
+        ],
+        vec![
+            REMAINING_BALANCE_BIT_LENGTH,
+            TRANSFER_AMOUNT_LO_BITS,
+            TRANSFER_AMOUNT_HI_BITS,
+            DELTA_BIT_LENGTH,
+            DELTA_BIT_LENGTH,
+            FEE_AMOUNT_LO_BITS,
+            FEE_AMOUNT_HI_BITS,
+        ],
+        vec![
+            &new_source_opening,
+            &transfer_amount_opening_lo,
+            &transfer_amount_opening_hi,
+            &claimed_opening,
+            &claimed_complement_opening,
+            &fee_opening_lo,
+            &fee_opening_hi,
+        ],
+    )
+    .map_err(TokenProofGenerationError::from)?;
+
+    Ok((
+        equality_proof_data,
+        transfer_amount_ciphertext_validity_proof_data,
+        percentage_with_cap_proof_data,
+        fee_ciphertext_validity_proof_data,
+        range_proof_data,
+    ))
+}
+
+fn calculate_fee(transfer_amount: u64, fee_rate_basis_points: u16) -> Option<(u64, u64)> {
+    let numerator = (transfer_amount as u128).checked_mul(fee_rate_basis_points as u128)?;
+
+    // Warning: Division may involve CPU opcodes that have variable execution times.
+    // This non-constant-time execution of the fee calculation can theoretically
+    // reveal information about the transfer amount. For transfers that involve
+    // extremely sensitive data, additional care should be put into how the fees
+    // are calculated.
+    let fee = numerator
+        .checked_add(ONE_IN_BASIS_POINTS)?
+        .checked_sub(1)?
+        .checked_div(ONE_IN_BASIS_POINTS)?;
+
+    let delta_fee = fee
+        .checked_mul(ONE_IN_BASIS_POINTS)?
+        .checked_sub(numerator)?;
+
+    Some((fee as u64, delta_fee as u64))
+}
+
+#[allow(clippy::arithmetic_side_effects)]
+fn compute_delta_commitment_and_opening(
+    (combined_commitment, combined_opening): (&PedersenCommitment, &PedersenOpening),
+    (combined_fee_commitment, combined_fee_opening): (&PedersenCommitment, &PedersenOpening),
+    fee_rate_basis_points: u16,
+) -> (PedersenCommitment, PedersenOpening) {
+    let fee_rate_scalar = Scalar::from(fee_rate_basis_points);
+    let delta_commitment = combined_fee_commitment * Scalar::from(MAX_FEE_BASIS_POINTS)
+        - combined_commitment * fee_rate_scalar;
+    let delta_opening = combined_fee_opening * Scalar::from(MAX_FEE_BASIS_POINTS)
+        - combined_opening * fee_rate_scalar;
+
+    (delta_commitment, delta_opening)
+}


Test patch:
diff --git a/token/confidential-transfer/proof-tests/tests/proof_test.rs b/token/confidential-transfer/proof-tests/tests/proof_test.rs
index b18ed3ac5..f1d52ff1a 100644
--- a/token/confidential-transfer/proof-tests/tests/proof_test.rs
+++ b/token/confidential-transfer/proof-tests/tests/proof_test.rs
@@ -3,21 +3,25 @@ use {
         encryption::{auth_encryption::AeKey, elgamal::ElGamalKeypair},
         zk_elgamal_proof_program::proof_data::ZkProofData,
     },
-    spl_token_confidential_transfer_proof_extraction::transfer::TransferProofContext,
-    spl_token_confidential_transfer_proof_generation::transfer::transfer_split_proof_data,
+    spl_token_confidential_transfer_proof_extraction::{
+        transfer::TransferProofContext, transfer_with_fee::TransferWithFeeProofContext,
+    },
+    spl_token_confidential_transfer_proof_generation::{
+        transfer::transfer_split_proof_data, transfer_with_fee::transfer_with_fee_split_proof_data,
+    },
 };
 
 #[test]
 fn test_transfer_correctness() {
-    test_proof_validity(0, 0);
-    test_proof_validity(1, 0);
-    test_proof_validity(1, 1);
-    test_proof_validity(65535, 65535); // 2^16 - 1
-    test_proof_validity(65536, 65536); // 2^16
-    test_proof_validity(281474976710655, 281474976710655); // 2^48 - 1
+    test_transfer_proof_validity(0, 0);
+    test_transfer_proof_validity(1, 0);
+    test_transfer_proof_validity(1, 1);
+    test_transfer_proof_validity(65535, 65535); // 2^16 - 1
+    test_transfer_proof_validity(65536, 65536); // 2^16
+    test_transfer_proof_validity(281474976710655, 281474976710655); // 2^48 - 1
 }
 
-fn test_proof_validity(spendable_balance: u64, transfer_amount: u64) {
+fn test_transfer_proof_validity(spendable_balance: u64, transfer_amount: u64) {
     let source_keypair = ElGamalKeypair::new_rand();
 
     let aes_key = AeKey::new_rand();
@@ -53,3 +57,86 @@ fn test_proof_validity(spendable_balance: u64, transfer_amount: u64) {
     )
     .unwrap();
 }
+
+#[test]
+fn test_transfer_with_fee_correctness() {
+    test_transfer_with_fee_proof_validity(0, 0, 0, 0);
+    test_transfer_with_fee_proof_validity(0, 0, 0, 1);
+    test_transfer_with_fee_proof_validity(0, 0, 1, 0);
+    test_transfer_with_fee_proof_validity(0, 0, 1, 1);
+    test_transfer_with_fee_proof_validity(1, 0, 0, 0);
+    test_transfer_with_fee_proof_validity(1, 1, 0, 0);
+
+    test_transfer_with_fee_proof_validity(100, 100, 5, 10);
+    test_transfer_with_fee_proof_validity(100, 100, 5, 1);
+
+    test_transfer_with_fee_proof_validity(65535, 65535, 5, 10);
+    test_transfer_with_fee_proof_validity(65535, 65535, 5, 1);
+
+    test_transfer_with_fee_proof_validity(65536, 65536, 5, 10);
+    test_transfer_with_fee_proof_validity(65536, 65536, 5, 1);
+
+    test_transfer_with_fee_proof_validity(281474976710655, 281474976710655, 5, 10); // 2^48 - 1
+    test_transfer_with_fee_proof_validity(281474976710655, 281474976710655, 5, 1);
+}
+
+fn test_transfer_with_fee_proof_validity(
+    spendable_balance: u64,
+    transfer_amount: u64,
+    fee_rate_basis_points: u16,
+    maximum_fee: u64,
+) {
+    let source_keypair = ElGamalKeypair::new_rand();
+    let aes_key = AeKey::new_rand();
+
+    let destination_keypair = ElGamalKeypair::new_rand();
+    let destination_pubkey = destination_keypair.pubkey();
+
+    let auditor_keypair = ElGamalKeypair::new_rand();
+    let auditor_pubkey = auditor_keypair.pubkey();
+
+    let withdraw_withheld_authority_keyupair = ElGamalKeypair::new_rand();
+    let withdraw_withheld_authority_pubkey = withdraw_withheld_authority_keyupair.pubkey();
+
+    let spendable_ciphertext = source_keypair.pubkey().encrypt(spendable_balance);
+    let decryptable_balance = aes_key.encrypt(spendable_balance);
+
+    let (
+        equality_proof_data,
+        transfer_amount_ciphertext_validity_proof_data,
+        percentage_with_cap_proof_data,
+        fee_ciphertext_validity_proof_data,
+        range_proof_data,
+    ) = transfer_with_fee_split_proof_data(
+        &spendable_ciphertext,
+        &decryptable_balance,
+        transfer_amount,
+        &source_keypair,
+        &aes_key,
+        destination_pubkey,
+        Some(auditor_pubkey),
+        withdraw_withheld_authority_pubkey,
+        fee_rate_basis_points,
+        maximum_fee,
+    )
+    .unwrap();
+
+    equality_proof_data.verify_proof().unwrap();
+    transfer_amount_ciphertext_validity_proof_data
+        .verify_proof()
+        .unwrap();
+    percentage_with_cap_proof_data.verify_proof().unwrap();
+    fee_ciphertext_validity_proof_data.verify_proof().unwrap();
+    range_proof_data.verify_proof().unwrap();
+
+    TransferWithFeeProofContext::verify_and_extract(
+        equality_proof_data.context_data(),
+        transfer_amount_ciphertext_validity_proof_data.context_data(),
+        percentage_with_cap_proof_data.context_data(),
+        fee_ciphertext_validity_proof_data.context_data(),
+        range_proof_data.context_data(),
+        fee_rate_basis_points,
+        maximum_fee,
+    )
+    .unwrap();
+}