feat: add audit in context in connect interceptor (#1309)

Author: whoAbhishekSah

core/audit/context.go

@@ -2,6 +2,7 @@ package audit
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/raystack/frontier/pkg/server/consts"
 )
@@ -11,6 +12,7 @@ import (
 func GetService(ctx context.Context) *Service {
 	u, ok := ctx.Value(consts.AuditServiceContextKey).(*Service)
 	if !ok {
+		fmt.Println("err: no audit log service found inside context")
 		return NewService("default", NewNoopRepository(), NewNoopWebhookService())
 	}
 	return u

internal/api/v1beta1connect/organization.go

@@ -158,10 +158,13 @@ func (h *ConnectHandler) CreateOrganization(ctx context.Context, request *connec
 		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
 	}
 
-	audit.GetAuditor(ctx, newOrg.ID).LogWithAttrs(audit.OrgCreatedEvent, audit.OrgTarget(newOrg.ID), map[string]string{
+	if err := audit.GetAuditor(ctx, newOrg.ID).LogWithAttrs(audit.OrgCreatedEvent, audit.OrgTarget(newOrg.ID), map[string]string{
 		"title": newOrg.Title,
 		"name":  newOrg.Name,
-	})
+	}); err != nil {
+		errorLogger.LogServiceError(ctx, request, "CreateOrganization.AuditLog", err,
+			zap.String("org_id", newOrg.ID))
+	}
 	return connect.NewResponse(&frontierv1beta1.CreateOrganizationResponse{Organization: orgPB}), nil
 }
 

pkg/server/connect_interceptors/audit.go

@@ -0,0 +1,39 @@
+package connectinterceptors
+
+import (
+	"context"
+
+	"connectrpc.com/connect"
+	"github.com/raystack/frontier/core/audit"
+)
+
+type AuditInterceptor struct {
+	service *audit.Service
+}
+
+func NewAuditInterceptor(service *audit.Service) *AuditInterceptor {
+	return &AuditInterceptor{
+		service: service,
+	}
+}
+
+func (a *AuditInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc {
+	return connect.UnaryFunc(func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
+		ctx = audit.SetContextWithService(ctx, a.service)
+		return next(ctx, req)
+	})
+}
+
+func (a *AuditInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc {
+	return connect.StreamingClientFunc(func(ctx context.Context, spec connect.Spec) connect.StreamingClientConn {
+		conn := next(ctx, spec)
+		return conn
+	})
+}
+
+func (a *AuditInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
+	return connect.StreamingHandlerFunc(func(ctx context.Context, conn connect.StreamingHandlerConn) error {
+		ctx = audit.SetContextWithService(ctx, a.service)
+		return next(ctx, conn)
+	})
+}

pkg/server/server.go

@@ -197,6 +197,7 @@ func ServeConnect(ctx context.Context, logger log.Logger, cfg Config, deps api.D
 	authNInterceptor := connectinterceptors.NewAuthenticationInterceptor(frontierService, cfg.Authentication.Session.Headers)
 	authZInterceptor := connectinterceptors.NewAuthorizationInterceptor(frontierService)
 	sessionInterceptor := connectinterceptors.NewSessionInterceptor(sessionCookieCutter, cfg.Authentication.Session, frontierService)
+	auditInterceptor := connectinterceptors.NewAuditInterceptor(deps.AuditService)
 
 	interceptors := connect.WithInterceptors(
 		otelInterceptor,
@@ -205,6 +206,7 @@ func ServeConnect(ctx context.Context, logger log.Logger, cfg Config, deps api.D
 		sessionInterceptor,
 		authNInterceptor,
 		authZInterceptor,
+		auditInterceptor,
 		sessionInterceptor.UnaryConnectResponseInterceptor())
 
 	frontierPath, frontierHandler := frontierv1beta1connect.NewFrontierServiceHandler(frontierService, interceptors, connect.WithCodec(connectCodec{}))