Skip to content

Commit 974d1be

Browse files
ncodencode
committed
fixing test cases
1 parent d57e12c commit 974d1be

1 file changed

Lines changed: 171 additions & 146 deletions

File tree

‎pkg/auditserver/server_test.go‎

Lines changed: 171 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,15 @@ package auditserver
33
import (
44
"bytes"
55
"encoding/json"
6-
"github.com/panjf2000/gnet"
7-
"github.com/redis/go-redis/v9"
8-
"log/slog"
6+
"io/ioutil"
97
"net"
108
"os"
119
"testing"
10+
"time"
11+
12+
"github.com/expr-lang/expr"
13+
"github.com/spf13/viper"
14+
"log/slog"
1215
)
1316

1417
// mockConn is a mock implementation of gnet.Conn
@@ -37,194 +40,216 @@ func (m *mockConn) Peek(n int) (buf []byte, err error) { return nil, nil }
3740
func (m *mockConn) Next(n int) (buf []byte, err error) { return nil, nil }
3841

3942
func TestAuditServer_React(t *testing.T) {
40-
tests := []struct {
41-
name string
42-
input AuditLog
43-
expectedAction gnet.Action
44-
expectedLog bool
45-
}{
43+
// Create a temporary directory for log files
44+
tempDir := t.TempDir()
45+
46+
// Define rule group configurations
47+
ruleGroupConfigs := []RuleGroupConfig{
4648
{
47-
name: "Valid KV update operation",
48-
input: AuditLog{
49-
Type: "audit",
50-
Time: "2023-07-31T12:34:56Z",
51-
Auth: Auth{
52-
PolicyResults: struct {
53-
Allowed bool `json:"allowed"`
54-
}{Allowed: true},
55-
},
56-
Request: Request{
57-
Operation: "update",
58-
MountType: "kv",
59-
Path: "/secret/data/test",
60-
},
61-
Response: Response{
62-
MountType: "kv",
63-
},
64-
RemoteAddr: "192.168.1.1",
49+
Name: "normal_operations",
50+
Rules: []string{
51+
`Request.Operation in ["read", "update"] && Request.Path startsWith "secret/data/" && Auth.PolicyResults.Allowed == true`,
6552
},
66-
expectedAction: gnet.None,
67-
expectedLog: true,
68-
},
69-
{
70-
name: "Valid KV create operation",
71-
input: AuditLog{
72-
Type: "audit",
73-
Time: "2023-07-31T12:34:56Z",
74-
Auth: Auth{
75-
PolicyResults: struct {
76-
Allowed bool `json:"allowed"`
77-
}{Allowed: true},
78-
},
79-
Request: Request{
80-
Operation: "create",
81-
MountType: "kv",
82-
Path: "/secret/data/test",
83-
},
84-
Response: Response{
85-
MountType: "kv",
86-
},
87-
RemoteAddr: "192.168.1.1",
53+
LogFile: LogFileConfig{
54+
FilePath: tempDir + "/normal_operations.log",
55+
MaxSize: 1,
56+
MaxBackups: 1,
57+
MaxAge: 1,
58+
Compress: false,
8859
},
89-
expectedAction: gnet.None,
90-
expectedLog: true,
9160
},
9261
{
93-
name: "Valid KV delete operation",
94-
input: AuditLog{
95-
Type: "audit",
96-
Time: "2023-07-31T12:34:56Z",
97-
Auth: Auth{
98-
PolicyResults: struct {
99-
Allowed bool `json:"allowed"`
100-
}{Allowed: true},
101-
},
102-
Request: Request{
103-
Operation: "delete",
104-
MountType: "kv",
105-
Path: "/secret/data/test",
106-
},
107-
Response: Response{
108-
MountType: "kv",
109-
},
110-
RemoteAddr: "192.168.1.1",
62+
Name: "critical_events",
63+
Rules: []string{
64+
`Request.Operation == "delete" && Auth.PolicyResults.Allowed == true`,
65+
`Request.Path startsWith "secret/metadata/" && Auth.PolicyResults.Allowed == true`,
11166
},
112-
expectedAction: gnet.None,
113-
expectedLog: true,
114-
},
115-
{
116-
name: "Non-KV operation",
117-
input: AuditLog{
118-
Type: "audit",
119-
Time: "2023-07-31T12:34:56Z",
120-
Auth: Auth{
121-
PolicyResults: struct {
122-
Allowed bool `json:"allowed"`
123-
}{Allowed: true},
124-
},
125-
Request: Request{
126-
Operation: "update",
127-
MountType: "transit",
128-
Path: "/transit/keys/test",
129-
},
130-
Response: Response{
131-
MountType: "transit",
132-
},
133-
RemoteAddr: "192.168.1.1",
67+
LogFile: LogFileConfig{
68+
FilePath: tempDir + "/critical_events.log",
69+
MaxSize: 1,
70+
MaxBackups: 1,
71+
MaxAge: 1,
72+
Compress: false,
13473
},
135-
expectedAction: gnet.Close,
136-
expectedLog: false,
13774
},
75+
}
76+
77+
// Initialize viper with the rule group configurations
78+
viper.Set("rule_groups", ruleGroupConfigs)
79+
80+
// Create the AuditServer
81+
logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
82+
as := New(logger)
83+
84+
tests := []struct {
85+
name string
86+
input AuditLog
87+
expectedLogs map[string]bool // Map of log file names to whether they should contain the log
88+
}{
13889
{
139-
name: "Disallowed operation",
90+
name: "Normal operation - update",
14091
input: AuditLog{
141-
Type: "audit",
142-
Time: "2023-07-31T12:34:56Z",
92+
Type: "request",
93+
Time: "2024-09-17T13:00:00Z",
14394
Auth: Auth{
95+
DisplayName: "user1",
96+
Policies: []string{"default", "writer"},
14497
PolicyResults: struct {
145-
Allowed bool `json:"allowed"`
146-
}{Allowed: false},
98+
Allowed bool `json:"allowed"`
99+
GrantingPolicies []struct {
100+
Name string `json:"name"`
101+
NamespaceID string `json:"namespace_id"`
102+
Type string `json:"type"`
103+
} `json:"granting_policies"`
104+
}{
105+
Allowed: true,
106+
},
147107
},
148108
Request: Request{
149109
Operation: "update",
150-
MountType: "kv",
151-
Path: "/secret/data/test",
110+
Path: "secret/data/myapp/config",
152111
},
153-
Response: Response{
154-
MountType: "kv",
155-
},
156-
RemoteAddr: "192.168.1.1",
157112
},
158-
expectedAction: gnet.Close,
159-
expectedLog: false,
113+
expectedLogs: map[string]bool{
114+
tempDir + "/normal_operations.log": true,
115+
tempDir + "/critical_events.log": false,
116+
},
160117
},
118+
// Add more test cases as needed
161119
}
162120

163121
for _, tt := range tests {
164122
t.Run(tt.name, func(t *testing.T) {
165-
var logBuffer bytes.Buffer
166-
logger := slog.New(slog.NewJSONHandler(&logBuffer, &slog.HandlerOptions{Level: slog.LevelInfo}))
167-
168-
server := New(logger, nil)
169-
170-
inputJSON, err := json.Marshal(tt.input)
123+
// Serialize the audit log to JSON
124+
frame, err := json.Marshal(tt.input)
171125
if err != nil {
172-
t.Fatalf("Failed to marshal input: %v", err)
126+
t.Fatalf("Failed to marshal audit log: %v", err)
173127
}
174128

175-
_, action := server.React(inputJSON, &mockConn{})
176-
177-
t.Logf("Test case: %s", tt.name)
178-
t.Logf("Input: %s", string(inputJSON))
179-
t.Logf("Action: %v", action)
180-
t.Logf("Log buffer: %s", logBuffer.String())
129+
// Call React
130+
as.React(frame, &mockConn{})
181131

182-
if action != tt.expectedAction {
183-
t.Errorf("Expected action %v, but got %v", tt.expectedAction, action)
184-
}
132+
// Give some time for the log to be written
133+
time.Sleep(100 * time.Millisecond)
185134

186-
if tt.expectedLog {
187-
if logBuffer.Len() == 0 {
188-
t.Errorf("Expected log output, but got none")
189-
} else {
190-
var logEntry map[string]interface{}
191-
err := json.Unmarshal(logBuffer.Bytes(), &logEntry)
192-
if err != nil {
193-
t.Fatalf("Failed to parse log output: %v", err)
135+
// Check log files
136+
for logFile, shouldContain := range tt.expectedLogs {
137+
content, err := ioutil.ReadFile(logFile)
138+
if err != nil {
139+
if os.IsNotExist(err) && !shouldContain {
140+
// File doesn't exist as expected
141+
continue
194142
}
143+
t.Fatalf("Failed to read log file '%s': %v", logFile, err)
144+
}
195145

196-
expectedFields := []string{"operation", "path"}
197-
for _, field := range expectedFields {
198-
if _, ok := logEntry[field]; !ok {
199-
t.Errorf("Expected '%s' field in log, but it was missing", field)
200-
}
146+
if shouldContain {
147+
if !bytes.Contains(content, frame) {
148+
t.Errorf("Expected log file '%s' to contain the audit log", logFile)
149+
}
150+
} else {
151+
if len(content) > 0 {
152+
t.Errorf("Expected log file '%s' to be empty", logFile)
201153
}
202154
}
203-
} else if logBuffer.Len() > 0 {
204-
t.Errorf("Expected no log output, but got: %s", logBuffer.String())
155+
}
156+
157+
// Clean up log files for next test
158+
for logFile := range tt.expectedLogs {
159+
os.Remove(logFile)
205160
}
206161
})
207162
}
208163
}
209164

210165
func TestNew(t *testing.T) {
211-
// Test with nil logger and nil publisher
212-
server := New(nil, nil)
166+
// Define rule group configurations
167+
ruleGroupConfigs := []RuleGroupConfig{
168+
{
169+
Name: "test_group",
170+
Rules: []string{
171+
`Request.Operation == "update"`,
172+
},
173+
LogFile: LogFileConfig{
174+
FilePath: "test.log",
175+
MaxSize: 1,
176+
MaxBackups: 1,
177+
MaxAge: 1,
178+
Compress: false,
179+
},
180+
},
181+
}
182+
183+
// Initialize viper with the rule group configurations
184+
viper.Set("rule_groups", ruleGroupConfigs)
185+
186+
// Test with nil logger
187+
server := New(nil)
213188
if server.logger == nil {
214189
t.Errorf("Expected non-nil logger when initialized with nil")
215190
}
216-
if server.publisher == nil {
217-
t.Errorf("Expected non-nil publisher when initialized with nil")
191+
192+
if len(server.ruleGroups) != len(ruleGroupConfigs) {
193+
t.Errorf("Expected %d rule groups, got %d", len(ruleGroupConfigs), len(server.ruleGroups))
218194
}
219195

220-
// Test with custom logger and publisher
196+
// Test with custom logger
221197
customLogger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
222-
customPublisher := redis.NewClient(&redis.Options{Addr: "localhost:6379"})
223-
server = New(customLogger, customPublisher)
198+
server = New(customLogger)
224199
if server.logger != customLogger {
225200
t.Errorf("Expected custom logger to be used")
226201
}
227-
if server.publisher != customPublisher {
228-
t.Errorf("Expected custom publisher to be used")
202+
}
203+
204+
func TestRuleGroup_shouldLog(t *testing.T) {
205+
// Define a sample audit log
206+
auditLog := &AuditLog{
207+
Type: "request",
208+
Time: "2024-09-17T13:00:00Z",
209+
Auth: Auth{
210+
DisplayName: "user1",
211+
Policies: []string{"default", "writer"},
212+
PolicyResults: struct {
213+
Allowed bool `json:"allowed"`
214+
GrantingPolicies []struct {
215+
Name string `json:"name"`
216+
NamespaceID string `json:"namespace_id"`
217+
Type string `json:"type"`
218+
} `json:"granting_policies"`
219+
}{
220+
Allowed: true,
221+
},
222+
},
223+
Request: Request{
224+
Operation: "update",
225+
Path: "secret/data/myapp/config",
226+
},
227+
}
228+
229+
// Compile a rule
230+
ruleStr := `Request.Operation == "update" && Request.Path startsWith "secret/data/" && Auth.PolicyResults.Allowed == true`
231+
program, err := expr.Compile(ruleStr, expr.Env(&AuditLog{}))
232+
if err != nil {
233+
t.Fatalf("Failed to compile rule: %v", err)
234+
}
235+
236+
// Create a RuleGroup
237+
rg := &RuleGroup{
238+
Name: "test_group",
239+
CompiledRules: []CompiledRule{
240+
{Program: program},
241+
},
242+
}
243+
244+
// Test shouldLog
245+
if !rg.shouldLog(auditLog) {
246+
t.Errorf("Expected shouldLog to return true, got false")
247+
}
248+
249+
// Modify audit log to not match
250+
auditLog.Request.Operation = "read"
251+
252+
if rg.shouldLog(auditLog) {
253+
t.Errorf("Expected shouldLog to return false, got true")
229254
}
230255
}

0 commit comments

Comments
 (0)