From 684019c2e3be57128562199533aab33136eae8a6 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Fri, 8 Dec 2023 23:23:33 +0700
Subject: [PATCH] all: force re-bootstrapping with timeout error

---
 cmd/cli/dns_proxy.go | 6 ++++++
 config.go            | 5 +++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/cmd/cli/dns_proxy.go b/cmd/cli/dns_proxy.go
index 080bebc..0a68071 100644
--- a/cmd/cli/dns_proxy.go
+++ b/cmd/cli/dns_proxy.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/rand"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"net"
 	"net/netip"
@@ -438,6 +439,11 @@ func (p *prog) proxy(ctx context.Context, req *proxyRequest) *dns.Msg {
 					go p.um.checkUpstream(upstreams[n], upstreamConfig)
 				}
 			}
+			// For timeout error (i.e: context deadline exceed), force re-bootstrapping.
+			var e net.Error
+			if errors.As(err, &e) && e.Timeout() {
+				upstreamConfig.ReBootstrap()
+			}
 			return nil
 		}
 		return answer
diff --git a/config.go b/config.go
index d3509be..aeb6e27 100644
--- a/config.go
+++ b/config.go
@@ -426,8 +426,9 @@ func (uc *UpstreamConfig) ReBootstrap() {
 		return
 	}
 	_, _, _ = uc.g.Do("ReBootstrap", func() (any, error) {
-		ProxyLogger.Load().Debug().Msg("re-bootstrapping upstream ip")
-		uc.rebootstrap.Store(true)
+		if uc.rebootstrap.CompareAndSwap(false, true) {
+			ProxyLogger.Load().Debug().Msg("re-bootstrapping upstream ip")
+		}
 		return true, nil
 	})
 }