8 files changed, 124 insertions, 55 deletions

diff --git a/src/message/handshake/clientinit.rs b/src/message/handshake/clientinit.rs
index 9d35243..3011821 100644
--- a/src/message/handshake/clientinit.rs
+++ b/src/message/handshake/clientinit.rs
@@ -63,13 +63,27 @@ impl HandshakeSerialize for ClientInit {
     }
 }
 
-impl From<VariantMap> for ClientInit {
-    fn from(mut input: VariantMap) -> Self {
-        ClientInit {
-            client_version: input.remove("ClientVersion").unwrap().try_into().unwrap(),
-            client_date: input.remove("ClientDate").unwrap().try_into().unwrap(),
-            client_features: input.remove("Features").unwrap().try_into().unwrap(),
-            feature_list: input.remove("FeatureList").unwrap().try_into().unwrap(),
-        }
+impl TryFrom<VariantMap> for ClientInit {
+    type Error = ProtocolError;
+
+    fn try_from(mut input: VariantMap) -> Result<Self, Self::Error> {
+        Ok(ClientInit {
+            client_version: input
+                .remove("ClientVersion")
+                .ok_or_else(|| ProtocolError::MissingField("ClientVersion".to_string()))?
+                .try_into()?,
+            client_date: input
+                .remove("ClientDate")
+                .ok_or_else(|| ProtocolError::MissingField("ClientDate".to_string()))?
+                .try_into()?,
+            client_features: input
+                .remove("Features")
+                .ok_or_else(|| ProtocolError::MissingField("Features".to_string()))?
+                .try_into()?,
+            feature_list: input
+                .remove("FeatureList")
+                .ok_or_else(|| ProtocolError::MissingField("FeatureList".to_string()))?
+                .try_into()?,
+        })
     }
 }
diff --git a/src/message/handshake/clientinitack.rs b/src/message/handshake/clientinitack.rs
index f3f4640..adb7425 100644
--- a/src/message/handshake/clientinitack.rs
+++ b/src/message/handshake/clientinitack.rs
@@ -45,16 +45,30 @@ impl HandshakeSerialize for ClientInitAck {
     }
 }
 
-impl From<VariantMap> for ClientInitAck {
-    fn from(input: VariantMap) -> Self {
-        ClientInitAck {
+impl TryFrom<VariantMap> for ClientInitAck {
+    type Error = ProtocolError;
+
+    fn try_from(input: VariantMap) -> Result<Self, Self::Error> {
+        Ok(ClientInitAck {
             // TODO make this compatible with older clients
             core_features: 0,
-            core_configured: input.get("Configured").unwrap().try_into().unwrap(),
-            storage_backends: input.get("StorageBackends").unwrap().try_into().unwrap(),
+            core_configured: input
+                .get("Configured")
+                .ok_or_else(|| ProtocolError::MissingField("Configured".to_string()))?
+                .try_into()?,
+            storage_backends: input
+                .get("StorageBackends")
+                .ok_or_else(|| ProtocolError::MissingField("StorageBackends".to_string()))?
+                .try_into()?,
             #[cfg(feature = "authenticators")]
-            authenticators: input.get("Authenticators").unwrap().try_into().unwrap(),
-            feature_list: input.get("FeatureList").unwrap().try_into().unwrap(),
-        }
+            authenticators: input
+                .get("Authenticators")
+                .ok_or_else(|| ProtocolError::MissingField("Authenticators".to_string()))?
+                .try_into()?,
+            feature_list: input
+                .get("FeatureList")
+                .ok_or_else(|| ProtocolError::MissingField("FeatureList".to_string()))?
+                .try_into()?,
+        })
     }
 }
diff --git a/src/message/handshake/clientinitreject.rs b/src/message/handshake/clientinitreject.rs
index fca5388..6e4622b 100644
--- a/src/message/handshake/clientinitreject.rs
+++ b/src/message/handshake/clientinitreject.rs
@@ -21,10 +21,15 @@ impl HandshakeSerialize for ClientInitReject {
     }
 }
 
-impl From<VariantMap> for ClientInitReject {
-    fn from(mut input: VariantMap) -> Self {
-        ClientInitReject {
-            error: input.remove("ErrorString").unwrap().try_into().unwrap(),
-        }
+impl TryFrom<VariantMap> for ClientInitReject {
+    type Error = ProtocolError;
+
+    fn try_from(mut input: VariantMap) -> Result<Self, Self::Error> {
+        Ok(ClientInitReject {
+            error: input
+                .remove("ErrorString")
+                .ok_or_else(|| ProtocolError::MissingField("ErrorString".to_string()))?
+                .try_into()?,
+        })
     }
 }
diff --git a/src/message/handshake/clientlogin.rs b/src/message/handshake/clientlogin.rs
index 4be4442..b1d9672 100644
--- a/src/message/handshake/clientlogin.rs
+++ b/src/message/handshake/clientlogin.rs
@@ -20,11 +20,19 @@ impl HandshakeSerialize for ClientLogin {
     }
 }
 
-impl From<VariantMap> for ClientLogin {
-    fn from(mut input: VariantMap) -> Self {
-        ClientLogin {
-            user: input.remove("User").unwrap().try_into().unwrap(),
-            password: input.remove("Password").unwrap().try_into().unwrap(),
-        }
+impl TryFrom<VariantMap> for ClientLogin {
+    type Error = ProtocolError;
+
+    fn try_from(mut input: VariantMap) -> Result<Self, Self::Error> {
+        Ok(ClientLogin {
+            user: input
+                .remove("User")
+                .ok_or_else(|| ProtocolError::MissingField("User".to_string()))?
+                .try_into()?,
+            password: input
+                .remove("Password")
+                .ok_or_else(|| ProtocolError::MissingField("Password".to_string()))?
+                .try_into()?,
+        })
     }
 }
diff --git a/src/message/handshake/clientloginack.rs b/src/message/handshake/clientloginack.rs
index de26274..949ce55 100644
--- a/src/message/handshake/clientloginack.rs
+++ b/src/message/handshake/clientloginack.rs
@@ -22,7 +22,10 @@ impl HandshakeDeserialize for ClientLoginAck {
     fn parse(b: &[u8]) -> Result<(usize, Self), ProtocolError> {
         let (len, mut values): (usize, VariantMap) = HandshakeDeserialize::parse(b)?;
 
-        let msgtype: String = values.remove("MsgType").unwrap().try_into().unwrap();
+        let msgtype: String = values
+            .remove("MsgType")
+            .ok_or_else(|| ProtocolError::MissingField("MsgType".to_string()))?
+            .try_into()?;
 
         if msgtype == "ClientLogin" {
             Ok((len, Self {}))
diff --git a/src/message/handshake/clientloginreject.rs b/src/message/handshake/clientloginreject.rs
index c4f09cc..300bc2e 100644
--- a/src/message/handshake/clientloginreject.rs
+++ b/src/message/handshake/clientloginreject.rs
@@ -21,10 +21,15 @@ impl HandshakeSerialize for ClientLoginReject {
     }
 }
 
-impl From<VariantMap> for ClientLoginReject {
-    fn from(mut input: VariantMap) -> Self {
-        ClientLoginReject {
-            error: input.remove("ErrorString").unwrap().try_into().unwrap(),
-        }
+impl TryFrom<VariantMap> for ClientLoginReject {
+    type Error = ProtocolError;
+
+    fn try_from(mut input: VariantMap) -> Result<Self, Self::Error> {
+        Ok(ClientLoginReject {
+            error: input
+                .remove("ErrorString")
+                .ok_or_else(|| ProtocolError::MissingField("ErrorString".to_string()))?
+                .try_into()?,
+        })
     }
 }
diff --git a/src/message/handshake/mod.rs b/src/message/handshake/mod.rs
index 9bcbe6e..1b02bb4 100644
--- a/src/message/handshake/mod.rs
+++ b/src/message/handshake/mod.rs
@@ -58,15 +58,18 @@ impl HandshakeDeserialize for HandshakeMessage {
     fn parse(b: &[u8]) -> Result<(usize, Self), ProtocolError> {
         let (size, mut res) = VariantMap::parse(b)?;
 
-        let msgtype: String = res.remove("MsgType").unwrap().try_into().unwrap();
+        let msgtype: String = res
+            .remove("MsgType")
+            .ok_or_else(|| ProtocolError::MissingField("MsgType".to_string()))?
+            .try_into()?;
         match msgtype.as_str() {
-            "ClientInit" => Ok((size, HandshakeMessage::ClientInit(res.into()))),
-            "ClientInitAck" => Ok((size, HandshakeMessage::ClientInitAck(res.into()))),
-            "ClientInitReject" => Ok((size, HandshakeMessage::ClientInitReject(res.into()))),
-            "ClientLogin" => Ok((size, HandshakeMessage::ClientLogin(res.into()))),
+            "ClientInit" => Ok((size, HandshakeMessage::ClientInit(res.try_into()?))),
+            "ClientInitAck" => Ok((size, HandshakeMessage::ClientInitAck(res.try_into()?))),
+            "ClientInitReject" => Ok((size, HandshakeMessage::ClientInitReject(res.try_into()?))),
+            "ClientLogin" => Ok((size, HandshakeMessage::ClientLogin(res.try_into()?))),
             "ClientLoginAck" => Ok((size, HandshakeMessage::ClientLoginAck)),
-            "ClientLoginReject" => Ok((size, HandshakeMessage::ClientLoginReject(res.into()))),
-            "SessionInit" => Ok((size, HandshakeMessage::SessionInit(res.into()))),
+            "ClientLoginReject" => Ok((size, HandshakeMessage::ClientLoginReject(res.try_into()?))),
+            "SessionInit" => Ok((size, HandshakeMessage::SessionInit(res.try_into()?))),
             _ => unimplemented!(),
         }
     }
diff --git a/src/message/handshake/sessioninit.rs b/src/message/handshake/sessioninit.rs
index b3a0932..18c026d 100644
--- a/src/message/handshake/sessioninit.rs
+++ b/src/message/handshake/sessioninit.rs
@@ -15,33 +15,50 @@ pub struct SessionInit {
     pub network_ids: Vec<NetworkId>,
 }
 
-impl From<VariantMap> for SessionInit {
-    fn from(input: VariantMap) -> Self {
-        let mut state: VariantMap = input.get("SessionState").unwrap().try_into().unwrap();
+impl TryFrom<VariantMap> for SessionInit {
+    type Error = ProtocolError;
+
+    fn try_from(input: VariantMap) -> Result<Self, Self::Error> {
+        let mut state: VariantMap = input
+            .get("SessionState")
+            .ok_or_else(|| ProtocolError::MissingField("SessionState".to_string()))?
+            .try_into()?;
 
         log::trace!("sessionstate: {:#?}", state);
 
-        let identities: VariantList = state.remove("Identities").unwrap().try_into().unwrap();
-        let buffers: VariantList = state.remove("BufferInfos").unwrap().try_into().unwrap();
-        let network_ids: VariantList = state.remove("NetworkIds").unwrap().try_into().unwrap();
+        let identities: VariantList = state
+            .remove("Identities")
+            .ok_or_else(|| ProtocolError::MissingField("Identities".to_string()))?
+            .try_into()?;
+        let buffers: VariantList = state
+            .remove("BufferInfos")
+            .ok_or_else(|| ProtocolError::MissingField("BufferInfos".to_string()))?
+            .try_into()?;
+        let network_ids: VariantList = state
+            .remove("NetworkIds")
+            .ok_or_else(|| ProtocolError::MissingField("NetworkIds".to_string()))?
+            .try_into()?;
 
-        SessionInit {
-            identities: identities.into_iter().map(|x| x.try_into().unwrap()).collect(),
+        Ok(SessionInit {
+            identities: identities
+                .into_iter()
+                .map(|x| x.try_into())
+                .collect::<Result<Vec<_>, _>>()?,
             buffers: buffers
                 .iter()
                 .map(|buffer| match buffer {
-                    Variant::BufferInfo(buffer) => buffer.clone(),
-                    _ => unimplemented!(),
+                    Variant::BufferInfo(buffer) => Ok(buffer.clone()),
+                    _ => Err(ProtocolError::WrongVariant),
                 })
-                .collect(),
+                .collect::<Result<Vec<_>, _>>()?,
             network_ids: network_ids
                 .iter()
                 .map(|network| match network {
-                    Variant::NetworkId(network) => *network,
-                    _ => unimplemented!(),
+                    Variant::NetworkId(network) => Ok(*network),
+                    _ => Err(ProtocolError::WrongVariant),
                 })
-                .collect(),
-        }
+                .collect::<Result<Vec<_>, _>>()?,
+        })
     }
 }