Skip to content

Commit

Permalink
Use buffer slice in encoder and protect from stack over/underflow
Browse files Browse the repository at this point in the history
  • Loading branch information
kaiburjack committed Feb 14, 2022
1 parent 80c1b49 commit 2e346b4
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 18 deletions.
3 changes: 3 additions & 0 deletions decoder.go
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,9 @@ func (thiz *decoder) decodeAttributes(b byte) ([]Attr, error) {
return nil, err
}
b, err = thiz.readByte()
if err != nil {
return nil, err
}
thiz.numAttributes[thiz.top]++
}
}
Expand Down
22 changes: 9 additions & 13 deletions encoder.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,17 +44,14 @@ type EncoderMiddleware interface {
// Encoder encodes Token values to an io.Writer.
type Encoder struct {
// buffers writes to the underlying io.Writer
buf [2048]byte
buf []byte

// middlewares can modify encoded tokens before encoding.
middlewares []EncoderMiddleware

// The io.Writer we encode/write into.
wr io.Writer

// the current write position into buf
w int

// Whether the last token was of type TokenTypeStartElement.
// This is used to delay encoding the ending ">" or "/>" string
// based on whether the element is immediately closed afterwards.
Expand All @@ -64,6 +61,7 @@ type Encoder struct {
// NewEncoder creates a new Encoder with the given middlewares and returns a pointer to it.
func NewEncoder(w io.Writer, middlewares ...EncoderMiddleware) *Encoder {
return &Encoder{
buf: make([]byte, 0, 2048),
wr: w,
middlewares: middlewares,
}
Expand All @@ -73,41 +71,39 @@ func NewEncoder(w io.Writer, middlewares ...EncoderMiddleware) *Encoder {
// It must be called after token encoding is done in order
// to write all remaining bytes into the io.Writer.
func (thiz *Encoder) Flush() error {
_, err := thiz.wr.Write(thiz.buf[:thiz.w])
thiz.w = 0
_, err := thiz.wr.Write(thiz.buf)
thiz.buf = thiz.buf[:0]
return err
}

func (thiz *Encoder) write(b byte) error {
if thiz.w >= len(thiz.buf) {
if len(thiz.buf)+1 >= cap(thiz.buf) {
err := thiz.Flush()
if err != nil {
return err
}
}
thiz.buf[thiz.w] = b
thiz.w++
thiz.buf = append(thiz.buf, b)
return nil
}

func (thiz *Encoder) writeBytes(bs []byte) error {
l := len(bs)
if thiz.w+l > len(thiz.buf) {
if len(thiz.buf)+l >= cap(thiz.buf) {
err := thiz.Flush()
if err != nil {
return err
}
}
copy(thiz.buf[thiz.w:], bs)
thiz.w += l
thiz.buf = append(thiz.buf, bs...)
return nil
}

// Reset resets this Encoder to write into the provided io.Writer
// and resets all middlewares.
func (thiz *Encoder) Reset(w io.Writer) {
thiz.wr = w
thiz.w = 0
thiz.buf = thiz.buf[:0]
thiz.lastStartElement = false
for _, middleware := range thiz.middlewares {
middleware.Reset()
Expand Down
27 changes: 22 additions & 5 deletions namespaceModifier.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
package gosaxml

import "bytes"
import (
"bytes"
"errors"
)

// NamespaceModifier can be used to obtain information about the
// effective namespace of a decoded Token via NamespaceOfToken
Expand Down Expand Up @@ -44,13 +47,19 @@ func (thiz *NamespaceModifier) Reset() {
// by the EncoderMiddleware.
func (thiz *NamespaceModifier) EncodeToken(t *Token) error {
if t.Kind == TokenTypeStartElement {
thiz.pushFrame()
err := thiz.pushFrame()
if err != nil {
return err
}
thiz.processNamespaces(t)
thiz.processElementName(t)
thiz.openNames[thiz.top] = t.Name
} else if t.Kind == TokenTypeEndElement {
thiz.processElementName(t)
thiz.popFrame()
err := thiz.popFrame()
if err != nil {
return err
}
}
return nil
}
Expand Down Expand Up @@ -115,18 +124,26 @@ func (thiz *NamespaceModifier) findPrefixAlias(prefix []byte) []byte {
return nil
}

func (thiz *NamespaceModifier) pushFrame() {
func (thiz *NamespaceModifier) pushFrame() error {
if thiz.top >= 255 {
return errors.New("stack overflow")
}
thiz.top++
thiz.nsOffs[thiz.top] = thiz.nsOffs[thiz.top-1]
thiz.prefixAliasesOffs[thiz.top] = thiz.prefixAliasesOffs[thiz.top-1]
return nil
}

func (thiz *NamespaceModifier) popFrame() {
func (thiz *NamespaceModifier) popFrame() error {
if thiz.top <= 0 {
return errors.New("stack underflow")
}
thiz.top--
off := thiz.nsOffs[thiz.top]
thiz.namespaces = thiz.namespaces[:off*2]
off = thiz.prefixAliasesOffs[thiz.top]
thiz.prefixAliases = thiz.prefixAliases[:off*2]
return nil
}

// processNamespaces scans the attributes of the given token for namespace declarations,
Expand Down

0 comments on commit 2e346b4

Please sign in to comment.