Skip to main content

aivpn_server/
key_rotation.rs

1//! Automatic Key Rotation (Phase 3)
2//!
3//! Implements automatic session key rotation for enhanced security
4//!
5//! Features:
6//! - Time-based rotation (every 2 minutes)
7//! - Data-based rotation (every 1 MB)
8//! - In-band CONTROL message signaling
9//! - Zero-downtime key transition
10
11use std::time::{Duration, Instant};
12use tracing::{debug, info};
13
14use aivpn_common::crypto::{KeyPair, X25519_PUBLIC_KEY_SIZE};
15use aivpn_common::error::Result;
16use aivpn_common::protocol::ControlPayload;
17
18/// Key rotation configuration
19#[derive(Debug, Clone)]
20pub struct KeyRotationConfig {
21    /// Rotate keys every N seconds
22    pub time_interval_secs: u64,
23
24    /// Rotate keys every N bytes
25    pub data_interval_bytes: u64,
26
27    /// Enable automatic rotation
28    pub enable_auto_rotation: bool,
29}
30
31impl Default for KeyRotationConfig {
32    fn default() -> Self {
33        Self {
34            time_interval_secs: 120,        // 2 minutes
35            data_interval_bytes: 1_000_000, // 1 MB
36            enable_auto_rotation: true,
37        }
38    }
39}
40
41/// Key rotation state
42#[derive(Debug)]
43pub struct KeyRotator {
44    config: KeyRotationConfig,
45
46    /// Current keypair
47    current_keypair: KeyPair,
48
49    /// Next keypair (pre-generated for smooth transition)
50    next_keypair: Option<KeyPair>,
51
52    /// Last rotation time
53    last_rotation: Instant,
54
55    /// Bytes transferred since last rotation
56    bytes_since_rotation: u64,
57
58    /// Rotation counter
59    rotation_count: u64,
60}
61
62impl KeyRotator {
63    /// Create new key rotator
64    pub fn new(config: KeyRotationConfig) -> Result<Self> {
65        let current_keypair = KeyPair::generate();
66
67        Ok(Self {
68            config,
69            current_keypair,
70            next_keypair: None,
71            last_rotation: Instant::now(),
72            bytes_since_rotation: 0,
73            rotation_count: 0,
74        })
75    }
76
77    /// Check if rotation is needed
78    pub fn needs_rotation(&self) -> bool {
79        if !self.config.enable_auto_rotation {
80            return false;
81        }
82
83        // Time-based rotation
84        let time_expired =
85            self.last_rotation.elapsed() >= Duration::from_secs(self.config.time_interval_secs);
86
87        // Data-based rotation
88        let data_expired = self.bytes_since_rotation >= self.config.data_interval_bytes;
89
90        time_expired || data_expired
91    }
92
93    /// Perform key rotation
94    pub fn rotate_keys(&mut self) -> Result<KeyRotationEvent> {
95        info!(
96            "Rotating session keys (rotation #{}).",
97            self.rotation_count + 1
98        );
99
100        // Generate new keypair
101        let new_keypair = KeyPair::generate();
102
103        // Prepare rotation event
104        let event = KeyRotationEvent {
105            new_eph_pub: new_keypair.public_key_bytes(),
106            rotation_count: self.rotation_count,
107        };
108
109        // Update state
110        self.next_keypair = Some(new_keypair);
111        self.last_rotation = Instant::now();
112        self.bytes_since_rotation = 0;
113        self.rotation_count += 1;
114
115        Ok(event)
116    }
117
118    /// Commit rotation (after ACK received)
119    pub fn commit_rotation(&mut self) {
120        if let Some(next) = self.next_keypair.take() {
121            self.current_keypair = next;
122            debug!("Key rotation committed successfully");
123        }
124    }
125
126    /// Record bytes transferred
127    pub fn record_bytes(&mut self, bytes: u64) {
128        self.bytes_since_rotation += bytes;
129    }
130
131    /// Get current public key
132    pub fn current_public_key(&self) -> [u8; X25519_PUBLIC_KEY_SIZE] {
133        self.current_keypair.public_key_bytes()
134    }
135
136    /// Get next public key (if pre-generated)
137    pub fn next_public_key(&self) -> Option<[u8; X25519_PUBLIC_KEY_SIZE]> {
138        self.next_keypair.as_ref().map(|k| k.public_key_bytes())
139    }
140
141    /// Generate a CONTROL message for key rotation.
142    ///
143    /// Returns `None` when no next key has been pre-generated yet — callers must
144    /// call `rotate_keys()` before this to avoid sending the current key as "new",
145    /// which would silently produce a no-op rotation on the peer.
146    pub fn create_rotation_message(&self) -> Option<ControlPayload> {
147        self.next_public_key()
148            .map(|new_eph_pub| ControlPayload::KeyRotate { new_eph_pub })
149    }
150
151    /// Get rotation statistics
152    pub fn stats(&self) -> KeyRotationStats {
153        KeyRotationStats {
154            rotation_count: self.rotation_count,
155            bytes_since_rotation: self.bytes_since_rotation,
156            time_since_rotation: self.last_rotation.elapsed(),
157            next_rotation_in: Duration::from_secs(self.config.time_interval_secs)
158                .saturating_sub(self.last_rotation.elapsed()),
159        }
160    }
161}
162
163/// Key rotation event
164#[derive(Debug, Clone)]
165pub struct KeyRotationEvent {
166    pub new_eph_pub: [u8; X25519_PUBLIC_KEY_SIZE],
167    pub rotation_count: u64,
168}
169
170/// Key rotation statistics
171#[derive(Debug, Clone)]
172pub struct KeyRotationStats {
173    pub rotation_count: u64,
174    pub bytes_since_rotation: u64,
175    pub time_since_rotation: Duration,
176    pub next_rotation_in: Duration,
177}