Refactor indefinite-length encoding and improve chunk validation during encoding

stream.go

@@ -180,9 +180,10 @@ func (e *IndefiniteLengthMapOddItemCountError) Error() string {
 
 // Encoder writes CBOR values to io.Writer.
 type Encoder struct {
-	w      io.Writer
-	em     *encMode
-	indefs []indefDataItem
+	w       io.Writer
+	em      *encMode
+	indefs  []indefDataItem
+	scratch [1]byte // reused for single-byte writes (indefinite-length head and break code)
 }
 
 // NewEncoder returns a new encoder that writes to w using the default encoding options.
@@ -193,40 +194,39 @@ func NewEncoder(w io.Writer) *Encoder {
 // Encode writes the CBOR encoding of v.
 func (enc *Encoder) Encode(v any) error {
 	if len(enc.indefs) > 0 {
-		switch enc.indefs[len(enc.indefs)-1].typ {
-		case cborTypeTextString:
-			if v == nil {
-				return errors.New("cbor: cannot encode nil for indefinite-length text string")
-			}
-			k := reflect.TypeOf(v).Kind()
-			if k != reflect.String {
-				return errors.New("cbor: cannot encode item type " + k.String() + " for indefinite-length text string")
-			}
-		case cborTypeByteString:
-			if v == nil {
-				return errors.New("cbor: cannot encode nil for indefinite-length byte string")
-			}
-			t := reflect.TypeOf(v)
-			k := t.Kind()
-			if (k != reflect.Array && k != reflect.Slice) || t.Elem().Kind() != reflect.Uint8 {
-				return errors.New("cbor: cannot encode item type " + k.String() + " for indefinite-length byte string")
-			}
+		err := validateIndefiniteLengthChunkByType(enc.indefs[len(enc.indefs)-1].typ, v)
+		if err != nil {
+			return err
 		}
 	}
 
 	buf := getEncodeBuffer()
 
 	err := encode(buf, enc.em, reflect.ValueOf(v))
+
+	// Validate the encoded chunk against the indefinite-length data item using a byte-based check.
+	// This reliably detects chunks from cbor.Marshaler, registered tags, StringToByteString, etc.,
+	// which may produce a chunk inconsistent with the parent's major type.
+	// Applies only to indefinite-length byte/text string parents (RFC 8949 Section 3.2.3).
+	if err == nil && len(enc.indefs) > 0 {
+		err = validateIndefiniteLengthChunkByData(enc.indefs[len(enc.indefs)-1].typ, buf.Bytes(), v)
+	}
+
 	if err == nil {
 		_, err = enc.w.Write(buf.Bytes())
 	}
 
 	putEncodeBuffer(buf)
 
-	if err == nil && len(enc.indefs) > 0 {
+	if err != nil {
+		return err
+	}
+
+	if len(enc.indefs) > 0 {
 		enc.indefs[len(enc.indefs)-1].count++
 	}
-	return err
+
+	return nil
 }
 
 // StartIndefiniteByteString starts indefinite-length byte string encoding.
@@ -266,49 +266,93 @@ func (enc *Encoder) EndIndefinite() error {
 	if len(enc.indefs) == 0 {
 		return errors.New("cbor: cannot encode \"break\" code outside indefinite length values")
 	}
+
+	// Verify that indefinite-length map has even number of elements
 	top := enc.indefs[len(enc.indefs)-1]
 	if top.typ == cborTypeMap && top.count%2 != 0 {
 		return &IndefiniteLengthMapOddItemCountError{count: top.count}
 	}
-	_, err := enc.w.Write([]byte{cborBreakFlag})
-	if err == nil {
-		enc.indefs = enc.indefs[:len(enc.indefs)-1]
-		// Increment parent container's item count because the child
-		// (indefinite-length data item) is fully written to the stream.
-		if len(enc.indefs) > 0 {
-			enc.indefs[len(enc.indefs)-1].count++
-		}
+
+	// Write break code
+	enc.scratch[0] = cborBreakFlag
+	_, err := enc.w.Write(enc.scratch[:])
+	if err != nil {
+		return err
 	}
-	return err
+
+	enc.indefs = enc.indefs[:len(enc.indefs)-1]
+
+	// Increment parent container's item count because the child
+	// (indefinite-length data item) is fully written to the stream.
+	if len(enc.indefs) > 0 {
+		enc.indefs[len(enc.indefs)-1].count++
+	}
+
+	return nil
 }
 
 func (enc *Encoder) startIndefinite(typ cborType) error {
 	if enc.em.indefLength == IndefLengthForbidden {
 		return &IndefiniteLengthError{typ}
 	}
+
+	// Verify that new indefinite-length data item is not a chunk in indefinite-length byte/text string.
 	if len(enc.indefs) > 0 {
 		parent := enc.indefs[len(enc.indefs)-1].typ
 		if parent == cborTypeByteString || parent == cborTypeTextString {
 			return errors.New("cbor: cannot encode indefinite-length " + typ.String() +
 				" as chunk of indefinite-length " + parent.String())
 		}
 	}
-	var head byte
-	switch typ {
-	case cborTypeByteString:
-		head = cborByteStringWithIndefiniteLengthHead
-	case cborTypeTextString:
-		head = cborTextStringWithIndefiniteLengthHead
-	case cborTypeArray:
-		head = cborArrayWithIndefiniteLengthHead
-	case cborTypeMap:
-		head = cborMapWithIndefiniteLengthHead
+
+	// Write indefinite-length head.
+	enc.scratch[0] = byte(typ) | additionalInformationAsIndefiniteLengthFlag
+	_, err := enc.w.Write(enc.scratch[:])
+	if err != nil {
+		return err
 	}
-	_, err := enc.w.Write([]byte{head})
-	if err == nil {
-		enc.indefs = append(enc.indefs, indefDataItem{typ: typ})
+
+	enc.indefs = append(enc.indefs, indefDataItem{typ: typ})
+	return nil
+}
+
+// validateIndefiniteLengthChunkByType rejects chunks based solely on their Go type.
+func validateIndefiniteLengthChunkByType(indefiniteLengthCborType cborType, v any) error {
+	if indefiniteLengthCborType != cborTypeByteString &&
+		indefiniteLengthCborType != cborTypeTextString {
+		return nil
 	}
-	return err
+	if v == nil {
+		return errors.New("cbor: cannot encode nil for indefinite-length " + indefiniteLengthCborType.String())
+	}
+	return nil
+}
+
+// validateIndefiniteLengthChunkByData checks that chunk is a definite-length
+// CBOR data item with a matching major type.
+// No-op for indefinite-length array/map, where any data item is valid.
+func validateIndefiniteLengthChunkByData(indefiniteLengthCborType cborType, chunk []byte, v any) error {
+	if indefiniteLengthCborType != cborTypeByteString &&
+		indefiniteLengthCborType != cborTypeTextString {
+		return nil
+	}
+
+	if len(chunk) == 0 {
+		return errors.New("cbor: cannot encode item type " + reflect.TypeOf(v).Kind().String() +
+			" for indefinite-length " + indefiniteLengthCborType.String())
+	}
+
+	t, ai := parseInitialByte(chunk[0])
+	if t != indefiniteLengthCborType {
+		return errors.New("cbor: cannot encode item type " + reflect.TypeOf(v).Kind().String() +
+			" for indefinite-length " + indefiniteLengthCborType.String())
+	}
+
+	if ai == additionalInformationAsIndefiniteLengthFlag {
+		return errors.New("cbor: cannot encode indefinite-length " + indefiniteLengthCborType.String() +
+			" as chunk of indefinite-length " + indefiniteLengthCborType.String())
+	}
+	return nil
 }
 
 // RawMessage is a raw encoded CBOR value.