Merge remote-tracking branch 'gitlab/mr/187'

2 files changed, 33 insertions, 12 deletions

diff --git a/broker/ipc.go b/broker/ipc.go
index 7cbdd06..132dba7 100644
--- a/broker/ipc.go
+++ b/broker/ipc.go
@@ -233,22 +233,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