271 lines
8.3 KiB
Go
271 lines
8.3 KiB
Go
package mysqlproxy
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"errors"
|
|
"linkport/internal/model"
|
|
"math/big"
|
|
"strconv"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
"github.com/go-mysql-org/go-mysql/mysql"
|
|
"github.com/pingcap/tidb/pkg/parser/charset"
|
|
)
|
|
|
|
// Query responses are written as streaming NDJSON events. We never accumulate
|
|
// a complete result set, and audit persistence must succeed before forwarding.
|
|
type queryResultAudit struct {
|
|
ctx context.Context
|
|
output model.CommandOutput
|
|
startedAt time.Time
|
|
binary bool
|
|
fields []*mysql.Field
|
|
index, rows, setRows uint64
|
|
affected big.Int
|
|
databaseError bool
|
|
}
|
|
|
|
func (a *queryResultAudit) start(command model.Command) error {
|
|
a.startedAt = time.Now()
|
|
a.output = model.CommandOutput{CommandID: command.AuditID, RecordedAt: command.RecordedAt, Stream: "result",
|
|
Result: model.CommandResult{Scope: "mysql", State: "open", Rows: "0", AffectedRows: "0"}}
|
|
return a.persist(nil)
|
|
}
|
|
|
|
func (a *queryResultAudit) persist(event any) error {
|
|
a.output.Result.UpdatedAt = time.Now().UTC()
|
|
a.output.Result.Rows = strconv.FormatUint(a.rows, 10)
|
|
a.output.Result.AffectedRows = a.affected.String()
|
|
if event != nil {
|
|
data, err := json.Marshal(event)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
a.output.Data = append(data, '\n')
|
|
}
|
|
err := model.AuditCommandOutput(a.ctx, a.output)
|
|
a.output.Data = nil
|
|
return err
|
|
}
|
|
|
|
func (a *queryResultAudit) finish(complete bool) error {
|
|
if a == nil {
|
|
return nil
|
|
}
|
|
a.output.Result.State = "interrupted"
|
|
if complete {
|
|
a.output.Result.State = "completed"
|
|
if a.databaseError {
|
|
a.output.Result.State = "error"
|
|
}
|
|
}
|
|
if !a.startedAt.IsZero() {
|
|
duration := float64(time.Since(a.startedAt)) / float64(time.Millisecond)
|
|
a.output.Result.DurationMS = &duration
|
|
if complete {
|
|
model.ObserveExecution(a.ctx, "mysql", "query", a.startedAt)
|
|
}
|
|
}
|
|
return a.persist(nil)
|
|
}
|
|
|
|
func (a *queryResultAudit) observeStatus(status uint16) {
|
|
a.output.Result.NoIndexUsed = a.output.Result.NoIndexUsed || status&mysql.SERVER_STATUS_NO_INDEX_USED != 0
|
|
a.output.Result.NoGoodIndexUsed = a.output.Result.NoGoodIndexUsed || status&mysql.SERVER_STATUS_NO_GOOD_INDEX_USED != 0
|
|
}
|
|
|
|
func (a *queryResultAudit) observe(phase string, packet []byte) error {
|
|
if a == nil {
|
|
return nil
|
|
}
|
|
event, err := a.decode(phase, packet)
|
|
if err != nil {
|
|
// A future/unknown wire type must not silently lose values. Preserve the
|
|
// entire payload, with an explicit decode warning, without guessing.
|
|
a.output.Result.DecodeIncomplete = true
|
|
event = map[string]any{"type": "raw", "index": a.index, "phase": phase, "base64": packet, "reason": "unsupported or malformed result packet"}
|
|
}
|
|
return a.persist(event)
|
|
}
|
|
|
|
func (a *queryResultAudit) decode(phase string, b []byte) (event any, err error) {
|
|
defer func() {
|
|
if recover() != nil {
|
|
err = errors.New("malformed result packet")
|
|
}
|
|
}()
|
|
if phase == "first" {
|
|
a.index++
|
|
a.fields = nil
|
|
a.setRows = 0
|
|
}
|
|
if len(b) == 0 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
if b[0] == mysql.ERR_HEADER {
|
|
if len(b) < 3 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
code, state, offset := binary.LittleEndian.Uint16(b[1:3]), "", 3
|
|
if len(b) > 3 && b[3] == '#' {
|
|
if len(b) < 9 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
state, offset = string(b[4:9]), 9
|
|
}
|
|
a.databaseError, a.output.Result.ErrorCode = true, code
|
|
return map[string]any{"type": "error", "index": a.index, "code": code, "sqlState": state, "message": resultValue(b[offset:], false)}, nil
|
|
}
|
|
switch phase {
|
|
case "first":
|
|
if b[0] == mysql.OK_HEADER {
|
|
pos := 1
|
|
affected, err := resultInteger(b, &pos)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
insertID, err := resultInteger(b, &pos)
|
|
if err != nil || len(b)-pos < 4 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
var count big.Int
|
|
count.SetUint64(affected)
|
|
a.affected.Add(&a.affected, &count)
|
|
status := binary.LittleEndian.Uint16(b[pos:])
|
|
a.observeStatus(status)
|
|
return map[string]any{"type": "ok", "index": a.index, "affectedRows": strconv.FormatUint(affected, 10), "insertId": strconv.FormatUint(insertID, 10), "status": status, "warnings": binary.LittleEndian.Uint16(b[pos+2:]), "message": resultValue(b[pos+4:], false)}, nil
|
|
}
|
|
pos := 0
|
|
n, err := resultInteger(b, &pos)
|
|
if err != nil || pos != len(b) || n == 0 || n > 65535 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
a.output.Result.ResultSets++
|
|
return map[string]any{"type": "begin", "index": a.index, "columnCount": n, "binary": a.binary}, nil
|
|
case "column":
|
|
// Field.Parse retains slices into its input. Copy before the packet
|
|
// library reuses the buffer for the next row.
|
|
a.fields = append(a.fields, nil)
|
|
field, err := mysql.FieldData(append([]byte(nil), b...)).Parse()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
a.fields[len(a.fields)-1] = field
|
|
return map[string]any{"type": "column", "index": a.index, "ordinal": len(a.fields), "name": resultValue(field.Name, false), "table": resultValue(field.Table, false), "mysqlType": field.Type, "charset": field.Charset, "flags": field.Flag}, nil
|
|
case "metadata":
|
|
if !isEOF(b) || len(b) < 5 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
status := binary.LittleEndian.Uint16(b[3:])
|
|
a.observeStatus(status)
|
|
return map[string]any{"type": "metadata", "index": a.index, "warnings": binary.LittleEndian.Uint16(b[1:]), "status": status}, nil
|
|
case "row":
|
|
if isEOF(b) {
|
|
if len(b) < 5 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
status := binary.LittleEndian.Uint16(b[3:])
|
|
a.observeStatus(status)
|
|
return map[string]any{"type": "end", "index": a.index, "rows": strconv.FormatUint(a.setRows, 10), "warnings": binary.LittleEndian.Uint16(b[1:]), "status": status}, nil
|
|
}
|
|
a.rows++
|
|
a.setRows++
|
|
values, err := a.row(b)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return map[string]any{"type": "row", "index": a.index, "values": values}, nil
|
|
}
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
|
|
func resultInteger(b []byte, pos *int) (uint64, error) {
|
|
if *pos >= len(b) {
|
|
return 0, mysql.ErrMalformPacket
|
|
}
|
|
v, null, n := mysql.LengthEncodedInt(b[*pos:])
|
|
if null || n == 0 || n > len(b)-*pos || b[*pos] == 0xff {
|
|
return 0, mysql.ErrMalformPacket
|
|
}
|
|
*pos += n
|
|
return v, nil
|
|
}
|
|
|
|
func resultValue(data []byte, binaryValue bool) any {
|
|
if binaryValue || !utf8.Valid(data) {
|
|
return map[string]any{"base64": base64.StdEncoding.EncodeToString(data)}
|
|
}
|
|
return string(data)
|
|
}
|
|
|
|
func binaryColumn(f *mysql.Field) bool {
|
|
switch f.Type {
|
|
case mysql.MYSQL_TYPE_BIT, mysql.MYSQL_TYPE_GEOMETRY, mysql.MYSQL_TYPE_VECTOR:
|
|
return true
|
|
case mysql.MYSQL_TYPE_VARCHAR, mysql.MYSQL_TYPE_VAR_STRING, mysql.MYSQL_TYPE_STRING, mysql.MYSQL_TYPE_ENUM, mysql.MYSQL_TYPE_SET, mysql.MYSQL_TYPE_TINY_BLOB, mysql.MYSQL_TYPE_MEDIUM_BLOB, mysql.MYSQL_TYPE_LONG_BLOB, mysql.MYSQL_TYPE_BLOB:
|
|
collation, err := charset.GetCollationByID(int(f.Charset))
|
|
if err != nil {
|
|
return true
|
|
}
|
|
return collation.CharsetName != "utf8" && collation.CharsetName != "utf8mb3" && collation.CharsetName != "utf8mb4" && collation.CharsetName != "ascii"
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (a *queryResultAudit) row(b []byte) ([]any, error) {
|
|
values := make([]any, len(a.fields))
|
|
for _, field := range a.fields {
|
|
if field == nil {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
}
|
|
if !a.binary {
|
|
pos := 0
|
|
for i, field := range a.fields {
|
|
if pos >= len(b) {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
value, null, n, err := mysql.LengthEncodedString(b[pos:])
|
|
if err != nil || n == 0 {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
pos += n
|
|
if !null {
|
|
values[i] = resultValue(value, binaryColumn(field))
|
|
}
|
|
}
|
|
if pos != len(b) {
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
return values, nil
|
|
}
|
|
parsed, err := mysql.RowData(b).ParseBinary(a.fields, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for i, value := range parsed {
|
|
switch value.Type {
|
|
case mysql.FieldValueTypeNull:
|
|
case mysql.FieldValueTypeUnsigned:
|
|
values[i] = strconv.FormatUint(value.AsUint64(), 10)
|
|
case mysql.FieldValueTypeSigned:
|
|
values[i] = strconv.FormatInt(value.AsInt64(), 10)
|
|
case mysql.FieldValueTypeFloat:
|
|
bits := 64
|
|
if a.fields[i].Type == mysql.MYSQL_TYPE_FLOAT {
|
|
bits = 32
|
|
}
|
|
values[i] = strconv.FormatFloat(value.AsFloat64(), 'g', -1, bits)
|
|
case mysql.FieldValueTypeString:
|
|
values[i] = resultValue(value.AsString(), binaryColumn(a.fields[i]))
|
|
default:
|
|
return nil, mysql.ErrMalformPacket
|
|
}
|
|
}
|
|
return values, nil
|
|
}
|