163 lines
5.4 KiB
Python
163 lines
5.4 KiB
Python
# Copyright 2017 The Chromium OS Authors. All rights reserved.
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
|
|
import logging
|
|
import os
|
|
import pyudev
|
|
import re
|
|
import select
|
|
import struct
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from autotest_lib.client.common_lib import error
|
|
|
|
|
|
JAIL_CONTROL_PATH = '/dev/jail-control'
|
|
JAIL_REQUEST_PATH = '/dev/jail-request'
|
|
|
|
# From linux/device_jail.h.
|
|
REQUEST_ALLOW = 0
|
|
REQUEST_ALLOW_WITH_LOCKDOWN = 1
|
|
REQUEST_ALLOW_WITH_DETACH = 2
|
|
REQUEST_DENY = 3
|
|
|
|
|
|
class OSFile:
|
|
"""Simple context manager for file descriptors."""
|
|
def __init__(self, path, flag):
|
|
self._fd = os.open(path, flag)
|
|
|
|
def close(self):
|
|
os.close(self._fd)
|
|
|
|
def __enter__(self):
|
|
"""Returns the fd so it can be used in with-blocks."""
|
|
return self._fd
|
|
|
|
def __exit__(self, exc_type, exc_val, traceback):
|
|
self.close()
|
|
|
|
|
|
class ConcurrentFunc:
|
|
"""Simple context manager that starts and joins a thread."""
|
|
def __init__(self, target_func, timeout_func):
|
|
self._thread = threading.Thread(target=target_func)
|
|
self._timeout_func = timeout_func
|
|
self._target_name = target_func.__name__
|
|
|
|
def __enter__(self):
|
|
self._thread.start()
|
|
|
|
def __exit__(self, exc_type, exc_val, traceback):
|
|
self._thread.join(self._timeout_func())
|
|
if self._thread.is_alive() and not exc_val:
|
|
raise error.TestError('Function %s timed out' % self._target_name)
|
|
|
|
|
|
class JailDevice:
|
|
TIMEOUT_SEC = 3
|
|
PATH_MAX = 4096
|
|
|
|
def __init__(self, path_to_jail):
|
|
self._path_to_jail = path_to_jail
|
|
|
|
|
|
def __enter__(self):
|
|
"""
|
|
Creates a jail device for the device located at self._path_to_jail.
|
|
If the jail already exists, don't take ownership of it.
|
|
"""
|
|
try:
|
|
output = subprocess.check_output(
|
|
['device_jail_utility',
|
|
'--add={0}'.format(self._path_to_jail)],
|
|
stderr=subprocess.STDOUT)
|
|
|
|
match = re.search('created jail at (.*)', output)
|
|
if match:
|
|
self._path = match.group(1)
|
|
self._owns_device = True
|
|
return self
|
|
|
|
match = re.search('jail already exists at (.*)', output)
|
|
if match:
|
|
self._path = match.group(1)
|
|
self._owns_device = False
|
|
return self
|
|
|
|
raise error.TestError('Failed to create device jail')
|
|
except subprocess.CalledProcessError as e:
|
|
raise error.TestError('Failed to call device_jail_utility')
|
|
|
|
|
|
def expect_open(self, verdict):
|
|
"""
|
|
Tries to open the jail device. This method mocks out the
|
|
device_jail request server which is normally run by permission_broker.
|
|
This allows us to set the verdict we want to test. Since the open
|
|
call will block until we return the verdict, we have to use a
|
|
separate thread to perform the open call, as well.
|
|
"""
|
|
# Python 2 does not support "nonlocal" so this closure can't
|
|
# set the values of identifiers it closes over unless they
|
|
# are in global scope. Work around this by using a list and
|
|
# value-mutation.
|
|
dev_file_wrapper = [None]
|
|
def open_device():
|
|
try:
|
|
dev_file_wrapper[0] = OSFile(self._path, os.O_RDWR)
|
|
except OSError as e:
|
|
# We don't throw an error because this might be intentional,
|
|
# such as when the verdict is REQUEST_DENY.
|
|
logging.info("Failed to open jail device: %s", e.strerror)
|
|
|
|
# timeout_sec should be used for the timeouts below.
|
|
# This ensures we don't spend much longer than TIMEOUT_SEC in
|
|
# this method.
|
|
deadline = time.time() + self.TIMEOUT_SEC
|
|
def timeout_sec():
|
|
return max(deadline - time.time(), 0.01)
|
|
|
|
# We have to use FDs because polling works with FDs and
|
|
# buffering is silly.
|
|
try:
|
|
req_f = OSFile(JAIL_REQUEST_PATH, os.O_RDWR)
|
|
except OSError as e:
|
|
raise error.TestError(
|
|
'Failed to open request device: %s' % e.strerror)
|
|
|
|
with req_f as req_fd:
|
|
poll_obj = select.poll()
|
|
poll_obj.register(req_fd, select.POLLIN)
|
|
|
|
# Starting open_device should ensure we have a request waiting
|
|
# on the request device.
|
|
with ConcurrentFunc(open_device, timeout_sec):
|
|
ready_fds = poll_obj.poll(timeout_sec() * 1000)
|
|
if not ready_fds:
|
|
raise error.TestError('Timed out waiting for jail-request')
|
|
|
|
# Sanity check the request.
|
|
path = os.read(req_fd, self.PATH_MAX)
|
|
logging.info('Received jail-request for path %s', path)
|
|
if path != self._path_to_jail:
|
|
raise error.TestError('Got request for the wrong path')
|
|
|
|
os.write(req_fd, struct.pack('I', verdict))
|
|
logging.info('Responded to jail-request')
|
|
|
|
return dev_file_wrapper[0]
|
|
|
|
|
|
def __exit__(self, exc_type, exc_val, traceback):
|
|
if self._owns_device:
|
|
subprocess.call(['device_jail_utility',
|
|
'--remove={0}'.format(self._path)])
|
|
|
|
|
|
def get_usb_devices():
|
|
context = pyudev.Context()
|
|
return [device for device in context.list_devices()
|
|
if device.device_node and device.device_node.startswith('/dev/bus/usb')]
|