From c605f6852d8138d5d8deb3ff7b1888c00e1e4349 Mon Sep 17 00:00:00 2001
From: MegaBrutal <code+git@megabrutal.com>
Date: Sat, 24 Dec 2022 00:45:00 +0100
Subject: [PATCH] Stricten color specification

---
 src/main.rs | 39 +++++++++++++++++++++++++++++++++++----
 1 file changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index f621684..db66481 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -37,6 +37,21 @@ macro_rules! extract_data {
     }
 }
 
+macro_rules! extract_color {
+    ($req:expr,$n:expr) => {
+        {
+            let data_raw = extract_data!($req, $n);
+            let mut data_decoded = percent_decode_str(data_raw);
+            let mut rgb = ToRgbIter::new(&mut data_decoded);
+            let color = rgb.next().ok_or(ColorError(data_raw.to_owned()))?;
+            match rgb.next() {
+                Some(_) => Err(ColorError(data_raw.to_owned())),
+                _ => Ok(color)
+            }
+        }
+    }
+}
+
 macro_rules! extract_header {
     ($req:expr,$header:expr) => {
         match $req.headers().get($header) {
@@ -46,6 +61,25 @@ macro_rules! extract_header {
     }
 }
 
+#[derive(Debug)]
+struct ColorError(String);
+
+impl Display for ColorError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Invalid color specification: {}", &self.0)
+    }
+}
+
+impl ResponseError for ColorError {
+    fn status_code(&self) -> StatusCode {
+        StatusCode::INTERNAL_SERVER_ERROR
+    }
+
+    fn error_response(&self) -> HttpResponse<BoxBody> {
+        HttpResponse::InternalServerError().body(format!("{}\n", &self))
+    }
+}
+
 #[derive(Debug)]
 struct ImageError(image::ImageError);
 
@@ -331,10 +365,7 @@ async fn img_gen3(req: HttpRequest, path: web::Path<u32>) -> Result<impl Respond
 async fn img_gen4(req: HttpRequest, path: web::Path<(u32, u32, u32, u32, u32, f32, f32, u32)>) -> Result<impl Responder> {
     let (dim_x, dim_y, scale, wave_offset_x, wave_offset_y, amplitude, freq_base, run) = path.into_inner();
     let freq_base = freq_base / 1000.0;
-    let wave_color = extract_data!(req, 11);
-    debug!("Wave color input: {:?}", wave_color);
-    let wave_color = ToRgbIter::new(&mut percent_decode_str(wave_color)).next().unwrap_or(Rgb([255, 255, 255]));
-    debug!("Decoded color: {:?}", wave_color);
+    let wave_color = extract_color!(req, 11)?;
     let signal = &mut percent_decode_str(extract_data!(req, 12));
     let cursor = make_png(dim_x, dim_y, scale, signal,
         |img, signal| sine_encode(img, (wave_offset_x, wave_offset_y), amplitude, freq_base, run, wave_color, signal))?;
-- 
2.34.1