Get rid of SUID binary to restart docker containers

This patch removes a script with SUID flag,restart_docker_container,
which was used to restart docker containers and replaces it by
native python implementation. The aim of this replacmenet is to reduce
the security risk caused by SUID binaries, as well as to implement
more features and unit test coverages easily.

To allow swift user to manage docker conainers, now it is required
that swift user belong to docker group and have access to docker
unix domain socket.

Change-Id: I8103d4d826f5121e16f67f1ff49102ceecaf8a80

.gitignore

@@ -36,9 +36,6 @@ src/java/SBus/org_openstack_storlet_sbus_SBusJNI.h
 *.jar
 StorletSamples/java/*/bin
 
-# scripts build
-scripts/restart_docker_container
-
 # functional tests
 tests/functional/.ipynb_checkpoints/
 

devstack/plugin.sh

@@ -47,7 +47,6 @@ SWIFT_MEMBER_USER_PWD=member
 # Storlets install tunables
 STORLETS_DEFAULT_USER_DOMAIN_ID=${STORLETS_DEFAULT_USER_DOMAIN_ID:-default}
 STORLETS_DEFAULT_PROJECT_DOMAIN_ID=${STORLETS_DEFAULT_PROJECT_DOMAIN_ID:-default}
-STORLET_MANAGEMENT_USER=${STORLET_MANAGEMENT_USER:-$USER}
 STORLETS_DOCKER_DEVICE=${STORLETS_DOCKER_DEVICE:-/home/docker_device}
 STORLETS_DOCKER_BASE_IMG=${STORLETS_DOCKER_BASE_IMG:-ubuntu:20.04}
 STORLETS_DOCKER_BASE_IMG_NAME=${STORLETS_DOCKER_BASE_IMG_NAME:-ubuntu_20.04}
@@ -60,7 +59,7 @@ STORLETS_STORLET_CONTAINER_NAME=${STORLETS_STORLET_CONTAINER_NAME:-storlet}
 STORLETS_DEPENDENCY_CONTAINER_NAME=${STORLETS_DEPENDENCY_CONTAINER_NAME:-dependency}
 STORLETS_LOG_CONTAIER_NAME=${STORLETS_LOG_CONTAIER_NAME:-log}
 STORLETS_GATEWAY_MODULE=${STORLETS_GATEWAY_MODULE:-docker}
-STORLETS_GATEWAY_CONF_FILE=${STORLETS_GATEWAY_CONF_FILE:-/etc/swift/storlet_docker_gateway.conf}
+STORLETS_GATEWAY_CONF_FILE=${STORLETS_GATEWAY_CONF_FILE:-${SWIFT_CONF_DIR}/storlet_docker_gateway.conf}
 STORLETS_PROXY_EXECUTION_ONLY=${STORLETS_PROXY_EXECUTION_ONLY:-false}
 STORLETS_SCRIPTS_DIR=${STORLETS_SCRIPTS_DIR:-"$STORLETS_DOCKER_DEVICE"/scripts}
 STORLETS_STORLETS_DIR=${STORLETS_STORLETS_DIR:-"$STORLETS_DOCKER_DEVICE"/storlets/scopes}
@@ -76,18 +75,6 @@ TMP_REGISTRY_PREFIX=/tmp/registry
 # Functions
 # ---------
 
-function _storlets_swift_start {
-    swift-init --run-dir=${SWIFT_DATA_DIR}/run all start || true
-}
-
-function _storlets_swift_stop {
-    swift-init --run-dir=${SWIFT_DATA_DIR}/run all stop || true
-}
-
-function _storlets_swift_restart {
-    swift-init --run-dir=${SWIFT_DATA_DIR}/run all restart || true
-}
-
 function _export_os_vars {
     export OS_IDENTITY_API_VERSION=3
     export OS_AUTH_URL="http://$KEYSTONE_IP/identity/v3"
@@ -148,7 +135,7 @@ function configure_swift_and_keystone_for_storlets {
     rm /tmp/storlet-docker-gateway.conf
 
     # Create storlet related containers and set ACLs
-    _storlets_swift_start
+    start_swift
     _export_swift_os_vars
     openstack object store account set --property Storlet-Enabled=True
     swift post --read-acl $SWIFT_DEFAULT_PROJECT:$SWIFT_MEMBER_USER $STORLETS_STORLET_CONTAINER_NAME
@@ -160,46 +147,32 @@ function _install_docker {
     # TODO: Add other dirstors.
     # This one is geared towards Ubuntu
     # See other projects that install docker
-    DOCKER_UNIX_SOCKET=/var/run/docker.sock
-    DOCKER_SERVICE_TIMEOUT=5
 
-    install_package socat
     wget http://get.docker.com -O install_docker.sh
     sudo chmod 777 install_docker.sh
     sudo bash -x install_docker.sh
     sudo rm install_docker.sh
 
-    sudo killall docker || true
+    # Add swift user to docker group so that the user can manage docker
-
+    # containers without sudo
-    # systemd env doesn't require /etc/default/docker options
+    sudo grep -q docker /etc/group
-    if [[ ! -e /etc/default/docker ]]; then
+    if [ $? -ne 0 ]; then
-        sudo touch /etc/default/docker
+      sudo groupadd docker
-        sudo ls /lib/systemd/system
-        sudo sed -i '0,/[service]/a EnvironmentFile=-/etc/default/docker' /lib/systemd/system/docker.service
-        sudo cat /lib/systemd/system/docker.service
     fi
-    sudo cat /etc/default/docker
+    add_user_to_group $STORLETS_SWIFT_RUNTIME_USER docker
-    sudo sed -r 's#^.*DOCKER_OPTS=.*$#DOCKER_OPTS="--debug -g /home/docker_device/docker --storage-opt dm.override_udev_sync_check=true"#' /etc/default/docker
 
-    # Start the daemon - restart just in case the package ever auto-starts...
+    if [ $STORLETS_SWIFT_RUNTIME_USER == $USER ]; then
+      # NOTE(takashi): We need this workaroud because we can't reload
+      #                user-group relationship in bash scripts
+      DOCKER_UNIX_SOCKET=/var/run/docker.sock
+      sudo chown $USER:$USER $DOCKER_UNIX_SOCKET
+    fi
+
+    # Restart docker daemon
     restart_service docker
-
-    echo "Waiting for docker daemon to start..."
-    DOCKER_GROUP=$(groups | cut -d' ' -f1)
-    CONFIGURE_CMD="while ! /bin/echo -e 'GET /version HTTP/1.0\n\n' | socat - unix-connect:$DOCKER_UNIX_SOCKET 2>/dev/null | grep -q '200 OK'; do
-      # Set the right group on docker unix socket before retrying
-      sudo chgrp $DOCKER_GROUP $DOCKER_UNIX_SOCKET
-      sudo chmod g+rw $DOCKER_UNIX_SOCKET
-      sleep 1
-    done"
-    if ! timeout $DOCKER_SERVICE_TIMEOUT sh -c "$CONFIGURE_CMD"; then
-      die $LINENO "docker did not start"
-fi
 }
 
 function prepare_storlets_install {
-    sudo mkdir -p "$STORLETS_DOCKER_DEVICE"/docker
-    sudo chmod 777 $STORLETS_DOCKER_DEVICE
     _install_docker
 
     if is_ubuntu; then
@@ -236,11 +209,11 @@ EOF
 
 function create_base_jre_image {
     echo "Create base jre image"
-    docker pull $STORLETS_DOCKER_BASE_IMG
+    sudo docker pull $STORLETS_DOCKER_BASE_IMG
     mkdir -p ${TMP_REGISTRY_PREFIX}/repositories/"$STORLETS_DOCKER_BASE_IMG_NAME"_jre${STORLETS_JDK_VERSION}
     _generate_jre_dockerfile
     cd ${TMP_REGISTRY_PREFIX}/repositories/"$STORLETS_DOCKER_BASE_IMG_NAME"_jre${STORLETS_JDK_VERSION}
-    docker build -q -t ${STORLETS_DOCKER_BASE_IMG_NAME}_jre${STORLETS_JDK_VERSION} .
+    sudo docker build -q -t ${STORLETS_DOCKER_BASE_IMG_NAME}_jre${STORLETS_JDK_VERSION} .
     cd -
 }
 
@@ -292,7 +265,7 @@ function create_storlet_engine_image {
     _generate_logback_xml
     _generate_jre_storlet_dockerfile
     cd ${TMP_REGISTRY_PREFIX}/repositories/"$STORLETS_DOCKER_BASE_IMG_NAME"_jre${STORLETS_JDK_VERSION}_storlets
-    docker build -q -t ${STORLETS_DOCKER_BASE_IMG_NAME}_jre${STORLETS_JDK_VERSION}_storlets .
+    sudo docker build -q -t ${STORLETS_DOCKER_BASE_IMG_NAME}_jre${STORLETS_JDK_VERSION}_storlets .
     cd -
 }
 
@@ -317,12 +290,8 @@ function install_storlets_code {
         sudo cp `which ${bin_file}` /usr/local/libexec/storlets/
     done
 
-    sudo mkdir -p $STORLETS_DOCKER_DEVICE/scripts
+    sudo mkdir -p -m 0755 $STORLETS_DOCKER_DEVICE
-    sudo chown "$STORLETS_SWIFT_RUNTIME_USER":"$STORLETS_SWIFT_RUNTIME_GROUP" "$STORLETS_DOCKER_DEVICE"/scripts
+    sudo chown -R "$STORLETS_SWIFT_RUNTIME_USER":"$STORLETS_SWIFT_RUNTIME_GROUP" $STORLETS_DOCKER_DEVICE
-    sudo chmod 0755 "$STORLETS_DOCKER_DEVICE"/scripts
-    sudo cp scripts/restart_docker_container "$STORLETS_DOCKER_DEVICE"/scripts/
-    sudo chmod 04755 "$STORLETS_DOCKER_DEVICE"/scripts/restart_docker_container
-    sudo chown root:root "$STORLETS_DOCKER_DEVICE"/scripts/restart_docker_container
 
     # NOTE(takashi): We should cleanup egg-info directory here, otherwise it
     #                causes permission denined when installing package by tox.
@@ -334,13 +303,11 @@ function install_storlets_code {
 function _generate_swift_middleware_conf {
     cat <<EOF > /tmp/swift_middleware_conf
 [proxy-confs]
-proxy_server_conf_file = /etc/swift/proxy-server.conf
+proxy_server_conf_file = ${SWIFT_CONF_DIR}/proxy-server.conf
-storlet_proxy_server_conf_file = /etc/swift/storlet-proxy-server.conf
+storlet_proxy_server_conf_file = ${SWIFT_CONF_DIR}/storlet-proxy-server.conf
 
 [object-confs]
-object_server_conf_files = /etc/swift/object-server/1.conf
+object_server_conf_files = ${SWIFT_CONF_DIR}/object-server/1.conf
-#object_server_conf_files = /etc/swift/object-server/1.conf, /etc/swift/object-server/2.conf, /etc/swift/object-server/3.conf, /etc/swift/object-server/4.conf
-#object_server_conf_files = /etc/swift/object-server.conf
 
 [common-confs]
 storlet_middleware = $STORLETS_MIDDLEWARE_NAME
@@ -379,7 +346,7 @@ function create_default_tenant_image {
     mkdir -p ${TMP_REGISTRY_PREFIX}/repositories/$SWIFT_DEFAULT_PROJECT_ID
     _generate_default_tenant_dockerfile
     cd ${TMP_REGISTRY_PREFIX}/repositories/$SWIFT_DEFAULT_PROJECT_ID
-    docker build -q -t ${SWIFT_DEFAULT_PROJECT_ID:0:13} .
+    sudo docker build -q -t ${SWIFT_DEFAULT_PROJECT_ID:0:13} .
     cd -
 }
 
@@ -414,12 +381,12 @@ function install_storlets {
     create_test_config_file
 
     echo "restart swift"
-    _storlets_swift_restart
+    stop_swift
+    start_swift
 }
 
 function uninstall_storlets {
     sudo service docker stop
-    sudo sed -r 's#^.*DOCKER_OPTS=.*$#DOCKER_OPTS="--debug --storage-opt dm.override_udev_sync_check=true"#' /etc/default/docker
 
     echo "Cleaning all storlets runtime stuff..."
     sudo rm -fr ${STORLETS_DOCKER_DEVICE}

doc/source/engine_dev_installation.rst

@@ -130,6 +130,14 @@ We need the following for Docker
     sed -i '$acomplete -F _docker docker' /etc/bash_completion.d/docker
     update-rc.d docker defaults
 
+Also, add swift user to docker group so that the user can manage docker
+containers without sudo
+
+::
+
+    sudo usermod -aG docker swift
+
+
 Get and install the storlets code
 ---------------------------------
 
@@ -258,18 +266,6 @@ Create the run time directory
     sudo mkdir -p $STORLETS_HOME
     sudo chmod 777 $STORLETS_HOME
 
-Create the scripts directory and populate it.
-Note that these scripts are executed by the middleware but
-require root privileges.
-
-::
-
-    mkdir $STORLETS_HOME/scripts
-    cd STORLETS_HOME/scripts
-    cp $HOME/scripts/restart_docker_container .
-    sudo chown root:root restart_docker_container
-    sudo chmod 04755 restart_docker_container
-
 The run time directory will be later populated by the middleware with:
  #. storlets - Docker container mapped directories keeping storlet jars
  #. pipe - A Docker container mapped directories holding named pipes shared between the middleware and the containers.

install_libs.sh

@@ -6,12 +6,6 @@
 #       so you may need root privilege to execute this script
 set -x
 
-# build scripts
-cd scripts
-# TODO(takashi): also install them
-make
-cd -
-
 # install c library
 cd src/c/sbus
 make && make install

requirements.txt

@@ -7,3 +7,4 @@ setuptools>=17.1
 eventlet>=0.17.4 # MIT
 greenlet>=0.3.1
 stevedore>=1.16.0  # Apache-2.0
+docker

s2aio.sh

@@ -103,7 +103,6 @@ function uninstall_swift_using_devstack {
     sudo sed -i.bak '/swift.img/d'  /etc/fstab
 }
 
-
 function uninstall_s2aio {
     _prepare_devstack_env
 
@@ -131,6 +130,7 @@ case $COMMAND in
   "stop" )
     stop_s2aio
     ;;
+
   * )
     usage
 esac

scripts/Makefile

@@ -1,16 +0,0 @@
-CC = gcc
-CFLAGS =
-LDFLAGS =
-TARGET = restart_docker_container
-
-SRCS = restart_docker_container.c
-OBJS = $(SRCS:.c=.o)
-
-.PHONY: all
-all: ${TARGET}
-
-$(TARGET): $(OBJS)
-	$(CC) ${LDFLAGS} -o $@ $^
-
-clean:
-	rm ${TARGET} ${OBJS}

scripts/restart_docker_container.c

@@ -1,87 +0,0 @@
-/*----------------------------------------------------------------------------
- * Copyright IBM Corp. 2015, 2015 All Rights Reserved
- * Copyright (c) 2010-2016 OpenStack 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.
- * ---------------------------------------------------------------------------
-*/
-
-#define _GNU_SOURCE
-#include <stdlib.h>
-#include <stdio.h>
-#include <unistd.h>
-#include <sys/types.h>
-
- /*
-  * Stop and Run a docker container using:
-  * docker stop <container name>
-  * docker run --net=none --name <container name> -d -v /dev/log:/dev/log \
-  *     -v <mount dir 1> -v <mount dir 2> -v <mount dir 3> -v <mount dir 4> \
-  *     <image name>
-  *
-  * <container name> - The name of the container to stop / to start
-  * <image name>     - the name of the image from which to start the container
-  * <mount dir 1>    - The directory where the named pipes are placed.
-  *                    Typically mounted to /mnt/channels in the container
-  * <mount dir 2>    - The directory where the storlets are placed.
-  *                    Typically mounted to /home/swift in the container
-  * <mount dir 3>    - The directory where storlets library are placed.
-  *                    Typically mounted to /usr/local/lib/storlets
-  * <mount dir 4>    - The directory where storlets executables are placed.
-  *                    Typically mounted to /usr/local/libexec/storlets
-  */
-
-int main(int argc, char **argv) {
-    char command[4096];
-    char container_name[256];
-    char container_image[256];
-    char mount_dir1[512];
-    char mount_dir2[512];
-    char mount_dir3[512];
-    char mount_dir4[512];
-
-    if (argc != 7) {
-        fprintf(stderr, "Usage: %s container_name container_image mount_dir1 mount_dir2 mount_dir3 mount_dir4\n",
-            argv[0]);
-        return 1;
-    }
-
-    snprintf(container_name,(size_t)256,"%s",argv[1]);
-    snprintf(container_image,(size_t)256,"%s",argv[2]);
-    snprintf(mount_dir1,(size_t)512, "%s", argv[3]);
-    snprintf(mount_dir2,(size_t)512, "%s", argv[4]);
-    snprintf(mount_dir3,(size_t)512, "%s", argv[5]);
-    snprintf(mount_dir4,(size_t)512, "%s", argv[6]);
-
-    int ret;
-    setresuid(0, 0, 0);
-    setresgid(0, 0, 0);
-    sprintf(command, "/usr/bin/docker stop -t 1 %s", container_name);
-    ret = system(command);
-
-    sprintf(command, "/usr/bin/docker rm %s", container_name);
-    ret = system(command);
-
-    sprintf(command,
-            "/usr/bin/docker run --net=none --name %s -d -v /dev/log:/dev/log -v %s -v %s -v %s -v %s %s",
-            container_name,
-            mount_dir1,
-            mount_dir2,
-            mount_dir3,
-            mount_dir4,
-            container_image);
-    ret = system(command);
-    if(ret){
-        return EXIT_FAILURE;
-    }
-    return EXIT_SUCCESS;
-}

storlets/gateway/gateways/docker/runtime.py

@@ -17,10 +17,12 @@ import errno
 import os
 import select
 import stat
-import subprocess
 import sys
 import time
 import six
+import docker
+import docker.errors
+from docker.types import Mount as DockerMount
 
 import eventlet
 import json
@@ -290,38 +292,44 @@ class RunTimeSandbox(object):
         docker_container_name = '%s_%s' % (self.docker_image_name_prefix,
                                            self.scope)
 
-        pipe_mount = '%s:%s' % (self.paths.host_pipe_dir,
+        mounts = [
-                                self.paths.sandbox_pipe_dir)
+            DockerMount('/dev/log', '/dev/log', type='bind'),
-        storlet_mount = '%s:%s:ro' % (self.paths.host_storlet_base_dir,
+            DockerMount(self.paths.sandbox_pipe_dir,
-                                      self.paths.sandbox_storlet_base_dir)
+                        self.paths.host_pipe_dir,
-        storlet_native_lib_mount = '%s:%s:ro' % (
+                        type='bind'),
-            self.paths.host_storlet_native_lib_dir,
+            DockerMount(self.paths.sandbox_storlet_base_dir,
-            self.paths.sandbox_storlet_native_lib_dir)
+                        self.paths.host_storlet_base_dir,
-        storlet_native_bin_mount = '%s:%s:ro' % (
+                        type='bind'),
-            self.paths.host_storlet_native_bin_dir,
+            DockerMount(self.paths.sandbox_storlet_native_lib_dir,
-            self.paths.sandbox_storlet_native_bin_dir)
+                        self.paths.host_storlet_native_lib_dir,
+                        type='bind', read_only=True),
+            DockerMount(self.paths.sandbox_storlet_native_bin_dir,
+                        self.paths.host_storlet_native_bin_dir,
+                        type='bind', read_only=True)
+            ]
 
-        cmd = [os.path.join(self.paths.host_restart_script_dir,
+        try:
-                            'restart_docker_container'),
+            client = docker.from_env()
-               docker_container_name, docker_image_name, pipe_mount,
+            # Stop the existing storlet container
-               storlet_mount, storlet_native_lib_mount,
+            try:
-               storlet_native_bin_mount]
+                scontainer = client.containers.get(docker_container_name)
+            except docker.errors.NotFound:
+                # The container is not yet created
+                pass
+            else:
+                scontainer.stop(timeout=1)
+                scontainer.remove()
 
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
+            # Start the new one
-                                stderr=subprocess.PIPE)
+            client.containers.run(
-        stdout, stderr = proc.communicate()
+                docker_image_name, detach=True, name=docker_container_name,
-
+                network_disabled=True, mounts=mounts)
-        if stdout:
+        except docker.errors.ImageNotFound:
-            if not isinstance(stdout, str):
+            msg = "Image %s is not found" % docker_image_name
-                stdout = stdout.decode("utf-8")
+            raise StorletRuntimeException(msg)
-            self.logger.debug('STDOUT: %s' % stdout.replace('\n', '#012'))
+        except docker.errors.APIError:
-        if stderr:
+            self.logger.exception("Failed to manage docker containers")
-            if not isinstance(stderr, str):
+            raise StorletRuntimeException("Docker runtime error")
-                stderr = stderr.decode("utf-8")
-            self.logger.error('STDERR: %s' % stderr.replace('\n', '#012'))
-
-        if proc.returncode:
-            raise StorletRuntimeException('Failed to restart docker container')
 
     def restart(self):
         """

tests/unit/gateway/gateways/docker/test_runtime.py

@@ -21,6 +21,10 @@ import errno
 from contextlib import contextmanager
 from six import StringIO
 from stat import ST_MODE
+import docker.client
+import docker.errors
+import docker.models.containers
+
 
 from storlets.sbus.client import SBusResponse
 from storlets.sbus.client.exceptions import SBusClientIOError, \
@@ -244,7 +248,8 @@ class TestRunTimeSandbox(unittest.TestCase):
     def setUp(self):
         self.logger = FakeLogger()
         # TODO(takashi): take these values from config file
-        self.conf = {'docker_repo': 'localhost:5001'}
+        self.conf = {'docker_repo': 'localhost:5001',
+                     'default_docker_image_name': 'defaultimage'}
         self.scope = '0123456789abc'
         self.sbox = RunTimeSandbox(self.scope, self.conf, self.logger)
 
@@ -277,8 +282,8 @@ class TestRunTimeSandbox(unittest.TestCase):
     def test_wait(self):
         with mock.patch('storlets.gateway.gateways.docker.runtime.'
                         'SBusClient.ping') as ping, \
-            mock.patch('storlets.gateway.gateways.docker.runtime.'
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-                       'time.sleep') as sleep:
+                           'time.sleep') as sleep:
             ping.return_value = SBusResponse(True, 'OK')
             self.sbox.wait()
             self.assertEqual(sleep.call_count, 0)
@@ -294,87 +299,161 @@ class TestRunTimeSandbox(unittest.TestCase):
 
         # TODO(takashi): should test timeout case
 
+    def test__restart(self):
+        # storlet container is not running
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_client.containers = mock_containers
+            mock_containers.get.side_effect = \
+                docker.errors.NotFound('container is not found')
+            docker_from_env.return_value = mock_client
+
+            self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(1, mock_containers.run.call_count)
+
+        # storlet container is running
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_client.containers = mock_containers
+            mock_container = \
+                mock.MagicMock(spec_set=docker.models.containers.Container)
+            mock_containers.get.return_value = mock_container
+            docker_from_env.return_value = mock_client
+
+            self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(1, mock_container.stop.call_count)
+            self.assertEqual(1, mock_container.remove.call_count)
+            self.assertEqual(1, mock_containers.run.call_count)
+
+        # get failed
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_client.containers = mock_containers
+            mock_containers.get.side_effect = \
+                docker.errors.APIError('api error')
+            docker_from_env.return_value = mock_client
+
+            with self.assertRaises(StorletRuntimeException):
+                self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(0, mock_containers.run.call_count)
+
+        # stop failed
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_client.containers = mock_containers
+            mock_container = \
+                mock.MagicMock(spec_set=docker.models.containers.Container)
+            mock_containers.get.return_value = mock_container
+            mock_container.stop.side_effect = \
+                docker.errors.APIError('api error')
+            docker_from_env.return_value = mock_client
+
+            with self.assertRaises(StorletRuntimeException):
+                self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(1, mock_container.stop.call_count)
+            self.assertEqual(0, mock_container.remove.call_count)
+            self.assertEqual(0, mock_containers.run.call_count)
+
+        # remove failed
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_client.containers = mock_containers
+            mock_container = \
+                mock.MagicMock(spec_set=docker.models.containers.Container)
+            mock_containers.get.return_value = mock_container
+            mock_container.remove.side_effect = \
+                docker.errors.APIError('api error')
+            docker_from_env.return_value = mock_client
+
+            with self.assertRaises(StorletRuntimeException):
+                self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(1, mock_container.stop.call_count)
+            self.assertEqual(1, mock_container.remove.call_count)
+            self.assertEqual(0, mock_containers.run.call_count)
+
+        # run failed
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+                        'docker.from_env') as docker_from_env:
+            mock_client = mock.MagicMock(spec_set=docker.client.DockerClient)
+            mock_containers = mock.MagicMock(
+                spec_set=docker.models.containers.ContainerCollection)
+            mock_containers.run.side_effect = \
+                docker.errors.APIError('api error')
+            mock_client.containers = mock_containers
+            mock_container = \
+                mock.MagicMock(spec_set=docker.models.containers.Container)
+            mock_containers.get.return_value = mock_container
+            docker_from_env.return_value = mock_client
+
+            with self.assertRaises(StorletRuntimeException):
+                self.sbox._restart('storlet_image')
+            self.assertEqual(1, mock_containers.get.call_count)
+            self.assertEqual(1, mock_container.stop.call_count)
+            self.assertEqual(1, mock_container.remove.call_count)
+            self.assertEqual(1, mock_containers.run.call_count)
+
     def test_restart(self):
-
-        class FakeProc(object):
-            def __init__(self, stdout, stderr, code):
-                self.stdout = stdout
-                self.stderr = stderr
-                self.returncode = code
-
-            def communicate(self):
-                return (self.stdout, self.stderr)
-
         with mock.patch('storlets.gateway.gateways.docker.runtime.'
-                        'RunTimePaths.create_host_pipe_dir'), \
+                        'RunTimePaths.create_host_pipe_dir') as pipe_dir, \
-            mock.patch('storlets.gateway.gateways.docker.runtime.'
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-                       'subprocess.Popen') as popen:
+                           'RunTimeSandbox._restart') as _restart, \
-            _wait = self.sbox.wait
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-
+                           'RunTimeSandbox.wait') as wait:
-            def dummy_wait_success(*args, **kwargs):
-                return 1
-
-            self.sbox.wait = dummy_wait_success
-
-            # Test that popen is called successfully
-            popen.return_value = FakeProc('Try to restart\nOK', '', 0)
             self.sbox.restart()
-            self.assertEqual(1, popen.call_count)
+            self.assertEqual(1, pipe_dir.call_count)
-            self.sbox.wait = _wait
+            self.assertEqual(1, _restart.call_count)
+            self.assertEqual((self.scope,), _restart.call_args.args)
+            self.assertEqual(1, wait.call_count)
 
         with mock.patch('storlets.gateway.gateways.docker.runtime.'
-                        'RunTimePaths.create_host_pipe_dir'), \
+                        'RunTimePaths.create_host_pipe_dir') as pipe_dir, \
-            mock.patch('storlets.gateway.gateways.docker.runtime.'
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-                       'subprocess.Popen') as popen:
+                           'RunTimeSandbox._restart') as _restart, \
-            _wait = self.sbox.wait
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
+                           'RunTimeSandbox.wait') as wait:
+            _restart.side_effect = [StorletRuntimeException(), None]
+            self.sbox.restart()
+            self.assertEqual(1, pipe_dir.call_count)
+            self.assertEqual(2, _restart.call_count)
+            self.assertEqual((self.scope,),
+                             _restart.call_args_list[0].args)
+            self.assertEqual(('defaultimage',),
+                             _restart.call_args_list[1].args)
+            self.assertEqual(1, wait.call_count)
 
-            def dummy_wait_success(*args, **kwargs):
+        with mock.patch('storlets.gateway.gateways.docker.runtime.'
-                return 1
+                        'RunTimePaths.create_host_pipe_dir') as pipe_dir, \
-
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-            self.sbox.wait = dummy_wait_success
+                           'RunTimeSandbox._restart') as _restart, \
-
+                mock.patch('storlets.gateway.gateways.docker.runtime.'
-            # Test double failure to restart the container
+                           'RunTimeSandbox.wait') as wait:
-            # for both the tenant image and generic image
+            _restart.side_effect = StorletTimeout()
-            popen.return_value = FakeProc('Try to restart', 'Some error', 1)
             with self.assertRaises(StorletRuntimeException):
                 self.sbox.restart()
-            self.assertEqual(2, popen.call_count)
+            self.assertEqual(1, pipe_dir.call_count)
-            self.sbox.wait = _wait
+            self.assertEqual(1, _restart.call_count)
-
+            self.assertEqual((self.scope,), _restart.call_args.args)
-        with mock.patch('storlets.gateway.gateways.docker.runtime.'
+            self.assertEqual(0, wait.call_count)
-                        'RunTimePaths.create_host_pipe_dir'), \
-            mock.patch('storlets.gateway.gateways.docker.runtime.'
-                       'subprocess.Popen') as popen:
-            _wait = self.sbox.wait
-
-            def dummy_wait_success(*args, **kwargs):
-                return 1
-
-            self.sbox.wait = dummy_wait_success
-
-            # Test failure to restart the container for the tenant image
-            # success for the generic image
-            popen.side_effect = [FakeProc('Try to restart', 'Some error', 1),
-                                 FakeProc('Try to restart\nOK', '', 0)]
-            self.sbox.restart()
-            self.assertEqual(2, popen.call_count)
-            self.sbox.wait = _wait
-
-        with mock.patch('storlets.gateway.gateways.docker.runtime.'
-                        'RunTimePaths.create_host_pipe_dir'), \
-            mock.patch('storlets.gateway.gateways.docker.runtime.'
-                       'subprocess.Popen') as popen:
-            _wait = self.sbox.wait
-
-            def dummy_wait_failure(*args, **kwargs):
-                raise StorletTimeout()
-
-            self.sbox.wait = dummy_wait_failure
-
-            popen.return_value = FakeProc('OK', '', 0)
-            with self.assertRaises(StorletRuntimeException):
-                self.sbox.restart()
-            self.sbox.wait = _wait
 
     def test_get_storlet_classpath(self):
         storlet_id = 'Storlet.jar'