增加命令行参数，修改对应go test

2 files changed, 409 insertions, 0 deletions

diff --git a/coredump-handler/coredump-handler.go b/coredump-handler/coredump-handler.go
new file mode 100644
index 0000000..22fbe5d
--- /dev/null
+++ b/coredump-handler/coredump-handler.go
@@ -0,0 +1,243 @@
+package main
+
+import (
+	"archive/zip"
+	"context"
+	"coredump-handler/config"
+	"coredump-handler/types"
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"regexp"
+	"strconv"
+	"strings"
+	"syscall"
+
+	"github.com/containerd/containerd"
+	"github.com/containerd/containerd/namespaces"
+	"github.com/coreos/go-systemd/v22/journal"
+)
+
+const chunkSize = 1024 * 1024 * 1024 // 1GB
+var coredump_config types.Coredump_config
+
+func argsJudge() error {
+	if coredump_config.Initial_ns_pid == "" || coredump_config.Process_ns_pid == "" || coredump_config.Corepipe_config_path == "" || coredump_config.Timestamp == 0 || coredump_config.Process_exe_path == "" || coredump_config.GID == "" || coredump_config.Hostname == "" || coredump_config.UID == "" || coredump_config.Signal == -1 {
+		err := fmt.Sprintf("Failed to initialize command line parameters. -P=%s -p=%s -E=%s -configpath=%s -t=%d -g=%s -h=%s -s=%d -u=%s", coredump_config.Initial_ns_pid, coredump_config.Process_ns_pid, coredump_config.Process_exe_path, coredump_config.Corepipe_config_path, coredump_config.Timestamp, coredump_config.GID, coredump_config.Hostname, coredump_config.Signal, coredump_config.UID)
+		return errors.New(err)
+	}
+	return nil
+}
+func isDiskSufficient(pipe_config types.Pipeconfig) (bool, error) {
+	percent, err := strconv.ParseFloat(pipe_config.Total_file_mem_limit[:len(pipe_config.Total_file_mem_limit)-1], 64)
+	if err != nil {
+		return false, err
+	}
+	percent = percent / 100.0
+	var stat syscall.Statfs_t
+	wd, err := syscall.Getwd()
+	if err != nil {
+		fmt.Println(err)
+		return false, err
+	}
+	syscall.Statfs(wd, &stat)
+	// 剩余空间的大小为块的数量 * 每个块的大小
+	// stat.Bfree表示可用的块的数量，stat.Bsize表示每个块的大小
+	usedSpace := (int64(stat.Blocks) - int64(stat.Bfree))
+	totalSpace := int64(stat.Blocks)
+	usage := float64(usedSpace) / float64(totalSpace)
+	if usage >= percent {
+		return false, nil
+	}
+	return true, nil
+}
+func createCoreDumpDir(pipe_config *types.Pipeconfig, args types.Coredump_config) error {
+	pipe_config.File_base_path = fmt.Sprintf("%s/%s_%s_%d", pipe_config.File_base_path, args.Initial_ns_pid, args.Process_ns_pid, args.Timestamp)
+	dirName := pipe_config.File_base_path
+	if _, err := os.Stat(dirName); os.IsNotExist(err) {
+		// 目录不存在，创建目录
+		if err := os.MkdirAll(dirName, os.ModePerm); err != nil {
+			return err
+		}
+	} else {
+		return errors.New("directory already exists")
+	}
+	return nil
+}
+func changeDirectory(dir string) error {
+	if err := os.Chdir(dir); err != nil {
+		return err
+	}
+	return nil
+}
+func getContainerId(pid string) (string, error) {
+	cgroup_path := fmt.Sprintf("/proc/%s/cgroup", pid)
+	content, err := ioutil.ReadFile(cgroup_path)
+	if err != nil {
+		return "", err
+	}
+	re := regexp.MustCompile(`([a-f\d]{64})`)
+	match := re.FindStringSubmatch(string(content))
+	if len(match) < 2 {
+		return "", errors.New("failed to extract container ID from cgroup file")
+	}
+	containerID := match[1]
+	return containerID, nil
+}
+func getImageId(container_id string, sock_path string) (string, error) {
+	// 连接 containerd daemon
+	client, err := containerd.New(sock_path)
+	if err != nil {
+		return "", err
+	}
+	defer client.Close()
+	// 根据容器 ID 获取容器信息
+	ctx := namespaces.WithNamespace(context.Background(), "k8s.io")
+	container, err := client.LoadContainer(ctx, container_id)
+	if err != nil {
+		return "", err
+	}
+	// 获取容器关联的镜像信息
+	imageRef, err := container.Image(ctx)
+	if err != nil {
+		return "", err
+	}
+	return imageRef.Name(), nil
+}
+func writeCoreConfig(config types.Coredump_config) error {
+	filename := fmt.Sprintf("%s_%s_%d_coredump.config", config.Initial_ns_pid, config.Process_ns_pid, config.Timestamp)
+	file, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	encoder := json.NewEncoder(file)
+	err = encoder.Encode(&config)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+func writeCoreDumpToFile(config types.Coredump_config) error {
+	filename := fmt.Sprintf("%s_%s_%d_coredump.info", config.Initial_ns_pid, config.Process_ns_pid, config.Timestamp)
+	file, err := os.Create(filename)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	// 创建一个管道
+	reader, writer := io.Pipe()
+	// 将标准输入连接到管道的写入端
+	go func() {
+		defer writer.Close()
+		io.Copy(writer, os.Stdin)
+	}()
+
+	// 从管道的读取端将内容写入文件
+	io.Copy(file, reader)
+	return nil
+}
+func compress(config types.Coredump_config) error {
+	// Create a new zip archive.
+	filename := fmt.Sprintf("%s_%s_%d_coredump.info", config.Initial_ns_pid, config.Process_ns_pid, config.Timestamp)
+	zipfile, err := os.Create(filename + ".zip")
+	if err != nil {
+		return err
+	}
+	defer zipfile.Close()
+
+	// Create a new zip writer.
+	zipwriter := zip.NewWriter(zipfile)
+	defer zipwriter.Close()
+
+	// Create a zip file header.
+	header := &zip.FileHeader{
+		Name:   filename,
+		Method: zip.Deflate,
+	}
+
+	// Write the header to the zip file.
+	writer, err := zipwriter.CreateHeader(header)
+	if err != nil {
+		return err
+	}
+	io.Copy(writer, os.Stdin)
+
+	return nil
+}
+func main() {
+	flag.StringVar(&coredump_config.Initial_ns_pid, "P", "", "initial ns pid")
+	flag.StringVar(&coredump_config.Process_ns_pid, "p", "", "process ns pid")
+	flag.StringVar(&coredump_config.Process_exe_path, "E", "", "pathname of executable process")
+	flag.StringVar(&coredump_config.Corepipe_config_path, "configpath", "", "configfile's path")
+	flag.Int64Var(&coredump_config.Timestamp, "t", 0, "the time of coredump")
+	flag.StringVar(&coredump_config.GID, "g", "", "Numeric real GID of dumped process.")
+	flag.StringVar(&coredump_config.Hostname, "h", "", "Hostname (same as nodename returned by uname).")
+	flag.IntVar(&coredump_config.Signal, "s", -1, "Number of signal causing dump")
+	flag.StringVar(&coredump_config.UID, "u", "", "Numeric real UID of dumped process.")
+	flag.Parse()
+	info := fmt.Sprintf("start handle coredump initialize command line parameters. -P=%s -p=%s -E=%s -configpath=%s -t=%d -g=%s -h=%s -s=%d -u=%s", coredump_config.Initial_ns_pid, coredump_config.Process_ns_pid, coredump_config.Process_exe_path, coredump_config.Corepipe_config_path, coredump_config.Timestamp, coredump_config.GID, coredump_config.Hostname, coredump_config.Signal, coredump_config.UID)
+	journal.Print(journal.PriInfo, info)
+	coredump_config.Process_exe_path = strings.Replace(coredump_config.Process_exe_path, "!", "/", -1)
+	//判断参数读取是否正确
+	err := argsJudge()
+	if err != nil {
+		journal.Print(journal.PriErr, err.Error())
+		return
+	}
+	//读取config文件并初始化
+	pipe_config, err := config.PipeInit(coredump_config.Corepipe_config_path)
+	if err != nil {
+		journal.Print(journal.PriErr, err.Error())
+		return
+	}
+	//判断硬盘使用率
+	flag, err := isDiskSufficient(pipe_config)
+	if err != nil && !flag {
+		journal.Print(journal.PriErr, err.Error())
+		return
+	}
+	//创建存储coredump内容文件夹
+	err = createCoreDumpDir(&pipe_config, coredump_config)
+	if err != nil {
+		journal.Print(journal.PriErr, err.Error())
+		return
+	}
+	coredump_config.Storage = pipe_config.File_base_path
+	//切换至存储coredump目录
+	err = changeDirectory(pipe_config.File_base_path)
+	if err != nil {
+		journal.Print(journal.PriErr, err.Error())
+		return
+	}
+	//查找发生coredump进程对应的container id
+	container_id, err := getContainerId(coredump_config.Initial_ns_pid)
+	//根据查找到的container id查找对应的image name
+	if err == nil && len(container_id) != 0 {
+		coredump_config.Image_id, err = getImageId(container_id, pipe_config.Containered_sock_path)
+		if err != nil {
+			journal.Print(journal.PriInfo, err.Error())
+		}
+	}
+	//将image name写入coredump config
+	err = writeCoreConfig(coredump_config)
+	if err != nil {
+		journal.Print(journal.PriInfo, err.Error())
+	}
+	//根据配置项选择存储coredump文件方式
+	if pipe_config.Save_model == 0 {
+		err = writeCoreDumpToFile(coredump_config)
+		if err != nil {
+			journal.Print(journal.PriErr, err.Error())
+		}
+	} else if pipe_config.Save_model == 1 {
+		err = compress(coredump_config)
+		if err != nil {
+			journal.Print(journal.PriErr, err.Error())
+		}
+	}
+}
diff --git a/coredump-handler/coredump-handler_test.go b/coredump-handler/coredump-handler_test.go
new file mode 100644
index 0000000..b8aef16
--- /dev/null
+++ b/coredump-handler/coredump-handler_test.go
@@ -0,0 +1,166 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"coredump-handler/types"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"os/exec"
+	"testing"
+)
+
+func TestMain(m *testing.M) {
+	// 初始化测试数据
+	coredump_config.Initial_ns_pid = "123"
+	coredump_config.Process_ns_pid = "456"
+	coredump_config.Process_exe_path = "/usr/bin/test"
+	coredump_config.Corepipe_config_path = "/etc/test/config.json"
+	coredump_config.Timestamp = 1615478400
+	coredump_config.GID = "1000"
+	coredump_config.Hostname = "localhost"
+	coredump_config.Signal = 11
+	coredump_config.UID = "1000"
+
+	// 执行测试
+	retCode := m.Run()
+
+	// 清理测试数据
+	os.Remove("123_456_1615478400_coredump.info")
+	os.Remove("123_456_1615478400_coredump.info.zip")
+	os.Remove("123_456_1615478400_coredump.config")
+
+	os.Exit(retCode)
+}
+
+func TestArgsJudge(t *testing.T) {
+	err := argsJudge()
+	if err == nil {
+		t.Errorf("argsJudge() error = %v; want err", err)
+	}
+}
+
+func TestIsDiskSufficient(t *testing.T) {
+	pipe_config := types.Pipeconfig{
+		Total_file_mem_limit: "50%",
+	}
+	res, err := isDiskSufficient(pipe_config)
+	if err != nil || !res {
+		t.Errorf("isDiskSufficient() error = %v; want res = true", err)
+	}
+}
+
+func TestCreateCoreDumpDir(t *testing.T) {
+	pipe_config := types.Pipeconfig{
+		File_base_path: "/tmp",
+	}
+	err := createCoreDumpDir(&pipe_config, coredump_config)
+	if err != nil {
+		t.Errorf("createCoreDumpDir() error = %v; want err = nil", err)
+	}
+}
+
+func TestChangeDirectory(t *testing.T) {
+	err := changeDirectory("/tmp")
+	if err != nil {
+		t.Errorf("changeDirectory() error = %v; want err = nil", err)
+	}
+}
+
+func TestGetContainerId(t *testing.T) {
+	container_id, err := getContainerId("1")
+	if err == nil && len(container_id) == 0 {
+		t.Errorf("getContainerId() error = %v; want container_id != \"\"", err)
+	}
+}
+
+func TestGetImageId(t *testing.T) {
+	image_id, err := getImageId("1234567890abcdef", "/var/run/containerd.sock")
+	if err != nil || len(image_id) == 0 {
+		t.Errorf("getImageId() error = %v; want image_id != \"\"", err)
+	}
+}
+
+func TestWriteCoreConfig(t *testing.T) {
+	err := writeCoreConfig(coredump_config)
+	if err != nil {
+		t.Errorf("writeCoreConfig() error = %v; want err = nil", err)
+	}
+}
+
+func TestWriteCoreDumpToFile(t *testing.T) {
+	config := types.Coredump_config{
+		Initial_ns_pid:       "1",
+		Process_ns_pid:       "2",
+		Process_exe_path:     "/bin/bash",
+		Corepipe_config_path: "/tmp/config.yaml",
+		Timestamp:            12345678,
+	}
+	cmd := exec.Command("echo", "test")
+	cmdReader, err := cmd.StdoutPipe()
+	if err != nil {
+		t.Errorf("WriteCoreDumpToFile() error = %v", err)
+	}
+	scanner := bufio.NewScanner(cmdReader)
+	go func() {
+		for scanner.Scan() {
+			fmt.Printf("reading output: %s\n", scanner.Text())
+		}
+	}()
+	err = cmd.Start()
+	if err != nil {
+		t.Errorf("WriteCoreDumpToFile() error = %v", err)
+	}
+	err = writeCoreDumpToFile(config)
+	if err != nil {
+		t.Errorf("WriteCoreDumpToFile() error = %v", err)
+	}
+	b, err := ioutil.ReadFile("coredump.info")
+	if err != nil {
+		t.Errorf("WriteCoreDumpToFile() error = %v", err)
+	}
+	expected := "test\n"
+	if string(b) != expected {
+		t.Errorf("WriteCoreDumpToFile() got = %v, want = %v", string(b), expected)
+	}
+	defer os.Remove("coredump.info")
+}
+
+func TestCompress(t *testing.T) {
+	config := types.Coredump_config{
+		Initial_ns_pid:       "1",
+		Process_ns_pid:       "2",
+		Process_exe_path:     "/bin/bash",
+		Corepipe_config_path: "/tmp/config.yaml",
+		Timestamp:            12345678,
+	}
+	cmd := exec.Command("echo", "test")
+	cmdReader, err := cmd.StdoutPipe()
+	if err != nil {
+		t.Errorf("compress() error = %v", err)
+	}
+	scanner := bufio.NewScanner(cmdReader)
+	go func() {
+		for scanner.Scan() {
+			fmt.Printf("reading output: %s\n", scanner.Text())
+		}
+	}()
+	err = cmd.Start()
+	if err != nil {
+		t.Errorf("compress() error = %v", err)
+	}
+	err = compress(config)
+	if err != nil {
+		t.Errorf("compress() error = %v", err)
+	}
+	b, err := ioutil.ReadFile("coredump.info.zip")
+	if err != nil {
+		t.Errorf("compress() error = %v", err)
+	}
+	expected := []byte{80, 75, 3, 4, 20, 0, 8, 8, 8, 0, 170, 114, 38, 99, 102, 0, 0, 0, 102, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 116, 101, 115, 116, 10, 0, 82, 105, 105, 90, 85, 50, 84, 110, 88, 70, 90, 65, 48, 108, 108, 81, 0, 116, 157, 80, 10, 128, 163, 239, 120, 53, 2, 190, 127, 7, 98, 131, 137, 51, 189, 206, 149, 14, 54, 126, 254, 152, 119, 177, 209, 155, 201, 23, 37, 4, 72, 113, 39, 46, 179, 144, 106, 184, 44, 251, 47, 88, 97, 2, 141, 11, 129, 71, 109, 187, 124, 32, 63, 22, 111, 181, 59, 30, 58, 184, 40, 203, 205, 3, 113, 165, 117, 232, 6, 228, 240, 132, 94, 137, 43, 95, 218, 221, 90, 203, 173, 43, 92, 216, 226, 65, 118, 222, 208, 59, 185, 250, 56, 70, 63, 135, 72, 44, 250, 215, 59, 36, 139, 74, 75, 112, 0, 0, 0}
+	if !bytes.Equal(b, expected) {
+		t.Errorf("compress() got = %v, want = %v", b, expected)
+	}
+	defer os.Remove("coredump.info.zip")
+}