From ae107002f1d526df084b078cd7f66876eb9aefb0 Mon Sep 17 00:00:00 2001
From: Patrick Beza <pab@matterlabs.dev>
Date: Thu, 5 Sep 2024 11:34:12 +0200
Subject: [PATCH] Fixes

---
 core/lib/basic_types/src/tee_types.rs        | 34 +++-----------------
 core/lib/dal/src/tee_proof_generation_dal.rs | 11 ++++---
 2 files changed, 11 insertions(+), 34 deletions(-)

diff --git a/core/lib/basic_types/src/tee_types.rs b/core/lib/basic_types/src/tee_types.rs
index 4665c07ac0a..9fd93a164ff 100644
--- a/core/lib/basic_types/src/tee_types.rs
+++ b/core/lib/basic_types/src/tee_types.rs
@@ -1,5 +1,3 @@
-use std::{fmt, str::FromStr};
-
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
@@ -9,23 +7,6 @@ pub enum TeeType {
     Sgx,
 }
 
-impl fmt::Display for TeeType {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self {
-            TeeType::Sgx => write!(f, "sgx"),
-        }
-    }
-}
-
-impl FromStr for TeeType {
-    type Err = String;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        serde_json::from_str(&format!("\"{}\"", s.to_lowercase()))
-            .map_err(|_| format!("Invalid value for TeeType: {}", s))
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use serde_json;
@@ -33,7 +14,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_deserialize_teetype() {
+    fn test_serialize_teetype() {
         let json_str = "\"sgx\"";
         let tee_type: TeeType = serde_json::from_str(json_str).unwrap();
         assert_eq!(tee_type, TeeType::Sgx);
@@ -45,14 +26,9 @@ mod tests {
     }
 
     #[test]
-    fn test_enumstring_teetype() {
-        assert_eq!(TeeType::from_str("sgx").unwrap(), TeeType::Sgx);
-        assert_eq!(TeeType::from_str("Sgx").unwrap(), TeeType::Sgx);
-        assert_eq!(TeeType::from_str("SGX").unwrap(), TeeType::Sgx);
-    }
-
-    #[test]
-    fn test_display_teetype() {
-        assert_eq!(TeeType::Sgx.to_string(), "sgx");
+    fn test_deserialize_teetype() {
+        let tee_type = TeeType::Sgx;
+        let json_str = serde_json::to_string(&tee_type).unwrap();
+        assert_eq!(json_str, "\"sgx\"");
     }
 }
diff --git a/core/lib/dal/src/tee_proof_generation_dal.rs b/core/lib/dal/src/tee_proof_generation_dal.rs
index cc6b87a07ac..7faa8f63fb6 100644
--- a/core/lib/dal/src/tee_proof_generation_dal.rs
+++ b/core/lib/dal/src/tee_proof_generation_dal.rs
@@ -75,7 +75,7 @@ impl TeeProofGenerationDal<'_, '_> {
                 tee_proof_generation_details.l1_batch_number
             "#,
             TeeProofGenerationJobStatus::PickedByProver.to_string(),
-            tee_type.to_string(),
+            serde_json::to_string(&tee_type).unwrap(),
             TeeVerifierInputProducerJobStatus::Successful as TeeVerifierInputProducerJobStatus,
             TeeProofGenerationJobStatus::Unpicked.to_string(),
             processing_timeout,
@@ -112,7 +112,7 @@ impl TeeProofGenerationDal<'_, '_> {
             "#,
             TeeProofGenerationJobStatus::Unpicked.to_string(),
             batch_number,
-            tee_type.to_string()
+            serde_json::to_string(&tee_type).unwrap()
         )
         .instrument("unlock_batch")
         .with_arg("l1_batch_number", &batch_number)
@@ -145,7 +145,7 @@ impl TeeProofGenerationDal<'_, '_> {
             WHERE
                 l1_batch_number = $6
             "#,
-            tee_type.to_string(),
+            serde_json::to_string(&tee_type).unwrap(),
             TeeProofGenerationJobStatus::Generated.to_string(),
             pubkey,
             signature,
@@ -189,7 +189,7 @@ impl TeeProofGenerationDal<'_, '_> {
             ON CONFLICT (l1_batch_number, tee_type) DO NOTHING
             "#,
             batch_number,
-            tee_type.to_string(),
+            serde_json::to_string(&tee_type).unwrap(),
             TeeProofGenerationJobStatus::Unpicked.to_string(),
         );
         let instrumentation = Instrumented::new("insert_tee_proof_generation_job")
@@ -259,7 +259,8 @@ impl TeeProofGenerationDal<'_, '_> {
             .bind(TeeProofGenerationJobStatus::Generated.to_string());
 
         if let Some(tee_type) = tee_type {
-            query = query.bind(tee_type.to_string());
+            let tee_type = serde_json::to_string(&tee_type).unwrap();
+            query = query.bind(tee_type);
         }
 
         let proofs: Vec<StorageTeeProof> = query.fetch_all(self.storage.conn()).await.unwrap();