1#!/usr/bin/python
2# Copyright 2017 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import os
7import tempfile
8import shutil
9import unittest
10from contextlib import contextmanager
11
12import common
13from autotest_lib.client.bin import utils
14from autotest_lib.client.common_lib import error
15from autotest_lib.site_utils import lxc
16from autotest_lib.site_utils.lxc import constants
17from autotest_lib.site_utils.lxc import unittest_http
18from autotest_lib.site_utils.lxc import unittest_setup
19from autotest_lib.site_utils.lxc import utils as lxc_utils
20
21
22@unittest.skipIf(lxc.IS_MOBLAB, 'Zygotes are not supported on moblab.')
23class ZygoteTests(unittest.TestCase):
24    """Unit tests for the Zygote class."""
25
26    @classmethod
27    def setUpClass(cls):
28        cls.test_dir = tempfile.mkdtemp(dir=lxc.DEFAULT_CONTAINER_PATH,
29                                        prefix='zygote_unittest_')
30
31        # Check if a base container exists on this machine and download one if
32        # necessary.
33        image = lxc.BaseImage()
34        try:
35            cls.base_container = image.get()
36            cls.cleanup_base_container = False
37        except error.ContainerError:
38            image.setup()
39            cls.base_container = image.get()
40            cls.cleanup_base_container = True
41        assert(cls.base_container is not None)
42
43        # Set up the zygote host path.
44        cls.shared_host_dir = lxc.SharedHostDir(
45                os.path.join(cls.test_dir, 'host'))
46
47
48    @classmethod
49    def tearDownClass(cls):
50        cls.base_container = None
51        if not unittest_setup.config.skip_cleanup:
52            if cls.cleanup_base_container:
53                lxc.BaseImage().cleanup()
54            cls.shared_host_dir.cleanup()
55            shutil.rmtree(cls.test_dir)
56
57
58    def testCleanup(self):
59        """Verifies that the zygote cleans up after itself."""
60        with self.createZygote() as zygote:
61            host_path = zygote.host_path
62
63            self.assertTrue(os.path.isdir(host_path))
64
65            # Start/stop the zygote to exercise the host mounts.
66            zygote.start(wait_for_network=False)
67            zygote.stop()
68
69        # After the zygote is destroyed, verify that the host path is cleaned
70        # up.
71        self.assertFalse(os.path.isdir(host_path))
72
73
74    def testCleanupWithUnboundHostDir(self):
75        """Verifies that cleanup works when the host dir is unbound."""
76        with self.createZygote() as zygote:
77            host_path = zygote.host_path
78
79            self.assertTrue(os.path.isdir(host_path))
80            # Don't start the zygote, so the host mount is not bound.
81
82        # After the zygote is destroyed, verify that the host path is cleaned
83        # up.
84        self.assertFalse(os.path.isdir(host_path))
85
86
87    def testCleanupWithNoHostDir(self):
88        """Verifies that cleanup works when the host dir is missing."""
89        with self.createZygote() as zygote:
90            host_path = zygote.host_path
91
92            utils.run('sudo rmdir %s' % zygote.host_path)
93            self.assertFalse(os.path.isdir(host_path))
94        # Zygote destruction should yield no errors if the host path is
95        # missing.
96
97
98    def testHostDir(self):
99        """Verifies that the host dir on the container is created, and correctly
100        bind-mounted."""
101        with self.createZygote() as zygote:
102            self.assertIsNotNone(zygote.host_path)
103            self.assertTrue(os.path.isdir(zygote.host_path))
104
105            zygote.start(wait_for_network=False)
106
107            self.verifyBindMount(
108                zygote,
109                container_path=lxc.CONTAINER_HOST_DIR,
110                host_path=zygote.host_path)
111
112
113    def testHostDirExists(self):
114        """Verifies that the host dir is just mounted if it already exists."""
115        # Pre-create the host dir and put a file in it.
116        test_host_path = os.path.join(self.shared_host_dir.path,
117                                      'testHostDirExists')
118        test_filename = 'test_file'
119        test_host_file = os.path.join(test_host_path, test_filename)
120        test_string = 'jackdaws love my big sphinx of quartz.'
121        os.makedirs(test_host_path)
122        with open(test_host_file, 'w') as f:
123            f.write(test_string)
124
125        # Sanity check
126        self.assertTrue(lxc_utils.path_exists(test_host_file))
127
128        with self.createZygote(host_path=test_host_path) as zygote:
129            zygote.start(wait_for_network=False)
130
131            self.verifyBindMount(
132                zygote,
133                container_path=lxc.CONTAINER_HOST_DIR,
134                host_path=zygote.host_path)
135
136            # Verify that the old directory contents was preserved.
137            cmd = 'cat %s' % os.path.join(lxc.CONTAINER_HOST_DIR,
138                                          test_filename)
139            test_output = zygote.attach_run(cmd).stdout.strip()
140            self.assertEqual(test_string, test_output)
141
142
143    def testInstallSsp(self):
144        """Verifies that installing the ssp in the container works."""
145        # Hard-coded path to some golden data for this test.
146        test_ssp = os.path.join(
147                common.autotest_dir,
148                'site_utils', 'lxc', 'test', 'test_ssp.tar.bz2')
149        # Create a container, install the self-served ssp, then check that it is
150        # installed into the container correctly.
151        with self.createZygote() as zygote:
152            # Note: start the zygote first, then install the SSP.  This mimics
153            # the way things would work in the production environment.
154            zygote.start(wait_for_network=False)
155            with unittest_http.serve_locally(test_ssp) as url:
156                zygote.install_ssp(url)
157
158            # The test ssp just contains a couple of text files, in known
159            # locations.  Verify the location and content of those files in the
160            # container.
161            cat = lambda path: zygote.attach_run('cat %s' % path).stdout
162            test0 = cat(os.path.join(constants.CONTAINER_AUTOTEST_DIR,
163                                     'test.0'))
164            test1 = cat(os.path.join(constants.CONTAINER_AUTOTEST_DIR,
165                                     'dir0', 'test.1'))
166            self.assertEquals('the five boxing wizards jumped quickly',
167                              test0)
168            self.assertEquals('the quick brown fox jumps over the lazy dog',
169                              test1)
170
171
172    def testInstallControlFile(self):
173        """Verifies that installing a control file in the container works."""
174        _unused, tmpfile = tempfile.mkstemp()
175        with self.createZygote() as zygote:
176            # Note: start the zygote first.  This mimics the way things would
177            # work in the production environment.
178            zygote.start(wait_for_network=False)
179            zygote.install_control_file(tmpfile)
180            # Verify that the file is found in the zygote.
181            zygote.attach_run(
182                'test -f %s' % os.path.join(lxc.CONTROL_TEMP_PATH,
183                                            os.path.basename(tmpfile)))
184
185
186    def testCopyFile(self):
187        """Verifies that files are correctly copied into the container."""
188        control_string = 'amazingly few discotheques provide jukeboxes'
189        with tempfile.NamedTemporaryFile() as tmpfile:
190            tmpfile.write(control_string)
191            tmpfile.flush()
192
193            with self.createZygote() as zygote:
194                dst = os.path.join(constants.CONTAINER_AUTOTEST_DIR,
195                                   os.path.basename(tmpfile.name))
196                zygote.start(wait_for_network=False)
197                zygote.copy(tmpfile.name, dst)
198                # Verify the file content.
199                test_string = zygote.attach_run('cat %s' % dst).stdout
200                self.assertEquals(control_string, test_string)
201
202
203    def testCopyDirectory(self):
204        """Verifies that directories are correctly copied into the container."""
205        control_string = 'pack my box with five dozen liquor jugs'
206        with lxc_utils.TempDir() as tmpdir:
207            fd, tmpfile = tempfile.mkstemp(dir=tmpdir)
208            f = os.fdopen(fd, 'w')
209            f.write(control_string)
210            f.close()
211
212            with self.createZygote() as zygote:
213                dst = os.path.join(constants.CONTAINER_AUTOTEST_DIR,
214                                   os.path.basename(tmpdir))
215                zygote.start(wait_for_network=False)
216                zygote.copy(tmpdir, dst)
217                # Verify the file content.
218                test_file = os.path.join(dst, os.path.basename(tmpfile))
219                test_string = zygote.attach_run('cat %s' % test_file).stdout
220                self.assertEquals(control_string, test_string)
221
222
223    def testFindHostMount(self):
224        """Verifies that zygotes pick up the correct host dirs."""
225        with self.createZygote() as zygote0:
226            # Not a clone, this just instantiates zygote1 on top of the LXC
227            # container created by zygote0.
228            zygote1 = lxc.Zygote(container_path=zygote0.container_path,
229                                 name=zygote0.name,
230                                 attribute_values={})
231            # Verify that the new zygote picked up the correct host path
232            # from the existing LXC container.
233            self.assertEquals(zygote0.host_path, zygote1.host_path)
234            self.assertEquals(zygote0.host_path_ro, zygote1.host_path_ro)
235
236
237    def testDetectExistingMounts(self):
238        """Verifies that host mounts are properly reconstructed.
239
240        When a Zygote is instantiated on top of an already-running container,
241        any previously-created bind mounts have to be detected.  This enables
242        proper cleanup later.
243        """
244        with lxc_utils.TempDir() as tmpdir, self.createZygote() as zygote0:
245            zygote0.start(wait_for_network=False)
246            # Create a bind mounted directory.
247            zygote0.mount_dir(tmpdir, 'foo')
248            # Create another zygote on top of the existing container.
249            zygote1 = lxc.Zygote(container_path=zygote0.container_path,
250                                 name=zygote0.name,
251                                 attribute_values={})
252            # Verify that the new zygote contains the same bind mounts.
253            self.assertEqual(zygote0.mounts, zygote1.mounts)
254
255
256    def testMountDirectory(self):
257        """Verifies that read-write mounts work."""
258        with lxc_utils.TempDir() as tmpdir, self.createZygote() as zygote:
259            dst = '/testMountDirectory/testMount'
260            zygote.start(wait_for_network=False)
261            zygote.mount_dir(tmpdir, dst, readonly=False)
262
263            # Verify that the mount point is correctly bound, and is read-write.
264            self.verifyBindMount(zygote, dst, tmpdir)
265            zygote.attach_run('test -r {0} -a -w {0}'.format(dst))
266
267
268    def testMountDirectoryReadOnly(self):
269        """Verifies that read-only mounts are mounted, and read-only."""
270        with lxc_utils.TempDir() as tmpdir, self.createZygote() as zygote:
271            dst = '/testMountDirectoryReadOnly/testMount'
272            zygote.start(wait_for_network=False)
273            zygote.mount_dir(tmpdir, dst, readonly=True)
274
275            # Verify that the mount point is correctly bound, and is read-only.
276            self.verifyBindMount(zygote, dst, tmpdir)
277            try:
278                zygote.attach_run('test -r {0} -a ! -w {0}'.format(dst))
279            except error.CmdError:
280                self.fail('Bind mount is not read-only')
281
282
283    def testMountDirectoryRelativePath(self):
284        """Verifies that relative-path mounts work."""
285        with lxc_utils.TempDir() as tmpdir, self.createZygote() as zygote:
286            dst = 'testMountDirectoryRelativePath/testMount'
287            zygote.start(wait_for_network=False)
288            zygote.mount_dir(tmpdir, dst, readonly=True)
289
290            # Verify that the mount points is correctly bound..
291            self.verifyBindMount(zygote, dst, tmpdir)
292
293
294    @contextmanager
295    def createZygote(self,
296                     name = None,
297                     attribute_values = None,
298                     snapshot = True,
299                     host_path = None):
300        """Clones a zygote from the test base container.
301        Use this to ensure that zygotes got properly cleaned up after each test.
302
303        @param container_path: The LXC path for the new container.
304        @param host_path: The host path for the new container.
305        @param name: The name of the new container.
306        @param attribute_values: Any attribute values for the new container.
307        @param snapshot: Whether to create a snapshot clone.
308        """
309        if name is None:
310            name = self.id().split('.')[-1]
311        if host_path is None:
312            host_path = os.path.join(self.shared_host_dir.path, name)
313        if attribute_values is None:
314            attribute_values = {}
315        zygote = lxc.Zygote(self.test_dir,
316                            name,
317                            attribute_values,
318                            self.base_container,
319                            snapshot,
320                            host_path)
321        try:
322            yield zygote
323        finally:
324            if not unittest_setup.config.skip_cleanup:
325                zygote.destroy()
326
327
328    def verifyBindMount(self, container, container_path, host_path):
329        """Verifies that a given path in a container is bind-mounted to a given
330        path in the host system.
331
332        @param container: The Container instance to be tested.
333        @param container_path: The path in the container to compare.
334        @param host_path: The path in the host system to compare.
335        """
336        container_inode = (container.attach_run('ls -id %s' % container_path)
337                           .stdout.split()[0])
338        host_inode = utils.run('ls -id %s' % host_path).stdout.split()[0]
339        # Compare the container and host inodes - they should match.
340        self.assertEqual(container_inode, host_inode)
341
342
343if __name__ == '__main__':
344    unittest_setup.setup()
345    unittest.main()
346