add set_psk() test and update README

README.md

@@ -41,7 +41,7 @@ Snow is currently based off of Noise revision 32.
 
 - [x] Rekey()
 - [x] `pskN` modifier
-- [ ] specifying PSKs after building `Session`
+- [x] specifying PSKs after building `Session`
 - [ ] `fallback` modifier
 
 ## Crypto

tests/general.rs

@@ -557,3 +557,39 @@ fn test_get_remote_static() {
     assert_eq!(h_i.get_remote_static().unwrap(), static_r.pubkey());
     assert_eq!(h_r.get_remote_static().unwrap(), static_i.pubkey());
 }
+
+#[test]
+fn test_set_psk() {
+    let params: NoiseParams = "Noise_XXpsk3_25519_AESGCM_SHA256".parse().unwrap();
+    let mut static_i: Dh25519 = Default::default();
+    let mut static_r: Dh25519 = Default::default();
+    static_i.set(&get_inc_key(0));
+    static_r.set(&get_inc_key(1));
+    let mut h_i = NoiseBuilder::new(params.clone())
+        .local_private_key(static_i.privkey())
+        .build_initiator().unwrap();
+    let mut h_r = NoiseBuilder::new(params)
+        .local_private_key(static_r.privkey())
+        .build_responder().unwrap();
+
+    let mut buf  = [0u8; 1024];
+    let mut buf2 = [0u8; 1024];
+
+    let psk = get_inc_key(3);
+
+    // XX(s, rs):
+    // -> e
+    let len = h_i.write_message(&[], &mut buf).unwrap();
+    let _   = h_r.read_message(&buf[..len], &mut buf2).unwrap();
+
+    // <- e, ee s, es
+    let len = h_r.write_message(&[], &mut buf).unwrap();
+    let _   = h_i.read_message(&buf[..len], &mut buf2).unwrap();
+
+    h_i.set_psk(3, &psk).unwrap();
+    h_r.set_psk(3, &psk).unwrap();
+
+    // -> s, se, psk
+    let len = h_i.write_message(&[], &mut buf).unwrap();
+    let _   = h_r.read_message(&buf[..len], &mut buf2).unwrap();
+}