Added L0 testcases and logging

Dockerfile

@@ -11,5 +11,5 @@ RUN go build -o main .
 
 COPY agentpods/* agentpods/
 
-EXPOSE 8082
+EXPOSE 8080
 CMD ["/app/main"]

agent-custom/Dockerfile

@@ -32,8 +32,14 @@ RUN BUILDKITVERSION=$(curl -s "https://api.github.com/repos/moby/buildkit/releas
      && mv ./bin/buildctl /usr/local/bin/ && mv ./bin/buildkitd /usr/local/bin/ && mv ./bin/buildkit-runc /usr/local/bin/
 
 # setting work directory
-WORKDIR /azp
+WORKDIR /azp/agent
 
+RUN VSTSVERSION=$(curl -s "https://api.github.com/repos/microsoft/azure-pipelines-agent/releases/latest" | jq -r .tag_name[1:]) \
+     && echo $VSTSVERSION \
+     && curl -LO https://vstsagentpackage.azureedge.net/agent/${VSTSVERSION}/vsts-agent-linux-x64-${VSTSVERSION}.tar.gz \
+     && tar -zxvf vsts-agent-linux-x64-${VSTSVERSION}.tar.gz
+
+WORKDIR /azp
 COPY ./start.sh .
 RUN chmod +x start.sh
 

agent-custom/start.sh

@@ -2,41 +2,52 @@
 
 set -e
 
-rm -rf /azp/agent
-mkdir /azp/agent
-cd /azp/agent
+AGENT_FOLDER="agent"
+AZP_AGENT_VERSION="$(cat /vsts/agent/.agentVersion)"
+AGENTVERSION="$(curl -s "https://api.github.com/repos/microsoft/azure-pipelines-agent/releases/latest" | jq -r .tag_name[1:])"
 
-# Read download URL from the secret
-AZP_DOWNLOAD_URL="$(cat /vsts/agent/.url)"
+echo "Latest agent release version $AGENTVERSION" 
+echo "Agent Version from request $AZP_AGENT_VERSION"
 
-# Download the requested agent, else fall back to the version that was default when this was released
-if [ -z "$AZP_DOWNLOAD_URL" ]; then
-  AZP_DOWNLOAD_URL="https://vstsagentpackage.azureedge.net/agent/2.158.0/vsts-agent-linux-x64-2.158.0.tar.gz"
+if [ $AGENTVERSION != $AZP_AGENT_VERSION ]; then
+
+  echo "Downloading the Azure Pipelines agent from agent request....."
+
+  # Deleting previous downloaded folder
+  rm -rf /azp/$AGENT_FOLDER
+
+  # setting the agent folder name to agent version from request
+  AGENT_FOLDER=AZP_AGENT_VERSION
+  mkdir /azp/$AGENT_FOLDER
+  cd /azp/$AGENT_FOLDER
+
+  # Read download URL from the secret
+  AZP_DOWNLOAD_URL="$(cat /vsts/agent/.url)"
+
+  # Download the requested agent, else fall back to the version that was default when this was released
+  if [ -z "$AZP_DOWNLOAD_URL" ]; then
+    AZP_DOWNLOAD_URL="https://vstsagentpackage.azureedge.net/agent/2.158.0/vsts-agent-linux-x64-2.158.0.tar.gz"
+  fi
+  
+  echo "Installing Azure Pipelines agent from the agent request..."
+
+  curl -LsS $AZP_DOWNLOAD_URL | tar -xz & wait $!
+
+else
+  echo "Using Azure Pipelines agent downloaded from the agent custom image..."
+  cd /azp/$AGENT_FOLDER
 fi
 
-print_header() {
-  lightcyan='\033[1;36m'
-  nocolor='\033[0m'
-  echo -e "${lightcyan}$1${nocolor}"
-}
-
-# Let the agent ignore the token env variables
-
-print_header "1. Downloading and installing Azure Pipelines agent from the agent request..."
-
-curl -LsS $AZP_DOWNLOAD_URL | tar -xz & wait $!
-
 source ./env.sh
 
 trap 'cleanup; exit 130' INT
 trap 'cleanup; exit 143' TERM
 
-print_header "2. Mounting the .agent and .credentials files from a different source"
-ln -fs /vsts/agent/.agent /azp/agent/.agent
-ln -fs /vsts/agent/.credentials /azp/agent/.credentials
-# ln -fs /vsts/agent/.credentials_rsaparams /azp/agent/.credentials_rsaparams
+echo "Mounting the .agent and .credentials files from a different source"
+ln -fs /vsts/agent/.agent /azp/$AGENT_FOLDER/.agent
+ln -fs /vsts/agent/.credentials /azp/$AGENT_FOLDER/.credentials
 
-print_header "3. Running Azure Pipelines agent..."
+echo "Running Azure Pipelines agent..."
 
 # `exec` the node runtime so it's aware of TERM and INT signals
 # AgentService.js understands how to handle agent self-update and restart

agentpods/agent-lean-dind.yaml

@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
   - name: vsts-agent
-    image: prebansa/myagent:v5.9
+    image: prebansa/myagent:v5.16
     env:
     - name: DOCKER_HOST
       value: tcp://localhost:2375
@@ -17,14 +17,4 @@ spec:
     volumeMounts:
     - name: agent-creds
       mountPath: "/vsts/agent"
-      readOnly: true
-  - name: dind-daemon
-    image: docker:18.09.6-dind
-    securityContext:
-      privileged: true
-    volumeMounts:
-    - name: agent-pv-storage
-      mountPath: /var/lib/docker
-  volumes:
-  - name: agent-pv-storage
-    emptyDir: {}
+      readOnly: true

consistent_hash_test.go

@@ -0,0 +1,29 @@
+package main
+
+import (
+	"testing"
+)
+
+func TestConsistentHashShouldReturnCorrectValueBasedOnKey(t *testing.T) {
+	nodes := []string{"buildkitd-0", "buildkitd-1"}
+
+	chosen:= ComputeConsistentHash(nodes, "microsoft/k8spoolprovider.")
+	
+	chosen2:= ComputeConsistentHash(nodes, "microsoft/k8spoolproviderroot/src.")
+	if (chosen != "buildkitd-1" || chosen2 != "buildkitd-0"){
+		t.Errorf("Consistent Hashing failed")
+	}
+
+}
+
+func TestConsistentHashShouldReturnSameValueBasedOnKey(t *testing.T) {
+	nodes := []string{"buildkitd-0", "buildkitd-1"}
+
+	chosen:= ComputeConsistentHash(nodes, "microsoft/k8spoolprovider.")
+	
+	chosen2:= ComputeConsistentHash(nodes, "microsoft/k8spoolprovider.")
+	if (chosen != "buildkitd-1" || chosen2 != "buildkitd-1"){
+		t.Errorf("Consistent Hashing failed")
+	}
+
+}

consistenthashtest.go

@@ -1,15 +0,0 @@
-package main
-
-import (
-	"testing"
-    "fmt"
-)
-
-func TestConsistentHash(t *testing.T) {
-	nodes := []string{"buildkitd-0", "buildkitd-1"}
-
-	chosen:= ComputeConsistentHash(nodes, "microsoft/k8spoolprovider.")
-	fmt.Println("Selected Node", chosen)
-	chosen2:= ComputeConsistentHash(nodes, "microsoft/k8spoolproviderroot/src.")
-	fmt.Println("Selected Node", chosen2)
-}

helm/k8s-poolprovider/templates/deployment.yaml

@@ -16,5 +16,11 @@ spec:
       - image: {{ .Values.image.repository }}:{{ .Values.image.tag }}
         name: {{ .Values.app.name }}
         command: ["/app/main"]
+        env:
+        - name: VSTS_SECRET
+          valueFrom:
+            secretKeyRef:
+              name: vsts
+              key: VSTS_SECRET
         ports:
         - containerPort: 8080

hmac.go

@@ -3,19 +3,21 @@ package main
 import (
 	"os"
 	"hash"
-
+    "encoding/hex"
 	"crypto/sha512"
 	"crypto/hmac"
+	"log"
 )
 
-func ComputeHash(message string) [] byte {
+func ComputeHash(message string) string {
 	hashAlgorithm := GetHashAlgorithm()
 	if(hashAlgorithm == nil) {
-		return nil;
+		return "";
 	}
 
 	hashAlgorithm.Write([]byte(message))
-	hashedMessage := hashAlgorithm.Sum(nil)
+	
+	hashedMessage := hex.EncodeToString(hashAlgorithm.Sum(nil))
 
 	return hashedMessage
 }
@@ -28,8 +30,13 @@ func ValidateHash(message, inputHmac string) bool {
 
 	hashAlgorithm.Write([]byte(message))
 	expectedMAC := hashAlgorithm.Sum(nil)
-	
-	return hmac.Equal([]byte(inputHmac), expectedMAC)
+
+	//return hmac.Equal(inputHMACbytearr, expectedMAC)
+	inputHMacArr, err := hex.DecodeString(inputHmac)
+	if err != nil {
+		log.Fatal(err)
+	}
+	return hmac.Equal(inputHMacArr, expectedMAC)
 }
 
 func GetHashAlgorithm() hash.Hash {

hmac_test.go

@@ -0,0 +1,32 @@
+package main
+
+import (
+	"testing"
+	"os"
+)
+
+func TestValidateAndComputeHash(t *testing.T) {
+
+	os.Setenv("VSTS_SECRET", "sharedsecret1234")
+	str := ComputeHash("teststring")
+	
+	check := ValidateHash("teststring",str)
+	if (check == false){
+		t.Errorf("Hmac validation failed")
+	}
+
+}
+func TestValidateHashShouldReturnFalse(t *testing.T) {
+
+	os.Setenv("VSTS_SECRET", "sharedsecret12345")
+	str := ComputeHash("teststring")
+
+	// changing secret value
+	os.Setenv("VSTS_SECRET", "sharedsecret1234")
+	
+	check := ValidateHash("teststring",str)
+	if (check == true){
+		t.Errorf("Hmac validation failed")
+	}
+
+}

kubernetes-auth.go

@@ -10,6 +10,8 @@ import (
 	"k8s.io/client-go/kubernetes/fake"
 )
 
+var client k8s
+
 // Gets the application client set. This can be used to initialize the various Kubernetes clients.
 // Uses in cliuster configuration when app is running inside the cluster, or kubeconfig file from
 // home directory when running in development mode.
@@ -63,12 +65,15 @@ func homeDir() string {
 }
 
 func CreateClientSet() *k8s {
-	var client k8s
-	//client = k8s{}
-	testingMode := os.Getenv("TESTING")
 	
-	if testingMode == "True" {
-		client.clientset = fake.NewSimpleClientset()
+	//client = k8s{}
+	testingMode := os.Getenv("COUNTTEST")
+	
+	if testingMode == "1" || testingMode == "2" {
+		if testingMode == "1" {
+			client.clientset = fake.NewSimpleClientset()
+			os.Setenv("COUNTTEST","2")
+		}
 	} else {
 	  	cs, _ := GetClientSet()
 	  	client.clientset = cs

kubernetes.go

@@ -6,7 +6,7 @@ import (
 	"io/ioutil"
 
 	"github.com/ghodss/yaml"
-
+    "log"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
@@ -28,25 +28,35 @@ const agentIdLabel = "AgentId"
 func CreatePod(agentRequest AgentRequest) AgentProvisionResponse {
 	cs := CreateClientSet()
 
+	log.Println("Starting pod creation")
 	var response AgentProvisionResponse
 
+	log.Println("create secret called")
 	secret := createSecret(cs, agentRequest)
-	pod, err := getAgentSpecification(agentRequest.AgentId)
+
+	// parse agent version
+	agentVersion := agentRequest.AgentConfiguration.AgentVersion
+
+	log.Println("Fetch the agent yaml to be applied")
+	pod, err := getAgentSpecification(agentRequest.AgentId, agentVersion)
+
 	if err != nil {
 		return getFailureResponse(response, err)
 	}
 
 	// Mount the secrets as a volume
 	pod.Spec.Volumes = append(pod.Spec.Volumes, *getSecretVolume(secret.Name))
-
-	//append(p1.Spec.Containers[0].Env, v1.EnvVar{Name: "VSTS_TOKEN", Value: token})
+	log.Println("secrets mounted as volume")
 
 	podClient := cs.clientset.CoreV1().Pods("azuredevops")
+
 	_, err2 := podClient.Create(pod)
 	if err2 != nil {
 		return getFailureResponse(response, err)
 	}
 
+	log.Println("Pod creation done")
+
 	response.Accepted = true
 	response.ResponseType = "Success"
 	return response
@@ -66,7 +76,7 @@ func GetBuildKitPod(key string) PodResponse {
 	if err2 != nil {
 		return getFailure(response, err2)
 	}
-
+	log.Println("Fetched list of pods configured as buildkit stateful pods")
 	var nodes []string
 
 	for _, items := range podlist.Items {
@@ -75,7 +85,8 @@ func GetBuildKitPod(key string) PodResponse {
 			nodes = append(nodes, s)
 		}
 	}
-
+	
+	log.Println("Fetching the target pod using consistent hash")
 	chosen:= ComputeConsistentHash(nodes, key)
 	response.Status = "success"
 	response.Message = chosen
@@ -93,31 +104,33 @@ func DeletePodWithAgentId(agentId string) PodResponse {
 	// Get the secret with this agentId
 	secrets, _ := secretClient.List(metav1.ListOptions{LabelSelector: agentIdLabel + "=" + agentId})
 	if secrets == nil || len(secrets.Items) == 0 {
-		return getFailure(response, errors.New("Could not find running pod with AgentId"+agentId))
+		return getFailure(response, errors.New("Could not find secret with AgentId "+agentId))
 	}
 	
 	// Get the pod with this agentId
 	pods, _ := podClient.List(metav1.ListOptions{LabelSelector: agentIdLabel + "=" + agentId})
 	if pods == nil || len(pods.Items) == 0 {
-		return getFailure(response, errors.New("Could not find running pod with AgentId"+agentId))
+		return getFailure(response, errors.New("Could not find running pod with AgentId "+agentId))
 	}
 
 	err1 := secretClient.Delete(secrets.Items[0].GetName(), &metav1.DeleteOptions{})
 	if err1 != nil {
 		return getFailure(response, err1)
 	}
+	log.Println("Delete agent secret done")
 
 	err2 := podClient.Delete(pods.Items[0].GetName(), &metav1.DeleteOptions{})
 	if err2 != nil {
 		return getFailure(response, err2)
 	}
+	log.Println("Delete agent pod done")
 
 	response.Status = "success"
 	response.Message = "Deleted " + pods.Items[0].GetName() + " and secret "+ secrets.Items[0].GetName()
 	return response
 }
 
-func getAgentSpecification(agentId string) (*v1.Pod, error) {
+func getAgentSpecification(agentId string, agentVersion string) (*v1.Pod, error) {
 	// Defaulting to use the DIND image, the podname can be exposed as a parameter and the user can then select which
 	// image will be used to create the agent.
 	podname := "agent-lean-dind"
@@ -138,12 +151,17 @@ func getAgentSpecification(agentId string) (*v1.Pod, error) {
 		})
 	}
 
+	/* uncomment this when using CI pipeline synched with agent releas
+	p1.Spec.Containers[0].Image = "prebansa/myagent:"+agentVersion
+	*/
+
 	return &p1, nil
 }
 
 func getAgentSecret() *v1.Secret {
 	var secret v1.Secret
 
+	log.Println("Reading agent-secret.yaml")
 	dat, _ := ioutil.ReadFile("agentpods/agent-secret.yaml")
 	var secretYaml = string(dat)
 	yaml.Unmarshal([]byte(secretYaml), &secret)
@@ -153,6 +171,8 @@ func getAgentSecret() *v1.Secret {
 
 func createSecret(cs *k8s, request AgentRequest) *v1.Secret {
 	secret := getAgentSecret()
+
+	log.Println("parsing secret data from agent request")
 	agentSettings, _ := json.Marshal(request.AgentConfiguration.AgentSettings)
 	agentCredentials, _ := json.Marshal(request.AgentConfiguration.AgentCredentials)
 
@@ -166,6 +186,7 @@ func createSecret(cs *k8s, request AgentRequest) *v1.Secret {
 	secret.Data[".agent"] = ([]byte(string(agentSettings)))
 	secret.Data[".credentials"] = ([]byte(string(agentCredentials)))
 	secret.Data[".url"] = ([]byte(request.AgentConfiguration.AgentDownloadUrls["linux-x64"]))
+	secret.Data[".agentVersion"] = ([]byte(request.AgentConfiguration.AgentVersion))
 
 	secretClient := cs.clientset.CoreV1().Secrets("azuredevops")
 	secret2, err := secretClient.Create(secret)
@@ -173,7 +194,7 @@ func createSecret(cs *k8s, request AgentRequest) *v1.Secret {
 	if err != nil {
 		secret2.Name = "newname"
 	}
-
+	log.Println("Secret creation done")
 	return secret2
 }
 

kubernetes_test.go

@@ -1,23 +1,127 @@
 package main
 
 import (
-	"fmt"
 	"testing"
 	"os"
+	v1 "k8s.io/api/core/v1"
+	"io/ioutil"
+
+	"github.com/ghodss/yaml"
 )
 
 func TestCreatePod(t *testing.T) {
 	var agentrequest AgentRequest
 	agentrequest.AgentId = "1"
-
-	os.Setenv("TESTING","True")
+	SetTestingEnvironmentVariables()
 
 	//clientset := fake.NewSimpleClientset()
-	pod := CreatePod(agentrequest)
-
+	testPod := CreatePod(agentrequest)
 	//secret := createSecret(&clientset,agentrequest)
-	if (pod.Accepted == true){
-		fmt.Println("Pod created", pod)
+	if (testPod.Accepted != true){
+		t.Errorf("Pod creation failed")
 	}
 
+}
+
+func TestCreateSecret(t *testing.T) {
+	var agentrequest AgentRequest
+	agentrequest.AgentId = "1"
+	SetTestingEnvironmentVariables()
+	cs := CreateClientSet()
+
+	testSecret := createSecret(cs,agentrequest)
+
+	if (testSecret.Name == "newname"){
+		t.Errorf("Secret creation failed")
+	}
+
+}
+
+func TestDeletePodShouldPassIfMatchingAgentIdinAgentRequest(t *testing.T) {
+	var agentrequest AgentRequest
+	agentrequest.AgentId = "1"
+	SetTestingEnvironmentVariables()
+
+	testPod := CreatePod(agentrequest)
+
+	if (testPod.Accepted != true){
+		t.Errorf("Pod creation failed")
+	}
+
+	testDeletepod := DeletePodWithAgentId(agentrequest.AgentId);
+	if (testDeletepod.Status != "success"){
+		t.Errorf("Pod deletion failed")
+	}
+
+}
+
+func TestDeletePodShouldFailIfNotMatchingAgentIdinAgentRequest(t *testing.T) {
+	var agentrequest AgentRequest
+	agentrequest.AgentId = "1"
+	SetTestingEnvironmentVariables()
+
+	testPod := CreatePod(agentrequest)
+	if (testPod.Accepted != true){
+		t.Errorf("Pod creation failed")
+	}
+
+	//Trying to delete pod with AgentId = 2
+	testDeletepod := DeletePodWithAgentId("2");
+	if (testDeletepod.Status != "fail"){
+		t.Errorf("Pod deletion passed but should have failed")
+	}
+
+}
+
+func TestGetBuildPodShouldReturnEmptyStringIfNoBuildKitPodPresent(t *testing.T) {
+	var agentrequest AgentRequest
+	agentrequest.AgentId = "1"
+	SetTestingEnvironmentVariables()
+
+	testPod := CreatePod(agentrequest)
+	if (testPod.Accepted != true){
+		t.Errorf("Pod creation failed")
+	}
+
+	testDeletepod := GetBuildKitPod("test");
+	if (testDeletepod.Message != ""){
+		t.Errorf("Test failed")
+	}
+
+}
+
+func TestGetBuildPodShouldReturnBuildKitPodNameIfPresent(t *testing.T) {
+	var agentrequest AgentRequest
+	agentrequest.AgentId = "1"
+	SetTestingEnvironmentVariables()
+
+    CreateDummyBuildKitPod()	
+	
+	testGetBuildpod := GetBuildKitPod("test");
+	if (testGetBuildpod.Message == ""){
+		t.Errorf("Test failed")
+	}
+
+}
+
+func CreateDummyBuildKitPod() {
+	cs := CreateClientSet()
+	var p1 v1.Pod
+
+	podname := "agent-lean-dind"
+
+	dat, _ := ioutil.ReadFile("agentpods/" + podname + ".yaml")
+	var podYaml = string(dat)
+	_ = yaml.Unmarshal([]byte(podYaml), &p1)
+		p1.SetLabels(map[string]string{
+		"role": "buildkit",
+		})
+	
+	p1.ObjectMeta.Name = "buildkitd-0"
+	podClient := cs.clientset.CoreV1().Pods("azuredevops")
+	_, _ = podClient.Create(&p1)
+}
+
+func SetTestingEnvironmentVariables() {
+	os.Setenv("COUNTTEST","1")
 }

main.go

@@ -26,7 +26,9 @@ func main() {
 func AcquireAgentHandler(resp http.ResponseWriter, req *http.Request) {
 	// HTTP method should be POST and the HMAC header should be valid
 	if req.Method == http.MethodPost {
+		log.Println("Recieved agent acquire request ....")
 		if isRequestHmacValid(req) {
+			log.Println("Hmac Validated")
 			var agentRequest AgentRequest
 
 			requestBody, err := ioutil.ReadAll(req.Body)
@@ -37,6 +39,7 @@ func AcquireAgentHandler(resp http.ResponseWriter, req *http.Request) {
 			} else if agentRequest.AgentId == "" {
 				writeJsonResponse(resp, http.StatusBadRequest, GetError(NoAgentIdError))
 			} else {
+				log.Println("Calling create pod")
 				var pods = CreatePod(agentRequest)
 				writeJsonResponse(resp, http.StatusCreated, pods)
 			}
@@ -49,8 +52,11 @@ func AcquireAgentHandler(resp http.ResponseWriter, req *http.Request) {
 }
 
 func ReleaseAgentHandler(resp http.ResponseWriter, req *http.Request) {
+
 	if req.Method == http.MethodPost {
+		log.Println("Recieved release agent request ....")
 		if isRequestHmacValid(req) {
+			log.Println("Hmac Validated")
 			var agentRequest ReleaseAgentRequest
 			requestBody, _ := ioutil.ReadAll(req.Body)
 			json.Unmarshal(requestBody, &agentRequest)
@@ -58,6 +64,7 @@ func ReleaseAgentHandler(resp http.ResponseWriter, req *http.Request) {
 			if agentRequest.AgentId == "" {
 				writeJsonResponse(resp, http.StatusBadRequest, GetError(NoAgentIdError))
 			} else {
+				log.Println("Calling delete pod")
 				var pods = DeletePodWithAgentId(agentRequest.AgentId)
 				writeJsonResponse(resp, http.StatusCreated, pods)
 			}
@@ -76,9 +83,11 @@ func EmptyResponeHandler(resp http.ResponseWriter, req *http.Request) {
 
 func GetBuildPodHandler(resp http.ResponseWriter, req *http.Request) {
 
+	log.Println("Recieved GetBuildPod request ....")
 	if req.Method == http.MethodGet {
 		keyHeader := "key"
 		headerVal := req.Header.Get(keyHeader)
+		log.Println("Calling getbuildkit pod")
 		var pods = GetBuildKitPod(headerVal)
 		writeJsonResponse(resp, http.StatusCreated, pods)
 	} else {

manifests/kubernetes.yaml

@@ -59,9 +59,15 @@ spec:
         tier: frontend
     spec:
       containers:
-      - image: prebansa/k8spoolprovider:v6.2
+      - image: prebansa/k8s-poolprovider:v1.23
         name: k8s-poolprovider
         command: ["/app/main"]
+        env:
+        - name: VSTS_SECRET
+          valueFrom:
+            secretKeyRef:
+              name: vsts
+              key: VSTS_SECRET
         ports:
         - containerPort: 8080
 ---