1use armature_h1::{ByteStr, HeaderId, header as header_id};
37use bytes::Bytes;
38use smallvec::SmallVec;
39use std::collections::HashMap;
40use std::fmt;
41
42pub const INLINE_HEADERS: usize = 12;
46
47enum Needle<'a> {
56 Known(HeaderId),
57 Custom(&'a str),
58}
59
60impl<'a> Needle<'a> {
61 #[inline]
62 fn new(name: &'a str) -> Self {
63 match HeaderId::from_bytes(name.as_bytes()) {
64 Some(id) => Needle::Known(id),
65 None => Needle::Custom(name),
66 }
67 }
68
69 #[inline]
74 fn matches(&self, id: &HeaderId) -> bool {
75 match self {
76 Needle::Known(known) => known == id,
77 Needle::Custom(name) => id.as_str().eq_ignore_ascii_case(name),
78 }
79 }
80}
81
82#[derive(Clone, PartialEq, Eq)]
84pub struct Header {
85 pub id: HeaderId,
87 pub value: Bytes,
89}
90
91impl Header {
92 #[inline]
94 pub fn new(name: impl AsRef<str>, value: impl HeaderValueInput) -> Self {
95 Self {
96 id: header_id::intern(name.as_ref()),
97 value: value.into_value(),
98 }
99 }
100
101 #[inline]
103 pub fn name(&self) -> &str {
104 self.id.as_str()
105 }
106
107 #[inline]
109 pub fn value_str(&self) -> Option<&str> {
110 std::str::from_utf8(&self.value).ok()
111 }
112}
113
114impl fmt::Debug for Header {
115 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116 match self.value_str() {
117 Some(v) => write!(f, "{}: {}", self.name(), v),
118 None => write!(f, "{}: <{} non-utf8 bytes>", self.name(), self.value.len()),
119 }
120 }
121}
122
123pub trait HeaderValueInput {
130 fn into_value(self) -> Bytes;
132}
133
134impl HeaderValueInput for Bytes {
135 #[inline]
136 fn into_value(self) -> Bytes {
137 self
138 }
139}
140
141impl HeaderValueInput for &str {
142 #[inline]
143 fn into_value(self) -> Bytes {
144 Bytes::copy_from_slice(self.as_bytes())
145 }
146}
147
148impl HeaderValueInput for &String {
149 #[inline]
150 fn into_value(self) -> Bytes {
151 Bytes::copy_from_slice(self.as_bytes())
152 }
153}
154
155impl HeaderValueInput for String {
156 #[inline]
157 fn into_value(self) -> Bytes {
158 Bytes::from(self.into_bytes())
159 }
160}
161
162impl HeaderValueInput for &[u8] {
163 #[inline]
164 fn into_value(self) -> Bytes {
165 Bytes::copy_from_slice(self)
166 }
167}
168
169impl HeaderValueInput for Vec<u8> {
170 #[inline]
171 fn into_value(self) -> Bytes {
172 Bytes::from(self)
173 }
174}
175
176impl HeaderValueInput for ByteStr {
177 #[inline]
178 fn into_value(self) -> Bytes {
179 self.into_bytes()
180 }
181}
182
183impl HeaderValueInput for std::borrow::Cow<'_, str> {
184 #[inline]
185 fn into_value(self) -> Bytes {
186 match self {
187 std::borrow::Cow::Borrowed(s) => Bytes::copy_from_slice(s.as_bytes()),
188 std::borrow::Cow::Owned(s) => Bytes::from(s.into_bytes()),
189 }
190 }
191}
192
193#[derive(Clone, Default)]
209pub struct HeaderMap {
210 inner: SmallVec<[Header; INLINE_HEADERS]>,
211}
212
213impl HeaderMap {
214 #[inline]
216 pub const fn new() -> Self {
217 Self {
218 inner: SmallVec::new_const(),
219 }
220 }
221
222 #[inline]
226 pub fn with_capacity(capacity: usize) -> Self {
227 Self {
228 inner: SmallVec::with_capacity(capacity),
229 }
230 }
231
232 #[inline]
234 pub fn is_inline(&self) -> bool {
235 !self.inner.spilled()
236 }
237
238 #[inline]
240 pub fn len(&self) -> usize {
241 self.inner.len()
242 }
243
244 #[inline]
246 pub fn is_empty(&self) -> bool {
247 self.inner.is_empty()
248 }
249
250 #[inline]
255 pub fn get(&self, name: &str) -> Option<&str> {
256 self.get_bytes(name)
257 .and_then(|v| std::str::from_utf8(v).ok())
258 }
259
260 #[inline]
262 pub fn get_bytes(&self, name: &str) -> Option<&Bytes> {
263 let needle = Needle::new(name);
264 self.inner
265 .iter()
266 .find(|h| needle.matches(&h.id))
267 .map(|h| &h.value)
268 }
269
270 #[inline]
275 pub fn get_id(&self, id: &HeaderId) -> Option<&Bytes> {
276 self.inner.iter().find(|h| &h.id == id).map(|h| &h.value)
277 }
278
279 #[inline]
283 pub fn get_ignore_case(&self, name: &str) -> Option<&str> {
284 self.get(name)
285 }
286
287 #[inline]
289 pub fn contains(&self, name: &str) -> bool {
290 self.get_bytes(name).is_some()
291 }
292
293 #[inline]
297 pub fn contains_key(&self, name: &str) -> bool {
298 self.contains(name)
299 }
300
301 #[inline]
305 pub fn insert(&mut self, name: impl AsRef<str>, value: impl HeaderValueInput) -> Option<Bytes> {
306 let id = header_id::intern(name.as_ref());
307 let value = value.into_value();
308 if let Some(existing) = self.inner.iter_mut().find(|h| h.id == id) {
309 return Some(std::mem::replace(&mut existing.value, value));
310 }
311 self.inner.push(Header { id, value });
312 None
313 }
314
315 #[inline]
320 pub fn append(&mut self, name: impl AsRef<str>, value: impl HeaderValueInput) {
321 self.inner.push(Header {
322 id: header_id::intern(name.as_ref()),
323 value: value.into_value(),
324 });
325 }
326
327 #[inline]
329 pub fn remove(&mut self, name: &str) -> Option<Bytes> {
330 let needle = Needle::new(name);
331 let pos = self.inner.iter().position(|h| needle.matches(&h.id))?;
332 Some(self.inner.remove(pos).value)
333 }
334
335 #[inline]
337 pub fn remove_all(&mut self, name: &str) -> usize {
338 let needle = Needle::new(name);
339 let before = self.inner.len();
340 self.inner.retain(|h| !needle.matches(&h.id));
341 before - self.inner.len()
342 }
343
344 #[inline]
346 pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
347 self.inner
348 .iter()
349 .filter_map(|h| h.value_str().map(|v| (h.name(), v)))
350 }
351
352 #[inline]
354 pub fn iter_raw(&self) -> impl Iterator<Item = (&HeaderId, &Bytes)> {
355 self.inner.iter().map(|h| (&h.id, &h.value))
356 }
357
358 #[inline]
360 pub fn names(&self) -> impl Iterator<Item = &str> {
361 self.inner.iter().map(|h| h.name())
362 }
363
364 #[inline]
368 pub fn keys(&self) -> impl Iterator<Item = &str> {
369 self.names()
370 }
371
372 #[inline]
374 pub fn values(&self) -> impl Iterator<Item = &str> {
375 self.inner.iter().filter_map(|h| h.value_str())
376 }
377
378 #[inline]
380 pub fn get_all(&self, name: &str) -> Vec<&str> {
381 let needle = Needle::new(name);
382 self.inner
383 .iter()
384 .filter(|h| needle.matches(&h.id))
385 .filter_map(|h| h.value_str())
386 .collect()
387 }
388
389 #[inline]
391 pub fn clear(&mut self) {
392 self.inner.clear();
393 }
394
395 #[inline]
397 pub fn extend<I, K, V>(&mut self, iter: I)
398 where
399 I: IntoIterator<Item = (K, V)>,
400 K: AsRef<str>,
401 V: HeaderValueInput,
402 {
403 for (k, v) in iter {
404 self.insert(k, v);
405 }
406 }
407
408 #[inline]
414 pub fn to_hash_map(&self) -> HashMap<String, String> {
415 self.iter()
416 .map(|(k, v)| (k.to_owned(), v.to_owned()))
417 .collect()
418 }
419
420 #[inline]
422 pub fn from_hash_map(map: HashMap<String, String>) -> Self {
423 let mut headers = Self::with_capacity(map.len());
424 for (k, v) in map {
425 headers.insert(k, v);
426 }
427 headers
428 }
429
430 #[inline]
436 pub fn content_type(&self) -> Option<&str> {
437 self.str_of(&HeaderId::ContentType)
438 }
439
440 #[inline]
442 pub fn content_length(&self) -> Option<usize> {
443 self.str_of(&HeaderId::ContentLength)?.parse().ok()
444 }
445
446 #[inline]
448 pub fn accept(&self) -> Option<&str> {
449 self.str_of(&HeaderId::Accept)
450 }
451
452 #[inline]
454 pub fn authorization(&self) -> Option<&str> {
455 self.str_of(&HeaderId::Authorization)
456 }
457
458 #[inline]
460 pub fn user_agent(&self) -> Option<&str> {
461 self.str_of(&HeaderId::UserAgent)
462 }
463
464 #[inline]
466 pub fn host(&self) -> Option<&str> {
467 self.str_of(&HeaderId::Host)
468 }
469
470 #[inline]
472 pub fn cookie(&self) -> Option<&str> {
473 self.str_of(&HeaderId::Cookie)
474 }
475
476 #[inline]
478 pub fn is_keep_alive(&self) -> bool {
479 self.str_of(&HeaderId::Connection)
480 .map(|v| v.eq_ignore_ascii_case("keep-alive"))
481 .unwrap_or(true) }
483
484 #[inline]
486 pub fn is_chunked(&self) -> bool {
487 self.str_of(&HeaderId::TransferEncoding)
488 .map(|v| v.contains("chunked"))
489 .unwrap_or(false)
490 }
491
492 #[inline]
494 pub fn set_content_type(&mut self, value: impl HeaderValueInput) {
495 self.insert("content-type", value);
496 }
497
498 #[inline]
500 pub fn set_content_length(&mut self, len: usize) {
501 self.insert("content-length", len.to_string());
502 }
503
504 #[inline]
506 fn str_of(&self, id: &HeaderId) -> Option<&str> {
507 self.get_id(id).and_then(|v| std::str::from_utf8(v).ok())
508 }
509}
510
511impl fmt::Debug for HeaderMap {
512 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
513 f.debug_map()
514 .entries(
515 self.inner
516 .iter()
517 .map(|h| (h.name(), h.value_str().unwrap_or("<non-utf8>"))),
518 )
519 .finish()
520 }
521}
522
523impl<K, V> FromIterator<(K, V)> for HeaderMap
524where
525 K: AsRef<str>,
526 V: HeaderValueInput,
527{
528 fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
529 let iter = iter.into_iter();
530 let (min, max) = iter.size_hint();
531 let mut map = HeaderMap::with_capacity(max.unwrap_or(min));
532 for (k, v) in iter {
533 map.insert(k, v);
534 }
535 map
536 }
537}
538
539impl Extend<(String, String)> for HeaderMap {
540 fn extend<I: IntoIterator<Item = (String, String)>>(&mut self, iter: I) {
541 for (k, v) in iter {
542 self.insert(k, v);
543 }
544 }
545}
546
547fn utf8_pair(h: &Header) -> Option<(&str, &str)> {
552 h.value_str().map(|v| (h.name(), v))
553}
554
555fn owned_utf8_pair(h: Header) -> Option<(String, String)> {
557 let name = h.name().to_owned();
558 String::from_utf8(h.value.to_vec())
559 .ok()
560 .map(|value| (name, value))
561}
562
563impl<'a> IntoIterator for &'a HeaderMap {
564 type Item = (&'a str, &'a str);
565 type IntoIter = std::iter::FilterMap<
566 std::slice::Iter<'a, Header>,
567 fn(&'a Header) -> Option<(&'a str, &'a str)>,
568 >;
569
570 fn into_iter(self) -> Self::IntoIter {
571 self.inner.iter().filter_map(utf8_pair as _)
572 }
573}
574
575impl IntoIterator for HeaderMap {
576 type Item = (String, String);
577 type IntoIter = std::iter::FilterMap<
578 smallvec::IntoIter<[Header; INLINE_HEADERS]>,
579 fn(Header) -> Option<(String, String)>,
580 >;
581
582 fn into_iter(self) -> Self::IntoIter {
583 self.inner.into_iter().filter_map(owned_utf8_pair as _)
584 }
585}
586
587impl std::ops::Index<&str> for HeaderMap {
589 type Output = str;
590
591 fn index(&self, name: &str) -> &Self::Output {
592 self.get(name).expect("header not found")
593 }
594}
595
596impl From<HashMap<String, String>> for HeaderMap {
601 fn from(map: HashMap<String, String>) -> Self {
602 Self::from_hash_map(map)
603 }
604}
605
606impl From<HeaderMap> for HashMap<String, String> {
607 fn from(map: HeaderMap) -> Self {
608 map.to_hash_map()
609 }
610}
611
612#[cfg(test)]
617mod tests {
618 use super::*;
619
620 #[test]
621 fn test_new_is_inline() {
622 let headers = HeaderMap::new();
623 assert!(headers.is_inline());
624 assert!(headers.is_empty());
625 }
626
627 #[test]
628 fn get_returns_str_and_well_known_names_are_interned() {
629 let mut h = HeaderMap::new();
630 h.insert("Content-Type", "application/json");
631 h.insert("X-Tenant-Id", "acme".to_string());
632
633 assert_eq!(h.get("content-type"), Some("application/json"));
635 assert_eq!(h.get("CONTENT-TYPE"), Some("application/json"));
636 assert_eq!(h.get("x-tenant-id"), Some("acme"));
637 assert_eq!(h.get("absent"), None);
638
639 assert_eq!(
641 h.get_id(&HeaderId::ContentType).map(|b| &b[..]),
642 Some(&b"application/json"[..])
643 );
644 }
645
646 #[test]
647 fn custom_names_stay_case_insensitive_through_the_borrowed_needle() {
648 let mut h = HeaderMap::new();
652 h.insert("X-Request-ID", "abc123");
653 h.append("x-request-id", "def456");
654
655 assert_eq!(h.get("x-request-id"), Some("abc123"));
656 assert_eq!(h.get("X-REQUEST-ID"), Some("abc123"));
657 assert!(h.contains("X-Request-Id"));
658 assert_eq!(h.get_all("X-Request-Id"), vec!["abc123", "def456"]);
659
660 h.insert("Content-Type", "text/plain");
662 assert_eq!(h.get("x-content-type"), None);
663
664 assert_eq!(h.remove_all("X-Request-ID"), 2);
665 assert_eq!(h.get("x-request-id"), None);
666 }
667
668 #[test]
669 fn non_utf8_value_is_invisible_to_get_but_reachable_as_bytes() {
670 let mut h = HeaderMap::new();
671 h.insert("x-raw", Bytes::from_static(&[0xff, 0x00]));
672 assert_eq!(h.get("x-raw"), None);
675 assert_eq!(h.get_bytes("x-raw").map(|b| b.len()), Some(2));
676 assert_eq!(h.len(), 1);
678 assert_eq!(h.iter().count(), 0);
679 assert_eq!(h.iter_raw().count(), 1);
680 }
681
682 #[test]
683 fn test_insert_and_get() {
684 let mut headers = HeaderMap::new();
685 headers.insert("Content-Type", "application/json");
686 headers.insert("Accept", "text/html");
687
688 assert_eq!(headers.len(), 2);
689 assert_eq!(headers.get("Content-Type"), Some("application/json"));
690 assert_eq!(headers.get("content-type"), Some("application/json"));
691 }
692
693 #[test]
694 fn test_insert_replaces() {
695 let mut headers = HeaderMap::new();
696 headers.insert("Content-Type", "text/plain");
697 let old = headers.insert("Content-Type", "application/json");
698
699 assert_eq!(old.as_deref(), Some(&b"text/plain"[..]));
700 assert_eq!(headers.len(), 1);
701 assert_eq!(headers.get("Content-Type"), Some("application/json"));
702 }
703
704 #[test]
705 fn test_append_duplicates() {
706 let mut headers = HeaderMap::new();
707 headers.append("Set-Cookie", "session=abc");
708 headers.append("Set-Cookie", "user=123");
709
710 assert_eq!(headers.len(), 2);
711 assert_eq!(
712 headers.get_all("set-cookie"),
713 vec!["session=abc", "user=123"]
714 );
715 }
716
717 #[test]
718 fn test_remove() {
719 let mut headers = HeaderMap::new();
720 headers.insert("Content-Type", "application/json");
721 headers.insert("Accept", "text/html");
722
723 let removed = headers.remove("Content-Type");
724 assert_eq!(removed.as_deref(), Some(&b"application/json"[..]));
725 assert_eq!(headers.len(), 1);
726 assert!(!headers.contains("Content-Type"));
727 }
728
729 #[test]
730 fn test_remove_all() {
731 let mut headers = HeaderMap::new();
732 headers.append("Set-Cookie", "a=1");
733 headers.append("set-cookie", "b=2");
734 headers.insert("Accept", "*/*");
735
736 assert_eq!(headers.remove_all("Set-Cookie"), 2);
737 assert_eq!(headers.len(), 1);
738 }
739
740 #[test]
741 fn test_inline_capacity() {
742 let mut headers = HeaderMap::new();
743
744 for i in 0..INLINE_HEADERS {
745 headers.insert(format!("Header-{i}"), format!("Value-{i}"));
746 }
747 assert!(headers.is_inline());
748
749 headers.insert("Extra-Header", "Extra-Value");
750 assert!(!headers.is_inline());
751 }
752
753 #[test]
754 fn test_iter() {
755 let mut headers = HeaderMap::new();
756 headers.insert("A", "1");
757 headers.insert("B", "2");
758
759 let pairs: Vec<_> = headers.iter().collect();
760 assert_eq!(pairs.len(), 2);
761 }
762
763 #[test]
764 fn iter_yields_lowercased_names_for_custom_headers() {
765 let mut h = HeaderMap::new();
766 h.insert("X-A", "1");
767 assert_eq!(h.iter().collect::<Vec<_>>(), vec![("x-a", "1")]);
770 }
771
772 #[test]
773 fn test_common_accessors() {
774 let mut headers = HeaderMap::new();
775 headers.insert("Content-Type", "application/json");
776 headers.insert("Content-Length", "100");
777 headers.insert("Connection", "keep-alive");
778 headers.insert("Transfer-Encoding", "chunked");
779
780 assert_eq!(headers.content_type(), Some("application/json"));
781 assert_eq!(headers.content_length(), Some(100));
782 assert!(headers.is_keep_alive());
783 assert!(headers.is_chunked());
784 }
785
786 #[test]
787 fn test_from_hash_map() {
788 let mut map = HashMap::new();
789 map.insert("Content-Type".to_string(), "application/json".to_string());
790 map.insert("Accept".to_string(), "text/html".to_string());
791
792 let headers = HeaderMap::from_hash_map(map);
793 assert_eq!(headers.len(), 2);
794 assert!(headers.contains("Content-Type"));
795 }
796
797 #[test]
798 fn test_to_hash_map_normalizes_names_to_lowercase() {
799 let mut headers = HeaderMap::new();
800 headers.insert("Content-Type", "application/json");
801
802 let map = headers.to_hash_map();
803 assert_eq!(
807 map.get("content-type").map(String::as_str),
808 Some("application/json")
809 );
810 assert_eq!(map.get("Content-Type"), None);
811 }
812
813 #[test]
814 fn test_from_iterator() {
815 let headers: HeaderMap = [
816 ("Content-Type", "application/json"),
817 ("Accept", "text/html"),
818 ]
819 .into_iter()
820 .collect();
821
822 assert_eq!(headers.len(), 2);
823 }
824
825 #[test]
826 fn test_indexing() {
827 let mut headers = HeaderMap::new();
828 headers.insert("Content-Type", "application/json");
829
830 assert_eq!(&headers["Content-Type"], "application/json");
831 }
832
833 #[test]
834 fn test_contains_key() {
835 let mut headers = HeaderMap::new();
836 headers.insert("Content-Type", "application/json");
837
838 assert!(headers.contains_key("Content-Type"));
839 assert!(headers.contains_key("content-type"));
840 assert!(!headers.contains_key("Accept"));
841 }
842
843 #[test]
844 fn test_keys() {
845 let mut headers = HeaderMap::new();
846 headers.insert("Content-Type", "application/json");
847 headers.insert("Accept", "text/html");
848
849 let keys: Vec<_> = headers.keys().collect();
850 assert_eq!(keys.len(), 2);
851 assert!(keys.contains(&"content-type"));
852 assert!(keys.contains(&"accept"));
853 }
854
855 #[test]
856 fn test_values() {
857 let mut headers = HeaderMap::new();
858 headers.insert("Content-Type", "application/json");
859 headers.insert("Accept", "text/html");
860
861 let values: Vec<_> = headers.values().collect();
862 assert_eq!(values.len(), 2);
863 assert!(values.contains(&"application/json"));
864 assert!(values.contains(&"text/html"));
865 }
866
867 #[test]
868 fn test_is_empty() {
869 let mut headers = HeaderMap::new();
870 assert!(headers.is_empty());
871 headers.insert("Content-Type", "application/json");
872 assert!(!headers.is_empty());
873 }
874
875 #[test]
876 fn test_default() {
877 let headers = HeaderMap::default();
878 assert!(headers.is_empty());
879 assert!(headers.is_inline());
880 }
881
882 #[test]
883 fn test_extend_trait() {
884 let mut headers = HeaderMap::new();
885 headers.insert("Existing", "1");
886
887 let extra: Vec<(String, String)> = vec![
888 ("Content-Type".to_string(), "application/json".to_string()),
889 ("Accept".to_string(), "text/html".to_string()),
890 ];
891 Extend::extend(&mut headers, extra);
892
893 assert_eq!(headers.len(), 3);
894 assert_eq!(headers.get("Content-Type"), Some("application/json"));
895 }
896
897 #[test]
898 fn test_into_iterator_owned() {
899 let mut headers = HeaderMap::new();
900 headers.insert("A", "1");
901 headers.insert("B", "2");
902
903 let collected: Vec<(String, String)> = headers.into_iter().collect();
904 assert_eq!(collected.len(), 2);
905 }
906
907 #[test]
908 fn test_into_iterator_ref() {
909 let mut headers = HeaderMap::new();
910 headers.insert("A", "1");
911
912 let collected: Vec<(&str, &str)> = (&headers).into_iter().collect();
913 assert_eq!(collected, vec![("a", "1")]);
914 }
915
916 #[test]
917 fn test_hashmap_roundtrip() {
918 let mut map = HashMap::new();
919 map.insert("Content-Type".to_string(), "application/json".to_string());
920
921 let headers: HeaderMap = map.clone().into();
922 assert!(headers.contains_key("content-type"));
923 let back: HashMap<String, String> = headers.into();
924 assert_eq!(back.get("content-type"), map.get("Content-Type"));
927 }
928
929 #[test]
930 fn cloning_a_value_does_not_copy_it() {
931 let mut headers = HeaderMap::new();
932 let big = Bytes::from(vec![b'x'; 4096]);
933 headers.insert("x-big", big.clone());
934 let copy = headers.clone();
935 assert_eq!(
936 copy.get_bytes("x-big").map(|b| b.as_ptr()),
937 Some(big.as_ptr())
938 );
939 }
940}