cleanup buf interfaces

This commit is contained in:
Darien Raymond 2018-07-31 13:43:27 +02:00
parent b3cf1f70d7
commit 7baa6977d3
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
9 changed files with 28 additions and 32 deletions

View File

@ -40,13 +40,7 @@ func ReadFullFrom(reader io.Reader, size int32) Supplier {
}
}
// ReadAtLeastFrom create a Supplier to read at least size bytes from the given io.Reader.
func ReadAtLeastFrom(reader io.Reader, size int) Supplier {
return func(b []byte) (int, error) {
return io.ReadAtLeast(reader, b, size)
}
}
// WriteAllBytes ensures all bytes are written into the given writer.
func WriteAllBytes(writer io.Writer, payload []byte) error {
for len(payload) > 0 {
n, err := writer.Write(payload)
@ -89,10 +83,3 @@ func NewWriter(writer io.Writer) Writer {
Writer: writer,
}
}
// NewSequentialWriter returns a Writer that write Buffers in a MultiBuffer sequentially.
func NewSequentialWriter(writer io.Writer) Writer {
return &SequentialWriter{
Writer: writer,
}
}

View File

@ -171,10 +171,12 @@ func (w *BufferedWriter) Close() error {
return common.Close(w.writer)
}
// SequentialWriter is a Writer that writes MultiBuffer sequentially into the underlying io.Writer.
type SequentialWriter struct {
io.Writer
}
// WriteMultiBuffer implements Writer.
func (w *SequentialWriter) WriteMultiBuffer(mb MultiBuffer) error {
defer mb.Release()

View File

@ -73,14 +73,21 @@ func TestDiscardBytesMultiBuffer(t *testing.T) {
}
func TestWriterInterface(t *testing.T) {
assert := With(t)
{
var writer interface{} = (*BufferToBytesWriter)(nil)
switch writer.(type) {
case Writer, io.Writer, io.ReaderFrom:
default:
t.Error("BufferToBytesWriter is not Writer, io.Writer or io.ReaderFrom")
}
}
assert((*BufferToBytesWriter)(nil), Implements, (*Writer)(nil))
assert((*BufferToBytesWriter)(nil), Implements, (*io.Writer)(nil))
assert((*BufferToBytesWriter)(nil), Implements, (*io.ReaderFrom)(nil))
assert((*BufferedWriter)(nil), Implements, (*Writer)(nil))
assert((*BufferedWriter)(nil), Implements, (*io.Writer)(nil))
assert((*BufferedWriter)(nil), Implements, (*io.ReaderFrom)(nil))
assert((*BufferedWriter)(nil), Implements, (*io.ByteWriter)(nil))
{
var writer interface{} = (*BufferedWriter)(nil)
switch writer.(type) {
case Writer, io.Writer, io.ReaderFrom, io.ByteWriter:
default:
t.Error("BufferedWriter is not Writer, io.Writer, io.ReaderFrom or io.ByteWriter")
}
}
}

View File

@ -110,14 +110,14 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in
} else {
//if we are in TPROXY mode, use linux's udp forging functionality
if !d.config.FollowRedirect {
writer = buf.NewSequentialWriter(conn)
writer = &buf.SequentialWriter{Writer: conn}
} else {
srca := net.UDPAddr{IP: dest.Address.IP(), Port: int(dest.Port.Value())}
origsend, err := udp.TransmitSocket(&srca, conn.RemoteAddr())
if err != nil {
return err
}
writer = buf.NewSequentialWriter(origsend)
writer = &buf.SequentialWriter{Writer: origsend}
}
}

View File

@ -128,7 +128,7 @@ func (h *Handler) Process(ctx context.Context, link *core.Link, dialer proxy.Dia
if destination.Network == net.Network_TCP {
writer = buf.NewWriter(conn)
} else {
writer = buf.NewSequentialWriter(conn)
writer = &buf.SequentialWriter{Writer: conn}
}
if err := buf.Copy(input, writer, buf.UpdateActivity(timer)); err != nil {
return newError("failed to process request").Base(err)

View File

@ -133,10 +133,10 @@ func (c *Client) Process(ctx context.Context, link *core.Link, dialer proxy.Dial
if request.Command == protocol.RequestCommandUDP {
writer := buf.NewSequentialWriter(&UDPWriter{
writer := &buf.SequentialWriter{Writer: &UDPWriter{
Writer: conn,
Request: request,
})
}}
requestDone := func() error {
defer timer.SetTimeout(sessionPolicy.Timeouts.DownlinkOnly)

View File

@ -142,7 +142,7 @@ func TestUDPReaderWriter(t *testing.T) {
}),
}
cache := buf.New()
writer := buf.NewSequentialWriter(&UDPWriter{
writer := &buf.SequentialWriter{Writer: &UDPWriter{
Writer: cache,
Request: &protocol.RequestHeader{
Version: Version,
@ -151,7 +151,7 @@ func TestUDPReaderWriter(t *testing.T) {
User: user,
Option: RequestOptionOneTimeAuth,
},
})
}}
reader := &UDPReader{
Reader: cache,

View File

@ -123,7 +123,7 @@ func (c *Client) Process(ctx context.Context, link *core.Link, dialer proxy.Dial
defer udpConn.Close() // nolint: errcheck
requestFunc = func() error {
defer timer.SetTimeout(p.Timeouts.DownlinkOnly)
return buf.Copy(link.Reader, buf.NewSequentialWriter(NewUDPWriter(request, udpConn)), buf.UpdateActivity(timer))
return buf.Copy(link.Reader, &buf.SequentialWriter{Writer: NewUDPWriter(request, udpConn)}, buf.UpdateActivity(timer))
}
responseFunc = func() error {
defer timer.SetTimeout(p.Timeouts.UplinkOnly)

View File

@ -20,7 +20,7 @@ func TestUDPEncoding(t *testing.T) {
Address: net.IPAddress([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}),
Port: 1024,
}
writer := buf.NewSequentialWriter(NewUDPWriter(request, b))
writer := &buf.SequentialWriter{Writer: NewUDPWriter(request, b)}
content := []byte{'a'}
payload := buf.New()