bagas/tunnel-please
Public Access
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

fix: resolve resource exhaustion with high connection counts
All checks were successful
Docker Build and Push / build-and-push (push) Successful in 3m39s
Details

Browse Source
This commit is contained in:
bagas
2025-07-22 12:55:02 +07:00
parent 4b21541668
commit 28f0caced1
1 changed files with 8 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
session/handler.go
View File

@ -3,7 +3,6 @@ package session
import ( import (
"bufio" "bufio"
"bytes" "bytes"
"context"
"encoding/binary" "encoding/binary"
"errors" "errors"
"fmt" "fmt"
@ -557,60 +556,21 @@ func (s *Session) HandleForwardedConnection(conn UserConnection, sshConn *ssh.Se
conn.Reader = bufio.NewReader(conn.Writer) conn.Reader = bufio.NewReader(conn.Writer)
} }
ctx, cancel := context.WithCancel(context.Background()) go io.Copy(channel, conn.Reader)
defer cancel()
go func() {
defer func() {
if r := recover(); r != nil {
log.Printf("Panic in reader copy: %v", r)
}
cancel()
}()
_, err := io.Copy(channel, conn.Reader)
if err != nil && !errors.Is(err, io.EOF) {
log.Printf("Error copying from conn.Reader to channel: %v", err)
}
cancel()
}()
reader := bufio.NewReader(channel) reader := bufio.NewReader(channel)
_, err = reader.Peek(1)
peekChan := make(chan error, 1) if err == io.EOF {
go func() { s.sendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"%s\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.TunnelType, "Could not forward request to the tunnel addr"))
_, err := reader.Peek(1)
peekChan <- err
}()
select {
case err := <-peekChan:
if err == io.EOF {
s.sendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.TunnelType))
sendBadGatewayResponse(conn.Writer)
return
}
if err != nil {
log.Printf("Error peeking channel data: %v", err)
s.sendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.TunnelType))
sendBadGatewayResponse(conn.Writer)
return
}
case <-time.After(5 * time.Second):
log.Printf("Timeout waiting for channel data from %s", conn.Writer.RemoteAddr())
s.sendMessage(fmt.Sprintf("\033[33m%s -> [%s] WARNING -- \"Could not forward request to the tunnel addr\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.TunnelType))
sendBadGatewayResponse(conn.Writer) sendBadGatewayResponse(conn.Writer)
return conn.Writer.Close()
case <-ctx.Done(): channel.Close()
return return
} }
s.sendMessage(fmt.Sprintf("\033[32m%s -> [%s] TUNNEL ADDRESS -- \"%s\"\033[0m\r\n", conn.Writer.RemoteAddr().String(), s.TunnelType, timestamp)) s.sendMessage(fmt.Sprintf("\033[32m %s -> [%s] TUNNEL ADDRESS -- \"%s\" \r\n \033[0m", conn.Writer.RemoteAddr().String(), s.TunnelType, timestamp))
_, err = io.Copy(conn.Writer, reader) io.Copy(conn.Writer, reader)
if err != nil && !errors.Is(err, io.EOF) {
log.Printf("Error copying from channel to conn.Writer: %v", err)
}
} }
func sendBadGatewayResponse(writer io.Writer) { func sendBadGatewayResponse(writer io.Writer) {