From e64e563e113984fc7d559bb314031e89f2a4bc50 Mon Sep 17 00:00:00 2001
From: Artemis Tosini <me@artem.ist>
Date: Tue, 10 Sep 2024 02:39:23 +0000
Subject: [PATCH] lib: Add utility functions for checking response

---
 lib/src/error.rs     |  5 +++--
 lib/src/message.rs   | 33 +++++++++++++++++++++++++++++++++
 lib/src/selection.rs | 11 +++++++----
 struct/src/lib.rs    |  8 ++++----
 4 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/lib/src/error.rs b/lib/src/error.rs
index 52bd954..9a0a310 100644
--- a/lib/src/error.rs
+++ b/lib/src/error.rs
@@ -1,10 +1,11 @@
 use snafu::prelude::*;
+use vapore_struct::eresult::EResult;
 
 #[derive(Debug, Snafu)]
 #[snafu(visibility(pub(crate)))]
 pub enum ClientError {
-    #[snafu(display("Valve returned bad result {eresult} with message `{message}`"))]
-    EResult { eresult: u32, message: String },
+    #[snafu(display("Valve returned bad result {eresult:?} with message `{message}`"))]
+    BadResult { eresult: EResult, message: String },
 
     #[snafu(display("Request failure"))]
     Reqwest { source: reqwest::Error },
diff --git a/lib/src/message.rs b/lib/src/message.rs
index 66f0f3d..4ade665 100644
--- a/lib/src/message.rs
+++ b/lib/src/message.rs
@@ -7,6 +7,7 @@ use vapore_proto::{
     enums_clientserver::EMsg,
     steammessages_base::{CMsgMulti, CMsgProtoBufHeader},
 };
+use vapore_struct::eresult::EResult;
 
 use crate::error::{
     ClientError, DecompressionInvalidSnafu, DecompressionSnafu, InvalidActionSnafu,
@@ -53,6 +54,22 @@ impl<T: protobuf::Message> CMProtoBufMessage<T> {
             body,
         })
     }
+
+    pub fn eresult(&self) -> Result<EResult, ClientError> {
+        Ok((self.header.eresult() as u32).try_into()?)
+    }
+
+    pub fn ok(&self) -> Result<(), ClientError> {
+        let eresult = self.eresult()?;
+        if eresult == EResult::OK {
+            Ok(())
+        } else {
+            Err(ClientError::BadResult {
+                eresult,
+                message: self.header.error_message().to_string(),
+            })
+        }
+    }
 }
 
 /// A message sent over the socket, but the body is still serialized
@@ -151,4 +168,20 @@ impl CMRawProtoBufMessage {
 
         Ok(items)
     }
+
+    pub fn eresult(&self) -> Result<EResult, ClientError> {
+        Ok((self.header.eresult() as u32).try_into()?)
+    }
+
+    pub fn ok(&self) -> Result<(), ClientError> {
+        let eresult = self.eresult()?;
+        if eresult == EResult::OK {
+            Ok(())
+        } else {
+            Err(ClientError::BadResult {
+                eresult,
+                message: self.header.error_message().to_string(),
+            })
+        }
+    }
 }
diff --git a/lib/src/selection.rs b/lib/src/selection.rs
index 7a900c7..6831a14 100644
--- a/lib/src/selection.rs
+++ b/lib/src/selection.rs
@@ -1,10 +1,12 @@
 use std::collections::BTreeMap;
 
+use num_enum::TryFromPrimitive;
 use serde::Deserialize;
 use snafu::prelude::*;
+use vapore_struct::eresult::EResult;
 
 use crate::{
-    error::{EResultSnafu, ReqwestSnafu, VdfSnafu},
+    error::{BadResultSnafu, ReqwestSnafu, VdfSnafu},
     ClientError,
 };
 
@@ -39,10 +41,11 @@ pub async fn bootstrap_find_servers() -> Result<Vec<String>, ClientError> {
     let result: GetCMListForConnectResponse =
         keyvalues_serde::from_str(&response).context(VdfSnafu {})?;
 
+    let eresult = EResult::try_from_primitive(result.success)?;
     ensure!(
-        result.success == 1,
-        EResultSnafu {
-            eresult: result.success,
+        eresult == EResult::OK,
+        BadResultSnafu {
+            eresult,
             message: result.message
         }
     );
diff --git a/struct/src/lib.rs b/struct/src/lib.rs
index 3970fa1..517c7c0 100644
--- a/struct/src/lib.rs
+++ b/struct/src/lib.rs
@@ -4,8 +4,8 @@
 pub mod eresult {
     use num_enum::{IntoPrimitive, TryFromPrimitive};
     #[repr(u32)]
-    #[derive(Debug, Copy, Clone, TryFromPrimitive, IntoPrimitive)]
-    enum EResult {
+    #[derive(Debug, Copy, Clone, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)]
+    pub enum EResult {
         Invalid = 0,
 
         OK = 1,
@@ -166,7 +166,7 @@ pub mod enums {
         }
     }
 
-    #[derive(Debug, Copy, Clone, TryFromPrimitive, IntoPrimitive)]
+    #[derive(Debug, Copy, Clone, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)]
     #[repr(u32)]
     pub enum ELicenseType {
         NoLicense = 0,
@@ -180,7 +180,7 @@ pub mod enums {
     }
 
     #[repr(u32)]
-    #[derive(Debug, Copy, Clone, TryFromPrimitive, IntoPrimitive)]
+    #[derive(Debug, Copy, Clone, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)]
     pub enum EPaymentMethod {
         None = 0,
         ActivationCode = 1,