316 lines
7.7 KiB
Go
316 lines
7.7 KiB
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) {
|
||
buf := &bytes.Buffer{}
|
||
|
||
// 写入占位符长度（4字节）
|
||
buf.Write(make([]byte, 4))
|
||
|
||
// 遍历文档元素
|
||
for key, value := range doc {
|
||
// 写入元素类型和键名
|
||
elementType, err := getBSONType(value)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
buf.WriteByte(elementType)
|
||
buf.WriteString(key)
|
||
buf.WriteByte(0x00) // 键名终止符
|
||
|
||
// 写入值数据
|
||
if err := writeBSONValue(buf, elementType, value); err != nil {
|
||
return nil, err
|
||
}
|
||
}
|
||
|
||
// 写入文档结束符
|
||
buf.WriteByte(0x00)
|
||
|
||
// 回填文档长度
|
||
length := uint32(buf.Len())
|
||
binary.LittleEndian.PutUint32(buf.Bytes(), length)
|
||
|
||
return buf.Bytes(), nil
|
||
}
|
||
|
||
// getBSONType 根据Go类型获取BSON元素类型
|
||
func getBSONType(value interface{}) (byte, error) {
|
||
switch value.(type) {
|
||
case int32:
|
||
return 0x10, nil // Int32
|
||
case int64:
|
||
return 0x12, nil // Int64
|
||
case float64:
|
||
return 0x01, nil // Double
|
||
case string:
|
||
return 0x02, nil // String
|
||
case bool:
|
||
return 0x08, nil // Boolean
|
||
case nil:
|
||
return 0x0A, nil // Null
|
||
case map[string]interface{}:
|
||
return 0x03, nil // EmbeddedDocument
|
||
default:
|
||
return 0x00, fmt.Errorf("unsupported BSON type: %T", value)
|
||
}
|
||
}
|
||
|
||
// writeBSONValue 写入BSON值数据
|
||
func writeBSONValue(buf *bytes.Buffer, elementType byte, value interface{}) error {
|
||
switch elementType {
|
||
case 0x10: // Int32
|
||
v, ok := value.(int32)
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for Int32")
|
||
}
|
||
b := make([]byte, 4)
|
||
binary.LittleEndian.PutUint32(b, uint32(v))
|
||
buf.Write(b)
|
||
|
||
case 0x12: // Int64
|
||
v, ok := value.(int64)
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for Int64")
|
||
}
|
||
b := make([]byte, 8)
|
||
binary.LittleEndian.PutUint64(b, uint64(v))
|
||
buf.Write(b)
|
||
|
||
case 0x01: // Double
|
||
v, ok := value.(float64)
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for Double")
|
||
}
|
||
b := make([]byte, 8)
|
||
binary.LittleEndian.PutUint64(b, math.Float64bits(v))
|
||
buf.Write(b)
|
||
|
||
case 0x02: // String
|
||
v, ok := value.(string)
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for String")
|
||
}
|
||
strBytes := []byte(v)
|
||
// 写入字符串长度（包含终止符）
|
||
length := uint32(len(strBytes) + 1)
|
||
b := make([]byte, 4)
|
||
binary.LittleEndian.PutUint32(b, length)
|
||
buf.Write(b)
|
||
// 写入字符串内容和终止符
|
||
buf.Write(strBytes)
|
||
buf.WriteByte(0x00)
|
||
|
||
case 0x08: // Boolean
|
||
v, ok := value.(bool)
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for Boolean")
|
||
}
|
||
if v {
|
||
buf.WriteByte(0x01)
|
||
} else {
|
||
buf.WriteByte(0x00)
|
||
}
|
||
|
||
case 0x0A: // Null
|
||
// 不需要写入任何数据
|
||
|
||
case 0x03: // EmbeddedDocument
|
||
v, ok := value.(map[string]interface{})
|
||
if !ok {
|
||
return fmt.Errorf("invalid type for EmbeddedDocument")
|
||
}
|
||
subDoc, err := BsonMarshal(v)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
// 直接写入子文档数据（包含完整的长度信息）
|
||
buf.Write(subDoc)
|
||
|
||
default:
|
||
return fmt.Errorf("unsupported BSON element type: 0x%02X", elementType)
|
||
}
|
||
return nil
|
||
}
|