Merge pull request #99 from deputinizer/fix-optimize

Optimizations + target DNS injection

go.mod

@@ -5,6 +5,7 @@ go 1.17
 require (
 	github.com/miekg/dns v1.1.46
 	github.com/peterbourgon/ff/v3 v3.1.2
+	github.com/valyala/fastrand v1.1.0
 )
 
 require (

go.sum

@@ -16,6 +16,8 @@ github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6Kllzaw
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
 github.com/valyala/fasthttp v1.33.0 h1:mHBKd98J5NcXuBddgjvim1i3kWzlng1SzLhrnBOU9g8=
 github.com/valyala/fasthttp v1.33.0/go.mod h1:KJRK/MXx0J+yd0c5hlR+s1tIHD72sniU8ZJjl97LIw4=
+github.com/valyala/fastrand v1.1.0 h1:f+5HkLW4rsgzdNoleUOB69hyT9IlD2ZQh9GyDMfb5G8=
+github.com/valyala/fastrand v1.1.0/go.mod h1:HWqCzkrkg6QXT8V2EXWvXCoow7vLwOFN002oeRzjapQ=
 github.com/valyala/tcplisten v1.0.0/go.mod h1:T0xQ8SeCZGxckz9qRXTfG43PvQ/mcWh7FwZEA7Ioqkc=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=

internal/stoppropaganda/http.go → internal/stoppropaganda/apihttp.go

@@ -4,18 +4,21 @@ import (
 	"encoding/json"
 	"sync"
 
+	"github.com/erkexzcx/stoppropaganda/internal/stoppropaganda/customresolver"
 	"github.com/valyala/fasthttp"
 )
 
 func fasthttpRequestHandler(ctx *fasthttp.RequestCtx) {
 	switch string(ctx.Path()) {
 	case "/status":
 		fasthttpStatusResponseHandler(ctx)
+	case "/dnscache":
+		fasthttpDnsCacheResponseHandler(ctx)
 	}
 }
 
 type StatusStruct struct {
-	DNS      map[string]*DNSServerStatus `json:"DNS"`
+	DNS      map[string]*DNSTargetStatus `json:"DNS"`
 	Websites map[string]*WebsiteStatus   `json:"Websites"`
 }
 type StatusService struct {
@@ -27,17 +30,17 @@ type StatusService struct {
 func fasthttpStatusResponseHandler(ctx *fasthttp.RequestCtx) {
 	statusService := StatusService{
 		AllStatus: StatusStruct{
-			DNS:      make(map[string]*DNSServerStatus, len(dnsServers)),
+			DNS:      make(map[string]*DNSTargetStatus, len(dnsTargets)),
 			Websites: make(map[string]*WebsiteStatus, len(websites)),
 		},
 	}
 
 	wg := sync.WaitGroup{}
-	wg.Add(len(dnsServers))
+	wg.Add(len(dnsTargets))
 	wg.Add(len(websites))
 
-	for endpoint, ds := range dnsServers {
-		go func(endpoint string, ds *DNSServer) {
+	for endpoint, ds := range dnsTargets {
+		go func(endpoint string, ds *DNSTarget) {
 			ds.mux.Lock()
 			dnsStatus := ds.Status
 			ds.mux.Unlock()
@@ -76,3 +79,17 @@ func fasthttpStatusResponseHandler(ctx *fasthttp.RequestCtx) {
 	}
 	ctx.Write(content)
 }
+
+func fasthttpDnsCacheResponseHandler(ctx *fasthttp.RequestCtx) {
+
+	cache := customresolver.DnsCache
+
+	dnsCacheItems := cache.Items()
+	content, err := json.MarshalIndent(dnsCacheItems, "", "    ")
+	if err != nil {
+		ctx.SetStatusCode(500)
+		ctx.WriteString("failed to marshal data")
+		return
+	}
+	ctx.Write(content)
+}

internal/stoppropaganda/customresolver/customgetip.go

@@ -0,0 +1,48 @@
+package customresolver
+
+import (
+	"context"
+	"net"
+)
+
+var getIPcachedResolver = &CustomResolver{
+	ParentResolver: InjectionGoResolver,
+}
+
+// Modified to use stoppropaganda's CustomResolver
+// so that it caches DNS records
+func CustomLookupIP(host string, helperIPBuf []net.IP) ([]net.IP, error) {
+	addrs, err := getIPcachedResolver.LookupIPAddr(context.Background(), host)
+	if err != nil {
+		return nil, err
+	}
+	ips := helperIPBuf[:0]
+	for _, ia := range addrs {
+		ips = append(ips, ia.IP)
+	}
+	return ips, nil
+}
+
+func GetIPs(host string, helperIPBuf []net.IP) (ips []net.IP, err error) {
+	addr := net.ParseIP(host)
+	if addr == nil {
+		ips, err := CustomLookupIP(host, helperIPBuf)
+		if err != nil {
+			return nil, err
+		}
+		for i := 0; i < len(ips); i++ {
+			ip := ips[i]
+			if ipv4 := ip.To4(); ipv4 == nil {
+				// Remove inplace trick
+				// - swap i and last element
+				ips[i] = ips[len(ips)-1]
+				// - pop last element
+				ips[len(ips)-1] = nil
+				ips = ips[:len(ips)-1]
+			}
+		}
+		return ips, nil
+	}
+	helperIPBuf[0] = addr
+	return helperIPBuf[:1], nil
+}

internal/stoppropaganda/customresolver/customresolver.go

@@ -8,7 +8,7 @@ import (
 	"github.com/patrickmn/go-cache"
 )
 
-var dnscache *cache.Cache
+var DnsCache *cache.Cache
 
 type CustomResolver struct {
 	ParentResolver *net.Resolver
@@ -19,17 +19,17 @@ type Resolver interface {
 }
 
 func (cr *CustomResolver) LookupIPAddr(ctx context.Context, host string) (names []net.IPAddr, err error) {
-	if c, found := dnscache.Get(host); found {
+	if c, found := DnsCache.Get(host); found {
 		return c.([]net.IPAddr), nil
 	}
 
 	names, err = cr.ParentResolver.LookupIPAddr(ctx, host)
 	if err == nil {
-		dnscache.SetDefault(host, names)
+		DnsCache.SetDefault(host, names)
 	}
 	return
 }
 
 func init() {
-	dnscache = cache.New(5*time.Minute, 10*time.Minute)
+	DnsCache = cache.New(5*time.Minute, 10*time.Minute)
 }

internal/stoppropaganda/customresolver/injectionresolver.go

@@ -0,0 +1,30 @@
+package customresolver
+
+import (
+	"context"
+	"net"
+	"time"
+
+	"github.com/erkexzcx/stoppropaganda/internal/stoppropaganda/targets"
+)
+
+var InjectionGoResolver = &net.Resolver{
+	PreferGo: true,
+	Dial: func(ctx context.Context, network, address string) (conn net.Conn, err error) {
+		d := net.Dialer{
+			Timeout: time.Millisecond * time.Duration(1000),
+		}
+
+		// use DNS targets from dns.go
+		for _, dnsTarget := range targets.ReferenceDNSServersForHTTP {
+			// eg. d.DialContext(ctx, "tcp", "194.54.14.186:53")
+			conn, err = d.DialContext(ctx, network, dnsTarget)
+			if err == nil {
+				// return first working conn to DNS
+				return
+			}
+		}
+
+		return d.DialContext(ctx, network, address)
+	},
+}

internal/stoppropaganda/customtcpdial/customtcpdial.go

@@ -58,9 +58,7 @@ type CustomTCPDialer struct {
 
 	once sync.Once
 
-	// stoppropaganda start
 	ParentDialer sockshttp.Dialer
-	// stoppropaganda end
 }
 
 // Dial dials the given TCP addr using tcp4.
@@ -197,15 +195,6 @@ func (d *CustomTCPDialer) dial(addr string, dualStack bool, timeout time.Duratio
 		return nil, errors.New("CustomTCPDialer: " + checkErr.Error())
 	}
 
-	ticketC := d.DialTicketsC
-	if ticketC != nil {
-		select {
-		case <-ticketC:
-		case <-time.After(1 * time.Second):
-			return nil, ErrTooFastDialSpam
-		}
-	}
-
 	var conn net.Conn
 	n := uint32(len(addrs))
 	deadline := time.Now().Add(timeout)
@@ -248,7 +237,17 @@ func (d *CustomTCPDialer) tryDial(network string, addr *net.TCPAddr, deadline ti
 		}
 		defer func() { <-concurrencyCh }()
 	}
-	// stoppropaganda start - add parent dialer
+	ticketC := d.DialTicketsC
+	if ticketC != nil {
+		select {
+		// either we catch the ticket instantly
+		case <-ticketC:
+		// or maybe let's wait until we have a green light
+		case <-time.After(timeout / 2):
+			// time passed, we didn't get a ticket :(
+			return nil, ErrTooFastDialSpam
+		}
+	}
 
 	dialer := d.ParentDialer
 	if dialer == nil {
@@ -260,7 +259,6 @@ func (d *CustomTCPDialer) tryDial(network string, addr *net.TCPAddr, deadline ti
 	}
 
 	conn, err := dialer.Dial(network, addr.String())
-	// stoppropaganda end
 	return conn, err
 }