552 lines
11 KiB
Go
552 lines
11 KiB
Go
package multiplex
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"encoding/binary"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
logging "github.com/ipfs/go-log"
|
|
pool "github.com/libp2p/go-buffer-pool"
|
|
"github.com/multiformats/go-varint"
|
|
)
|
|
|
|
var log = logging.Logger("mplex")
|
|
|
|
var MaxMessageSize = 1 << 20
|
|
|
|
// Max time to block waiting for a slow reader to read from a stream before
|
|
// resetting it. Preferably, we'd have some form of back-pressure mechanism but
|
|
// we don't have that in this protocol.
|
|
var ReceiveTimeout = 5 * time.Second
|
|
|
|
// ErrShutdown is returned when operating on a shutdown session
|
|
var ErrShutdown = errors.New("session shut down")
|
|
|
|
// ErrTwoInitiators is returned when both sides think they're the initiator
|
|
var ErrTwoInitiators = errors.New("two initiators")
|
|
|
|
// ErrInvalidState is returned when the other side does something it shouldn't.
|
|
// In this case, we close the connection to be safe.
|
|
var ErrInvalidState = errors.New("received an unexpected message from the peer")
|
|
|
|
var errTimeout = timeout{}
|
|
var errStreamClosed = errors.New("stream closed")
|
|
|
|
var (
|
|
ResetStreamTimeout = 2 * time.Minute
|
|
|
|
WriteCoalesceDelay = 100 * time.Microsecond
|
|
)
|
|
|
|
type timeout struct{}
|
|
|
|
func (_ timeout) Error() string {
|
|
return "i/o deadline exceeded"
|
|
}
|
|
|
|
func (_ timeout) Temporary() bool {
|
|
return true
|
|
}
|
|
|
|
func (_ timeout) Timeout() bool {
|
|
return true
|
|
}
|
|
|
|
// +1 for initiator
|
|
const (
|
|
newStreamTag = 0
|
|
messageTag = 2
|
|
closeTag = 4
|
|
resetTag = 6
|
|
)
|
|
|
|
// Multiplex is a mplex session.
|
|
type Multiplex struct {
|
|
con net.Conn
|
|
buf *bufio.Reader
|
|
nextID uint64
|
|
initiator bool
|
|
|
|
closed chan struct{}
|
|
shutdown chan struct{}
|
|
shutdownErr error
|
|
shutdownLock sync.Mutex
|
|
|
|
writeCh chan []byte
|
|
writeTimer *time.Timer
|
|
writeTimerFired bool
|
|
|
|
nstreams chan *Stream
|
|
|
|
channels map[streamID]*Stream
|
|
chLock sync.Mutex
|
|
}
|
|
|
|
// NewMultiplex creates a new multiplexer session.
|
|
func NewMultiplex(con net.Conn, initiator bool) *Multiplex {
|
|
mp := &Multiplex{
|
|
con: con,
|
|
initiator: initiator,
|
|
buf: bufio.NewReader(con),
|
|
channels: make(map[streamID]*Stream),
|
|
closed: make(chan struct{}),
|
|
shutdown: make(chan struct{}),
|
|
writeCh: make(chan []byte, 16),
|
|
writeTimer: time.NewTimer(0),
|
|
nstreams: make(chan *Stream, 16),
|
|
}
|
|
|
|
go mp.handleIncoming()
|
|
go mp.handleOutgoing()
|
|
|
|
return mp
|
|
}
|
|
|
|
func (mp *Multiplex) newStream(id streamID, name string) (s *Stream) {
|
|
s = &Stream{
|
|
id: id,
|
|
name: name,
|
|
dataIn: make(chan []byte, 8),
|
|
rDeadline: makePipeDeadline(),
|
|
wDeadline: makePipeDeadline(),
|
|
mp: mp,
|
|
writeCancel: make(chan struct{}),
|
|
readCancel: make(chan struct{}),
|
|
}
|
|
return
|
|
}
|
|
|
|
// Accept accepts the next stream from the connection.
|
|
func (m *Multiplex) Accept() (*Stream, error) {
|
|
select {
|
|
case s, ok := <-m.nstreams:
|
|
if !ok {
|
|
return nil, errors.New("multiplex closed")
|
|
}
|
|
return s, nil
|
|
case <-m.closed:
|
|
return nil, m.shutdownErr
|
|
}
|
|
}
|
|
|
|
// Close closes the session.
|
|
func (mp *Multiplex) Close() error {
|
|
mp.closeNoWait()
|
|
|
|
// Wait for the receive loop to finish.
|
|
<-mp.closed
|
|
|
|
return nil
|
|
}
|
|
|
|
func (mp *Multiplex) closeNoWait() {
|
|
mp.shutdownLock.Lock()
|
|
select {
|
|
case <-mp.shutdown:
|
|
default:
|
|
mp.con.Close()
|
|
close(mp.shutdown)
|
|
}
|
|
mp.shutdownLock.Unlock()
|
|
}
|
|
|
|
// IsClosed returns true if the session is closed.
|
|
func (mp *Multiplex) IsClosed() bool {
|
|
select {
|
|
case <-mp.closed:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) sendMsg(timeout, cancel <-chan struct{}, header uint64, data []byte) error {
|
|
buf := pool.Get(len(data) + 20)
|
|
|
|
n := 0
|
|
n += binary.PutUvarint(buf[n:], header)
|
|
n += binary.PutUvarint(buf[n:], uint64(len(data)))
|
|
n += copy(buf[n:], data)
|
|
|
|
select {
|
|
case mp.writeCh <- buf[:n]:
|
|
return nil
|
|
case <-mp.shutdown:
|
|
return ErrShutdown
|
|
case <-timeout:
|
|
return errTimeout
|
|
case <-cancel:
|
|
return ErrStreamClosed
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) handleOutgoing() {
|
|
for {
|
|
select {
|
|
case <-mp.shutdown:
|
|
return
|
|
|
|
case data := <-mp.writeCh:
|
|
// FIXME: https://github.com/libp2p/go-libp2p/issues/644
|
|
// write coalescing disabled until this can be fixed.
|
|
//err := mp.writeMsg(data)
|
|
err := mp.doWriteMsg(data)
|
|
pool.Put(data)
|
|
if err != nil {
|
|
// the connection is closed by this time
|
|
log.Warnf("error writing data: %s", err.Error())
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) writeMsg(data []byte) error {
|
|
if len(data) >= 512 {
|
|
err := mp.doWriteMsg(data)
|
|
pool.Put(data)
|
|
return err
|
|
}
|
|
|
|
buf := pool.Get(4096)
|
|
defer pool.Put(buf)
|
|
|
|
n := copy(buf, data)
|
|
pool.Put(data)
|
|
|
|
if !mp.writeTimerFired {
|
|
if !mp.writeTimer.Stop() {
|
|
<-mp.writeTimer.C
|
|
}
|
|
}
|
|
mp.writeTimer.Reset(WriteCoalesceDelay)
|
|
mp.writeTimerFired = false
|
|
|
|
for {
|
|
select {
|
|
case data = <-mp.writeCh:
|
|
wr := copy(buf[n:], data)
|
|
if wr < len(data) {
|
|
// we filled the buffer, send it
|
|
err := mp.doWriteMsg(buf)
|
|
if err != nil {
|
|
pool.Put(data)
|
|
return err
|
|
}
|
|
|
|
if len(data)-wr >= 512 {
|
|
// the remaining data is not a small write, send it
|
|
err := mp.doWriteMsg(data[wr:])
|
|
pool.Put(data)
|
|
return err
|
|
}
|
|
|
|
n = copy(buf, data[wr:])
|
|
|
|
// we've written some, reset the timer to coalesce the rest
|
|
if !mp.writeTimer.Stop() {
|
|
<-mp.writeTimer.C
|
|
}
|
|
mp.writeTimer.Reset(WriteCoalesceDelay)
|
|
} else {
|
|
n += wr
|
|
}
|
|
|
|
pool.Put(data)
|
|
|
|
case <-mp.writeTimer.C:
|
|
mp.writeTimerFired = true
|
|
return mp.doWriteMsg(buf[:n])
|
|
|
|
case <-mp.shutdown:
|
|
return ErrShutdown
|
|
}
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) doWriteMsg(data []byte) error {
|
|
if mp.isShutdown() {
|
|
return ErrShutdown
|
|
}
|
|
|
|
_, err := mp.con.Write(data)
|
|
if err != nil {
|
|
mp.closeNoWait()
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
func (mp *Multiplex) nextChanID() uint64 {
|
|
out := mp.nextID
|
|
mp.nextID++
|
|
return out
|
|
}
|
|
|
|
// NewStream creates a new stream.
|
|
func (mp *Multiplex) NewStream(ctx context.Context) (*Stream, error) {
|
|
return mp.NewNamedStream(ctx, "")
|
|
}
|
|
|
|
// NewNamedStream creates a new named stream.
|
|
func (mp *Multiplex) NewNamedStream(ctx context.Context, name string) (*Stream, error) {
|
|
mp.chLock.Lock()
|
|
|
|
// We could call IsClosed but this is faster (given that we already have
|
|
// the lock).
|
|
if mp.channels == nil {
|
|
mp.chLock.Unlock()
|
|
return nil, ErrShutdown
|
|
}
|
|
|
|
sid := mp.nextChanID()
|
|
header := (sid << 3) | newStreamTag
|
|
|
|
if name == "" {
|
|
name = fmt.Sprint(sid)
|
|
}
|
|
s := mp.newStream(streamID{
|
|
id: sid,
|
|
initiator: true,
|
|
}, name)
|
|
mp.channels[s.id] = s
|
|
mp.chLock.Unlock()
|
|
|
|
err := mp.sendMsg(ctx.Done(), nil, header, []byte(name))
|
|
if err != nil {
|
|
if err == errTimeout {
|
|
return nil, ctx.Err()
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
return s, nil
|
|
}
|
|
|
|
func (mp *Multiplex) cleanup() {
|
|
mp.closeNoWait()
|
|
|
|
// Take the channels.
|
|
mp.chLock.Lock()
|
|
channels := mp.channels
|
|
mp.channels = nil
|
|
mp.chLock.Unlock()
|
|
|
|
// Cancel any reads/writes
|
|
for _, msch := range channels {
|
|
msch.cancelRead(ErrStreamReset)
|
|
msch.cancelWrite(ErrStreamReset)
|
|
}
|
|
|
|
// And... shutdown!
|
|
if mp.shutdownErr == nil {
|
|
mp.shutdownErr = ErrShutdown
|
|
}
|
|
close(mp.closed)
|
|
}
|
|
|
|
func (mp *Multiplex) handleIncoming() {
|
|
defer mp.cleanup()
|
|
|
|
recvTimeout := time.NewTimer(0)
|
|
defer recvTimeout.Stop()
|
|
|
|
if !recvTimeout.Stop() {
|
|
<-recvTimeout.C
|
|
}
|
|
|
|
for {
|
|
chID, tag, err := mp.readNextHeader()
|
|
if err != nil {
|
|
mp.shutdownErr = err
|
|
return
|
|
}
|
|
|
|
remoteIsInitiator := tag&1 == 0
|
|
ch := streamID{
|
|
// true if *I'm* the initiator.
|
|
initiator: !remoteIsInitiator,
|
|
id: chID,
|
|
}
|
|
// Rounds up the tag:
|
|
// 0 -> 0
|
|
// 1 -> 2
|
|
// 2 -> 2
|
|
// 3 -> 4
|
|
// etc...
|
|
tag += (tag & 1)
|
|
|
|
b, err := mp.readNext()
|
|
if err != nil {
|
|
mp.shutdownErr = err
|
|
return
|
|
}
|
|
|
|
mp.chLock.Lock()
|
|
msch, ok := mp.channels[ch]
|
|
mp.chLock.Unlock()
|
|
|
|
switch tag {
|
|
case newStreamTag:
|
|
if ok {
|
|
log.Debugf("received NewStream message for existing stream: %d", ch)
|
|
mp.shutdownErr = ErrInvalidState
|
|
return
|
|
}
|
|
|
|
name := string(b)
|
|
pool.Put(b)
|
|
|
|
msch = mp.newStream(ch, name)
|
|
mp.chLock.Lock()
|
|
mp.channels[ch] = msch
|
|
mp.chLock.Unlock()
|
|
select {
|
|
case mp.nstreams <- msch:
|
|
case <-mp.shutdown:
|
|
return
|
|
}
|
|
|
|
case resetTag:
|
|
if !ok {
|
|
// This is *ok*. We forget the stream on reset.
|
|
continue
|
|
}
|
|
|
|
// Cancel any ongoing reads/writes.
|
|
msch.cancelRead(ErrStreamReset)
|
|
msch.cancelWrite(ErrStreamReset)
|
|
case closeTag:
|
|
if !ok {
|
|
// may have canceled our reads already.
|
|
continue
|
|
}
|
|
|
|
// unregister and throw away future data.
|
|
mp.chLock.Lock()
|
|
delete(mp.channels, ch)
|
|
mp.chLock.Unlock()
|
|
|
|
// close data channel, there will be no more data.
|
|
close(msch.dataIn)
|
|
|
|
// We intentionally don't cancel any deadlines, cancel reads, cancel
|
|
// writes, etc. We just deliver the EOF by closing the
|
|
// data channel, and unregister the channel so we don't
|
|
// receive any more data. The user still needs to call
|
|
// `Close()` or `Reset()`.
|
|
case messageTag:
|
|
if !ok {
|
|
// We're not accepting data on this stream, for
|
|
// some reason. It's likely that we reset it, or
|
|
// simply canceled reads (e.g., called Close).
|
|
pool.Put(b)
|
|
continue
|
|
}
|
|
|
|
recvTimeout.Reset(ReceiveTimeout)
|
|
select {
|
|
case msch.dataIn <- b:
|
|
case <-msch.readCancel:
|
|
// the user has canceled reading. walk away.
|
|
pool.Put(b)
|
|
case <-recvTimeout.C:
|
|
pool.Put(b)
|
|
log.Warnf("timed out receiving message into stream queue.")
|
|
// Do not do this asynchronously. Otherwise, we
|
|
// could drop a message, then receive a message,
|
|
// then reset.
|
|
msch.Reset()
|
|
continue
|
|
case <-mp.shutdown:
|
|
pool.Put(b)
|
|
return
|
|
}
|
|
if !recvTimeout.Stop() {
|
|
<-recvTimeout.C
|
|
}
|
|
default:
|
|
log.Debugf("message with unknown header on stream %s", ch)
|
|
if ok {
|
|
msch.Reset()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) isShutdown() bool {
|
|
select {
|
|
case <-mp.shutdown:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) sendResetMsg(header uint64, hard bool) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), ResetStreamTimeout)
|
|
defer cancel()
|
|
|
|
err := mp.sendMsg(ctx.Done(), nil, header, nil)
|
|
if err != nil && !mp.isShutdown() {
|
|
if hard {
|
|
log.Warnf("error sending reset message: %s; killing connection", err.Error())
|
|
mp.Close()
|
|
} else {
|
|
log.Debugf("error sending reset message: %s", err.Error())
|
|
}
|
|
}
|
|
}
|
|
|
|
func (mp *Multiplex) readNextHeader() (uint64, uint64, error) {
|
|
h, err := varint.ReadUvarint(mp.buf)
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
|
|
// get channel ID
|
|
ch := h >> 3
|
|
|
|
rem := h & 7
|
|
|
|
return ch, rem, nil
|
|
}
|
|
|
|
func (mp *Multiplex) readNext() ([]byte, error) {
|
|
// get length
|
|
l, err := varint.ReadUvarint(mp.buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if l > uint64(MaxMessageSize) {
|
|
return nil, fmt.Errorf("message size too large!")
|
|
}
|
|
|
|
if l == 0 {
|
|
return nil, nil
|
|
}
|
|
|
|
buf := pool.Get(int(l))
|
|
n, err := io.ReadFull(mp.buf, buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return buf[:n], nil
|
|
}
|
|
|
|
func isFatalNetworkError(err error) bool {
|
|
nerr, ok := err.(net.Error)
|
|
if ok {
|
|
return !(nerr.Timeout() || nerr.Temporary())
|
|
}
|
|
return false
|
|
}
|