Add support for ulimit flag (#291)

* Barebones extension
* Add expected format to help message
* Add tests for ulimit extension

---------

Co-authored-by: Tully Foote <[email protected]>

Author: fpadula

setup.py

@@ -43,7 +43,7 @@
         'console_scripts': [
             'rocker = rocker.cli:main',
             'detect_docker_image_os = rocker.cli:detect_image_os',
-	    ],
+        ],
         'rocker.extensions': [
             'cuda = rocker.nvidia_extension:Cuda',
             'devices = rocker.extensions:Devices',
@@ -63,11 +63,12 @@
             'pulse = rocker.extensions:PulseAudio',
             'rmw = rocker.rmw_extension:RMW',
             'ssh = rocker.ssh_extension:Ssh',
+            'ulimit = rocker.ulimit_extension:Ulimit',
             'user = rocker.extensions:User',
             'volume = rocker.volume_extension:Volume',
             'x11 = rocker.nvidia_extension:X11',
         ]
-	},
+    },
     'author': 'Tully Foote',
     'author_email': '[email protected]',
     'keywords': ['Docker'],
@@ -91,4 +92,3 @@
 }
 
 setup(**kwargs)
-

src/rocker/ulimit_extension.py

@@ -0,0 +1,67 @@
+# Copyright 2019 Open Source Robotics Foundation
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from argparse import ArgumentTypeError
+import re
+from rocker.extensions import RockerExtension, name_to_argument
+
+
+class Ulimit(RockerExtension):
+    """
+    A RockerExtension to handle ulimit settings for Docker containers.
+
+    This extension allows specifying ulimit options in the format TYPE=SOFT_LIMIT[:HARD_LIMIT]
+    and validates the format before passing them as Docker arguments.
+    """
+    EXPECTED_FORMAT = "TYPE=SOFT_LIMIT[:HARD_LIMIT]"
+
+    @staticmethod
+    def get_name():
+        return 'ulimit'
+
+    def get_docker_args(self, cliargs):
+        args = ['']
+        ulimits = [x for sublist in cliargs[Ulimit.get_name()] for x in sublist]
+        for ulimit in ulimits:
+            if self.arg_format_is_valid(ulimit):
+                args.append(f"--ulimit {ulimit}")
+            else:
+                raise ArgumentTypeError(
+                    f"Error processing {Ulimit.get_name()} flag '{ulimit}': expected format"
+                    f" {Ulimit.EXPECTED_FORMAT}")
+        return ' '.join(args)
+
+    def arg_format_is_valid(self, arg: str):
+        """
+        Validate the format of the ulimit argument.
+
+        Args:
+            arg (str): The ulimit argument to validate.
+
+        Returns:
+            bool: True if the format is valid, False otherwise.
+        """
+        ulimit_format = r'(\w+)=(\w+)(:\w+)?$'
+        match = re.match(ulimit_format, arg)
+        return match is not None
+
+    @staticmethod
+    def register_arguments(parser, defaults):
+        parser.add_argument(name_to_argument(Ulimit.get_name()),
+                            type=str,
+                            nargs='+',
+                            action='append',
+                            metavar=Ulimit.EXPECTED_FORMAT,
+                            default=defaults.get(Ulimit.get_name(), None),
+                            help='ulimit options to add into the container.')

test/test_ulimit.py

@@ -0,0 +1,100 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import unittest
+from argparse import ArgumentTypeError
+
+from rocker.ulimit_extension import Ulimit
+
+
+class UlimitTest(unittest.TestCase):
+    """Unit tests for the Ulimit class."""
+
+    def setUp(self):
+        self._instance = Ulimit()
+
+    def _is_arg_translation_ok(self, mock_cliargs, expected):
+        is_ok = False
+        message_string = ""
+        try:
+            docker_args = self._instance.get_docker_args(
+                {self._instance.get_name(): [mock_cliargs]})
+            is_ok = docker_args == expected
+            message_string = f"Expected: '{expected}', got: '{docker_args}'"
+        except ArgumentTypeError:
+            message_string = "Incorrect argument format"
+        return (is_ok, message_string)
+
+    def test_args_single_soft(self):
+        """Test single soft limit argument."""
+        mock_cliargs = ["rtprio=99"]
+        expected = " --ulimit rtprio=99"
+        self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_multiple_soft(self):
+        """Test multiple soft limit arguments."""
+        mock_cliargs = ["rtprio=99", "memlock=102400"]
+        expected = " --ulimit rtprio=99 --ulimit memlock=102400"
+        self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_single_hard(self):
+        """Test single hard limit argument."""
+        mock_cliargs = ["nofile=1024:524288"]
+        expected = " --ulimit nofile=1024:524288"
+        self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_multiple_hard(self):
+        """Test multiple hard limit arguments."""
+        mock_cliargs = ["nofile=1024:524288", "rtprio=90:99"]
+        expected = " --ulimit nofile=1024:524288 --ulimit rtprio=90:99"
+        self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_multiple_mix(self):
+        """Test multiple mixed limit arguments."""
+        mock_cliargs = ["rtprio=99", "memlock=102400", "nofile=1024:524288"]
+        expected = " --ulimit rtprio=99 --ulimit memlock=102400 --ulimit nofile=1024:524288"
+        self.assertTrue(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_wrong_single_soft(self):
+        """Test if single soft limit argument is wrong."""
+        mock_cliargs = ["rtprio99"]
+        expected = " --ulimit rtprio99"
+        self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_wrong_multiple_soft(self):
+        """Test if multiple soft limit arguments are wrong."""
+        mock_cliargs = ["rtprio=99", "memlock102400"]
+        expected = " --ulimit rtprio=99 --ulimit memlock=102400"
+        self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_wrong_single_hard(self):
+        """Test if single hard limit arguments are wrong."""
+        mock_cliargs = ["nofile=1024:524288:"]
+        expected = " --ulimit nofile=1024:524288"
+        self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_wrong_multiple_hard(self):
+        """Test if multiple hard limit arguments are wrong."""
+        mock_cliargs = ["nofile1024524288", "rtprio=90:99"]
+        expected = " --ulimit nofile=1024:524288 --ulimit rtprio=90:99"
+        self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))
+
+    def test_args_wrong_multiple_mix(self):
+        """Test if multiple mixed limit arguments are wrong."""
+        mock_cliargs = ["rtprio=:", "memlock102400", "nofile1024:524288:"]
+        expected = " --ulimit rtprio=99 --ulimit memlock=102400 --ulimit nofile=1024:524288"
+        self.assertFalse(*self._is_arg_translation_ok(mock_cliargs, expected))