Merge pull request #51 from yihong0618/postgres_probe

feat: support postgres query hook

.gitignore

@@ -14,3 +14,6 @@
 # Dependency directories (remove the comment below to include it)
 # vendor/
 /assets/ebpf_probe.go
+
+# VSCode
+.vscode/settings.json

Makefile

@@ -196,6 +196,7 @@ TARGETS += kern/bash
 TARGETS += kern/gnutls
 TARGETS += kern/nspr
 TARGETS += kern/mysqld
+TARGETS += kern/postgres
 
 # Generate file name-scheme based on TARGETS
 KERN_SOURCES = ${TARGETS:=_kern.c}

cli/cmd/postgres.go

@@ -0,0 +1,78 @@
+/*
+Copyright © 2022 CFC4N <cfc4n.cs@gmail.com>
+
+*/
+package cmd
+
+import (
+	"context"
+	"ecapture/user"
+	"log"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/spf13/cobra"
+)
+
+var postgresConfig = user.NewPostgresConfig()
+
+//postgres Cmd represents the postgres command
+var postgresCmd = &cobra.Command{
+	Use:   "postgres",
+	Short: "capture sql queries from postgres 10+.",
+	Run:   postgresCommandFunc,
+}
+
+func init() {
+	postgresCmd.PersistentFlags().StringVarP(&postgresConfig.PostgresPath, "postgres", "m", "/usr/bin/postgres", "postgres binary file path, use to hook")
+	postgresCmd.PersistentFlags().StringVarP(&postgresConfig.FuncName, "funcname", "f", "", "function name to hook")
+	rootCmd.AddCommand(postgresCmd)
+}
+
+// postgres CommandFunc executes the "psql" command.
+func postgresCommandFunc(command *cobra.Command, args []string) {
+	stopper := make(chan os.Signal, 1)
+	signal.Notify(stopper, os.Interrupt, syscall.SIGTERM)
+	ctx, cancelFun := context.WithCancel(context.TODO())
+
+	mod := user.GetModuleByName(user.MODULE_NAME_POSTGRES)
+
+	logger := log.Default()
+
+	logger.Printf("start to run %s module", mod.Name())
+
+	// save global config
+	gConf, e := getGlobalConf(command)
+	if e != nil {
+		logger.Fatal(e)
+		os.Exit(1)
+	}
+	postgresConfig.Pid = gConf.Pid
+	postgresConfig.Debug = gConf.Debug
+	postgresConfig.IsHex = gConf.IsHex
+
+	log.Printf("pid info: %d", os.Getpid())
+	//bc.Pid = globalFlags.Pid
+	if e := postgresConfig.Check(); e != nil {
+		logger.Fatal(e)
+		os.Exit(1)
+	}
+	// init
+	err := mod.Init(ctx, logger, postgresConfig)
+	if err != nil {
+		logger.Fatal(err)
+		os.Exit(1)
+	}
+
+	// 加载ebpf，挂载到hook点上，开始监听
+	go func(module user.IModule) {
+		err := module.Run()
+		if err != nil {
+			logger.Fatalf("%v", err)
+		}
+	}(mod)
+	<-stopper
+	cancelFun()
+	os.Exit(0)
+}

kern/common.h

@@ -14,6 +14,7 @@
 #define TASK_COMM_LEN 16
 #define MAX_DATA_SIZE_OPENSSL 1024 * 4
 #define MAX_DATA_SIZE_MYSQL 256
+#define MAX_DATA_SIZE_POSTGRES 256
 
 // enum_server_command, via
 // https://dev.mysql.com/doc/internals/en/com-query.html COM_QUERT command 03

kern/postgres_kern.c

@@ -0,0 +1,41 @@
+#include "ecapture.h"
+
+struct data_t {
+    u64 pid;
+    u64 timestamp;
+    char query[MAX_DATA_SIZE_POSTGRES];
+    char comm[TASK_COMM_LEN];
+};
+
+struct
+{
+    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
+} events SEC(".maps");
+
+// https://github.com/postgres/postgres/blob/7b7ed046cb2ad9f6efac90380757d5977f0f563f/src/backend/tcop/postgres.c#L987-L992
+// hook function exec_simple_query
+// versions 10 - now
+// static void exec_simple_query(const char *query_string)
+SEC("uprobe/exec_simple_query")
+int postgres_query(struct pt_regs *ctx) {
+
+    u64 current_pid_tgid = bpf_get_current_pid_tgid();
+    u32 pid = current_pid_tgid >> 32;
+
+#ifndef KERNEL_LESS_5_2
+    // if target_ppid is 0 then we target all pids
+    if (target_pid != 0 && target_pid != pid) {
+        return 0;
+    }
+#endif
+
+    struct data_t data = {};
+    data.pid = pid;   // only process id
+    data.timestamp = bpf_ktime_get_ns();
+
+    char *sql_string= (char *)PT_REGS_PARM1(ctx);
+    bpf_get_current_comm(&data.comm, sizeof(data.comm));
+    bpf_probe_read(&data.query, sizeof(data.query), sql_string);
+    bpf_perf_event_output(ctx, &events, BPF_F_CURRENT_CPU, &data, sizeof(data));
+    return 0;
+}

user/config_postgres.go

@@ -0,0 +1,38 @@
+/*
+Copyright © 2022 CFC4N <cfc4n.cs@gmail.com>
+
+*/
+package user
+
+import (
+	"os"
+	"strings"
+
+	"github.com/pkg/errors"
+)
+
+type PostgresConfig struct {
+	eConfig
+	PostgresPath string `json:"postgresPath"`
+	FuncName     string `json:"funcName"`
+}
+
+func NewPostgresConfig() *PostgresConfig {
+	config := &PostgresConfig{}
+	return config
+}
+
+func (this *PostgresConfig) Check() error {
+
+	if this.PostgresPath == "" || len(strings.TrimSpace(this.PostgresPath)) <= 0 {
+		return errors.New("Postgres path cant be null.")
+	}
+
+	_, e := os.Stat(this.PostgresPath)
+	if e != nil {
+		return e
+	}
+
+	this.FuncName = "exec_simple_query"
+	return nil
+}

user/const.go

@@ -8,11 +8,12 @@ const (
 )
 
 const (
-	MODULE_NAME_BASH    = "EBPFProbeBash"
-	MODULE_NAME_MYSQLD  = "EBPFProbeMysqld"
-	MODULE_NAME_OPENSSL = "EBPFProbeOPENSSL"
-	MODULE_NAME_GNUTLS  = "EBPFProbeGNUTLS"
-	MODULE_NAME_NSPR    = "EBPFProbeNSPR"
+	MODULE_NAME_BASH     = "EBPFProbeBash"
+	MODULE_NAME_MYSQLD   = "EBPFProbeMysqld"
+	MODULE_NAME_POSTGRES = "EBPFProbePostgres"
+	MODULE_NAME_OPENSSL  = "EBPFProbeOPENSSL"
+	MODULE_NAME_GNUTLS   = "EBPFProbeGNUTLS"
+	MODULE_NAME_NSPR     = "EBPFProbeNSPR"
 )
 
 const (

user/event_bash.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+
 	"golang.org/x/sys/unix"
 )
 

user/event_mysqld.go

@@ -8,16 +8,17 @@ import (
 	"bytes"
 	"encoding/binary"
 	"fmt"
+
 	"golang.org/x/sys/unix"
 )
 
 /*
-	u64 pid;
-    u64 timestamp;
-    char query[MAX_DATA_SIZE];
-    u64 alllen;
-    u64 len;
-    char comm[TASK_COMM_LEN];
+   u64 pid;
+   u64 timestamp;
+   char query[MAX_DATA_SIZE];
+   u64 alllen;
+   u64 len;
+   char comm[TASK_COMM_LEN];
 */
 const MYSQLD_MAX_DATA_SIZE = 256
 

user/event_postgres.go

@@ -0,0 +1,68 @@
+/*
+Copyright © 2022 CFC4N <cfc4n.cs@gmail.com>
+
+*/
+package user
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+
+	"golang.org/x/sys/unix"
+)
+
+/*
+   u64 pid;
+   u64 timestamp;
+   char query[MAX_DATA_SIZE];
+   char comm[TASK_COMM_LEN];
+*/
+const POSTGRES_MAX_DATA_SIZE = 256
+
+type postgresEvent struct {
+	module    IModule
+	Pid       uint64
+	Timestamp uint64
+	query     [POSTGRES_MAX_DATA_SIZE]uint8
+	comm      [16]uint8
+}
+
+func (this *postgresEvent) Decode(payload []byte) (err error) {
+	buf := bytes.NewBuffer(payload)
+	if err = binary.Read(buf, binary.LittleEndian, &this.Pid); err != nil {
+		return
+	}
+	if err = binary.Read(buf, binary.LittleEndian, &this.Timestamp); err != nil {
+		return
+	}
+	if err = binary.Read(buf, binary.LittleEndian, &this.query); err != nil {
+		return
+	}
+	if err = binary.Read(buf, binary.LittleEndian, &this.comm); err != nil {
+		return
+	}
+	return nil
+}
+
+func (this *postgresEvent) String() string {
+	s := fmt.Sprintf(fmt.Sprintf(" PID: %d, Comm: %s, Time: %d, Query: %s", this.Pid, this.comm, this.Timestamp, unix.ByteSliceToString((this.query[:]))))
+	return s
+}
+
+func (this *postgresEvent) StringHex() string {
+	s := fmt.Sprintf(fmt.Sprintf(" PID: %d, Comm: %s, Time: %d, Query: %s", this.Pid, this.comm, this.Timestamp, unix.ByteSliceToString((this.query[:]))))
+	return s
+}
+
+func (this *postgresEvent) SetModule(module IModule) {
+	this.module = module
+}
+
+func (this *postgresEvent) Module() IModule {
+	return this.module
+}
+
+func (this *postgresEvent) Clone() IEventStruct {
+	return new(postgresEvent)
+}

user/probe_openssl.go

@@ -4,13 +4,14 @@ import (
 	"bytes"
 	"context"
 	"ecapture/assets"
+	"log"
+	"math"
+	"os"
+
 	"github.com/cilium/ebpf"
 	manager "github.com/ehids/ebpfmanager"
 	"github.com/pkg/errors"
 	"golang.org/x/sys/unix"
-	"log"
-	"math"
-	"os"
 )
 
 const CONN_NOT_FOUND = "[ADDR_NOT_FOUND]"