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