【GO语言小案例手记】基于GIN的简易代理网关

背景

正好最近对GO也有点兴趣，搞个小项目练练手。

目标

1. 网关需要能够根据路由自动映射到服务
2. 支持轮询、加权轮询、随机轮询三种算法
3. 简单好理解好使用，最好一个配置文件就能跑起来
4. 网关本身需要能够健康检查
5. 能够劫持请求修改请求头

开工

依赖

module test-gin

go 1.22.0

require (
	github.com/bytedance/sonic v1.9.1 // indirect
	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
	github.com/gin-contrib/sse v0.1.0 // indirect
	github.com/gin-gonic/gin v1.9.1 // indirect
	github.com/go-playground/locales v0.14.1 // indirect
	github.com/go-playground/universal-translator v0.18.1 // indirect
	github.com/go-playground/validator/v10 v10.14.0 // indirect
	github.com/goccy/go-json v0.10.2 // indirect
	github.com/json-iterator/go v1.1.12 // indirect
	github.com/klauspost/cpuid/v2 v2.2.4 // indirect
	github.com/leodido/go-urn v1.2.4 // indirect
	github.com/mattn/go-isatty v0.0.19 // indirect
	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
	github.com/modern-go/reflect2 v1.0.2 // indirect
	github.com/pelletier/go-toml/v2 v2.0.8 // indirect
	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
	github.com/ugorji/go/codec v1.2.11 // indirect
	golang.org/x/arch v0.3.0 // indirect
	golang.org/x/crypto v0.9.0 // indirect
	golang.org/x/net v0.10.0 // indirect
	golang.org/x/sys v0.8.0 // indirect
	golang.org/x/text v0.9.0 // indirect
	google.golang.org/protobuf v1.30.0 // indirect
	gopkg.in/yaml.v3 v3.0.1 // indirect
)

主体代码

package main

import (
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
)

type Target struct {
	URL    string `json:"url"`
	Weight int    `json:"weight"`
}

type RouteConfig struct {
	Path     string   `json:"path"`
	Targets  []Target `json:"targets"`
	Strategy string   `json:"strategy"`
}

type Config struct {
	Routes []RouteConfig `json:"routes"`
}

// 更新Balancer结构体
type Balancer struct {
	targets     []*url.URL
	weights     []int
	totalWeight int
	strategy    string
	counter     uint64
	rnd         *rand.Rand
	weightIndex int
	weightAccum int
}

// 更新NewBalancer函数
func NewBalancer(route RouteConfig, rnd *rand.Rand) (*Balancer, error) {
	var urls []*url.URL
	var weights []int
	totalWeight := 0

	for _, t := range route.Targets {
		u, err := url.Parse(t.URL)
		if err != nil {
			return nil, err
		}
		urls = append(urls, u)
		weights = append(weights, t.Weight)
		totalWeight += t.Weight
	}

	return &Balancer{
		targets:     urls,
		weights:     weights,
		totalWeight: totalWeight,
		strategy:    route.Strategy,
		rnd:         rnd,
	}, nil
}

// 修改Next方法中的随机选择逻辑
func (b *Balancer) Next() *url.URL {
	switch b.strategy {
	case "random":
		return b.targets[b.rnd.Intn(len(b.targets))]
	case "weighted-round-robin":
		for {
			b.weightIndex = (b.weightIndex + 1) % len(b.targets)
			if b.weightIndex == 0 {
				b.weightAccum -= b.totalWeight
				if b.weightAccum < 0 {
					b.weightAccum = 0
				}
			}
			if b.weights[b.weightIndex] > b.weightAccum {
				b.weightAccum += 1
				return b.targets[b.weightIndex]
			}
		}
	default: // 普通round-robin
		idx := atomic.AddUint64(&b.counter, 1) % uint64(len(b.targets))
		return b.targets[idx]
	}
}

func main() {
	rnd := rand.New(rand.NewSource(time.Now().UnixNano()))

	// 读取配置文件
	configFile, err := os.ReadFile("config.json")
	if err != nil {
		panic("无法读取配置文件: " + err.Error())
	}

	var config Config
	if err := json.Unmarshal(configFile, &config); err != nil {
		panic("解析配置文件失败: " + err.Error())
	}

	r := gin.Default()

	// 添加请求日志中间件
	r.Use(func(c *gin.Context) {
		start := time.Now()
		c.Next()
		log.Printf("[%s] %s %s %d %s",
			c.Request.Method,
			c.Request.URL.Path,
			c.ClientIP(),
			c.Writer.Status(),
			time.Since(start),
		)
	})

	// 为每个路由配置创建反向代理
	for _, route := range config.Routes {
		balancer, err := NewBalancer(route, rnd) // 传入随机数生成器
		if err != nil {
			panic("创建负载均衡器失败: " + err.Error())
		}

		r.Any(route.Path+"/*any", func(c *gin.Context) {
			target := balancer.Next()
			proxy := httputil.NewSingleHostReverseProxy(target)

			// 修改请求头
			originalDirector := proxy.Director
			proxy.Director = func(req *http.Request) {
				originalDirector(req)
				req.Header.Set("X-Forwarded-For", c.ClientIP())
				req.Header.Set("X-Proxy", "gin-reverse-proxy")
				req.Header.Set("X-Request-Id", time.Now().Format("20060102150405"))
			}

			// 转发请求
			proxy.ServeHTTP(c.Writer, c.Request)
		})
	}

	// 健康检查端点
	r.GET("/health", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"status": "ok"})
	})

	// 启动服务
	r.Run(":8080")
}

配置文件

{
    "routes": [
        {
            "path": "/api/users",
            "targets": [
                {
                    "url": "http://user-service-1:10031",
                    "weight": 3
                },
                {
                    "url": "http://user-service-2:10031",
                    "weight": 2
                },
                {
                    "url": "http://user-service-3:10031",
                    "weight": 1
                }
            ],
            "strategy": "weighted-round-robin"
        },
        {
            "path": "/api/products",
            "targets": [
                {
                    "url": "http://product-service-1:21010",
                    "weight": 3
                },
                {
                    "url": "http://product-service-2:21010",
                    "weight": 2
                }
            ],
            "strategy": "random"
        }
    ]
}

后记

GO语言起手是真的快，唯一不安全的地方在于GO完全由阿美莉卡控制，如果未来脱钩越来越狠，那GO语言作为生产要素的可持续获得性有点让人怀疑