[v15] Add a fallback for EmitAuditEvents failure due to event conflicts (DynamoDB backend)

lib/events/dynamoevents/dynamoevents.go

@@ -112,6 +112,9 @@ type Config struct {
 	UIDGenerator utils.UID
 	// Endpoint is an optional non-AWS endpoint
 	Endpoint string
+	// DisableConflictCheck disables conflict checks when emitting an event.
+	// Disabling it can cause events to be lost due to them being overwritten.
+	DisableConflictCheck bool
 
 	// ReadMaxCapacity is the maximum provisioned read capacity.
 	ReadMaxCapacity int64
@@ -146,6 +149,10 @@ func (cfg *Config) SetFromURL(in *url.URL) error {
 		cfg.Endpoint = endpoint
 	}
 
+	if disableConflictCheck := in.Query().Get("disable_conflict_check"); disableConflictCheck != "" {
+		cfg.DisableConflictCheck = true
+	}
+
 	const boolErrorTemplate = "failed to parse URI %q flag %q - %q, supported values are 'true', 'false', or any other" +
 		"supported boolean in https://pkg.go.dev/strconv#ParseBool"
 	if val := in.Query().Get(events.UseFIPSQueryParam); val != "" {
@@ -366,24 +373,19 @@ const (
 func (l *Log) EmitAuditEvent(ctx context.Context, in apievents.AuditEvent) error {
 	ctx = context.WithoutCancel(ctx)
 	sessionID := getSessionID(in)
-	if err := l.putAuditEvent(ctx, sessionID, in); err != nil {
-		switch {
-		case isAWSValidationError(err):
-			// In case of ValidationException: Item size has exceeded the maximum allowed size
-			// sanitize event length and retry upload operation.
-			return trace.Wrap(l.handleAWSValidationError(ctx, err, sessionID, in))
-		}
-		return trace.Wrap(err)
-	}
-	return nil
+	return trace.Wrap(l.putAuditEvent(ctx, sessionID, in))
 }
 
 func (l *Log) handleAWSValidationError(ctx context.Context, err error, sessionID string, in apievents.AuditEvent) error {
+	if alreadyTrimmed := ctx.Value(largeEventHandledContextKey); alreadyTrimmed != nil {
+		return err
+	}
+
 	se, ok := trimEventSize(in)
 	if !ok {
 		return trace.BadParameter(err.Error())
 	}
-	if err := l.putAuditEvent(ctx, sessionID, se); err != nil {
+	if err := l.putAuditEvent(context.WithValue(ctx, largeEventHandledContextKey, true), sessionID, se); err != nil {
 		return trace.BadParameter(err.Error())
 	}
 	fields := log.Fields{"event_id": in.GetID(), "event_type": in.GetType()}
@@ -392,6 +394,22 @@ func (l *Log) handleAWSValidationError(ctx context.Context, err error, sessionID
 	return nil
 }
 
+func (l *Log) handleConditionError(ctx context.Context, err error, sessionID string, in apievents.AuditEvent) error {
+	if alreadyUpdated := ctx.Value(conflictHandledContextKey); alreadyUpdated != nil {
+		return err
+	}
+
+	// Update index using the current system time instead of event time to
+	// ensure the value is always set.
+	in.SetIndex(l.Clock.Now().UnixNano())
+
+	if err := l.putAuditEvent(context.WithValue(ctx, conflictHandledContextKey, true), sessionID, in); err != nil {
+		return trace.Wrap(err)
+	}
+	l.WithFields(log.Fields{"event_id": in.GetID(), "event_type": in.GetType()}).Debug("Event index overwritten")
+	return nil
+}
+
 // getSessionID if set returns event ID obtained from metadata or generates a new one.
 func getSessionID(in apievents.AuditEvent) string {
 	s, ok := in.(events.SessionMetadataGetter)
@@ -415,13 +433,48 @@ func trimEventSize(event apievents.AuditEvent) (apievents.AuditEvent, bool) {
 	return m.TrimToMaxSize(maxItemSize), true
 }
 
+// putAuditEventContextKey represents context keys of putAuditEvent.
+type putAuditEventContextKey int
+
+const (
+	// conflictHandledContextKey if present on the context, the conflict error
+	// was already handled.
+	conflictHandledContextKey putAuditEventContextKey = iota
+	// largeEventHandledContextKey if present on the context, the large event
+	// error was already handled.
+	largeEventHandledContextKey
+)
+
 func (l *Log) putAuditEvent(ctx context.Context, sessionID string, in apievents.AuditEvent) error {
 	input, err := l.createPutItem(sessionID, in)
 	if err != nil {
 		return trace.Wrap(err)
 	}
-	_, err = l.svc.PutItemWithContext(ctx, input)
-	return convertError(err)
+
+	if _, err = l.svc.PutItemWithContext(ctx, input); err != nil {
+		err = convertError(err)
+
+		switch {
+		case isAWSValidationError(err):
+			// In case of ValidationException: Item size has exceeded the maximum allowed size
+			// sanitize event length and retry upload operation.
+			return trace.Wrap(l.handleAWSValidationError(ctx, err, sessionID, in))
+		case trace.IsAlreadyExists(err):
+			// Condition errors are directly related to the uniqueness of the
+			// item event index/session id. Since we can't change the session
+			// id, update the event index with a new value and retry the put
+			// item.
+			l.
+				WithError(err).
+				WithFields(log.Fields{"event_type": in.GetType(), "session_id": sessionID, "event_index": in.GetIndex()}).
+				Error("Conflict on event session_id and event_index")
+			return trace.Wrap(l.handleConditionError(ctx, err, sessionID, in))
+		}
+
+		return err
+	}
+
+	return nil
 }
 
 func (l *Log) createPutItem(sessionID string, in apievents.AuditEvent) (*dynamodb.PutItemInput, error) {
@@ -443,11 +496,17 @@ func (l *Log) createPutItem(sessionID string, in apievents.AuditEvent) (*dynamod
 	if err != nil {
 		return nil, trace.Wrap(err)
 	}
-	return &dynamodb.PutItemInput{
-		Item:                av,
-		TableName:           aws.String(l.Tablename),
-		ConditionExpression: aws.String("attribute_not_exists(SessionID) AND attribute_not_exists(EventIndex)"),
-	}, nil
+
+	input := &dynamodb.PutItemInput{
+		Item:      av,
+		TableName: aws.String(l.Tablename),
+	}
+
+	if !l.Config.DisableConflictCheck {
+		input.ConditionExpression = aws.String("attribute_not_exists(SessionID) AND attribute_not_exists(EventIndex)")
+	}
+
+	return input, nil
 }
 
 type messageSizeTrimmer interface {

lib/events/dynamoevents/dynamoevents_test.go

@@ -35,7 +35,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport"
-	apidefaults "github.com/gravitational/teleport/api/defaults"
 	"github.com/gravitational/teleport/api/types"
 	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/lib/events"
@@ -68,6 +67,7 @@ func setupDynamoContext(t *testing.T) *dynamoContext {
 		Tablename:    fmt.Sprintf("teleport-test-%v", uuid.New().String()),
 		Clock:        fakeClock,
 		UIDGenerator: utils.NewFakeUID(),
+		Endpoint:     "http://localhost:8000",
 	})
 	require.NoError(t, err)
 
@@ -437,38 +437,43 @@ func TestConfig_CheckAndSetDefaults(t *testing.T) {
 }
 
 // TestEmitSessionEventsSameIndex given events that share the same session ID
-// and index, the emit should fail, avoiding any event to get overwritten.
+// and index, the emit should succeed.
 func TestEmitSessionEventsSameIndex(t *testing.T) {
 	ctx := context.Background()
 	tt := setupDynamoContext(t)
 	sessionID := session.NewID()
 
-	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 0)))
-	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 1)))
-	require.Error(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 1)))
+	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 0, "")))
+	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 1, "")))
+	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 1, "")))
 }
 
-func generateEvent(sessionID session.ID, index int64) apievents.AuditEvent {
-	return &apievents.AppSessionChunk{
+// TestValidationErrorsHandling given events that return validation
+// errors (large event size and already exists), the emit should handle them
+// and succeed on emitting the event when it does support trimming.
+func TestValidationErrorsHandling(t *testing.T) {
+	ctx := context.Background()
+	tt := setupDynamoContext(t)
+	sessionID := session.NewID()
+	largeQuery := strings.Repeat("A", maxItemSize)
+
+	// First write should only trigger the large event size
+	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 0, largeQuery)))
+	// Second should trigger both errors.
+	require.NoError(t, tt.log.EmitAuditEvent(ctx, generateEvent(sessionID, 0, largeQuery)))
+}
+
+func generateEvent(sessionID session.ID, index int64, query string) apievents.AuditEvent {
+	return &apievents.DatabaseSessionQuery{
 		Metadata: apievents.Metadata{
-			Type:        events.AppSessionChunkEvent,
-			Code:        events.AppSessionChunkCode,
+			Type:        events.DatabaseSessionQueryEvent,
 			ClusterName: "root",
 			Index:       index,
 		},
-		ServerMetadata: apievents.ServerMetadata{
-			ServerID:        uuid.New().String(),
-			ServerNamespace: apidefaults.Namespace,
-		},
 		SessionMetadata: apievents.SessionMetadata{
 			SessionID: sessionID.String(),
 		},
-		AppMetadata: apievents.AppMetadata{
-			AppURI:        "nginx",
-			AppPublicAddr: "https://nginx",
-			AppName:       "nginx",
-		},
-		SessionChunkID: uuid.New().String(),
+		DatabaseQuery: query,
 	}
 }