Instrument remainder of file operations within os sink

instrumentation/sinks/os/disabled_test.go

@@ -13,33 +13,38 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestExamine_Disabled(t *testing.T) {
-	originalDisabled := config.IsZenDisabled()
-	defer config.SetZenDisabled(originalDisabled)
+func setupDisabledTest(t *testing.T) *testutil.MockCloudClient {
+	t.Helper()
 
+	originalDisabled := config.IsZenDisabled()
 	originalClient := agent.GetCloudClient()
 	originalBlocking := config.IsBlockingEnabled()
+
 	config.SetBlocking(true)
-	defer func() {
-		config.SetBlocking(originalBlocking)
-		agent.SetCloudClient(originalClient)
-	}()
+	config.SetZenDisabled(true)
 
 	mockClient := testutil.NewMockCloudClient()
 	agent.SetCloudClient(mockClient)
 
-	config.SetZenDisabled(true)
+	t.Cleanup(func() {
+		config.SetZenDisabled(originalDisabled)
+		config.SetBlocking(originalBlocking)
+		agent.SetCloudClient(originalClient)
+	})
+
+	return mockClient
+}
 
-	maliciousPath := "../../etc/passwd"
+func TestExamineOp_Disabled(t *testing.T) {
+	mockClient := setupDisabledTest(t)
 
-	err := os.Examine(maliciousPath)
+	err := os.ExamineOp("os.Chmod", "../../etc/passwd")
 
-	require.NoError(t, err, "Examine should return early with no error when zen is disabled")
+	require.NoError(t, err, "ExamineOp should return early with no error when zen is disabled")
 
 	select {
 	case <-mockClient.AttackDetectedEventSent:
 		t.Fatal("No attack should be detected when zen is disabled")
 	case <-time.After(50 * time.Millisecond):
-		// Expected: no attack detected
 	}
 }

instrumentation/sinks/os/examine.go

@@ -10,22 +10,18 @@ import (
 	"github.com/AikidoSec/firewall-go/zen"
 )
 
-func Examine(path string) error {
+func ExamineOp(op, path string) error {
 	if zen.IsDisabled() {
 		return nil
 	}
 
-	hooks.OnOperationCall("os.OpenFile", operation.KindFileSystem)
+	hooks.OnOperationCall(op, operation.KindFileSystem)
 
-	// The error that the vulnerability scan returns is NOT deferred with os.OpenFile
-	// We block and report immediately
-	err := vulnerabilities.ScanWithOptions(context.Background(), "os.OpenFile", pathtraversal.PathTraversalVulnerability, &pathtraversal.ScanArgs{
+	return vulnerabilities.ScanWithOptions(context.Background(), op, pathtraversal.PathTraversalVulnerability, &pathtraversal.ScanArgs{
 		FilePath:       path,
 		CheckPathStart: true,
 	}, vulnerabilities.ScanOptions{
 		DeferReporting: false,
 		Module:         "os",
 	})
-
-	return err
 }

instrumentation/sinks/os/examine_test.go

@@ -12,7 +12,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestExamine_TracksOperationStats(t *testing.T) {
+func TestExamineOp_TracksOperationStats(t *testing.T) {
 	originalDisabled := zen.IsDisabled()
 	defer zen.SetDisabled(originalDisabled)
 
@@ -24,16 +24,13 @@ func TestExamine_TracksOperationStats(t *testing.T) {
 	mockClient := testutil.NewMockCloudClient()
 	agent.SetCloudClient(mockClient)
 
-	// Clear stats before test
 	agent.Stats().GetAndClear()
 
-	// Open multiple files
-	_ = os.Examine("/tmp/file1.txt")
-	_ = os.Examine("/tmp/file2.txt")
-	_ = os.Examine("/var/log/test.log")
+	_ = os.ExamineOp("os.OpenFile", "/tmp/file1.txt")
+	_ = os.ExamineOp("os.OpenFile", "/tmp/file2.txt")
+	_ = os.ExamineOp("os.Chmod", "/tmp/test.txt")
 
-	// Get stats and verify operations were tracked
 	stats := agent.Stats().GetAndClear()
-	require.Contains(t, stats.Operations, "os.OpenFile")
-	require.Equal(t, 3, stats.Operations["os.OpenFile"].Total, "should track 3 file operations")
+	require.Equal(t, 2, stats.Operations["os.OpenFile"].Total)
+	require.Equal(t, 1, stats.Operations["os.Chmod"].Total)
 }

instrumentation/sinks/os/integration_test.go

@@ -54,6 +54,49 @@ func (m *mockCloudClient) SendAttackDetectedEvent(agentInfo cloud.AgentInfo, req
 	m.attackDetectedEventSent <- struct{}{}
 }
 
+func setupBlockingTest(t *testing.T) (*mockCloudClient, context.Context) {
+	t.Helper()
+	require.NoError(t, zen.Protect())
+
+	originalClient := agent.GetCloudClient()
+	original := config.IsBlockingEnabled()
+	config.SetBlocking(true)
+
+	t.Cleanup(func() {
+		config.SetBlocking(original)
+		agent.SetCloudClient(originalClient)
+	})
+
+	client := &mockCloudClient{
+		attackDetectedEventSent: make(chan struct{}),
+	}
+	agent.SetCloudClient(client)
+
+	req := httptest.NewRequest("GET", "/route?path=../test.txt", http.NoBody)
+	ip := "127.0.0.1"
+	ctx := request.SetContext(context.Background(), req, request.ContextData{
+		Source:        "test",
+		Route:         "/route",
+		RemoteAddress: &ip,
+	})
+
+	return client, ctx
+}
+
+func assertAttackDetected(t *testing.T, client *mockCloudClient, operation string) {
+	t.Helper()
+	select {
+	case <-client.attackDetectedEventSent:
+		assert.Equal(t, "path_traversal", client.capturedAttack.Kind)
+		assert.True(t, client.capturedAttack.Blocked)
+		assert.Equal(t, "../test.txt", client.capturedAttack.Payload)
+		assert.Equal(t, operation, client.capturedAttack.Operation)
+		assert.Equal(t, "os", client.capturedAttack.Module)
+	case <-time.After(1 * time.Second):
+		t.Fatal("timeout waiting for attack event")
+	}
+}
+
 func TestOpenFileIsAutomaticallyInstrumented(t *testing.T) {
 	require.NoError(t, zen.Protect())
 
@@ -181,6 +224,79 @@ func TestOpenFileIsNotBlockedWhenInMonitoringMode(t *testing.T) {
 	}
 }
 
+// TestRemoveIsAutomaticallyInstrumented covers single-path ops that return error (Chmod, Chown, etc.)
+func TestRemoveIsAutomaticallyInstrumented(t *testing.T) {
+	client, ctx := setupBlockingTest(t)
+
+	request.WrapWithGLS(ctx, func() {
+		err := os.Remove("/tmp/" + "../test.txt")
+
+		var detectedErr *vulnerabilities.AttackDetectedError
+		require.ErrorAs(t, err, &detectedErr)
+	})
+
+	assertAttackDetected(t, client, "os.Remove")
+}
+
+// TestRenameIsAutomaticallyInstrumented covers two-path ops, verifying both arguments are checked.
+func TestRenameIsAutomaticallyInstrumented(t *testing.T) {
+	t.Run("first path", func(t *testing.T) {
+		client, ctx := setupBlockingTest(t)
+
+		request.WrapWithGLS(ctx, func() {
+			err := os.Rename("/tmp/"+"../test.txt", "/tmp/newname.txt")
+
+			var detectedErr *vulnerabilities.AttackDetectedError
+			require.ErrorAs(t, err, &detectedErr)
+		})
+
+		assertAttackDetected(t, client, "os.Rename")
+	})
+
+	t.Run("second path", func(t *testing.T) {
+		client, ctx := setupBlockingTest(t)
+
+		request.WrapWithGLS(ctx, func() {
+			err := os.Rename("/tmp/oldname.txt", "/tmp/"+"../test.txt")
+
+			var detectedErr *vulnerabilities.AttackDetectedError
+			require.ErrorAs(t, err, &detectedErr)
+		})
+
+		assertAttackDetected(t, client, "os.Rename")
+	})
+}
+
+// TestReadlinkIsAutomaticallyInstrumented covers single-path ops that return (string, error).
+func TestReadlinkIsAutomaticallyInstrumented(t *testing.T) {
+	client, ctx := setupBlockingTest(t)
+
+	request.WrapWithGLS(ctx, func() {
+		result, err := os.Readlink("/tmp/" + "../test.txt")
+
+		require.Empty(t, result)
+		var detectedErr *vulnerabilities.AttackDetectedError
+		require.ErrorAs(t, err, &detectedErr)
+	})
+
+	assertAttackDetected(t, client, "os.Readlink")
+}
+
+// TestReadDirIsAutomaticallyInstrumented covers ReadDir which returns ([]DirEntry, error).
+func TestReadDirIsAutomaticallyInstrumented(t *testing.T) {
+	client, ctx := setupBlockingTest(t)
+
+	request.WrapWithGLS(ctx, func() {
+		entries, err := os.ReadDir("/tmp/" + "../test.txt")
+
+		require.Nil(t, entries)
+		var detectedErr *vulnerabilities.AttackDetectedError
+		require.ErrorAs(t, err, &detectedErr)
+	})
+
+	assertAttackDetected(t, client, "os.ReadDir")
+}
+
 func (m *mockCloudClient) SendAttackWaveDetectedEvent(agentInfo cloud.AgentInfo, request cloud.AttackWaveRequestInfo, attack cloud.AttackWaveDetails) {
 	panic("not implemented")
 }

instrumentation/sinks/os/zen.instrument.yml

@@ -3,24 +3,86 @@ meta:
   description: Protection from Local File Inclusion (LFI) Attacks.
 
 rules:
-  # Inject go:linkname declaration for Examine
-  - id: os.Examine.linkname
+  # Inject go:linkname declaration for ExamineOp
+  - id: os.ExamineOp.linkname
     type: inject-decl
     package: os
     anchor: Getpid
     links:
       - github.com/AikidoSec/firewall-go/instrumentation/sinks/os
     template: |
-      //go:linkname __aikido_os_Examine github.com/AikidoSec/firewall-go/instrumentation/sinks/os.Examine
-      func __aikido_os_Examine(string) error
+      //go:linkname __aikido_os_ExamineOp github.com/AikidoSec/firewall-go/instrumentation/sinks/os.ExamineOp
+      func __aikido_os_ExamineOp(string, string) error
 
   # Prepend security check to os.OpenFile
   - id: os.OpenFile
     type: prepend
     package: os
     function: OpenFile
     template: |
-      _aikido_block := __aikido_os_Examine({{ .Function.Argument 0 }})
+      _aikido_block := __aikido_os_ExamineOp("os.OpenFile", {{ .Function.Argument 0 }})
+      if _aikido_block != nil {
+        return nil, _aikido_block
+      }
+
+  # Single-path operations that return error
+  - id: os.SinglePathErrorOps
+    type: prepend
+    package: os
+    functions:
+      - Chmod
+      - Chown
+      - Lchown
+      - Chtimes
+      - Mkdir
+      - MkdirAll
+      - Remove
+      - RemoveAll
+      - Truncate
+    template: |
+      _aikido_block := __aikido_os_ExamineOp("os.{{ .Function.Name }}", {{ .Function.Argument 0 }})
+      if _aikido_block != nil {
+        return _aikido_block
+      }
+
+  # Two-path operations that return error
+  - id: os.TwoPathErrorOps
+    type: prepend
+    package: os
+    functions:
+      - Rename
+      - Symlink
+      - Link
+    template: |
+      _aikido_block := __aikido_os_ExamineOp("os.{{ .Function.Name }}", {{ .Function.Argument 0 }})
+      if _aikido_block != nil {
+        return _aikido_block
+      }
+      _aikido_block = __aikido_os_ExamineOp("os.{{ .Function.Name }}", {{ .Function.Argument 1 }})
+      if _aikido_block != nil {
+        return _aikido_block
+      }
+
+  # Single-path operations that return (string, error)
+  - id: os.SinglePathStringErrorOps
+    type: prepend
+    package: os
+    functions:
+      - Readlink
+      - MkdirTemp
+    template: |
+      _aikido_block := __aikido_os_ExamineOp("os.{{ .Function.Name }}", {{ .Function.Argument 0 }})
+      if _aikido_block != nil {
+        return "", _aikido_block
+      }
+
+  # ReadDir returns ([]DirEntry, error)
+  - id: os.ReadDir
+    type: prepend
+    package: os
+    function: ReadDir
+    template: |
+      _aikido_block := __aikido_os_ExamineOp("os.ReadDir", {{ .Function.Argument 0 }})
       if _aikido_block != nil {
         return nil, _aikido_block
       }