Refactor shared DIFC decisions and logger level wrappers

Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/97151d73-796a-42a0-9fb4-c0de9f160c32

Co-authored-by: lpcox <[email protected]>

Author: Copilot

internal/difc/pipeline_decisions.go

@@ -0,0 +1,26 @@
+package difc
+
+// ShouldBypassCoarseDeny returns true when a coarse-grained deny should still
+// proceed to backend execution so Phase 5 can enforce per-item policy.
+func ShouldBypassCoarseDeny(operation OperationType) bool {
+	return operation == OperationRead
+}
+
+// ShouldCallLabelResponse returns true when guards should label response data
+// for possible fine-grained filtering.
+func ShouldCallLabelResponse(operation OperationType, enforcementMode EnforcementMode) bool {
+	isPureWrite := operation == OperationWrite
+	return !isPureWrite && (operation != OperationReadWrite || enforcementMode != EnforcementStrict)
+}
+
+// ShouldBlockFilteredResponse returns true when filtered items should block the
+// whole response instead of returning a partially filtered result.
+func ShouldBlockFilteredResponse(enforcementMode EnforcementMode, filteredCount int) bool {
+	return enforcementMode == EnforcementStrict && filteredCount > 0
+}
+
+// ShouldAccumulateReadLabels returns true when read labels should be
+// accumulated back into the agent label set.
+func ShouldAccumulateReadLabels(operation OperationType, enforcementMode EnforcementMode) bool {
+	return operation != OperationWrite && enforcementMode == EnforcementPropagate
+}

internal/difc/pipeline_decisions_test.go

@@ -0,0 +1,36 @@
+package difc
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestShouldBypassCoarseDeny(t *testing.T) {
+	assert.True(t, ShouldBypassCoarseDeny(OperationRead))
+	assert.False(t, ShouldBypassCoarseDeny(OperationWrite))
+	assert.False(t, ShouldBypassCoarseDeny(OperationReadWrite))
+}
+
+func TestShouldCallLabelResponse(t *testing.T) {
+	assert.False(t, ShouldCallLabelResponse(OperationWrite, EnforcementStrict))
+	assert.False(t, ShouldCallLabelResponse(OperationReadWrite, EnforcementStrict))
+	assert.True(t, ShouldCallLabelResponse(OperationRead, EnforcementStrict))
+	assert.True(t, ShouldCallLabelResponse(OperationReadWrite, EnforcementFilter))
+	assert.True(t, ShouldCallLabelResponse(OperationReadWrite, EnforcementPropagate))
+}
+
+func TestShouldBlockFilteredResponse(t *testing.T) {
+	assert.True(t, ShouldBlockFilteredResponse(EnforcementStrict, 1))
+	assert.False(t, ShouldBlockFilteredResponse(EnforcementStrict, 0))
+	assert.False(t, ShouldBlockFilteredResponse(EnforcementFilter, 3))
+	assert.False(t, ShouldBlockFilteredResponse(EnforcementPropagate, 2))
+}
+
+func TestShouldAccumulateReadLabels(t *testing.T) {
+	assert.True(t, ShouldAccumulateReadLabels(OperationRead, EnforcementPropagate))
+	assert.True(t, ShouldAccumulateReadLabels(OperationReadWrite, EnforcementPropagate))
+	assert.False(t, ShouldAccumulateReadLabels(OperationWrite, EnforcementPropagate))
+	assert.False(t, ShouldAccumulateReadLabels(OperationRead, EnforcementStrict))
+	assert.False(t, ShouldAccumulateReadLabels(OperationRead, EnforcementFilter))
+}

internal/logger/common.go

@@ -215,6 +215,24 @@ import (
 // (markdown_logger.go) and logWithLevelAndServer (server_file_logger.go).
 // When adding a new LogLevel constant, add a corresponding entry here so
 // that all dispatch sites automatically support the new level.
+func makeLevelLogger(
+	dispatch func(level LogLevel, category, format string, args ...interface{}),
+	level LogLevel,
+) func(category, format string, args ...interface{}) {
+	return func(category, format string, args ...interface{}) {
+		dispatch(level, category, format, args...)
+	}
+}
+
+func makeServerLevelLogger(
+	dispatch func(serverID string, level LogLevel, category, format string, args ...interface{}),
+	level LogLevel,
+) func(serverID, category, format string, args ...interface{}) {
+	return func(serverID, category, format string, args ...interface{}) {
+		dispatch(serverID, level, category, format, args...)
+	}
+}
+
 var logFuncs = map[LogLevel]func(string, string, ...interface{}){
 	LogLevelInfo:  LogInfo,
 	LogLevelWarn:  LogWarn,

internal/logger/file_logger.go

@@ -113,25 +113,16 @@ func logWithLevel(level LogLevel, category, format string, args ...interface{})
 	})
 }
 
-// LogInfo logs an informational message
-func LogInfo(category, format string, args ...interface{}) {
-	logWithLevel(LogLevelInfo, category, format, args...)
-}
-
-// LogWarn logs a warning message
-func LogWarn(category, format string, args ...interface{}) {
-	logWithLevel(LogLevelWarn, category, format, args...)
-}
-
-// LogError logs an error message
-func LogError(category, format string, args ...interface{}) {
-	logWithLevel(LogLevelError, category, format, args...)
-}
-
-// LogDebug logs a debug message
-func LogDebug(category, format string, args ...interface{}) {
-	logWithLevel(LogLevelDebug, category, format, args...)
-}
+var (
+	// LogInfo logs an informational message.
+	LogInfo = makeLevelLogger(logWithLevel, LogLevelInfo)
+	// LogWarn logs a warning message.
+	LogWarn = makeLevelLogger(logWithLevel, LogLevelWarn)
+	// LogError logs an error message.
+	LogError = makeLevelLogger(logWithLevel, LogLevelError)
+	// LogDebug logs a debug message.
+	LogDebug = makeLevelLogger(logWithLevel, LogLevelDebug)
+)
 
 // CloseGlobalLogger closes the global file logger
 func CloseGlobalLogger() error {

internal/logger/markdown_logger.go

@@ -180,25 +180,16 @@ func logWithMarkdown(level LogLevel, category, format string, args ...interface{
 	})
 }
 
-// LogInfoMd logs to both regular and markdown loggers
-func LogInfoMd(category, format string, args ...interface{}) {
-	logWithMarkdown(LogLevelInfo, category, format, args...)
-}
-
-// LogWarnMd logs to both regular and markdown loggers
-func LogWarnMd(category, format string, args ...interface{}) {
-	logWithMarkdown(LogLevelWarn, category, format, args...)
-}
-
-// LogErrorMd logs to both regular and markdown loggers
-func LogErrorMd(category, format string, args ...interface{}) {
-	logWithMarkdown(LogLevelError, category, format, args...)
-}
-
-// LogDebugMd logs to both regular and markdown loggers
-func LogDebugMd(category, format string, args ...interface{}) {
-	logWithMarkdown(LogLevelDebug, category, format, args...)
-}
+var (
+	// LogInfoMd logs to both regular and markdown loggers.
+	LogInfoMd = makeLevelLogger(logWithMarkdown, LogLevelInfo)
+	// LogWarnMd logs to both regular and markdown loggers.
+	LogWarnMd = makeLevelLogger(logWithMarkdown, LogLevelWarn)
+	// LogErrorMd logs to both regular and markdown loggers.
+	LogErrorMd = makeLevelLogger(logWithMarkdown, LogLevelError)
+	// LogDebugMd logs to both regular and markdown loggers.
+	LogDebugMd = makeLevelLogger(logWithMarkdown, LogLevelDebug)
+)
 
 // CloseMarkdownLogger closes the global markdown logger
 func CloseMarkdownLogger() error {

internal/logger/server_file_logger.go

@@ -155,25 +155,16 @@ func logWithLevelAndServer(serverID string, level LogLevel, category, format str
 	}
 }
 
-// LogInfoWithServer logs an informational message to the server-specific log file
-func LogInfoWithServer(serverID, category, format string, args ...interface{}) {
-	logWithLevelAndServer(serverID, LogLevelInfo, category, format, args...)
-}
-
-// LogWarnWithServer logs a warning message to the server-specific log file
-func LogWarnWithServer(serverID, category, format string, args ...interface{}) {
-	logWithLevelAndServer(serverID, LogLevelWarn, category, format, args...)
-}
-
-// LogErrorWithServer logs an error message to the server-specific log file
-func LogErrorWithServer(serverID, category, format string, args ...interface{}) {
-	logWithLevelAndServer(serverID, LogLevelError, category, format, args...)
-}
-
-// LogDebugWithServer logs a debug message to the server-specific log file
-func LogDebugWithServer(serverID, category, format string, args ...interface{}) {
-	logWithLevelAndServer(serverID, LogLevelDebug, category, format, args...)
-}
+var (
+	// LogInfoWithServer logs an informational message to the server-specific log file.
+	LogInfoWithServer = makeServerLevelLogger(logWithLevelAndServer, LogLevelInfo)
+	// LogWarnWithServer logs a warning message to the server-specific log file.
+	LogWarnWithServer = makeServerLevelLogger(logWithLevelAndServer, LogLevelWarn)
+	// LogErrorWithServer logs an error message to the server-specific log file.
+	LogErrorWithServer = makeServerLevelLogger(logWithLevelAndServer, LogLevelError)
+	// LogDebugWithServer logs a debug message to the server-specific log file.
+	LogDebugWithServer = makeServerLevelLogger(logWithLevelAndServer, LogLevelDebug)
+)
 
 // CloseServerFileLogger closes the global server file logger
 func CloseServerFileLogger() error {

internal/proxy/handler.go

@@ -178,7 +178,7 @@ func (h *proxyHandler) handleWithDIFC(w http.ResponseWriter, r *http.Request, pa
 	evalResult := s.evaluator.Evaluate(agentLabels.Secrecy, agentLabels.Integrity, resource, operation)
 
 	if !evalResult.IsAllowed() {
-		if operation == difc.OperationRead {
+		if difc.ShouldBypassCoarseDeny(operation) {
 			// Read in filter mode: skip coarse block, proceed to fine-grained filtering
 			logHandler.Printf("[DIFC] Phase 2: coarse check failed for read, proceeding to Phase 3")
 		} else {
@@ -266,7 +266,7 @@ func (h *proxyHandler) handleWithDIFC(w http.ResponseWriter, r *http.Request, pa
 			}
 
 			// Strict mode: block entire response if any item filtered
-			if s.enforcementMode == difc.EnforcementStrict && filtered.GetFilteredCount() > 0 {
+			if difc.ShouldBlockFilteredResponse(s.enforcementMode, filtered.GetFilteredCount()) {
 				logHandler.Printf("[DIFC] STRICT: blocking response — %d filtered items", filtered.GetFilteredCount())
 				writeDIFCForbidden(w, fmt.Sprintf("DIFC policy violation: %d of %d items not accessible",
 					filtered.GetFilteredCount(), filtered.TotalCount))
@@ -318,7 +318,7 @@ func (h *proxyHandler) handleWithDIFC(w http.ResponseWriter, r *http.Request, pa
 	}
 
 	// **Phase 6: Label accumulation (propagate mode)**
-	if s.enforcementMode == difc.EnforcementPropagate && labeledData != nil {
+	if labeledData != nil && difc.ShouldAccumulateReadLabels(operation, s.enforcementMode) {
 		overall := labeledData.Overall()
 		agentLabels.AccumulateFromRead(overall)
 		logHandler.Printf("[DIFC] Phase 6: accumulated labels")

internal/server/unified.go

@@ -531,7 +531,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 	// For read operations in any mode, we skip the coarse-grained block
 	// and let the request proceed. Fine-grained filtering at Phase 5 will filter
 	// individual items from the response based on their actual labels from LabelResponse().
-	isReadOperation := (operation == difc.OperationRead)
+	isReadOperation := difc.ShouldBypassCoarseDeny(operation)
 	result := requestEvaluator.Evaluate(agentLabels.Secrecy, agentLabels.Integrity, resource, operation)
 
 	if !result.IsAllowed() {
@@ -603,8 +603,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 	// Per spec: LabelResponse() is only called for read operations in all modes,
 	// and for read-write operations in filter/propagate modes.
 	// For write operations and read-write in strict mode, skip LabelResponse().
-	isPureWrite := (operation == difc.OperationWrite)
-	shouldCallLabelResponse := !isPureWrite && (operation != difc.OperationReadWrite || enforcementMode != difc.EnforcementStrict)
+	shouldCallLabelResponse := difc.ShouldCallLabelResponse(operation, enforcementMode)
 
 	var labeledData difc.LabeledData
 	if shouldCallLabelResponse {
@@ -631,7 +630,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 				filtered.GetAccessibleCount(), filtered.TotalCount)
 
 			// **Strict mode: block entire response if ANY item is filtered**
-			if enforcementMode == difc.EnforcementStrict && filtered.GetFilteredCount() > 0 {
+			if difc.ShouldBlockFilteredResponse(enforcementMode, filtered.GetFilteredCount()) {
 				logger.LogWarn("difc", "STRICT MODE: Blocking entire response - %d/%d items violate DIFC policy",
 					filtered.GetFilteredCount(), filtered.TotalCount)
 				blockErr := fmt.Errorf("DIFC policy violation: %d of %d items in response are not accessible to agent %s",
@@ -664,7 +663,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 		// **Phase 6: Accumulate labels from this operation (for reads in PROPAGATE mode only)**
 		// Label accumulation should only happen when mode is EnforcementPropagate
 		// Filter mode does NOT accumulate - it just filters what the agent can see
-		if !isPureWrite && enforcementMode == difc.EnforcementPropagate {
+		if difc.ShouldAccumulateReadLabels(operation, enforcementMode) {
 			overall := labeledData.Overall()
 			agentLabels.AccumulateFromRead(overall)
 			logUnified.Printf("[DIFC] Agent %s accumulated labels (propagate mode) | Secrecy: %v | Integrity: %v",
@@ -675,7 +674,7 @@ func (us *UnifiedServer) callBackendTool(ctx context.Context, serverID, toolName
 		finalResult = backendResult
 
 		// **Phase 6: Accumulate labels from resource (for reads in PROPAGATE mode only)**
-		if !isPureWrite && enforcementMode == difc.EnforcementPropagate {
+		if difc.ShouldAccumulateReadLabels(operation, enforcementMode) {
 			agentLabels.AccumulateFromRead(resource)
 			logUnified.Printf("[DIFC] Agent %s accumulated labels (propagate mode) | Secrecy: %v | Integrity: %v",
 				agentID, agentLabels.GetSecrecyTags(), agentLabels.GetIntegrityTags())