RemoteDesk/client/crates/client-core/src/pipe_peer.rs

use std::fmt;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HelperProcessIdentity {
    pub pid: u32,
    /// An opaque token obtained from the OS process creation time/identity API.
    pub creation_token: u64,
    pub build_hash: String,
}

impl HelperProcessIdentity {
    pub fn new(pid: u32, creation_token: u64, build_hash: impl Into<String>) -> Self {
        Self {
            pid,
            creation_token,
            build_hash: build_hash.into(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChallengeMacVerification {
    Verified,
    Rejected,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PipePeerError {
    NotAuthenticated,
    PidMismatch { expected: u32, actual: u32 },
    CreationTokenMismatch,
    BuildHashMismatch,
    ChallengeMacRejected,
}

impl fmt::Display for PipePeerError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotAuthenticated => formatter.write_str("pipe peer is not authenticated"),
            Self::PidMismatch { expected, actual } => {
                write!(
                    formatter,
                    "pipe peer PID mismatch: expected {expected}, got {actual}"
                )
            }
            Self::CreationTokenMismatch => {
                formatter.write_str("pipe peer process creation token mismatch")
            }
            Self::BuildHashMismatch => formatter.write_str("pipe peer build hash mismatch"),
            Self::ChallengeMacRejected => {
                formatter.write_str("pipe challenge MAC was not verified")
            }
        }
    }
}

impl std::error::Error for PipePeerError {}

#[derive(Clone, PartialEq, Eq)]
pub struct SessionMaterial(Vec<u8>);

impl fmt::Debug for SessionMaterial {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("SessionMaterial")
            .field("bytes", &"[REDACTED]")
            .finish()
    }
}

impl SessionMaterial {
    pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
        Self(bytes.into())
    }

    #[must_use]
    pub fn expose_to_verified_peer<'a>(&'a self, _peer: &VerifiedPipePeer) -> &'a [u8] {
        &self.0
    }
}

#[derive(Debug, Clone)]
pub struct PipePeerVerifier {
    expected: HelperProcessIdentity,
    verified_peer: Option<VerifiedPipePeer>,
}

impl PipePeerVerifier {
    #[must_use]
    pub fn new(expected: HelperProcessIdentity) -> Self {
        Self {
            expected,
            verified_peer: None,
        }
    }

    /// Authenticates the connected pipe peer against process and challenge evidence.
    ///
    /// # Errors
    ///
    /// Returns a [`PipePeerError`] when any expected identity field differs or
    /// the challenge MAC was not verified. A failed attempt clears prior peer
    /// authentication state.
    pub fn authenticate(
        &mut self,
        actual: HelperProcessIdentity,
        challenge_mac: ChallengeMacVerification,
    ) -> Result<&VerifiedPipePeer, PipePeerError> {
        self.verified_peer = None;

        if actual.pid != self.expected.pid {
            return Err(PipePeerError::PidMismatch {
                expected: self.expected.pid,
                actual: actual.pid,
            });
        }
        if actual.creation_token != self.expected.creation_token {
            return Err(PipePeerError::CreationTokenMismatch);
        }
        if actual.build_hash != self.expected.build_hash {
            return Err(PipePeerError::BuildHashMismatch);
        }
        if challenge_mac != ChallengeMacVerification::Verified {
            return Err(PipePeerError::ChallengeMacRejected);
        }

        self.verified_peer = Some(VerifiedPipePeer { identity: actual });
        self.verified_peer()
    }

    /// Returns the authenticated peer for this verifier.
    ///
    /// # Errors
    ///
    /// Returns [`PipePeerError::NotAuthenticated`] until authentication has
    /// completed successfully, and after any failed reauthentication attempt.
    pub fn verified_peer(&self) -> Result<&VerifiedPipePeer, PipePeerError> {
        self.verified_peer
            .as_ref()
            .ok_or(PipePeerError::NotAuthenticated)
    }

    /// Exposes session material only after successful pipe peer authentication.
    ///
    /// # Errors
    ///
    /// Returns [`PipePeerError::NotAuthenticated`] when no authenticated peer
    /// is currently bound to the verifier.
    pub fn session_material<'a>(
        &self,
        material: &'a SessionMaterial,
    ) -> Result<&'a [u8], PipePeerError> {
        let peer = self.verified_peer()?;
        Ok(material.expose_to_verified_peer(peer))
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VerifiedPipePeer {
    identity: HelperProcessIdentity,
}

impl VerifiedPipePeer {
    #[must_use]
    pub fn identity(&self) -> &HelperProcessIdentity {
        &self.identity
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn identity() -> HelperProcessIdentity {
        HelperProcessIdentity::new(41, 9001, "build-a")
    }

    #[test]
    fn session_material_is_unavailable_before_authentication() {
        let verifier = PipePeerVerifier::new(identity());
        let material = SessionMaterial::new(b"secret".to_vec());
        assert_eq!(
            verifier.session_material(&material),
            Err(PipePeerError::NotAuthenticated)
        );
    }

    #[test]
    fn all_identity_fields_and_mac_must_match() {
        let variants = [
            (
                HelperProcessIdentity::new(42, 9001, "build-a"),
                ChallengeMacVerification::Verified,
            ),
            (
                HelperProcessIdentity::new(41, 9002, "build-a"),
                ChallengeMacVerification::Verified,
            ),
            (
                HelperProcessIdentity::new(41, 9001, "build-b"),
                ChallengeMacVerification::Verified,
            ),
            (identity(), ChallengeMacVerification::Rejected),
        ];

        for (actual, mac) in variants {
            let mut verifier = PipePeerVerifier::new(identity());
            assert!(verifier.authenticate(actual, mac).is_err());
            assert_eq!(
                verifier.verified_peer(),
                Err(PipePeerError::NotAuthenticated)
            );
        }
    }

    #[test]
    fn verified_peer_can_receive_session_material() {
        let mut verifier = PipePeerVerifier::new(identity());
        verifier
            .authenticate(identity(), ChallengeMacVerification::Verified)
            .unwrap();
        let material = SessionMaterial::new(b"session-ticket".to_vec());
        assert_eq!(
            verifier.session_material(&material).unwrap(),
            b"session-ticket"
        );
        assert!(!format!("{material:?}").contains("session-ticket"));
    }
}