Remove bytes.Buffer usage in pk.NBT()

This commit is contained in:
Tnze
2023-04-05 19:39:04 +08:00
parent a42267ba31
commit 91e23602a5

View File

@ -1,7 +1,6 @@
package packet package packet
import ( import (
"bytes"
"errors" "errors"
"io" "io"
"math" "math"
@ -480,23 +479,49 @@ type nbtField struct {
} }
func (n nbtField) WriteTo(w io.Writer) (int64, error) { func (n nbtField) WriteTo(w io.Writer) (int64, error) {
var buf bytes.Buffer
if n.V == nil { if n.V == nil {
buf.WriteByte(nbt.TagEnd) n, err := w.Write([]byte{nbt.TagEnd})
} else if err := nbt.NewEncoder(&buf).Encode(n.V, n.FieldName); err != nil { return int64(n), err
return 0, err
} }
return buf.WriteTo(w) // nbt Encode method does not count written bytes,
// so we warp the writer to count it.
cw := countingWriter{w: w}
err := nbt.NewEncoder(&cw).Encode(n.V, n.FieldName)
return cw.n, err
} }
func (n nbtField) ReadFrom(r io.Reader) (int64, error) { func (n nbtField) ReadFrom(r io.Reader) (int64, error) {
// LimitReader is used to count reader length // LimitReader is used to count reader length
lr := &io.LimitedReader{R: r, N: math.MaxInt64} cr := countingReader{r: r}
_, err := nbt.NewDecoder(lr).Decode(n.V) _, err := nbt.NewDecoder(&cr).Decode(n.V)
if err != nil && errors.Is(err, nbt.ErrEND) { if err != nil && errors.Is(err, nbt.ErrEND) {
err = nil err = nil
} }
return math.MaxInt64 - lr.N, err return cr.n, err
}
// countingWriter is a wrapper of io.Writer to externally count written bytes
type countingWriter struct {
n int64
w io.Writer
}
func (c *countingWriter) Write(p []byte) (n int, err error) {
n, err = c.w.Write(p)
c.n += int64(n)
return
}
// countingReader is a wrapper of io.Reader to externally count read bytes
type countingReader struct {
n int64
r io.Reader
}
func (c *countingReader) Read(p []byte) (n int, err error) {
n, err = c.r.Read(p)
c.n += int64(n)
return
} }
func (b ByteArray) WriteTo(w io.Writer) (n int64, err error) { func (b ByteArray) WriteTo(w io.Writer) (n int64, err error) {