From 9cb16ba01d86f14ffab33db661f1b41e347cd7b1 Mon Sep 17 00:00:00 2001
From: Yong Wen Chua <lawliet89@users.noreply.github.com>
Date: Thu, 1 Nov 2018 14:44:28 +0800
Subject: [PATCH] Use request local cache for fairing implementation (#50)

---
 src/fairing.rs | 56 ++++++++------------------------------------------
 src/lib.rs     | 10 +--------
 2 files changed, 10 insertions(+), 56 deletions(-)

diff --git a/src/fairing.rs b/src/fairing.rs
index c33dd6b..6277e26 100644
--- a/src/fairing.rs
+++ b/src/fairing.rs
@@ -1,47 +1,19 @@
 //! Fairing implementation
-use std::str::FromStr;
 
 use log::{error, info, log};
-use rocket::http::{self, uri::Origin, Header, Status};
+use rocket::http::{self, uri::Origin, Status};
 use rocket::{self, error_, info_, log_, Outcome, Request};
 
 use crate::{
     actual_request_response, origin, preflight_response, request_headers, validate, Cors, Error,
 };
 
-/// An injected header to quickly give the result of CORS
-static CORS_HEADER: &str = "ROCKET-CORS";
-enum InjectedHeader {
+/// Request Local State to store CORS validation results
+enum CorsValidation {
     Success,
     Failure,
 }
 
-impl InjectedHeader {
-    fn to_str(&self) -> &'static str {
-        match *self {
-            InjectedHeader::Success => "Success",
-            InjectedHeader::Failure => "Failure",
-        }
-    }
-}
-
-impl FromStr for InjectedHeader {
-    type Err = Error;
-    fn from_str(s: &str) -> Result<Self, Error> {
-        match s {
-            "Success" => Ok(InjectedHeader::Success),
-            "Failure" => Ok(InjectedHeader::Failure),
-            other => {
-                error_!(
-                    "Unknown injected header encountered: {}\nThis is probably a bug.",
-                    other
-                );
-                Err(Error::UnknownInjectedHeader)
-            }
-        }
-    }
-}
-
 /// Route for Fairing error handling
 pub(crate) fn fairing_error_route<'r>(
     request: &'r Request,
@@ -71,11 +43,6 @@ fn route_to_fairing_error_handler(options: &Cors, status: u16, request: &mut Req
     request.set_uri(origin);
 }
 
-/// Inject a header into the Request with result
-fn inject_request_header(header: &InjectedHeader, request: &mut Request) {
-    request.replace_header(Header::new(CORS_HEADER, header.to_str()));
-}
-
 fn on_response_wrapper(
     options: &Cors,
     request: &Request,
@@ -89,14 +56,9 @@ fn on_response_wrapper(
         Some(origin) => origin,
     };
 
-    // Get validation result from injected header
-    let injected_header = request
-        .headers()
-        .get_one(CORS_HEADER)
-        .ok_or_else(|| Error::MissingInjectedHeader)?;
-    let result = InjectedHeader::from_str(injected_header)?;
+    let result = request.local_cache(|| unreachable!("This should not be executed so late"));
 
-    if let InjectedHeader::Failure = result {
+    if let &CorsValidation::Failure = result {
         // Nothing else for us to do
         return Ok(());
     }
@@ -151,17 +113,17 @@ impl rocket::fairing::Fairing for Cors {
     }
 
     fn on_request(&self, request: &mut Request, _: &rocket::Data) {
-        let injected_header = match validate(self, request) {
-            Ok(_) => InjectedHeader::Success,
+        let result = match validate(self, request) {
+            Ok(_) => CorsValidation::Success,
             Err(err) => {
                 error_!("CORS Error: {}", err);
                 let status = err.status();
                 route_to_fairing_error_handler(self, status.code, request);
-                InjectedHeader::Failure
+                CorsValidation::Failure
             }
         };
 
-        inject_request_header(&injected_header, request);
+        let _ = request.local_cache(|| result);
     }
 
     fn on_response(&self, request: &Request, response: &mut rocket::Response) {
diff --git a/src/lib.rs b/src/lib.rs
index 4d56907..57c8529 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -637,9 +637,6 @@ pub enum Error {
     /// The `on_response` handler of Fairing could not find the injected header from the Request.
     /// Either some other fairing has removed it, or this is a bug.
     MissingInjectedHeader,
-    /// The `on_response` handler of Fairing found an unknown injected header value from the
-    /// Request. Either some other fairing has modified it, or this is a bug.
-    UnknownInjectedHeader,
 }
 
 impl Error {
@@ -651,8 +648,7 @@ impl Error {
             | Error::HeadersNotAllowed => Status::Forbidden,
             Error::CredentialsWithWildcardOrigin
             | Error::MissingCorsInRocketState
-            | Error::MissingInjectedHeader
-            | Error::UnknownInjectedHeader => Status::InternalServerError,
+            | Error::MissingInjectedHeader => Status::InternalServerError,
             _ => Status::BadRequest,
         }
     }
@@ -688,10 +684,6 @@ impl error::Error for Error {
                 "The `on_response` handler of Fairing could not find the injected header from the \
                  Request. Either some other fairing has removed it, or this is a bug."
             }
-            Error::UnknownInjectedHeader => {
-                "The `on_response` handler of Fairing found an unknown injected header value from \
-                 the Request. Either some other fairing has modified it, or this is a bug."
-            }
         }
     }