go-gost/x
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

improve tunnel handler

Browse Source
This commit is contained in:
ginuerzh
2023-10-27 22:11:11 +08:00
parent a4e3c25b22
commit 0bef7c0cdf
14 changed files with 158 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
handler/tunnel/bind.go
View File

@ -4,9 +4,11 @@ import (
"context"
"crypto/md5"
"encoding/hex"
"fmt"
"net"
"github.com/go-gost/core/logger"
"github.com/go-gost/core/recorder"
"github.com/go-gost/relay"
"github.com/go-gost/x/internal/util/mux"
"github.com/google/uuid"
@ -57,8 +59,15 @@ func (h *tunnelHandler) handleBind(ctx context.Context, conn net.Conn, network,
if h.md.ingress != nil {
h.md.ingress.Set(ctx, addr, tunnelID.String())
}
if h.recorder.Recorder != nil {
h.recorder.Recorder.Record(ctx, tunnelID[:])
if h.recorder != nil {
h.recorder.Record(ctx,
[]byte(fmt.Sprintf("%s:%s", tunnelID, connectorID)),
recorder.MetadataReocrdOption(connectorMetadata{
Op: "add",
Network: network,
Server: conn.LocalAddr().String(),
}),
)
}
log.Debugf("%s/%s: tunnel=%s, connector=%s established", addr, network, tunnelID, connectorID)

9
handler/tunnel/handler.go
View File

@ -39,7 +39,7 @@ type tunnelHandler struct {
md metadata
options handler.Options
pool *ConnectorPool
recorder recorder.RecorderObject
recorder recorder.Recorder
epSvc service.Service
ep *entrypoint
}
@ -52,7 +52,6 @@ func NewHandler(opts ...handler.Option) handler.Handler {
return &tunnelHandler{
options: options,
pool: NewConnectorPool(),
}
}
@ -68,12 +67,14 @@ func (h *tunnelHandler) Init(md md.Metadata) (err error) {
if opts := h.router.Options(); opts != nil {
for _, ro := range opts.Recorders {
if ro.Record == xrecorder.RecorderServiceHandlerTunnelEndpoint {
h.recorder = ro
if ro.Record == xrecorder.RecorderServiceHandlerTunnelConnector {
h.recorder = ro.Recorder
break
}
}
}
h.pool = NewConnectorPool()
h.pool.WithRecorder(h.recorder)
h.ep = &entrypoint{
pool: h.pool,

121
handler/tunnel/tunnel.go
View File

@ -1,6 +1,7 @@
package tunnel
import (
"context"
"fmt"
"net"
"sync"
@ -8,11 +9,18 @@ import (
"time"
"github.com/go-gost/core/logger"
"github.com/go-gost/core/recorder"
"github.com/go-gost/relay"
"github.com/go-gost/x/internal/util/mux"
"github.com/google/uuid"
)
type connectorMetadata struct {
Op string
Network string
Server string
}
type Connector struct {
id relay.ConnectorID
t time.Time
@ -54,18 +62,25 @@ type Tunnel struct {
connectors []*Connector
t time.Time
n uint64
close chan struct{}
mu sync.RWMutex
recorder recorder.Recorder
}
func NewTunnel(id relay.TunnelID) *Tunnel {
t := &Tunnel{
id: id,
t: time.Now(),
id: id,
t: time.Now(),
close: make(chan struct{}),
}
go t.clean()
return t
}
func (t *Tunnel) WithRecorder(recorder recorder.Recorder) {
t.recorder = recorder
}
func (t *Tunnel) ID() relay.TunnelID {
return t.id
}
@ -87,6 +102,9 @@ func (t *Tunnel) GetConnector(network string) *Connector {
var connectors []*Connector
for _, c := range t.connectors {
if c.Session().IsClosed() {
continue
}
if network == "udp" && c.id.IsUDP() ||
network != "udp" && !c.id.IsUDP() {
connectors = append(connectors, c)
@ -99,34 +117,83 @@ func (t *Tunnel) GetConnector(network string) *Connector {
return connectors[n%uint64(len(connectors))]
}
func (t *Tunnel) CloseOnIdle() bool {
t.mu.RLock()
defer t.mu.RUnlock()
select {
case <-t.close:
default:
if len(t.connectors) == 0 {
close(t.close)
return true
}
}
return false
}
func (t *Tunnel) clean() {
ticker := time.NewTicker(30 * time.Second)
for range ticker.C {
t.mu.Lock()
var connectors []*Connector
for _, c := range t.connectors {
if c.Session().IsClosed() {
logger.Default().Debugf("remove tunnel %s connector %s", t.id, c.id)
continue
ticker := time.NewTicker(1 * time.Minute)
defer ticker.Stop()
for {
select {
case <-ticker.C:
t.mu.Lock()
if len(t.connectors) == 0 {
t.mu.Unlock()
}
connectors = append(connectors, c)
var connectors []*Connector
for _, c := range t.connectors {
if c.Session().IsClosed() {
logger.Default().Debugf("remove tunnel: %s, connector: %s", t.id, c.id)
if t.recorder != nil {
t.recorder.Record(context.Background(),
[]byte(fmt.Sprintf("%s:%s", t.id, c.id)),
recorder.MetadataReocrdOption(connectorMetadata{
Op: "del",
}),
)
}
continue
}
connectors = append(connectors, c)
if t.recorder != nil {
t.recorder.Record(context.Background(),
[]byte(fmt.Sprintf("%s:%s", t.id, c.id)),
recorder.MetadataReocrdOption(connectorMetadata{
Op: "set",
}),
)
}
}
if len(connectors) != len(t.connectors) {
t.connectors = connectors
}
t.mu.Unlock()
case <-t.close:
return
}
if len(connectors) != len(t.connectors) {
t.connectors = connectors
}
t.mu.Unlock()
}
}
type ConnectorPool struct {
tunnels map[string]*Tunnel
mu sync.RWMutex
tunnels map[string]*Tunnel
mu sync.RWMutex
recorder recorder.Recorder
}
func NewConnectorPool() *ConnectorPool {
return &ConnectorPool{
p := &ConnectorPool{
tunnels: make(map[string]*Tunnel),
}
go p.closeIdles()
return p
}
func (p *ConnectorPool) WithRecorder(recorder recorder.Recorder) {
p.recorder = recorder
}
func (p *ConnectorPool) Add(tid relay.TunnelID, c *Connector) {
@ -138,6 +205,8 @@ func (p *ConnectorPool) Add(tid relay.TunnelID, c *Connector) {
t := p.tunnels[s]
if t == nil {
t = NewTunnel(tid)
t.WithRecorder(p.recorder)
p.tunnels[s] = t
}
t.AddConnector(c)
@ -159,6 +228,22 @@ func (p *ConnectorPool) Get(network string, tid relay.TunnelID) *Connector {
return t.GetConnector(network)
}
func (p *ConnectorPool) closeIdles() {
ticker := time.NewTicker(1 * time.Hour)
defer ticker.Stop()
for range ticker.C {
p.mu.Lock()
for k, v := range p.tunnels {
if v.CloseOnIdle() {
delete(p.tunnels, k)
logger.Default().Debugf("remove idle tunnel: %s", k)
}
}
p.mu.Unlock()
}
}
func parseTunnelID(s string) (tid relay.TunnelID) {
if s == "" {
return