From 47483bf875a53bce9121f7f5bfd5f34527dd4b89 Mon Sep 17 00:00:00 2001
From: ChengyuZhu6 <chengyu.zhu@intel.com>
Date: Wed, 23 Aug 2023 13:22:17 +0800
Subject: [PATCH] pre-install: Configure nydus snapshotter in containerd

Signed-off-by: ChengyuZhu6 <chengyu.zhu@intel.com>
---
 .../scripts/container-engine-for-cc-deploy.sh | 51 +++++++++++++++++--
 1 file changed, 46 insertions(+), 5 deletions(-)

diff --git a/install/pre-install-payload/scripts/container-engine-for-cc-deploy.sh b/install/pre-install-payload/scripts/container-engine-for-cc-deploy.sh
index cd4141d7..d16d5e36 100755
--- a/install/pre-install-payload/scripts/container-engine-for-cc-deploy.sh
+++ b/install/pre-install-payload/scripts/container-engine-for-cc-deploy.sh
@@ -4,6 +4,8 @@ set -o errexit
 set -o pipefail
 set -o nounset
 
+containerd_config="/etc/containerd/config.toml"
+
 die() {
 	msg="$*"
 	echo "ERROR: $msg" >&2
@@ -20,7 +22,7 @@ function get_container_engine() {
 		die "${container_engine} is not yet supported"
 	fi
 
-	echo "$container_engine"	
+	echo "$container_engine"
 }
 
 function set_container_engine() {
@@ -36,6 +38,9 @@ function install_artifacts() {
 
 	install -D -m 755 ${artifacts_dir}/opt/confidential-containers/bin/containerd /opt/confidential-containers/bin/containerd
 	install -D -m 644 ${artifacts_dir}/etc/systemd/system/containerd.service.d/containerd-for-cc-override.conf /etc/systemd/system/containerd.service.d/containerd-for-cc-override.conf
+	if [ "$SNAPSHOTTER" == "nydus" ]; then
+		install -D -m 755 ${artifacts_dir}/opt/confidential-containers/bin/containerd-nydus-grpc /opt/confidential-containers/bin/containerd-nydus-grpc
+	fi
 }
 
 function uninstall_artifacts() {
@@ -45,7 +50,14 @@ function uninstall_artifacts() {
 	rm -f /etc/systemd/system/${container_engine}.service.d/${container_engine}-for-cc-override.conf
 	echo "Removing the systemd drop-in file's directory, if empty"
 	[ -d /etc/systemd/system/${container_engine}.service.d ] && rmdir --ignore-fail-on-non-empty /etc/systemd/system/${container_engine}.service.d
-	
+
+	if [ "$SNAPSHOTTER" == "nydus" ]; then
+		echo "Removing the nydus-snapshotter binary"
+		rm -f /opt/confidential-containers/bin/containerd-nydus-grpc
+		echo "unplugging the nydus-snapshotter from containerd"
+		remove_snapshotter_from_containerd
+	fi
+
 	restart_systemd_service
 
 	echo "Removing the containerd binary"
@@ -61,6 +73,34 @@ function restart_systemd_service() {
 	host_systemctl restart "${container_engine}"
 }
 
+function config_containerd() {
+	echo "configure snapshotter in containerd"
+
+	if [ ! -f "$containerd_config" ]; then
+		die "failed to find containerd config"
+	fi
+
+	if [ "$SNAPSHOTTER" == "nydus" ]; then
+		echo "Plug nydus snapshotter into containerd"
+		snapshotter_socket="/run/containerd-nydus/containerd-nydus-grpc.sock"
+	fi
+	proxy_config="  [proxy_plugins.$SNAPSHOTTER]\n    type = \"snapshot\"\n    address = ${snapshotter_socket}"
+
+	if grep -q "\[proxy_plugins\]" "$containerd_config"; then
+		sed -i '/\[proxy_plugins\]/a\'"$proxy_config" "$containerd_config"
+	else
+		echo -e "[proxy_plugins]" >>"$containerd_config"
+		echo -e "$proxy_config" >>"$containerd_config"
+	fi
+}
+function remove_snapshotter_from_containerd() {
+	if [ ! -f "$containerd_config" ]; then
+		die "failed to find containerd config"
+	fi
+	if [ "$SNAPSHOTTER" == "nydus" ]; then
+		sed -i '/\[proxy_plugins.nydus\]/,/address = "\/run\/containerd-nydus\/containerd-nydus-grpc\.sock"/d' "$containerd_config"
+	fi
+}
 label_node() {
 	case "${1}" in
 	install)
@@ -69,8 +109,7 @@ label_node() {
 	uninstall)
 		kubectl label node "${NODE_NAME}" cc-postuninstall/done=true
 		;;
-	*)
-		;;
+	*) ;;
 	esac
 }
 
@@ -95,6 +134,9 @@ function main() {
 	case "${action}" in
 	install)
 		install_artifacts
+		if [ "$SNAPSHOTTER" == "nydus" ]; then
+			config_containerd
+		fi
 		restart_systemd_service
 		;;
 	uninstall)
@@ -114,7 +156,6 @@ function main() {
 
 	label_node "${action}"
 
-
 	# It is assumed this script will be called as a daemonset. As a result, do
 	# not return, otherwise the daemon will restart and reexecute the script.
 	sleep infinity