Can restrict updates by ip

src/main.rs

@@ -38,6 +38,10 @@ struct Args {
     /// Handle for not found files. By default, a basic message is returned
     #[clap(short, long, env, default_value = "")]
     not_found_file: String,
+
+    /// Optional proxy IP
+    #[clap(short, long, env)]
+    proxy_ip: Option<String>,
 }
 
 impl Args {
@@ -60,9 +64,57 @@ struct NewFile {
     bytes: Vec<u8>,
 }
 
+// Check if two ips matches
+pub fn match_ip(pattern: &str, ip: &str) -> bool {
+    if pattern.eq(ip) {
+        return true;
+    }
+
+    if pattern.ends_with('*') && ip.starts_with(&pattern.replace('*', "")) {
+        return true;
+    }
+
+    false
+}
+
+
+/// Get the remote IP address
+fn get_remote_ip(req: &HttpRequest, args: &Args) -> String {
+    let mut ip = req.peer_addr().unwrap().ip().to_string();
+
+    // We check if the request comes from a trusted reverse proxy
+    if let Some(proxy) = args.proxy_ip.as_ref() {
+        if match_ip(proxy, &ip) {
+            if let Some(header) = req.headers().get("X-Forwarded-For") {
+                let header: Vec<String> = header
+                    .to_str()
+                    .unwrap()
+                    .to_string()
+                    .split(",")
+                    .map(|f| f.to_string())
+                    .collect();
+
+                if header.len() > 0 {
+                    ip = header[0].to_string();
+                }
+            }
+        }
+    }
+
+    ip
+}
+
+
 /// Replace all the files of the website
 async fn replace_files(args: web::Data<Args>, req: HttpRequest, mut payload: Multipart) -> Result<HttpResponse, Error> {
 
+    // Validate remote IP
+    let remote_ip = get_remote_ip(&req, &args);
+    if !match_ip(&args.allowed_ips_for_update, &remote_ip) {
+        log::warn!("Block unauthorized attempt to perform site update from {}", remote_ip);
+        return Err(ErrorUnauthorized("You are not allowed to perform updates!"));
+    }
+
     // Check token
     let token = match req.headers().get("Token") {
         None => {