Added sticky mode

Author: buglloc

README.md

@@ -9,6 +9,7 @@ DNS server parses requested name to extract the requested mode, IP or CNAME by t
     <proxy-name>.p.<zone>  -> resolve proxy name and returns it
     <ip1>.<ip2>.r.<zone>  -> pick random <ip1> or <ip2>
     <ip1>.<ip2>.l.<zone>  -> loop over <ip1> and <ip2>
+    <ip1>.<ip2>.s.<zone>  -> "sticky" - <ip1> for first request, then <ip2> in sticky TTL (30 sec by default)
     <cname>.c.<zone>  -> return CNAME record with <cname>
     <any-name>.<zone>  -> returns default address
     [(<IPv4>.4|<IPv6>.6)...(<IPv4>.4|<IPv6>.6)].m.<zone>  -> returns multiple address according to order and type
@@ -55,6 +56,9 @@ When requesting it, we should get the following responses:
 # Loop
     8ba299a7.8ba299a8.l.example.com ->  loop over 139.162.153.167 and 139.162.153.168
 
+# Sticky
+    8ba299a7.8ba299a8.s.example.com ->  139.162.153.167 then 139.162.153.168, then 139.162.153.168 and so on
+
 # Cname
     ya.ru.c.example.com ->  canonical name ya.ru
     google.com.c.example.com ->  canonical name google.com

commands/ns.go

@@ -3,6 +3,7 @@ package commands
 import (
 	"errors"
 	"net"
+	"time"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -29,6 +30,8 @@ func init() {
 		"default ipv4 address")
 	flags.Uint32("ttl", cfg.TTL,
 		"DNS records TTL")
+	flags.Uint32("sticky-ttl", 30,
+		"sticky record TTL in seconds")
 	flags.String("ipv6", "::1",
 		"default ipv6 address")
 	flags.String("upstream", "77.88.8.8:53",
@@ -61,5 +64,6 @@ func parseServerConfig(cmd *cobra.Command, args []string) error {
 	cfg.UseDefault = viper.GetBool("UseDefault")
 	cfg.Upstream = viper.GetString("Upstream")
 	cfg.TTL = uint32(viper.GetInt("Ttl"))
+	cfg.StickyTTL = time.Duration(viper.GetInt("StickyTtl")) * time.Second
 	return nil
 }

pkg/cfg/cfg.go

@@ -1,6 +1,9 @@
 package cfg
 
-import "net"
+import (
+	"net"
+	"time"
+)
 
 const (
 	// RIP version
@@ -21,5 +24,6 @@ var (
 	// Enable "strict" mode
 	UseDefault bool
 	AllowProxy bool
-	TTL        uint32 = 0
+	TTL        uint32        = 0
+	StickyTTL  time.Duration = 30
 )

pkg/handlers/handler.go

@@ -15,6 +15,7 @@ var handlers = map[string]Handler{
 	"c": CnameHandler,
 	"p": ProxyHandler,
 	"l": LoopHandler,
+	"s": StickyHandler,
 	"r": RandomHandler,
 	"4": Ipv4Handler,
 	"6": Ipv6Handler,

pkg/handlers/sticky.go

@@ -0,0 +1,36 @@
+package handlers
+
+import (
+	"strings"
+
+	"github.com/buglloc/simplelog"
+	"github.com/miekg/dns"
+
+	"github.com/buglloc/rip/pkg/ip_stick"
+)
+
+func StickyHandler(question dns.Question, name string, l log.Logger) (rrs []dns.RR) {
+	ips := strings.Split(name, ".")
+	if len(ips) < 2 {
+		log.Error("failed to parse loop annotation")
+		return
+	}
+
+	var key string
+	if question.Qtype == dns.TypeA {
+		key = name + "A"
+	} else {
+		key = name + "AAAA"
+	}
+
+	ips = ips[len(ips)-2:]
+	targetIp := ip_stick.GetCurrent(key, ips)
+	ip := parseIp(question.Qtype, targetIp)
+	if ip == nil {
+		return
+	}
+
+	rrs = createIpsRR(question, ip)
+	l.Info("cooking response", "mode", "sticky", "ip", ip.String())
+	return
+}

pkg/ip_stick/ip_stick.go

@@ -0,0 +1,22 @@
+package ip_stick
+
+import (
+	"github.com/karlseguin/ccache"
+
+	"github.com/buglloc/rip/pkg/cfg"
+)
+
+var lruCache *ccache.Cache
+
+func init() {
+	lruCache = ccache.New(ccache.Configure().MaxSize(1000))
+}
+
+func GetCurrent(key string, ips []string) string {
+	if item := lruCache.Get(key); item != nil && !item.Expired() && item.Value() != nil {
+		return ips[1]
+	}
+
+	lruCache.Set(key, true, cfg.StickyTTL)
+	return ips[0]
+}