sing-box/common/mux/session.go

92 lines
1.9 KiB
Go
Raw Normal View History

2022-08-03 13:51:34 +00:00
package mux
import (
"io"
"net"
"github.com/sagernet/sing/common"
"github.com/sagernet/sing/common/buf"
"github.com/sagernet/sing/common/bufio"
2022-08-12 08:49:25 +00:00
N "github.com/sagernet/sing/common/network"
2022-08-12 04:13:57 +00:00
"github.com/sagernet/smux"
2022-08-03 13:51:34 +00:00
)
type abstractSession interface {
Open() (net.Conn, error)
Accept() (net.Conn, error)
NumStreams() int
Close() error
IsClosed() bool
}
var _ abstractSession = (*smuxSession)(nil)
type smuxSession struct {
*smux.Session
}
func (s *smuxSession) Open() (net.Conn, error) {
return s.OpenStream()
}
func (s *smuxSession) Accept() (net.Conn, error) {
return s.AcceptStream()
}
type protocolConn struct {
net.Conn
protocol Protocol
protocolWritten bool
}
func (c *protocolConn) Write(p []byte) (n int, err error) {
if c.protocolWritten {
return c.Conn.Write(p)
}
_buffer := buf.StackNewSize(2 + len(p))
defer common.KeepAlive(_buffer)
buffer := common.Dup(_buffer)
defer buffer.Release()
EncodeRequest(buffer, Request{
Protocol: c.protocol,
})
common.Must(common.Error(buffer.Write(p)))
n, err = c.Conn.Write(buffer.Bytes())
if err == nil {
n--
}
c.protocolWritten = true
return n, err
}
func (c *protocolConn) ReadFrom(r io.Reader) (n int64, err error) {
if !c.protocolWritten {
return bufio.ReadFrom0(c, r)
}
return bufio.Copy(c.Conn, r)
}
func (c *protocolConn) Upstream() any {
return c.Conn
}
2022-08-12 08:49:25 +00:00
type vectorisedProtocolConn struct {
protocolConn
N.VectorisedWriter
}
func (c *vectorisedProtocolConn) WriteVectorised(buffers []*buf.Buffer) error {
if c.protocolWritten {
return c.VectorisedWriter.WriteVectorised(buffers)
}
c.protocolWritten = true
_buffer := buf.StackNewSize(2)
defer common.KeepAlive(_buffer)
buffer := common.Dup(_buffer)
defer buffer.Release()
EncodeRequest(buffer, Request{
Protocol: c.protocol,
})
return c.VectorisedWriter.WriteVectorised(append([]*buf.Buffer{buffer}, buffers...))
}