From 2e346b468bda41f989295af7f62c8a8db1fa48b6 Mon Sep 17 00:00:00 2001
From: kaiburjack <kb@hbt.de>
Date: Mon, 14 Feb 2022 19:58:32 +0100
Subject: [PATCH] Use buffer slice in encoder and protect from stack
 over/underflow

---
 decoder.go           |  3 +++
 encoder.go           | 22 +++++++++-------------
 namespaceModifier.go | 27 ++++++++++++++++++++++-----
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/decoder.go b/decoder.go
index 73b73b6..6451d11 100644
--- a/decoder.go
+++ b/decoder.go
@@ -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]++
 		}
 	}
diff --git a/encoder.go b/encoder.go
index e60e52c..8e034f5 100644
--- a/encoder.go
+++ b/encoder.go
@@ -44,7 +44,7 @@ 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
@@ -52,9 +52,6 @@ type Encoder struct {
 	// 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.
@@ -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,
 	}
@@ -73,33 +71,31 @@ 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
 }
 
@@ -107,7 +103,7 @@ func (thiz *Encoder) writeBytes(bs []byte) error {
 // 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()
diff --git a/namespaceModifier.go b/namespaceModifier.go
index 26fa78a..e8c5cb2 100644
--- a/namespaceModifier.go
+++ b/namespaceModifier.go
@@ -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
@@ -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
 }
@@ -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,