x/handler/relay/entrypoint.go
2023-11-14 22:35:46 +08:00

125 lines
2.9 KiB
Go

package relay
import (
"context"
"net"
"time"
"github.com/go-gost/core/handler"
"github.com/go-gost/core/listener"
md "github.com/go-gost/core/metadata"
"github.com/go-gost/relay"
admission "github.com/go-gost/x/admission/wrapper"
xnet "github.com/go-gost/x/internal/net"
"github.com/go-gost/x/internal/net/proxyproto"
"github.com/go-gost/x/internal/util/mux"
climiter "github.com/go-gost/x/limiter/conn/wrapper"
limiter "github.com/go-gost/x/limiter/traffic/wrapper"
metrics "github.com/go-gost/x/metrics/wrapper"
)
type tcpListener struct {
ln net.Listener
options listener.Options
}
func newTCPListener(ln net.Listener, opts ...listener.Option) listener.Listener {
options := listener.Options{}
for _, opt := range opts {
opt(&options)
}
return &tcpListener{
ln: ln,
options: options,
}
}
func (l *tcpListener) Init(md md.Metadata) (err error) {
// l.logger.Debugf("pp: %d", l.options.ProxyProtocol)
ln := l.ln
ln = proxyproto.WrapListener(l.options.ProxyProtocol, ln, 10*time.Second)
ln = metrics.WrapListener(l.options.Service, ln)
ln = admission.WrapListener(l.options.Admission, ln)
ln = limiter.WrapListener(l.options.TrafficLimiter, ln)
ln = climiter.WrapListener(l.options.ConnLimiter, ln)
l.ln = ln
return
}
func (l *tcpListener) Accept() (conn net.Conn, err error) {
return l.ln.Accept()
}
func (l *tcpListener) Addr() net.Addr {
return l.ln.Addr()
}
func (l *tcpListener) Close() error {
return l.ln.Close()
}
type tcpHandler struct {
session *mux.Session
options handler.Options
}
func newTCPHandler(session *mux.Session, opts ...handler.Option) handler.Handler {
options := handler.Options{}
for _, opt := range opts {
opt(&options)
}
return &tcpHandler{
session: session,
options: options,
}
}
func (h *tcpHandler) Init(md md.Metadata) (err error) {
return
}
func (h *tcpHandler) Handle(ctx context.Context, conn net.Conn, opts ...handler.HandleOption) error {
defer conn.Close()
start := time.Now()
log := h.options.Logger.WithFields(map[string]any{
"remote": conn.RemoteAddr().String(),
"local": conn.LocalAddr().String(),
})
log.Infof("%s <> %s", conn.RemoteAddr(), conn.LocalAddr())
defer func() {
log.WithFields(map[string]any{
"duration": time.Since(start),
}).Infof("%s >< %s", conn.RemoteAddr(), conn.LocalAddr())
}()
cc, err := h.session.GetConn()
if err != nil {
log.Error(err)
return err
}
defer cc.Close()
af := &relay.AddrFeature{}
af.ParseFrom(conn.RemoteAddr().String())
resp := relay.Response{
Version: relay.Version1,
Status: relay.StatusOK,
Features: []relay.Feature{af},
}
if _, err := resp.WriteTo(cc); err != nil {
log.Error(err)
return err
}
t := time.Now()
log.Debugf("%s <-> %s", conn.RemoteAddr(), cc.RemoteAddr())
xnet.Transport(conn, cc)
log.WithFields(map[string]any{"duration": time.Since(t)}).
Debugf("%s >-< %s", conn.RemoteAddr(), cc.RemoteAddr())
return nil
}