service_matcher-go/postgresql/postgresql.go

package postgresql

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

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

const (
	RESPONSE_TYPE_ERR uint8 = 0x45
)

type pgsql struct {
	Len              uint32
	MessageType      uint16
	_                uint16
	Name             [5]byte
	NameValue        byte
	Db               [9]byte
	DBValue          byte
	Encoding         [16]byte
	EncodingValue    [5]byte
	DateStyle        [10]byte
	DateStyleValue   [4]byte
	TimeZone         [9]byte
	TimeZoneValue    [11]byte
	ExtraDigits      [19]byte
	ExtraDigitsValue uint16
	End              byte
}

type pgsqlErrResponse struct {
	ResponseType uint8
	Len          [4]byte
	Data         [128]byte
}

type PostgreSQLMatcher struct {
	osm.Matchers
}

func (m *PostgreSQLMatcher) Key() string {
	return "POSTGRESQL"
}

func (m *PostgreSQLMatcher) Name(matchCtx *osm.MatchCtx) string {
	return "PostgreSQL"
}

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

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

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

func (m *PostgreSQLMatcher) 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)

	pg := pgsqlErrResponse{}
	if err := binary.Read(reader, binary.BigEndian, &pg); err != nil {
		return err
	}

	if pg.ResponseType != RESPONSE_TYPE_ERR {
		return osm.NotMatchedError()
	}

	length := binary.BigEndian.Uint32(pg.Len[:])
	if length+1 != uint32(packet.Len) {
		return osm.NotMatchedError()
	}

	data := string(pg.Data[:])
	splits := strings.Split(data, "\x00")

	findSeverity := false
	findErrorCode := false
	for _, s := range splits {
		if strings.Contains(s, "FATAL") {
			findSeverity = true
		}
		if strings.Contains(s, "28000") {
			findErrorCode = true
		}
	}

	if !findSeverity || !findErrorCode {
		return osm.NotMatchedError()
	}

	return nil

}

func NewMatcher() osm.Matcher {

	m := &PostgreSQLMatcher{}

	pg := pgsql{}

	pg.Len = 0x00000065
	pg.MessageType = 0x0003

	var name [5]byte
	copy(name[:], "user")
	pg.Name = name
	pg.NameValue = 0x00

	var db [9]byte
	copy(db[:], "database")
	pg.Db = db
	pg.DBValue = 0x00

	var encoding [16]byte
	copy(encoding[:], "client_encoding")
	pg.Encoding = encoding

	var encodingValue [5]byte
	copy(encodingValue[:], "UTF8")
	pg.EncodingValue = encodingValue

	var dateStyle [10]byte
	copy(dateStyle[:], "DateStyle")
	pg.DateStyle = dateStyle

	var dateStyleValue [4]byte
	copy(dateStyleValue[:], "ISO")
	pg.DateStyleValue = dateStyleValue

	var timeZone [9]byte
	copy(timeZone[:], "TimeZone")
	pg.TimeZone = timeZone

	var timeZoneValue [11]byte
	copy(timeZoneValue[:], "Asia/Seoul")
	pg.TimeZoneValue = timeZoneValue

	var extraDigit [19]byte
	copy(extraDigit[:], "extra_float_digits")
	pg.ExtraDigits = extraDigit

	pg.ExtraDigitsValue = 0x3200

	writer := new(bytes.Buffer)
	binary.Write(writer, binary.BigEndian, pg)

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

	return m
}