From e25ab6d0fd7f64097706793b0209a59161dc9cb2 Mon Sep 17 00:00:00 2001
From: Younes <3153107+younes-io@users.noreply.github.com>
Date: Sat, 19 Oct 2024 10:08:07 +0200
Subject: [PATCH] Implement prost::Message for google.protobuf.Timestamp

Fixes #1171
---
 prost-build/src/config.rs     |  1 +
 prost-types/src/timestamp.rs  | 47 +++++++++++++++++++++++++++++++++++
 tests/src/well_known_types.rs | 18 ++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/prost-build/src/config.rs b/prost-build/src/config.rs
index af0e2e34d..35fd39bd2 100644
--- a/prost-build/src/config.rs
+++ b/prost-build/src/config.rs
@@ -382,6 +382,7 @@ impl Config {
     /// types, and instead generate Protobuf well-known types from their `.proto` definitions.
     pub fn compile_well_known_types(&mut self) -> &mut Self {
         self.prost_types = false;
+        self.extern_path(".google.protobuf.Timestamp", "::prost_types::Timestamp");
         self
     }
 
diff --git a/prost-types/src/timestamp.rs b/prost-types/src/timestamp.rs
index 1d7e609f4..7e498283f 100644
--- a/prost-types/src/timestamp.rs
+++ b/prost-types/src/timestamp.rs
@@ -427,3 +427,50 @@ mod tests {
         }
     }
 }
+
+impl prost::Message for Timestamp {
+    fn encode_raw<B>(&self, buf: &mut B)
+    where
+        B: bytes::BufMut,
+    {
+        if self.seconds != 0 {
+            prost::encoding::int64::encode(1, &self.seconds, buf);
+        }
+        if self.nanos != 0 {
+            prost::encoding::int32::encode(2, &self.nanos, buf);
+        }
+    }
+
+    fn merge_field<B>(
+        &mut self,
+        tag: u32,
+        wire_type: prost::encoding::WireType,
+        buf: &mut B,
+        ctx: prost::encoding::DecodeContext,
+    ) -> Result<(), prost::DecodeError>
+    where
+        B: bytes::Buf,
+    {
+        match tag {
+            1 => prost::encoding::int64::merge(wire_type, &mut self.seconds, buf, ctx),
+            2 => prost::encoding::int32::merge(wire_type, &mut self.nanos, buf, ctx),
+            _ => prost::encoding::skip_field(wire_type, tag, buf, ctx),
+        }
+    }
+
+    fn encoded_len(&self) -> usize {
+        let mut len = 0;
+        if self.seconds != 0 {
+            len += prost::encoding::int64::encoded_len(1, &self.seconds);
+        }
+        if self.nanos != 0 {
+            len += prost::encoding::int32::encoded_len(2, &self.nanos);
+        }
+        len
+    }
+
+    fn clear(&mut self) {
+        self.seconds = 0;
+        self.nanos = 0;
+    }
+}
diff --git a/tests/src/well_known_types.rs b/tests/src/well_known_types.rs
index 88fa1ec52..fe9ee375f 100644
--- a/tests/src/well_known_types.rs
+++ b/tests/src/well_known_types.rs
@@ -69,3 +69,21 @@ mod include {
         "/wellknown_include/wellknown_include.rs"
     ));
 }
+
+#[test]
+fn test_prost_message_for_timestamp() {
+    use prost::Message;
+
+    let timestamp = ::prost_types::Timestamp {
+        seconds: 100,
+        nanos: 42,
+    };
+
+    let mut buf = Vec::new();
+    timestamp.encode(&mut buf).expect("Failed to encode Timestamp");
+
+    let decoded_timestamp = ::prost_types::Timestamp::decode(&buf[..])
+        .expect("Failed to decode Timestamp");
+
+    assert_eq!(timestamp, decoded_timestamp, "Encoded and decoded Timestamp should be equal");
+}