Improve multiplex

common/mux/client.go

@@ -41,7 +41,7 @@ func NewClient(ctx context.Context, dialer N.Dialer, protocol Protocol, maxConne
 	}
 }
 
-func NewClientWithOptions(ctx context.Context, dialer N.Dialer, options option.MultiplexOptions) (N.Dialer, error) {
+func NewClientWithOptions(ctx context.Context, dialer N.Dialer, options option.MultiplexOptions) (*Client, error) {
 	if !options.Enabled {
 		return nil, nil
 	}
@@ -120,17 +120,16 @@ func (c *Client) offer() (abstractSession, error) {
 		sessions = append(sessions, element.Value)
 		element = element.Next()
 	}
-	sLen := len(sessions)
-	if sLen == 0 {
+	session := common.MinBy(common.Filter(sessions, abstractSession.CanTakeNewRequest), abstractSession.NumStreams)
+	if session == nil {
 		return c.offerNew()
 	}
-	session := common.MinBy(sessions, abstractSession.NumStreams)
 	numStreams := session.NumStreams()
 	if numStreams == 0 {
 		return session, nil
 	}
 	if c.maxConnections > 0 {
-		if sLen >= c.maxConnections || numStreams < c.minStreams {
+		if len(sessions) >= c.maxConnections || numStreams < c.minStreams {
 			return session, nil
 		}
 	} else {
@@ -159,6 +158,15 @@ func (c *Client) offerNew() (abstractSession, error) {
 	return session, nil
 }
 
+func (c *Client) Reset() {
+	c.access.Lock()
+	defer c.access.Unlock()
+	for _, session := range c.connections.Array() {
+		session.Close()
+	}
+	c.connections.Init()
+}
+
 func (c *Client) Close() error {
 	c.access.Lock()
 	defer c.access.Unlock()

common/mux/h2mux.go

@@ -0,0 +1,235 @@
+package mux
+
+import (
+	"context"
+	"crypto/tls"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"os"
+	"time"
+
+	"github.com/sagernet/sing-box/transport/v2rayhttp"
+	"github.com/sagernet/sing/common/atomic"
+	"github.com/sagernet/sing/common/buf"
+	"github.com/sagernet/sing/common/bufio"
+	E "github.com/sagernet/sing/common/exceptions"
+	N "github.com/sagernet/sing/common/network"
+
+	"golang.org/x/net/http2"
+)
+
+const idleTimeout = 30 * time.Second
+
+var _ abstractSession = (*H2MuxServerSession)(nil)
+
+type H2MuxServerSession struct {
+	server  http2.Server
+	active  atomic.Int32
+	conn    net.Conn
+	inbound chan net.Conn
+	done    chan struct{}
+}
+
+func NewH2MuxServer(conn net.Conn) *H2MuxServerSession {
+	session := &H2MuxServerSession{
+		conn:    conn,
+		inbound: make(chan net.Conn),
+		done:    make(chan struct{}),
+		server: http2.Server{
+			IdleTimeout: idleTimeout,
+		},
+	}
+	go func() {
+		session.server.ServeConn(conn, &http2.ServeConnOpts{
+			Handler: session,
+		})
+		_ = session.Close()
+	}()
+	return session
+}
+
+func (s *H2MuxServerSession) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
+	s.active.Add(1)
+	defer s.active.Add(-1)
+	writer.WriteHeader(http.StatusOK)
+	conn := newHTTP2Wrapper(&v2rayhttp.ServerHTTPConn{
+		HTTP2Conn: v2rayhttp.NewHTTPConn(request.Body, writer),
+		Flusher:   writer.(http.Flusher),
+	})
+	s.inbound <- conn
+	select {
+	case <-conn.done:
+	case <-s.done:
+	}
+}
+
+func (s *H2MuxServerSession) Open() (net.Conn, error) {
+	return nil, os.ErrInvalid
+}
+
+func (s *H2MuxServerSession) Accept() (net.Conn, error) {
+	select {
+	case conn := <-s.inbound:
+		return conn, nil
+	case <-s.done:
+		return nil, os.ErrClosed
+	}
+}
+
+func (s *H2MuxServerSession) NumStreams() int {
+	return int(s.active.Load())
+}
+
+func (s *H2MuxServerSession) Close() error {
+	select {
+	case <-s.done:
+	default:
+		close(s.done)
+	}
+	return s.conn.Close()
+}
+
+func (s *H2MuxServerSession) IsClosed() bool {
+	select {
+	case <-s.done:
+		return true
+	default:
+		return false
+	}
+}
+
+func (s *H2MuxServerSession) CanTakeNewRequest() bool {
+	return false
+}
+
+type h2MuxConnWrapper struct {
+	N.ExtendedConn
+	done chan struct{}
+}
+
+func newHTTP2Wrapper(conn net.Conn) *h2MuxConnWrapper {
+	return &h2MuxConnWrapper{
+		ExtendedConn: bufio.NewExtendedConn(conn),
+		done:         make(chan struct{}),
+	}
+}
+
+func (w *h2MuxConnWrapper) Write(p []byte) (n int, err error) {
+	select {
+	case <-w.done:
+		return 0, net.ErrClosed
+	default:
+	}
+	return w.ExtendedConn.Write(p)
+}
+
+func (w *h2MuxConnWrapper) WriteBuffer(buffer *buf.Buffer) error {
+	select {
+	case <-w.done:
+		return net.ErrClosed
+	default:
+	}
+	return w.ExtendedConn.WriteBuffer(buffer)
+}
+
+func (w *h2MuxConnWrapper) Close() error {
+	select {
+	case <-w.done:
+	default:
+		close(w.done)
+	}
+	return w.ExtendedConn.Close()
+}
+
+func (w *h2MuxConnWrapper) Upstream() any {
+	return w.ExtendedConn
+}
+
+var _ abstractSession = (*H2MuxClientSession)(nil)
+
+type H2MuxClientSession struct {
+	transport  *http2.Transport
+	clientConn *http2.ClientConn
+	done       chan struct{}
+}
+
+func NewH2MuxClient(conn net.Conn) (*H2MuxClientSession, error) {
+	session := &H2MuxClientSession{
+		transport: &http2.Transport{
+			DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) {
+				return conn, nil
+			},
+			ReadIdleTimeout: idleTimeout,
+		},
+		done: make(chan struct{}),
+	}
+	session.transport.ConnPool = session
+	clientConn, err := session.transport.NewClientConn(conn)
+	if err != nil {
+		return nil, err
+	}
+	session.clientConn = clientConn
+	return session, nil
+}
+
+func (s *H2MuxClientSession) GetClientConn(req *http.Request, addr string) (*http2.ClientConn, error) {
+	return s.clientConn, nil
+}
+
+func (s *H2MuxClientSession) MarkDead(conn *http2.ClientConn) {
+	s.Close()
+}
+
+func (s *H2MuxClientSession) Open() (net.Conn, error) {
+	pipeInReader, pipeInWriter := io.Pipe()
+	request := &http.Request{
+		Method: http.MethodConnect,
+		Body:   pipeInReader,
+		URL:    &url.URL{Scheme: "https", Host: "localhost"},
+	}
+	conn := v2rayhttp.NewLateHTTPConn(pipeInWriter)
+	go func() {
+		response, err := s.transport.RoundTrip(request)
+		if err != nil {
+			conn.Setup(nil, err)
+		} else if response.StatusCode != 200 {
+			response.Body.Close()
+			conn.Setup(nil, E.New("unexpected status: ", response.StatusCode, " ", response.Status))
+		} else {
+			conn.Setup(response.Body, nil)
+		}
+	}()
+	return conn, nil
+}
+
+func (s *H2MuxClientSession) Accept() (net.Conn, error) {
+	return nil, os.ErrInvalid
+}
+
+func (s *H2MuxClientSession) NumStreams() int {
+	return s.clientConn.State().StreamsActive
+}
+
+func (s *H2MuxClientSession) Close() error {
+	select {
+	case <-s.done:
+	default:
+		close(s.done)
+	}
+	return s.clientConn.Close()
+}
+
+func (s *H2MuxClientSession) IsClosed() bool {
+	select {
+	case <-s.done:
+		return true
+	default:
+	}
+	return s.clientConn.State().Closed
+}
+
+func (s *H2MuxClientSession) CanTakeNewRequest() bool {
+	return s.clientConn.CanTakeNewRequest()
+}

common/mux/protocol.go

@@ -25,6 +25,7 @@ var Destination = M.Socksaddr{
 const (
 	ProtocolSMux Protocol = iota
 	ProtocolYAMux
+	ProtocolH2Mux
 )
 
 type Protocol byte
@@ -35,8 +36,10 @@ func ParseProtocol(name string) (Protocol, error) {
 		return ProtocolSMux, nil
 	case "yamux":
 		return ProtocolYAMux, nil
+	case "h2mux":
+		return ProtocolH2Mux, nil
 	default:
-		return ProtocolYAMux, E.New("unknown multiplex protocol: ", name)
+		return ProtocolSMux, E.New("unknown multiplex protocol: ", name)
 	}
 }
 
@@ -49,7 +52,13 @@ func (p Protocol) newServer(conn net.Conn) (abstractSession, error) {
 		}
 		return &smuxSession{session}, nil
 	case ProtocolYAMux:
-		return yamux.Server(conn, yaMuxConfig())
+		session, err := yamux.Server(conn, yaMuxConfig())
+		if err != nil {
+			return nil, err
+		}
+		return &yamuxSession{session}, nil
+	case ProtocolH2Mux:
+		return NewH2MuxServer(conn), nil
 	default:
 		panic("unknown protocol")
 	}
@@ -64,7 +73,13 @@ func (p Protocol) newClient(conn net.Conn) (abstractSession, error) {
 		}
 		return &smuxSession{session}, nil
 	case ProtocolYAMux:
-		return yamux.Client(conn, yaMuxConfig())
+		session, err := yamux.Client(conn, yaMuxConfig())
+		if err != nil {
+			return nil, err
+		}
+		return &yamuxSession{session}, nil
+	case ProtocolH2Mux:
+		return NewH2MuxClient(conn)
 	default:
 		panic("unknown protocol")
 	}
@@ -90,6 +105,8 @@ func (p Protocol) String() string {
 		return "smux"
 	case ProtocolYAMux:
 		return "yamux"
+	case ProtocolH2Mux:
+		return "h2mux"
 	default:
 		return "unknown"
 	}
@@ -115,9 +132,6 @@ func ReadRequest(reader io.Reader) (*Request, error) {
 	if err != nil {
 		return nil, err
 	}
-	if protocol > byte(ProtocolYAMux) {
-		return nil, E.New("unsupported protocol: ", protocol)
-	}
 	return &Request{Protocol: Protocol(protocol)}, nil
 }
 

common/mux/session.go

@@ -9,6 +9,8 @@ import (
 	"github.com/sagernet/sing/common/bufio"
 	N "github.com/sagernet/sing/common/network"
 	"github.com/sagernet/smux"
+
+	"github.com/hashicorp/yamux"
 )
 
 type abstractSession interface {
@@ -17,6 +19,7 @@ type abstractSession interface {
 	NumStreams() int
 	Close() error
 	IsClosed() bool
+	CanTakeNewRequest() bool
 }
 
 var _ abstractSession = (*smuxSession)(nil)
@@ -33,6 +36,18 @@ func (s *smuxSession) Accept() (net.Conn, error) {
 	return s.AcceptStream()
 }
 
+func (s *smuxSession) CanTakeNewRequest() bool {
+	return true
+}
+
+type yamuxSession struct {
+	*yamux.Session
+}
+
+func (y *yamuxSession) CanTakeNewRequest() bool {
+	return true
+}
+
 type protocolConn struct {
 	net.Conn
 	protocol        Protocol