Update tg-ws-proxy.go

v1.0.1
This commit is contained in:
amurcanov 2026-03-25 05:03:25 +03:00 committed by GitHub
parent 7f34c0123b
commit b67f45e497
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 36 additions and 7 deletions

View File

@ -54,12 +54,12 @@ const (
defaultRecvBuf = 256 * 1024
defaultSendBuf = 256 * 1024
wsPoolSize = 4
wsPoolMaxAge = 90.0 // seconds — reduced from 120 to catch stale faster
wsPoolMaxAge = 60.0 // seconds — reduced from 90 to 60 for VPNs
wsBridgeIdle = 120.0 // seconds — max idle time before bridge considers WS dead
dcFailCooldown = 30.0 // seconds
wsFailTimeout = 2.0 // seconds
poolMaintainInterval = 30 // seconds — how often pool evicts stale and refills
poolMaintainInterval = 15 // seconds — frequent to send pings/keep-alives
)
var (
@ -169,6 +169,7 @@ var ipToDC = map[string]dcInfo{
// DC2
"149.154.167.41": {2, false}, "149.154.167.50": {2, false},
"149.154.167.51": {2, false}, "149.154.167.220": {2, false},
"149.154.167.35": {2, false}, "149.154.167.36": {2, false},
"95.161.76.100": {2, false},
"149.154.167.151": {2, true}, "149.154.167.222": {2, true},
"149.154.167.223": {2, true}, "149.154.162.123": {2, true},
@ -242,7 +243,7 @@ func (s *Stats) Summary() string {
ph := s.poolHits.Load()
pm := s.poolMisses.Load()
return fmt.Sprintf(
"total=%d ws=%d tcp_fb=%d http_skip=%d pass=%d err=%d pool=%d/%d up=%s down=%s",
"total=%d ws=%d tcp_fb=%d http_skip=%d pass=%d err=%d pool_hits=%d/%d up=%s down=%s",
s.connectionsTotal.Load(),
s.connectionsWs.Load(),
s.connectionsTcpFallback.Load(),
@ -554,6 +555,17 @@ func (ws *RawWebSocket) Close() {
_ = ws.conn.Close()
}
func (ws *RawWebSocket) SendPing() error {
if ws.closed.Load() {
return fmt.Errorf("WebSocket closed")
}
frame := buildFrame(opPing, []byte{}, true)
ws.mu.Lock()
defer ws.mu.Unlock()
_, err := ws.conn.Write(frame)
return err
}
func buildFrame(opcode int, data []byte, mask bool) []byte {
length := len(data)
fb := byte(0x80 | opcode)
@ -992,6 +1004,12 @@ func (p *WsPool) Maintain(ctx context.Context, dcOptMap map[int]string) {
if age > wsPoolMaxAge || e.ws.closed.Load() {
go e.ws.Close()
} else {
// Send ping to keep VPN NAT session alive, and actively close if network is dead
go func(ws *RawWebSocket) {
if err := ws.SendPing(); err != nil {
ws.Close()
}
}(e.ws)
fresh = append(fresh, e)
}
}
@ -1016,6 +1034,16 @@ func (p *WsPool) Maintain(ctx context.Context, dcOptMap map[int]string) {
}
}
func (p *WsPool) IdleCount() int {
p.mu.Lock()
defer p.mu.Unlock()
count := 0
for _, bucket := range p.idle {
count += len(bucket)
}
return count
}
var wsPool = newWsPool()
// ---------------------------------------------------------------------------
@ -1426,7 +1454,7 @@ func handleClient(ctx context.Context, conn net.Conn) {
dcOptMu.RUnlock()
if !dcOk || !dcConfigured {
logWarn.Printf("[%s] unknown DC%d for %s:%d -> TCP passthrough", label, dc, dst, port)
logDebug.Printf("[%s] unknown DC%d for %s:%d -> TCP passthrough", label, dc, dst, port)
tcpFallback(ctx, conn, dst, port, init, label, dc, isMedia)
return
}
@ -1565,7 +1593,7 @@ func handleClient(ctx context.Context, conn net.Conn) {
// Send init packet
if err := ws.Send(init); err != nil {
logWarn.Printf("[%s] failed to send init via WS: %v", label, err)
logDebug.Printf("[%s] reconnecting via TCP fallback (WS broken by NAT): %v", label, err)
ws.Close()
tcpFallback(ctx, conn, dst, port, init, label, dc, isMedia)
return
@ -1651,7 +1679,8 @@ func runProxy(ctx context.Context, host string, port int, dcOptMap map[int]strin
if len(blParts) > 0 {
bl = strings.Join(blParts, ", ")
}
logInfo.Printf("stats: %s | ws_bl: %s", stats.Summary(), bl)
idleCount := wsPool.IdleCount()
logInfo.Printf("stats: %s idle_conns=%d | ws_bl: %s", stats.Summary(), idleCount, bl)
}
}
}()
@ -1951,4 +1980,4 @@ func main() {
// Ensure mrand is used to avoid import errors
var _ = mrand.Int
var _ = big.NewInt
var _ = big.NewInt