Redis 调用

version: '3.9'
services:
  envoy:
    image: higress-registry.cn-hangzhou.cr.aliyuncs.com/higress/gateway:v1.4.1
    entrypoint: /usr/local/bin/envoy
    # 注意这里对wasm开启了debug级别日志，正式部署时则默认info级别
    command: -c /etc/envoy/envoy.yaml --log-level info --log-path /etc/envoy/envoy.log --component-log-level wasm:debug
    depends_on:
      - echo-server
    networks:
      - wasmtest
    ports:
      - "10000:10000"
      - "9901:9901"
    volumes:
      - ./envoy.yaml:/etc/envoy/envoy.yaml
      - ./build/plugin.wasm:/etc/envoy/plugin.wasm
      - ./envoy.log:/etc/envoy/envoy.log
  echo-server:
    image: higress-registry.cn-hangzhou.cr.aliyuncs.com/higress/echo-server:1.3.0
    networks:
      - wasmtest
    ports:
      - "3000:3000"
  redis:
    image: registry.cn-hangzhou.aliyuncs.com/2456868764/redis:latest
    environment:
      - ALLOW_EMPTY_PASSWORD=yes
    networks:
      wasmtest:
        ipv4_address: 172.20.0.100
    ports:
      - "6379:6379"
networks:
  wasmtest:
    ipam:
      config:
        - subnet: 172.20.0.0/24

admin:
  address:
    socket_address:
      protocol: TCP
      address: 0.0.0.0
      port_value: 9901
static_resources:
  listeners:
    - name: listener_0
      address:
        socket_address:
          protocol: TCP
          address: 0.0.0.0
          port_value: 10000
      filter_chains:
        - filters:
            - name: envoy.filters.network.http_connection_manager
              typed_config:
                "@type": type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager
                scheme_header_transformation:
                  scheme_to_overwrite: https
                stat_prefix: ingress_http
                route_config:
                  name: local_route
                  virtual_hosts:
                    - name: local_service
                      domains: ["*"]
                      routes:
                        - match:
                            prefix: "/"
                          route:
                            cluster: echo-server
                http_filters:
                  - name: wasmdemo
                    typed_config:
                      "@type": type.googleapis.com/udpa.type.v1.TypedStruct
                      type_url: type.googleapis.com/envoy.extensions.filters.http.wasm.v3.Wasm
                      value:
                        config:
                          name: wasmdemo
                          vm_config:
                            runtime: envoy.wasm.runtime.v8
                            code:
                              local:
                                filename: /etc/envoy/plugin.wasm
                          configuration:
                            "@type": "type.googleapis.com/google.protobuf.StringValue"
                            value: |-
                              {
                                "keys": [
                                  "authorization"
                                ],
                                "in_header": true,
                                "limits": [
                                  {
                                    "name": "credential1",
                                    "consumer": "Bearer credential1",
                                    "rate": 2,
                                    "capacity": 4
                                  },
                                  {
                                    "name": "all",
                                    "consumer": "*",
                                    "rate": 1,
                                    "capacity": 2
                                  }
                                ],
                                "rejected_code": 429,
                                "rejected_msg": "Too Many Requests",
                                "show_limit_quota_header": true,
                                "redis":{
                                  "service_name": "redis.static",
                                  "service_port": 6379,
                                  "timeout": 2000
                                }
                              }
                  - name: envoy.filters.http.router
                    typed_config:
                      "@type": type.googleapis.com/envoy.extensions.filters.http.router.v3.Router
  clusters:
    - name: echo-server
      connect_timeout: 30s
      type: LOGICAL_DNS
      # Comment out the following line to test on v6 networks
      dns_lookup_family: V4_ONLY
      lb_policy: ROUND_ROBIN
      load_assignment:
        cluster_name: echo-server
        endpoints:
          - lb_endpoints:
              - endpoint:
                  address:
                    socket_address:
                      address: echo-server
                      port_value: 3000
    - name: outbound|6379||redis.static
      connect_timeout: 30s
      type: STATIC
      load_assignment:
        cluster_name: outbound|6379||redis.static
        endpoints:
          - lb_endpoints:
              - endpoint:
                  address:
                    socket_address:
                      address: 172.20.0.100
                      port_value: 6379

// LimitConfig 定义了限流插件的配置结构。
type LimitConfig struct {
  Keys                 []string            `yaml:"keys"`                    // 定义了用于提取限流信息的HTTP请求头字段名称。
  InQuery              bool                `yaml:"in_query,omitempty"`      // 标识是否从查询参数中获取限流信息。
  InHeader             bool                `yaml:"in_header,omitempty"`     // 标识是否从请求头中获取限流信息。
  Limits               []LimitItem         `yaml:"limits"`                  // 包含具体的限流规则项。
  RejectedCode         uint32              `yaml:"rejected_code"`           // 请求超过阈值被拒绝时返回的 HTTP 状态码。
  RejectedMsg          string              `yaml:"rejected_msg"`            // 请求超过阈值被拒绝时返回的响应体。
  ShowLimitQuotaHeader bool                `yaml:"show_limit_quota_header"` // 标识是否在响应头中显示限流配额信息。
  RedisInfo            RedisInfo           `yaml:"redis"`                   // 定义了与 Redis 交互所需的信息。
  RedisClient          wrapper.RedisClient // Redis客户端，用于执行Redis命令。
}

// LimitItem 定义了具体的限流项，包括限流名称、消费者标识、请求速率和容量。
type LimitItem struct {
  Name     string `yaml:"name"`     // 限流项的名称。
  Consumer string `yaml:"consumer"` // 限流项关联的消费者标识，用于匹配特定的请求。
  Rate     int    `yaml:"rate"`     // 每秒放入桶内的令牌数量。
  Capacity int    `yaml:"capacity"` // 限流桶的最大容量。
}

// RedisInfo 定义了连接Redis所需的详细信息。
type RedisInfo struct {
  ServiceName string `required:"true" yaml:"service_name" json:"service_name"` // Redis 服务名或地址。
  ServicePort int    `yaml:"service_port" json:"service_port"`                 // Redis 服务端口。
  Username    string `yaml:"username" json:"username"`                         // 连接 Redis 的用户名，如果需要。
  Password    string `yaml:"password" json:"password"`                         // 连接 Redis 的密码，如果需要。
  Timeout     int    `yaml:"timeout" json:"timeout"`                           // 连接 Redis 的超时时间（毫秒）。
}

// LimitContext 定义了限流上下文，存储了限流相关的信息。
type LimitContext struct {
  Allowed   int // 表示当前请求是否被允许。
  Count     int // 限流桶当前的计数。
  Remaining int // 限流桶剩余的容量。
  Reset     int // 限流桶重置时间（秒）。
}

// 解析配置，这里忽略插件配置解析的细节。主要显示 Redis 部分解析配置。
func parseConfig(json gjson.Result, config *LimitConfig, log wrapper.Log) error {
  // keys
  names := json.Get("keys")
  ...
  // in_query and in_header
  in_query := json.Get("in_query")
  in_header := json.Get("in_header")
  ...
  // parse limit
  limits := json.Get("limits")
  ...
  config.ShowLimitQuotaHeader = json.Get("show_limit_quota_header").Bool()

  // parse redis
  redisConfig := json.Get("redis")
  if !redisConfig.Exists() {
    return errors.New("missing redis in config")
  }
  serviceName := redisConfig.Get("service_name").String()
  if serviceName == "" {
    return errors.New("redis service name must not be empty")
  }
  servicePort := int(redisConfig.Get("service_port").Int())
  if servicePort == 0 {
    if strings.HasSuffix(serviceName, ".static") {
      // use default logic port which is 80 for static service
      servicePort = 80
    } else {
      servicePort = 6379
    }
  }
  username := redisConfig.Get("username").String()
  password := redisConfig.Get("password").String()
  timeout := int(redisConfig.Get("timeout").Int())
  if timeout == 0 {
    timeout = 1000
  }
  config.RedisInfo.ServiceName = serviceName
  config.RedisInfo.ServicePort = servicePort
  config.RedisInfo.Username = username
  config.RedisInfo.Password = password
  config.RedisInfo.Timeout = timeout
  config.RedisClient = wrapper.NewRedisClusterClient(wrapper.FQDNCluster{
    FQDN: serviceName,
    Port: int64(servicePort),
  })
  log.Debugf("parseConfig()+%+v", config)
  return config.RedisClient.Init(username, password, int64(timeout))
}

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
local unit = tonumber(ARGV[5])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2*unit)

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

local delta = math.max(0, (now-last_refreshed)/unit)
local filled_tokens = math.min(capacity, last_tokens + math.floor(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

if ttl > 0 then
  redis.call("setex", tokens_key, ttl, new_tokens)
  redis.call("setex", timestamp_key, ttl, now)
end

return { allowed_num, new_tokens, capacity, ttl}

// onHttpRequestHeaders 函数在处理 HTTP 请求头时被调用，用于执行限流逻辑。
func onHttpRequestHeaders(ctx wrapper.HttpContext, config LimitConfig, log wrapper.Log) types.Action {
  log.Debugf("onHttpRequestHeaders()")
  // 初始化 tokens 切片，用于存储从请求头或查询参数中提取的 tokens 信息
  var tokens []string

  // 如果配置指定从请求头中获取 tokens 信息
  if config.InHeader {
    // 遍历配置中定义的所有键（key），尝试从请求头中获取每个键的值
    for _, key := range config.Keys {
      // 使用 proxywasm 库的 GetHttpRequestHeader 函数获取请求头中的值
      value, err := proxywasm.GetHttpRequestHeader(key)
      // 如果没有错误且值不为空，则将值添加到 tokens 切片中
      if err == nil && value != "" {
        tokens = append(tokens, value)
      }
    }
  } else if config.InQuery {
    // 如果配置指定从查询参数中获取 tokens 信息
    // 获取 ":path" 请求头以获取请求 URL
    requestUrl, _ := proxywasm.GetHttpRequestHeader(":path")
    // 解析 URL 并获取查询参数
    url, _ := url.Parse(requestUrl)
    queryValues := url.Query()

    // 遍历配置中定义的所有键
    for _, key := range config.Keys {
      // 从查询参数中获取每个键的值
      values, ok := queryValues[key]
      // 如果查询参数存在且有值，则将值添加到 tokens 切片中
      if ok && len(values) > 0 {
        tokens = append(tokens, values...)
      }
    }
  }

  // 如果从请求中提取了多于一个的 tokens，返回错误处理
  if len(tokens) > 1 {
    return deniedMultiKeyAuthData()
  } else if len(tokens) <= 0 {
    // 如果没有提取到 tokens，返回错误处理
    return deniedNoKeyAuthData()
  }

  // 提取第一个 token 作为主要的令牌信息
  limitKey := strings.Split(tokens[0], ",")[0]
  log.Debugf("limitKey:%s", limitKey)

  // 根据提取的 limitKey 查找对应的限流项
  limitItem := findLimitItem(config, limitKey)
  log.Debugf("limitItem:%+v", limitItem)

  // 如果没有找到对应的限流项，继续处理请求
  if limitItem.Consumer == "" {
    return types.ActionContinue
  }

  // 构建 Redis 脚本需要的键，用于操作令牌桶
  tokenKey := fmt.Sprintf(ClusterRateLimitFormat, limitKey, "token")
  expireKey := fmt.Sprintf(ClusterRateLimitFormat, limitKey, "expire")

  // 获取当前时间，用于计算令牌桶的填充状态
  now := time.Now()
  // 将当前时间转换为 Unix 时间戳（秒）
  unixTimestamp := now.UnixNano()
  milliseconds := unixTimestamp / 1e6
  seconds := milliseconds / 1000

  // 构建调用 Redis 脚本所需的参数
  keys := []interface{}{tokenKey, expireKey}
  args := []interface{}{limitItem.Rate, limitItem.Capacity, seconds, 1, 1}

  // 调用 Redis 脚本执行限流逻辑
  err := config.RedisClient.Eval(BucketTokenScript, 2, keys, args, func(response resp.Value) {
    log.Debugf("RedisClient.Eval()，keys:%+v,args:%+v", keys, args)
    // 检查脚本返回的结果是否包含 4 个元素
    resultArray := response.Array()
    if len(resultArray) != 4 {
      log.Errorf("redis response parse error, response: %v", response)
            proxywasm.ResumeHttpRequest()
      return
    }
    // 根据脚本返回的结果创建 LimitContext 对象
    context := LimitContext{
      Allowed:   resultArray[0].Integer(),
      Remaining: resultArray[1].Integer(),
      Count:     resultArray[2].Integer(),
      Reset:     resultArray[3].Integer(),
    }
    log.Debugf("context:%+v", context)
    // 如果请求未被允许（Allowed <= 0），触发限流逻辑
    if context.Allowed <= 0 {
      log.Debugf("request rejected")
      rejected(config, context)
      return
    } else {
      // 将限流上下文存储在 HttpContext 中，供后续处理使用
      ctx.SetContext(LimitContextKey, context)
    }
    // 恢复 HTTP 请求处理
    proxywasm.ResumeHttpRequest()
  })
  // 如果调用 Redis 脚本时出现错误，记录错误并继续处理请求
  if err != nil {
    log.Errorf("redis call failed: %v", err)
    return types.ActionContinue
  }
  // 暂停处理当前请求头，等待 Redis 脚本调用完成
  return types.HeaderStopAllIterationAndWatermark
}

// findLimitItem 函数用于在给定的配置中查找与特定消费者匹配的限流项。如果没有找到具体的匹配项，它将返回一个默认的 LimitItem 结构。
func findLimitItem(config LimitConfig, key string) LimitItem {
  // 遍历配置中的所有限流项
  for _, limitItem := range config.Limits {
    // 检查当前限流项的消费者字段是否与提供的 key 匹配，且消费者不是通配符"*"
    if limitItem.Consumer == key && limitItem.Consumer != "*" {
      // 如果找到匹配的限流项，返回这个限流项
      return limitItem
    }
  }
  // 再次遍历配置中的所有限流项，这次是为了查找通配符"*"的消费者
  // 通配符"*"表示这个限流项适用于所有消费者
  for _, limitItem := range config.Limits {
    if limitItem.Consumer == "*" {
      // 如果找到通配符限流项，返回这个限流项
      return limitItem
    }
  }
  // 则返回一个空的 LimitItem 结构，表示没有找到任何适用的限流规则
  return LimitItem{}
}