From dd9dd1a88f12fab0b1f187732d168a9ae2fd5a2c Mon Sep 17 00:00:00 2001
From: Sergii Putko <mail@psvlan.com>
Date: Thu, 7 Mar 2024 07:26:44 +0200
Subject: [PATCH] ttl

---
 README.md               | 1 +
 pkg/discovery/docker.go | 9 +++++++++
 pkg/handlers/dns.go     | 8 ++++----
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 92bdd9c..4af4a82 100644
--- a/README.md
+++ b/README.md
@@ -29,6 +29,7 @@ Usage of /disconter:
 - **disconter.service.priority** - dns discovery priority, default 1
 - **disconter.service.weight** - dns discovery weight, default 1
 - **disconter.service.port** - dns discovery port, default 80
+- **disconter.service.ttl** - dns discovery ttl in seconds, default 0
 
 ### run container with label `disconter.service`
 ```
diff --git a/pkg/discovery/docker.go b/pkg/discovery/docker.go
index 6b3900a..57b0e07 100644
--- a/pkg/discovery/docker.go
+++ b/pkg/discovery/docker.go
@@ -36,6 +36,7 @@ type ContainerInfo struct {
 		Priority uint16
 		Weight   uint16
 		Port     uint16
+		TTL      uint16
 	}
 }
 
@@ -182,6 +183,7 @@ func ListContainers(socket string) (containers []ContainerInfo, err error) {
 		c.DisconterService.Priority = 1
 		c.DisconterService.Weight = 1
 		c.DisconterService.Port = 80
+		c.DisconterService.TTL = 0
 
 		for _, v := range resp.NetworkSettings.Networks {
 			if v.IPAddress != "" {
@@ -214,6 +216,13 @@ func ListContainers(socket string) (containers []ContainerInfo, err error) {
 				c.DisconterService.Port = uint16(port)
 			}
 		}
+
+		if resp.Labels["disconter.service.ttl"] != "" {
+			ttl, err := strconv.Atoi(resp.Labels["disconter.service.ttl"])
+			if err == nil {
+				c.DisconterService.TTL = uint16(ttl)
+			}
+		}
 		containers[i] = c
 	}
 
diff --git a/pkg/handlers/dns.go b/pkg/handlers/dns.go
index b3f5fb5..36e9248 100644
--- a/pkg/handlers/dns.go
+++ b/pkg/handlers/dns.go
@@ -34,7 +34,7 @@ func HandleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
 			if (dns.TypeToString[q.Qtype] == "A" || dns.TypeToString[q.Qtype] == "ANY") && strings.HasSuffix(q.Name, "container.disconter.") {
 				for _, c := range discovery.ServiceContainers {
 					if q.Name == fmt.Sprintf("%s.container.disconter.", c.Name) {
-						rr, err := dns.NewRR(fmt.Sprintf("%s 0 A %s", q.Name, c.IP))
+						rr, err := dns.NewRR(fmt.Sprintf("%s %d A %s", q.Name, c.DisconterService.TTL, c.IP))
 						if err == nil {
 							m.Answer = append(m.Answer, rr)
 						}
@@ -44,7 +44,7 @@ func HandleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
 			if (dns.TypeToString[q.Qtype] == "A" || dns.TypeToString[q.Qtype] == "ANY") && strings.HasSuffix(q.Name, "service.disconter.") {
 				for _, c := range discovery.ServiceContainers {
 					if q.Name == fmt.Sprintf("%s.service.disconter.", c.DisconterService.Name) {
-						rr, err := dns.NewRR(fmt.Sprintf("%s 0 A %s", q.Name, c.IP))
+						rr, err := dns.NewRR(fmt.Sprintf("%s %d A %s", q.Name, c.DisconterService.TTL, c.IP))
 						if err == nil {
 							m.Answer = append(m.Answer, rr)
 						}
@@ -54,11 +54,11 @@ func HandleDNSRequest(w dns.ResponseWriter, r *dns.Msg) {
 			if dns.TypeToString[q.Qtype] == "SRV" && strings.HasSuffix(q.Name, "service.disconter.") {
 				for _, c := range discovery.ServiceContainers {
 					if q.Name == fmt.Sprintf("%s.service.disconter.", c.DisconterService.Name) || q.Name == fmt.Sprintf("_%s._tcp.service.disconter.", c.DisconterService.Name) {
-						rr, err := dns.NewRR(fmt.Sprintf("%s 0 SRV %d %d %d %s.container.disconter", q.Name, c.DisconterService.Priority, c.DisconterService.Weight, c.DisconterService.Port, c.Name))
+						rr, err := dns.NewRR(fmt.Sprintf("%s %d SRV %d %d %d %s.container.disconter", q.Name, c.DisconterService.TTL, c.DisconterService.Priority, c.DisconterService.Weight, c.DisconterService.Port, c.Name))
 						if err == nil {
 							m.Answer = append(m.Answer, rr)
 						}
-						rrA, err := dns.NewRR(fmt.Sprintf("%s.container.disconter 0 A %s", c.Name, c.IP))
+						rrA, err := dns.NewRR(fmt.Sprintf("%s.container.disconter %d A %s", c.Name, c.DisconterService.TTL, c.IP))
 						if err == nil {
 							m.Extra = append(m.Extra, rrA)
 						}