From e44314b03e3cae40728843e5d01d696acbf20615 Mon Sep 17 00:00:00 2001 From: Takeshi Yoneda Date: Wed, 5 Feb 2025 18:05:49 -0800 Subject: [PATCH] extproc: eliminates context.Background (#296) **Commit Message** This refactors the extproc package and eliminates all context.Background usages. Signed-off-by: Takeshi Yoneda Signed-off-by: Eric Mariasis --- internal/extproc/backendauth/api_key.go | 3 ++- internal/extproc/backendauth/api_key_test.go | 3 ++- internal/extproc/backendauth/auth.go | 3 ++- internal/extproc/backendauth/aws.go | 6 +++--- internal/extproc/backendauth/aws_test.go | 5 +++-- internal/extproc/processor.go | 4 ++-- 6 files changed, 14 insertions(+), 10 deletions(-) diff --git a/internal/extproc/backendauth/api_key.go b/internal/extproc/backendauth/api_key.go index 3a94fb0e5..e7c36869d 100644 --- a/internal/extproc/backendauth/api_key.go +++ b/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"])}, diff --git a/internal/extproc/backendauth/api_key_test.go b/internal/extproc/backendauth/api_key_test.go index cfb3ce135..6f08f3882 100644 --- a/internal/extproc/backendauth/api_key_test.go +++ b/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"] diff --git a/internal/extproc/backendauth/auth.go b/internal/extproc/backendauth/auth.go index e8375843f..8dd4fbc8b 100644 --- a/internal/extproc/backendauth/auth.go +++ b/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. diff --git a/internal/extproc/backendauth/aws.go b/internal/extproc/backendauth/aws.go index 65fbf098c..e01ecc661 100644 --- a/internal/extproc/backendauth/aws.go +++ b/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) diff --git a/internal/extproc/backendauth/aws_test.go b/internal/extproc/backendauth/aws_test.go index 3a3da56cb..e0d922b6f 100644 --- a/internal/extproc/backendauth/aws_test.go +++ b/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) }) } diff --git a/internal/extproc/processor.go b/internal/extproc/processor.go index 7af784240..bf1610e07 100644 --- a/internal/extproc/processor.go +++ b/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) } }