use go-cache for tracked containers

go.mod

@@ -2,7 +2,10 @@ module github.com/shinebayar-g/ufw-docker-automated
 
 go 1.19
 
-require github.com/docker/docker v20.10.21+incompatible
+require (
+	github.com/docker/docker v20.10.21+incompatible
+	github.com/patrickmn/go-cache v2.1.0+incompatible
+)
 
 require (
 	github.com/Microsoft/go-winio v0.6.0 // indirect

go.sum

@@ -27,6 +27,8 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0-rc2 h1:2zx/Stx4Wc5pIPDvIxHXvXtQFW/7XWJGmnM7r3wg034=
 github.com/opencontainers/image-spec v1.1.0-rc2/go.mod h1:3OVijpioIKYWTqjiG0zfF6wvoJ4fAXGbjdZuI2NgsRQ=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

main.go

@@ -10,6 +10,7 @@ import (
 	"github.com/docker/docker/api/types/events"
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/docker/client"
+	"github.com/patrickmn/go-cache"
 	"github.com/shinebayar-g/ufw-docker-automated/ufwhandler"
 )
 
@@ -55,13 +56,12 @@ func main() {
 	}
 	createChannel := make(chan *types.ContainerJSON)
 	deleteChannel := make(chan string)
-
-	trackedContainers := make(map[string]*ufwhandler.TrackedContainer)
+	trackedContainers := cache.New(cache.NoExpiration, 0)
 
 	go ufwhandler.CreateUfwRule(createChannel, trackedContainers)
 	go ufwhandler.DeleteUfwRule(deleteChannel, trackedContainers)
-	go ufwhandler.Cleanup(client, ctx)
-	go ufwhandler.Sync(createChannel, client, ctx)
+	// go ufwhandler.Cleanup(client, ctx)
+	// go ufwhandler.Sync(createChannel, client, ctx)
 
 	messages, errors := streamEvents(ctx, client)
 	for {

ufwhandler/create.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/docker/docker/api/types"
+	"github.com/patrickmn/go-cache"
 )
 
 func checkIP(ip string) bool {
@@ -21,7 +22,7 @@ func checkCIDR(cidr string) bool {
 	return err == nil
 }
 
-func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]*TrackedContainer) {
+func CreateUfwRule(ch <-chan *types.ContainerJSON, c *cache.Cache) {
 	for container := range ch {
 		containerName := strings.Replace(container.Name, "/", "", 1) // container name appears with prefix "/"
 		containerIP := container.NetworkSettings.IPAddress
@@ -37,13 +38,13 @@ func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]
 			}
 		}
 
-		trackedContainers[containerID] = &TrackedContainer{
+		cachedContainer := TrackedContainer{
 			Name:      containerName,
 			IPAddress: containerIP,
 			Labels:    container.Config.Labels,
 		}
 
-		c := trackedContainers[containerID]
+		c.Set(containerID, &cachedContainer, cache.NoExpiration)
 
 		// Handle inbound rules
 		for port, portMaps := range container.HostConfig.PortBindings {
@@ -101,7 +102,7 @@ func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]
 					}
 				}
 
-				c.UfwInboundRules = append(c.UfwInboundRules, ufwRules...)
+				cachedContainer.UfwInboundRules = append(cachedContainer.UfwInboundRules, ufwRules...)
 				// ufw route allow proto tcp from any to 172.17.0.2 port 80 comment "Comment"
 				// ufw route allow proto <tcp|udp> <source> to <container_ip> port <port> comment <comment>
 				// ufw route delete allow proto tcp from any to 172.17.0.2 port 80 comment "Comment"
@@ -167,7 +168,7 @@ func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]
 					}
 				}
 
-				c.UfwOutboundRules = append(c.UfwOutboundRules, ufwRules...)
+				cachedContainer.UfwOutboundRules = append(cachedContainer.UfwOutboundRules, ufwRules...)
 			}
 
 			// Handle deny all out
@@ -185,6 +186,5 @@ func CreateUfwRule(ch <-chan *types.ContainerJSON, trackedContainers map[string]
 				log.Println("ufw:", stdout.String())
 			}
 		}
-
 	}
 }

ufwhandler/delete.go

@@ -4,15 +4,18 @@ import (
 	"bytes"
 	"log"
 	"os/exec"
+
+	"github.com/patrickmn/go-cache"
 )
 
-func DeleteUfwRule(containerID <-chan string, trackedContainers map[string]*TrackedContainer) {
+func DeleteUfwRule(containerID <-chan string, c *cache.Cache) {
 	for id := range containerID {
 
-		if c, ok := trackedContainers[id]; ok {
+		if cachedContainer, found := c.Get(id); found {
+			container := cachedContainer.(*TrackedContainer)
 			// Handle inbound rules
-			for _, rule := range c.UfwInboundRules {
-				cmd := exec.Command("sudo", "ufw", "route", "delete", "allow", "proto", rule.Proto, "from", rule.CIDR, "to", c.IPAddress, "port", rule.Port, "comment", c.Name+":"+id+rule.Comment)
+			for _, rule := range container.UfwInboundRules {
+				cmd := exec.Command("sudo", "ufw", "route", "delete", "allow", "proto", rule.Proto, "from", rule.CIDR, "to", container.IPAddress, "port", rule.Port, "comment", container.Name+":"+id+rule.Comment)
 				log.Println("ufw-docker-automated: Deleting rule:", cmd)
 
 				var stdout, stderr bytes.Buffer
@@ -27,12 +30,12 @@ func DeleteUfwRule(containerID <-chan string, trackedContainers map[string]*Trac
 				}
 			}
 			// Handle outbound rules
-			for _, rule := range c.UfwOutboundRules {
+			for _, rule := range container.UfwOutboundRules {
 				var cmd *exec.Cmd
 				if rule.Port == "" {
-					cmd = exec.Command("sudo", "ufw", "route", "delete", "allow", "from", c.IPAddress, "to", rule.CIDR, "comment", c.Name+":"+id+rule.Comment)
+					cmd = exec.Command("sudo", "ufw", "route", "delete", "allow", "from", container.IPAddress, "to", rule.CIDR, "comment", container.Name+":"+id+rule.Comment)
 				} else {
-					cmd = exec.Command("sudo", "ufw", "route", "delete", "allow", "from", c.IPAddress, "to", rule.CIDR, "port", rule.Port, "comment", c.Name+":"+id+rule.Comment)
+					cmd = exec.Command("sudo", "ufw", "route", "delete", "allow", "from", container.IPAddress, "to", rule.CIDR, "port", rule.Port, "comment", container.Name+":"+id+rule.Comment)
 				}
 				log.Println("ufw-docker-automated: Deleting rule:", cmd)
 
@@ -48,7 +51,7 @@ func DeleteUfwRule(containerID <-chan string, trackedContainers map[string]*Trac
 				}
 			}
 			// Handle deny all out
-			cmd := exec.Command("sudo", "ufw", "route", "delete", "deny", "from", c.IPAddress, "to", "any", "comment", c.Name+":"+id)
+			cmd := exec.Command("sudo", "ufw", "route", "delete", "deny", "from", container.IPAddress, "to", "any", "comment", container.Name+":"+id)
 			log.Println("ufw-docker-automated: Deleting rule:", cmd)
 
 			var stdout, stderr bytes.Buffer