Golang

haproxy_stick_table.go

package main

import (
        "context"
        "fmt"
        "io"
        "log"
        "net"
        "net/netip"
        "regexp"
        "strconv"
        "strings"
        "time"
)

const unixsocket string = "/var/lib/haproxy/stats"
const timeout = 5 * time.Second

// Send command to HAProxy UNIX socket and return the response
func sendCommand(table string, socketFile string, storeType string, minRequestRate int) (string, error) {
        if table == "" || socketFile == "" || storeType == "" {
                return "", fmt.Errorf("Invalid parameters: table='%s', socketFile='%s', storeType='%s'", table, socketFile, storeType)
        }
        if minRequestRate < 0 {
                return "", fmt.Errorf("Invalid minRequestRate: %d", minRequestRate)
        }

        var d net.Dialer
        ctx, cancel := context.WithTimeout(context.Background(), timeout)
        defer cancel()

        d.LocalAddr = nil
        raddr := net.UnixAddr{Name: socketFile}
        conn, err := d.DialContext(ctx, "unix", raddr.String())
        if err != nil {
                return "", fmt.Errorf("Failed to connect to %s UNIX socket: %v", socketFile, err)
        }
        defer conn.Close()

        if err := conn.SetDeadline(time.Now().Add(timeout)); err != nil {
                return "", err
        }
        cmd := fmt.Sprintf("show table %s data.%s gt %d\n", table, storeType, minRequestRate)
        if _, err := conn.Write([]byte(cmd)); err != nil {
                return "", fmt.Errorf("Failed to send command to socket: %v", err)
        }

        buf := make([]byte, 1024)
        var data strings.Builder
        for {
                n, err := conn.Read(buf)
                if err != nil {
                        if err == io.EOF {
                                break
                        }
                        return "", fmt.Errorf("Error reading from socket: %v", err)
                }
                data.Write(buf[0:n])
        }
        r := strings.TrimSuffix(data.String(), "\n> ")
        r = strings.TrimSuffix(r, "\n")
        r = strings.TrimSpace(r)

        return r, nil
}

func parse(response string, expectedStoreDataType string) (map[netip.Addr]int, error) {

        requests := make(map[netip.Addr]int)
        if response == "" {
                return nil, fmt.Errorf("Response is empty or malformed")
        }

        lines := strings.Split(response, "\n")
        if len(lines) < 2 {
                return requests, nil
        }

        // Find the data type of stick table.
        // Read http://docs.haproxy.org/dev/configuration.html#4.2-stick-table%20type

        // **Please note** that someone can define multiple data types and that it will change
        // the entries we get.
        // For instance, let's say we have below configuration where we set conn_cnt and
        // http_req_rate data types for the value of entries in the stick table:
        // backend table_requests_limiter_src_ip
        // stick-table type ip size 1m expire 60s store http_req_rate(60s),conn_cnt

        // the response will contain lines like the one below:
        // 0x7fcf0c057200: key=127.0.0.1 use=0 exp=58330 shard=0 conn_cnt=3 http_req_rate(60000)=3

        // The regex we use we does not parse this, therefore we can only support a stick
        // table where we set only one data type and for increment rate types.
        exampleLine := lines[1]
        // Matches a line as the below:
        // 0x7fcf0c057200: key=127.0.0.1 use=0 exp=58330 shard=0 http_req_rate(60000)=3
        e := regexp.MustCompile(
                `^` +
                        `\s*0x[[:alnum:]]+: ` + // Match the entry start with a hexadecimal address
                        `key=(?P<ip>[0-9a-fA-F:.]+) ` + // Match and capture the IP address; 1st group
                        `use=[[:digit:]]+ ` + // Match the use count
                        `exp=[[:digit:]]+ ` + // Match the expiration time
                        `shard=[[:digit:]]+` + // Match the shard value
                        `(?: gpc\d=\d+)? ` + // Optionally match gpc field
                        `(?P<storeType>[[:alnum:]_]+)` + // Match and capture the store type; 2nd group
                        `\([[:digit:]]+\)=(?P<rate>[[:digit:]]+)$`, // Match and capture the rate; 3rd group
        )
        m := e.FindStringSubmatch(exampleLine)

        if len(m) != 4 {
                return requests, fmt.Errorf("Failed to parse entries in the stick table")
        }
        storeType := m[2]
        if storeType != expectedStoreDataType {
                return requests, fmt.Errorf("Store type mismatch: expected '%s', but found '%s'", expectedStoreDataType, storeType)
        }

        for i := 0; i < len(lines); i++ {
                m := e.FindStringSubmatch(lines[i])

                if len(m) == 4 {
                        groups := make(map[string]string)
                        for i, name := range e.SubexpNames() {
                                if name != "" {
                                        groups[name] = m[i]
                                }
                        }

                        ip, err := netip.ParseAddr(groups["ip"])
                        if err != nil {
                                return nil, fmt.Errorf("Failed to parse IP address: %v", err)
                        }

                        rate, err := strconv.Atoi(groups["rate"])
                        if err != nil {
                                return nil, fmt.Errorf("Failed to parse rate: %v", err)
                        }
                        if _, ok := requests[ip]; ok {
                                return nil, fmt.Errorf("Duplicate key detected: %s", ip)
                        }

                        requests[ip] = rate
                }
        }

        return requests, nil
}

// Check if the response is a stick-table of ip type and of expected name
func validateHeader(response string, expectedTableName string) error {
        lines := strings.Split(response, "\n")

        if len(lines) == 0 {
                return fmt.Errorf("Response is empty or malformed")
        }

        header := lines[0]
        // The first line must look like the one below, yes it starts with a #
        // # table: table_requests_limiter_src_ip, type: ip, size:1048576, used:2
        r := regexp.MustCompile(`^#\s+table:\s*(?P<tableName>[\w\-.]+)\s*,\s*type:\s*(?P<tableType>[[:alpha:]]+),`)
        m := r.FindStringSubmatch(header)

        if len(m) != 3 {
                return fmt.Errorf("Failed to parse table header, got '%s'", header)
        }

        tableName := m[1]
        tableType := m[2]
        if tableName != expectedTableName {
                return fmt.Errorf("Table name mismatch. Expected '%s', got '%s'", expectedTableName, tableName)
        }
        if tableType != "ip" {
                return fmt.Errorf("Unsupported table type '%s'. Only 'ip' type is supported", tableType)
        }

        return nil
}
func main() {

        response, err := sendCommand("table_requests_limiter_src_ip", unixsocket, "http_req_rate", 1)
        if err != nil {
                log.Fatal(err)
        }
        fmt.Printf("Received: %s\n", response)
        if err := validateHeader(response, "table_requests_limiter_src_ip"); err != nil {
                fmt.Println(err)
        }
        requests, err := parse(response, "http_req_rate")
        if err != nil {
                fmt.Println(err)
        }
        fmt.Println(requests)
}