jailer: make resource limits configurable

Signed-off-by: Luminita Voicu <[email protected]>

Author: luminitavoicu

src/jailer/src/env.rs

@@ -12,7 +12,7 @@ use std::process::{Command, Stdio};
 use crate::cgroup;
 use crate::cgroup::Cgroup;
 use crate::chroot::chroot;
-use crate::resource_limits::ResourceLimits;
+use crate::resource_limits::{ResourceLimits, FSIZE_ARG, NO_FILE_ARG};
 use crate::{Error, Result};
 use std::io;
 use std::io::Write;
@@ -169,7 +169,10 @@ impl Env {
             }
         }
 
-        let resource_limits = ResourceLimits::default();
+        let mut resource_limits = ResourceLimits::default();
+        if let Some(args) = arguments.multiple_values("resource-limit") {
+            Env::parse_resource_limits(&mut resource_limits, args)?;
+        }
 
         Ok(Env {
             id: id.to_owned(),
@@ -201,6 +204,24 @@ impl Env {
         self.uid
     }
 
+    fn parse_resource_limits(resource_limits: &mut ResourceLimits, args: &[String]) -> Result<()> {
+        for arg in args {
+            let (name, value) = arg
+                .split_once('=')
+                .ok_or_else(|| Error::ResLimitFormat(arg.to_string()))?;
+
+            let limit_value = value
+                .parse::<u64>()
+                .map_err(|err| Error::ResLimitValue(value.to_string(), err.to_string()))?;
+            match name {
+                FSIZE_ARG => resource_limits.set_file_size(limit_value),
+                NO_FILE_ARG => resource_limits.set_no_file(limit_value),
+                _ => return Err(Error::ResLimitArgument(name.to_string())),
+            }
+        }
+        Ok(())
+    }
+
     fn exec_into_new_pid_ns(&mut self, chroot_exec_file: PathBuf) -> Result<()> {
         // Unshare into a new PID namespace.
         // The current process will not be moved into the newly created namespace, but its first
@@ -562,6 +583,7 @@ mod tests {
         pub daemonize: bool,
         pub new_pid_ns: bool,
         pub cgroups: Vec<&'a str>,
+        pub resource_limits: Vec<&'a str>,
     }
 
     impl ArgVals<'_> {
@@ -577,6 +599,7 @@ mod tests {
                 daemonize: true,
                 new_pid_ns: true,
                 cgroups: vec!["cpu.shares=2", "cpuset.mems=0"],
+                resource_limits: vec!["no-file=1024", "fsize=1048575"],
             }
         }
     }
@@ -607,6 +630,12 @@ mod tests {
             arg_vec.push((*cg).to_string());
         }
 
+        // Append limits arguments
+        for limit in &arg_vals.resource_limits {
+            arg_vec.push("--resource-limit".to_string());
+            arg_vec.push((*limit).to_string());
+        }
+
         if let Some(s) = arg_vals.netns {
             arg_vec.push("--netns".to_string());
             arg_vec.push(s.to_string());
@@ -702,6 +731,16 @@ mod tests {
         args.parse(&make_args(&invalid_cgroup_arg_vals)).unwrap();
         assert!(Env::new(&args, 0, 0).is_err());
 
+        let invalid_res_limit_arg_vals = ArgVals {
+            resource_limits: vec!["zzz"],
+            ..base_invalid_arg_vals.clone()
+        };
+
+        let arg_parser = build_arg_parser();
+        args = arg_parser.arguments().clone();
+        args.parse(&make_args(&invalid_res_limit_arg_vals)).unwrap();
+        assert!(Env::new(&args, 0, 0).is_err());
+
         let invalid_id_arg_vals = ArgVals {
             id: "/ad./sa12",
             ..base_invalid_arg_vals.clone()
@@ -884,6 +923,7 @@ mod tests {
             daemonize: false,
             new_pid_ns: false,
             cgroups: Vec::new(),
+            resource_limits: Vec::new(),
         };
         fs::write(some_file_path, "some_content").unwrap();
         args.parse(&make_args(&some_arg_vals)).unwrap();
@@ -1006,6 +1046,71 @@ mod tests {
         assert!(Env::new(&args, 0, 0).is_ok());
     }
 
+    #[test]
+    fn test_parse_resource_limits() {
+        let mut resource_limits = ResourceLimits::default();
+
+        // Cases that should fail
+
+        // Check invalid formats
+        let invalid_formats = ["", "foo"];
+        for format in invalid_formats.iter() {
+            let arg = vec![format.to_string()];
+            assert_eq!(
+                format!(
+                    "{:?}",
+                    Env::parse_resource_limits(&mut resource_limits, &*arg)
+                        .err()
+                        .unwrap()
+                ),
+                format!("{:?}", Error::ResLimitFormat(format.to_string()))
+            );
+        }
+
+        // Check invalid resource arguments
+        let invalid_resources = ["foo", "", " "];
+        for res in invalid_resources.iter() {
+            let arg = format!("{}=2", res);
+            assert_eq!(
+                format!(
+                    "{:?}",
+                    Env::parse_resource_limits(&mut resource_limits, &*vec![arg])
+                        .err()
+                        .unwrap()
+                ),
+                format!("{:?}", Error::ResLimitArgument(res.to_string()))
+            );
+        }
+
+        // Check invalid limit values
+        let invalid_values = ["foo", "2.3", "2-3", " "];
+        for val in invalid_values.iter() {
+            let arg = format!("fsize={}", val);
+            assert_eq!(
+                format!(
+                    "{:?}",
+                    Env::parse_resource_limits(&mut resource_limits, &*vec![arg])
+                        .err()
+                        .unwrap()
+                ),
+                format!(
+                    "{:?}",
+                    Error::ResLimitValue(
+                        val.to_string(),
+                        "invalid digit found in string".to_string()
+                    )
+                )
+            );
+        }
+
+        // Check valid cases
+        let resources = [FSIZE_ARG, NO_FILE_ARG];
+        for resource in resources.iter() {
+            let arg = vec![resource.to_string() + &"=4098".to_string()];
+            Env::parse_resource_limits(&mut resource_limits, &*arg).unwrap();
+        }
+    }
+
     #[test]
     #[cfg(target_arch = "aarch64")]
     fn test_copy_cache_info() {

src/jailer/src/main.rs

@@ -57,6 +57,9 @@ pub enum Error {
     ReadLine(PathBuf, io::Error),
     ReadToString(PathBuf, io::Error),
     RegEx(regex::Error),
+    ResLimitArgument(String),
+    ResLimitFormat(String),
+    ResLimitValue(String, String),
     RmOldRootDir(io::Error),
     SetCurrentDir(io::Error),
     SetNetNs(io::Error),
@@ -191,6 +194,11 @@ impl fmt::Display for Error {
                 format!("Failed to read file {:?} into a string: {}", path, err).replace("\"", "")
             ),
             RegEx(ref err) => write!(f, "Regex failed: {:?}", err),
+            ResLimitArgument(ref arg) => write!(f, "Invalid resource argument: {}", arg,),
+            ResLimitFormat(ref arg) => write!(f, "Invalid format for resources limits: {}", arg,),
+            ResLimitValue(ref arg, ref err) => {
+                write!(f, "Invalid limit value for resource: {}: {}", arg, err)
+            }
             RmOldRootDir(ref err) => write!(f, "Failed to remove old jail root directory: {}", err),
             SetCurrentDir(ref err) => write!(f, "Failed to change current directory: {}", err),
             SetNetNs(ref err) => write!(f, "Failed to join network namespace: netns: {}", err),
@@ -282,6 +290,15 @@ pub fn build_arg_parser() -> ArgParser<'static> {
              <cgroup_file>=<value> (e.g cpu.shares=10). This argument can be used \
              multiple times to add multiple cgroups.",
         ))
+        .arg(Argument::new("resource-limit").allow_multiple(true).help(
+            "Resource limit values to be set by the jailer. It must follow this format: \
+             <resource>=<value> (e.g no-file=1024). This argument can be used \
+             multiple times to add multiple resource limits. \
+             Current available resource values are:\n\
+             \t\tfsize: The maximum size in bytes for files created by the process.\n\
+             \t\tno-file: Specifies a value one greater than the maximum file descriptor number \
+             that can be opened by this process.",
+        ))
         .arg(
             Argument::new("version")
                 .takes_value(false)
@@ -645,6 +662,21 @@ mod tests {
             format!("{}", Error::RegEx(err_regex.clone())),
             format!("Regex failed: {:?}", err_regex),
         );
+        assert_eq!(
+            format!("{}", Error::ResLimitArgument("foo".to_string())),
+            "Invalid resource argument: foo",
+        );
+        assert_eq!(
+            format!("{}", Error::ResLimitFormat("foo".to_string())),
+            "Invalid format for resources limits: foo",
+        );
+        assert_eq!(
+            format!(
+                "{}",
+                Error::ResLimitValue("foo".to_string(), "bar".to_string())
+            ),
+            "Invalid limit value for resource: foo: bar",
+        );
         assert_eq!(
             format!("{}", Error::RmOldRootDir(io::Error::from_raw_os_error(42))),
             "Failed to remove old jail root directory: No message of desired type (os error 42)",

src/jailer/src/resource_limits.rs

@@ -1,14 +1,17 @@
 // Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-#![allow(unused)]
 use super::{Error, Result};
 use std::fmt;
 use std::fmt::{Display, Formatter};
 use utils::syscall::SyscallReturnCode;
 
 // Default limit for the maximum number of file descriptors open at a time.
 const NO_FILE: u64 = 2048;
+// File size resource argument name.
+pub(crate) const FSIZE_ARG: &str = "fsize";
+// Number of files resource argument name.
+pub(crate) const NO_FILE_ARG: &str = "no-file";
 
 #[derive(Clone, Copy)]
 pub enum Resource {