chore: add additional trace spans

authz/authz.go

@@ -21,6 +21,8 @@ import (
 	"github.com/gogo/googleapis/google/rpc"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
 	"golang.org/x/net/http2"
 	rpcstatus "google.golang.org/genproto/googleapis/rpc/status"
 
@@ -31,6 +33,10 @@ import (
 
 const ServiceName = "envoy-authz"
 
+var (
+	tracer = otel.Tracer("authz")
+)
+
 type Service struct {
 	authv3connect.UnimplementedAuthorizationHandler
 
@@ -135,6 +141,9 @@ func (s *Service) Check(ctx context.Context, req *connect.Request[auth.CheckRequ
 }
 
 func (s *Service) authProcess(ctx context.Context, req *auth.AttributeContext_HttpRequest, provider *OIDCProvider) (*auth.CheckResponse, error) {
+	ctx, span := tracer.Start(ctx, "authProcess")
+	defer span.End()
+
 	var headers []*core.HeaderValueOption
 	var sessionCookieName = provider.CookieNamePrefix + "-" + ServiceName
 
@@ -147,18 +156,24 @@ func (s *Service) authProcess(ctx context.Context, req *auth.AttributeContext_Ht
 		slog.Debug("session data not found in cookie, creating new")
 		headers, err := s.newSession(ctx, requestedURL, sessionCookieName, provider)
 		if err != nil {
+			span.RecordError(err)
 			return nil, err
 		}
 		// set downstream headers and redirect to Idp
 		return s.authResponse(false, envoy_type.StatusCode_Found, headers, nil, "redirect to Idp"), nil
 	}
 
 	slog.Debug("session data found in cookie", slog.String("session_id", sessionId), slog.String("url", requestedURL))
+	span.SetAttributes(
+		attribute.String("session_id", sessionId),
+		attribute.String("requested_url", requestedURL),
+	)
 
 	// If the request is for the callback URI, then we need to exchange the code for tokens
 	if strings.HasPrefix(requestedURL, provider.CallbackURI+"?") && sessionData.AccessToken == "" {
 		err := s.retriveTokens(ctx, provider, sessionData, requestedURL, sessionCookieName, sessionId)
 		if err != nil {
+			span.RecordError(err)
 			return nil, err
 		}
 		// set downstream headers and redirect client to requested URL from session cookie
@@ -172,6 +187,7 @@ func (s *Service) authProcess(ctx context.Context, req *auth.AttributeContext_Ht
 		slog.Warn("couldn't validating tokens", slog.String("err", err.Error()))
 		headers, err := s.newSession(ctx, requestedURL, sessionCookieName, provider)
 		if err != nil {
+			span.RecordError(err)
 			return nil, err
 		}
 		return s.authResponse(false, envoy_type.StatusCode_Found, headers, nil, "redirect to Idp"), nil

oidc/oidc.go

@@ -13,6 +13,8 @@ import (
 	"github.com/zitadel/oidc/v3/pkg/oidc"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/httptrace/otelhttptrace"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
 )
 
 // Create auth provicer interface
@@ -30,6 +32,10 @@ type OIDCProvider struct {
 	isPKCE   bool
 }
 
+var (
+	tracer = otel.Tracer("oidc")
+)
+
 // NewOIDCProvider creates a new oidc provider
 func NewOIDCProvider(clientID, clientSecret, redirectURI, issuer string, scopes []string) (*OIDCProvider, error) {
 	ctx := context.Background()
@@ -74,20 +80,31 @@ func (o *OIDCProvider) IdpAuthURL(codeChallenge string) string {
 
 func (o *OIDCProvider) VerifyTokens(ctx context.Context, accessToken, idToken string) (bool, error) {
 	var expired bool
-	_, err := rp.VerifyTokens[*oidc.IDTokenClaims](ctx, accessToken, idToken, o.provider.IDTokenVerifier())
+	ctx, span := tracer.Start(ctx, "VerifyTokens")
+	defer span.End()
+
+	t, err := rp.VerifyTokens[*oidc.IDTokenClaims](ctx, accessToken, idToken, o.provider.IDTokenVerifier())
 	if err != nil {
 		if err == oidc.ErrExpired {
 			expired = true
 		} else {
+			span.RecordError(err)
 			return false, err
 		}
 	}
+	span.SetAttributes(
+		attribute.String("issuer", t.GetIssuer()),
+		attribute.String("expire", t.GetExpiration().String()),
+		attribute.Bool("has_expired", expired),
+	)
 	return expired, nil
 }
 
 // RetriveTokens retrieves the tokens from the idp callback redirect and returns them
 // `code` is the `code` query parameter from the idp callback redirect
 func (o *OIDCProvider) RetriveTokens(ctx context.Context, code, codeVerifier string) (*oidc.Tokens[*oidc.IDTokenClaims], error) {
+	ctx, span := tracer.Start(ctx, "RetriveTokens")
+	defer span.End()
 	slog.Debug("retriving tokens", slog.String("authorization_code", code), slog.String("code_verifier", codeVerifier))
 	var opts []rp.CodeExchangeOpt
 

session/encrypt.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"errors"
 
+	"go.opentelemetry.io/otel"
 	"golang.org/x/crypto/nacl/secretbox"
 	"google.golang.org/protobuf/proto"
 
@@ -13,9 +14,13 @@ import (
 
 var (
 	errInvalid = errors.New("invalid encrypted data")
+	tracer     = otel.Tracer("session")
 )
 
 func EncryptSession(ctx context.Context, key [32]byte, sessionData *pb.SessionData) ([]byte, error) {
+	_, span := tracer.Start(ctx, "EncryptSession")
+	defer span.End()
+
 	message, err := proto.Marshal(sessionData)
 	if err != nil {
 		return nil, err
@@ -33,6 +38,8 @@ func EncryptSession(ctx context.Context, key [32]byte, sessionData *pb.SessionDa
 }
 
 func DecryptSession(ctx context.Context, key [32]byte, box []byte) (*pb.SessionData, error) {
+	_, span := tracer.Start(ctx, "DecryptSession")
+	defer span.End()
 	if len(box) < 24 {
 		return nil, errInvalid
 	}