From 7da1b85503bae02ac11f8a5df835f705de158d0e Mon Sep 17 00:00:00 2001
From: Pan Luo <pan.luo@ubc.ca>
Date: Wed, 15 Feb 2023 23:37:40 -0800
Subject: [PATCH]  Add cli parameter for kubeconfig and node name

This is to add support for out-cluster development and testing.
Also removed Titlefile and deployment manifest as this can be tested
outside the cluster
---
 images/taint-manager/README.md       |  15 +-
 images/taint-manager/Tiltfile        | 105 ----------
 images/taint-manager/deployment.yaml |  30 ---
 images/taint-manager/go.mod          |   2 +
 images/taint-manager/go.sum          |   3 +
 images/taint-manager/taintmanager.go | 275 ++++++++++++++++-----------
 6 files changed, 177 insertions(+), 253 deletions(-)
 delete mode 100644 images/taint-manager/Tiltfile
 delete mode 100644 images/taint-manager/deployment.yaml

diff --git a/images/taint-manager/README.md b/images/taint-manager/README.md
index 0e8dad27c1..022580a978 100644
--- a/images/taint-manager/README.md
+++ b/images/taint-manager/README.md
@@ -1,6 +1,8 @@
 # In Cluster Taint Manager
 
-To add or remove taint of a node from a in-cluster pod.
+To add or remove taint within a pod. The default mode is to run as a pod in a Kubernetes cluster. It will try to
+authenticate to K8S API server with in-cluster config (the service account token mounted inside pod). It also uses
+downward API to retrieve the node name where the pod is running on in order to change the taint.
 
 ## Compile
 
@@ -10,9 +12,16 @@ GOOS=linux GOARCH=amd64 go build -o taintmanager taintmanager.go
 
 ## Development and Debug
 
-The dev/debug environment is setup by `tilt`. To start, run `tilt up`.
+To test outside a cluster, commandline parameter `-kubeconfig` and `-node` can be specified. If `-kubeconfig` is not
+specified, taintmanager will try to load the kubeconfig from default path `HOME/.kube/config`
 
-## Test
+Using default kubeconfig:
+
+```
+taintmanager -node f6.workers.ctlt.ubc.ca -remove hub.jupyter.org/imagepulling:NoExecute
+```
+
+## Test in Cluster
 
 The `test` directory contains YAML files for deploy a pod with required permissions to run taintmanager.
 Please change `namespace` field in `clusterrolebinding.yaml` before deploying to a cluster.
diff --git a/images/taint-manager/Tiltfile b/images/taint-manager/Tiltfile
deleted file mode 100644
index 4f783ce202..0000000000
--- a/images/taint-manager/Tiltfile
+++ /dev/null
@@ -1,105 +0,0 @@
-# Welcome to Tilt!
-#   To get you started as quickly as possible, we have created a
-#   starter Tiltfile for you.
-#
-#   Uncomment, modify, and delete any commands as needed for your
-#   project's configuration.
-
-cluster_name = 'k8scluster'
-project_name = 'taintmanager'
-
-# Allow the cluster to avoid problems while having kubectl configured to talk to a remote cluster.
-allow_k8s_contexts(cluster_name)
-
-# Use custom Dockerfile for Tilt builds, which only takes locally built binary for live reloading.
-dockerfile = '''
-    FROM golang:1.19-alpine
-    RUN go install github.com/go-delve/delve/cmd/dlv@latest
-    COPY %s /usr/local/bin/%s
-    ''' % (project_name, project_name)# Build Docker image
-
-#   Tilt will automatically associate image builds with the resource(s)
-#   that reference them (e.g. via Kubernetes or Docker Compose YAML).
-#
-#   More info: https://docs.tilt.dev/api.html#api.docker_build
-#
-# docker_build('registry.example.com/my-image',
-#              context='.',
-#              # (Optional) Use a custom Dockerfile path
-#              dockerfile='./deploy/app.dockerfile',
-#              # (Optional) Filter the paths used in the build
-#              only=['./app'],
-#              # (Recommended) Updating a running container in-place
-#              # https://docs.tilt.dev/live_update_reference.html
-#              live_update=[
-#                 # Sync files from host to container
-#                 sync('./app', '/src/'),
-#                 # Execute commands inside the container when certain
-#                 # paths change
-#                 run('/src/codegen.sh', trigger=['./app/api'])
-#              ]
-# )
-
-
-# Run local commands
-#   Local commands can be helpful for one-time tasks like installing
-#   project prerequisites. They can also manage long-lived processes
-#   for non-containerized services or dependencies.
-#
-#   More info: https://docs.tilt.dev/local_resource.html
-#
-# local_resource('install-helm',
-#                cmd='which helm > /dev/null || brew install helm',
-#                # `cmd_bat`, when present, is used instead of `cmd` on Windows.
-#                cmd_bat=[
-#                    'powershell.exe',
-#                    '-Noninteractive',
-#                    '-Command',
-#                    '& {if (!(Get-Command helm -ErrorAction SilentlyContinue)) {scoop install helm}}'
-#                ]
-# )
-
-# Building binary locally.
-local_resource('%s-binary' % project_name,
-    'GOOS=linux GOARCH=amd64 go build -gcflags "all=-N -l" -o taintmanager taintmanager.go',
-    deps=[
-        './taintmanager.go',
-    ],
-)
-
-# Extensions are open-source, pre-packaged functions that extend Tilt
-#
-#   More info: https://github.com/tilt-dev/tilt-extensions
-#
-load('ext://git_resource', 'git_checkout')
-
-# Load the restart_process extension with the docker_build_with_restart func for live reloading.
-load('ext://restart_process', 'docker_build_with_restart')
-
-# Wrap a docker_build to restart the given entrypoint after a Live Update.
-docker_build(
-    'jupyterhub/' + project_name,
-    '.',
-    dockerfile_contents=dockerfile,
-    entrypoint='/go/bin/dlv --listen=0.0.0.0:50100 --api-version=2 --headless=true --only-same-user=false --accept-multiclient --check-go-version=false exec -- /usr/local/bin/taintmanager -remove jupyterhub:NoSchedule',
-    live_update=[
-        # Copy the binary so it gets restarted.
-        sync(project_name, '/usr/local/bin/%s' % project_name),
-    ],
-)
-
-# Apply Kubernetes manifests
-#   Tilt will build & push any necessary images, re-deploying your
-#   resources as they change.
-#
-#   More info: https://docs.tilt.dev/api.html#api.k8s_yaml
-#
-# k8s_yaml(['k8s/deployment.yaml', 'k8s/service.yaml'])
-# Create the deployment from YAML file path.
-k8s_yaml('deployment.yaml')
-
-# Configure the resource.
-k8s_resource(
-        project_name,
-        port_forwards = ["50100:50100"] # Set up the K8s port-forward to be able to connect to it locally.
-)
diff --git a/images/taint-manager/deployment.yaml b/images/taint-manager/deployment.yaml
deleted file mode 100644
index af82e8620a..0000000000
--- a/images/taint-manager/deployment.yaml
+++ /dev/null
@@ -1,30 +0,0 @@
-apiVersion: apps/v1
-kind: Deployment
-metadata:
-  name: taintmanager
-spec:
-  replicas: 1
-  selector:
-    matchLabels:
-      app: taintmanager
-  template:
-    metadata:
-      labels:
-        app: taintmanager
-    spec:
-      containers:
-        - image: jupyterhub/k8s-taint-manager
-          imagePullPolicy: IfNotPresent
-          name: taintmanager
-          env:
-            - name: MY_POD_NAME
-              valueFrom:
-                fieldRef:
-                  fieldPath: metadata.name
-            - name: MY_NODE_NAME
-              valueFrom:
-                fieldRef:
-                  fieldPath: spec.nodeName
-          ports:
-            - containerPort: 8080
-      restartPolicy: Always
diff --git a/images/taint-manager/go.mod b/images/taint-manager/go.mod
index aacda2c1f8..35bdcff638 100644
--- a/images/taint-manager/go.mod
+++ b/images/taint-manager/go.mod
@@ -20,12 +20,14 @@ require (
 	github.com/google/gnostic v0.6.9 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/gofuzz v1.2.0 // indirect
+	github.com/imdario/mergo v0.3.6 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/net v0.2.0 // indirect
 	golang.org/x/oauth2 v0.2.0 // indirect
 	golang.org/x/sys v0.2.0 // indirect
diff --git a/images/taint-manager/go.sum b/images/taint-manager/go.sum
index 2d904547ec..86b9d35d86 100644
--- a/images/taint-manager/go.sum
+++ b/images/taint-manager/go.sum
@@ -69,6 +69,8 @@ github.com/google/gofuzz v1.2.0 h1:xRy4A+RhZaiKjJ1bPfwQ8sedCA+YS2YcCHW6ec7JMi0=
 github.com/google/gofuzz v1.2.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
+github.com/imdario/mergo v0.3.6 h1:xTNEAn+kxVO7dTZGu0CegyqKZmoWFI0rF8UxjlB2d28=
+github.com/imdario/mergo v0.3.6/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
 github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
 github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
@@ -100,6 +102,7 @@ github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stoewer/go-strcase v1.2.0/go.mod h1:IBiWB2sKIp3wVVQ3Y035++gc+knqhUQag1KpM8ahLw8=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
diff --git a/images/taint-manager/taintmanager.go b/images/taint-manager/taintmanager.go
index 5e536acb50..941564bee0 100644
--- a/images/taint-manager/taintmanager.go
+++ b/images/taint-manager/taintmanager.go
@@ -1,132 +1,177 @@
 package main
 
 import (
-    "context"
-    "flag"
-    "fmt"
-    v1 "k8s.io/api/core/v1"
-    "k8s.io/apimachinery/pkg/api/errors"
-    metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-    "k8s.io/apimachinery/pkg/util/validation"
-    "k8s.io/client-go/kubernetes"
-    "k8s.io/client-go/rest"
-    "log"
-    "os"
-    "strings"
+	"context"
+	"flag"
+	"fmt"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/validation"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
+	"k8s.io/client-go/util/homedir"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
 )
 
+func buildConfig(kubeconfig string) (*rest.Config, error) {
+	// load kubeconfig from command line
+	if kubeconfig != "" {
+		cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err != nil {
+			return nil, err
+		}
+		log.Printf("Using kubeconfig %v.", kubeconfig)
+		return cfg, nil
+	}
+
+	// try the default location HOME/.kube/config
+	if home := homedir.HomeDir(); home != "" {
+		kubeconfig = filepath.Join(home, ".kube", "config")
+		cfg, err := clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err == nil {
+			log.Printf("Using kubeconfig %v.", kubeconfig)
+			return cfg, nil
+		}
+	}
+
+	// try in-cluster auth
+	cfg, err := rest.InClusterConfig()
+	if err != nil {
+		return nil, err
+	}
+	log.Printf("Using in-cluster config.")
+	return cfg, nil
+}
+
+func usage() {
+	log.Println("Usage: taintmanager [-add TAINT] [-remove TAINT] [-node NODE_NAME] [-kubeconfig /PATH/TO/KUBECONFIG]")
+	flag.PrintDefaults()
+	os.Exit(1)
+}
+
 func main() {
-    taintAdd := flag.String("add", "", "The taint to add")
-    taintRemove := flag.String("remove", "", "The taint to remove")
-    flag.Parse()
-    tAdd, errAdd := parseTaint(*taintAdd)
-    tRemove, errRemove := parseTaint(*taintRemove)
-
-    if errAdd != nil && errRemove != nil {
-        log.Println("Please specify at least one option -add or -remove")
-        log.Println("Usage: taintmanager [-add TAINT] [-remove TAINT]")
-        flag.PrintDefaults()
-        os.Exit(1)
-    }
-
-    // creates the in-cluster config
-    config, err := rest.InClusterConfig()
-    if err != nil {
-        panic(err.Error())
-    }
-    // creates the clientset
-    clientset, err := kubernetes.NewForConfig(config)
-    if err != nil {
-        panic(err.Error())
-    }
-
-    podName := os.Getenv("MY_POD_NAME")
-    nodeName := os.Getenv("MY_NODE_NAME")
-    log.Printf("Pod name: %v and node name %v\n", podName, nodeName)
-
-    // Get node info from API server
-    node, err := clientset.CoreV1().Nodes().Get(context.TODO(), nodeName, metav1.GetOptions{})
-    if errors.IsNotFound(err) {
-        log.Printf("Node %v not found\n", nodeName)
-    } else if statusError, isStatus := err.(*errors.StatusError); isStatus {
-        log.Printf("Error getting node %v\n", statusError.ErrStatus.Message)
-    } else if err != nil {
-        panic(err.Error())
-    } else {
-        log.Printf("Found node %v\n", node.GetName())
-        // Get existing taints
-        for k, v := range node.Spec.Taints {
-            log.Printf("%v: %v\n", k, v)
-        }
-        // Add taint
-        if errAdd == nil {
-            node.Spec.Taints = append(node.Spec.Taints, tAdd)
-            clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
-            log.Printf("Taint %v is added to node %v.", tAdd.ToString(), nodeName)
-        }
-        // Remove taint
-        if errRemove == nil {
-            for i, taint := range node.Spec.Taints {
-                if taint.Key == tRemove.Key {
-                    node.Spec.Taints = append(node.Spec.Taints[:i], node.Spec.Taints[i+1:]...)
-                    log.Printf("Taint %v is removed from node %v.", tRemove.ToString(), nodeName)
-                    clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
-                }
-            }
-        }
-    }
+	kubeconfig := flag.String("kubeconfig", "", "(optional) absolute path to the kubeconfig file")
+	nodeName := flag.String("node", "", "(optional) The name of the node to add/remove taints. Can be passed through environment variable MY_NODE_NAME")
+	taintAdd := flag.String("add", "", "The taint to add")
+	taintRemove := flag.String("remove", "", "The taint to remove")
+	flag.Parse()
+
+	// load kubeconfig
+	config, err := buildConfig(*kubeconfig)
+	if err != nil {
+		log.Printf(err.Error())
+		log.Printf("Failed to load kubeconfig")
+		usage()
+	}
+
+	// parse taints
+	tAdd, errAdd := parseTaint(*taintAdd)
+	tRemove, errRemove := parseTaint(*taintRemove)
+	if errAdd != nil && errRemove != nil {
+		log.Println("Please specify at least one option -add or -remove")
+		usage()
+	}
+
+	// check node name
+	if *nodeName == "" {
+		// try from env var
+		*nodeName = os.Getenv("MY_NODE_NAME")
+	}
+	if *nodeName == "" {
+		log.Println("Please specify the node name")
+		usage()
+	}
+	log.Printf("Node name %v\n", *nodeName)
+
+	// creates the clientset
+	clientset, err := kubernetes.NewForConfig(config)
+	if err != nil {
+		panic(err.Error())
+	}
+	node, err := clientset.CoreV1().Nodes().Get(context.TODO(), *nodeName, metav1.GetOptions{})
+	if errors.IsNotFound(err) {
+		log.Printf("Node %v not found in default namespace\n", *nodeName)
+	} else if statusError, isStatus := err.(*errors.StatusError); isStatus {
+		log.Printf("Error getting node %v\n", statusError.ErrStatus.Message)
+	} else if err != nil {
+		panic(err.Error())
+	} else {
+		log.Printf("Found node %v\n", node.GetName())
+		for k, v := range node.Spec.Taints {
+			log.Printf("%v: %v\n", k, v)
+		}
+		if errAdd == nil {
+			node.Spec.Taints = append(node.Spec.Taints, tAdd)
+			clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+			log.Printf("Taint %v is added to node %v.", tAdd.ToString(), *nodeName)
+		}
+		if errRemove == nil {
+			for i, taint := range node.Spec.Taints {
+				if taint.Key == tRemove.Key {
+					node.Spec.Taints = append(node.Spec.Taints[:i], node.Spec.Taints[i+1:]...)
+					log.Printf("Taint %v is removed from node %v.", tRemove.ToString(), *nodeName)
+					clientset.CoreV1().Nodes().Update(context.TODO(), node, metav1.UpdateOptions{})
+				}
+			}
+		}
+	}
 }
 
 // copied from https://github.com/kubernetes/kubernetes/blob/v1.25.4/pkg/util/taints/taints.go
 // parseTaint parses a taint from a string, whose form must be either
 // '<key>=<value>:<effect>', '<key>:<effect>', or '<key>'.
 func parseTaint(st string) (v1.Taint, error) {
-    var taint v1.Taint
-
-    var key string
-    var value string
-    var effect v1.TaintEffect
-
-    parts := strings.Split(st, ":")
-    switch len(parts) {
-    case 1:
-        key = parts[0]
-    case 2:
-        effect = v1.TaintEffect(parts[1])
-        if err := validateTaintEffect(effect); err != nil {
-            return taint, err
-        }
-
-        partsKV := strings.Split(parts[0], "=")
-        if len(partsKV) > 2 {
-            return taint, fmt.Errorf("invalid taint spec: %v", st)
-        }
-        key = partsKV[0]
-        if len(partsKV) == 2 {
-            value = partsKV[1]
-            if errs := validation.IsValidLabelValue(value); len(errs) > 0 {
-                return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
-            }
-        }
-    default:
-        return taint, fmt.Errorf("invalid taint spec: %v", st)
-    }
-
-    if errs := validation.IsQualifiedName(key); len(errs) > 0 {
-        return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
-    }
-
-    taint.Key = key
-    taint.Value = value
-    taint.Effect = effect
-
-    return taint, nil
+	var taint v1.Taint
+
+	var key string
+	var value string
+	var effect v1.TaintEffect
+
+	parts := strings.Split(st, ":")
+	switch len(parts) {
+	case 1:
+		key = parts[0]
+	case 2:
+		effect = v1.TaintEffect(parts[1])
+		if err := validateTaintEffect(effect); err != nil {
+			return taint, err
+		}
+
+		partsKV := strings.Split(parts[0], "=")
+		if len(partsKV) > 2 {
+			return taint, fmt.Errorf("invalid taint spec: %v", st)
+		}
+		key = partsKV[0]
+		if len(partsKV) == 2 {
+			value = partsKV[1]
+			if errs := validation.IsValidLabelValue(value); len(errs) > 0 {
+				return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
+			}
+		}
+	default:
+		return taint, fmt.Errorf("invalid taint spec: %v", st)
+	}
+
+	if errs := validation.IsQualifiedName(key); len(errs) > 0 {
+		return taint, fmt.Errorf("invalid taint spec: %v, %s", st, strings.Join(errs, "; "))
+	}
+
+	taint.Key = key
+	taint.Value = value
+	taint.Effect = effect
+
+	return taint, nil
 }
 
 func validateTaintEffect(effect v1.TaintEffect) error {
-    if effect != v1.TaintEffectNoSchedule && effect != v1.TaintEffectPreferNoSchedule && effect != v1.TaintEffectNoExecute {
-        return fmt.Errorf("invalid taint effect: %v, unsupported taint effect", effect)
-    }
+	if effect != v1.TaintEffectNoSchedule && effect != v1.TaintEffectPreferNoSchedule && effect != v1.TaintEffectNoExecute {
+		return fmt.Errorf("invalid taint effect: %v, unsupported taint effect", effect)
+	}
 
-    return nil
+	return nil
 }