goaidb/protocol/parser.go

// Package protocol 实现MongoDB协议解析
package protocol

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

// UpdateFlags are the flags for OP_UPDATE
const (
	Upsert      = 1 << iota
	MultiUpdate // 标志位用于多文档更新
)

// Header 消息头
type Header struct {
	MessageLength int32
	RequestID     int32
	ResponseTo    int32
	OpCode        OpCode
}

// OpCode 操作码
type OpCode int32

const (
	OP_REPLY         OpCode = 1
	OP_MSG           OpCode = 2
	OP_UPDATE        OpCode = 2001
	OP_INSERT        OpCode = 2002
	RESERVED         OpCode = 2003
	OP_QUERY         OpCode = 2004
	OP_GET_MORE      OpCode = 2005
	OP_DELETE        OpCode = 2006
	OP_KILL_CURSORS  OpCode = 2007
	OP_COMMAND       OpCode = 2010
	OP_COMMAND_REPLY OpCode = 2011
	OP_COMPRESSED    OpCode = 2012
	OP_ENCRYPTED     OpCode = 2013
)

// Message MongoDB协议消息结构
type Message struct {
	Header       Header
	OpCode       OpCode
	OriginalBody []byte      // 原始消息体（解析前）
	Body         interface{} // 解析后的消息体
}

// 解析更新请求
func parseUpdate(data []byte) (interface{}, error) {
	// 检查数据长度（最小需要4字节的flags + 1字节的数据库名终止符）
	if len(data) < 5 {
		return nil, fmt.Errorf("update message data too short")
	}

	// 读取flags
	flags := UpdateFlags(binary.LittleEndian.Uint32(data[0:4]))

	// 查找数据库名结束位置（C风格字符串，以0终止）
	dbEnd := bytes.IndexByte(data[4:], 0)
	if dbEnd == -1 {
		return nil, fmt.Errorf("database name not null terminated")
	}
	dbEnd += 4 // 调整到正确的位置

	// 提取数据库名称
	dbName := string(data[4:dbEnd])

	// 剩余数据包含集合名、查询文档和更新文档
	remaining := data[dbEnd+1:]

	// 查找集合名结束位置
	collEnd := bytes.IndexByte(remaining, 0)
	if collEnd == -1 {
		return nil, fmt.Errorf("collection name not null terminated")
	}

	// 提取集合名
	collName := string(remaining[:collEnd])

	// 剩余数据包含查询文档和更新文档
	bsonData := remaining[collEnd+1:]

	// 解析BSON文档
	queryDoc, _, err := parseBSON(bsonData)
	if err != nil {
		return nil, fmt.Errorf("failed to parse query document: %v", err)
	}

	// 解析更新文档
	updateDoc, _, err := parseBSON(bsonData)
	if err != nil {
		return nil, fmt.Errorf("failed to parse update document: %v", err)
	}

	return &UpdateMessage{
		Body: struct {
			Flags        UpdateFlags
			DatabaseName string
			CollName     string
			Query        map[string]interface{}
			UpdateSpec   map[string]interface{}
		}{
			Flags:        flags,
			DatabaseName: dbName,
			CollName:     collName,
			Query:        queryDoc,
			UpdateSpec:   updateDoc,
		},
	}, nil
}

// ParseMessage 解析MongoDB协议消息
func ParseMessage(data []byte) (*Message, error) {
	// 最小消息长度为16字节（消息头长度）
	if len(data) < 16 {
		return nil, fmt.Errorf("data too short for message header")
	}

	// 验证消息长度是否完整
	messageLength := int(binary.LittleEndian.Uint32(data[0:4]))
	if len(data) < messageLength {
		return nil, fmt.Errorf("data too short for complete message")
	}

	// 截取实际的消息数据（可能有多条消息）
	actualData := data[:messageLength]

	// 解析消息头
	header := &Header{
		MessageLength: int32(binary.LittleEndian.Uint32(actualData[0:4])),
		RequestID:     int32(binary.LittleEndian.Uint32(actualData[4:8])),
		ResponseTo:    int32(binary.LittleEndian.Uint32(actualData[8:12])),
		OpCode:        OpCode(binary.LittleEndian.Uint32(actualData[12:16])),
	}

	// 获取消息体
	body := data[16:]

	// 解析特定操作码的消息体
	var parsedBody interface{}
	switch header.OpCode {
	case OP_QUERY:
		query, err := parseQuery(body)
		if err != nil {
			return nil, err
		}
		parsedBody = query
	case OP_INSERT:
		insert, err := parseInsert(body)
		if err != nil {
			return nil, err
		}
		parsedBody = insert
	case OP_UPDATE:
		update, err := parseUpdate(body)
		if err != nil {
			return nil, err
		}
		parsedBody = update
	// 这里可以添加更多操作码的解析逻辑
	default:
		// 未知操作码，保留原始数据
		parsedBody = body
	}

	return &Message{
		Header:       *header,
		OpCode:       header.OpCode,
		OriginalBody: body,
		Body:         parsedBody,
	}, nil
}

// 解析查询请求
func parseQuery(data []byte) (interface{}, error) {
	// 实现具体的查询消息解析逻辑
	// 这里返回原始数据作为占位符
	if len(data) < 4 {
		return nil, fmt.Errorf("query data too short")
	}

	// 示例：读取查询标志
	flags := binary.LittleEndian.Uint32(data[0:4])

	// 提取数据库名称
	dbEnd := bytes.IndexByte(data[4:], 0)
	if dbEnd == -1 {
		return nil, fmt.Errorf("database name not null terminated")
	}
	dbcolName := string(data[4 : dbEnd+4])
	dcnames := strings.Split(dbcolName, ".")
	dbName := dcnames[0]
	collName := dcnames[1]

	// 剩余数据包含集合名和查询条件
	remaining := data[dbEnd+5:] // 跳过终止符
	var numberToSkip, numberToReturn int32
	binary.Read(bytes.NewReader(remaining[0:4]), binary.LittleEndian, &numberToSkip)
	binary.Read(bytes.NewReader(remaining[4:8]), binary.LittleEndian, &numberToReturn)
	// 提取集合名
	// collEnd := bytes.IndexByte(remaining, 0)
	// if collEnd == -1 {
	// 	return nil, fmt.Errorf("collection name not null terminated")
	// }
	// if collEnd == 0 {
	// 	collEnd = bytes.IndexFunc(remaining, func(r rune) bool {
	// 		return r != 0
	// 	})
	// }
	// collName := string(remaining[:collEnd])

	// 解析查询条件
	queryDoc, _, err := parseBSON(remaining[8:])
	if err != nil {
		return nil, fmt.Errorf("failed to parse query conditions: %v", err)
	}

	return &QueryMessage{
		Flags:          flags,
		DatabaseName:   dbName,
		CollName:       collName,
		NumberToSkip:   numberToSkip,
		NumberToReturn: numberToReturn,
		Query:          queryDoc,
	}, nil
}

// 解析插入请求
func parseInsert(data []byte) (interface{}, error) {
	// 实现具体的插入消息解析逻辑
	if len(data) < 4 {
		return nil, fmt.Errorf("insert data too short")
	}

	// 示例：读取插入标志
	flags := binary.LittleEndian.Uint32(data[0:4])

	// 提取数据库名称
	dbEnd := bytes.IndexByte(data[4:], 0)
	if dbEnd == -1 {
		return nil, fmt.Errorf("database name not null terminated")
	}
	dbName := string(data[4 : dbEnd+4])

	// 剩余数据包含集合名和文档数据
	remaining := data[dbEnd+5:] // 跳过终止符

	// 提取集合名
	collEnd := bytes.IndexByte(remaining, 0)
	if collEnd == -1 {
		return nil, fmt.Errorf("collection name not null terminated")
	}
	collName := string(remaining[:collEnd])

	// 解析文档数据
	documents := make([]map[string]interface{}, 0)
	rest := remaining[collEnd+1:]
	for len(rest) > 0 {
		doc, remainingData, err := parseBSON(rest)
		if err != nil {
			return nil, fmt.Errorf("failed to parse document: %v", err)
		}
		documents = append(documents, doc)
		rest = remainingData
	}

	return &InsertMessage{
		Flags:        flags,
		DatabaseName: dbName,
		CollName:     collName,
		Documents:    documents,
	}, nil
}

// QueryMessage OP_QUERY消息体结构
type QueryMessage struct {
	Flags          uint32 // 查询标志
	DatabaseName   string // 数据库名称
	CollName       string // 集合名称
	NumberToSkip   int32
	NumberToReturn int32
	Query          map[string]interface{} // 查询条件
}

// InsertMessage OP_INSERT消息体结构
type InsertMessage struct {
	Flags        uint32                   // 插入标志
	DatabaseName string                   // 数据库名称
	CollName     string                   // 集合名称
	Documents    []map[string]interface{} // 要插入的文档
}

// UpdateMessage represents an OP_UPDATE message
type UpdateMessage struct {
	Header Header
	Body   struct {
		Flags        UpdateFlags
		DatabaseName string
		CollName     string
		Query        map[string]interface{}
		UpdateSpec   map[string]interface{}
	}
}

// Update applies the update operation to a document
func (u *UpdateMessage) Update(doc map[string]interface{}) (map[string]interface{}, error) {
	result := make(map[string]interface{})
	for key, value := range doc {
		result[key] = value
	}

	// Process $set operator
	if setOp, ok := u.Body.UpdateSpec["$set"].(map[string]interface{}); ok {
		for key, value := range setOp {
			result[key] = value
		}
	}

	// Process $inc operator
	if incOp, ok := u.Body.UpdateSpec["$inc"].(map[string]interface{}); ok {
		for key, value := range incOp {
			if num, ok := value.(int32); ok {
				if current, exists := result[key]; exists {
					if currentNum, ok := current.(int32); ok {
						result[key] = currentNum + num
					} else {
						return nil, fmt.Errorf("cannot apply $inc to non-integer field")
					}
				} else {
					result[key] = num
				}
			} else {
				return nil, fmt.Errorf("invalid increment value")
			}
		}
	}

	return result, nil
}