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

refactor(forwarder): remove CreateForwardedTCPIPPayload method

Browse Source 
- OpenForwardedChannel now privately calls CreateForwardedTCPIPPayload
- Removed an unused function
This commit is contained in:
bagas
2026-01-25 20:49:12 +07:00
parent 79fd292a77
commit 65df01fee5
6 changed files with 20 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/transport/httphandler.go
+1 -2
View File
@@ -152,10 +152,9 @@ func (hh *httpHandler) handlePingRequest(slug string, conn net.Conn) bool {
}
func (hh *httpHandler) forwardRequest(hw stream.HTTP, initialRequest header.RequestHeader, sshSession registry.Session) {
payload := sshSession.Forwarder().CreateForwardedTCPIPPayload(hw.RemoteAddr())
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
defer cancel()
channel, reqs, err := sshSession.Forwarder().OpenForwardedChannel(ctx, payload)
channel, reqs, err := sshSession.Forwarder().OpenForwardedChannel(ctx, hw.RemoteAddr())
if err != nil {
log.Printf("Failed to open forwarded-tcpip channel: %v", err)
return
internal/transport/httphandler_test.go
+8 -17
View File
@@ -176,15 +176,6 @@ type MockForwarder struct {
mock.Mock
}
func (m *MockForwarder) CreateForwardedTCPIPPayload(origin net.Addr) []byte {
args := m.Called(origin)
return args.Get(0).([]byte)
}
func (m *MockForwarder) WriteBadGatewayResponse(dst io.Writer) {
m.Called(dst)
}
func (m *MockForwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel) {
m.Called(dst, src)
}
@@ -221,8 +212,8 @@ func (m *MockForwarder) Listener() net.Listener {
return args.Get(0).(net.Listener)
}
func (m *MockForwarder) OpenForwardedChannel(ctx context.Context, payload []byte) (ssh.Channel, <-chan *ssh.Request, error) {
args := m.Called(ctx, payload)
func (m *MockForwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error) {
args := m.Called(ctx, origin)
if args.Get(0) == nil {
return nil, args.Get(1).(<-chan *ssh.Request), args.Error(2)
}
@@ -413,7 +404,7 @@ func TestHandler(t *testing.T) {
mockSession.On("Forwarder").Return(mockForwarder)
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Return((ssh.Channel)(nil), (<-chan *ssh.Request)(nil), fmt.Errorf("open channel failed"))
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return((ssh.Channel)(nil), (<-chan *ssh.Request)(nil), fmt.Errorf("open channel failed"))
},
},
{
@@ -437,7 +428,7 @@ func TestHandler(t *testing.T) {
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
reqCh := make(chan *ssh.Request)
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockSSHChannel.On("Write", mock.Anything).Return(0, fmt.Errorf("write error"))
mockSSHChannel.On("Close").Return(nil)
@@ -469,7 +460,7 @@ func TestHandler(t *testing.T) {
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
reqCh := make(chan *ssh.Request)
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockSSHChannel.On("Write", mock.Anything).Return(0, nil)
mockSSHChannel.On("Close").Return(nil)
@@ -577,7 +568,7 @@ func TestHandler(t *testing.T) {
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
reqCh := make(chan *ssh.Request)
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockSSHChannel.On("Write", mock.Anything).Return(0, nil)
mockSSHChannel.On("Close").Return(nil)
@@ -640,7 +631,7 @@ func TestHandler(t *testing.T) {
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
reqCh := make(chan *ssh.Request)
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
mockSSHChannel.On("Write", mock.Anything).Return(0, nil)
mockSSHChannel.On("Close").Return(fmt.Errorf("close error"))
@@ -666,7 +657,7 @@ func TestHandler(t *testing.T) {
mockForwarder.On("CreateForwardedTCPIPPayload", mock.Anything).Return([]byte("payload"))
mockForwarder.On("OpenForwardedChannel", mock.Anything, []byte("payload")).Run(func(args mock.Arguments) {
mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
ctx := args.Get(0).(context.Context)
<-ctx.Done()
}).Return((ssh.Channel)(nil), (<-chan *ssh.Request)(nil), context.DeadlineExceeded)
internal/transport/tcp.go
+2 -4
View File
@@ -18,8 +18,7 @@ type tcp struct {
}
type Forwarder interface {
CreateForwardedTCPIPPayload(origin net.Addr) []byte
OpenForwardedChannel(ctx context.Context, payload []byte) (ssh.Channel, <-chan *ssh.Request, error)
OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error)
HandleConnection(dst io.ReadWriter, src ssh.Channel)
}
@@ -55,10 +54,9 @@ func (tt *tcp) handleTcp(conn net.Conn) {
log.Printf("Failed to close connection: %v", err)
}
}()
payload := tt.forwarder.CreateForwardedTCPIPPayload(conn.RemoteAddr())
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
defer cancel()
channel, reqs, err := tt.forwarder.OpenForwardedChannel(ctx, payload)
channel, reqs, err := tt.forwarder.OpenForwardedChannel(ctx, conn.RemoteAddr())
if err != nil {
log.Printf("Failed to open forwarded-tcpip channel: %v", err)
return
internal/transport/tcp_test.go
+4 -13
View File
@@ -72,10 +72,8 @@ func TestTCPServer_Serve_Success(t *testing.T) {
assert.NoError(t, err)
port := listener.Addr().(*net.TCPAddr).Port
payload := []byte("test-payload")
mf.On("CreateForwardedTCPIPPayload", mock.Anything).Return(payload)
reqs := make(chan *ssh.Request)
mf.On("OpenForwardedChannel", mock.Anything, payload).Return(new(MockSSHChannel), (<-chan *ssh.Request)(reqs), nil)
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(new(MockSSHChannel), (<-chan *ssh.Request)(reqs), nil)
mf.On("HandleConnection", mock.Anything, mock.Anything).Return()
go func() {
@@ -99,12 +97,9 @@ func TestTCPServer_handleTcp_Success(t *testing.T) {
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
payload := []byte("test-payload")
mf.On("CreateForwardedTCPIPPayload", mock.Anything).Return(payload)
reqs := make(chan *ssh.Request)
mockChannel := new(MockSSHChannel)
mf.On("OpenForwardedChannel", mock.Anything, payload).Return(mockChannel, (<-chan *ssh.Request)(reqs), nil)
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockChannel, (<-chan *ssh.Request)(reqs), nil)
mf.On("HandleConnection", serverConn, mockChannel).Return()
@@ -121,9 +116,7 @@ func TestTCPServer_handleTcp_CloseError(t *testing.T) {
mc.On("Close").Return(errors.New("close error"))
mc.On("RemoteAddr").Return(&net.TCPAddr{})
payload := []byte("test-payload")
mf.On("CreateForwardedTCPIPPayload", mock.Anything).Return(payload)
mf.On("OpenForwardedChannel", mock.Anything, payload).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
srv.handleTcp(mc)
mc.AssertExpectations(t)
@@ -136,9 +129,7 @@ func TestTCPServer_handleTcp_OpenChannelError(t *testing.T) {
serverConn, clientConn := net.Pipe()
defer clientConn.Close()
payload := []byte("test-payload")
mf.On("CreateForwardedTCPIPPayload", mock.Anything).Return(payload)
mf.On("OpenForwardedChannel", mock.Anything, payload).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
mf.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(nil, (<-chan *ssh.Request)(nil), errors.New("open error"))
srv.handleTcp(serverConn)