2 Commits

Author SHA1 Message Date
7716eb7f29 perf: optimize header parsing with zero-copy ReadSlice
All checks were successful
renovate / renovate (push) Successful in 35s
Docker Build and Push / build-and-push-branches (push) Successful in 4m39s
Docker Build and Push / build-and-push-tags (push) Successful in 4m52s
- Replace ReadString with ReadSlice to eliminate allocations
- Use bytes operations instead of strings
- Add FromBytes variant for in-memory parsing
2025-12-31 23:18:53 +07:00
b115369913 fix: wait for both goroutines before cleanup in HandleConnection
All checks were successful
renovate / renovate (push) Successful in 1m42s
Docker Build and Push / build-and-push-branches (push) Successful in 4m46s
Docker Build and Push / build-and-push-tags (push) Successful in 4m51s
Only waited for one of two copy goroutines, leaking the second. Now waits
for both to complete before closing connections.

Fixes file descriptor leak causing 'too many open files' under load.

Fixes: #56
2025-12-31 22:22:51 +07:00
3 changed files with 123 additions and 35 deletions

View File

@@ -4,7 +4,6 @@ import (
"bufio"
"bytes"
"fmt"
"strings"
)
type HeaderManager interface {
@@ -44,43 +43,132 @@ type requestHeaderFactory struct {
headers map[string]string
}
func NewRequestHeaderFactory(br *bufio.Reader) (RequestHeaderManager, error) {
func NewRequestHeaderFactory(r interface{}) (RequestHeaderManager, error) {
switch v := r.(type) {
case []byte:
return parseHeadersFromBytes(v)
case *bufio.Reader:
return parseHeadersFromReader(v)
default:
return nil, fmt.Errorf("unsupported type: %T", r)
}
}
func parseHeadersFromBytes(headerData []byte) (RequestHeaderManager, error) {
header := &requestHeaderFactory{
headers: make(map[string]string),
headers: make(map[string]string, 16),
}
startLine, err := br.ReadString('\n')
if err != nil {
return nil, err
lineEnd := bytes.IndexByte(headerData, '\n')
if lineEnd == -1 {
return nil, fmt.Errorf("invalid request: no newline found")
}
startLine = strings.TrimRight(startLine, "\r\n")
header.startLine = []byte(startLine)
parts := strings.Split(startLine, " ")
startLine := bytes.TrimRight(headerData[:lineEnd], "\r\n")
header.startLine = make([]byte, len(startLine))
copy(header.startLine, startLine)
parts := bytes.Split(startLine, []byte{' '})
if len(parts) < 3 {
return nil, fmt.Errorf("invalid request line")
}
header.method = parts[0]
header.path = parts[1]
header.version = parts[2]
header.method = string(parts[0])
header.path = string(parts[1])
header.version = string(parts[2])
for {
line, err := br.ReadString('\n')
if err != nil {
return nil, err
remaining := headerData[lineEnd+1:]
for len(remaining) > 0 {
lineEnd = bytes.IndexByte(remaining, '\n')
if lineEnd == -1 {
lineEnd = len(remaining)
}
line = strings.TrimRight(line, "\r\n")
if line == "" {
line := bytes.TrimRight(remaining[:lineEnd], "\r\n")
if len(line) == 0 {
break
}
kv := strings.SplitN(line, ":", 2)
if len(kv) != 2 {
colonIdx := bytes.IndexByte(line, ':')
if colonIdx != -1 {
key := bytes.TrimSpace(line[:colonIdx])
value := bytes.TrimSpace(line[colonIdx+1:])
header.headers[string(key)] = string(value)
}
if lineEnd == len(remaining) {
break
}
remaining = remaining[lineEnd+1:]
}
return header, nil
}
func parseHeadersFromReader(br *bufio.Reader) (RequestHeaderManager, error) {
header := &requestHeaderFactory{
headers: make(map[string]string, 16),
}
startLineBytes, err := br.ReadSlice('\n')
if err != nil {
if err == bufio.ErrBufferFull {
var startLine string
startLine, err = br.ReadString('\n')
if err != nil {
return nil, err
}
startLineBytes = []byte(startLine)
} else {
return nil, err
}
}
startLineBytes = bytes.TrimRight(startLineBytes, "\r\n")
header.startLine = make([]byte, len(startLineBytes))
copy(header.startLine, startLineBytes)
parts := bytes.Split(startLineBytes, []byte{' '})
if len(parts) < 3 {
return nil, fmt.Errorf("invalid request line")
}
header.method = string(parts[0])
header.path = string(parts[1])
header.version = string(parts[2])
for {
lineBytes, err := br.ReadSlice('\n')
if err != nil {
if err == bufio.ErrBufferFull {
var line string
line, err = br.ReadString('\n')
if err != nil {
return nil, err
}
lineBytes = []byte(line)
} else {
return nil, err
}
}
lineBytes = bytes.TrimRight(lineBytes, "\r\n")
if len(lineBytes) == 0 {
break
}
colonIdx := bytes.IndexByte(lineBytes, ':')
if colonIdx == -1 {
continue
}
header.headers[strings.TrimSpace(kv[0])] = strings.TrimSpace(kv[1])
key := bytes.TrimSpace(lineBytes[:colonIdx])
value := bytes.TrimSpace(lineBytes[colonIdx+1:])
header.headers[string(key)] = string(value)
}
return header, nil

View File

@@ -99,8 +99,7 @@ func (cw *customWriter) Read(p []byte) (int, error) {
}
}
headerReader := bufio.NewReader(bytes.NewReader(header))
reqhf, err := NewRequestHeaderFactory(headerReader)
reqhf, err := NewRequestHeaderFactory(header)
if err != nil {
return 0, err
}

View File

@@ -152,25 +152,26 @@ func (f *Forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel, remoteA
log.Printf("Handling new forwarded connection from %s", remoteAddr)
done := make(chan struct{}, 2)
go func() {
_, err := copyWithBuffer(src, dst)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
log.Printf("Error copying from conn.Reader to channel: %v", err)
}
done <- struct{}{}
}()
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
_, err := copyWithBuffer(dst, src)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
log.Printf("Error copying from channel to conn.Writer: %v", err)
log.Printf("Error copying src→dst: %v", err)
}
done <- struct{}{}
}()
<-done
go func() {
defer wg.Done()
_, err := copyWithBuffer(src, dst)
if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
log.Printf("Error copying dst→src: %v", err)
}
}()
wg.Wait()
}
func (f *Forwarder) SetType(tunnelType types.TunnelType) {