diff --git a/go.mod b/go.mod index 9853842..0e36fb0 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/fumiama/WireGold go 1.18 require ( + github.com/FloatTech/zbputils v1.3.2 github.com/fumiama/blake2b-simd v0.0.0-20220412110131-4481822068bb github.com/fumiama/go-base16384 v1.4.0 github.com/fumiama/go-x25519 v1.0.0 diff --git a/go.sum b/go.sum index d21e2f0..78db297 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/FloatTech/zbputils v1.3.2 h1:TxQ6SjT/f+iHLoz+IlQ3Hk2dUb2P9SIs4DvjGWLjZYs= +github.com/FloatTech/zbputils v1.3.2/go.mod h1:V2Y8jkxjwACWuJsDCyNVKc+qABfMAUh/mgBYdWBxOPo= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/gold/link/me.go b/gold/link/me.go index 754f95f..a1ff0de 100644 --- a/gold/link/me.go +++ b/gold/link/me.go @@ -117,17 +117,17 @@ func (m *Me) Close() error { } func (m *Me) Write(packet []byte) (n int, err error) { - m.writer.Write(packet) - packet = m.writer.Bytes() + remain := m.writer.Len() + if remain > 0 { + m.writer.Write(packet) + packet = m.writer.Bytes() + } logrus.Debugln("[me] writer eating", len(packet), "bytes...") - n, packet = m.sendAllSameDst(packet) - if len(packet) > 0 { - w := helper.SelectWriter() - w.Write(packet) - helper.PutWriter(m.writer) - m.writer = w - logrus.Debugln("[me] writer remain", w.Len(), "bytes") - } else if n > 0 { + n = m.sendAllSameDst(packet) + if len(packet) > n { + _, _ = m.writer.Skip(remain + n - len(packet)) + logrus.Debugln("[me] writer remain", m.writer.Len(), "bytes") + } else if n > 0 && remain > 0 { m.writer.Reset() logrus.Debugln("[me] writer becomes empty") } @@ -149,8 +149,8 @@ func (p packetID) issame(packet []byte) bool { return p == waterutil.IPv4Identification(packet) } -func (m *Me) sendAllSameDst(packet []byte) (n int, rem []byte) { - rem = packet +func (m *Me) sendAllSameDst(packet []byte) (n int) { + rem := packet if !waterutil.IsIPv4(packet) { for len(rem) > 20 && waterutil.IsIPv6(rem) { pktl := int(binary.BigEndian.Uint16(packet[4:6])) + 40 @@ -163,7 +163,7 @@ func (m *Me) sendAllSameDst(packet []byte) (n int, rem []byte) { } if len(rem) == 0 || !waterutil.IsIPv4(rem) { logrus.Warnln("[me] skip to send", len(packet), "bytes full packet") - return len(packet), nil + return len(packet) } } p := newpacketid(rem) diff --git a/helper/writer.go b/helper/writer.go index f36a9fd..3de9218 100644 --- a/helper/writer.go +++ b/helper/writer.go @@ -6,6 +6,10 @@ import ( "bytes" "encoding/binary" "encoding/hex" + "io" + "unsafe" + + "github.com/FloatTech/zbputils/math" ) // Writer 写入 @@ -118,6 +122,50 @@ func (w *Writer) Grow(n int) { (*bytes.Buffer)(w).Grow(n) } +func (w *Writer) Skip(n int) (int, error) { + b := (*buffer)(unsafe.Pointer(w)) + b.lastRead = opInvalid + if len(b.buf) <= b.off { + // Buffer is empty, reset to recover space. + w.Reset() + if n == 0 { + return 0, nil + } + return 0, io.EOF + } + n = math.Min(n, len(b.buf[b.off:])) + b.off += n + if n > 0 { + b.lastRead = opRead + } + return n, nil +} + func (w *Writer) put() { PutWriter(w) } + +// A Buffer is a variable-sized buffer of bytes with Read and Write methods. +// The zero value for Buffer is an empty buffer ready to use. +type buffer struct { + buf []byte // contents are the bytes buf[off : len(buf)] + off int // read at &buf[off], write at &buf[len(buf)] + lastRead readOp // last read operation, so that Unread* can work correctly. +} + +// The readOp constants describe the last action performed on +// the buffer, so that UnreadRune and UnreadByte can check for +// invalid usage. opReadRuneX constants are chosen such that +// converted to int they correspond to the rune size that was read. +type readOp int8 + +// Don't use iota for these, as the values need to correspond with the +// names and comments, which is easier to see when being explicit. +const ( + opRead readOp = -1 // Any other read operation. + opInvalid readOp = 0 // Non-read operation. + opReadRune1 readOp = 1 // Read rune of size 1. + opReadRune2 readOp = 2 // Read rune of size 2. + opReadRune3 readOp = 3 // Read rune of size 3. + opReadRune4 readOp = 4 // Read rune of size 4. +)