third_party/rust/glean-core/src/upload/request.rs

   1 // This Source Code Form is subject to the terms of the Mozilla Public
   2 // License, v. 2.0. If a copy of the MPL was not distributed with this
   3 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
   4
   5 //! Ping request representation.
   6
   7 use std::collections::HashMap;
   8
   9 use chrono::prelude::{DateTime, Utc};
  10 use flate2::{read::GzDecoder, write::GzEncoder, Compression};
  11 use serde_json::{self, Value as JsonValue};
  12 use std::io::prelude::*;
  13
  14 use crate::error::{ErrorKind, Result};
  15 use crate::system;
  16
  17 /// A representation for request headers.
  18 pub type HeaderMap = HashMap<String, String>;
  19
  20 /// Creates a formatted date string that can be used with Date headers.
  21 pub(crate) fn create_date_header_value(current_time: DateTime<Utc>) -> String {
  22     // Date headers are required to be in the following format:
  23     //
  24     // <day-name>, <day> <month> <year> <hour>:<minute>:<second> GMT
  25     //
  26     // as documented here:
  27     // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Date
  28     // Unfortunately we can't use `current_time.to_rfc2822()` as it
  29     // formats as "Mon, 22 Jun 2020 10:40:34 +0000", with an ending
  30     // "+0000" instead of "GMT". That's why we need to go with manual
  31     // formatting.
  32     current_time.format("%a, %d %b %Y %T GMT").to_string()
  33 }
  34
  35 fn create_x_telemetry_agent_header_value(
  36     version: &str,
  37     language_binding_name: &str,
  38     system: &str,
  39 ) -> String {
  40     format!(
  41         "Glean/{} ({} on {})",
  42         version, language_binding_name, system
  43     )
  44 }
  45
  46 /// Attempt to gzip the contents of a ping.
  47 fn gzip_content(path: &str, content: &[u8]) -> Option<Vec<u8>> {
  48     let mut gzipper = GzEncoder::new(Vec::new(), Compression::default());
  49
  50     // Attempt to add the content to the gzipper.
  51     if let Err(e) = gzipper.write_all(content) {
  52         log::warn!("Failed to write to the gzipper: {} - {:?}", path, e);
  53         return None;
  54     }
  55
  56     gzipper.finish().ok()
  57 }
  58
  59 pub struct Builder {
  60     document_id: Option<String>,
  61     path: Option<String>,
  62     body: Option<Vec<u8>>,
  63     headers: HeaderMap,
  64     body_max_size: usize,
  65     body_has_info_sections: Option<bool>,
  66     ping_name: Option<String>,
  67 }
  68
  69 impl Builder {
  70     /// Creates a new builder for a PingRequest.
  71     pub fn new(language_binding_name: &str, body_max_size: usize) -> Self {
  72         let mut headers = HashMap::new();
  73         headers.insert(
  74             "X-Telemetry-Agent".to_string(),
  75             create_x_telemetry_agent_header_value(
  76                 crate::GLEAN_VERSION,
  77                 language_binding_name,
  78                 system::OS,
  79             ),
  80         );
  81         headers.insert(
  82             "Content-Type".to_string(),
  83             "application/json; charset=utf-8".to_string(),
  84         );
  85
  86         Self {
  87             document_id: None,
  88             path: None,
  89             body: None,
  90             headers,
  91             body_max_size,
  92             body_has_info_sections: None,
  93             ping_name: None,
  94         }
  95     }
  96
  97     /// Sets the document_id for this request.
  98     pub fn document_id<S: Into<String>>(mut self, value: S) -> Self {
  99         self.document_id = Some(value.into());
 100         self
 101     }
 102
 103     /// Sets the path for this request.
 104     pub fn path<S: Into<String>>(mut self, value: S) -> Self {
 105         self.path = Some(value.into());
 106         self
 107     }
 108
 109     /// Sets the body for this request.
 110     ///
 111     /// This method will also attempt to gzip the body contents
 112     /// and add headers related to the body that was just added.
 113     ///
 114     /// Namely these headers are the "Content-Length" with the length of the body
 115     /// and in case we are successfull on gzipping the contents, the "Content-Encoding"="gzip".
 116     ///
 117     /// **Important**
 118     /// If we are unable to gzip we don't panic and instead just set the uncompressed body.
 119     ///
 120     /// # Panics
 121     ///
 122     /// This method will panic in case we try to set the body before setting the path.
 123     pub fn body<S: Into<String>>(mut self, value: S) -> Self {
 124         // Attempt to gzip the body contents.
 125         let original_as_string = value.into();
 126         let gzipped_content = gzip_content(
 127             self.path
 128                 .as_ref()
 129                 .expect("Path must be set before attempting to set the body"),
 130             original_as_string.as_bytes(),
 131         );
 132         let add_gzip_header = gzipped_content.is_some();
 133         let body = gzipped_content.unwrap_or_else(|| original_as_string.into_bytes());
 134
 135         // Include headers related to body
 136         self = self.header("Content-Length", &body.len().to_string());
 137         if add_gzip_header {
 138             self = self.header("Content-Encoding", "gzip");
 139         }
 140
 141         self.body = Some(body);
 142         self
 143     }
 144
 145     /// Sets whether the request body has {client|ping}_info sections.
 146     pub fn body_has_info_sections(mut self, body_has_info_sections: bool) -> Self {
 147         self.body_has_info_sections = Some(body_has_info_sections);
 148         self
 149     }
 150
 151     /// Sets the ping's name aka doctype.
 152     pub fn ping_name<S: Into<String>>(mut self, ping_name: S) -> Self {
 153         self.ping_name = Some(ping_name.into());
 154         self
 155     }
 156
 157     /// Sets a header for this request.
 158     pub fn header<S: Into<String>>(mut self, key: S, value: S) -> Self {
 159         self.headers.insert(key.into(), value.into());
 160         self
 161     }
 162
 163     /// Sets multiple headers for this request at once.
 164     pub fn headers(mut self, values: HeaderMap) -> Self {
 165         self.headers.extend(values);
 166         self
 167     }
 168
 169     /// Consumes the builder and create a PingRequest.
 170     ///
 171     /// # Panics
 172     ///
 173     /// This method will panic if any of the required fields are missing:
 174     /// `document_id`, `path` and `body`.
 175     pub fn build(self) -> Result<PingRequest> {
 176         let body = self
 177             .body
 178             .expect("body must be set before attempting to build PingRequest");
 179
 180         if body.len() > self.body_max_size {
 181             return Err(ErrorKind::PingBodyOverflow(body.len()).into());
 182         }
 183
 184         Ok(PingRequest {
 185             document_id: self
 186                 .document_id
 187                 .expect("document_id must be set before attempting to build PingRequest"),
 188             path: self
 189                 .path
 190                 .expect("path must be set before attempting to build PingRequest"),
 191             body,
 192             headers: self.headers,
 193             body_has_info_sections: self.body_has_info_sections.expect(
 194                 "body_has_info_sections must be set before attempting to build PingRequest",
 195             ),
 196             ping_name: self
 197                 .ping_name
 198                 .expect("ping_name must be set before attempting to build PingRequest"),
 199         })
 200     }
 201 }
 202
 203 /// Represents a request to upload a ping.
 204 #[derive(PartialEq, Eq, Debug, Clone)]
 205 pub struct PingRequest {
 206     /// The Job ID to identify this request,
 207     /// this is the same as the ping UUID.
 208     pub document_id: String,
 209     /// The path for the server to upload the ping to.
 210     pub path: String,
 211     /// The body of the request, as a byte array. If gzip encoded, then
 212     /// the `headers` list will contain a `Content-Encoding` header with
 213     /// the value `gzip`.
 214     pub body: Vec<u8>,
 215     /// A map with all the headers to be sent with the request.
 216     pub headers: HeaderMap,
 217     /// Whether the body has {client|ping}_info sections.
 218     pub body_has_info_sections: bool,
 219     /// The ping's name. Likely also somewhere in `path`.
 220     pub ping_name: String,
 221 }
 222
 223 impl PingRequest {
 224     /// Creates a new builder-style structure to help build a PingRequest.
 225     ///
 226     /// # Arguments
 227     ///
 228     /// * `language_binding_name` - The name of the language used by the binding that instantiated this Glean instance.
 229     ///                             This is used to build the X-Telemetry-Agent header value.
 230     /// * `body_max_size` - The maximum size in bytes the compressed ping body may have to be eligible for upload.
 231     pub fn builder(language_binding_name: &str, body_max_size: usize) -> Builder {
 232         Builder::new(language_binding_name, body_max_size)
 233     }
 234
 235     /// Verifies if current request is for a deletion-request ping.
 236     pub fn is_deletion_request(&self) -> bool {
 237         self.ping_name == "deletion-request"
 238     }
 239
 240     /// Decompresses and pretty-format the ping payload
 241     ///
 242     /// Should be used for logging when required.
 243     /// This decompresses the payload in memory.
 244     pub fn pretty_body(&self) -> Option<String> {
 245         let mut gz = GzDecoder::new(&self.body[..]);
 246         let mut s = String::with_capacity(self.body.len());
 247
 248         gz.read_to_string(&mut s)
 249             .ok()
 250             .map(|_| &s[..])
 251             .or_else(|| std::str::from_utf8(&self.body).ok())
 252             .and_then(|payload| serde_json::from_str::<JsonValue>(payload).ok())
 253             .and_then(|json| serde_json::to_string_pretty(&json).ok())
 254     }
 255 }
 256
 257 #[cfg(test)]
 258 mod test {
 259     use super::*;
 260     use chrono::offset::TimeZone;
 261
 262     #[test]
 263     fn date_header_resolution() {
 264         let date: DateTime<Utc> = Utc.ymd(2018, 2, 25).and_hms(11, 10, 37);
 265         let test_value = create_date_header_value(date);
 266         assert_eq!("Sun, 25 Feb 2018 11:10:37 GMT", test_value);
 267     }
 268
 269     #[test]
 270     fn x_telemetry_agent_header_resolution() {
 271         let test_value = create_x_telemetry_agent_header_value("0.0.0", "Rust", "Windows");
 272         assert_eq!("Glean/0.0.0 (Rust on Windows)", test_value);
 273     }
 274
 275     #[test]
 276     fn correctly_builds_ping_request() {
 277         let request = PingRequest::builder(/* language_binding_name */ "Rust", 1024 * 1024)
 278             .document_id("woop")
 279             .path("/random/path/doesnt/matter")
 280             .body("{}")
 281             .body_has_info_sections(false)
 282             .ping_name("whatevs")
 283             .build()
 284             .unwrap();
 285
 286         assert_eq!(request.document_id, "woop");
 287         assert_eq!(request.path, "/random/path/doesnt/matter");
 288         assert!(!request.body_has_info_sections);
 289         assert_eq!(request.ping_name, "whatevs");
 290
 291         // Make sure all the expected headers were added.
 292         assert!(request.headers.contains_key("X-Telemetry-Agent"));
 293         assert!(request.headers.contains_key("Content-Type"));
 294         assert!(request.headers.contains_key("Content-Length"));
 295
 296         // the `Date` header is added by the `get_upload_task` just before
 297         // returning the upload request
 298     }
 299
 300     #[test]
 301     fn errors_when_request_body_exceeds_max_size() {
 302         // Create a new builder with an arbitrarily small value,
 303         // se we can test that the builder errors when body max size exceeds the expected.
 304         let request = Builder::new(
 305             /* language_binding_name */ "Rust", /* body_max_size */ 1,
 306         )
 307         .document_id("woop")
 308         .path("/random/path/doesnt/matter")
 309         .body("{}")
 310         .build();
 311
 312         assert!(request.is_err());
 313     }
 314 }