mirror of
https://gitee.com/samwaf/SamWaf.git
synced 2026-11-01 02:31:39 +08:00
103 lines
2.5 KiB
Go
103 lines
2.5 KiB
Go
package wafenginecore
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestWafEngine_Start_WAF(t *testing.T) {
|
|
|
|
respContentType := "text/html; charset=utf-8"
|
|
respContentType = strings.Replace(respContentType, "; charset=utf-8", "", -1)
|
|
println(respContentType)
|
|
// 创建被测试的HTTP服务器
|
|
testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
fmt.Fprintln(w, "Hello, world!")
|
|
}))
|
|
defer testServer.Close()
|
|
|
|
transport := &http.Transport{
|
|
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
|
// 在这里可以自定义 DNS 解析过程
|
|
host, port, err := net.SplitHostPort(addr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
// 使用解析后的 IP 地址进行连接
|
|
dialer := net.Dialer{
|
|
Timeout: 30 * time.Second,
|
|
KeepAlive: 30 * time.Second,
|
|
}
|
|
//通过自定义nameserver获取域名解析的IP
|
|
//ips, _ := dialer.Resolver.LookupHost(ctx, host)
|
|
//for _, s := range ips {
|
|
// log.Println(s)
|
|
//}
|
|
|
|
// 创建链接
|
|
if host == "www.qdbinet.com" {
|
|
ip := "127.0.0.1"
|
|
log.Println(ip)
|
|
log.Println(port)
|
|
conn, err := dialer.DialContext(ctx, network, ip+":"+"81")
|
|
if err == nil {
|
|
return conn, nil
|
|
}
|
|
}
|
|
|
|
return dialer.DialContext(ctx, network, host+":80")
|
|
},
|
|
}
|
|
// 创建反向代理
|
|
// 创建反向代理
|
|
proxy := httputil.NewSingleHostReverseProxy(&url.URL{
|
|
Scheme: "http",
|
|
Host: "www.qdbinet.com",
|
|
})
|
|
proxy.Transport = transport
|
|
proxy.ModifyResponse = modifyResponse()
|
|
// 创建测试请求
|
|
req, err := http.NewRequest("GET", "/", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// 使用httptest来捕获响应
|
|
recorder := httptest.NewRecorder()
|
|
|
|
// 调用被测试的处理函数
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
proxy.ServeHTTP(w, r)
|
|
})
|
|
handler.ServeHTTP(recorder, req)
|
|
|
|
// 验证代理是否正常工作
|
|
expectedResponse := "Hello, world!\n"
|
|
if recorder.Body.String() != expectedResponse {
|
|
t.Errorf("Proxy response mismatch: got %v want %v",
|
|
recorder.Body.String(), expectedResponse)
|
|
}
|
|
|
|
}
|
|
func modifyResponse() func(*http.Response) error {
|
|
return func(resp *http.Response) error {
|
|
resp.Header.Set("WAF", "SamWAF")
|
|
// 遍历头部数据
|
|
fmt.Println("Response Headers:")
|
|
for key, values := range resp.Header {
|
|
fmt.Printf("%s: %s\n", key, values)
|
|
}
|
|
fmt.Printf("%s: %d\n", "header长度", resp.ContentLength)
|
|
return nil
|
|
}
|
|
}
|