add redirect for handler recorder object

This commit is contained in:
ginuerzh
2024-10-31 19:14:21 +08:00
parent d8126e8e30
commit 88ba01b3dd
5 changed files with 127 additions and 199 deletions
+19 -41
View File
@@ -11,6 +11,7 @@ import (
"fmt"
"hash/crc32"
"io"
"math"
"net"
"net/http"
"net/http/httputil"
@@ -44,6 +45,7 @@ import (
xrecorder "github.com/go-gost/x/recorder"
"github.com/go-gost/x/registry"
"golang.org/x/net/http/httpguts"
"golang.org/x/time/rate"
)
func init() {
@@ -679,44 +681,33 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read
if sampleRate == 0 {
sampleRate = sniffing.DefaultSampleRate
}
d := time.Duration(1 / sampleRate * 1e9)
if sampleRate < 0 {
sampleRate = math.MaxFloat64
}
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(cc, rw, buf, "client", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
@@ -725,37 +716,24 @@ func (h *httpHandler) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(rw, cc, buf, "server", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
+44 -51
View File
@@ -7,6 +7,7 @@ import (
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"net/http/httputil"
@@ -37,6 +38,12 @@ import (
stats_wrapper "github.com/go-gost/x/observer/stats/wrapper"
xrecorder "github.com/go-gost/x/recorder"
"golang.org/x/net/http/httpguts"
"golang.org/x/time/rate"
)
const (
httpHeaderSID = "Gost-Sid"
httpHeaderForwardedNode = "Gost-Forwarded-Node"
)
type entrypoint struct {
@@ -108,17 +115,17 @@ func (ep *entrypoint) Handle(ctx context.Context, conn net.Conn) (err error) {
return ep.handleHTTP(ctx, conn, ro, log)
}
func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.Conn, err error) {
func (ep *entrypoint) dial(ctx context.Context, network, addr string) (conn net.Conn, err error) {
var tunnelID relay.TunnelID
if h.ingress != nil {
if rule := h.ingress.GetRule(ctx, addr); rule != nil {
if ep.ingress != nil {
if rule := ep.ingress.GetRule(ctx, addr); rule != nil {
tunnelID = parseTunnelID(rule.Endpoint)
}
}
log := ctxvalue.LoggerFromContext(ctx)
if log == nil {
log = h.log
log = ep.log
}
log.Debugf("dial: new connection to host %s", addr)
@@ -126,9 +133,8 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C
return nil, fmt.Errorf("%w %s", ErrTunnelRoute, addr)
}
if ro := ctx_internal.RecorderObjectFromContext(ctx); ro != nil {
ro := ctx_internal.RecorderObjectFromContext(ctx)
ro.ClientID = tunnelID.String()
}
if tunnelID.IsPrivate() {
return nil, fmt.Errorf("%w: tunnel %s is private for host %s", ErrPrivateTunnel, tunnelID, addr)
@@ -139,9 +145,9 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C
})
d := &Dialer{
node: h.node,
pool: h.pool,
sd: h.sd,
node: ep.node,
pool: ep.pool,
sd: ep.sd,
retry: 3,
timeout: 15 * time.Second,
log: log,
@@ -152,7 +158,9 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C
}
log.Debugf("dial: connected to host %s, tunnel: %s, connector: %s", addr, tunnelID, cid)
if node == h.node {
if node == ep.node {
ro.Redirect = ""
clientAddr := ctxvalue.ClientAddrFromContext(ctx)
var features []relay.Feature
af := &relay.AddrFeature{}
@@ -171,6 +179,8 @@ func (h *entrypoint) dial(ctx context.Context, network, addr string) (conn net.C
conn.Close()
return nil, err
}
} else {
ro.Redirect = node
}
return conn, nil
@@ -237,6 +247,13 @@ func (ep *entrypoint) httpRoundTrip(ctx context.Context, rw io.ReadWriter, req *
*ro2 = *ro
ro = ro2
if sid := req.Header.Get(httpHeaderSID); sid != "" {
ro.SID = sid
} else {
req.Header.Set(httpHeaderSID, ro.SID)
}
req.Header.Set(httpHeaderForwardedNode, ro.Node)
host := req.Host
if _, port, _ := net.SplitHostPort(host); port == "" {
host = net.JoinHostPort(strings.Trim(host, "[]"), "80")
@@ -419,44 +436,33 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read
if sampleRate == 0 {
sampleRate = sniffing.DefaultSampleRate
}
d := time.Duration(1 / sampleRate * 1e9)
if sampleRate < 0 {
sampleRate = math.MaxFloat64
}
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro)
select {
case <-ticker.C:
if err == nil {
if err := ep.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, ep.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
@@ -465,37 +471,24 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro)
select {
case <-ticker.C:
if err == nil {
if err := ep.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, ep.recorder.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
+19 -41
View File
@@ -11,6 +11,7 @@ import (
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"net/http/httputil"
@@ -37,6 +38,7 @@ import (
xrecorder "github.com/go-gost/x/recorder"
"golang.org/x/net/http/httpguts"
"golang.org/x/net/http2"
"golang.org/x/time/rate"
)
var (
@@ -564,44 +566,33 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit
if sampleRate == 0 {
sampleRate = sniffing.DefaultSampleRate
}
d := time.Duration(1 / sampleRate * 1e9)
if sampleRate < 0 {
sampleRate = math.MaxFloat64
}
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(cc, rw, buf, "client", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
@@ -610,37 +601,24 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(rw, cc, buf, "server", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
+19 -41
View File
@@ -11,6 +11,7 @@ import (
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"net/http/httputil"
@@ -33,6 +34,7 @@ import (
xrecorder "github.com/go-gost/x/recorder"
"golang.org/x/net/http/httpguts"
"golang.org/x/net/http2"
"golang.org/x/time/rate"
)
const (
@@ -427,44 +429,33 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit
if sampleRate == 0 {
sampleRate = DefaultSampleRate
}
d := time.Duration(1 / sampleRate * 1e9)
if sampleRate < 0 {
sampleRate = math.MaxFloat64
}
go func() {
ro2 := &xrecorder.HandlerRecorderObject{}
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(cc, rw, buf, "client", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(cc, rw, buf, "client", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
@@ -473,37 +464,24 @@ func (h *Sniffer) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWrit
*ro2 = *ro
ro := ro2
ticker := &time.Ticker{}
if d > 0 {
ticker = time.NewTicker(d)
defer ticker.Stop()
} else {
tc := make(chan time.Time)
close(tc)
ticker.C = tc
}
limiter := rate.NewLimiter(rate.Limit(sampleRate), int(sampleRate))
buf := &bytes.Buffer{}
for {
start := time.Now()
err := h.copyWebsocketFrame(rw, cc, buf, "server", ro)
select {
case <-ticker.C:
if err == nil {
if err := h.copyWebsocketFrame(rw, cc, buf, "server", ro); err != nil {
errc <- err
return
}
if limiter.Allow() {
ro.Duration = time.Since(start)
ro.Time = time.Now()
if err := ro.Record(ctx, h.Recorder); err != nil {
log.Errorf("record: %v", err)
}
}
default:
}
if err != nil {
errc <- err
return
}
}
}()
+1
View File
@@ -89,6 +89,7 @@ type HandlerRecorderObject struct {
Route string `json:"route,omitempty"`
InputBytes uint64 `json:"inputBytes"`
OutputBytes uint64 `json:"outputBytes"`
Redirect string `json:"redirect,omitempty"`
Err string `json:"err,omitempty"`
SID string `json:"sid"`
Duration time.Duration `json:"duration"`