x/limiter/traffic/generator.go
2022-09-14 00:15:19 +08:00

62 lines
1.0 KiB
Go

package traffic
import (
limiter "github.com/go-gost/core/limiter/traffic"
)
type TrafficLimitGenerator interface {
In() limiter.Limiter
Out() limiter.Limiter
}
type trafficLimitGenerator struct {
in int
out int
}
func NewTrafficLimitGenerator(in, out int) TrafficLimitGenerator {
return &trafficLimitGenerator{
in: in,
out: out,
}
}
func (p *trafficLimitGenerator) In() limiter.Limiter {
if p == nil || p.in <= 0 {
return nil
}
return NewLimiter(p.in)
}
func (p *trafficLimitGenerator) Out() limiter.Limiter {
if p == nil || p.out <= 0 {
return nil
}
return NewLimiter(p.out)
}
type trafficLimitSingleGenerator struct {
in limiter.Limiter
out limiter.Limiter
}
func NewTrafficLimitSingleGenerator(in, out int) TrafficLimitGenerator {
p := &trafficLimitSingleGenerator{}
if in > 0 {
p.in = NewLimiter(in)
}
if out > 0 {
p.out = NewLimiter(out)
}
return p
}
func (p *trafficLimitSingleGenerator) In() limiter.Limiter {
return p.in
}
func (p *trafficLimitSingleGenerator) Out() limiter.Limiter {
return p.out
}