remove usage of event message when creating rules

shinebayar-g · Dec 27, 2021 · 62f87a3 · 62f87a3
1 parent 6f0a30e
commit 62f87a3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 27 deletions.
diff --git a/main.go b/main.go
@@ -33,9 +33,8 @@ func addFilters(client *client.Client, ctx *context.Context) (<-chan events.Mess
 func main() {
 	client := createClient()
 	ctx := context.Background()
-	messages, errors := addFilters(client, &ctx)
 
-	createChannel := make(chan *ufwhandler.UfwEvent)
+	createChannel := make(chan *types.ContainerJSON)
 	deleteChannel := make(chan string)
 
 	trackedContainers := make(map[string]*ufwhandler.TrackedContainer)
@@ -44,6 +43,7 @@ func main() {
 	go ufwhandler.DeleteUfwRule(deleteChannel, trackedContainers)
 	go ufwhandler.Cleanup(client, &ctx)
 
+	messages, errors := addFilters(client, &ctx)
 	for {
 		select {
 		case msg := <-messages:
@@ -54,7 +54,7 @@ func main() {
 						log.Println("ufw-docker-automated: Couldn't inspect container:", err)
 						continue
 					}
-					createChannel <- &ufwhandler.UfwEvent{Container: &container, Msg: &msg}
+					createChannel <- &container
 				}
 				if msg.Action == "kill" {
 					deleteChannel <- msg.ID[:12]

diff --git a/ufwhandler/create.go b/ufwhandler/create.go
@@ -8,6 +8,8 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
+
+	"github.com/docker/docker/api/types"
 )
 
 func checkIP(ip string) bool {
@@ -19,15 +21,15 @@ func checkCIDR(cidr string) bool {
 	return err == nil
 }
 
-func CreateUfwRule(ch <-chan *UfwEvent, trackedContainers map[string]*TrackedContainer) {
-	for event := range ch {
-		containerName := event.Msg.Actor.Attributes["name"]
-		containerIP := event.Container.NetworkSettings.IPAddress
-		containerID := event.Msg.ID[:12]
+func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]*TrackedContainer) {
+	for container := range ch {
+		containerName := strings.Replace(container.Name, "/", "", 1) // container name appears with prefix "/"
+		containerIP := container.NetworkSettings.IPAddress
+		containerID := container.ID[:12]
 		// If docker-compose, container IP is defined here
 		if containerIP == "" {
-			networkMode := event.Container.HostConfig.NetworkMode.NetworkName()
-			if ip, ok := event.Container.NetworkSettings.Networks[networkMode]; ok {
+			networkMode := container.HostConfig.NetworkMode.NetworkName()
+			if ip, ok := container.NetworkSettings.Networks[networkMode]; ok {
 				containerIP = ip.IPAddress
 			} else {
 				log.Println("ufw-docker-automated: Couldn't detect the container IP address.")
@@ -38,18 +40,18 @@ func CreateUfwRule(ch <-chan *UfwEvent, trackedContainers map[string]*TrackedCon
 		trackedContainers[containerID] = &TrackedContainer{
 			Name:      containerName,
 			IPAddress: containerIP,
-			Labels:    event.Container.Config.Labels,
+			Labels:    container.Config.Labels,
 		}
 
 		c := trackedContainers[containerID]
 
 		// Handle inbound rules
-		for port, portMaps := range event.Container.HostConfig.PortBindings {
+		for port, portMaps := range container.HostConfig.PortBindings {
 			// List is non empty if port is published
 			if len(portMaps) > 0 {
 				ufwRules := []UfwRule{}
-				if event.Msg.Actor.Attributes["UFW_ALLOW_FROM"] != "" {
-					ufwAllowFromLabelParsed := strings.Split(event.Msg.Actor.Attributes["UFW_ALLOW_FROM"], ";")
+				if container.Config.Labels["UFW_ALLOW_FROM"] != "" {
+					ufwAllowFromLabelParsed := strings.Split(container.Config.Labels["UFW_ALLOW_FROM"], ";")
 
 					for _, allowFrom := range ufwAllowFromLabelParsed {
 						ip := strings.Split(allowFrom, "-")
@@ -108,11 +110,11 @@ func CreateUfwRule(ch <-chan *UfwEvent, trackedContainers map[string]*TrackedCon
 		}
 
 		// Handle outbound rules
-		if strings.ToUpper(event.Msg.Actor.Attributes["UFW_DENY_OUT"]) == "TRUE" {
+		if container.Config.Labels["UFW_DENY_OUT"] == "TRUE" {
 
-			if event.Msg.Actor.Attributes["UFW_ALLOW_TO"] != "" {
+			if container.Config.Labels["UFW_ALLOW_TO"] != "" {
 				ufwRules := []UfwRule{}
-				ufwAllowToLabelParsed := strings.Split(event.Msg.Actor.Attributes["UFW_ALLOW_TO"], ";")
+				ufwAllowToLabelParsed := strings.Split(container.Config.Labels["UFW_ALLOW_TO"], ";")
 
 				for _, allowTo := range ufwAllowToLabelParsed {
 					ip := strings.Split(allowTo, "-")

diff --git a/ufwhandler/types.go b/ufwhandler/types.go
@@ -1,10 +1,5 @@
 package ufwhandler
 
-import (
-	"github.com/docker/docker/api/types"
-	"github.com/docker/docker/api/types/events"
-)
-
 type TrackedContainer struct {
 	Name             string
 	IPAddress        string
@@ -19,8 +14,3 @@ type UfwRule struct {
 	Proto   string
 	Comment string
 }
-
-type UfwEvent struct {
-	Container *types.ContainerJSON
-	Msg       *events.Message
-}