From 4ff36e3f071bd21cd2d2292b952fdc9430c99934 Mon Sep 17 00:00:00 2001
From: WofWca <wofwca@protonmail.com>
Date: Mon, 2 Oct 2023 21:26:13 +0400
Subject: [PATCH] improvement(broker): don't reject unrestricted client if
 there are no restricted proxies

I.e. match it with an unrestricted proxy (if there is one).

The old behavior exists since the inception of the restricted vs
unrestricted feature, i.e. 0052c0e10c
---
 broker/ipc.go                   | 21 +++++++++------------
 broker/snowflake-broker_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/broker/ipc.go b/broker/ipc.go
index b77c579..94ea53e 100644
--- a/broker/ipc.go
+++ b/broker/ipc.go
@@ -234,22 +234,19 @@ func (i *IPC) ClientOffers(arg messages.Arg, response *[]byte) error {
 }
 
 func (i *IPC) matchSnowflake(natType string) *Snowflake {
-	// Only hand out known restricted snowflakes to unrestricted clients
-	var snowflakeHeap *SnowflakeHeap
-	if natType == NATUnrestricted {
-		snowflakeHeap = i.ctx.restrictedSnowflakes
-	} else {
-		snowflakeHeap = i.ctx.snowflakes
-	}
-
 	i.ctx.snowflakeLock.Lock()
 	defer i.ctx.snowflakeLock.Unlock()
 
-	if snowflakeHeap.Len() > 0 {
-		return heap.Pop(snowflakeHeap).(*Snowflake)
-	} else {
-		return nil
+	// Proiritize known restricted snowflakes for unrestricted clients
+	if natType == NATUnrestricted && i.ctx.restrictedSnowflakes.Len() > 0 {
+		return heap.Pop(i.ctx.restrictedSnowflakes).(*Snowflake)
 	}
+
+	if i.ctx.snowflakes.Len() > 0 {
+		return heap.Pop(i.ctx.snowflakes).(*Snowflake)
+	}
+
+	return nil
 }
 
 func (i *IPC) ProxyAnswers(arg messages.Arg, response *[]byte) error {
diff --git a/broker/snowflake-broker_test.go b/broker/snowflake-broker_test.go
index b42e094..e335258 100644
--- a/broker/snowflake-broker_test.go
+++ b/broker/snowflake-broker_test.go
@@ -167,6 +167,30 @@ func TestBroker(t *testing.T) {
 				So(w.Code, ShouldEqual, http.StatusOK)
 			})
 
+			Convey("with unrestricted proxy to unrestricted client if there are no restricted proxies", func() {
+				snowflake := ctx.AddSnowflake("test", "", NATUnrestricted, 0)
+				offerData, err := createClientOffer(sdp, NATUnrestricted, "")
+				So(err, ShouldBeNil)
+				r, err := http.NewRequest("POST", "snowflake.broker/client", offerData)
+
+				done := make(chan bool)
+				go func() {
+					clientOffers(i, w, r)
+					done <- true
+				}()
+
+				select {
+				case <-snowflake.offerChannel:
+				case <-time.After(250 * time.Millisecond):
+					So(false, ShouldBeTrue)
+					return
+				}
+				snowflake.answerChannel <- "test answer"
+
+				<-done
+				So(w.Body.String(), ShouldEqual, `{"answer":"test answer"}`)
+			})
+
 			Convey("Times out when no proxy responds.", func() {
 				if testing.Short() {
 					return