refactor(webhook): refactor sidecar injection

pisa-controller/pkg/webhook/injection.go

@@ -15,10 +15,9 @@
 package webhook
 
 import (
+	"errors"
 	"fmt"
 	"net/http"
-	"os"
-	"strconv"
 	"strings"
 
 	"github.com/gin-gonic/gin"
@@ -28,176 +27,141 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/serializer"
-	"k8s.io/apimachinery/pkg/util/json"
+	"k8s.io/apimachinery/pkg/types"
 )
 
 var (
-	pisaProxyImage, pisaControllerService, pisaControllerNamespace, pisaProxyAdminListenHost, pisaProxyLoglevel string
-	pisaProxyAdminListenPort                                                                                    uint32
-)
-
-const (
-	SidecarNamePisaProxy            = "pisa-proxy"
-	EnvPisaProxyAdminListenHost     = "PISA_PROXY_ADMIN_LISTEN_HOST"
-	EnvPisaProxyAdminListenPort     = "PISA_PROXY_ADMIN_LISTEN_PORT"
-	EnvPisaProxyLoglevel            = "PISA_PROXY_ADMIN_LOG_LEVEL"
-	DefaultPisaProxyAdminListenHost = "0.0.0.0"
-	DefaultPisaProxyAdminListenPort = 5591
-	DefaultPisaProxyLoglevel        = "INFO"
+	runtimeScheme = runtime.NewScheme()
+	codecs        = serializer.NewCodecFactory(runtimeScheme)
+	deserializer  = codecs.UniversalDeserializer()
 )
 
-func init() {
-	pisaProxyImage = os.Getenv("PISA_PROXY_IMAGE")
-	if pisaProxyImage == "" {
-		pisaProxyImage = "pisanixio/proxy:latest"
-	}
-	pisaControllerService = os.Getenv("PISA_CONTROLLER_SERVICE")
-	pisaControllerNamespace = os.Getenv("PISA_CONTROLLER_NAMESPACE")
-	if host := os.Getenv(EnvPisaProxyAdminListenHost); host == "" {
-		pisaProxyAdminListenHost = DefaultPisaProxyAdminListenHost
-	} else {
-		pisaProxyAdminListenHost = host
+func Injection(ctx *gin.Context) {
+	rawData, err := ctx.GetRawData()
+	if err != nil || len(rawData) == 0 {
+		log.Error("get body raw data error.")
+		ctx.JSON(http.StatusBadRequest, NewV1AdmissionResponseFromError(err))
+		return
 	}
-	if port, err := strconv.Atoi(os.Getenv(EnvPisaProxyAdminListenPort)); port <= 0 || err != nil {
-		pisaProxyAdminListenPort = DefaultPisaProxyAdminListenPort
-	} else {
-		pisaProxyAdminListenPort = uint32(port)
+
+	ar := &v1.AdmissionReview{}
+	if _, _, err := deserializer.Decode(rawData, nil, ar); err != nil {
+		log.Errorf("can not decode body to AdmissionReview: %v", err)
+		ctx.JSON(http.StatusBadRequest, NewV1AdmissionResponseFromError(err))
+		return
 	}
-	if lv := os.Getenv(EnvPisaProxyLoglevel); lv == "" {
-		pisaProxyLoglevel = DefaultPisaProxyLoglevel
-	} else {
-		pisaProxyLoglevel = lv
+
+	if err = injection(ar); err != nil {
+		log.Errorf("injection error: %v", err)
+		ctx.JSON(http.StatusInternalServerError, NewV1AdmissionResponseFromError(err))
+		return
 	}
+
+	log.Infof("mutating Success %s/%s", ar.Request.Namespace, ar.Request.Name)
+	ctx.JSON(http.StatusOK, ar)
 }
 
-const (
-	podsSidecarPatch = `[
-		{
-			"op":"add", 
-			"path":"/spec/containers/-",
-			"value":{
-				"image":"%v",
-				"name":"%s",
-				"args": ["sidecar"],
-				"ports": [
-					{
-						"containerPort": %d,
-						"name": "pisa-admin",
-						"protocol": "TCP"
-					}	
-				],
-				"resources":{},
-				"env": [
-					{
-						"name": "PISA_CONTROLLER_SERVICE",
-						"value": "%s"
-					},{
-						"name": "PISA_CONTROLLER_NAMESPACE",
-						"value": "%s"
-					},{
-						"name": "PISA_DEPLOYED_NAMESPACE",
-						"value": "%s"
-					},{
-						"name": "PISA_DEPLOYED_NAME",
-						"value": "%s"
-					},{
-						"name": "PISA_PROXY_ADMIN_LISTEN_HOST",
-						"value": "%s"
-					},{
-						"name": "PISA_PROXY_ADMIN_LISTEN_PORT",
-						"value": "%d"
-					},{
-						"name": "PISA_PROXY_ADMIN_LOG_LEVEL",
-						"value": "%s"
-					}
-				]
-			}
-		}
-	]`
-)
+func injection(ar *v1.AdmissionReview) error {
+	pod := retrievePodFromAdmissionRequest(ar.Request)
+	if pod == nil {
+		return errors.New("retrieve pod from admission request error")
+	}
 
-var (
-	runtimeScheme = runtime.NewScheme()
-	codecs        = serializer.NewCodecFactory(runtimeScheme)
-	deserializer  = codecs.UniversalDeserializer()
-)
+	patch := buildPatch(pod)
 
-type PodInfo struct {
-	Metadata struct {
-		GenerateName string `json:"generateName"`
-	} `json:"metadata"`
+	resp := buildPodPatchResponse(ar.Request.UID, pod, patch)
+	if resp != nil {
+		return errors.New("build pod patch response error")
+	}
+
+	ar.Response = resp
+	return nil
 }
 
-func InjectSidecar(ctx *gin.Context) {
-	rawData, err := ctx.GetRawData()
-	if err != nil {
-		log.Error("get body raw data error.")
-		ctx.JSON(http.StatusBadRequest, toV1AdmissionResponse(err))
-	}
-	if len(rawData) == 0 {
-		log.Errorf("get a empty body!")
-		return
+func retrievePodFromAdmissionRequest(req *v1.AdmissionRequest) *corev1.Pod {
+	gvr := metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"}
+	if req.Resource != gvr {
+		log.Errorf("expect resource to be %s", gvr)
+		return nil
 	}
 
-	ar := v1.AdmissionReview{}
-	if _, _, err := deserializer.Decode(rawData, nil, &ar); err != nil {
-		log.Errorf("can't decode body to AdmissionReview: %v", err)
-		ctx.JSON(http.StatusBadRequest, toV1AdmissionResponse(err))
-		return
+	pod := &corev1.Pod{}
+	if _, _, err := deserializer.Decode(req.Object.Raw, nil, pod); err != nil {
+		log.Errorf("retrieve from raw object error: %s", err)
+		return nil
 	}
-	shouldPatchPod := func(pod *corev1.Pod) bool {
-		return !hasContainer(pod.Spec.Containers, SidecarNamePisaProxy)
+
+	return pod
+}
+
+func buildPatch(pod *corev1.Pod) string {
+	var pisaProxyDeployedName string
+
+	if pod.OwnerReferences == nil || pod.OwnerReferences[0].Kind == "" {
+		pisaProxyDeployedName = pod.ObjectMeta.Name
+	} else {
+		pisaProxyDeployedName = getPisaProxyDeployedNameFromPod(pod.OwnerReferences[0].Kind, pod.ObjectMeta.GenerateName)
 	}
-	podinfo := &PodInfo{}
-	_ = json.Unmarshal(ar.Request.Object.Raw, podinfo)
-	podSlice := strings.Split(podinfo.Metadata.GenerateName, "-")
-	podSlice = podSlice[:len(podSlice)-2]
 
-	patch := fmt.Sprintf(podsSidecarPatch,
+	return fmt.Sprintf(podsSidecarPatch,
 		pisaProxyImage,
-		SidecarNamePisaProxy,
+		pisaProxyContainerName,
 		pisaProxyAdminListenPort,
 		pisaControllerService,
 		pisaControllerNamespace,
-		ar.Request.Namespace,
-		strings.Join(podSlice, "-"),
+		pisaProxyDeployedName,
 		pisaProxyAdminListenHost,
 		pisaProxyAdminListenPort,
 		pisaProxyLoglevel,
 	)
-	ar.Response = applyPodPatch(ar, shouldPatchPod, patch)
-	log.Infof("mutating Success %v", patch)
-
-	ctx.JSON(http.StatusOK, ar)
 }
 
-func applyPodPatch(ar v1.AdmissionReview, shouldPatchPod func(*corev1.Pod) bool, patch string) *v1.AdmissionResponse {
-	log.Info("mutating pods")
-	podResource := metav1.GroupVersionResource{Group: "", Version: "v1", Resource: "pods"}
-	if ar.Request.Resource != podResource {
-		log.Errorf("expect resource to be %s", podResource)
-		return nil
+func getPisaProxyDeployedNameFromPod(kind, generatedName string) string {
+	var name string
+
+	switch kind {
+	case "ReplicaSet":
+		fallthrough
+	case "Job":
+		name = getPodNameFromGeneratedName(generatedName, 2)
+	case "StatefulSet":
+		fallthrough
+	case "DaemonSet":
+		name = getPodNameFromGeneratedName(generatedName, 1)
 	}
+	return name
+}
+
+func getPodNameFromGeneratedName(generatedName string, offset uint) string {
+	podSlice := strings.Split(generatedName, "-")
+	l := len(podSlice)
 
-	raw := ar.Request.Object.Raw
-	pod := corev1.Pod{}
-	if _, _, err := deserializer.Decode(raw, nil, &pod); err != nil {
-		log.Error(err)
-		return toV1AdmissionResponse(err)
+	if l > int(offset) {
+		podSlice = podSlice[:l-int(offset)]
+		return strings.Join(podSlice, "-")
 	}
-	log.Infof("pod %v", pod)
-	reviewResponse := v1.AdmissionResponse{}
-	reviewResponse.UID = ar.Request.UID
-	reviewResponse.Allowed = true
-	if shouldPatchPod(&pod) {
-		reviewResponse.Patch = []byte(patch)
+	return generatedName
+}
+
+func buildPodPatchResponse(UID types.UID, pod *corev1.Pod, patch string) *v1.AdmissionResponse {
+	log.Info("mutating pods")
+
+	if !hasContainer(pod.Spec.Containers, pisaProxyContainerName) {
+		resp := &v1.AdmissionResponse{
+			UID:     UID,
+			Allowed: true,
+			Patch:   []byte(patch),
+		}
+
 		pt := v1.PatchTypeJSONPatch
-		reviewResponse.PatchType = &pt
+		resp.PatchType = &pt
+		return resp
 	}
-	return &reviewResponse
+
+	return nil
 }
 
-func toV1AdmissionResponse(err error) *v1.AdmissionResponse {
+func NewV1AdmissionResponseFromError(err error) *v1.AdmissionResponse {
 	return &v1.AdmissionResponse{
 		Result: &metav1.Status{
 			Message: err.Error(),

pisa-controller/pkg/webhook/injection_test.go

@@ -0,0 +1,50 @@
+// Copyright 2022 SphereEx Authors
+//
+// 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.
+
+package webhook
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_GetPisaProxyDeployedNameFromPod(t *testing.T) {
+	cases := []struct {
+		name          string
+		generatedName string
+		kind          string
+		exp           string
+		msg           string
+	}{
+		{
+			name:          "a-b-c-6668ddbf55-abcde",
+			generatedName: "a-b-c-6668ddbf55-",
+			kind:          "ReplicaSet",
+			exp:           "a-b-c",
+			msg:           "Deployment or CronJob name should be equal",
+		},
+		{
+			name:          "a-b-0",
+			generatedName: "a-b-",
+			kind:          "StatefulSet",
+			exp:           "a-b",
+			msg:           "StatefulSet, Job or DaemonSet name should be equal",
+		},
+	}
+
+	for _, c := range cases {
+		assert.Equal(t, c.exp, getPisaProxyDeployedNameFromPod(c.kind, c.generatedName), c.msg)
+	}
+}

pisa-controller/pkg/webhook/server.go

@@ -43,7 +43,7 @@ func Handler() http.Handler {
 	// NOTE: there is not API path in this request
 	// TODO: test trailing slash
 	g.GET("", ApiCheck)
-	g.POST("/mutate", InjectSidecar)
+	g.POST("/mutate", Injection)
 
 	return r
 }