Add CBOR support for etcd encoding

Up until now we would write CBOR values back as etcd, after this commit
we preserve the original encoding, so CBOR values are written back as
CBOR.

This required tracking the encoding type in the in-memory etcd cache,
and modifying the encode function to handle CBOR.

Author: omertuc

src/cluster_crypto/cert_key_pair.rs

@@ -357,7 +357,7 @@ impl CertKeyPair {
                 .as_bytes()
                 .to_vec(),
             )
-            .await;
+            .await.context("putting in etcd")?;
 
         Ok(())
     }

src/cluster_crypto/distributed_jwt.rs

@@ -86,7 +86,7 @@ impl DistributedJwt {
                 &k8slocation.resource_location.as_etcd_key(),
                 serde_json::to_string(&resource)?.as_bytes().to_vec(),
             )
-            .await;
+            .await.context("putting in etcd")?;
 
         Ok(())
     }

src/cluster_crypto/distributed_private_key.rs

@@ -96,7 +96,7 @@ impl DistributedPrivateKey {
                 .as_bytes()
                 .to_vec(),
             )
-            .await;
+            .await.context("putting in etcd")?;
 
         Ok(())
     }

src/cluster_crypto/distributed_public_key.rs

@@ -129,7 +129,7 @@ impl DistributedPublicKey {
                 .as_bytes()
                 .to_vec(),
             )
-            .await;
+            .await.context("putting in etcd")?;
 
         Ok(())
     }

src/etcd_encoding.rs

@@ -65,17 +65,24 @@ k8s_type!(OAuthClientWithMeta, OAuthClient);
 
 mod k8s_cbor;
 
-pub(crate) async fn decode(data: &[u8]) -> Result<Vec<u8>> {
+#[derive(Clone)]
+pub(crate) enum Encoding {
+    Protobuf,
+    Cbor,
+    Json,
+}
+
+pub(crate) async fn decode(data: &[u8]) -> Result<(Vec<u8>, Encoding)> {
     if !data.starts_with("k8s\x00".as_bytes()) {
         // k8s uses CBOR with the self-describing tag 55799, we can use its bytes to detect CBOR
         if data.starts_with([0xd9, 0xd9, 0xf7].as_ref()) {
             // It's CBOR, just convert to JSON
             let json_value = k8s_cbor::k8s_cbor_bytes_to_json(data).context("converting CBOR to JSON")?;
-            return Ok(serde_json::to_vec(&json_value)?);
+            return Ok((serde_json::to_vec(&json_value)?, Encoding::Cbor));
         }
 
         // Not CBOR, not protobuf, it's probably just raw JSON, return as-is
-        return Ok(data.to_vec());
+        return Ok((data.to_vec(), Encoding::Json));
     }
 
     let data = &data[4..];
@@ -89,7 +96,7 @@ pub(crate) async fn decode(data: &[u8]) -> Result<Vec<u8>> {
         .context("missing kind")?
         .as_str();
 
-    Ok(match kind {
+    let decoded_data = match kind {
         "Route" => serde_json::to_vec(&RouteWithMeta::try_from(unknown)?)?,
         "Deployment" => serde_json::to_vec(&DeploymentWithMeta::try_from(unknown)?)?,
         "ControllerRevision" => serde_json::to_vec(&ControllerRevisionWithMeta::try_from(unknown)?)?,
@@ -105,11 +112,20 @@ pub(crate) async fn decode(data: &[u8]) -> Result<Vec<u8>> {
         "MutatingWebhookConfiguration" => serde_json::to_vec(&MutatingWebhookConfigurationWithMeta::try_from(unknown)?)?,
         "OAuthClient" => serde_json::to_vec(&OAuthClientWithMeta::try_from(unknown)?)?,
         _ => bail!("unknown kind {}", kind),
-    })
+    };
+
+    Ok((decoded_data, Encoding::Protobuf))
 }
 
-pub(crate) async fn encode(data: &[u8]) -> Result<Vec<u8>> {
+pub(crate) async fn encode(data: &[u8], encoding: Encoding) -> Result<Vec<u8>> {
     let value: Value = serde_json::from_slice(data)?;
+
+    if matches!(encoding, Encoding::Cbor) {
+        return k8s_cbor::json_to_k8s_cbor_bytes(value).context("converting JSON to CBOR");
+    }
+
+    // If kind is a known protobuf kind, write it back as protobuf, otherwise return raw JSON
+    // TODO: Just look at the new encoding param?
     let kind = value
         .pointer("/kind")
         .context("missing kind")?

src/etcd_encoding/k8s_cbor.rs

@@ -90,3 +90,15 @@ pub(crate) fn k8s_cbor_bytes_to_json(cbor_bytes: &[u8]) -> Result<JsonValue> {
 
     cbor_to_json(v)
 }
+
+pub(crate) fn json_to_k8s_cbor_bytes(json: JsonValue) -> Result<Vec<u8>> {
+    let cbor = json_to_cbor(json)?;
+
+    // Put back the self-describing CBOR tag that we stripped
+    let tagged_cbor = CborValue::Tag(SELF_DESCRIBING_CBOR_TAG, Box::new(cbor));
+
+    let mut bytes = Vec::new();
+    ciborium::ser::into_writer(&tagged_cbor, &mut bytes)?;
+
+    Ok(bytes)
+}

src/k8s_etcd.rs

@@ -3,6 +3,7 @@ use crate::encrypt::ResourceTransformers;
 use crate::etcd_encoding;
 use anyhow::{bail, ensure, Context, Result};
 use etcd_client::{Client as EtcdClient, GetOptions};
+use etcd_encoding::Encoding;
 use futures_util::future::join_all;
 use serde_json::Value;
 use std::collections::{HashMap, HashSet};
@@ -22,7 +23,7 @@ pub(crate) struct EtcdResult {
 /// have to go through etcd for every single edit.
 pub(crate) struct InMemoryK8sEtcd {
     pub(crate) etcd_client: Option<Arc<EtcdClient>>,
-    etcd_keyvalue_hashmap: Mutex<HashMap<String, Vec<u8>>>,
+    etcd_keyvalue_hashmap: Mutex<HashMap<String, (Encoding, Vec<u8>)>>,
     edited: Mutex<HashMap<String, Vec<u8>>>,
     deleted_keys: Mutex<HashSet<String>>,
     decrypt_resource_transformers: Option<ResourceTransformers>,
@@ -105,10 +106,10 @@ impl InMemoryK8sEtcd {
                 continue;
             }
             let key = key.clone();
-            let value = value.clone();
+            let (encoding, value) = value.clone();
             let etcd_client = Arc::clone(etcd_client);
 
-            let mut value = etcd_encoding::encode(value.as_slice()).await.context("encoding value")?;
+            let mut value = etcd_encoding::encode(value.as_slice(), encoding).await.context("encoding value")?;
 
             if let Some(resource_transformers) = &self.encrypt_resource_transformers {
                 // https://github.com/kubernetes/apiserver/blob/3423727e46efe7dfa40dcdb1a9c5c5027b07303d/pkg/storage/value/transformer.go#L172
@@ -184,7 +185,7 @@ impl InMemoryK8sEtcd {
 
         {
             let hashmap = self.etcd_keyvalue_hashmap.lock().await;
-            if let Some(value) = hashmap.get(&key) {
+            if let Some((_encoding, value)) = hashmap.get(&key) {
                 result.value.clone_from(value);
                 return Ok(Some(result));
             }
@@ -195,7 +196,7 @@ impl InMemoryK8sEtcd {
         if let Some(value) = get_result.kvs().first() {
             let raw_etcd_value = value.value();
 
-            let mut decoded_value = etcd_encoding::decode(raw_etcd_value).await.context("decoding value")?;
+            let (mut decoded_value, mut encoding) = etcd_encoding::decode(raw_etcd_value).await.context("decoding value")?;
 
             if let Some(resource_transformers) = &self.decrypt_resource_transformers {
                 // https://github.com/kubernetes/apiserver/blob/3423727e46efe7dfa40dcdb1a9c5c5027b07303d/pkg/storage/value/transformer.go#L110
@@ -209,7 +210,7 @@ impl InMemoryK8sEtcd {
                                 .decrypt(key.to_string(), raw_etcd_value.to_vec())
                                 .await
                                 .context("decrypting etcd value")?;
-                            decoded_value = etcd_encoding::decode(&plaintext_value).await.context("decoding value")?;
+                            (decoded_value, encoding) = etcd_encoding::decode(&plaintext_value).await.context("decoding value")?;
                             break;
                         }
                     }
@@ -219,7 +220,7 @@ impl InMemoryK8sEtcd {
             self.etcd_keyvalue_hashmap
                 .lock()
                 .await
-                .insert(key.to_string(), decoded_value.clone());
+                .insert(key.to_string(), (encoding, decoded_value.clone()));
 
             result.value = decoded_value;
             return Ok(Some(result));
@@ -228,10 +229,18 @@ impl InMemoryK8sEtcd {
         Ok(None)
     }
 
-    pub(crate) async fn put(&self, key: &str, value: Vec<u8>) {
-        self.etcd_keyvalue_hashmap.lock().await.insert(key.to_string(), value.clone());
+    pub(crate) async fn put(&self, key: &str, value: Vec<u8>) -> Result<()> {
+        let mut hashmap = self.etcd_keyvalue_hashmap.lock().await;
+
+        // Only put if the key already exists in the cache, preserving the encoding
+        let (encoding, _) = hashmap.get(key).context(format!("key '{}' not found in cache", key))?;
+        let encoding = encoding.clone(); // Clone the encoding
+        hashmap.insert(key.to_string(), (encoding, value.clone()));
+        drop(hashmap); // Release the lock early
+
         self.deleted_keys.lock().await.remove(key);
         self.edited.lock().await.insert(key.to_string(), value);
+        Ok(())
     }
 
     pub(crate) async fn list_keys(&self, resource_kind: &str) -> Result<Vec<String>> {
@@ -337,6 +346,6 @@ pub(crate) async fn get_etcd_json(client: &InMemoryK8sEtcd, k8slocation: &K8sRes
 pub(crate) async fn put_etcd_yaml(client: &InMemoryK8sEtcd, k8slocation: &K8sResourceLocation, value: Value) -> Result<()> {
     client
         .put(&k8slocation.as_etcd_key(), serde_json::to_string(&value)?.as_bytes().into())
-        .await;
+        .await.context("putting in etcd")?;
     Ok(())
 }

src/ocp_postprocess/encryption_config/etcd_rename.rs

@@ -177,7 +177,7 @@ async fn update_encryption_key(component: &str, etcd_client: &Arc<InMemoryK8sEtc
                                 )),
                                 serde_json::to_string(&secret).context("serializing value")?.as_bytes().to_vec(),
                             )
-                            .await;
+                            .await.context("putting in etcd")?;
                         etcd_client.delete(&key).await.context(format!("deleting {}", key))?;
                     }
                 }

src/ocp_postprocess/hostname_rename/etcd_rename.rs

@@ -227,7 +227,7 @@ pub(crate) async fn fix_etcd_secrets(etcd_client: &Arc<InMemoryK8sEtcd>, origina
                             &(format!("/kubernetes.io/secrets/openshift-etcd/{new_secret_name}")),
                             serde_json::to_string(&etcd_value).context("serializing value")?.as_bytes().to_vec(),
                         )
-                        .await;
+                        .await.context("putting in etcd")?;
 
                     etcd_client.delete(&key).await.context(format!("deleting {}", key))?;