diff --git a/common/mux/client.go b/common/mux/client.go index 0463be5fc..1d0180eef 100644 --- a/common/mux/client.go +++ b/common/mux/client.go @@ -314,12 +314,10 @@ func (m *ClientWorker) Dispatch(ctx context.Context, link *transport.Link) bool } sm := m.sessionManager - s := sm.Allocate(&m.strategy) + s := sm.Allocate(&m.strategy, link.Reader, link.Writer) if s == nil { return false } - s.input = link.Reader - s.output = link.Writer go fetchInput(ctx, s, m.link.Writer) if _, ok := link.Reader.(*pipe.Reader); !ok { select { diff --git a/common/mux/session.go b/common/mux/session.go index 66b9674cf..53bc53569 100644 --- a/common/mux/session.go +++ b/common/mux/session.go @@ -51,7 +51,7 @@ func (m *SessionManager) Count() int { return int(m.count) } -func (m *SessionManager) Allocate(Strategy *ClientStrategy) *Session { +func (m *SessionManager) Allocate(Strategy *ClientStrategy, input buf.Reader, output buf.Writer) *Session { m.Lock() defer m.Unlock() @@ -64,6 +64,8 @@ func (m *SessionManager) Allocate(Strategy *ClientStrategy) *Session { m.count++ s := &Session{ + input: input, + output: output, ID: m.count, parent: m, done: done.New(), diff --git a/common/mux/session_test.go b/common/mux/session_test.go index 8ef278778..7c47cb052 100644 --- a/common/mux/session_test.go +++ b/common/mux/session_test.go @@ -9,7 +9,7 @@ import ( func TestSessionManagerAdd(t *testing.T) { m := NewSessionManager() - s := m.Allocate(&ClientStrategy{}) + s := m.Allocate(&ClientStrategy{}, nil, nil) if s.ID != 1 { t.Error("id: ", s.ID) } @@ -17,7 +17,7 @@ func TestSessionManagerAdd(t *testing.T) { t.Error("size: ", m.Size()) } - s = m.Allocate(&ClientStrategy{}) + s = m.Allocate(&ClientStrategy{}, nil, nil) if s.ID != 2 { t.Error("id: ", s.ID) } @@ -39,7 +39,7 @@ func TestSessionManagerAdd(t *testing.T) { func TestSessionManagerClose(t *testing.T) { m := NewSessionManager() - s := m.Allocate(&ClientStrategy{}) + s := m.Allocate(&ClientStrategy{}, nil, nil) if m.CloseIfNoSessionAndIdle(m.Size(), m.Count()) { t.Error("able to close")