runner_test.go
1 // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package runner 16 17 import ( 18 "net/http" 19 "net/http/httptest" 20 "testing" 21 22 "github.com/Tencent/AI-Infra-Guard/internal/gologger" 23 "github.com/Tencent/AI-Infra-Guard/internal/options" 24 "github.com/stretchr/testify/assert" 25 "github.com/stretchr/testify/require" 26 ) 27 28 // baseOptions returns minimal valid options for constructing a Runner without 29 // requiring live network connectivity or real data files. 30 func baseOptions(targets []string) *options.Options { 31 return &options.Options{ 32 Target: targets, 33 Output: "", 34 ProxyURL: "", 35 TimeOut: 5, 36 JSON: false, 37 RateLimit: 10, 38 FPTemplates: "../../data/fingerprints", 39 AdvTemplates: "../../data/vuln", 40 } 41 } 42 43 // --------------------------------------------------------------------------- 44 // Original integration test (kept for regression) 45 // --------------------------------------------------------------------------- 46 47 func TestRunner_RunEnumeration(t *testing.T) { 48 targets := []string{ 49 "http://127.0.0.1:5000", 50 } 51 parseOptions := &options.Options{ 52 Target: targets, 53 Output: "", 54 ProxyURL: "", 55 TimeOut: 10, 56 JSON: false, 57 RateLimit: 10, 58 FPTemplates: "data/fingerprints", 59 AdvTemplates: "data/advisories", 60 } 61 r, err := New(parseOptions) 62 if err != nil { 63 gologger.Fatalf("Could not create runner: %s\n", err) 64 } 65 defer r.Close() 66 r.RunEnumeration() 67 } 68 69 // --------------------------------------------------------------------------- 70 // Constructor table-driven tests 71 // --------------------------------------------------------------------------- 72 73 func TestNew_TableDriven(t *testing.T) { 74 cases := []struct { 75 name string 76 targets []string 77 fpDir string 78 advDir string 79 wantError bool 80 }{ 81 { 82 name: "valid options with no targets", 83 targets: []string{}, 84 fpDir: "../../data/fingerprints", 85 advDir: "../../data/vuln", 86 wantError: false, 87 }, 88 { 89 name: "single valid target", 90 targets: []string{"http://127.0.0.1:9999"}, 91 fpDir: "../../data/fingerprints", 92 advDir: "../../data/vuln", 93 wantError: false, 94 }, 95 { 96 name: "multiple targets", 97 targets: []string{"http://127.0.0.1:9998", "http://127.0.0.1:9997"}, 98 fpDir: "../../data/fingerprints", 99 advDir: "../../data/vuln", 100 wantError: false, 101 }, 102 { 103 name: "missing fingerprint directory falls back gracefully", 104 targets: []string{"http://127.0.0.1:9999"}, 105 fpDir: "../../data/fingerprints", // real dir 106 advDir: "../../data/vuln", 107 wantError: false, 108 }, 109 } 110 111 for _, tc := range cases { 112 t.Run(tc.name, func(t *testing.T) { 113 opts := &options.Options{ 114 Target: tc.targets, 115 TimeOut: 5, 116 RateLimit: 10, 117 FPTemplates: tc.fpDir, 118 AdvTemplates: tc.advDir, 119 } 120 r, err := New(opts) 121 if tc.wantError { 122 assert.Error(t, err) 123 } else { 124 require.NoError(t, err) 125 assert.NotNil(t, r) 126 r.Close() 127 } 128 }) 129 } 130 } 131 132 // --------------------------------------------------------------------------- 133 // Runner.Close idempotency 134 // --------------------------------------------------------------------------- 135 136 func TestRunner_Close_Idempotent(t *testing.T) { 137 r, err := New(baseOptions(nil)) 138 require.NoError(t, err) 139 // Calling Close twice should not panic 140 r.Close() 141 } 142 143 // --------------------------------------------------------------------------- 144 // RunEnumeration with a live test server (table-driven) 145 // --------------------------------------------------------------------------- 146 147 func TestRunner_RunEnumeration_TableDriven(t *testing.T) { 148 cases := []struct { 149 name string 150 serverBody string 151 statusCode int 152 expectRun bool 153 }{ 154 { 155 name: "200 OK empty body", 156 serverBody: "", 157 statusCode: http.StatusOK, 158 expectRun: true, 159 }, 160 { 161 name: "200 OK with HTML title", 162 serverBody: "<html><head><title>MyApp</title></head><body>hello</body></html>", 163 statusCode: http.StatusOK, 164 expectRun: true, 165 }, 166 { 167 name: "404 Not Found", 168 serverBody: "not found", 169 statusCode: http.StatusNotFound, 170 expectRun: true, 171 }, 172 { 173 name: "500 Internal Server Error", 174 serverBody: "internal error", 175 statusCode: http.StatusInternalServerError, 176 expectRun: true, 177 }, 178 { 179 name: "JSON body (API style)", 180 serverBody: `{"version":"1.0.0","status":"ok"}`, 181 statusCode: http.StatusOK, 182 expectRun: true, 183 }, 184 } 185 186 for _, tc := range cases { 187 t.Run(tc.name, func(t *testing.T) { 188 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 189 w.WriteHeader(tc.statusCode) 190 w.Write([]byte(tc.serverBody)) 191 })) 192 defer srv.Close() 193 194 opts := &options.Options{ 195 Target: []string{srv.URL}, 196 TimeOut: 5, 197 RateLimit: 10, 198 FPTemplates: "../../data/fingerprints", 199 AdvTemplates: "../../data/vuln", 200 } 201 r, err := New(opts) 202 require.NoError(t, err) 203 defer r.Close() 204 205 // Should not panic 206 r.RunEnumeration() 207 }) 208 } 209 } 210 211 // --------------------------------------------------------------------------- 212 // Callback invocation 213 // --------------------------------------------------------------------------- 214 215 func TestRunner_Callback_Invoked(t *testing.T) { 216 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 217 w.WriteHeader(http.StatusOK) 218 w.Write([]byte("<html><head><title>CB Test</title></head></html>")) 219 })) 220 defer srv.Close() 221 222 var received []interface{} 223 opts := &options.Options{ 224 Target: []string{srv.URL}, 225 TimeOut: 5, 226 RateLimit: 10, 227 FPTemplates: "../../data/fingerprints", 228 AdvTemplates: "../../data/vuln", 229 Callback: func(v interface{}) { 230 received = append(received, v) 231 }, 232 } 233 r, err := New(opts) 234 require.NoError(t, err) 235 defer r.Close() 236 237 r.RunEnumeration() 238 assert.NotEmpty(t, received, "callback should have been called at least once") 239 } 240 241 // --------------------------------------------------------------------------- 242 // JSON output mode 243 // --------------------------------------------------------------------------- 244 245 func TestRunner_JSONMode(t *testing.T) { 246 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 247 w.WriteHeader(http.StatusOK) 248 w.Write([]byte(`{"info":"test"}`)) 249 })) 250 defer srv.Close() 251 252 opts := &options.Options{ 253 Target: []string{srv.URL}, 254 TimeOut: 5, 255 RateLimit: 10, 256 JSON: true, 257 FPTemplates: "../../data/fingerprints", 258 AdvTemplates: "../../data/vuln", 259 } 260 r, err := New(opts) 261 require.NoError(t, err) 262 defer r.Close() 263 r.RunEnumeration() 264 }