use crate::transport::h2::hpack_impl::{Decoder, Encoder};
use bytes::Bytes;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PseudoHeaderOrder {
#[default]
Chrome,
Firefox,
Safari,
Standard,
Custom([u8; 4]),
}
impl PseudoHeaderOrder {
fn order(&self) -> [usize; 4] {
match self {
Self::Chrome => [0, 2, 1, 3], Self::Firefox => [0, 3, 1, 2], Self::Safari => [0, 2, 3, 1], Self::Standard => [0, 1, 2, 3], Self::Custom(order) => [
order[0] as usize,
order[1] as usize,
order[2] as usize,
order[3] as usize,
],
}
}
pub fn akamai_string(&self) -> &'static str {
match self {
Self::Chrome => "m,s,a,p",
Self::Firefox => "m,p,a,s",
Self::Safari => "m,s,p,a",
Self::Standard => "m,a,s,p",
Self::Custom(_) => "custom",
}
}
}
pub struct HpackEncoder {
encoder: Encoder,
pseudo_order: PseudoHeaderOrder,
}
impl HpackEncoder {
pub fn new(pseudo_order: PseudoHeaderOrder) -> Self {
Self {
encoder: Encoder::new(),
pseudo_order,
}
}
pub fn chrome() -> Self {
Self::new(PseudoHeaderOrder::Chrome)
}
pub fn set_max_table_size(&mut self, size: usize) {
self.encoder.set_max_table_size(size);
}
pub fn encode_request(
&mut self,
method: &str,
scheme: &str,
authority: &str,
path: &str,
headers: &[(String, String)],
) -> Bytes {
let pseudo_headers: [(&[u8], &[u8]); 4] = [
(b":method", method.as_bytes()),
(b":authority", authority.as_bytes()),
(b":scheme", scheme.as_bytes()),
(b":path", path.as_bytes()),
];
let mut all_headers: Vec<(&[u8], &[u8])> = Vec::new();
let mut valid_headers: Vec<(String, &str)> = Vec::with_capacity(headers.len());
for (name, value) in headers {
if name.starts_with(':') {
continue;
}
if name.is_empty() {
continue;
}
if name
.as_bytes()
.iter()
.any(|&b| b < 0x21 || (b > 0x7E && b != 0x7F))
{
continue;
}
let name_lower = name.to_lowercase();
if name_lower == "connection"
|| name_lower == "keep-alive"
|| name_lower == "proxy-connection"
|| name_lower == "transfer-encoding"
|| name_lower == "upgrade"
{
continue;
}
if name_lower == "te" && value.to_lowercase() != "trailers" {
continue;
}
valid_headers.push((name_lower, value));
}
let order = self.pseudo_order.order();
for &idx in &order {
all_headers.push(pseudo_headers[idx]);
}
for (n, v) in &valid_headers {
all_headers.push((n.as_bytes(), v.as_bytes()));
}
let encoded = self.encoder.encode(&all_headers);
Bytes::from(encoded)
}
pub fn chunk_encoded(encoded: Bytes, max_frame_size: usize) -> (Bytes, Vec<Bytes>) {
if encoded.len() <= max_frame_size {
return (encoded, Vec::new());
}
let mut chunks: Vec<Bytes> = encoded
.chunks(max_frame_size)
.map(Bytes::copy_from_slice)
.collect();
let first = chunks.remove(0);
(first, chunks)
}
}
pub struct HpackDecoder {
decoder: Decoder,
}
impl HpackDecoder {
pub fn new() -> Self {
Self {
decoder: Decoder::new(),
}
}
pub fn set_max_table_size(&mut self, size: usize) {
self.decoder.set_max_table_size(size);
}
pub fn decode(&mut self, data: &[u8]) -> Result<Vec<(String, String)>, String> {
let mut headers = Vec::new();
self.decoder
.decode_with_cb(data, |name, value| {
let name_str = String::from_utf8_lossy(name).into_owned();
let value_str = String::from_utf8_lossy(value).into_owned();
headers.push((name_str, value_str));
})
.map_err(|e| format!("HPACK decode error: {:?}", e))?;
Ok(headers)
}
}
impl Default for HpackDecoder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pseudo_order_chrome() {
let order = PseudoHeaderOrder::Chrome;
assert_eq!(order.akamai_string(), "m,s,a,p");
}
#[test]
fn test_pseudo_order_standard() {
let order = PseudoHeaderOrder::Standard;
assert_eq!(order.akamai_string(), "m,a,s,p");
}
#[test]
fn test_encoder_creates_valid_block() {
let mut encoder = HpackEncoder::chrome();
let block = encoder.encode_request(
"GET",
"https",
"example.com",
"/",
&[("user-agent".to_string(), "test".to_string())],
);
assert!(!block.is_empty());
let mut decoder = HpackDecoder::new();
let headers = decoder.decode(&block).unwrap();
assert_eq!(headers.len(), 5);
assert_eq!(headers[0].0, ":method");
assert_eq!(headers[0].1, "GET");
assert_eq!(headers[1].0, ":scheme");
assert_eq!(headers[1].1, "https");
assert_eq!(headers[2].0, ":authority");
assert_eq!(headers[2].1, "example.com");
assert_eq!(headers[3].0, ":path");
assert_eq!(headers[3].1, "/");
assert_eq!(headers[4].0, "user-agent");
assert_eq!(headers[4].1, "test");
}
#[test]
fn test_encoder_standard_order() {
let mut encoder = HpackEncoder::new(PseudoHeaderOrder::Standard);
let block = encoder.encode_request("GET", "https", "example.com", "/", &[]);
let mut decoder = HpackDecoder::new();
let headers = decoder.decode(&block).unwrap();
assert_eq!(headers[0].0, ":method");
assert_eq!(headers[1].0, ":authority");
assert_eq!(headers[2].0, ":scheme");
assert_eq!(headers[3].0, ":path");
}
#[test]
fn test_encoder_filters_connection_headers() {
let mut encoder = HpackEncoder::chrome();
let block = encoder.encode_request(
"GET",
"https",
"example.com",
"/",
&[
("connection".to_string(), "keep-alive".to_string()),
("keep-alive".to_string(), "timeout=5".to_string()),
("user-agent".to_string(), "test".to_string()),
],
);
let mut decoder = HpackDecoder::new();
let headers = decoder.decode(&block).unwrap();
assert_eq!(headers.len(), 5);
assert_eq!(headers[4].0, "user-agent");
}
}