diff --git a/admission/wrapper/conn.go b/admission/wrapper/conn.go deleted file mode 100644 index 2cb527f..0000000 --- a/admission/wrapper/conn.go +++ /dev/null @@ -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 -} diff --git a/admission/wrapper/listener.go b/admission/wrapper/listener.go deleted file mode 100644 index 05d4dee..0000000 --- a/admission/wrapper/listener.go +++ /dev/null @@ -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 - } -} diff --git a/limiter/limiter.go b/limiter/limiter.go new file mode 100644 index 0000000..d1b0f83 --- /dev/null +++ b/limiter/limiter.go @@ -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 +} diff --git a/limiter/rate.go b/limiter/rate.go deleted file mode 100644 index 1edd939..0000000 --- a/limiter/rate.go +++ /dev/null @@ -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 -} diff --git a/limiter/wrapper/conn.go b/limiter/wrapper/conn.go deleted file mode 100644 index 29d044c..0000000 --- a/limiter/wrapper/conn.go +++ /dev/null @@ -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) -} diff --git a/limiter/wrapper/listener.go b/limiter/wrapper/listener.go deleted file mode 100644 index 42a2c8a..0000000 --- a/limiter/wrapper/listener.go +++ /dev/null @@ -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 -} diff --git a/metrics/wrapper/conn.go b/metrics/wrapper/conn.go deleted file mode 100644 index aa7a4f4..0000000 --- a/metrics/wrapper/conn.go +++ /dev/null @@ -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 -} diff --git a/metrics/wrapper/listener.go b/metrics/wrapper/listener.go deleted file mode 100644 index f663cb8..0000000 --- a/metrics/wrapper/listener.go +++ /dev/null @@ -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 -}