rpn: impl exclude country codes

Author: ignoramous

intra/backend/ipn_proxies.go

@@ -6,7 +6,11 @@
 
 package backend
 
-import "fmt"
+import (
+	"fmt"
+	"slices"
+	"strings"
+)
 
 const ( // see ipn/proxies.go
 	// IDs for default proxies
@@ -90,15 +94,16 @@ type RpnOps struct {
 	newPort           uint16 // fixed WireGuard port; 0 = random per wsRandomPort()
 	dnsConfig         string // csv of DNS filter presets: "family", "security", "social", "privacy", "all", "none", "default"
 	forceInit         bool   // when false, skips expensive ops unless absolutely required.
+	excludeCCs        string // csv of (sorted) country codes to exclude from selection.
 }
 
 func NewRpnOps() *RpnOps {
 	return &RpnOps{}
 }
 
 func (o *RpnOps) String() string {
-	return fmt.Sprintf("rotate: %t; perma: %t; forceFetchServers: %t; port: %d; dns: %s; forceInit: %t",
-		o.rotateCreds, o.permaCreds, o.forceFetchServers, o.newPort, o.dnsConfig, o.forceInit)
+	return fmt.Sprintf("rotate: %t; perma: %t; forceFetchServers: %t; port: %d; dns: %s; forceInit: %t; excludeCCs: %v",
+		o.rotateCreds, o.permaCreds, o.forceFetchServers, o.newPort, o.dnsConfig, o.forceInit, o.excludeCCs)
 }
 
 // SetRotateCreds forces generation of a new WireGuard keypair on the next Update.
@@ -130,6 +135,23 @@ func (o *RpnOps) SetDNSConfig(v string) {
 // expensive ops are skipped if called within some threshold of the previous call.
 func (o *RpnOps) SetForceInit(v bool) { o.forceInit = v }
 
+// SetExcludeCCs sets a CSV of country codes to exclude from CC selection.
+// Empty or whitespace entries are ignored; codes are normalised to upper-case.
+func (o *RpnOps) SetExcludeCCs(v string) {
+	parts := slices.Sorted(strings.SplitSeq(v, ","))
+	out := make([]string, 0, len(parts))
+	for _, p := range parts {
+		p = strings.ToUpper(strings.TrimSpace(p))
+		if len(p) > 0 {
+			out = append(out, p)
+		}
+	}
+	o.excludeCCs = strings.Join(out, ",")
+}
+
+// ExcludeCCs returns the CSV of country codes excluded from CC selection.
+func (o RpnOps) ExcludeCCs() (csv string) { return o.excludeCCs }
+
 // Rotate reports whether a new WG keypair should be generated.
 func (o RpnOps) Rotate() bool { return o.rotateCreds }
 
@@ -152,12 +174,19 @@ func (o RpnOps) DNSConfig() string {
 // ForceInit returns false if expensive update ops may be skipped if approp.
 func (o RpnOps) ForceInit() bool { return o.forceInit }
 
+// ExcludeCCs returns csv of (to be) sorted excluded country codes.
+func (o RpnOps) GetExcludeCCs() string { return o.excludeCCs }
+
 // ChangesConfig reports whether this RpnOps would cause a change in wg config
 // if applied to override "other".
 func (o RpnOps) ChangesConfig(other RpnOps) bool {
 	return o.rotateCreds != other.rotateCreds ||
 		o.permaCreds != other.permaCreds ||
 		o.newPort != other.newPort
+	// TODO: asses if excludeccs would cause change in wg config (RegionalWgConfs)
+	// a empty excludeCCs means existing ones, if any, are retained.
+	// excludedccs is a sorted csv and so string equality should work.
+	// len(o.excludeCCs) > 0 && o.excludeCCs != other.excludeCCs
 }
 
 type Rpn interface {
@@ -412,8 +441,10 @@ type RpnServer struct {
 	Link int32
 	// Number of active servers in this CC+City.
 	Count int32
-	// Premium
+	// Premium server
 	Premium bool
+	// Excluded by end-user
+	Excluded bool
 }
 
 type IPMetadata struct {

intra/ipn/rpn.go

@@ -364,21 +364,45 @@ func (r *rpnp) forkMain() error {
 	return err
 }
 
+// ccCsvAsSet mods a comma-separated list of country codes into a lookup set.
+func ccCsvAsSet(csv string) map[string]struct{} {
+	if len(csv) <= 0 {
+		return nil
+	}
+	parts := strings.Split(csv, ",")
+	out := make(map[string]struct{}, len(parts))
+	for _, p := range parts {
+		p = strings.ToUpper(strings.TrimSpace(p))
+		if len(p) > 0 {
+			out[p] = struct{}{}
+		}
+	}
+	return out
+}
+
 func (r *rpnp) forkAll() error {
 	provider := r.RpnAcc.ProviderID()
 	kids := r.flattenKids()
-	log.I("proxy: rpn: forkAll: %s[%v]", provider, kids)
 
 	errs := make([]error, 0) // may contain nil errors
 
-	e := r.forkMain()
+	ops := r.RpnAcc.Ops()
+	excludedSet := ccCsvAsSet(ops.ExcludeCCs())
 
+	log.I("proxy: rpn: forkAll: %s [%v] incl: %d / excl: %d", provider, kids, len(kids), len(excludedSet))
+
+	e := r.forkMain()
 	errs = append(errs, e)
 
 	for _, cc := range kids {
-		_, e := r.fork(cc)
-		loged(e)("proxy: rpn: forkAll: forked %s[%s]; err? %v", provider, cc, e)
-		errs = append(errs, e)
+		if _, excluded := excludedSet[cc]; excluded {
+			r.purge(cc) // remove excluded kid
+			log.I("proxy: rpn: forkAll: %s[%s] excluded; purged", provider, cc)
+		} else {
+			_, e := r.fork(cc)
+			loged(e)("proxy: rpn: forkAll: forked %s[%s]; err? %v", provider, cc, e)
+			errs = append(errs, e)
+		}
 	}
 	return core.JoinErr(errs...)
 }

intra/ipn/rpn/yegor.go

@@ -178,6 +178,7 @@ var (
 	errWsRetryUpdate    = errors.New("ws: retry update")
 	errWsNoCcConfig     = errors.New("ws: not available in that location")
 	errWsNoFilters      = errors.New("ws: no filter list")
+	errWsCCExcluded     = errors.New("ws: cc excluded")
 )
 
 /*
@@ -1054,6 +1055,7 @@ func (a *WsClient) Locations() (x.RpnServers, error) {
 	if len(c.Configs) <= 0 {
 		return nil, errWsNoCcConfig
 	}
+	excl := ccCsvAsSet(a.Ops().ExcludeCCs())
 	visited := make(map[string]bool, len(c.Configs))
 	s := make([]x.RpnServer, 0, len(c.Configs)/maxPerRegionWgConfs)
 	for i, rc := range c.Configs {
@@ -1070,14 +1072,16 @@ func (a *WsClient) Locations() (x.RpnServers, error) {
 			continue
 		}
 		if !visited[rc.Name] {
+			_, isExcluded := excl[rc.CC]
 			s = append(s, x.RpnServer{
-				CC:      rc.CC,
-				City:    rc.City,
-				Name:    rc.Name,
-				Load:    rc.Load,
-				Link:    rc.Link,
-				Count:   rc.Count,
-				Premium: rc.Premium,
+				CC:       rc.CC,
+				City:     rc.City,
+				Name:     rc.Name,
+				Load:     rc.Load,
+				Link:     rc.Link,
+				Count:    rc.Count,
+				Premium:  rc.Premium,
+				Excluded: isExcluded,
 				// cc is always suffixed; see proxy.go:proxifier.postAddRpnProxy
 				Key:   strings.Join([]string{rc.City, rc.CC}, confKeySep),
 				Addrs: strings.Join([]string{rc.ServerDomainPort, rc.addrCsv()}, ","),
@@ -1100,9 +1104,15 @@ func (a *WsClient) Update(ops *x.RpnOps) (newstate []byte, err error) {
 	curops := a.Ops()
 	if ops == nil {
 		ops = curops
-	} else if len(ops.DNSConfig()) <= 0 {
-		// retain existing dns config
-		ops.SetDNSConfig(curops.DNSConfig())
+	} else {
+		if len(ops.DNSConfig()) <= 0 {
+			// retain existing dns config
+			ops.SetDNSConfig(curops.DNSConfig())
+		}
+		// retain existing excludeCCs if not set by incoming ops
+		if len(ops.ExcludeCCs()) <= 0 {
+			ops.SetExcludeCCs(curops.ExcludeCCs())
+		}
 	}
 	start := time.Now()
 	b, refreshed, needsRedo, err := makeWsWgFrom(a.http, c, *ops, true /*updating*/, ops.ChangesConfig(*curops))
@@ -1159,12 +1169,28 @@ func (a *WsClient) Conf(cc string) (string, error) {
 		city = cccsv[0]
 		cc = cccsv[1]
 	}
-	visited := make(map[string]struct{}, 0)
 	// in sync with anyCountryCode / noCountryForOldMen vars in proxy.go
 	chooseAny := cc == "**" || len(cc) <= 0
 	hasCity := len(city) > 0
-	tot := 0
-	c := 0
+	cc = strings.ToUpper(cc)
+
+	excl := ccCsvAsSet(a.Ops().ExcludeCCs())
+	// if a specific (non-wildcard) CC is explicitly excluded, bail early
+	if !chooseAny && len(excl) > 0 {
+		if _, excluded := excl[cc]; excluded {
+			log.W("ws: conf: cc %s is excluded...", cc)
+			return "", errWsCCExcluded
+		}
+	}
+
+	retried := false
+reconf:
+	tot := 0  // total seen
+	c := 0    // good cc conf
+	badc := 0 // bad cc conf
+	x := 0    // total excluded
+	v := 0    // total visited
+	visited := make(map[string]struct{}, len(cfg.Configs))
 	out := make([]string, 0, maxPerRegionWgConfs)
 	ids := make([]string, 0, maxPerRegionWgConfs)
 	for _, rc := range cfg.Configs {
@@ -1175,8 +1201,14 @@ func (a *WsClient) Conf(cc string) (string, error) {
 					continue
 				}
 				visited[rc.CC] = struct{}{}
+				v++
+				// skip CCs the user has excluded
+				if _, excluded := excl[rc.CC]; excluded {
+					x++
+					continue
+				}
 				if c > 2 {
-					// choose only low load and high link speed servers
+					// after a couple random servers, prefer low load and high link speed servers
 					gbps10 := rc.Link >= 10000
 					healthy50 := rc.Load <= 50
 					gbps1 := rc.Link >= 1000
@@ -1211,6 +1243,8 @@ func (a *WsClient) Conf(cc string) (string, error) {
 				out = append(out, confstr)
 				ids = append(ids, strings.Join([]string{rc.CC, rc.City, rc.Name}, "/"))
 				c++
+			} else {
+				badc++
 			}
 		}
 		tot++
@@ -1220,6 +1254,16 @@ func (a *WsClient) Conf(cc string) (string, error) {
 		log.I("ws: conf: cc %s(%s): %d/%d => chosen (any? %t): %d[%s] (port: %s)", cc, city, c, len(out), chooseAny, r, ids[r], portstr)
 		return out[r], nil
 	}
+	if tot == 0 || v <= x { // fail open if all CCs excluded
+		logew(retried)("ws: conf: cc %s(%s): all visited(%d) / excluded(%d) / bad(%d); tot: %d / excl: %d; retry?",
+			cc, city, v, x, badc, tot, len(excl), !retried)
+		if !retried {
+			clear(excl) // fail open; excluded none
+			clear(visited)
+			retried = true
+			goto reconf
+		}
+	}
 	log.E("ws: conf: cc %s(%s) not found (tot: %d)", cc, city, tot)
 	return "", errWsNoCcConfig
 }
@@ -2555,6 +2599,22 @@ func createPermaCreds(h *http.Client, ent *WsEntitlement, bearer, pubkey string)
 	return &cfg, nil
 }
 
+// ccCsvAsSet mods a csv of country codes into a set.
+func ccCsvAsSet(csv string) map[string]struct{} {
+	if len(csv) <= 0 {
+		return nil
+	}
+	parts := strings.Split(csv, ",")
+	out := make(map[string]struct{}, len(parts))
+	for _, p := range parts {
+		p = strings.ToUpper(strings.TrimSpace(p))
+		if len(p) > 0 {
+			out[p] = struct{}{}
+		}
+	}
+	return out
+}
+
 func wsBriefPauseBeforeRetry() {
 	time.Sleep(2200 * time.Millisecond)
 }
@@ -2566,6 +2626,13 @@ func loge(err error) log.LogFn {
 	return log.E
 }
 
+func logew(cond bool) log.LogFn {
+	if cond {
+		return log.E
+	}
+	return log.W
+}
+
 func sha(p string) []byte {
 	return shab([]byte(p))
 }