use bytes::{BufMut, Bytes, BytesMut};
use std::ops::Deref;
pub const DEFAULT_RESPONSE_CAPACITY: usize = 512;
pub const MEDIUM_RESPONSE_CAPACITY: usize = 4096;
pub const LARGE_RESPONSE_CAPACITY: usize = 65536;
#[derive(Debug)]
pub struct ResponseBuffer {
inner: BytesMut,
}
impl ResponseBuffer {
#[inline]
pub fn new() -> Self {
Self {
inner: BytesMut::with_capacity(DEFAULT_RESPONSE_CAPACITY),
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self {
inner: BytesMut::with_capacity(capacity),
}
}
#[inline]
pub fn medium() -> Self {
Self::with_capacity(MEDIUM_RESPONSE_CAPACITY)
}
#[inline]
pub fn large() -> Self {
Self::with_capacity(LARGE_RESPONSE_CAPACITY)
}
#[inline]
pub fn len(&self) -> usize {
self.inner.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[inline]
pub fn remaining_capacity(&self) -> usize {
self.inner.capacity() - self.inner.len()
}
#[inline]
pub fn capacity(&self) -> usize {
self.inner.capacity()
}
#[inline]
pub fn write(&mut self, data: &[u8]) {
self.inner.extend_from_slice(data);
}
#[inline]
pub fn write_str(&mut self, s: &str) {
self.inner.extend_from_slice(s.as_bytes());
}
#[inline]
pub fn put(&mut self, data: &[u8]) {
self.inner.put_slice(data);
}
#[inline]
pub fn reserve(&mut self, additional: usize) {
self.inner.reserve(additional);
}
#[inline]
pub fn clear(&mut self) {
self.inner.clear();
}
#[inline]
pub fn freeze(self) -> Bytes {
self.inner.freeze()
}
#[inline]
pub fn as_slice(&self) -> &[u8] {
&self.inner
}
#[inline]
pub fn inner_mut(&mut self) -> &mut BytesMut {
&mut self.inner
}
#[inline]
pub fn into_inner(self) -> BytesMut {
self.inner
}
#[inline]
pub fn split(&mut self) -> Bytes {
self.inner.split().freeze()
}
}
impl Default for ResponseBuffer {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl Deref for ResponseBuffer {
type Target = [u8];
#[inline]
fn deref(&self) -> &Self::Target {
&self.inner
}
}
impl AsRef<[u8]> for ResponseBuffer {
#[inline]
fn as_ref(&self) -> &[u8] {
&self.inner
}
}
impl AsMut<BytesMut> for ResponseBuffer {
#[inline]
fn as_mut(&mut self) -> &mut BytesMut {
&mut self.inner
}
}
impl From<ResponseBuffer> for Bytes {
#[inline]
fn from(buf: ResponseBuffer) -> Self {
buf.freeze()
}
}
impl From<ResponseBuffer> for BytesMut {
#[inline]
fn from(buf: ResponseBuffer) -> Self {
buf.inner
}
}
impl std::io::Write for ResponseBuffer {
#[inline]
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.inner.extend_from_slice(buf);
Ok(buf.len())
}
#[inline]
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[derive(Debug)]
pub struct ResponseBuilder {
status: u16,
headers: Vec<(String, String)>,
body: ResponseBuffer,
}
impl ResponseBuilder {
#[inline]
pub fn new() -> Self {
Self {
status: 200,
headers: Vec::with_capacity(8), body: ResponseBuffer::new(),
}
}
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self {
status: 200,
headers: Vec::with_capacity(8),
body: ResponseBuffer::with_capacity(capacity),
}
}
#[inline]
pub fn status(mut self, status: u16) -> Self {
self.status = status;
self
}
#[inline]
pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.push((name.into(), value.into()));
self
}
#[inline]
pub fn content_type(self, value: impl Into<String>) -> Self {
self.header("Content-Type", value)
}
#[inline]
pub fn body(mut self, data: &[u8]) -> Self {
self.body.write(data);
self
}
#[inline]
pub fn body_str(mut self, s: &str) -> Self {
self.body.write_str(s);
self
}
#[inline]
pub fn body_json<T: serde::Serialize>(mut self, value: &T) -> Result<Self, crate::Error> {
let json =
crate::json::to_vec(value).map_err(|e| crate::Error::Serialization(e.to_string()))?;
self.body.write(&json);
self.headers
.push(("Content-Type".to_string(), "application/json".to_string()));
Ok(self)
}
#[inline]
pub fn body_mut(&mut self) -> &mut ResponseBuffer {
&mut self.body
}
#[inline]
pub fn build(self) -> crate::HttpResponse {
let mut response = crate::HttpResponse::new(self.status);
for (name, value) in self.headers {
response.headers.insert(name, value);
}
if !self.body.is_empty() {
response = response.with_bytes_body(self.body.freeze());
}
response
}
}
impl Default for ResponseBuilder {
#[inline]
fn default() -> Self {
Self::new()
}
}
#[inline]
pub fn response_builder() -> ResponseBuilder {
ResponseBuilder::new()
}
#[inline]
pub fn response_builder_with_capacity(capacity: usize) -> ResponseBuilder {
ResponseBuilder::with_capacity(capacity)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_response_buffer_new() {
let buf = ResponseBuffer::new();
assert_eq!(buf.capacity(), DEFAULT_RESPONSE_CAPACITY);
assert!(buf.is_empty());
}
#[test]
fn test_response_buffer_write() {
let mut buf = ResponseBuffer::new();
buf.write(b"Hello, ");
buf.write(b"World!");
assert_eq!(buf.as_slice(), b"Hello, World!");
}
#[test]
fn test_response_buffer_no_realloc_small() {
let mut buf = ResponseBuffer::new();
let initial_cap = buf.capacity();
buf.write(b"Small response");
assert_eq!(buf.capacity(), initial_cap);
}
#[test]
fn test_response_buffer_freeze() {
let mut buf = ResponseBuffer::new();
buf.write(b"test data");
let bytes = buf.freeze();
assert_eq!(&*bytes, b"test data");
}
#[test]
fn test_response_buffer_medium() {
let buf = ResponseBuffer::medium();
assert_eq!(buf.capacity(), MEDIUM_RESPONSE_CAPACITY);
}
#[test]
fn test_response_buffer_large() {
let buf = ResponseBuffer::large();
assert_eq!(buf.capacity(), LARGE_RESPONSE_CAPACITY);
}
#[test]
fn test_response_builder() {
let response = ResponseBuilder::new()
.status(201)
.header("X-Custom", "value")
.body_str("Created")
.build();
assert_eq!(response.status, 201);
assert_eq!(response.headers.get("X-Custom"), Some(&"value".to_string()));
}
#[test]
fn test_response_builder_json() {
#[derive(serde::Serialize)]
struct Data {
message: &'static str,
}
let response = ResponseBuilder::new()
.status(200)
.body_json(&Data { message: "ok" })
.unwrap()
.build();
assert_eq!(response.status, 200);
assert!(
response
.headers
.get("Content-Type")
.unwrap()
.contains("json")
);
}
#[test]
fn test_http_response_with_capacity() {
let response = crate::HttpResponse::with_capacity(200, 1024);
assert_eq!(response.status, 200);
assert!(response.body.capacity() >= 1024);
}
#[test]
fn test_http_response_preallocated() {
let response = crate::HttpResponse::ok_preallocated();
assert_eq!(response.status, 200);
assert!(response.body.capacity() >= 512);
}
#[test]
fn test_response_builder_standalone() {
let response = response_builder()
.status(404)
.content_type("text/plain")
.body_str("Not Found")
.build();
assert_eq!(response.status, 404);
}
#[test]
fn test_response_buffer_io_write() {
use std::io::Write;
let mut buf = ResponseBuffer::new();
write!(buf, "Hello, World!").unwrap();
assert_eq!(buf.as_slice(), b"Hello, World!");
}
}