From a87a062e8293f11966dcc78ffbe86bf830eb605c Mon Sep 17 00:00:00 2001
From: Dinculescu <mihai.dinculescu@outlook.com>
Date: Wed, 20 Dec 2023 13:51:43 +0000
Subject: [PATCH] Properly handle several Klap Protocol edge cases

Fixes #140.
---
 CHANGELOG.md                           |  8 +++++
 tapo/src/api/protocol/klap_protocol.rs | 46 ++++++++++++++++----------
 tapo/src/error.rs                      |  2 ++
 3 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cdd10ca..943a12d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,12 +10,20 @@ file. This change log follows the conventions of
 
 - The `anyhow::anyhow!("Local hash does not match server hash")` error has been replaced with the more specific `tapo::TapoResponseError::InvalidCredentials` error.
 
+### Fixed
+
+- A handful of edge cases around the Klap Protocol that were causing panics have been fixed to return `tapo::TapoResponseError::SessionTimeout` or `tapo::TapoResponseError::InvalidResponse` errors instead.
+
 ## [Python Unreleased][Unreleased]
 
 ### Changed
 
 - The `anyhow::anyhow!("Local hash does not match server hash")` error has been replaced with the more specific `tapo::TapoResponseError::InvalidCredentials` error.
 
+### Fixed
+
+- A handful of edge cases around the Klap Protocol that were causing panics have been fixed to return `tapo::TapoResponseError::SessionTimeout` or `tapo::TapoResponseError::InvalidResponse` errors instead.
+
 ## [Rust v0.7.6][v0.7.6] - 2023-11-25
 
 ### Added
diff --git a/tapo/src/api/protocol/klap_protocol.rs b/tapo/src/api/protocol/klap_protocol.rs
index f6b339b..d8771aa 100644
--- a/tapo/src/api/protocol/klap_protocol.rs
+++ b/tapo/src/api/protocol/klap_protocol.rs
@@ -62,15 +62,24 @@ impl TapoProtocolExt for KlapProtocol {
             .body(payload)
             .map_err(isahc::Error::from)?;
 
-        let response = self
-            .client
-            .send_async(request)
-            .await?
-            .bytes()
-            .await
-            .map_err(anyhow::Error::from)?;
-
-        let response_decrypted = cipher.decrypt(seq, response)?;
+        let mut response = self.client.send_async(request).await?;
+
+        if !response.status().is_success() {
+            warn!("Response error: {}", response.status());
+
+            let error = match response.status() {
+                isahc::http::StatusCode::UNAUTHORIZED | isahc::http::StatusCode::FORBIDDEN => {
+                    TapoResponseError::SessionTimeout
+                }
+                _ => TapoResponseError::InvalidResponse,
+            };
+
+            return Err(Error::Tapo(error));
+        }
+
+        let response_body = response.bytes().await.map_err(anyhow::Error::from)?;
+
+        let response_decrypted = cipher.decrypt(seq, response_body)?;
         debug!("Device responded with: {response_decrypted:?}");
 
         let inner_response: TapoResponse<R> = serde_json::from_str(&response_decrypted)?;
@@ -144,16 +153,16 @@ impl KlapProtocol {
             .body(local_seed)
             .map_err(isahc::Error::from)?;
 
-        let response = self
-            .client
-            .send_async(request)
-            .await?
-            .bytes()
-            .await
-            .map_err(anyhow::Error::from)?;
+        let mut response = self.client.send_async(request).await?;
+
+        if !response.status().is_success() {
+            warn!("Handshake1 error: {}", response.status());
+            return Err(Error::Tapo(TapoResponseError::InvalidResponse));
+        }
 
-        let (remote_seed, server_hash) = response.split_at(16);
+        let response_body = response.bytes().await.map_err(anyhow::Error::from)?;
 
+        let (remote_seed, server_hash) = response_body.split_at(16);
         let local_hash = KlapCipher::sha256(&[local_seed, remote_seed, auth_hash].concat());
 
         if local_hash != server_hash {
@@ -186,7 +195,8 @@ impl KlapProtocol {
         let response = self.client.send_async(request).await?;
 
         if !response.status().is_success() {
-            return Err(anyhow::anyhow!("").into());
+            warn!("Handshake2 error: {}", response.status());
+            return Err(Error::Tapo(TapoResponseError::InvalidResponse));
         }
 
         debug!("Handshake2 OK");
diff --git a/tapo/src/error.rs b/tapo/src/error.rs
index e7d979e..d16ba31 100644
--- a/tapo/src/error.rs
+++ b/tapo/src/error.rs
@@ -4,6 +4,8 @@
 pub enum TapoResponseError {
     /// Invalid request.
     InvalidRequest,
+    /// Invalid response.
+    InvalidResponse,
     /// Malformed request.
     MalformedRequest,
     /// Invalid public key.