Declare types and enums and add better error handling

Signed-off-by: AbdelrahmanElawady <[email protected]>

KubeArmor/core/hook_handler.go

@@ -5,6 +5,7 @@ package core
 
 import (
 	"encoding/json"
+	"errors"
 	"io"
 	"log"
 	"net"
@@ -27,7 +28,11 @@ func (dm *KubeArmorDaemon) ListenToHook() {
 	}
 
 	listenPath := filepath.Join(kubearmorDir, "ka.sock")
-	_ = os.Remove(listenPath) // in case kubearmor crashed and the socket wasn't removed
+	err := os.Remove(listenPath) // in case kubearmor crashed and the socket wasn't removed
+	if err != nil && !errors.Is(err, os.ErrNotExist) {
+		log.Fatal(err)
+	}
+
 	socket, err := net.Listen("unix", listenPath)
 	if err != nil {
 		log.Fatal(err)
@@ -66,11 +71,7 @@ func (dm *KubeArmorDaemon) handleConn(conn net.Conn, ready *atomic.Bool) {
 			log.Fatal(err)
 		}
 
-		data := struct {
-			Operation string          `json:"operation"`
-			Detached  bool            `json:"detached"`
-			Container types.Container `json:"container"`
-		}{}
+		data := types.HookRequest{}
 
 		err = json.Unmarshal(buf[:n], &data)
 		if err != nil {
@@ -100,10 +101,10 @@ func (dm *KubeArmorDaemon) handleConn(conn net.Conn, ready *atomic.Bool) {
 			return
 		}
 
-		if data.Operation == "create" {
+		if data.Operation == types.HookContainerCreate {
 			dm.handleContainerCreate(data.Container)
 		} else {
-			dm.handleContainerStop(data.Container.ContainerID)
+			dm.handleContainerDelete(data.Container.ContainerID)
 		}
 	}
 }
@@ -113,16 +114,15 @@ func (dm *KubeArmorDaemon) handleContainerCreate(container types.Container) {
 	dm.Logger.Printf("added %s", container.ContainerID)
 
 	dm.ContainersLock.Lock()
+	defer dm.ContainersLock.Unlock()
 	if _, ok := dm.Containers[container.ContainerID]; !ok {
 		dm.Containers[container.ContainerID] = container
-		dm.ContainersLock.Unlock()
 	} else if dm.Containers[container.ContainerID].PidNS == 0 && dm.Containers[container.ContainerID].MntNS == 0 {
 		c := dm.Containers[container.ContainerID]
 		c.MntNS = container.MntNS
 		c.PidNS = container.PidNS
 		c.AppArmorProfile = container.AppArmorProfile
 		dm.Containers[c.ContainerID] = c
-		dm.ContainersLock.Unlock()
 
 		dm.EndPointsLock.Lock()
 		for idx, endPoint := range dm.EndPoints {
@@ -143,9 +143,8 @@ func (dm *KubeArmorDaemon) handleContainerCreate(container types.Container) {
 			}
 		}
 		dm.EndPointsLock.Unlock()
-	} else {
-		dm.ContainersLock.Unlock()
 	}
+
 	if len(dm.OwnerInfo) > 0 {
 		container.Owner = dm.OwnerInfo[container.EndPointName]
 	}
@@ -163,7 +162,7 @@ func (dm *KubeArmorDaemon) handleContainerCreate(container types.Container) {
 		}
 	}
 }
-func (dm *KubeArmorDaemon) handleContainerStop(containerID string) {
+func (dm *KubeArmorDaemon) handleContainerDelete(containerID string) {
 	dm.ContainersLock.Lock()
 	container, ok := dm.Containers[containerID]
 	dm.Logger.Printf("deleted %s", containerID)

KubeArmor/types/types.go

@@ -641,3 +641,20 @@ type PidNode struct {
 
 // KubeArmorHostPolicyEventCallback Function
 type KubeArmorHostPolicyEventCallback func(K8sKubeArmorHostPolicyEvent) pb.PolicyStatus
+
+// =========== //
+// == Hooks == //
+// =========== //
+
+type HookRequest struct {
+	Operation HookOperation `json:"operation"`
+	Detached  bool          `json:"detached"`
+	Container Container     `json:"container"`
+}
+
+type HookOperation int
+
+const (
+	HookContainerCreate HookOperation = iota
+	HookContainerDelete
+)

pkg/KubeArmorOperator/hook/crio.go

@@ -14,6 +14,13 @@ import (
 	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
 )
 
+type containerInfo struct {
+	SandboxID   string     `json:"sandboxID"`
+	Pid         int        `json:"pid"`
+	RuntimeSpec specs.Spec `json:"runtimeSpec"`
+	Privileged  bool       `json:"privileged"`
+}
+
 type crioHandler struct {
 	client runtime.RuntimeServiceClient
 	conn   *grpc.ClientConn
@@ -32,8 +39,9 @@ func newCRIOHandler(socket string) (handler, error) {
 	return &crioHandler{client: client, conn: conn}, nil
 }
 
-func (h *crioHandler) close() {
-	_ = h.conn.Close()
+func (h *crioHandler) close() error {
+	return h.conn.Close()
+
 }
 
 func (h *crioHandler) listContainers(ctx context.Context) ([]types.Container, error) {
@@ -97,10 +105,3 @@ func containerFromContainerStatus(status *runtime.ContainerStatus, info string)
 
 	return container, nil
 }
-
-type containerInfo struct {
-	SandboxID   string     `json:"sandboxID"`
-	Pid         int        `json:"pid"`
-	RuntimeSpec specs.Spec `json:"runtimeSpec"`
-	Privileged  bool       `json:"privileged"`
-}

pkg/KubeArmorOperator/hook/main.go

@@ -88,11 +88,11 @@ func runDetached() error {
 	}
 
 	for _, container := range containers {
-		data := struct {
-			Operation string          `json:"operation"`
-			Detached  bool            `json:"detached"`
-			Container types.Container `json:"container"`
-		}{Operation: "create", Detached: true, Container: container}
+		data := types.HookRequest{
+			Operation: types.HookContainerCreate,
+			Detached:  true,
+			Container: container,
+		}
 
 		dataJSON, err := json.Marshal(data)
 		if err != nil {
@@ -118,18 +118,21 @@ func runDetached() error {
 
 func run(state specs.State) error {
 	var container types.Container
-	operation := "create"
+	operation := types.HookContainerCreate
 	// we try to connect to runtime here to make sure the socket is correct
 	// before spawning a detached process
 	handler, err := newCRIOHandler(runtimeSocket)
 	if err != nil {
 		return err
 	}
-	handler.close()
+	err = handler.close()
+	if err != nil {
+		log.Printf("failed to close runtime connection: %s", err.Error())
+	}
 
 	container.ContainerID = state.ID
 	if state.Status == specs.StateStopped {
-		operation = "delete"
+		operation = types.HookContainerDelete
 		return sendContainer(container, operation)
 	}
 
@@ -204,7 +207,7 @@ func getNS(pid int) (uint32, uint32) {
 	return pidNS, mntNS
 }
 
-func sendContainer(container types.Container, operation string) error {
+func sendContainer(container types.Container, operation types.HookOperation) error {
 	conn, err := net.Dial("unix", kubeArmorSocket)
 	if err != nil {
 		// not returning error here because this can happen in multiple cases
@@ -216,11 +219,11 @@ func sendContainer(container types.Container, operation string) error {
 
 	defer conn.Close()
 
-	data := struct {
-		Operation string          `json:"operation"`
-		Detached  bool            `json:"detached"`
-		Container types.Container `json:"container"`
-	}{Operation: operation, Detached: false, Container: container}
+	data := types.HookRequest{
+		Operation: operation,
+		Detached:  false,
+		Container: container,
+	}
 
 	dataJSON, err := json.Marshal(data)
 	if err != nil {

pkg/KubeArmorOperator/hook/types.go

@@ -11,5 +11,5 @@ import (
 
 type handler interface {
 	listContainers(ctx context.Context) ([]types.Container, error)
-	close()
+	close() error
 }