From e0eb5dfe189d2ccff546265177b5a7914e98c664 Mon Sep 17 00:00:00 2001 From: Brendan Zabarauskas Date: Sat, 15 Jul 2017 16:54:46 +1000 Subject: [PATCH] Add tests for forwarded for header --- src/lib.rs | 33 +++++++++++++++++------ tests/lib.rs | 75 +++++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 88 insertions(+), 20 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index a55caee..3980356 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,9 +9,8 @@ extern crate unicase; extern crate void; use futures::future::Future; -use hyper::{Body, Request, Response, StatusCode}; +use hyper::{Body, Headers, Request, Response, StatusCode}; use hyper::server::Service; -use hyper::header::Headers; use std::marker::PhantomData; use std::net::IpAddr; use void::Void; @@ -61,6 +60,24 @@ header! { /// * `203.0.113.195` /// * `203.0.113.195, 70.41.3.18, 150.172.238.178` /// + /// # Examples + /// + /// ``` + /// # extern crate hyper; + /// # extern crate hyper_reverse_proxy; + /// use hyper::Headers; + /// use hyper_reverse_proxy::XForwardedFor; + /// use std::net::{Ipv4Addr, Ipv6Addr}; + /// + /// # fn main() { + /// let mut headers = Headers::new(); + /// headers.set(XForwardedFor(vec![ + /// Ipv4Addr::new(127, 0, 0, 1).into(), + /// Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1).into(), + /// ])); + /// # } + /// ``` + /// /// # References /// /// - [MDN](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/X-Forwarded-For) @@ -93,16 +110,16 @@ fn create_proxied_response(mut response: Response) -> Response { /// [`httputil.ReverseProxy`]: https://golang.org/pkg/net/http/httputil/#ReverseProxy pub struct ReverseProxy { client: C, - remote_ip_addr: Option, + remote_ip: Option, _pantom_data: PhantomData, } impl ReverseProxy { /// Construct a reverse proxy that dispatches to the given client. - pub fn new(client: C, remote_ip_addr: Option) -> ReverseProxy { + pub fn new(client: C, remote_ip: Option) -> ReverseProxy { ReverseProxy { client, - remote_ip_addr, + remote_ip, _pantom_data: PhantomData, } } @@ -111,15 +128,15 @@ impl ReverseProxy { *request.headers_mut() = remove_hop_headers(request.headers()); // Add forwarding information in the headers - if let Some(ip_addr) = self.remote_ip_addr { + if let Some(ip) = self.remote_ip { // This is kind of ugly because of borrowing. Maybe hyper's `Headers` object // could use an entry API like `std::collections::HashMap`? if request.headers().has::() { if let Some(prior) = request.headers_mut().get_mut::() { - prior.0.push(ip_addr); + prior.0.push(ip); } } else { - let header = XForwardedFor(vec![ip_addr]); + let header = XForwardedFor(vec![ip]); request.headers_mut().set(header); } } diff --git a/tests/lib.rs b/tests/lib.rs index 8dc8e2d..dba7f49 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -22,10 +22,57 @@ impl Response> Service for MockService { } #[test] -#[ignore] -fn adds_forwarded_for_header() { - // TODO: https://github.com/hyperium/hyper/issues/1258 - unimplemented!() +fn begins_forwarded_for_header() { + use hyper_reverse_proxy::XForwardedFor; + use std::net::Ipv6Addr; + + let mut request = Request::new(Get, "/".parse().unwrap()); + request.set_body("request"); + + let remote_ip = Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8); + let client = MockService(|request| { + assert_eq!( + request.headers().get::(), + Some(&XForwardedFor( + vec![Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8).into()], + )) + ); + + Response::new() + }); + let service = ReverseProxy::new(client, Some(remote_ip.into())); + + service.call(request).wait().unwrap(); +} + +#[test] +fn continues_forwarded_for_header() { + use hyper_reverse_proxy::XForwardedFor; + use std::net::{Ipv4Addr, Ipv6Addr}; + + let mut request = Request::new(Get, "/".parse().unwrap()); + request.set_body("request"); + request.headers_mut().set(XForwardedFor(vec![ + Ipv4Addr::new(127, 0, 0, 1).into(), + Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1).into(), + ])); + + let remote_ip = Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8); + let client = MockService(|request| { + assert_eq!( + request.headers().get::(), + Some(&XForwardedFor(vec![ + Ipv4Addr::new(127, 0, 0, 1).into(), + Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1).into(), + Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8).into(), + ])) + ); + + Response::new() + }); + let service = ReverseProxy::new(client, Some(remote_ip.into())); + + service.call(request).wait().unwrap(); } #[test] @@ -35,12 +82,13 @@ fn forwards_the_bodies() { let mut request = Request::new(Get, "/".parse().unwrap()); request.set_body("request"); - let service = ReverseProxy::new(MockService(|request| { + let client = MockService(|request| { let body = request.body().concat2().wait().unwrap(); assert_eq!(body.as_ref(), b"request"); Response::new().with_body("response") - })); + }); + let service = ReverseProxy::new(client, None); let response = service.call(request).wait().unwrap(); let body = response.body().concat2().wait().unwrap(); @@ -56,13 +104,14 @@ fn clones_headers() { request.headers_mut().set(XTestHeader1("Test1".to_owned())); request.headers_mut().set(XTestHeader2("Test2".to_owned())); - let service = ReverseProxy::new(MockService(|request| { + let client = MockService(|request| { let header1 = request.headers().get::().unwrap(); let header2 = request.headers().get::().unwrap(); assert_eq!(header1, &XTestHeader1("Test1".to_owned())); assert_eq!(header2, &XTestHeader2("Test2".to_owned())); Response::new() - })); + }); + let service = ReverseProxy::new(client, None); service.call(request).wait().unwrap(); } @@ -81,7 +130,7 @@ fn removes_request_hop_headers() { request.headers_mut().set(TransferEncoding(vec![])); request.headers_mut().set(Upgrade(vec![])); - let service = ReverseProxy::new(MockService(|request| { + let client = MockService(|request| { assert_eq!(request.headers().get::(), None); assert_eq!(request.headers().get_raw("Keep-Alive"), None); assert_eq!(request.headers().get_raw("Proxy-Authenticate"), None); @@ -91,7 +140,8 @@ fn removes_request_hop_headers() { assert_eq!(request.headers().get::(), None); assert_eq!(request.headers().get::(), None); Response::new() - })); + }); + let service = ReverseProxy::new(client, None); service.call(request).wait().unwrap(); } @@ -102,7 +152,7 @@ fn removes_response_hop_headers() { let request = Request::new(Get, "/".parse().unwrap()); - let service = ReverseProxy::new(MockService(|_| { + let client = MockService(|_| { let mut response = Response::new(); response.headers_mut().set(Connection(vec![])); response.headers_mut().set_raw("Keep-Alive", ""); @@ -113,7 +163,8 @@ fn removes_response_hop_headers() { response.headers_mut().set(TransferEncoding(vec![])); response.headers_mut().set(Upgrade(vec![])); response - })); + }); + let service = ReverseProxy::new(client, None); let response = service.call(request).wait().unwrap(); assert_eq!(response.headers().get::(), None);