go-libp2p-core/helpers/stream.go

57 lines
1.5 KiB
Go

package helpers
import (
"errors"
"io"
"time"
"github.com/libp2p/go-libp2p-core/network"
)
// EOFTimeout is the maximum amount of time to wait to successfully observe an
// EOF on the stream. Defaults to 60 seconds.
var EOFTimeout = time.Second * 60
// ErrExpectedEOF is returned when we read data while expecting an EOF.
var ErrExpectedEOF = errors.New("read data when expecting EOF")
// FullClose closes the stream and waits to read an EOF from the other side.
//
// * If it reads any data *before* the EOF, it resets the stream.
// * If it doesn't read an EOF within EOFTimeout, it resets the stream.
//
// You'll likely want to invoke this as `go FullClose(stream)` to close the
// stream in the background.
func FullClose(s network.Stream) error {
if err := s.Close(); err != nil {
s.Reset()
return err
}
return AwaitEOF(s)
}
// AwaitEOF waits for an EOF on the given stream, returning an error if that
// fails. It waits at most EOFTimeout (defaults to 1 minute) after which it
// resets the stream.
func AwaitEOF(s network.Stream) error {
// So we don't wait forever
s.SetDeadline(time.Now().Add(EOFTimeout))
// We *have* to observe the EOF. Otherwise, we leak the stream.
// Now, technically, we should do this *before*
// returning from SendMessage as the message
// hasn't really been sent yet until we see the
// EOF but we don't actually *know* what
// protocol the other side is speaking.
n, err := s.Read([]byte{0})
if n > 0 || err == nil {
s.Reset()
return ErrExpectedEOF
}
if err != io.EOF {
s.Reset()
return err
}
return nil
}