add tests for context timeout

Signed-off-by: Anumita <ansheno@microsoft.com>
kubeguard · Jan 16, 2023 · 2954694 · 2954694
1 parent d5cf063
commit 2954694
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 28 deletions.
diff --git a/authz/providers/azure/azure_test.go b/authz/providers/azure/azure_test.go
@@ -71,7 +71,7 @@ func clientSetup(serverUrl, mode string) (*Authorizer, error) {
 	return c, nil
 }
 
-func serverSetup(loginResp, checkaccessResp string, loginStatus, checkaccessStatus int) (*httptest.Server, error) {
+func serverSetup(loginResp, checkaccessResp string, loginStatus, checkaccessStatus int, sleepFor time.Duration) (*httptest.Server, error) {
 	listener, err := net.Listen("tcp", "127.0.0.1:")
 	if err != nil {
 		return nil, err
@@ -85,6 +85,7 @@ func serverSetup(loginResp, checkaccessResp string, loginStatus, checkaccessStat
 	}))
 
 	m.Post("/arm/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(sleepFor)
 		w.WriteHeader(checkaccessStatus)
 		_, _ = w.Write([]byte(checkaccessResp))
 	}))
@@ -98,8 +99,8 @@ func serverSetup(loginResp, checkaccessResp string, loginStatus, checkaccessStat
 	return srv, nil
 }
 
-func getServerAndClient(t *testing.T, loginResp, checkaccessResp string, checkaccessStatus int) (*httptest.Server, *Authorizer, authz.Store) {
-	srv, err := serverSetup(loginResp, checkaccessResp, http.StatusOK, checkaccessStatus)
+func getServerAndClient(t *testing.T, loginResp, checkaccessResp string, checkaccessStatus int, sleepFor time.Duration) (*httptest.Server, *Authorizer, authz.Store) {
+	srv, err := serverSetup(loginResp, checkaccessResp, http.StatusOK, checkaccessStatus, sleepFor)
 	if err != nil {
 		t.Fatalf("Error when creating server, reason: %v", err)
 	}
@@ -129,7 +130,7 @@ func TestCheck(t *testing.T) {
 		"actionId":"Microsoft.Kubernetes/connectedClusters/pods/delete",
 		"isDataAction":true,"roleAssignment":null,"denyAssignment":null,"timeToLiveInMs":300000}]`
 
-		srv, client, store := getServerAndClient(t, loginResp, validBody, http.StatusOK)
+		srv, client, store := getServerAndClient(t, loginResp, validBody, http.StatusOK, 1*time.Second)
 		defer srv.Close()
 		defer store.Close()
 
@@ -150,7 +151,7 @@ func TestCheck(t *testing.T) {
 
 	t.Run("unsuccessful request", func(t *testing.T) {
 		validBody := `""`
-		srv, client, store := getServerAndClient(t, loginResp, validBody, http.StatusInternalServerError)
+		srv, client, store := getServerAndClient(t, loginResp, validBody, http.StatusInternalServerError, 1*time.Second)
 		defer srv.Close()
 		defer store.Close()
 
@@ -167,4 +168,24 @@ func TestCheck(t *testing.T) {
 		assert.NotNilf(t, err, "should get error")
 		assert.Contains(t, err.Error(), "Error occured during authorization check")
 	})
+
+	t.Run("context timeout request", func(t *testing.T) {
+		validBody := `""`
+		srv, client, store := getServerAndClient(t, loginResp, validBody, http.StatusInternalServerError, 25*time.Second)
+		defer srv.Close()
+		defer store.Close()
+
+		request := &authzv1.SubjectAccessReviewSpec{
+			User: "beta@bing.com",
+			ResourceAttributes: &authzv1.ResourceAttributes{
+				Namespace: "dev", Group: "", Resource: "pods",
+				Subresource: "status", Version: "v1", Name: "test", Verb: "delete",
+			}, Extra: map[string]authzv1.ExtraValue{"oid": {"00000000-0000-0000-0000-000000000000"}},
+		}
+
+		resp, err := client.Check(request, store)
+		assert.Nilf(t, resp, "response should be nil")
+		assert.NotNilf(t, err, "should get error")
+		assert.Contains(t, err.Error(), "Checkaccess requests have timed out")
+	})
 }
diff --git a/authz/providers/azure/rbac/rbac.go b/authz/providers/azure/rbac/rbac.go
@@ -316,35 +316,29 @@ func (a *AccessInfo) CheckAccess(request *authzv1.SubjectAccessReviewSpec) (*aut
 		body := checkAccessBody
 		eg.Go(func() error {
 			err := a.sendCheckAccessRequest(egCtx, checkAccessURL, body, ch)
-			return err
+			if err != nil {
+				return err
+			}
+			return nil
 		})
 	}
 
-	var waitErr error
-	waitEgCh := make(chan error)
-	go func() {
-		waitErr = eg.Wait()
-		waitEgCh <- waitErr
-	}()
-
-	select {
-	case <-ctx.Done():
-		klog.V(5).Infof("Checkaccess requests have timed out. Error: %v\n", ctx.Err())
-		actionsCount := 0
-		for i := 0; i < len(checkAccessBodies); i += 1 {
-			actionsCount = actionsCount + len(checkAccessBodies[i].Actions)
-		}
-		checkAccessContextTimedOutCount.WithLabelValues(azureutils.ConvertIntToString(len(checkAccessBodies)), azureutils.ConvertIntToString(actionsCount)).Inc()
-		close(ch)
-		return nil, errutils.WithCode(errors.Wrap(ctx.Err(), "Checkaccess requests have timed out."), http.StatusInternalServerError)
-	case err := <-waitEgCh:
-		if err != nil {
+	if err := eg.Wait(); err != nil {
+		if ctx.Err() == context.DeadlineExceeded {
+			klog.V(5).Infof("Checkaccess requests have timed out. Error: %v\n", ctx.Err())
+			actionsCount := 0
+			for i := 0; i < len(checkAccessBodies); i += 1 {
+				actionsCount = actionsCount + len(checkAccessBodies[i].Actions)
+			}
+			checkAccessContextTimedOutCount.WithLabelValues(azureutils.ConvertIntToString(len(checkAccessBodies)), azureutils.ConvertIntToString(actionsCount)).Inc()
+			close(ch)
+			return nil, errutils.WithCode(errors.Wrap(ctx.Err(), "Checkaccess requests have timed out."), http.StatusInternalServerError)
+		} else {
 			close(ch)
 			return nil, err
 		}
-	case <-ch:
-		// do nothing
 	}
+	close(ch)
 
 	var finalStatus *authzv1.SubjectAccessReviewStatus
 	for status := range ch {
@@ -355,7 +349,6 @@ func (a *AccessInfo) CheckAccess(request *authzv1.SubjectAccessReviewSpec) (*aut
 
 		finalStatus = status
 	}
-	close(ch)
 	return finalStatus, nil
 }