mirror of
https://github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin.git
synced 2026-07-21 19:48:59 +02:00
✨ redis included
This commit is contained in:
Max Lerebourg
+165
-342
@@ -1,375 +1,198 @@
|
||||
package redis
|
||||
|
||||
// My intepretation of the RESP protocol.
|
||||
// Referenced from https://redis.io/docs/reference/protocol-spec/
|
||||
// 8-7-2022
|
||||
package simpleredis
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/textproto"
|
||||
"net"
|
||||
"bufio"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/tehnerd/goUtils/netutils"
|
||||
|
||||
logger "github.com/maxlerebourg/crowdsec-bouncer-traefik-plugin/pkg/logger"
|
||||
)
|
||||
|
||||
type RespT byte
|
||||
|
||||
const (
|
||||
SimpleString RespT = '+'
|
||||
Error RespT = '-'
|
||||
Integer RespT = ':'
|
||||
BulkString RespT = '$'
|
||||
Array RespT = '*'
|
||||
)
|
||||
|
||||
type Message interface {
|
||||
String() string
|
||||
type RedisCmd struct {
|
||||
Command string
|
||||
Name string
|
||||
Data []byte
|
||||
Error error
|
||||
}
|
||||
|
||||
type MsgSimpleStr string
|
||||
type MsgError string
|
||||
type MsgBulkStr string
|
||||
type MsgInteger int64
|
||||
type MsgArray []Message
|
||||
|
||||
func (m MsgInteger) String() string {
|
||||
return fmt.Sprintf("%d", m)
|
||||
type SimpleRedis struct {
|
||||
redisChanRead chan RedisCmd
|
||||
redisChanWrite chan RedisCmd
|
||||
redisHost string
|
||||
redisCmd RedisCmd
|
||||
}
|
||||
|
||||
func (m MsgBulkStr) String() string {
|
||||
return string(m)
|
||||
}
|
||||
|
||||
func (m MsgError) String() string {
|
||||
return string(m)
|
||||
}
|
||||
|
||||
func (m MsgSimpleStr) String() string {
|
||||
return string(m)
|
||||
}
|
||||
|
||||
// func (m MsgArray) String() string {
|
||||
// return fmt.Sprintf("[%s]", strings.Join(Map(m, Message.String), ","))
|
||||
// }
|
||||
|
||||
func (m RedisMessage) String() string {
|
||||
return m.Choice.String()
|
||||
}
|
||||
|
||||
type RedisMessage struct {
|
||||
RedisType RespT
|
||||
Raw string
|
||||
Choice Message
|
||||
}
|
||||
|
||||
func NewParseError(got RespT, expected ...RespT) error {
|
||||
return fmt.Errorf("got %v but expected %v", got, expected)
|
||||
}
|
||||
|
||||
func (msg RedisMessage) Symbol() rune {
|
||||
return rune(msg.RedisType)
|
||||
}
|
||||
|
||||
func (msg RedisMessage) AsString() (string, error) {
|
||||
if conv, ok := msg.Choice.(MsgSimpleStr); ok {
|
||||
return string(conv), nil
|
||||
} else if conv, ok := msg.Choice.(MsgBulkStr); ok {
|
||||
return string(conv), nil
|
||||
} else {
|
||||
return "", NewParseError(msg.RedisType, SimpleString, BulkString)
|
||||
func GenRedisArray(params ...[]byte) []byte {
|
||||
CRLF := "\r\n"
|
||||
MSG := ""
|
||||
for cntr := 0; cntr < len(params); cntr++ {
|
||||
MSG = strings.Join([]string{MSG, string(params[cntr])}, " ")
|
||||
}
|
||||
MSG = strings.Trim(MSG, " ")
|
||||
MSG = strings.Join([]string{MSG, CRLF}, "")
|
||||
return []byte(MSG)
|
||||
}
|
||||
|
||||
func (msg RedisMessage) AsInteger() (int64, error) {
|
||||
if conv, ok := msg.Choice.(MsgInteger); ok {
|
||||
return int64(conv), nil
|
||||
} else {
|
||||
return 0, NewParseError(msg.RedisType, Integer)
|
||||
}
|
||||
func RedisSet(name string, data []byte) []byte {
|
||||
return GenRedisArray([]byte("SET"), []byte(name), data)
|
||||
}
|
||||
|
||||
func (msg RedisMessage) AsError() (error, error) {
|
||||
if conv, ok := msg.Choice.(MsgError); ok {
|
||||
return errors.New(string(conv)), nil
|
||||
} else {
|
||||
return nil, NewParseError(msg.RedisType, Error)
|
||||
}
|
||||
func RedisGet(name string) []byte {
|
||||
return GenRedisArray([]byte("GET"), []byte(name))
|
||||
}
|
||||
|
||||
// func (msg RedisMessage) AsArray() ([]Message, error) {
|
||||
// if conv, ok := (msg.Choice).(MsgArray); ok {
|
||||
// return conv, nil
|
||||
// } else {
|
||||
// return nil, NewParseError(msg.RedisType, Array)
|
||||
// }
|
||||
// }
|
||||
|
||||
type RedisReader struct {
|
||||
Rd *textproto.Reader
|
||||
Out chan RedisMessage
|
||||
}
|
||||
|
||||
// Reset the parsing state of the reader
|
||||
func (rr *RedisReader) Reset() {
|
||||
|
||||
}
|
||||
|
||||
func NewRespReader(tp *textproto.Reader, out chan RedisMessage) RedisReader {
|
||||
return RedisReader{
|
||||
Rd: tp,
|
||||
Out: out,
|
||||
func ParseRedisResponse(response []byte, dataBuf []byte, Len *int) ([]byte, []byte, error) {
|
||||
dataBuf = append(dataBuf, response...)
|
||||
lenCRLF := 2
|
||||
if *Len != 0 {
|
||||
if len(dataBuf) < *Len {
|
||||
return nil, dataBuf, nil
|
||||
} else {
|
||||
return dataBuf[:*Len], dataBuf[*Len:], nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fetch grabs the next incoming RESP type from the IO input
|
||||
func Fetch(rr *RedisReader) *RedisMessage {
|
||||
byteCode, err := rr.TryReadByte()
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
switch RespT(byteCode) {
|
||||
case BulkString:
|
||||
return RespBulkStr(rr)
|
||||
case SimpleString:
|
||||
return RespSimpleStr(rr)
|
||||
case Integer:
|
||||
return RespInt(rr)
|
||||
case Error:
|
||||
return RespError(rr)
|
||||
// case Array:
|
||||
// return RespArray(rr)
|
||||
default:
|
||||
panic(fmt.Sprint("Unknown bytecode: ", byteCode))
|
||||
}
|
||||
}
|
||||
|
||||
// Scan is the main API to send redis messages out to a channel.
|
||||
func (rr *RedisReader) Scan() {
|
||||
rr.Out <- *Fetch(rr)
|
||||
}
|
||||
|
||||
// RespSimpleStr attempts to parse a string from the resulting line.
|
||||
//
|
||||
// TODO: I don't think we need the high level facilities from net/textproto anymore
|
||||
// since the parsing is quite trivial.
|
||||
func RespSimpleStr(rr *RedisReader) *RedisMessage {
|
||||
s, err := rr.Rd.ReadLine()
|
||||
if err != nil {
|
||||
log.Printf("Fatal RespSimpleStr: %+v", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
return &RedisMessage{
|
||||
RedisType: RespT(SimpleString),
|
||||
Raw: s,
|
||||
Choice: MsgSimpleStr(s),
|
||||
}
|
||||
}
|
||||
|
||||
func RespError(rr *RedisReader) *RedisMessage {
|
||||
s, err := rr.Rd.ReadLine()
|
||||
if err != nil {
|
||||
log.Printf("Fatal RespError: %+v", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
return &RedisMessage{
|
||||
RedisType: RespT(Error),
|
||||
Raw: s,
|
||||
Choice: MsgError(s),
|
||||
}
|
||||
}
|
||||
|
||||
// RespSimpleStr attempts to parse a string from the resulting line.
|
||||
//
|
||||
// It first expects a length value parsed as an integer,
|
||||
// then directly copies the subsequent payload into a string
|
||||
// buffer
|
||||
func RespBulkStr(rr *RedisReader) *RedisMessage {
|
||||
len, err := loopReadInt(rr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Read the string data into buffer
|
||||
bulkStr := make([]byte, len)
|
||||
io.ReadFull(rr.Rd.R, bulkStr)
|
||||
|
||||
// Discard two bytes since we don't need the '\r\n'
|
||||
_, err = rr.Rd.R.Discard(2)
|
||||
if err != nil {
|
||||
log.Printf("Fatal: %+v", err)
|
||||
}
|
||||
|
||||
return &RedisMessage{
|
||||
RedisType: BulkString,
|
||||
Raw: string(bulkStr),
|
||||
Choice: MsgBulkStr(string(bulkStr)),
|
||||
}
|
||||
}
|
||||
|
||||
// RespInt reads an integer from the IO stream
|
||||
func RespInt(rr *RedisReader) *RedisMessage {
|
||||
val, err := loopReadInt(rr)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return &RedisMessage{
|
||||
RedisType: Integer,
|
||||
Raw: fmt.Sprintf("%d", val),
|
||||
Choice: MsgInteger(val),
|
||||
}
|
||||
}
|
||||
|
||||
// RespArray reads an array of RESP objects from the IO stream
|
||||
// func RespArray(rr *RedisReader) *RedisMessage {
|
||||
// len, err := loopReadInt(rr)
|
||||
// if err != nil {
|
||||
// return nil
|
||||
// }
|
||||
// fetched := make([]RedisMessage, 0, len)
|
||||
// for i := 0; i < len; i++ {
|
||||
// fetched[i] = *Fetch(rr)
|
||||
// }
|
||||
// return &RedisMessage{
|
||||
// RedisType: Array,
|
||||
// Raw: fmt.Sprint(Map(fetched, func(rm RedisMessage) string { return rm.Raw })),
|
||||
// Choice: MsgArray(Map(fetched, func(rm RedisMessage) Message { return rm.Choice })),
|
||||
// }
|
||||
// }
|
||||
|
||||
// loopReadInt reads from IO and returns an integer, discarding \r\n.
|
||||
func loopReadInt(rr *RedisReader) (int, error) {
|
||||
val := 0
|
||||
for b, _ := rr.TryReadByte(); b != '\r'; b, _ = rr.TryReadByte() {
|
||||
val = (val * 10) + int(b-byte('0'))
|
||||
}
|
||||
_, err := rr.Rd.R.Discard(1)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("Fatal loopReadInt: %+v", err)
|
||||
}
|
||||
return val, nil
|
||||
}
|
||||
|
||||
// TryReadByte attempts to read a single byte from IO and panics on any error
|
||||
func (rr *RedisReader) TryReadByte() (byte, error) {
|
||||
b, err := rr.Rd.R.ReadByte()
|
||||
if err != nil {
|
||||
return byte('0'), fmt.Errorf("Fatal TryReadByte: %+v", err)
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// Map. Your standard good old fashioned generic map function :)
|
||||
//
|
||||
// A la []T -> []K
|
||||
func Map[T any, K any](slice []T, fn func(T) K) []K {
|
||||
out := make([]K, 0, len(slice))
|
||||
for i := range out {
|
||||
out[i] = fn(slice[i])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
func loopScan(reader *textproto.Reader, msg_ch chan RedisMessage) {
|
||||
rr := NewRespReader(reader, msg_ch)
|
||||
for {
|
||||
rr.Scan()
|
||||
switch string(dataBuf[0]) {
|
||||
case "+", "-", ":":
|
||||
//simple strings, error,int. usually ther are in format (+|-|:)DATA\r\n"
|
||||
if len(dataBuf) < 3 {
|
||||
return nil, dataBuf, nil
|
||||
}
|
||||
cntr := 1
|
||||
for ; cntr < len(dataBuf); cntr++ {
|
||||
if dataBuf[cntr] == '\r' {
|
||||
break
|
||||
}
|
||||
}
|
||||
if cntr == len(dataBuf) {
|
||||
return nil, dataBuf, nil
|
||||
}
|
||||
response = dataBuf[1:cntr]
|
||||
return response, dataBuf[cntr+2:], nil
|
||||
case "$":
|
||||
//bulk string. format $LEN\r\nDATA\r\n. up to 512MB
|
||||
cntr := 1
|
||||
for ; cntr < len(dataBuf); cntr++ {
|
||||
if string(dataBuf[cntr]) == "\r" {
|
||||
break
|
||||
}
|
||||
}
|
||||
if cntr == len(dataBuf) || cntr+lenCRLF > len(dataBuf) {
|
||||
return nil, dataBuf, nil
|
||||
}
|
||||
dataLen, err := strconv.Atoi(string(dataBuf[1:cntr]))
|
||||
if err != nil {
|
||||
return nil, dataBuf[cntr:], nil
|
||||
}
|
||||
|
||||
if dataLen == -1 {
|
||||
return nil, dataBuf[cntr:], fmt.Errorf("NOT FOUND")
|
||||
}
|
||||
if cntr+lenCRLF > len(dataBuf)-lenCRLF {
|
||||
*Len = dataLen
|
||||
return nil, dataBuf[cntr+lenCRLF:], nil
|
||||
}
|
||||
if len(dataBuf[cntr+lenCRLF:len(dataBuf)-lenCRLF]) < dataLen {
|
||||
*Len = dataLen
|
||||
return nil, dataBuf[cntr+lenCRLF:], nil
|
||||
} else {
|
||||
return dataBuf[cntr+lenCRLF : cntr+dataLen+lenCRLF], dataBuf[cntr+dataLen+lenCRLF:], nil
|
||||
}
|
||||
case "*":
|
||||
panic("array")
|
||||
default:
|
||||
if len(dataBuf) > 1 {
|
||||
dataBuf = dataBuf[1:]
|
||||
} else {
|
||||
return nil, dataBuf, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, dataBuf, nil
|
||||
}
|
||||
|
||||
// repl starts an read input loop on the standard input.
|
||||
//
|
||||
// Commands are sent out after a new line is parsed.
|
||||
//
|
||||
// TODO: Improve ergonomics of the repl.
|
||||
func repl(input *textproto.Reader, wr *textproto.Writer) {
|
||||
for {
|
||||
print("> ")
|
||||
input_str, err := input.ReadLine()
|
||||
if len(input_str) > 0 && err == nil {
|
||||
// The redis server expects arrays of bulk strings for sending commands.
|
||||
commands := strings.Split(input_str, " ")
|
||||
// Specify the array length and create the bulk strings
|
||||
pipe := []string{fmt.Sprintf("*%d", len(commands))}
|
||||
pipe = append(pipe, createBulkStrings(commands)...)
|
||||
for _, out := range pipe {
|
||||
wr.PrintfLine(out)
|
||||
func RedisContext(hostnamePort string, redisCmdWrite, redisCmdRead chan RedisCmd) {
|
||||
tcpRemoteAddress, err := net.ResolveTCPAddr("tcp", hostnamePort)
|
||||
if err != nil {
|
||||
panic("cant resolve remote redis address")
|
||||
}
|
||||
var ladr *net.TCPAddr
|
||||
msgBuf := make([]byte, 65000)
|
||||
initMsg := []byte("*1\r\n$4\r\nPING\r\n")
|
||||
writeChan := make(chan []byte)
|
||||
readChan := make(chan []byte)
|
||||
flushChan := make(chan int)
|
||||
go netutils.AutoRecoonectedTCP(ladr, tcpRemoteAddress, msgBuf, initMsg, writeChan, readChan, flushChan)
|
||||
<-readChan
|
||||
dataBuf := make([]byte, 0)
|
||||
dataLen := 0
|
||||
for true {
|
||||
select {
|
||||
case cmd := <-redisCmdWrite:
|
||||
switch cmd.Command {
|
||||
case "SET":
|
||||
data := RedisSet(cmd.Name, cmd.Data)
|
||||
writeChan <- data
|
||||
case "GET":
|
||||
data := RedisGet(cmd.Name)
|
||||
writeChan <- data
|
||||
}
|
||||
case response := <-readChan:
|
||||
data, dataBuf, err := ParseRedisResponse(response, dataBuf, &dataLen)
|
||||
if dataLen != 0 {
|
||||
for data == nil {
|
||||
response = <-readChan
|
||||
data, dataBuf, err = ParseRedisResponse(response, dataBuf, &dataLen)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
select {
|
||||
case redisCmdRead <- RedisCmd{
|
||||
Error: err,
|
||||
}:
|
||||
case <-time.After(time.Second * 5):
|
||||
}
|
||||
}
|
||||
if data != nil && string(data) != "PONG" {
|
||||
select {
|
||||
case redisCmdRead <- RedisCmd{
|
||||
Data: data,
|
||||
}:
|
||||
case <-time.After(time.Second * 5):
|
||||
}
|
||||
dataLen = 0
|
||||
}
|
||||
case <-flushChan:
|
||||
dataBuf = dataBuf[:]
|
||||
dataLen = 0
|
||||
select {
|
||||
case redisCmdRead <- RedisCmd{}:
|
||||
case <-time.After(time.Second * 5):
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// createBulkStrings converts normal strings into bulk strings to send to redis.
|
||||
func createBulkStrings(commands []string) []string {
|
||||
cmds := make([]string, len(commands)*2)
|
||||
for i, v := range commands {
|
||||
cmds[i*2] = fmt.Sprintf("$%d", len(commands[i]))
|
||||
cmds[(i*2)+1] = v
|
||||
}
|
||||
return cmds
|
||||
func (sr *SimpleRedis) Init(redisHost string) {
|
||||
sr.redisHost = redisHost
|
||||
sr.redisChanWrite = make(chan RedisCmd)
|
||||
sr.redisChanRead = make(chan RedisCmd)
|
||||
go RedisContext(sr.redisHost, sr.redisChanWrite, sr.redisChanRead)
|
||||
}
|
||||
|
||||
|
||||
// cmd/client starts a simple Read-Send-Print-Loop (RSPL?)
|
||||
// with a redis server.
|
||||
//
|
||||
// TODO: Implement command flags
|
||||
func Init(host string, password string) (*textproto.Writer, *textproto.Reader) {
|
||||
log.Printf("yo")
|
||||
// Create the connection to the redis server
|
||||
// TODO: cli, rm hardcode
|
||||
conn, err := net.Dial("tcp", host)
|
||||
if err != nil {
|
||||
log.Printf("[net] unable to connect: %v", err)
|
||||
return nil, nil
|
||||
func (sr *SimpleRedis) Do(cmd, name string, data []byte) ([]byte, error) {
|
||||
sr.redisCmd.Command = cmd
|
||||
sr.redisCmd.Name = name
|
||||
sr.redisCmd.Data = data
|
||||
sr.redisChanWrite <- sr.redisCmd
|
||||
resp := <- sr.redisChanRead
|
||||
if resp.Error != nil {
|
||||
return nil, resp.Error
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// Create the readers and writers we will pass to our subroutines
|
||||
msg_channel := make(chan RedisMessage)
|
||||
write_sock := textproto.NewWriter(bufio.NewWriter(conn))
|
||||
read_sock := textproto.NewReader(bufio.NewReader(conn))
|
||||
|
||||
go runAsWaitGroup(
|
||||
// Subscribe to redis connection
|
||||
func() {
|
||||
loopScan(read_sock, msg_channel)
|
||||
},
|
||||
// Write to output
|
||||
func() {
|
||||
for {
|
||||
msg := <-msg_channel
|
||||
log.Printf("%v", msg)
|
||||
}
|
||||
},
|
||||
).Wait()
|
||||
log.Printf("yolo")
|
||||
return write_sock, read_sock
|
||||
}
|
||||
|
||||
// runAsWaitGroup runs closures within a sync.WaitGroup
|
||||
//
|
||||
// Calling .Wait() on the resultant WaitGroup is a blocking operation until all closures finish running.
|
||||
func runAsWaitGroup(closures ...func()) *sync.WaitGroup {
|
||||
wg := sync.WaitGroup{}
|
||||
for _, fn := range closures {
|
||||
wg.Add(1)
|
||||
go func(fn func()) {
|
||||
fn()
|
||||
wg.Done()
|
||||
}(fn)
|
||||
}
|
||||
return &wg
|
||||
return resp.Data, nil
|
||||
}
|
||||
Reference in New Issue
Block a user