41 lines
680 B
Go
41 lines
680 B
Go
package wrapper
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
|
|
"github.com/go-gost/core/admission"
|
|
"github.com/go-gost/core/logger"
|
|
)
|
|
|
|
type listener struct {
|
|
net.Listener
|
|
admission admission.Admission
|
|
log logger.Logger
|
|
}
|
|
|
|
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(context.Background(), c.RemoteAddr().String()) {
|
|
c.Close()
|
|
continue
|
|
}
|
|
return c, err
|
|
}
|
|
}
|