chore(core::writing): Check if it exists before setting the header.

salvo-rs · Oct 10, 2024 · 71af9b6 · 71af9b6
1 parent 8bb0e52
commit 71af9b6
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 12 deletions.
diff --git a/crates/core/src/writing/json.rs b/crates/core/src/writing/json.rs
@@ -3,7 +3,7 @@ use std::fmt::{self, Debug, Display, Formatter};
 use async_trait::async_trait;
 use serde::Serialize;
 
-use super::Scribe;
+use super::{try_set_header, Scribe};
 use crate::http::header::{HeaderValue, CONTENT_TYPE};
 use crate::http::{Response, StatusError};
 
@@ -36,7 +36,8 @@ where
     fn render(self, res: &mut Response) {
         match serde_json::to_vec(&self.0) {
             Ok(bytes) => {
-                let _ = res.headers.append(
+                try_set_header(
+                    &mut res.headers,
                     CONTENT_TYPE,
                     HeaderValue::from_static("application/json; charset=utf-8"),
                 );

diff --git a/crates/core/src/writing/mod.rs b/crates/core/src/writing/mod.rs
@@ -5,7 +5,8 @@ mod redirect;
 mod seek;
 mod text;
 
-use http::StatusCode;
+use http::header::{AsHeaderName, IntoHeaderName};
+use http::{HeaderMap, StatusCode};
 pub use json::Json;
 pub use redirect::Redirect;
 pub use seek::ReadSeeker;
@@ -97,7 +98,8 @@ impl Scribe for StatusCode {
 impl Scribe for &'static str {
     #[inline]
     fn render(self, res: &mut Response) {
-        let _ = res.headers.append(
+        try_set_header(
+            &mut res.headers,
             CONTENT_TYPE,
             HeaderValue::from_static("text/plain; charset=utf-8"),
         );
@@ -107,7 +109,8 @@ impl Scribe for &'static str {
 impl Scribe for &String {
     #[inline]
     fn render(self, res: &mut Response) {
-        let _ = res.headers.append(
+        try_set_header(
+            &mut res.headers,
             CONTENT_TYPE,
             HeaderValue::from_static("text/plain; charset=utf-8"),
         );
@@ -117,7 +120,8 @@ impl Scribe for &String {
 impl Scribe for String {
     #[inline]
     fn render(self, res: &mut Response) {
-        let _ = res.headers.append(
+        try_set_header(
+            &mut res.headers,
             CONTENT_TYPE,
             HeaderValue::from_static("text/plain; charset=utf-8"),
         );
@@ -149,6 +153,16 @@ macro_rules! writer_tuple_impls {
 
 crate::for_each_tuple!(writer_tuple_impls);
 
+fn try_set_header<K, V>(headers: &mut HeaderMap<V>, key: K, val: V)
+where
+    K: IntoHeaderName,
+    for<'a> &'a K: AsHeaderName,
+{
+    if !headers.contains_key(&key) {
+        let _ = headers.insert(key, val);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::prelude::*;

diff --git a/crates/core/src/writing/text.rs b/crates/core/src/writing/text.rs
@@ -1,6 +1,6 @@
 use std::fmt::{self, Debug, Display, Formatter};
 
-use super::Scribe;
+use super::{try_set_header, Scribe};
 use crate::http::header::{HeaderValue, CONTENT_TYPE};
 use crate::http::Response;
 
@@ -44,7 +44,7 @@ impl<C> Text<C>
 where
     C: AsRef<str>,
 {
-    fn append_header(self, res: &mut Response) -> C {
+    fn try_set_header(self, res: &mut Response) -> C {
         let (ctype, content) = match self {
             Self::Plain(content) => (
                 HeaderValue::from_static("text/plain; charset=utf-8"),
@@ -81,28 +81,28 @@ where
                 content,
             ),
         };
-        res.headers.append(CONTENT_TYPE, ctype);
+        try_set_header(&mut res.headers, CONTENT_TYPE, ctype);
         content
     }
 }
 impl Scribe for Text<&'static str> {
     #[inline]
     fn render(self, res: &mut Response) {
-        let content = self.append_header(res);
+        let content = self.try_set_header(res);
         let _ = res.write_body(content);
     }
 }
 impl Scribe for Text<String> {
     #[inline]
     fn render(self, res: &mut Response) {
-        let content = self.append_header(res);
+        let content = self.try_set_header(res);
         let _ = res.write_body(content);
     }
 }
 impl Scribe for Text<&String> {
     #[inline]
     fn render(self, res: &mut Response) {
-        let content = self.append_header(res);
+        let content = self.try_set_header(res);
         let _ = res.write_body(content.as_bytes().to_vec());
     }
 }