xml_decoder.go

package binxml

import (
	"bytes"
	"encoding/base64"
	"encoding/xml"
	"fmt"
	"io"
	"strconv"
	"strings"
	"unicode/utf16"

	"wcfproxy/assert"
	dt "wcfproxy/datatype"
	"wcfproxy/logging"
)

const XMLNSArrays = "http://schemas.microsoft.com/2003/10/Serialization/Arrays"

const (
	XmlDecodeModeTrucate = iota // truncate if value does not fit into type
	XmlDecodeModeUpsize         // upsize type (e.g. i8 -> i16 or i8 -> i32) if value does not fit into type
)

type XMLDecoder struct {
	dict                *StringDict
	compressEndElements bool
	logger              *logging.Logger
}

func (d *XMLDecoder) DecodeEnvelope(xmlData string) ([]IRecord, error) {
	var records []IRecord
	reader := xml.NewDecoder(bytes.NewReader([]byte(xmlData)))
	var prevRecord *IRecord

	for {
		token, err := reader.RawToken()
		if err == io.EOF {
			return records, nil
		}

		if err != nil {
			return nil, err
		}

		if chardata, ok := token.(xml.CharData); ok && isWhitespaceOnly(string(chardata)) {
			// skip char data that is white space only (likely stems from XML formatting)
			// because of type hints, there is no risk to strip legitimate text records
			continue
		}

		var decodedRecords []IRecord
		if itemType, ok := checkArrayStart(token); ok {
			arrayTokens, err := collectArrayTokens(reader, token)
			if err != nil {
				return nil, err
			}

			// first element is embedded as ElementRecord into the array
			startElement, ok := token.(xml.StartElement)
			if !ok {
				return nil, fmt.Errorf("expected xml.StartElement, got %v (%T) (this is a bug)", startElement, startElement)
			}

			arrayRecord, err := d.decodeArray(arrayTokens, itemType)
			if err != nil {
				return nil, err
			}

			assert.Assertf(len(arrayRecord) == 1, "expected to decode one array record but got %d", len(arrayRecord))

			decodedRecords = append(decodedRecords, arrayRecord...)
			endRecord := Record{Type: RTEndElement}
			decodedRecords = append(decodedRecords, &endRecord)
		} else if decodedRecords, err = d.recordsFromXMLToken(token, prevRecord); err != nil {
			return nil, err
		}

		// allow XML tokens to be ignored (useful to skip whitespace-only tokens)
		if decodedRecords == nil {
			continue
		}

		records = append(records, decodedRecords...)
		if len(decodedRecords) > 0 {
			prevRecord = &decodedRecords[len(decodedRecords)-1]
		}
	}
}

func (d *XMLDecoder) recordsFromXMLToken(token xml.Token, prevRecord *IRecord) ([]IRecord, error) {
	switch t := token.(type) {
	case xml.StartElement:
		se, err := d.decodeStartElement(t)
		if err != nil {
			return nil, err
		}
		return []IRecord{se}, nil
	case xml.EndElement:
		// if possible include end element in last (text) record
		if d.compressEndElements && prevRecord != nil {
			if rtEndElem, ok := EndElementRecordsMap[(*prevRecord).GetType()]; ok {
				(*prevRecord).SetType(rtEndElem)
				return nil, nil
			}
		}

		// end element was not included in text record, therefore emit regular end element
		return []IRecord{&Record{Type: RTEndElement}}, nil
	case xml.CharData:
		record, err := d.decodeTextRecord(string(t))
		if err != nil {
			return nil, err
		}

		// ignore whitespaceo only char data
		if record == nil {
			return nil, nil
		}

		return []IRecord{record}, err
	case xml.Comment:
		// TODO: implement XML comments
		return nil, fmt.Errorf("xml comments not implemented")
	default:
		return nil, fmt.Errorf("unhandled XML token: %v (%T)", token, token)
	}
}

func collectArrayTokens(reader *xml.Decoder, firstToken xml.Token) ([]xml.Token, error) {
	level := 1 // start at 1 since the first start element was already read
	arrayItems := []xml.Token{firstToken}
	for {
		token, err := reader.RawToken()
		if err == io.EOF {
			return nil, fmt.Errorf("illegal state: got EOF while reading array")
		}

		if err != nil {
			return nil, err
		}

		// we ignore whitespace only character data (which most likely stems from formatting)
		if chardata, ok := token.(xml.CharData); ok && isWhitespaceOnly(string(chardata)) {
			continue
		}

		switch token.(type) {
		case xml.StartElement:
			level++
		case xml.EndElement:
			// the end element after the array will decrement level to -1
			level--
		}

		if level > 1 || level < -1 {
			return nil, fmt.Errorf("unxpected level during array decoding: %d", level)
		}

		arrayItems = append(arrayItems, xml.CopyToken(token))
		if level < 0 {
			return arrayItems, nil
		}
	}
}

func (d *XMLDecoder) decodeArray(tokens []xml.Token, itemType byte) ([]IRecord, error) {
	if len(tokens) < 2 {
		return nil, fmt.Errorf("error decoding ArrayElement: not enough XML tokens (%d)", len(tokens))
	}

	startElement, ok := tokens[0].(xml.StartElement)
	if !ok {
		return nil, fmt.Errorf("array must start with an xml.StartElement, got %T (%v) (this is a bug)", tokens[0], tokens[0])
	}

	startElementRecord, err := d.decodeStartElement(startElement)
	if err != nil {
		return nil, err
	}

	if _, ok = tokens[len(tokens)-1].(xml.EndElement); !ok {
		return nil, fmt.Errorf("array must end with an xml.EndElement, got %T (%v) (this is a bug)", tokens[len(tokens)-1], tokens[len(tokens)-1])
	}

	level := 0
	var items []ITextRecord
	for _, e := range tokens {
		switch e := e.(type) {
		case xml.StartElement:
			level++
		case xml.EndElement:
			level--
			if level < -1 || level > 1 {
				return nil, fmt.Errorf("invalid array: unbalanced tags")
			}
		case xml.CharData:
			textRecord, err := d.decodeArrayElement(string(e), itemType)
			if err != nil {
				return nil, err
			}

			// skip white space only char data
			if textRecord == nil {
				continue
			}

			textRecord.SetType(itemType)
			items = append(items, textRecord)
		}
	}

	if level != -1 {
		return nil, fmt.Errorf("unexpected level %d at the end of array decoding", level)
	}

	arrayRecord := NewArrayRecord(startElementRecord, itemType, items)
	return []IRecord{&arrayRecord}, nil
}

func (d *XMLDecoder) decodeArrayElement(element string, arrayType byte) (ITextRecord, error) {
	switch arrayType {
	case RTBoolTextWithEndElement:
		return d.decodeBool(element)
	case RTInt16TextWithEndElement:
		return d.decodeInt16(element)
	case RTInt32TextWithEndElement:
		return d.decodeInt32(element)
	case RTInt64TextWithEndElement:
		return d.decodeInt64(element)
	case RTFloatTextWithEndElement:
		return d.decodeFloat(element)
	case RTDoubleTextWithEndElement:
		return d.decodeDouble(element)
	case RTDecimalTextWithEndElement:
		return d.decodeDecimal(element)
	case RTDateTimeTextWithEndElement:
		return d.decodeDateTime(element)
	case RTTimeSpanTextWithEndElement:
		return d.decodeTimeSpan(element)
	case RTUuidTextWithEndElement:
		return d.decodeUUID(element)
	default:
		return nil, fmt.Errorf("unknown array element type: %d", arrayType)
	}
}

func checkArrayStart(token xml.Token) (byte, bool) {
	se, ok := token.(xml.StartElement)
	if !ok {
		return 0, false
	}

	for _, a := range se.Attr {
		if a.Name.Space == HintXmlNamespace && a.Name.Local == HintXmlAttrName {
			if strings.HasPrefix(a.Value, HintXmlAttrValArrayPrefix) {
				rt, err := arrayTypeFromHint(a.Value)
				if err != nil {
					return 0, false
				}

				return rt, true
			}
		}
	}

	return 0, false
}

func (d *XMLDecoder) decodeStartElement(e xml.StartElement) (IElementRecord, error) {
	prefix := e.Name.Space
	name := e.Name.Local

	var attributes []IAttributeRecord
	for _, attr := range e.Attr {
		attrRecord, err := d.decodeAttribute(attr)
		if err != nil {
			return nil, err
		}

		// ingore attributes that decode to nil (e.g. type hints inserted during translation from binary XML to XML)
		if attrRecord != nil {
			attributes = append(attributes, attrRecord)
		}
	}

	switch {
	case prefix == "":
		// ShortElement or ShortDictionaryElement
		if i, ok := d.dict.ReverseLookup(name); ok {
			// ShortDictionaryElement
			r := NewShortDictElementRecord(i, attributes)
			return &r, nil
		} else {
			// ShortElement
			r := NewShortElementRecord(name, attributes)
			return &r, nil
		}

	case isAZPrefix(prefix):
		// PrefixElementAZ or PrefixDictionaryElementAZ
		prefixChar := prefix[0]
		if i, ok := d.dict.ReverseLookup(name); ok {
			// PrefixDictionaryElementAZ
			r := NewPrefixDictElementAZRecord(prefixChar, i, attributes)
			return &r, nil
		} else {
			// PrefixElementAZ
			r := NewPrefixElementAZ(prefixChar, name, attributes)
			return &r, nil
		}

	default:
		// ElementRecord or DictionaryElementRecord
		if i, ok := d.dict.ReverseLookup(name); ok {
			// DictionaryElementRecord
			r := NewDictElementRecord(prefix, i, attributes)
			return &r, nil
		} else {
			// ElementRecord
			r := NewElementRecord(prefix, name, attributes)
			return &r, nil
		}
	}
}

func (d *XMLDecoder) decodeAttribute(a xml.Attr) (IAttributeRecord, error) {
	if isWcfpHint(a) {
		return nil, nil
	}

	attrName := a.Name.Local
	prefix := a.Name.Space

	switch {
	case prefix == "":
		// ShortAttribute, ShortDictionaryAttribute, ShortXmlnsAttribute or ShortDictionaryXmlnsAttribute
		switch {
		case attrName == "xmlns":
			// ShortXmlnsAttribute or ShortDictionaryXmlnsAttribute
			if i, ok := d.dict.ReverseLookup(a.Value); ok {
				// ShortDictionaryXmlnsAttribute
				r := NewShortDictXmlnsAttributeRecord(i)
				return &r, nil
			} else {
				// ShortXmlnsAttribute
				r := NewShortXmlnsAttributeRecord(a.Value)
				return &r, nil
			}

		case d.dict.HasString(attrName):
			// ShortDictionaryAttribute
			attrValue, err := d.decodeTextRecord(a.Value)
			if err != nil {
				return nil, err
			}

			i, _ := d.dict.ReverseLookup(attrName)
			r := NewShortDictAttributeRecord(i, attrValue)
			return &r, nil

		default:
			// ShortAttribute
			attrValue, err := d.decodeTextRecord(a.Value)
			if err != nil {
				return nil, err
			}

			r := NewShortAttributeRecord(attrName, attrValue)
			return &r, nil
		}
	case prefix == "xmlns":
		// XmlnsAttribute or DictionaryXmlnsAttribute
		if i, ok := d.dict.ReverseLookup(a.Value); ok {
			// DictionaryXmlnsAttribute
			r := NewDictXmlnsAttributeRecord(attrName, i)
			return &r, nil
		} else {
			// XmlnsAttribute
			r := NewXmlnsAttributeRecord(attrName, a.Value)
			return &r, nil
		}

	case isAZPrefix(prefix):
		// PrefixAttriuteAZ or PrefixDictionaryAttributeAZ
		attrValue, err := d.decodeTextRecord(a.Value)
		if err != nil {
			return nil, err
		}

		prefixChar := prefix[0]
		if i, ok := d.dict.ReverseLookup(attrName); ok {
			// PrefixDictionaryAttributeAZ
			r := NewPrefixDictAttributeAZ(prefixChar, i, attrValue)
			return &r, nil
		} else {
			// PrefixAttriuteAZ
			r := NewPrefixAttributeAZ(prefixChar, attrName, attrValue)
			return &r, nil
		}

	default:
		// AttributeRecord or DictionaryAttribute
		attrVal, err := d.decodeTextRecord(a.Value)
		if err != nil {
			return nil, err
		}

		if i, ok := d.dict.ReverseLookup(attrName); ok {
			// DictionaryAttribute
			r := NewDictAttributeRecord(prefix, i, attrVal)
			return &r, nil
		} else {
			// AttributeRecord
			r := NewAttributeRecord(prefix, attrName, attrVal)
			return &r, nil
		}
	}
}

func isWcfpHint(attr xml.Attr) bool {
	return attr.Name.Space == HintXmlNamespace && attr.Name.Local == HintXmlAttrName
}

// ListTextRecords can contain mixed types, breaking simple type hinting.
// - For this reason, a ListTextRecord is tranlated
// - to a whole XML tree during decoding of binary XML and each list element is
// - annotated with a type hint.
// - This translation needs to be reversed during XML decoding, thus the special treatment.
func (d *XMLDecoder) decodeListTextRecord(tokens []xml.Token) (ITextRecord, error) {
	return nil, fmt.Errorf("list text records not implemented")
}

// Decode a text record.
func (d *XMLDecoder) decodeTextRecord(textData string) (ITextRecord, error) {
	parts := strings.SplitN(textData, HintSeparator, 2)
	if len(parts) < 2 {
		return nil, fmt.Errorf("no type hint given: %s (%x)", textData, []byte(textData))
	}

	hint := parts[0]
	data := parts[1]
	switch hint {
	case HintBool:
		return d.decodeBool(data)
	case HintBytes:
		return d.decodeBytes(data)
	case HintChars:
		return d.decodeChars(data)
	case HintConst:
		return d.decodeConst(data)
	case HintDateTime:
		return d.decodeDateTime(data)
	case HintDecimal:
		return d.decodeDecimal(data)
	case HintDouble:
		return d.decodeDouble(data)
	case HintFloat:
		return d.decodeFloat(data)
	case HintInt:
		return d.decodeInt(data)
	case HintQName:
		return d.decodeQName(data)
	case HintUint:
		return d.decodeUint(data)
	case HintUUID:
		return d.decodeUUID(data)
	case HintTimeSpan:
		return d.decodeTimeSpan(data)
	case HintUnicode:
		return d.decodeUnicodeChars(data)
	default:
		return nil, fmt.Errorf("invalid type hint: %s", hint)
	}

}

func (d *XMLDecoder) decodeBool(text string) (ITextRecord, error) {
	boolText := strings.ToLower(strings.TrimSpace(text))
	var record BoolTextRecord
	switch boolText {
	case "true":
		record = NewBoolTextRecord(true)
	case "false":
		record = NewBoolTextRecord(false)
	default:
		return nil, fmt.Errorf("invalid bool value: %s", text)
	}

	return &record, nil
}

func (d *XMLDecoder) decodeBytes(text string) (ITextRecord, error) {
	bytesText := strings.TrimSpace(text)
	bytes, err := base64.StdEncoding.DecodeString(bytesText)
	if err != nil {
		return nil, err
	}

	record := NewBytesTextRecord(bytes)
	return &record, nil
}

func (d *XMLDecoder) decodeChars(text string) (ITextRecord, error) {
	if i, ok := d.dict.ReverseLookup(text); ok {
		record := NewDictTextRecord(i)
		return &record, nil
	} else {
		record := NewCharsTextRecord(text)
		return &record, nil
	}
}

func (d *XMLDecoder) decodeUnicodeChars(text string) (ITextRecord, error) {
	utf16Str := utf16.Encode([]rune(text))
	record := NewUnicodeCharsTextRecord(utf16Str)
	return &record, nil
}

func (d *XMLDecoder) decodeConst(text string) (ITextRecord, error) {
	constText := strings.ToLower(strings.TrimSpace(text))
	var rec ConstTextRecord
	switch constText {
	case "":
		rec = NewConstTextRecord(RTEmptyText)
	case "0":
		rec = NewConstTextRecord(RTZeroText)
	case "1":
		rec = NewConstTextRecord(RTOneText)
	case "false":
		rec = NewConstTextRecord(RTFalseText)
	case "true":
		rec = NewConstTextRecord(RTTrueText)
	default:
		return nil, fmt.Errorf("invalid constant value: %s", text)
	}

	return &rec, nil
}

func (d *XMLDecoder) decodeDateTime(text string) (ITextRecord, error) {
	dateTimeText := strings.TrimSpace(text)
	dateTime, err := dt.DateTimeFromString(dateTimeText)
	if err != nil {
		return nil, err
	}

	record := NewDateTimeRecord(dateTime)
	return &record, nil
}

func (d *XMLDecoder) decodeDecimal(text string) (ITextRecord, error) {
	decimalText := strings.TrimSpace(text)
	decimal, err := dt.DecimalFromString(decimalText)
	if err != nil {
		return nil, err
	}

	record := NewDecimalTextRecord(decimal)
	return &record, nil
}

func (d *XMLDecoder) decodeDouble(text string) (ITextRecord, error) {
	doubleText := strings.TrimSpace(text)
	doubleVal, err := strconv.ParseFloat(doubleText, 64)
	if err != nil {
		return nil, err
	}

	record := NewDoubleTextRecord(doubleVal)
	return &record, nil
}

func (d *XMLDecoder) decodeFloat(text string) (ITextRecord, error) {
	floatText := strings.TrimSpace(text)
	floatVal, err := strconv.ParseFloat(floatText, 32)
	if err != nil {
		return nil, err
	}

	record := NewFloatTextRecord(float32(floatVal))
	return &record, nil
}

func (d *XMLDecoder) parseIntGeneric(text string) (int64, error) {
	intText := strings.TrimSpace(text)
	return strconv.ParseInt(intText, 10, 64)
}

func (d *XMLDecoder) decodeInt(text string) (ITextRecord, error) {
	i, err := d.parseIntGeneric(text)
	if err != nil {
		return nil, err
	}

	r := NewIntTextRecord(i)
	return &r, nil
}

func (d *XMLDecoder) decodeInt16(text string) (ITextRecord, error) {
	i, err := d.parseIntGeneric(text)
	if err != nil {
		return nil, err
	}

	i16 := int16(i)
	if !(i < (1<<15) && i >= -(1<<15)) {
		d.logger.Warn("[binxml] Truncating %d to %d (int16)", i, i16)
	}

	r := NewInt16TextRecord(i16)
	return &r, nil
}

func (d *XMLDecoder) decodeInt32(text string) (ITextRecord, error) {
	i, err := d.parseIntGeneric(text)
	if err != nil {
		return nil, err
	}

	i32 := int32(i)
	if !(i < (1<<31) && i >= -(1<<31)) {
		d.logger.Warn("[binxml] Truncating %d to %d (int32)", i, i32)
	}

	r := NewInt32TextRecord(i32)
	return &r, nil
}

func (d *XMLDecoder) decodeInt64(text string) (ITextRecord, error) {
	i, err := d.parseIntGeneric(text)
	if err != nil {
		return nil, err
	}

	r := NewInt64TextRecord(i)
	return &r, nil
}

func (d *XMLDecoder) decodeQName(text string) (ITextRecord, error) {
	parts := strings.SplitN(text, ":", 2)
	if len(parts) < 2 {
		return nil, fmt.Errorf("invalid QNameText: %s", text)
	}

	prefixStr := parts[0]
	nameStr := parts[1]
	if !isAZPrefix(prefixStr) {
		return nil, fmt.Errorf("invalid QNameText: %s", text)
	}

	prefix := prefixStr[0] - 'a'
	if i, ok := d.dict.ReverseLookup(nameStr); ok {
		record := NewQNameDictTextRecord(prefix, i)
		return &record, nil
	}

	return nil, fmt.Errorf("invalid QnameText: %s", text)
}

func (d *XMLDecoder) decodeUint(text string) (ITextRecord, error) {
	uintText := strings.TrimSpace(text)
	uintVal, err := strconv.ParseUint(uintText, 10, 64)
	if err != nil {
		return nil, err
	}

	record := NewUint64TextRecord(uintVal)
	return &record, nil
}

func (d *XMLDecoder) decodeUUID(text string) (ITextRecord, error) {
	uuidText := strings.TrimSpace(text)
	isUniqueId := strings.HasPrefix(uuidText, UniqueIDPrefix)
	uuidText = strings.TrimPrefix(uuidText, UniqueIDPrefix)
	uuid, err := dt.UUIDFromString(uuidText)
	if err != nil {
		return nil, err
	}

	var record UUIDTextRecord
	if isUniqueId {
		record = NewUniqueIDTextRecord(uuid)
	} else {
		record = NewUUIDTextRecord(uuid)
	}

	return &record, nil
}

func (d *XMLDecoder) decodeTimeSpan(text string) (ITextRecord, error) {
	timeSpanText := strings.TrimSpace(text)
	timeSpan, err := dt.TimeSpanFromString(timeSpanText)
	if err != nil {
		return nil, err
	}

	record := NewTimeSpanRecord(timeSpan)
	return &record, nil
}

func isAZPrefix(s string) bool {
	if len(s) != 1 {
		return false
	}

	c := s[0]
	return c <= 'z' && c >= 'a'
}

func isWhitespaceOnly(s string) bool {
	return strings.TrimSpace(s) == ""
}