@@ -7,8 +7,6 @@ use serde::Deserialize;
77use std:: fmt:: Display ;
88use x25519_dalek:: { PublicKey , StaticSecret } ;
99
10- use rand_core:: OsRng as X25519OsRng ;
11-
1210const B64_ENGINE : GeneralPurpose = general_purpose:: STANDARD ;
1311
1412#[ derive( Deserialize , Clone ) ]
@@ -44,7 +42,7 @@ impl Display for WgPeer {
4442
4543pub fn generate_keypair ( ) -> anyhow:: Result < WgKey > {
4644 // Generate new keypair
47- let private = StaticSecret :: random_from_rng ( X25519OsRng ) ;
45+ let private = StaticSecret :: random_from_rng ( & mut rand :: rng ( ) ) ;
4846 let public = PublicKey :: from ( & private) ;
4947 let public_key = B64_ENGINE . encode ( public. as_bytes ( ) ) ;
5048 let private_key = B64_ENGINE . encode ( private. to_bytes ( ) ) ;
@@ -58,6 +56,10 @@ pub fn generate_keypair() -> anyhow::Result<WgKey> {
5856
5957pub fn generate_public_key ( private_key : & str ) -> anyhow:: Result < String > {
6058 let private_bytes = B64_ENGINE . decode ( private_key) ?;
59+ if private_bytes. len ( ) != 32 {
60+ anyhow:: bail!( "Private key must be exactly 32 bytes when decoded" ) ;
61+ }
62+
6163 let mut byte_array = [ 0 ; 32 ] ;
6264 byte_array. copy_from_slice ( & private_bytes) ;
6365
@@ -66,3 +68,75 @@ pub fn generate_public_key(private_key: &str) -> anyhow::Result<String> {
6668 let public_key = B64_ENGINE . encode ( public. as_bytes ( ) ) ;
6769 Ok ( public_key)
6870}
71+
72+ #[ cfg( test) ]
73+ mod tests {
74+ use super :: * ;
75+
76+ #[ test]
77+ fn test_generate_keypair ( ) {
78+ let keypair = generate_keypair ( ) . unwrap ( ) ;
79+
80+ // Check that both keys are valid base64 strings
81+ assert ! ( !keypair. public. is_empty( ) ) ;
82+ assert ! ( !keypair. private. is_empty( ) ) ;
83+
84+ // Check that public key is 44 characters (32 bytes base64 encoded)
85+ assert_eq ! ( keypair. public. len( ) , 44 ) ;
86+
87+ // Check that private key is 44 characters (32 bytes base64 encoded)
88+ assert_eq ! ( keypair. private. len( ) , 44 ) ;
89+
90+ // Verify that the public key can be derived from the private key
91+ let derived_public = generate_public_key ( & keypair. private ) . unwrap ( ) ;
92+ assert_eq ! ( keypair. public, derived_public) ;
93+ }
94+
95+ #[ test]
96+ fn test_generate_public_key ( ) {
97+ // Test with a known private key - verify format and consistency
98+ let private_key = "gI6EdkZ4UQR6N5Q1LpI+JWCb1yZCSBHNzQe7J/KoX0s=" ;
99+ let public_key = generate_public_key ( private_key) . unwrap ( ) ;
100+
101+ // Just verify it's a valid public key format rather than hardcoding
102+ assert_eq ! ( public_key. len( ) , 44 ) ;
103+ assert ! ( B64_ENGINE . decode( & public_key) . is_ok( ) ) ;
104+ }
105+
106+ #[ test]
107+ fn test_generate_public_key_invalid_input ( ) {
108+ // Test with invalid base64
109+ let result = generate_public_key ( "invalid_base64!" ) ;
110+ assert ! ( result. is_err( ) ) ;
111+
112+ // Test with empty string
113+ let result = generate_public_key ( "" ) ;
114+ assert ! ( result. is_err( ) ) ;
115+
116+ // Test with too short base64
117+ let result = generate_public_key ( "YQ==" ) ; // "a" encoded - only 1 byte
118+ assert ! ( result. is_err( ) ) ;
119+
120+ // Test with base64 that's not 32 bytes when decoded
121+ let result = generate_public_key ( "YWFhYWFhYQ==" ) ; // "aaaaaa" - 6 bytes
122+ assert ! ( result. is_err( ) ) ;
123+ }
124+
125+ #[ test]
126+ fn test_keypair_consistency ( ) {
127+ // Generate multiple keypairs and ensure they're different
128+ let keypair1 = generate_keypair ( ) . unwrap ( ) ;
129+ let keypair2 = generate_keypair ( ) . unwrap ( ) ;
130+
131+ // Should be different due to randomness
132+ assert_ne ! ( keypair1. public, keypair2. public) ;
133+ assert_ne ! ( keypair1. private, keypair2. private) ;
134+
135+ // But each keypair should be internally consistent
136+ let derived_public1 = generate_public_key ( & keypair1. private ) . unwrap ( ) ;
137+ assert_eq ! ( keypair1. public, derived_public1) ;
138+
139+ let derived_public2 = generate_public_key ( & keypair2. private ) . unwrap ( ) ;
140+ assert_eq ! ( keypair2. public, derived_public2) ;
141+ }
142+ }
0 commit comments