[DTLS] add TestPSK and test_psk_hint_fail

rainliu · rainliu · commit 7b490d18344e · 2020-12-05T23:25:15.000-08:00
diff --git a/dtls/src/config.rs b/dtls/src/config.rs
@@ -45,7 +45,7 @@ pub struct Config {
     // psk sets the pre-shared key used by this DTLS connection
     // If psk is non-nil only psk cipher_suites will be used
     pub(crate) psk: Option<PSKCallback>,
-    pub(crate) psk_identity_hint: Vec<u8>,
+    pub(crate) psk_identity_hint: Option<Vec<u8>>,
 
     // insecure_skip_verify controls whether a client verifies the
     // server's certificate chain and host name.
@@ -151,7 +151,7 @@ pub(crate) fn validate_config(config: &Config) -> Result<(), Error> {
         return Err(ERR_PSK_AND_CERTIFICATE.clone());
     }
 
-    if !config.psk_identity_hint.is_empty() && config.psk.is_none() {
+    if config.psk_identity_hint.is_some() && config.psk.is_none() {
         return Err(ERR_IDENTITY_NO_PSK.clone());
     }
 
diff --git a/dtls/src/conn.rs b/dtls/src/conn.rs
@@ -172,7 +172,7 @@ impl Conn {
 
         let cfg = HandshakeConfig {
             local_psk_callback: config.psk.take(),
-            local_psk_identity_hint: config.psk_identity_hint.clone(),
+            local_psk_identity_hint: config.psk_identity_hint.take(),
             local_cipher_suites,
             local_signature_schemes,
             extended_master_secret: config.extended_master_secret,
@@ -326,6 +326,9 @@ impl Conn {
                                 srv_cli_str(is_client),
                                 err
                             );
+                            if err == *ERR_ALERT_FATAL_OR_CLOSE {
+                                break;
+                            }
                         }
                     }
                 }
@@ -950,13 +953,13 @@ impl Conn {
                 let mut reader = BufReader::new(out.as_slice());
                 let raw_handshake = match Handshake::unmarshal(&mut reader) {
                     Ok(rh) => {
-                        trace!(
+                        /*trace!(
                             "Recv [handshake:{}] -> {} (epoch: {}, seq: {})",
                             srv_cli_str(ctx.is_client),
                             rh.handshake_header.handshake_type.to_string(),
                             h.epoch,
                             rh.handshake_header.message_sequence
-                        );
+                        );*/
                         rh
                     }
                     Err(err) => {
diff --git a/dtls/src/conn/conn_test.rs b/dtls/src/conn/conn_test.rs
@@ -10,8 +10,18 @@ use crate::errors::*;
 use tokio::net::UdpSocket;
 
 use std::time::SystemTime;
+
 //use std::io::Write;
 
+lazy_static! {
+    pub static ref ERR_TEST_PSK_INVALID_IDENTITY: Error =
+        Error::new("TestPSK: Server got invalid identity".to_owned());
+    pub static ref ERR_PSK_REJECTED: Error = Error::new("PSK Rejected".to_owned());
+    pub static ref ERR_NOT_EXPECTED_CHAIN: Error = Error::new("not expected chain".to_owned());
+    pub static ref ERR_EXPECTED_CHAIN: Error = Error::new("expected chain".to_owned());
+    pub static ref ERR_WRONG_CERT: Error = Error::new("wrong cert".to_owned());
+}
+
 async fn build_pipe() -> Result<(Conn, Conn), Error> {
     let (ua, ub) = pipe().await?;
 
@@ -39,10 +49,13 @@ async fn pipe_conn(ca: UdpSocket, cb: UdpSocket) -> Result<(Conn, Conn), Error>
             ca,
             Config {
                 srtp_protection_profiles: vec![SRTPProtectionProfile::SRTP_AES128_CM_HMAC_SHA1_80],
+                //TODO: change PSK to cert
                 cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256],
+                psk: Some(psk_callback_client),
+                psk_identity_hint: Some("WebRTC.rs DTLS Server".as_bytes().to_vec()),
                 ..Default::default()
             },
-            false,
+            false, //TODO: use ceritificate
         )
         .await;
 
@@ -54,10 +67,13 @@ async fn pipe_conn(ca: UdpSocket, cb: UdpSocket) -> Result<(Conn, Conn), Error>
         cb,
         Config {
             srtp_protection_profiles: vec![SRTPProtectionProfile::SRTP_AES128_CM_HMAC_SHA1_80],
+            //TODO: change PSK to cert
             cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256],
+            psk: Some(psk_callback_server),
+            psk_identity_hint: Some("WebRTC.rs DTLS Client".as_bytes().to_vec()),
             ..Default::default()
         },
-        false,
+        false, //TODO: use ceritificate
     )
     .await?;
 
@@ -86,16 +102,17 @@ fn psk_callback_server(hint: &[u8]) -> Result<Vec<u8>, Error> {
     Ok(vec![0xAB, 0xC1, 0x23])
 }
 
+fn psk_callback_hint_fail(_hint: &[u8]) -> Result<Vec<u8>, Error> {
+    Err(ERR_PSK_REJECTED.clone())
+}
+
 async fn create_test_client(
     ca: UdpSocket,
     mut cfg: Config,
     generate_certificate: bool,
 ) -> Result<Conn, Error> {
     if generate_certificate {
         //TODO:
-    } else {
-        cfg.psk = Some(psk_callback_client);
-        cfg.psk_identity_hint = "WebRTC.rs DTLS Server".as_bytes().to_vec();
     }
 
     cfg.insecure_skip_verify = true;
@@ -104,15 +121,13 @@ async fn create_test_client(
 
 async fn create_test_server(
     cb: UdpSocket,
-    mut cfg: Config,
+    cfg: Config,
     generate_certificate: bool,
 ) -> Result<Conn, Error> {
     if generate_certificate {
         //TODO:
-    } else {
-        cfg.psk = Some(psk_callback_server);
-        cfg.psk_identity_hint = "WebRTC.rs DTLS Client".as_bytes().to_vec();
     }
+
     Conn::new(cb, cfg, false, None).await
 }
 
@@ -340,11 +355,11 @@ async fn test_handshake_with_alert() -> Result<(), Error> {
 
         let (ca, cb) = pipe().await?;
         tokio::spawn(async move {
-            let result = create_test_client(ca, config_client, false).await;
+            let result = create_test_client(ca, config_client, false).await; //TODO: use certificate
             let _ = client_err_tx.send(result).await;
         });
 
-        let result_server = create_test_server(cb, config_server, false).await;
+        let result_server = create_test_server(cb, config_server, false).await; //TODO: use certificate
         if let Err(err) = result_server {
             assert_eq!(
                 err, err_server,
@@ -495,3 +510,141 @@ async fn test_export_keying_material() -> Result<(), Error> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_psk() -> Result<(), Error> {
+    /*env_logger::Builder::new()
+    .format(|buf, record| {
+        writeln!(
+            buf,
+            "{}:{} [{}] {} - {}",
+            record.file().unwrap_or("unknown"),
+            record.line().unwrap_or(0),
+            record.level(),
+            chrono::Local::now().format("%H:%M:%S.%6f"),
+            record.args()
+        )
+    })
+    .filter(None, LevelFilter::Trace)
+    .init();*/
+
+    let tests = vec![
+        (
+            "Server identity specified",
+            Some("Test Identity".as_bytes().to_vec()),
+        ),
+        ("Server identity nil", None),
+    ];
+
+    for (name, server_identity) in tests {
+        let client_identity = "Client Identity".as_bytes();
+        let (client_res_tx, mut client_res_rx) = mpsc::channel(1);
+
+        let (ca, cb) = pipe().await?;
+        tokio::spawn(async move {
+            let conf = Config {
+                psk: Some(psk_callback_client),
+                psk_identity_hint: Some(client_identity.to_vec()),
+                cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256], //TODO: change it to TLS_PSK_WITH_AES_128_CCM_8
+                ..Default::default()
+            };
+
+            let result = create_test_client(ca, conf, false).await;
+            let _ = client_res_tx.send(result).await;
+        });
+
+        let config = Config {
+            psk: Some(psk_callback_server),
+            psk_identity_hint: server_identity,
+            cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256], //TODO: change it to TLS_PSK_WITH_AES_128_CCM_8
+            ..Default::default()
+        };
+
+        let mut server = create_test_server(cb, config, false).await?;
+
+        if let Some(result) = client_res_rx.recv().await {
+            if let Ok(mut client) = result {
+                client.close().await?;
+            } else {
+                assert!(
+                    false,
+                    "{}: Expected create_test_client successfully, but got error",
+                    name,
+                );
+            }
+        }
+
+        server.close().await?;
+    }
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_psk_hint_fail() -> Result<(), Error> {
+    /*env_logger::Builder::new()
+    .format(|buf, record| {
+        writeln!(
+            buf,
+            "{}:{} [{}] {} - {}",
+            record.file().unwrap_or("unknown"),
+            record.line().unwrap_or(0),
+            record.level(),
+            chrono::Local::now().format("%H:%M:%S.%6f"),
+            record.args()
+        )
+    })
+    .filter(None, LevelFilter::Trace)
+    .init();*/
+
+    let (client_res_tx, mut client_res_rx) = mpsc::channel(1);
+
+    let (ca, cb) = pipe().await?;
+    tokio::spawn(async move {
+        let conf = Config {
+            psk: Some(psk_callback_hint_fail),
+            psk_identity_hint: Some(vec![]),
+            cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256], //TODO: change it to TLS_PSK_WITH_AES_128_CCM_8
+            ..Default::default()
+        };
+
+        let result = create_test_client(ca, conf, false).await;
+        let _ = client_res_tx.send(result).await;
+    });
+
+    let config = Config {
+        psk: Some(psk_callback_hint_fail),
+        psk_identity_hint: Some(vec![]),
+        cipher_suites: vec![CipherSuiteID::TLS_PSK_WITH_AES_128_GCM_SHA256], //TODO: change it to TLS_PSK_WITH_AES_128_CCM_8
+        ..Default::default()
+    };
+
+    if let Err(server_err) = create_test_server(cb, config, false).await {
+        assert_eq!(
+            server_err,
+            ERR_ALERT_FATAL_OR_CLOSE.clone(),
+            "TestPSK: Server error exp({}) failed({})",
+            ERR_ALERT_FATAL_OR_CLOSE.clone(),
+            server_err,
+        );
+    } else {
+        assert!(false, "Expected server error, but got OK");
+    }
+
+    let result = client_res_rx.recv().await;
+    if let Some(client) = result {
+        if let Err(client_err) = client {
+            assert_eq!(
+                client_err,
+                ERR_PSK_REJECTED.clone(),
+                "TestPSK: Client error exp({}) failed({})",
+                ERR_PSK_REJECTED.clone(),
+                client_err,
+            );
+        } else {
+            assert!(false, "Expected client error, but got OK");
+        }
+    }
+
+    Ok(())
+}
diff --git a/dtls/src/flight/flight3.rs b/dtls/src/flight/flight3.rs
@@ -91,7 +91,7 @@ impl Flight for Flight3 {
             }
         }
 
-        let result = if cfg.local_psk_callback.is_none() {
+        let result = if cfg.local_psk_callback.is_some() {
             cache
                 .full_pull_map(
                     state.handshake_recv_sequence,
diff --git a/dtls/src/flight/flight4.rs b/dtls/src/flight/flight4.rs
@@ -660,7 +660,7 @@ impl Flight for Flight4 {
                     reset_local_sequence_number: false,
                 });
             }
-        } else if !cfg.local_psk_identity_hint.is_empty() {
+        } else if let Some(local_psk_identity_hint) = &cfg.local_psk_identity_hint {
             // To help the client in selecting which identity to use, the server
             // can provide a "PSK identity hint" in the ServerKeyExchange message.
             // If no hint is provided, the ServerKeyExchange message is omitted.
@@ -672,7 +672,7 @@ impl Flight for Flight4 {
                     0,
                     Content::Handshake(Handshake::new(HandshakeMessage::ServerKeyExchange(
                         HandshakeMessageServerKeyExchange {
-                            identity_hint: cfg.local_psk_identity_hint.clone(),
+                            identity_hint: local_psk_identity_hint.clone(),
                             elliptic_curve_type: EllipticCurveType::Unsupported,
                             named_curve: NamedCurve::Unsupported,
                             public_key: vec![],
diff --git a/dtls/src/flight/flight5.rs b/dtls/src/flight/flight5.rs
@@ -235,8 +235,8 @@ impl Flight for Flight5 {
             if let Some(local_keypair) = &state.local_keypair {
                 client_key_exchange.public_key = local_keypair.public_key.clone();
             }
-        } else {
-            client_key_exchange.identity_hint = cfg.local_psk_identity_hint.clone();
+        } else if let Some(local_psk_identity_hint) = &cfg.local_psk_identity_hint {
+            client_key_exchange.identity_hint = local_psk_identity_hint.clone();
         }
 
         pkts.push(Packet {
diff --git a/dtls/src/handshaker.rs b/dtls/src/handshaker.rs

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ pub struct Config {`
`45`	`45`	`// psk sets the pre-shared key used by this DTLS connection`
`46`	`46`	`// If psk is non-nil only psk cipher_suites will be used`
`47`	`47`	`pub(crate) psk: Option<PSKCallback>,`
`48`		`- pub(crate) psk_identity_hint: Vec<u8>,`
	`48`	`+ pub(crate) psk_identity_hint: Option<Vec<u8>>,`
`49`	`49`
`50`	`50`	`// insecure_skip_verify controls whether a client verifies the`
`51`	`51`	`// server's certificate chain and host name.`
`@@ -151,7 +151,7 @@ pub(crate) fn validate_config(config: &Config) -> Result<(), Error> {`
`151`	`151`	`return Err(ERR_PSK_AND_CERTIFICATE.clone());`
`152`	`152`	`}`
`153`	`153`
`154`		`- if !config.psk_identity_hint.is_empty() && config.psk.is_none() {`
	`154`	`+ if config.psk_identity_hint.is_some() && config.psk.is_none() {`
`155`	`155`	`return Err(ERR_IDENTITY_NO_PSK.clone());`
`156`	`156`	`}`
`157`	`157`
Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ impl Flight for Flight3 {`
`91`	`91`	`}`
`92`	`92`	`}`
`93`	`93`
`94`		`- let result = if cfg.local_psk_callback.is_none() {`
	`94`	`+ let result = if cfg.local_psk_callback.is_some() {`
`95`	`95`	`cache`
`96`	`96`	`.full_pull_map(`
`97`	`97`	`state.handshake_recv_sequence,`
Original file line number	Diff line number	Diff line change
`@@ -235,8 +235,8 @@ impl Flight for Flight5 {`
`235`	`235`	`if let Some(local_keypair) = &state.local_keypair {`
`236`	`236`	`client_key_exchange.public_key = local_keypair.public_key.clone();`
`237`	`237`	`}`
`238`		`- } else {`
`239`		`- client_key_exchange.identity_hint = cfg.local_psk_identity_hint.clone();`
	`238`	`+ } else if let Some(local_psk_identity_hint) = &cfg.local_psk_identity_hint {`
	`239`	`+ client_key_exchange.identity_hint = local_psk_identity_hint.clone();`
`240`	`240`	`}`
`241`	`241`
`242`	`242`	`pkts.push(Packet {`