add recorder
This commit is contained in:
parent
a117222cde
commit
03988fee0b
@ -10,6 +10,7 @@ import (
|
||||
"github.com/go-gost/core/connector"
|
||||
"github.com/go-gost/core/hosts"
|
||||
"github.com/go-gost/core/logger"
|
||||
"github.com/go-gost/core/recorder"
|
||||
"github.com/go-gost/core/resolver"
|
||||
)
|
||||
|
||||
@ -25,6 +26,7 @@ type Router struct {
|
||||
chain Chainer
|
||||
resolver resolver.Resolver
|
||||
hosts hosts.HostMapper
|
||||
recorders []recorder.RecorderObject
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
@ -70,16 +72,29 @@ func (r *Router) Hosts() hosts.HostMapper {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Router) WithRecorder(recorders ...recorder.RecorderObject) *Router {
|
||||
r.recorders = recorders
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Router) WithLogger(logger logger.Logger) *Router {
|
||||
r.logger = logger
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Router) Dial(ctx context.Context, network, address string) (conn net.Conn, err error) {
|
||||
host := address
|
||||
if h, _, _ := net.SplitHostPort(address); h != "" {
|
||||
host = h
|
||||
}
|
||||
r.record(ctx, recorder.RecorderServiceRouterDialAddress, []byte(host))
|
||||
|
||||
conn, err = r.dial(ctx, network, address)
|
||||
if err != nil {
|
||||
r.record(ctx, recorder.RecorderServiceRouterDialAddressError, []byte(host))
|
||||
return
|
||||
}
|
||||
|
||||
if network == "udp" || network == "udp4" || network == "udp6" {
|
||||
if _, ok := conn.(net.PacketConn); !ok {
|
||||
return &packetConn{conn}, nil
|
||||
@ -88,6 +103,23 @@ func (r *Router) Dial(ctx context.Context, network, address string) (conn net.Co
|
||||
return
|
||||
}
|
||||
|
||||
func (r *Router) record(ctx context.Context, name string, data []byte) error {
|
||||
if len(data) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, rec := range r.recorders {
|
||||
if rec.Record == name {
|
||||
err := rec.Recorder.Record(ctx, data)
|
||||
if err != nil {
|
||||
r.logger.Errorf("record %s: %v", name, err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Router) dial(ctx context.Context, network, address string) (conn net.Conn, err error) {
|
||||
count := r.retries + 1
|
||||
if count <= 0 {
|
||||
|
17
recorder/recorder.go
Normal file
17
recorder/recorder.go
Normal file
@ -0,0 +1,17 @@
|
||||
package recorder
|
||||
|
||||
import "context"
|
||||
|
||||
type Recorder interface {
|
||||
Record(ctx context.Context, b []byte) error
|
||||
}
|
||||
|
||||
type RecorderObject struct {
|
||||
Recorder Recorder
|
||||
Record string
|
||||
}
|
||||
|
||||
const (
|
||||
RecorderServiceRouterDialAddress = "recorder.service.router.dial.address"
|
||||
RecorderServiceRouterDialAddressError = "recorder.service.router.dial.address.error"
|
||||
)
|
@ -98,6 +98,7 @@ func (s *service) Serve() error {
|
||||
if s.options.admission != nil &&
|
||||
!s.options.admission.Admit(conn.RemoteAddr().String()) {
|
||||
conn.Close()
|
||||
s.options.logger.Debugf("admission: %s is denied", conn.RemoteAddr())
|
||||
continue
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user