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  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  package preload
 20  
 21  import (
 22  	"net/http"
 23  	"net/http/httptest"
 24  	"testing"
 25  	"time"
 26  
 27  	"github.com/Tencent/AI-Infra-Guard/common/fingerprints/parser"
 28  	"github.com/Tencent/AI-Infra-Guard/pkg/httpx"
 29  	"github.com/stretchr/testify/assert"
 30  )
 31  
 32  func newHTTPXForTest(t *testing.T) *httpx.HTTPX {
 33  	t.Helper()
 34  	httpOptions := &httpx.HTTPOptions{
 35  		Timeout:          5 * time.Second,
 36  		RetryMax:         1,
 37  		FollowRedirects:  false,
 38  		HTTPProxy:        "",
 39  		Unsafe:           false,
 40  		DefaultUserAgent: httpx.GetRandomUserAgent(),
 41  		Dialer:           nil,
 42  	}
 43  	hp, err := httpx.NewHttpx(httpOptions)
 44  	assert.NoError(t, err)
 45  	return hp
 46  }
 47  
 48  func TestEvalFpVersionExact(t *testing.T) {
 49  	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 50  		if r.URL.Path == "/exact" {
 51  			w.Header().Set("X-Version", "1.2.3")
 52  			_, _ = w.Write([]byte("ok"))
 53  			return
 54  		}
 55  		http.NotFound(w, r)
 56  	}))
 57  	defer server.Close()
 58  
 59  	data := `info:
 60    name: test
 61    author: test
 62    severity: info
 63    metadata:
 64      product: test
 65      vendor: test
 66  version:
 67    - method: GET
 68      path: '/exact'
 69      extractor:
 70        part: header
 71        group: '1'
 72        regex: 'X-Version:\s*([0-9.]+)'
 73  `
 74  	fp, err := parser.InitFingerPrintFromData([]byte(data))
 75  	assert.NoError(t, err)
 76  
 77  	hp := newHTTPXForTest(t)
 78  	version, err := EvalFpVersion(server.URL, hp, *fp)
 79  	assert.NoError(t, err)
 80  	assert.Equal(t, "1.2.3", version)
 81  }
 82  
 83  func TestEvalFpVersionFuzzyIntersection(t *testing.T) {
 84  	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 85  		switch r.URL.Path {
 86  		case "/fuzzy1":
 87  			_, _ = w.Write([]byte("range-one"))
 88  		case "/fuzzy2":
 89  			_, _ = w.Write([]byte("range-two"))
 90  		default:
 91  			http.NotFound(w, r)
 92  		}
 93  	}))
 94  	defer server.Close()
 95  
 96  	data := `info:
 97    name: test
 98    author: test
 99    severity: info
100    metadata:
101      product: test
102      vendor: test
103  version:
104    - method: GET
105      path: '/fuzzy1'
106      matchers:
107        - body="range-one"
108      versionrange: '>=1.0.0,<2.0.0'
109    - method: GET
110      path: '/fuzzy2'
111      matchers:
112        - body="range-two"
113      versionrange: '>=1.5.0,<3.0.0'
114  `
115  	fp, err := parser.InitFingerPrintFromData([]byte(data))
116  	assert.NoError(t, err)
117  
118  	hp := newHTTPXForTest(t)
119  	version, err := EvalFpVersion(server.URL, hp, *fp)
120  	assert.NoError(t, err)
121  	assert.Equal(t, ">=1.5.0,<2.0.0", version)
122  }
123  
124  func TestEvalFpVersionFuzzyHashIntersection(t *testing.T) {
125  	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
126  		switch r.URL.Path {
127  		case "/fuzzy1":
128  			_, _ = w.Write([]byte("range-one"))
129  		case "/fuzzy2":
130  			_, _ = w.Write([]byte("range-two"))
131  		default:
132  			http.NotFound(w, r)
133  		}
134  	}))
135  	defer server.Close()
136  
137  	data := `info:
138    name: test
139    author: test
140    severity: info
141    metadata:
142      product: test
143      vendor: test
144  version:
145    - method: GET
146      path: '/fuzzy1'
147      matchers:
148        - hash=="012c14d354f49d6e682efaa1e8d3f1433ff7da7093b2b5964aac1302303f52b4"
149      versionrange: '>=1.0.0,<2.0.0'
150    - method: GET
151      path: '/fuzzy2'
152      matchers:
153        - hash=="1773f6ec9285e9d638a94a19375353fa9c8c891c732d80a996724ed8017fe196"
154      versionrange: '>=1.5.0,<3.0.0'
155  `
156  	fp, err := parser.InitFingerPrintFromData([]byte(data))
157  	assert.NoError(t, err)
158  
159  	hp := newHTTPXForTest(t)
160  	version, err := EvalFpVersion(server.URL, hp, *fp)
161  	assert.NoError(t, err)
162  	assert.Equal(t, ">=1.5.0,<2.0.0", version)
163  }
164  
165  func TestEvalFpVersionFuzzyNoMatch(t *testing.T) {
166  	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
167  		_, _ = w.Write([]byte("no-match"))
168  	}))
169  	defer server.Close()
170  
171  	data := `info:
172    name: test
173    author: test
174    severity: info
175    metadata:
176      product: test
177      vendor: test
178  version:
179    - method: GET
180      path: '/fuzzy'
181      matchers:
182        - body="another"
183      versionrange: '>=1.0.0,<2.0.0'
184  `
185  	fp, err := parser.InitFingerPrintFromData([]byte(data))
186  	assert.NoError(t, err)
187  
188  	hp := newHTTPXForTest(t)
189  	version, err := EvalFpVersion(server.URL, hp, *fp)
190  	assert.NoError(t, err)
191  	assert.Equal(t, "", version)
192  }