refactor: extract function to inject project ids headergkg-inject-ids

2 files changed, 57 insertions, 44 deletions

diff --git a/pkg/authz/check_service.go b/pkg/authz/check_service.go
index 38e8b410..92f6da40 100644
--- a/pkg/authz/check_service.go
+++ b/pkg/authz/check_service.go
@@ -2,10 +2,8 @@ package authz
 
 import (
 	"context"
-	"io"
 	"net/http"
 	"path/filepath"
-	"strings"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 	authzed "github.com/authzed/authzed-go/v1"
@@ -35,7 +33,7 @@ func NewCheckService(client *authzed.Client) auth.AuthorizationServer {
 
 func (svc *CheckService) Check(ctx context.Context, request *auth.CheckRequest) (*auth.CheckResponse, error) {
 	if svc.isAuthorized(ctx, request) {
-		return svc.OK(ctx, svc.injectHeaders(ctx, request)), nil
+		return svc.OK(ctx, WithProjectIDs(ctx, svc.client, request)), nil
 	}
 	return svc.Denied(ctx), nil
 }
@@ -82,13 +80,13 @@ func (svc *CheckService) validRequest(ctx context.Context, r *auth.CheckRequest)
 		x.IsPresent(r.Attributes.Request.Http)
 }
 
-func (svc *CheckService) OK(ctx context.Context, f x.Option[*auth.CheckResponse_OkResponse]) *auth.CheckResponse {
+func (svc *CheckService) OK(ctx context.Context, option x.Option[*auth.CheckResponse_OkResponse]) *auth.CheckResponse {
 	log.WithFields(ctx, log.Fields{"authorized": true})
 	return &auth.CheckResponse{
 		Status: &status.Status{
 			Code: int32(codes.OK),
 		},
-		HttpResponse: f(&auth.CheckResponse_OkResponse{
+		HttpResponse: option(&auth.CheckResponse_OkResponse{
 			OkResponse: &auth.OkHttpResponse{
 				Headers:              []*core.HeaderValueOption{},
 				HeadersToRemove:      []string{},
@@ -114,42 +112,3 @@ func (svc *CheckService) Denied(ctx context.Context) *auth.CheckResponse {
 		},
 	}
 }
-
-func (svc *CheckService) injectHeaders(ctx context.Context, request *auth.CheckRequest) x.Option[*auth.CheckResponse_OkResponse] {
-	return x.With[*auth.CheckResponse_OkResponse](func(response *auth.CheckResponse_OkResponse) {
-		if x.IsZero(svc.client) {
-			return
-		}
-
-		stream, err := svc.client.LookupResources(ctx, &v1.LookupResourcesRequest{
-			ResourceObjectType: "project",
-			Permission:         "read_project",
-			Subject:            mapper.MapFrom[*auth.CheckRequest, *v1.SubjectReference](request),
-		})
-		if err != nil {
-			pls.LogError(ctx, err)
-			return
-		}
-
-		var projectIDs []string
-		for {
-			result, err := stream.Recv()
-			if err == io.EOF {
-				break
-			}
-			if err != nil {
-				pls.LogError(ctx, err)
-				break
-			}
-			projectIDs = append(projectIDs, result.ResourceObjectId)
-		}
-
-		response.OkResponse.Headers = append(response.OkResponse.Headers, &core.HeaderValueOption{
-			Header: &core.HeaderValue{
-				Key:   "x-project-ids",
-				Value: strings.Join(projectIDs, ","),
-			},
-			AppendAction: core.HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD,
-		})
-	})
-}
diff --git a/pkg/authz/option.go b/pkg/authz/option.go
new file mode 100644
index 00000000..585deedf
--- /dev/null
+++ b/pkg/authz/option.go
@@ -0,0 +1,54 @@
+package authz
+
+import (
+	"context"
+	"io"
+	"strings"
+
+	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+	authzed "github.com/authzed/authzed-go/v1"
+	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
+	auth "github.com/envoyproxy/go-control-plane/envoy/service/auth/v3"
+	"github.com/xlgmokha/x/pkg/mapper"
+	"github.com/xlgmokha/x/pkg/x"
+	"gitlab.com/gitlab-org/software-supply-chain-security/authorization/authzd.git/pkg/pls"
+)
+
+func WithProjectIDs(ctx context.Context, client *authzed.Client, request *auth.CheckRequest) x.Option[*auth.CheckResponse_OkResponse] {
+	return x.With[*auth.CheckResponse_OkResponse](func(response *auth.CheckResponse_OkResponse) {
+		if x.IsZero(client) {
+			return
+		}
+
+		stream, err := client.LookupResources(ctx, &v1.LookupResourcesRequest{
+			ResourceObjectType: "project",
+			Permission:         "read_project",
+			Subject:            mapper.MapFrom[*auth.CheckRequest, *v1.SubjectReference](request),
+		})
+		if err != nil {
+			pls.LogError(ctx, err)
+			return
+		}
+
+		var projectIDs []string
+		for {
+			result, err := stream.Recv()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				pls.LogError(ctx, err)
+				break
+			}
+			projectIDs = append(projectIDs, result.ResourceObjectId)
+		}
+
+		response.OkResponse.Headers = append(response.OkResponse.Headers, &core.HeaderValueOption{
+			Header: &core.HeaderValue{
+				Key:   "x-project-ids",
+				Value: strings.Join(projectIDs, ","),
+			},
+			AppendAction: core.HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD,
+		})
+	})
+}