ginuerzh
49ceda6cc2
Add UnwrapConn() to the admission, proxyproto, conn limiter, metrics, and observer/stats TCP connection wrappers, complementing the traffic limiter's limitConn added in the previous commit. This makes the unwrapConn() helper in handler/sshd able to peel through any combination of these wrappers uniformly, keeping the pattern complete and future-proof for handlers that assert on concrete connection types.
79 lines
1.7 KiB
Go
79 lines
1.7 KiB
Go
// Package wrapper provides net.Conn and net.Listener wrappers that enforce
|
|
// connection limits.
|
|
package wrapper
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"syscall"
|
|
|
|
limiter "github.com/go-gost/core/limiter/conn"
|
|
"github.com/go-gost/x/ctx"
|
|
xio "github.com/go-gost/x/internal/io"
|
|
)
|
|
|
|
var (
|
|
errUnsupport = errors.New("unsupported operation")
|
|
)
|
|
|
|
// serverConn is a server side Conn with metrics supported.
|
|
type serverConn struct {
|
|
net.Conn
|
|
limiter limiter.Limiter
|
|
}
|
|
|
|
// WrapConn wraps a net.Conn with a connection limiter. On Close, the
|
|
// limiter's Allow(-1) is called to decrement the current count. If limiter
|
|
// is nil, the original connection is returned unchanged.
|
|
func WrapConn(limiter limiter.Limiter, c net.Conn) net.Conn {
|
|
if limiter == nil {
|
|
return c
|
|
}
|
|
return &serverConn{
|
|
Conn: c,
|
|
limiter: limiter,
|
|
}
|
|
}
|
|
|
|
// UnwrapConn returns the underlying connection, allowing type assertions
|
|
// through wrapper layers.
|
|
func (c *serverConn) UnwrapConn() net.Conn {
|
|
return c.Conn
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
func (c *serverConn) Close() error {
|
|
c.limiter.Allow(-1)
|
|
return c.Conn.Close()
|
|
}
|
|
|
|
func (c *serverConn) CloseRead() error {
|
|
if sc, ok := c.Conn.(xio.CloseRead); ok {
|
|
return sc.CloseRead()
|
|
}
|
|
return xio.ErrUnsupported
|
|
}
|
|
|
|
func (c *serverConn) CloseWrite() error {
|
|
if sc, ok := c.Conn.(xio.CloseWrite); ok {
|
|
return sc.CloseWrite()
|
|
}
|
|
return xio.ErrUnsupported
|
|
}
|
|
|
|
func (c *serverConn) Context() context.Context {
|
|
if innerCtx, ok := c.Conn.(ctx.Context); ok {
|
|
return innerCtx.Context()
|
|
}
|
|
return context.Background()
|
|
}
|