From ff905615b8c25f35312fdca9eadf507e11bf1715 Mon Sep 17 00:00:00 2001
From: Brian Kassouf <briankassouf@users.noreply.github.com>
Date: Tue, 16 Feb 2021 16:39:53 -0800
Subject: [PATCH] Use the request context for Kubernetes API call (#104) (#105)

* Use the request context for kubernetes API call

* Upgrade go version
---
 .circleci/config.yml |  2 +-
 path_login.go        |  6 +++---
 path_login_test.go   | 27 +++++++++++++++++++++++++++
 token_review.go      | 13 +++++++++----
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 86ff7666..5409ec83 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -2,7 +2,7 @@ version: 2
 jobs:
   build:
     docker:
-      - image: circleci/golang:1.12
+      - image: docker.mirror.hashicorp.services/circleci/golang:1.15
     working_directory: /go/src/github.com/hashicorp/vault-plugin-auth-kubernetes
     steps:
       - checkout
diff --git a/path_login.go b/path_login.go
index 504700d8..9f3fc7ce 100644
--- a/path_login.go
+++ b/path_login.go
@@ -103,7 +103,7 @@ func (b *kubeAuthBackend) pathLogin(ctx context.Context, req *logical.Request, d
 	}
 
 	// look up the JWT token in the kubernetes API
-	err = serviceAccount.lookup(jwtStr, b.reviewFactory(config))
+	err = serviceAccount.lookup(ctx, jwtStr, b.reviewFactory(config))
 	if err != nil {
 		b.Logger().Error(`login unauthorized due to: ` + err.Error())
 		return nil, logical.ErrPermissionDenied
@@ -350,8 +350,8 @@ type projectedServiceAccountPod struct {
 
 // lookup calls the TokenReview API in kubernetes to verify the token and secret
 // still exist.
-func (s *serviceAccount) lookup(jwtStr string, tr tokenReviewer) error {
-	r, err := tr.Review(jwtStr, s.Audience)
+func (s *serviceAccount) lookup(ctx context.Context, jwtStr string, tr tokenReviewer) error {
+	r, err := tr.Review(ctx, jwtStr, s.Audience)
 	if err != nil {
 		return err
 	}
diff --git a/path_login_test.go b/path_login_test.go
index 661b53a0..f560151c 100644
--- a/path_login_test.go
+++ b/path_login_test.go
@@ -263,6 +263,33 @@ func TestLogin(t *testing.T) {
 	}
 }
 
+func TestLogin_ContextError(t *testing.T) {
+	b, storage := setupBackend(t, testDefaultPEMs, testName, testNamespace)
+
+	data := map[string]interface{}{
+		"role": "plugin-test",
+		"jwt":  jwtData,
+	}
+
+	req := &logical.Request{
+		Operation: logical.UpdateOperation,
+		Path:      "login",
+		Storage:   storage,
+		Data:      data,
+		Connection: &logical.Connection{
+			RemoteAddr: "127.0.0.1",
+		},
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := b.HandleRequest(ctx, req)
+	if err != context.Canceled {
+		t.Fatalf("expected context canceled error, got: %v", err)
+	}
+}
+
 func TestLogin_ECDSA_PEM(t *testing.T) {
 	b, storage := setupBackend(t, testNoPEMs, testName, testNamespace)
 
diff --git a/token_review.go b/token_review.go
index 75e86e09..6b36a95a 100644
--- a/token_review.go
+++ b/token_review.go
@@ -2,6 +2,7 @@ package kubeauth
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"encoding/json"
@@ -28,7 +29,7 @@ type tokenReviewResult struct {
 
 // This exists so we can use a mock TokenReview when running tests
 type tokenReviewer interface {
-	Review(string, []string) (*tokenReviewResult, error)
+	Review(context.Context, string, []string) (*tokenReviewResult, error)
 }
 
 type tokenReviewFactory func(*kubeConfig) tokenReviewer
@@ -44,7 +45,7 @@ func tokenReviewAPIFactory(config *kubeConfig) tokenReviewer {
 	}
 }
 
-func (t *tokenReviewAPI) Review(jwt string, aud []string) (*tokenReviewResult, error) {
+func (t *tokenReviewAPI) Review(ctx context.Context, jwt string, aud []string) (*tokenReviewResult, error) {
 
 	client := cleanhttp.DefaultClient()
 
@@ -75,7 +76,7 @@ func (t *tokenReviewAPI) Review(jwt string, aud []string) (*tokenReviewResult, e
 
 	// Build the request to the token review API
 	url := fmt.Sprintf("%s/apis/authentication.k8s.io/v1/tokenreviews", strings.TrimSuffix(t.config.Host, "/"))
-	req, err := http.NewRequest("POST", url, bytes.NewBuffer(trJSON))
+	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(trJSON))
 	if err != nil {
 		return nil, err
 	}
@@ -188,7 +189,11 @@ func mockTokenReviewFactory(name, namespace, UID string) tokenReviewFactory {
 	}
 }
 
-func (t *mockTokenReview) Review(jwt string, aud []string) (*tokenReviewResult, error) {
+func (t *mockTokenReview) Review(ctx context.Context, cjwt string, aud []string) (*tokenReviewResult, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
 	return &tokenReviewResult{
 		Name:      t.saName,
 		Namespace: t.saNamespace,