hysteria/app/internal/proxymux/mux_test.go
2024-04-11 23:21:32 +08:00

111 lines
2.2 KiB
Go

package proxymux
import (
"net"
"testing"
"time"
"github.com/apernet/hysteria/app/internal/proxymux/internal/mocks"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/assert"
)
//go:generate mockery
func testMockListener(t *testing.T, firstByte byte) net.Listener {
mockConn := mocks.NewMockConn(t)
mockConn.EXPECT().Read(mock.Anything).RunAndReturn(func(b []byte) (int, error) {
b[0] = firstByte
return 1, nil
})
mockConn.EXPECT().Close().Return(nil)
mockListener := mocks.NewMockListener(t)
mockListener.EXPECT().Accept().RunAndReturn(func() (net.Conn, error) {
// Wait for all listener set up
time.Sleep(200 * time.Millisecond)
return mockConn, nil
})
mockListener.EXPECT().Close().Return(nil)
return mockListener
}
func TestMuxHTTP(t *testing.T) {
mockListener := testMockListener(t, 'C')
mux := newMuxListener(mockListener, func() {})
hl, err := mux.ListenHTTP()
if !assert.NoError(t, err) {
return
}
sl, err := mux.ListenSOCKS()
if !assert.NoError(t, err) {
return
}
var socksConn, httpConn net.Conn
var socksErr, httpErr error
go func() {
socksConn, socksErr = sl.Accept()
}()
go func() {
httpConn, httpErr = hl.Accept()
}()
time.Sleep(1 * time.Second)
sl.Close()
hl.Close()
// Wait for unmatched handler error
time.Sleep(1 * time.Second)
assert.Nil(t, socksConn)
assert.ErrorIs(t, socksErr, net.ErrClosed)
assert.NotNil(t, httpConn)
httpConn.Close()
assert.NoError(t, httpErr)
// Wait for muxListener released
time.Sleep(time.Second)
}
func TestMuxSOCKS(t *testing.T) {
mockListener := testMockListener(t, '\x05')
mux := newMuxListener(mockListener, func() {})
hl, err := mux.ListenHTTP()
if !assert.NoError(t, err) {
return
}
sl, err := mux.ListenSOCKS()
if !assert.NoError(t, err) {
return
}
var socksConn, httpConn net.Conn
var socksErr, httpErr error
go func() {
socksConn, socksErr = sl.Accept()
}()
go func() {
httpConn, httpErr = hl.Accept()
}()
time.Sleep(1 * time.Second)
sl.Close()
hl.Close()
// Wait for unmatched handler error
time.Sleep(1 * time.Second)
assert.NotNil(t, socksConn)
socksConn.Close()
assert.NoError(t, socksErr)
assert.Nil(t, httpConn)
assert.ErrorIs(t, httpErr, net.ErrClosed)
// Wait for muxListener released
time.Sleep(time.Second)
}