goproxy/cmd/example/main.go

package main

import (
	"flag"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/darkit/goproxy/internal/cache"
	"github.com/darkit/goproxy/internal/config"
	"github.com/darkit/goproxy/internal/healthcheck"
	"github.com/darkit/goproxy/internal/loadbalance"
	"github.com/darkit/goproxy/internal/metrics"
	"github.com/darkit/goproxy/internal/proxy"
)

var (
	// 监听地址
	addr = flag.String("addr", ":8080", "代理服务器监听地址")
	// 上游代理服务器
	upstream = flag.String("upstream", "", "上游代理服务器地址，多个地址用逗号分隔")
	// 是否启用负载均衡
	enableLoadBalance = flag.Bool("enable-lb", false, "是否启用负载均衡")
	// 是否启用健康检查
	enableHealthCheck = flag.Bool("enable-hc", false, "是否启用健康检查")
	// 是否启用缓存
	enableCache = flag.Bool("enable-cache", false, "是否启用缓存")
	// 是否启用重试
	enableRetry = flag.Bool("enable-retry", false, "是否启用重试")
	// 是否启用监控
	enableMetrics = flag.Bool("enable-metrics", false, "是否启用监控")
	// 监控地址
	metricsAddr = flag.String("metrics-addr", ":8081", "监控服务器监听地址")
)

// 解析目标地址
func parseTargets(targets string) []string {
	if targets == "" {
		return nil
	}
	return strings.Split(targets, ",")
}

func main() {
	// 解析命令行参数
	flag.Parse()

	// 创建配置
	cfg := config.DefaultConfig()
	cfg.EnableLoadBalancing = *enableLoadBalance
	cfg.EnableHealthCheck = *enableHealthCheck
	cfg.EnableCache = *enableCache
	cfg.EnableRetry = *enableRetry

	// 创建选项
	opts := &proxy.Options{
		Config: cfg,
	}

	// 创建负载均衡器
	if *enableLoadBalance && *upstream != "" {
		lb := loadbalance.NewRoundRobinBalancer()
		for _, target := range parseTargets(*upstream) {
			lb.Add(target, 1)
		}
		opts.LoadBalancer = lb
	}

	// 创建健康检查器
	if *enableHealthCheck && opts.LoadBalancer != nil {
		hc := healthcheck.NewHealthChecker(&healthcheck.Config{
			Interval:   time.Second * 10,
			Timeout:    time.Second * 2,
			MaxFails:   3,
			MinSuccess: 2,
		})
		opts.HealthChecker = hc
	}

	// 创建缓存
	if *enableCache {
		c := cache.NewMemoryCache(time.Minute*5, time.Minute*5, 1000)
		opts.HTTPCache = c
	}

	// 创建监控
	if *enableMetrics {
		m := metrics.NewSimpleMetrics()
		opts.Metrics = m

		// 启动监控服务器
		go func() {
			mux := http.NewServeMux()
			handler := m.GetHandler()
			mux.Handle("/metrics", handler)
			log.Printf("监控服务器启动在 %s\n", *metricsAddr)
			if err := http.ListenAndServe(*metricsAddr, mux); err != nil {
				log.Fatalf("监控服务器启动失败: %v", err)
			}
		}()
	}

	// 创建自定义委托
	delegate := &CustomDelegate{}
	opts.Delegate = delegate

	// 创建代理
	p := proxy.New(opts)

	// 创建HTTP服务器
	server := &http.Server{
		Addr:    *addr,
		Handler: p,
	}

	// 启动HTTP服务器
	go func() {
		log.Printf("代理服务器启动在 %s\n", *addr)
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("代理服务器启动失败: %v", err)
		}
	}()

	// 等待信号
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	log.Println("正在关闭代理服务器...")
	server.Close()
	log.Println("代理服务器已关闭")
}

// CustomDelegate 自定义委托
type CustomDelegate struct {
	proxy.DefaultDelegate
}

// Connect 连接事件
func (d *CustomDelegate) Connect(ctx *proxy.Context, rw http.ResponseWriter) {
	log.Printf("收到连接: %s -> %s\n", ctx.Req.RemoteAddr, ctx.Req.URL.Host)
}

// BeforeRequest 请求前事件
func (d *CustomDelegate) BeforeRequest(ctx *proxy.Context) {
	log.Printf("请求: %s %s\n", ctx.Req.Method, ctx.Req.URL.String())
}

// BeforeResponse 响应前事件
func (d *CustomDelegate) BeforeResponse(ctx *proxy.Context, resp *http.Response, err error) {
	if err != nil {
		log.Printf("响应错误: %v\n", err)
		return
	}
	log.Printf("响应: %d %s\n", resp.StatusCode, resp.Status)
}

// ErrorLog 错误日志
func (d *CustomDelegate) ErrorLog(err error) {
	log.Printf("错误: %v\n", err)
}