goaidb/protocol/bson.go

package protocol

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"strings"
)

// BSONElement represents a single BSON element
type BSONElement struct {
	ElementType byte
	KeyName     string
	Value       interface{}
}

// ParseBSON 解析BSON文档并返回映射
func ParseBSON(data []byte) (map[string]interface{}, error) {
	result, _, err := parseBSON(data)
	return result, err
}

// parseBSON 内部使用的BSON解析函数
func parseBSON(data []byte) (map[string]interface{}, []byte, error) {
	// 检查数据长度（最小需要4字节的文档长度）
	if len(data) <= 4 {
		return nil, data, fmt.Errorf("data too short for BSON document length")
	}

	// 读取BSON文档长度
	length := int(binary.LittleEndian.Uint32(data[0:4]))
	if len(data) < length {
		return nil, data, fmt.Errorf("data too short for BSON document")
	}

	// 创建结果映射
	doc := make(map[string]interface{})

	// 指向当前解析位置
	pos := 4

	// 解析BSON元素，直到遇到结束符0x00
	for pos < length {
		// 检查是否有足够的数据读取元素类型
		if pos+1 > len(data) {
			return nil, data, fmt.Errorf("unexpected end of data reading element type")
		}

		// 读取元素类型
		elementType := data[pos]
		pos++

		// 如果是结束元素，跳出循环
		if elementType == 0x00 {
			break
		}

		// 读取键名（C风格字符串，以0终止）
		// keyStart := pos
		keyEnd := bytes.IndexByte(data[pos:], 0)
		if keyEnd == -1 {
			return nil, data, fmt.Errorf("key not null terminated")
		}
		keyEnd += pos // 调整到正确的位置

		keyName := string(data[pos:keyEnd])
		pos = keyEnd + 1

		// 根据元素类型解析值
		value, newPos, err := parseBSONValue(elementType, data[pos:], 0)
		if err != nil {
			return nil, data, fmt.Errorf("failed to parse value for key %s: %v", keyName, err)
		}

		doc[keyName] = value
		pos += newPos
	}

	return doc, data[length:], nil
}

// parseBSONValue 解析特定类型的BSON值
func parseBSONValue(elementType byte, data []byte, pos int) (interface{}, int, error) {
	switch elementType {
	case 0x10: // Int32
		if len(data) < 4 {
			return nil, 0, fmt.Errorf("data too short for Int32")
		}
		value := int32(binary.LittleEndian.Uint32(data[0:4]))
		return value, 4, nil

	case 0x12: // Int64
		if len(data) < 8 {
			return nil, 0, fmt.Errorf("data too short for Int64")
		}
		value := int64(binary.LittleEndian.Uint64(data[0:8]))
		return value, 8, nil

	case 0x01: // Double
		if len(data) < 8 {
			return nil, 0, fmt.Errorf("data too short for Double")
		}
		bits := binary.LittleEndian.Uint64(data[0:8])
		value := math.Float64frombits(bits)
		return value, 8, nil

	case 0x02: // String
		if len(data) < 4 {
			return nil, 0, fmt.Errorf("data too short for String length")
		}

		// 读取字符串长度
		strLength := int(binary.LittleEndian.Uint32(data[0:4]))
		if strLength < 1 {
			return nil, 0, fmt.Errorf("invalid string length %d", strLength)
		}

		// 检查剩余数据是否足够包含整个字符串
		if len(data)-4 < strLength {
			return nil, 0, fmt.Errorf("data too short for String content")
		}

		// 读取字符串内容（忽略最后的终止符）
		value := strings.Trim(string(data[4:4+strLength]), "\x00")
		return value, 4 + strLength, nil

	case 0x08: // Boolean
		if len(data) < 1 {
			return nil, 0, fmt.Errorf("data too short for Boolean")
		}
		value := data[0] != 0
		return value, 1, nil

	case 0x0A: // Null
		return nil, 0, nil

	case 0x03: // EmbeddedDocument
		// 解析嵌入文档
		if len(data) < 4 {
			return nil, 0, fmt.Errorf("data too short for EmbeddedDocument length")
		}

		// 读取嵌入文档长度
		docLength := int(binary.LittleEndian.Uint32(data[0:4]))
		if len(data) < docLength {
			return nil, 0, fmt.Errorf("data too short for EmbeddedDocument")
		}

		// 检查是否为更新操作符（以$开头的键名）
		if pos > 0 && data[pos-1] == 0x02 { // 前一个字节是字符串类型标记
			// 查找前一个键名
			for i := pos - 2; i >= 0; i-- {
				if data[i] == 0x02 { // 找到字符串类型标记
					keyLen := int(binary.LittleEndian.Uint32(data[i+1 : i+5]))
					keyStart := i + 5
					if keyStart+keyLen <= pos {
						key := string(data[keyStart : keyStart+keyLen-1])
						if len(key) > 0 && key[0] == '$' {
							// 这是一个更新操作符
							value := make(map[string]interface{})
							value[key] = "operator_placeholder"
							return value, docLength, nil
						}
					}
				}
			}
		}

		// 解析嵌入文档内容
		subDoc, _, err := parseBSON(data[0:docLength])
		if err != nil {
			return nil, 0, fmt.Errorf("failed to parse embedded document: %v", err)
		}

		return subDoc, docLength, nil

	default:
		return nil, 0, fmt.Errorf("unsupported BSON element type: 0x%02X", elementType)
	}
}

// BsonMarshal 将map转换为BSON格式的字节流
func BsonMarshal(doc map[string]interface{}) ([]byte, error) {
	// TODO: 实现实际的BSON序列化或使用现有库
	// 这里返回模拟实现
	return []byte{}, nil
}