fix(port, session): harden port registry and fix concurent allocation #155

Merged
bagas merged 4 commits from fix/harden-port into staging 2026-07-19 10:59:23 +07:00
2 changed files with 140 additions and 11 deletions
Showing only changes of commit 9d623ae3f9 - Show all commits
+22 -9
View File
@@ -1,6 +1,7 @@
package transport
import (
"bufio"
"bytes"
"context"
"errors"
@@ -52,25 +53,37 @@ func (hh *httpHandler) badRequest(conn net.Conn) error {
return nil
}
func readHTTPHeader(br *bufio.Reader, limit int) ([]byte, error) {
var headerBuf []byte
for {
line, err := br.ReadBytes('\n')
headerBuf = append(headerBuf, line...)
if err != nil {
return nil, err
}
if bytes.HasSuffix(headerBuf, []byte("\r\n\r\n")) {
return headerBuf, nil
}
if len(headerBuf) > limit {
return nil, fmt.Errorf("headers too large")
}
}
}
func (hh *httpHandler) Handler(conn net.Conn, isTLS bool) {
defer hh.closeConnection(conn)
_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
buf := make([]byte, hh.config.HeaderSize())
n, err := conn.Read(buf)
br := bufio.NewReaderSize(conn, hh.config.HeaderSize())
headerBuf, err := readHTTPHeader(br, hh.config.HeaderSize())
if err != nil {
_ = hh.badRequest(conn)
return
}
if idx := bytes.Index(buf[:n], []byte("\r\n\r\n")); idx == -1 {
_ = hh.badRequest(conn)
return
}
_ = conn.SetReadDeadline(time.Time{})
reqhf, err := header.NewRequest(buf[:n])
reqhf, err := header.NewRequest(headerBuf)
if err != nil {
log.Printf("Error creating request header: %v", err)
_ = hh.badRequest(conn)
@@ -101,7 +114,7 @@ func (hh *httpHandler) Handler(conn net.Conn, isTLS bool) {
return
}
hw := stream.New(conn, conn, conn.RemoteAddr())
hw := stream.New(conn, br, conn.RemoteAddr())
defer func(hw stream.HTTP) {
err = hw.Close()
if err != nil {
+118 -2
View File
@@ -321,8 +321,14 @@ func TestHandler(t *testing.T) {
isTLS: false,
redirectTLS: false,
request: []byte(""),
expected: []byte("HTTP/1.1 400 Bad Request\r\n\r\n"),
setupMocks: func(msr *MockSessionRegistry) {
expected: []byte(""),
setupConn: func() (net.Conn, net.Conn) {
mc := new(MockConn)
mc.ReadBuffer = bytes.NewBuffer(nil)
mc.On("SetReadDeadline", mock.Anything).Return(nil)
mc.On("Write", []byte("HTTP/1.1 400 Bad Request\r\n\r\n")).Return(0, nil)
mc.On("Close").Return(nil)
return mc, nil
},
},
{
@@ -715,3 +721,113 @@ func TestHandler(t *testing.T) {
})
}
}
func TestHandlerForwardsPostBody(t *testing.T) {
mockSessionRegistry := new(MockSessionRegistry)
mockConfig := &MockConfig{}
mockConfig.On("Domain").Return("example.com")
mockConfig.On("HTTPPort").Return("0")
mockConfig.On("HeaderSize").Return(4096)
mockConfig.On("TLSRedirect").Return(true)
hh := &httpHandler{
sessionRegistry: mockSessionRegistry,
config: mockConfig,
}
mockSession := new(MockSession)
mockForwarder := new(MockForwarder)
mockSSHChannel := new(MockSSHChannel)
mockSessionRegistry.On("Get", types.SessionKey{
Id: "test",
Type: types.TunnelTypeHTTP,
}).Return(mockSession, nil)
mockSession.On("Forwarder").Return(mockForwarder)
reqCh := make(chan *ssh.Request)
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
var mu sync.Mutex
var capturedHeaders []byte
mockSSHChannel.On("Write", mock.Anything).Run(func(args mock.Arguments) {
mu.Lock()
capturedHeaders = append(capturedHeaders, args.Get(0).([]byte)...)
mu.Unlock()
}).Return(0, nil)
mockSSHChannel.On("Close").Return(nil)
bodyChan := make(chan string, 1)
mockForwarder.On("HandleConnection", mock.Anything, mockSSHChannel).Run(func(args mock.Arguments) {
w := args.Get(0).(io.ReadWriter)
buf := make([]byte, len("hello=world"))
if _, err := io.ReadFull(w, buf); err != nil {
bodyChan <- ""
} else {
bodyChan <- string(buf)
}
_, _ = w.Write([]byte("HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok"))
})
go func() {
for range reqCh {
}
}()
serverConn, clientConn := net.Pipe()
defer func() {
_ = clientConn.Close()
}()
remoteAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:12345")
wrappedServerConn := &wrappedConn{Conn: serverConn, remoteAddr: remoteAddr}
go hh.Handler(wrappedServerConn, true)
request := []byte("POST / HTTP/1.1\r\nHost: test.domain\r\nContent-Type: application/x-www-form-urlencoded\r\nContent-Length: 11\r\n\r\nhello=world")
go func() {
_, _ = clientConn.Write(request)
}()
var response []byte
respDone := make(chan struct{})
go func() {
defer close(respDone)
buf := make([]byte, 4096)
for {
n, err := clientConn.Read(buf)
if err != nil {
break
}
response = append(response, buf[:n]...)
if bytes.Contains(response, []byte("\r\n\r\nok")) {
break
}
}
}()
select {
case body := <-bodyChan:
assert.Equal(t, "hello=world", body)
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for forwarded body")
}
select {
case <-respDone:
resStr := string(response)
assert.True(t, strings.HasPrefix(resStr, "HTTP/1.1 200 OK\r\n"))
assert.Contains(t, resStr, "Server: Tunnel Please\r\n")
assert.True(t, strings.HasSuffix(resStr, "\r\n\r\nok"))
case <-time.After(5 * time.Second):
t.Fatal("timed out waiting for response")
}
mu.Lock()
hdrStr := string(capturedHeaders)
mu.Unlock()
assert.Contains(t, hdrStr, "POST / HTTP/1.1\r\n")
assert.Contains(t, hdrStr, "Content-Length: 11\r\n")
assert.Contains(t, hdrStr, "X-Forwarded-For: 127.0.0.1\r\n")
mockSessionRegistry.AssertExpectations(t)
}