address comments

Signed-off-by: tison <wander4096@gmail.com>
GreptimeTeam · Mar 13, 2024 · 44dfd2a · 44dfd2a
1 parent 158d1c1
commit 44dfd2a
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 27 deletions.
diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs
@@ -281,8 +281,22 @@ pub enum Error {
     #[snafu(display("Unsupported http auth scheme, name: {}", name))]
     UnsupportedAuthScheme { name: String },
 
+    #[snafu(display("Invalid visibility ASCII chars"))]
+    InvalidAuthHeaderInvisibleASCII {
+        #[snafu(source)]
+        error: hyper::header::ToStrError,
+        location: Location,
+    },
+
+    #[snafu(display("Invalid utf-8 value"))]
+    InvalidAuthHeaderInvalidUtf8Value {
+        #[snafu(source)]
+        error: FromUtf8Error,
+        location: Location,
+    },
+
     #[snafu(display("Invalid http authorization header"))]
-    InvalidAuthorizationHeader { location: Location },
+    InvalidAuthHeader { location: Location },
 
     #[snafu(display("Invalid base64 value"))]
     InvalidBase64Value {
@@ -527,16 +541,20 @@ impl ErrorExt for Error {
             DescribeStatement { source } => source.status_code(),
 
             NotFoundAuthHeader { .. } | NotFoundInfluxAuth { .. } => StatusCode::AuthHeaderNotFound,
-            InvisibleASCII { .. }
+            InvalidAuthHeaderInvisibleASCII { .. }
             | UnsupportedAuthScheme { .. }
-            | InvalidAuthorizationHeader { .. }
+            | InvalidAuthHeader { .. }
             | InvalidBase64Value { .. }
-            | InvalidUtf8Value { .. } => StatusCode::InvalidAuthHeader,
+            | InvalidAuthHeaderInvalidUtf8Value { .. } => StatusCode::InvalidAuthHeader,
 
             DatabaseNotFound { .. } => StatusCode::DatabaseNotFound,
             #[cfg(feature = "mem-prof")]
             DumpProfileData { source, .. } => source.status_code(),
-            InvalidFlushArgument { .. } | InvalidGzip { .. } => StatusCode::InvalidArguments,
+
+            InvalidUtf8Value { .. }
+            | InvisibleASCII { .. }
+            | InvalidFlushArgument { .. }
+            | InvalidGzip { .. } => StatusCode::InvalidArguments,
 
             ReplacePreparedStmtParams { source, .. }
             | GetPreparedStmtParams { source, .. }

diff --git a/src/servers/src/http/authorize.rs b/src/servers/src/http/authorize.rs
@@ -35,7 +35,7 @@ use snafu::{ensure, OptionExt, ResultExt};
 use super::header::{GreptimeDbName, GREPTIME_TIMEZONE_HEADER_NAME};
 use super::{ResponseFormat, PUBLIC_APIS};
 use crate::error::{
-    self, InvalidAuthorizationHeaderSnafu, InvalidParameterSnafu, InvisibleASCIISnafu,
+    self, InvalidAuthHeaderInvisibleASCIISnafu, InvalidAuthHeaderSnafu, InvalidParameterSnafu,
     NotFoundInfluxAuthSnafu, Result, UnsupportedAuthSchemeSnafu, UrlDecodeSnafu,
 };
 use crate::http::error_result::ErrorResponse;
@@ -174,15 +174,13 @@ fn get_influxdb_credentials<B>(request: &Request<B>) -> Result<Option<(Username,
         // try header
         let (auth_scheme, credential) = header
             .to_str()
-            .context(InvisibleASCIISnafu)?
+            .context(InvalidAuthHeaderInvisibleASCIISnafu)?
             .split_once(' ')
-            .context(InvalidAuthorizationHeaderSnafu)?;
+            .context(InvalidAuthHeaderSnafu)?;
 
         let (username, password) = match auth_scheme.to_lowercase().as_str() {
             "token" => {
-                let (u, p) = credential
-                    .split_once(':')
-                    .context(InvalidAuthorizationHeaderSnafu)?;
+                let (u, p) = credential.split_once(':').context(InvalidAuthHeaderSnafu)?;
                 (u.to_string(), p.to_string().into())
             }
             "basic" => decode_basic(credential)?,
@@ -237,13 +235,10 @@ impl TryFrom<&str> for AuthScheme {
     type Error = error::Error;
 
     fn try_from(value: &str) -> Result<Self> {
-        let (scheme, encoded_credentials) = value
-            .split_once(' ')
-            .context(InvalidAuthorizationHeaderSnafu)?;
-        ensure!(
-            !encoded_credentials.contains(' '),
-            InvalidAuthorizationHeaderSnafu
-        );
+        let (scheme, encoded_credentials) =
+            value.split_once(' ').context(InvalidAuthHeaderSnafu)?;
+
+        ensure!(!encoded_credentials.contains(' '), InvalidAuthHeaderSnafu);
 
         match scheme.to_lowercase().as_str() {
             "basic" => decode_basic(encoded_credentials)
@@ -261,7 +256,7 @@ fn auth_header<B>(req: &Request<B>) -> Result<AuthScheme> {
         .get(http::header::AUTHORIZATION)
         .context(error::NotFoundAuthHeaderSnafu)?
         .to_str()
-        .context(InvisibleASCIISnafu)?;
+        .context(InvalidAuthHeaderInvisibleASCIISnafu)?;
 
     auth_header.try_into()
 }
@@ -270,13 +265,14 @@ fn decode_basic(credential: Credential) -> Result<(Username, Password)> {
     let decoded = BASE64_STANDARD
         .decode(credential)
         .context(error::InvalidBase64ValueSnafu)?;
-    let as_utf8 = String::from_utf8(decoded).context(error::InvalidUtf8ValueSnafu)?;
+    let as_utf8 =
+        String::from_utf8(decoded).context(error::InvalidAuthHeaderInvalidUtf8ValueSnafu)?;
 
     if let Some((user_id, password)) = as_utf8.split_once(':') {
         return Ok((user_id.to_string(), password.to_string().into()));
     }
 
-    InvalidAuthorizationHeaderSnafu {}.fail()
+    InvalidAuthHeaderSnafu {}.fail()
 }
 
 fn need_auth<B>(req: &Request<B>) -> bool {
@@ -395,10 +391,7 @@ mod tests {
 
         let wrong_req = mock_http_request(Some("Basic dXNlcm5hbWU6 cGFzc3dvcmQ="), None).unwrap();
         let res = auth_header(&wrong_req);
-        assert_matches!(
-            res.err(),
-            Some(error::Error::InvalidAuthorizationHeader { .. })
-        );
+        assert_matches!(res.err(), Some(error::Error::InvalidAuthHeader { .. }));
 
         let wrong_req = mock_http_request(Some("Digest dXNlcm5hbWU6cGFzc3dvcmQ="), None).unwrap();
         let res = auth_header(&wrong_req);

diff --git a/src/servers/src/http/influxdb.rs b/src/servers/src/http/influxdb.rs
@@ -15,7 +15,7 @@
 use std::collections::HashMap;
 
 use axum::extract::{Query, State};
-use axum::http::{HeaderMap, StatusCode};
+use axum::http::{header, HeaderMap, StatusCode};
 use axum::response::IntoResponse;
 use axum::Extension;
 use bytes::Bytes;
@@ -46,9 +46,11 @@ pub async fn influxdb_health() -> Result<impl IntoResponse> {
 
 fn is_gzip(headers: &HeaderMap) -> Result<bool> {
     match headers
-        .get("Content-Encoding")
+        .get(header::CONTENT_ENCODING)
         .map(|val| val.to_str().context(InvisibleASCIISnafu))
         .transpose()?
+        .map(|val| val.to_lowercase())
+        .as_deref()
     {
         None | Some("identity") => Ok(false),
         Some("gzip") => Ok(true),