service_matcher-go/mongodb/mongodb.go

package mongodb

import (
	"bytes"
	"encoding/binary"
	"math/rand"

	osm "git.loafle.net/overflow/service_matcher-go"
)

const (
	MONGO_OP_REQUEST uint32 = 2004
	MONGO_OP_REPLY   uint32 = 1
	MONGO_FCNAME     string = "admin.$cmd"
	MONGO_ELEMENT    string = "ismaster"
)

var MONGO_REQUEST_ID uint32

type MsgHeader struct {
	MessageLength uint32
	RequestId     uint32
	ResponseTo    uint32
	OpCode        uint32
}

type OP_request struct {
	Header             MsgHeader
	Flags              uint32
	FullCollectionName [11]byte
	NumberToSkip       uint32
	NumberToReturn     int32
	DocumentLength     uint32
	Type_              uint8
	Element            [9]byte
	Value              uint8
	_                  uint8
}

type OP_reply struct {
	Header         MsgHeader
	ResponseFlags  int32
	CursorID       int64
	StartingFrom   int32
	NumberReturned int32
	Documents      [512]byte
}

type MongoDBMatcher struct {
	osm.Matchers
}

func (m *MongoDBMatcher) Key() string {
	return "MONGODB"
}

func (m *MongoDBMatcher) Type() string {
	return "NOSQL"
}

func (m *MongoDBMatcher) Name(matchCtx *osm.MatchCtx) string {
	return "MongoDB"
}

func (m *MongoDBMatcher) IsPrePacket() bool {
	return false
}

func (m *MongoDBMatcher) HasResponse(matchCtx *osm.MatchCtx, index int) bool {
	return true
}

func (m *MongoDBMatcher) IsError(matchCtx *osm.MatchCtx, index int, packet *osm.Packet) bool {
	return false
}

func (m *MongoDBMatcher) Match(matchCtx *osm.MatchCtx, index int, packet *osm.Packet) error {

	if packet == nil || !packet.Valid() {
		return osm.NoPacketReceivedError()
	}

	reader := new(bytes.Buffer)
	reader.Write(packet.Buffer)

	reply := OP_reply{}
	if err := binary.Read(reader, binary.LittleEndian, &reply); err != nil {
		return err
	}

	if uint32(packet.Len) != reply.Header.MessageLength {
		return osm.NotMatchedError()
	}

	if reply.Header.ResponseTo != MONGO_REQUEST_ID {
		return osm.NotMatchedError()
	}

	if reply.Header.OpCode != MONGO_OP_REPLY {
		return osm.NotMatchedError()
	}

	return nil

}

func NewMatcher() osm.Matcher {

	mm := &MongoDBMatcher{}

	tempBuf := new(bytes.Buffer)
	binary.Write(tempBuf, binary.BigEndian, OP_request{})

	var fcn [11]byte
	copy(fcn[:], MONGO_FCNAME)

	var elem [9]byte
	copy(elem[:], MONGO_ELEMENT)

	MONGO_REQUEST_ID = rand.Uint32()
	m := OP_request{
		Header: MsgHeader{
			MessageLength: uint32(len(tempBuf.Bytes())),
			RequestId:     MONGO_REQUEST_ID,
			ResponseTo:    0,
			OpCode:        MONGO_OP_REQUEST,
		},
		Flags:              0,
		FullCollectionName: fcn,
		NumberToSkip:       0,
		NumberToReturn:     -1,
		DocumentLength:     16,
		Type_:              0x08,
		Element:            elem,
		Value:              1,
	}
	writer := new(bytes.Buffer)
	binary.Write(writer, binary.LittleEndian, m)

	mm.AddPacket(osm.NewPacket(writer.Bytes(), writer.Len()))

	return mm
}