133 lines
3.4 KiB
Go
133 lines
3.4 KiB
Go
package main
|
||
|
||
import (
|
||
"context"
|
||
"flag"
|
||
"log/slog"
|
||
"net/http"
|
||
"os"
|
||
"os/signal"
|
||
"syscall"
|
||
"time"
|
||
|
||
"github.com/darkit/goproxy"
|
||
"github.com/darkit/goproxy/pkg/dns"
|
||
)
|
||
|
||
func main() {
|
||
// 解析命令行参数
|
||
var (
|
||
listenAddr = flag.String("listen", ":8080", "代理服务器监听地址")
|
||
targetAddr = flag.String("target", "http://example.com", "目标服务器地址（反向代理模式下使用）")
|
||
proxyMode = flag.String("mode", "reverse", "代理模式: forward, reverse, transparent")
|
||
enableHTTPS = flag.Bool("https", false, "是否启用HTTPS拦截")
|
||
verbosity = flag.Int("v", 0, "日志详细级别 (0-3)")
|
||
host = flag.String("dns-host", "example.com", "要解析的主机名")
|
||
ip = flag.String("dns-ip", "127.0.0.1", "解析的IP地址")
|
||
)
|
||
flag.Parse()
|
||
|
||
// 配置日志
|
||
logLevel := slog.LevelInfo
|
||
switch *verbosity {
|
||
case 1:
|
||
logLevel = slog.LevelDebug
|
||
case 2:
|
||
logLevel = slog.LevelInfo
|
||
case 3:
|
||
logLevel = slog.LevelDebug - 3 // 更详细的调试级别
|
||
}
|
||
|
||
logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
|
||
Level: logLevel,
|
||
}))
|
||
slog.SetDefault(logger)
|
||
|
||
// 创建自定义DNS解析器
|
||
resolver := dns.NewResolver(
|
||
dns.WithFallback(true), // 如果找不到自定义规则，回退到系统DNS
|
||
dns.WithTTL(5*time.Minute), // 设置缓存TTL
|
||
)
|
||
|
||
// 添加自定义DNS解析规则
|
||
if err := resolver.Add(*host, *ip); err != nil {
|
||
slog.Error("添加DNS解析规则失败", "error", err)
|
||
return
|
||
}
|
||
|
||
// 添加通配符DNS解析规则示例
|
||
if err := resolver.AddWildcard("*.example.org", "192.168.1.2"); err != nil {
|
||
slog.Error("添加通配符DNS解析规则失败", "error", err)
|
||
return
|
||
}
|
||
|
||
slog.Info("已添加DNS解析规则", "host", *host, "ip", *ip)
|
||
|
||
// 创建代理选项
|
||
var (
|
||
proxy goproxy.UnifiedProxy
|
||
err error
|
||
)
|
||
|
||
// 根据代理模式创建不同类型的代理
|
||
switch *proxyMode {
|
||
case "forward":
|
||
proxy, err = goproxy.NewForwardProxy(*listenAddr, goproxy.WithUnifiedDNSResolver(resolver))
|
||
case "reverse":
|
||
proxy, err = goproxy.NewReverseProxy(*listenAddr, *targetAddr, goproxy.WithUnifiedDNSResolver(resolver))
|
||
case "transparent":
|
||
proxy, err = goproxy.NewTransparentProxy(*listenAddr, goproxy.WithUnifiedDNSResolver(resolver))
|
||
default:
|
||
slog.Error("不支持的代理模式", "mode", *proxyMode)
|
||
return
|
||
}
|
||
|
||
if err != nil {
|
||
slog.Error("创建代理失败", "error", err)
|
||
return
|
||
}
|
||
|
||
// 设置信号处理
|
||
sigCh := make(chan os.Signal, 1)
|
||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||
|
||
// 打印启动信息
|
||
slog.Info("代理服务器已启动",
|
||
"mode", *proxyMode,
|
||
"listen", *listenAddr,
|
||
"target", *targetAddr,
|
||
"https", *enableHTTPS)
|
||
|
||
// 启动HTTP服务器
|
||
server := &http.Server{
|
||
Addr: *listenAddr,
|
||
Handler: proxy,
|
||
}
|
||
|
||
// 监听和服务连接
|
||
go func() {
|
||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||
slog.Error("服务器运行失败", "error", err)
|
||
os.Exit(1)
|
||
}
|
||
}()
|
||
|
||
// 等待中断信号
|
||
<-sigCh
|
||
slog.Info("接收到信号，正在关闭...")
|
||
|
||
// 创建关闭上下文
|
||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
defer cancel()
|
||
|
||
// 优雅关闭HTTP服务器
|
||
if err := server.Shutdown(ctx); err != nil {
|
||
slog.Error("服务器关闭失败", "error", err)
|
||
}
|
||
|
||
// 关闭代理
|
||
if err := proxy.Close(); err != nil {
|
||
slog.Error("关闭代理时出错", "error", err)
|
||
}
|
||
}
|