extproc: eliminates context.Background (envoyproxy#296)

**Commit Message**

This refactors the extproc package and
eliminates all context.Background usages.

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>
Signed-off-by: Eric Mariasis <ericmariasis829@gmail.com>

internal/extproc/backendauth/api_key.go

@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -27,7 +28,7 @@ func newAPIKeyHandler(auth *filterapi.APIKeyAuth) (Handler, error) {
 // Do implements [Handler.Do].
 //
 // Extracts the api key from the local file and set it as an authorization header.
-func (a *apiKeyHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
+func (a *apiKeyHandler) Do(_ context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
 	requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", a.apiKey)
 	headerMut.SetHeaders = append(headerMut.SetHeaders, &corev3.HeaderValueOption{
 		Header: &corev3.HeaderValue{Key: "Authorization", RawValue: []byte(requestHeaders["Authorization"])},

internal/extproc/backendauth/api_key_test.go

@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -62,7 +63,7 @@ func TestApiKeyHandler_Do(t *testing.T) {
 			Body: []byte(`{"messages": [{"role": "user", "content": [{"text": "Say this is a test!"}]}]}`),
 		},
 	}
-	err = handler.Do(requestHeaders, headerMut, bodyMut)
+	err = handler.Do(context.Background(), requestHeaders, headerMut, bodyMut)
 	require.NoError(t, err)
 
 	bearerToken, ok := requestHeaders["Authorization"]

internal/extproc/backendauth/auth.go

@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"errors"
 
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -13,7 +14,7 @@ import (
 // TODO: maybe this can be just "post-transformation" handler, as it is not really only about auth.
 type Handler interface {
 	// Do performs the backend auth, and make changes to the request headers and body mutations.
-	Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error
+	Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error
 }
 
 // NewHandler returns a new implementation of [Handler] based on the configuration.

internal/extproc/backendauth/aws.go

@@ -27,7 +27,7 @@ type awsHandler struct {
 	region      string
 }
 
-func newAWSHandler(awsAuth *filterapi.AWSAuth) (*awsHandler, error) {
+func newAWSHandler(awsAuth *filterapi.AWSAuth) (Handler, error) {
 	var credentials aws.Credentials
 	var region string
 
@@ -60,7 +60,7 @@ func newAWSHandler(awsAuth *filterapi.AWSAuth) (*awsHandler, error) {
 //
 // This assumes that during the transformation, the path is set in the header mutation as well as
 // the body in the body mutation.
-func (a *awsHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error {
+func (a *awsHandler) Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error {
 	method := requestHeaders[":method"]
 	path := ""
 	if headerMut.SetHeaders != nil {
@@ -90,7 +90,7 @@ func (a *awsHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.H
 		return fmt.Errorf("cannot create request: %w", err)
 	}
 
-	err = a.signer.SignHTTP(context.Background(), a.credentials, req,
+	err = a.signer.SignHTTP(ctx, a.credentials, req,
 		hex.EncodeToString(payloadHash[:]), "bedrock", a.region, time.Now())
 	if err != nil {
 		return fmt.Errorf("cannot sign request: %w", err)

internal/extproc/backendauth/aws_test.go

@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -42,7 +43,7 @@ func TestAWSHandler_Do(t *testing.T) {
 
 	for _, tc := range []struct {
 		name    string
-		handler *awsHandler
+		handler Handler
 	}{
 		{
 			name:    "Using AWS Credential File",
@@ -64,7 +65,7 @@ func TestAWSHandler_Do(t *testing.T) {
 					Body: []byte(`{"messages": [{"role": "user", "content": [{"text": "Say this is a test!"}]}]}`),
 				},
 			}
-			err = tc.handler.Do(requestHeaders, headerMut, bodyMut)
+			err = tc.handler.Do(context.Background(), requestHeaders, headerMut, bodyMut)
 			require.NoError(t, err)
 		})
 	}

internal/extproc/processor.go

@@ -80,7 +80,7 @@ func (p *Processor) ProcessRequestHeaders(_ context.Context, headers *corev3.Hea
 }
 
 // ProcessRequestBody implements [Processor.ProcessRequestBody].
-func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
+func (p *Processor) ProcessRequestBody(ctx context.Context, rawBody *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
 	path := p.requestHeaders[":path"]
 	model, body, err := p.config.bodyParser(path, rawBody)
 	if err != nil {
@@ -122,7 +122,7 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt
 	})
 
 	if authHandler, ok := p.config.backendAuthHandlers[b.Name]; ok {
-		if err := authHandler.Do(p.requestHeaders, headerMutation, bodyMutation); err != nil {
+		if err := authHandler.Do(ctx, p.requestHeaders, headerMutation, bodyMutation); err != nil {
 			return nil, fmt.Errorf("failed to do auth request: %w", err)
 		}
 	}