Cloak/internal/multiplex/datagramBufferedPipe.go

163 lines
3.4 KiB
Go
Raw Normal View History

2019-08-14 09:04:27 +00:00
// This is base on https://github.com/golang/go/blob/0436b162397018c45068b47ca1b5924a3eafdee0/src/net/net_fake.go#L173
package multiplex
import (
2020-04-12 10:18:41 +00:00
"bytes"
2019-08-14 09:04:27 +00:00
"io"
"sync"
2020-04-09 17:56:17 +00:00
"time"
2019-08-14 09:04:27 +00:00
)
// datagramBufferedPipe is the same as streamBufferedPipe with the exception that it's message-oriented,
2019-08-20 21:43:04 +00:00
// instead of byte-oriented. The integrity of datagrams written into this buffer is preserved.
// it won't get chopped up into individual bytes
type datagramBufferedPipe struct {
pLens []int
// lazily allocated
2020-04-12 10:18:41 +00:00
buf *bytes.Buffer
closed bool
2020-04-09 17:56:17 +00:00
rwCond *sync.Cond
wtTimeout time.Duration
2020-04-09 17:56:17 +00:00
rDeadline time.Time
2019-08-14 09:04:27 +00:00
}
func NewDatagramBufferedPipe() *datagramBufferedPipe {
d := &datagramBufferedPipe{
2019-08-14 09:04:27 +00:00
rwCond: sync.NewCond(&sync.Mutex{}),
}
return d
}
func (d *datagramBufferedPipe) Read(target []byte) (int, error) {
2019-08-14 09:04:27 +00:00
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
2020-04-12 10:18:41 +00:00
if d.buf == nil {
d.buf = new(bytes.Buffer)
}
2019-08-14 09:04:27 +00:00
for {
if d.closed && len(d.pLens) == 0 {
2019-08-14 09:04:27 +00:00
return 0, io.EOF
}
2020-04-09 17:56:17 +00:00
if !d.rDeadline.IsZero() {
delta := time.Until(d.rDeadline)
if delta <= 0 {
return 0, ErrTimeout
}
time.AfterFunc(delta, d.rwCond.Broadcast)
}
2020-04-12 10:18:41 +00:00
if len(d.pLens) > 0 {
2019-08-14 09:04:27 +00:00
break
}
d.rwCond.Wait()
}
2020-04-12 10:18:41 +00:00
dataLen := d.pLens[0]
if len(target) < dataLen {
2020-04-08 13:59:09 +00:00
return 0, io.ErrShortBuffer
}
2020-04-12 10:18:41 +00:00
d.pLens = d.pLens[1:]
d.buf.Read(target[:dataLen])
2019-08-14 09:04:27 +00:00
// err will always be nil because we have already verified that buf.Len() != 0
d.rwCond.Broadcast()
2020-04-12 10:18:41 +00:00
return dataLen, nil
2019-08-14 09:04:27 +00:00
}
func (d *datagramBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
2020-04-12 10:33:11 +00:00
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
if d.buf == nil {
d.buf = new(bytes.Buffer)
}
for {
if d.closed && len(d.pLens) == 0 {
2020-04-12 10:33:11 +00:00
return 0, io.EOF
}
if !d.rDeadline.IsZero() {
delta := time.Until(d.rDeadline)
if delta <= 0 {
return 0, ErrTimeout
}
if d.wtTimeout == 0 {
// if there hasn't been a scheduled broadcast
time.AfterFunc(delta, d.rwCond.Broadcast)
}
}
if d.wtTimeout != 0 {
d.rDeadline = time.Now().Add(d.wtTimeout)
time.AfterFunc(d.wtTimeout, d.rwCond.Broadcast)
2020-04-12 10:33:11 +00:00
}
if len(d.pLens) > 0 {
var dataLen int
dataLen, d.pLens = d.pLens[0], d.pLens[1:]
written, er := w.Write(d.buf.Next(dataLen))
n += int64(written)
if er != nil {
d.rwCond.Broadcast()
return n, er
}
d.rwCond.Broadcast()
} else {
d.rwCond.Wait()
2020-04-12 10:33:11 +00:00
}
}
}
func (d *datagramBufferedPipe) Write(f Frame) (toBeClosed bool, err error) {
2019-08-14 09:04:27 +00:00
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
2020-04-12 10:18:41 +00:00
if d.buf == nil {
d.buf = new(bytes.Buffer)
}
2019-08-14 09:04:27 +00:00
for {
if d.closed {
return true, io.ErrClosedPipe
2019-08-14 09:04:27 +00:00
}
if d.buf.Len() <= recvBufferSizeLimit {
2019-08-14 09:04:27 +00:00
// if d.buf gets too large, write() will panic. We don't want this to happen
break
}
d.rwCond.Wait()
}
if f.Closing != C_NOOP {
d.closed = true
d.rwCond.Broadcast()
return true, nil
}
2020-04-12 10:18:41 +00:00
dataLen := len(f.Payload)
d.pLens = append(d.pLens, dataLen)
d.buf.Write(f.Payload)
2019-08-14 09:04:27 +00:00
// err will always be nil
d.rwCond.Broadcast()
return false, nil
2019-08-14 09:04:27 +00:00
}
func (d *datagramBufferedPipe) Close() error {
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
d.closed = true
2019-08-14 09:04:27 +00:00
d.rwCond.Broadcast()
return nil
}
2020-04-09 17:56:17 +00:00
func (d *datagramBufferedPipe) SetReadDeadline(t time.Time) {
2020-04-09 17:56:17 +00:00
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
d.rDeadline = t
d.rwCond.Broadcast()
}
func (d *datagramBufferedPipe) SetWriteToTimeout(t time.Duration) {
d.rwCond.L.Lock()
defer d.rwCond.L.Unlock()
d.wtTimeout = t
d.rwCond.Broadcast()
}