diff --git a/internal/client/piper.go b/internal/client/piper.go index 0cbf2d9..fd1746a 100644 --- a/internal/client/piper.go +++ b/internal/client/piper.go @@ -44,7 +44,7 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration streams[addr.String()] = stream proxyAddr := addr - go func() { + go func(stream *mux.Stream, localConn *net.UDPConn) { buf := make([]byte, 8192) for { n, err := stream.Read(buf) @@ -61,7 +61,7 @@ func RouteUDP(bindFunc func() (*net.UDPConn, error), streamTimeout time.Duration return } } - }() + }(stream, localConn) } _, err = stream.Write(data[:i]) @@ -85,7 +85,7 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu if sesh == nil || sesh.IsClosed() || sesh.Singleplex { sesh = newSeshFunc() } - go func() { + go func(sesh *mux.Session, localConn net.Conn) { data := make([]byte, 10240) i, err := io.ReadAtLeast(localConn, data, 1) if err != nil { @@ -121,7 +121,7 @@ func RouteTCP(listener net.Listener, streamTimeout time.Duration, newSeshFunc fu if _, err = common.Copy(stream, localConn); err != nil { log.Tracef("copying proxy client to stream: %v", err) } - }() + }(sesh, localConn) } } diff --git a/internal/test/integration_test.go b/internal/test/integration_test.go index 773227e..23f8a61 100644 --- a/internal/test/integration_test.go +++ b/internal/test/integration_test.go @@ -108,7 +108,9 @@ var basicTCPConfig = client.RawConfig{ RemotePort: "9999", LocalHost: "127.0.0.1", LocalPort: "9999", + BrowserSig: "firefox", } + var singleplexTCPConfig = client.RawConfig{ ServerName: "www.example.com", ProxyMethod: "shadowsocks", @@ -122,6 +124,7 @@ var singleplexTCPConfig = client.RawConfig{ RemotePort: "9999", LocalHost: "127.0.0.1", LocalPort: "9999", + BrowserSig: "chrome", } func generateClientConfigs(rawConfig client.RawConfig, state common.WorldState) (client.LocalConnConfig, client.RemoteConnConfig, client.AuthInfo) { @@ -327,14 +330,7 @@ func TestTCPSingleplex(t *testing.T) { if err != nil { t.Error(err) } - _, err = proxyConn1.Write([]byte("hello")) - if err != nil { - t.Error(err) - } - - // make sure the server has accepted the connection before fetching the server - proxyConn1.Read(make([]byte, 10)) - + runEchoTest(t, []net.Conn{proxyConn1}, 65536) user, err := sta.Panel.GetUser(ai.UID[:]) if err != nil { t.Fatalf("failed to fetch user: %v", err) @@ -348,19 +344,33 @@ func TestTCPSingleplex(t *testing.T) { if err != nil { t.Error(err) } - proxyConn2.Write([]byte("hello")) - // make sure the server has accepted the connection before fetching the server - proxyConn2.Read(make([]byte, 10)) + runEchoTest(t, []net.Conn{proxyConn2}, 65536) if user.NumSession() != 2 { t.Error("no extra session were made on second connection establishment") } + // Both conns should work + runEchoTest(t, []net.Conn{proxyConn1, proxyConn2}, 65536) + proxyConn1.Close() time.Sleep(delayBeforeTestingConnClose) if user.NumSession() != 1 { t.Error("first session was not closed on connection close") } + // conn2 should still work + runEchoTest(t, []net.Conn{proxyConn2}, 65536) + + var conns [numConns]net.Conn + for i := 0; i < numConns; i++ { + conns[i], err = proxyToCkClientD.Dial("", "") + if err != nil { + t.Error(err) + } + } + + runEchoTest(t, conns[:], 65536) + } func TestTCPMultiplex(t *testing.T) {