go-gost/core
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add rate limiter

Browse Source
This commit is contained in:
ginuerzh 2022-09-05 22:46:53 +08:00
parent b25f90c55e
commit 5c46613716
8 changed files with 15 additions and 771 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

223
admission/wrapper/conn.go
View File

@ -1,223 +0,0 @@
package wrapper
import (
"errors"
"io"
"net"
"syscall"
"github.com/go-gost/core/admission"
)
var (
errUnsupport = errors.New("unsupported operation")
)
type packetConn struct {
net.PacketConn
admission admission.Admission
}
func WrapPacketConn(admission admission.Admission, pc net.PacketConn) net.PacketConn {
if admission == nil {
return pc
}
return &packetConn{
PacketConn: pc,
admission: admission,
}
}
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
for {
n, addr, err = c.PacketConn.ReadFrom(p)
if err != nil {
return
}
if c.admission != nil &&
!c.admission.Admit(addr.String()) {
continue
}
return
}
}
type udpConn struct {
net.PacketConn
admission admission.Admission
}
func WrapUDPConn(admission admission.Admission, pc net.PacketConn) UDPConn {
return &udpConn{
PacketConn: pc,
admission: admission,
}
}
func (c *udpConn) RemoteAddr() net.Addr {
if nc, ok := c.PacketConn.(remoteAddr); ok {
return nc.RemoteAddr()
}
return nil
}
func (c *udpConn) SetReadBuffer(n int) error {
if nc, ok := c.PacketConn.(setBuffer); ok {
return nc.SetReadBuffer(n)
}
return errUnsupport
}
func (c *udpConn) SetWriteBuffer(n int) error {
if nc, ok := c.PacketConn.(setBuffer); ok {
return nc.SetWriteBuffer(n)
}
return errUnsupport
}
func (c *udpConn) Read(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Reader); ok {
n, err = nc.Read(b)
return
}
err = errUnsupport
return
}
func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
for {
n, addr, err = c.PacketConn.ReadFrom(p)
if err != nil {
return
}
if c.admission != nil &&
!c.admission.Admit(addr.String()) {
continue
}
return
}
}
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(readUDP); ok {
for {
n, addr, err = nc.ReadFromUDP(b)
if err != nil {
return
}
if c.admission != nil &&
!c.admission.Admit(addr.String()) {
continue
}
return
}
}
err = errUnsupport
return
}
func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(readUDP); ok {
for {
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
if err != nil {
return
}
if c.admission != nil &&
!c.admission.Admit(addr.String()) {
continue
}
return
}
}
err = errUnsupport
return
}
func (c *udpConn) Write(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Writer); ok {
n, err = nc.Write(b)
return
}
err = errUnsupport
return
}
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
n, err = c.PacketConn.WriteTo(p, addr)
return
}
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
if nc, ok := c.PacketConn.(writeUDP); ok {
n, err = nc.WriteToUDP(b, addr)
return
}
err = errUnsupport
return
}
func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
if nc, ok := c.PacketConn.(writeUDP); ok {
n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
return
}
err = errUnsupport
return
}
func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
if nc, ok := c.PacketConn.(syscallConn); ok {
return nc.SyscallConn()
}
err = errUnsupport
return
}
func (c *udpConn) SetDSCP(n int) error {
if nc, ok := c.PacketConn.(setDSCP); ok {
return nc.SetDSCP(n)
}
return nil
}
type UDPConn interface {
net.PacketConn
io.Reader
io.Writer
readUDP
writeUDP
setBuffer
syscallConn
remoteAddr
}
type setBuffer interface {
SetReadBuffer(bytes int) error
SetWriteBuffer(bytes int) error
}
type readUDP interface {
ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
}
type writeUDP interface {
WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
}
type syscallConn interface {
SyscallConn() (syscall.RawConn, error)
}
type remoteAddr interface {
RemoteAddr() net.Addr
}
// tcpraw.TCPConn
type setDSCP interface {
SetDSCP(int) error
}

37
admission/wrapper/listener.go
View File

@ -1,37 +0,0 @@
package wrapper
import (
"net"
"github.com/go-gost/core/admission"
)
type listener struct {
net.Listener
admission admission.Admission
}
func WrapListener(admission admission.Admission, ln net.Listener) net.Listener {
if admission == nil {
return ln
}
return &listener{
Listener: ln,
admission: admission,
}
}
func (ln *listener) Accept() (net.Conn, error) {
for {
c, err := ln.Listener.Accept()
if err != nil {
return nil, err
}
if ln.admission != nil &&
!ln.admission.Admit(c.RemoteAddr().String()) {
c.Close()
continue
}
return c, err
}
}

15
limiter/limiter.go Normal file
View File

@ -0,0 +1,15 @@
package limiter
import "context"
type Limiter interface {
// Wait blocks with the requested n and returns the result value,
// the returned value is less or equal to n.
Wait(ctx context.Context, n int) int
Limit() int
}
type RateLimiter interface {
In(key string) Limiter
Out(key string) Limiter
}

12
limiter/rate.go
View File

@ -1,12 +0,0 @@
package limiter
type Limiter interface {
// Limit checks the requested size b and returns the limit size,
// the returned value is less or equal to b.
Limit(b int) int
}
type RateLimiter interface {
Input() Limiter
Output() Limiter
}

130
limiter/wrapper/conn.go
View File

@ -1,130 +0,0 @@
package wrapper
import (
"bytes"
"errors"
"net"
"syscall"
"github.com/go-gost/core/limiter"
)
var (
errUnsupport = errors.New("unsupported operation")
)
// serverConn is a server side Conn with metrics supported.
type serverConn struct {
net.Conn
rlimiter limiter.RateLimiter
rbuf bytes.Buffer
}
func WrapConn(rlimiter limiter.RateLimiter, c net.Conn) net.Conn {
if rlimiter == nil {
return c
}
return &serverConn{
Conn: c,
rlimiter: rlimiter,
}
}
func (c *serverConn) Read(b []byte) (n int, err error) {
if c.rlimiter == nil || c.rlimiter.Input() == nil {
return c.Conn.Read(b)
}
burst := len(b)
if c.rbuf.Len() > 0 {
if c.rbuf.Len() < burst {
burst = c.rbuf.Len()
}
return c.rbuf.Read(b[:c.rlimiter.Input().Limit(burst)])
}
nn, err := c.Conn.Read(b)
if err != nil {
return nn, err
}
n = c.rlimiter.Input().Limit(nn)
if n < nn {
if _, err = c.rbuf.Write(b[n:nn]); err != nil {
return 0, err
}
}
return
}
func (c *serverConn) Write(b []byte) (n int, err error) {
if c.rlimiter == nil || c.rlimiter.Output() == nil {
return c.Conn.Write(b)
}
nn := 0
for len(b) > 0 {
nn, err = c.Conn.Write(b[:c.rlimiter.Output().Limit(len(b))])
n += nn
if err != nil {
return
}
b = b[nn:]
}
return
}
func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
if sc, ok := c.Conn.(syscall.Conn); ok {
rc, err = sc.SyscallConn()
return
}
err = errUnsupport
return
}
type packetConn struct {
net.PacketConn
rlimiter limiter.RateLimiter
}
func WrapPacketConn(rlimiter limiter.RateLimiter, pc net.PacketConn) net.PacketConn {
if rlimiter == nil {
return pc
}
return &packetConn{
PacketConn: pc,
rlimiter: rlimiter,
}
}
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
for {
n, addr, err = c.PacketConn.ReadFrom(p)
if err != nil {
return
}
if c.rlimiter == nil || c.rlimiter.Input() == nil {
return
}
if c.rlimiter.Input().Limit(n) < n {
continue
}
return
}
}
func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
if c.rlimiter != nil &&
c.rlimiter.Output() != nil &&
c.rlimiter.Output().Limit(len(p)) < len(p) {
n = len(p)
return
}
return c.PacketConn.WriteTo(p, addr)
}

32
limiter/wrapper/listener.go
View File

@ -1,32 +0,0 @@
package wrapper
import (
"net"
"github.com/go-gost/core/limiter"
)
type listener struct {
net.Listener
rlimiter limiter.RateLimiter
}
func WrapListener(rlimiter limiter.RateLimiter, ln net.Listener) net.Listener {
if rlimiter == nil {
return ln
}
return &listener{
rlimiter: rlimiter,
Listener: ln,
}
}
func (ln *listener) Accept() (net.Conn, error) {
c, err := ln.Listener.Accept()
if err != nil {
return nil, err
}
return WrapConn(ln.rlimiter, c), nil
}

305
metrics/wrapper/conn.go
View File

@ -1,305 +0,0 @@
package wrapper
import (
"errors"
"io"
"net"
"syscall"
"github.com/go-gost/core/metrics"
)
var (
errUnsupport = errors.New("unsupported operation")
)
// serverConn is a server side Conn with metrics supported.
type serverConn struct {
net.Conn
service string
}
func WrapConn(service string, c net.Conn) net.Conn {
if !metrics.IsEnabled() {
return c
}
return &serverConn{
service: service,
Conn: c,
}
}
func (c *serverConn) Read(b []byte) (n int, err error) {
n, err = c.Conn.Read(b)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
func (c *serverConn) Write(b []byte) (n int, err error) {
n, err = c.Conn.Write(b)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
func (c *serverConn) SyscallConn() (rc syscall.RawConn, err error) {
if sc, ok := c.Conn.(syscall.Conn); ok {
rc, err = sc.SyscallConn()
return
}
err = errUnsupport
return
}
type packetConn struct {
net.PacketConn
service string
}
func WrapPacketConn(service string, pc net.PacketConn) net.PacketConn {
if !metrics.IsEnabled() {
return pc
}
return &packetConn{
PacketConn: pc,
service: service,
}
}
func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
n, addr, err = c.PacketConn.ReadFrom(p)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
n, err = c.PacketConn.WriteTo(p, addr)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
type udpConn struct {
net.PacketConn
service string
}
func WrapUDPConn(service string, pc net.PacketConn) UDPConn {
return &udpConn{
PacketConn: pc,
service: service,
}
}
func (c *udpConn) RemoteAddr() net.Addr {
if nc, ok := c.PacketConn.(remoteAddr); ok {
return nc.RemoteAddr()
}
return nil
}
func (c *udpConn) SetReadBuffer(n int) error {
if nc, ok := c.PacketConn.(setBuffer); ok {
return nc.SetReadBuffer(n)
}
return errUnsupport
}
func (c *udpConn) SetWriteBuffer(n int) error {
if nc, ok := c.PacketConn.(setBuffer); ok {
return nc.SetWriteBuffer(n)
}
return errUnsupport
}
func (c *udpConn) Read(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Reader); ok {
n, err = nc.Read(b)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
n, addr, err = c.PacketConn.ReadFrom(p)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
func (c *udpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(readUDP); ok {
n, addr, err = nc.ReadFromUDP(b)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error) {
if nc, ok := c.PacketConn.(readUDP); ok {
n, oobn, flags, addr, err = nc.ReadMsgUDP(b, oob)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferInputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) Write(b []byte) (n int, err error) {
if nc, ok := c.PacketConn.(io.Writer); ok {
n, err = nc.Write(b)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
n, err = c.PacketConn.WriteTo(p, addr)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
func (c *udpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
if nc, ok := c.PacketConn.(writeUDP); ok {
n, err = nc.WriteToUDP(b, addr)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error) {
if nc, ok := c.PacketConn.(writeUDP); ok {
n, oobn, err = nc.WriteMsgUDP(b, oob, addr)
if counter := metrics.GetCounter(
metrics.MetricServiceTransferOutputBytesCounter,
metrics.Labels{
"service": c.service,
}); counter != nil {
counter.Add(float64(n))
}
return
}
err = errUnsupport
return
}
func (c *udpConn) SyscallConn() (rc syscall.RawConn, err error) {
if nc, ok := c.PacketConn.(syscall.Conn); ok {
return nc.SyscallConn()
}
err = errUnsupport
return
}
func (c *udpConn) SetDSCP(n int) error {
if nc, ok := c.PacketConn.(setDSCP); ok {
return nc.SetDSCP(n)
}
return nil
}
type UDPConn interface {
net.PacketConn
io.Reader
io.Writer
readUDP
writeUDP
setBuffer
remoteAddr
syscall.Conn
}
type setBuffer interface {
SetReadBuffer(bytes int) error
SetWriteBuffer(bytes int) error
}
type readUDP interface {
ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
}
type writeUDP interface {
WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
}
type remoteAddr interface {
RemoteAddr() net.Addr
}
// tcpraw.TCPConn
type setDSCP interface {
SetDSCP(int) error
}

32
metrics/wrapper/listener.go
View File

@ -1,32 +0,0 @@
package wrapper
import (
"net"
"github.com/go-gost/core/metrics"
)
type listener struct {
service string
net.Listener
}
func WrapListener(service string, ln net.Listener) net.Listener {
if !metrics.IsEnabled() {
return ln
}
return &listener{
service: service,
Listener: ln,
}
}
func (ln *listener) Accept() (net.Conn, error) {
c, err := ln.Listener.Accept()
if err != nil {
return nil, err
}
return WrapConn(ln.service, c), nil
}