authz: add more tracing (#5435)

authorize/databroker.go

@@ -9,6 +9,7 @@ import (
 	"github.com/pomerium/pomerium/pkg/grpc/user"
 	"github.com/pomerium/pomerium/pkg/grpcutil"
 	"github.com/pomerium/pomerium/pkg/storage"
+	octrace "go.opencensus.io/trace"
 )
 
 type sessionOrServiceAccount interface {
@@ -22,6 +23,14 @@ func getDataBrokerRecord(
 	recordID string,
 	lowestRecordVersion uint64,
 ) (*databroker.Record, error) {
+	ctx, span := trace.StartSpan(ctx, "authorize.getDataBrokerRecord")
+	span.AddAttributes(
+		octrace.StringAttribute("record_type", recordType),
+		octrace.StringAttribute("record_id", recordID),
+		octrace.Int64Attribute("lowest_record_version", int64(lowestRecordVersion)),
+	)
+	defer span.End()
+
 	q := storage.GetQuerier(ctx)
 
 	req := &databroker.QueryRequest{
@@ -32,14 +41,17 @@ func getDataBrokerRecord(
 
 	res, err := q.Query(ctx, req)
 	if err != nil {
+		span.SetStatus(octrace.Status{Code: octrace.StatusCodeInternal, Message: err.Error()})
 		return nil, err
 	}
 	if len(res.GetRecords()) == 0 {
+		span.SetStatus(octrace.Status{Code: octrace.StatusCodeNotFound})
 		return nil, storage.ErrNotFound
 	}
 
 	// if the current record version is less than the lowest we'll accept, invalidate the cache
-	if res.GetRecords()[0].GetVersion() < lowestRecordVersion {
+	if v := res.GetRecords()[0].GetVersion(); v < lowestRecordVersion {
+		span.AddAttributes(octrace.Int64Attribute("got_record_version", int64(v)))
 		q.InvalidateCache(ctx, req)
 	} else {
 		return res.GetRecords()[0], nil
@@ -48,9 +60,11 @@ func getDataBrokerRecord(
 	// retry with an up to date cache
 	res, err = q.Query(ctx, req)
 	if err != nil {
+		span.SetStatus(octrace.Status{Code: octrace.StatusCodeInternal, Message: err.Error()})
 		return nil, err
 	}
 	if len(res.GetRecords()) == 0 {
+		span.SetStatus(octrace.Status{Code: octrace.StatusCodeNotFound})
 		return nil, storage.ErrNotFound
 	}
 
@@ -96,10 +110,12 @@ func (a *Authorize) getDataBrokerUser(
 	userID string,
 ) (*user.User, error) {
 	ctx, span := trace.StartSpan(ctx, "authorize.getDataBrokerUser")
+	span.AddAttributes(octrace.StringAttribute("user_id", userID))
 	defer span.End()
 
 	record, err := getDataBrokerRecord(ctx, grpcutil.GetTypeURL(new(user.User)), userID, 0)
 	if err != nil {
+		span.SetStatus(octrace.Status{Code: octrace.StatusCodeInternal, Message: err.Error()})
 		return nil, err
 	}
 

authorize/grpc.go

@@ -3,12 +3,14 @@ package authorize
 import (
 	"context"
 	"encoding/pem"
+	"fmt"
 	"io"
 	"net/http"
 	"net/url"
 	"strings"
 
 	envoy_service_auth_v3 "github.com/envoyproxy/go-control-plane/envoy/service/auth/v3"
+	octrace "go.opencensus.io/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/structpb"
@@ -31,7 +33,6 @@ import (
 func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRequest) (*envoy_service_auth_v3.CheckResponse, error) {
 	ctx, span := trace.StartSpan(ctx, "authorize.grpc.Check")
 	defer span.End()
-
 	querier := storage.NewTracingQuerier(
 		storage.NewCachingQuerier(
 			storage.NewCachingQuerier(
@@ -48,6 +49,7 @@ func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRe
 	// convert the incoming envoy-style http request into a go-style http request
 	hreq := getHTTPRequestFromCheckRequest(in)
 	requestID := requestid.FromHTTPHeader(hreq.Header)
+	span.AddAttributes(octrace.StringAttribute("request_id", requestID))
 	ctx = requestid.WithValue(ctx, requestID)
 
 	sessionState, _ := state.sessionStore.LoadSessionStateAndCheckIDP(hreq)
@@ -59,7 +61,7 @@ func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRe
 		s, err = a.getDataBrokerSessionOrServiceAccount(ctx, sessionState.ID, sessionState.DatabrokerRecordVersion)
 		if status.Code(err) == codes.Unavailable {
 			log.Ctx(ctx).Debug().Str("request-id", requestID).Err(err).Msg("temporary error checking authorization: data broker unavailable")
-			return nil, err
+			return nil, fmt.Errorf("databroker unavailable")
 		} else if err != nil {
 			log.Ctx(ctx).Info().Err(err).Str("request-id", requestID).Msg("clearing session due to missing or invalid session or service account")
 			sessionState = nil
@@ -72,7 +74,7 @@ func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRe
 	req, err := a.getEvaluatorRequestFromCheckRequest(ctx, in, sessionState)
 	if err != nil {
 		log.Ctx(ctx).Error().Err(err).Str("request-id", requestID).Msg("error building evaluator request")
-		return nil, err
+		return nil, fmt.Errorf("build evaluator request: %w", err)
 	}
 
 	// take the state lock here so we don't update while evaluating
@@ -81,7 +83,7 @@ func (a *Authorize) Check(ctx context.Context, in *envoy_service_auth_v3.CheckRe
 	a.stateLock.RUnlock()
 	if err != nil {
 		log.Ctx(ctx).Error().Err(err).Str("request-id", requestID).Msg("error during OPA evaluation")
-		return nil, err
+		return nil, fmt.Errorf("evaluate: %w", err)
 	}
 
 	// if show error details is enabled, attach the policy evaluation traces

pkg/storage/querier.go

@@ -15,6 +15,7 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/pomerium/pomerium/internal/telemetry/trace"
 	"github.com/pomerium/pomerium/pkg/cryptutil"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 	"github.com/pomerium/pomerium/pkg/protoutil"
@@ -229,6 +230,9 @@ func NewCachingQuerier(q Querier, cache Cache) Querier {
 }
 
 func (q *cachingQuerier) InvalidateCache(ctx context.Context, in *databroker.QueryRequest) {
+	ctx, span := trace.StartSpan(ctx, "storage.cachingQuerier.InvalidateCache")
+	defer span.End()
+
 	key, err := (&proto.MarshalOptions{
 		Deterministic: true,
 	}).Marshal(in)
@@ -240,6 +244,9 @@ func (q *cachingQuerier) InvalidateCache(ctx context.Context, in *databroker.Que
 }
 
 func (q *cachingQuerier) Query(ctx context.Context, in *databroker.QueryRequest, opts ...grpc.CallOption) (*databroker.QueryResponse, error) {
+	ctx, span := trace.StartSpan(ctx, "storage.cachingQuerier.Query")
+	defer span.End()
+
 	key, err := (&proto.MarshalOptions{
 		Deterministic: true,
 	}).Marshal(in)
@@ -248,6 +255,9 @@ func (q *cachingQuerier) Query(ctx context.Context, in *databroker.QueryRequest,
 	}
 
 	rawResult, err := q.cache.GetOrUpdate(ctx, key, func(ctx context.Context) ([]byte, error) {
+		ctx, span := trace.StartSpan(ctx, "storage.cachingQuerier.Query.Update")
+		defer span.End()
+
 		res, err := q.q.Query(ctx, in, opts...)
 		if err != nil {
 			return nil, err