fix(agent/stream) TCP/UDP server now handle stream headers with read deadlines

This commit is contained in:
yusing
2026-01-09 21:22:51 +08:00
parent 69d04f1b76
commit 2e57ca7743
2 changed files with 6 additions and 0 deletions

View File

@@ -7,6 +7,7 @@ import (
"errors" "errors"
"io" "io"
"net" "net"
"time"
"github.com/rs/zerolog" "github.com/rs/zerolog"
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
@@ -147,9 +148,11 @@ func (s *TCPServer) handle(conn net.Conn) {
func (s *TCPServer) redirect(conn net.Conn) (net.Conn, error) { func (s *TCPServer) redirect(conn net.Conn) (net.Conn, error) {
// Read the stream header once as a handshake. // Read the stream header once as a handshake.
var headerBuf [headerSize]byte var headerBuf [headerSize]byte
_ = conn.SetReadDeadline(time.Now().Add(dialTimeout))
if _, err := io.ReadFull(conn, headerBuf[:]); err != nil { if _, err := io.ReadFull(conn, headerBuf[:]); err != nil {
return nil, err return nil, err
} }
_ = conn.SetReadDeadline(time.Time{})
header := ToHeader(&headerBuf) header := ToHeader(&headerBuf)
if !header.Validate() { if !header.Validate() {

View File

@@ -102,10 +102,13 @@ func (s *UDPServer) handleDTLSConnection(clientConn net.Conn) {
// Read the stream header once as a handshake. // Read the stream header once as a handshake.
var headerBuf [headerSize]byte var headerBuf [headerSize]byte
_ = clientConn.SetReadDeadline(time.Now().Add(dialTimeout))
if _, err := io.ReadFull(clientConn, headerBuf[:]); err != nil { if _, err := io.ReadFull(clientConn, headerBuf[:]); err != nil {
s.logger(clientConn).Err(err).Msg("failed to read stream header") s.logger(clientConn).Err(err).Msg("failed to read stream header")
return return
} }
_ = clientConn.SetReadDeadline(time.Time{})
header := ToHeader(&headerBuf) header := ToHeader(&headerBuf)
if !header.Validate() { if !header.Validate() {
s.logger(clientConn).Error().Bytes("header", headerBuf[:]).Msg("invalid stream header received") s.logger(clientConn).Error().Bytes("header", headerBuf[:]).Msg("invalid stream header received")