Have broker pass client NAT type to proxy

This will allow browser-based proxies that are unable to determine their
NAT type to conservatively label themselves as restricted NATs if they
fail to work with clients that have restricted NATs.

7 files changed, 61 insertions, 42 deletions

diff --git a/broker/broker.go b/broker/broker.go
index 9297980..99f6c69 100644
--- a/broker/broker.go
+++ b/broker/broker.go
@@ -111,17 +111,17 @@ type ProxyPoll struct {
 	id           string
 	proxyType    string
 	natType      string
-	offerChannel chan []byte
+	offerChannel chan *ClientOffer
 }
 
 // Registers a Snowflake and waits for some Client to send an offer,
 // as part of the polling logic of the proxy handler.
-func (ctx *BrokerContext) RequestOffer(id string, proxyType string, natType string) []byte {
+func (ctx *BrokerContext) RequestOffer(id string, proxyType string, natType string) *ClientOffer {
 	request := new(ProxyPoll)
 	request.id = id
 	request.proxyType = proxyType
 	request.natType = natType
-	request.offerChannel = make(chan []byte)
+	request.offerChannel = make(chan *ClientOffer)
 	ctx.proxyPolls <- request
 	// Block until an offer is available, or timeout which sends a nil offer.
 	offer := <-request.offerChannel
@@ -165,7 +165,7 @@ func (ctx *BrokerContext) AddSnowflake(id string, proxyType string, natType stri
 	snowflake.id = id
 	snowflake.clients = 0
 	snowflake.proxyType = proxyType
-	snowflake.offerChannel = make(chan []byte)
+	snowflake.offerChannel = make(chan *ClientOffer)
 	snowflake.answerChannel = make(chan []byte)
 	ctx.snowflakeLock.Lock()
 	if natType == NATRestricted {
@@ -213,7 +213,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 		ctx.metrics.proxyIdleCount++
 		ctx.metrics.lock.Unlock()
 
-		b, err = messages.EncodePollResponse("", false)
+		b, err = messages.EncodePollResponse("", false, "")
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
 			return
@@ -222,7 +222,7 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 		w.Write(b)
 		return
 	}
-	b, err = messages.EncodePollResponse(string(offer), true)
+	b, err = messages.EncodePollResponse(string(offer.sdp), true, offer.natType)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
@@ -232,28 +232,37 @@ func proxyPolls(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+// Client offer contains an SDP and the NAT type of the client
+type ClientOffer struct {
+	natType string
+	sdp     []byte
+}
+
 /*
 Expects a WebRTC SDP offer in the Request to give to an assigned
 snowflake proxy, which responds with the SDP answer to be sent in
 the HTTP response back to the client.
 */
 func clientOffers(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
+	var err error
+
 	startTime := time.Now()
-	offer, err := ioutil.ReadAll(http.MaxBytesReader(w, r.Body, readLimit))
+	offer := &ClientOffer{}
+	offer.sdp, err = ioutil.ReadAll(http.MaxBytesReader(w, r.Body, readLimit))
 	if nil != err {
 		log.Println("Invalid data.")
 		w.WriteHeader(http.StatusBadRequest)
 		return
 	}
 
-	natType := r.Header.Get("Snowflake-NAT-Type")
-	if natType == "" {
-		natType = NATUnknown
+	offer.natType = r.Header.Get("Snowflake-NAT-Type")
+	if offer.natType == "" {
+		offer.natType = NATUnknown
 	}
 
 	// Only hand out known restricted snowflakes to unrestricted clients
 	var snowflakeHeap *SnowflakeHeap
-	if natType == NATUnrestricted {
+	if offer.natType == NATUnrestricted {
 		snowflakeHeap = ctx.restrictedSnowflakes
 	} else {
 		snowflakeHeap = ctx.snowflakes
diff --git a/broker/snowflake-broker_test.go b/broker/snowflake-broker_test.go
index 91383a1..d03dca7 100644
--- a/broker/snowflake-broker_test.go
+++ b/broker/snowflake-broker_test.go
@@ -37,7 +37,7 @@ func TestBroker(t *testing.T) {
 		Convey("Broker goroutine matches clients with proxies", func() {
 			p := new(ProxyPoll)
 			p.id = "test"
-			p.offerChannel = make(chan []byte)
+			p.offerChannel = make(chan *ClientOffer)
 			go func(ctx *BrokerContext) {
 				ctx.proxyPolls <- p
 				close(ctx.proxyPolls)
@@ -45,23 +45,23 @@ func TestBroker(t *testing.T) {
 			ctx.Broker()
 			So(ctx.snowflakes.Len(), ShouldEqual, 1)
 			snowflake := heap.Pop(ctx.snowflakes).(*Snowflake)
-			snowflake.offerChannel <- []byte("test offer")
+			snowflake.offerChannel <- &ClientOffer{sdp: []byte("test offer")}
 			offer := <-p.offerChannel
 			So(ctx.idToSnowflake["test"], ShouldNotBeNil)
-			So(offer, ShouldResemble, []byte("test offer"))
+			So(offer.sdp, ShouldResemble, []byte("test offer"))
 			So(ctx.snowflakes.Len(), ShouldEqual, 0)
 		})
 
 		Convey("Request an offer from the Snowflake Heap", func() {
-			done := make(chan []byte)
+			done := make(chan *ClientOffer)
 			go func() {
 				offer := ctx.RequestOffer("test", "", NATUnknown)
 				done <- offer
 			}()
 			request := <-ctx.proxyPolls
-			request.offerChannel <- []byte("test offer")
+			request.offerChannel <- &ClientOffer{sdp: []byte("test offer")}
 			offer := <-done
-			So(offer, ShouldResemble, []byte("test offer"))
+			So(offer.sdp, ShouldResemble, []byte("test offer"))
 		})
 
 		Convey("Responds to client offers...", func() {
@@ -85,7 +85,7 @@ func TestBroker(t *testing.T) {
 					done <- true
 				}()
 				offer := <-snowflake.offerChannel
-				So(offer, ShouldResemble, []byte("test"))
+				So(offer.sdp, ShouldResemble, []byte("test"))
 				snowflake.answerChannel <- []byte("fake answer")
 				<-done
 				So(w.Body.String(), ShouldEqual, "fake answer")
@@ -104,7 +104,7 @@ func TestBroker(t *testing.T) {
 					done <- true
 				}()
 				offer := <-snowflake.offerChannel
-				So(offer, ShouldResemble, []byte("test"))
+				So(offer.sdp, ShouldResemble, []byte("test"))
 				<-done
 				So(w.Code, ShouldEqual, http.StatusGatewayTimeout)
 			})
@@ -125,10 +125,10 @@ func TestBroker(t *testing.T) {
 				// Pass a fake client offer to this proxy
 				p := <-ctx.proxyPolls
 				So(p.id, ShouldEqual, "ymbcCMto7KHNGYlp")
-				p.offerChannel <- []byte("fake offer")
+				p.offerChannel <- &ClientOffer{sdp: []byte("fake offer")}
 				<-done
 				So(w.Code, ShouldEqual, http.StatusOK)
-				So(w.Body.String(), ShouldEqual, `{"Status":"client match","Offer":"fake offer"}`)
+				So(w.Body.String(), ShouldEqual, `{"Status":"client match","Offer":"fake offer","NAT":""}`)
 			})
 
 			Convey("return empty 200 OK when no client offer is available.", func() {
@@ -141,7 +141,7 @@ func TestBroker(t *testing.T) {
 				// nil means timeout
 				p.offerChannel <- nil
 				<-done
-				So(w.Body.String(), ShouldEqual, `{"Status":"no match","Offer":""}`)
+				So(w.Body.String(), ShouldEqual, `{"Status":"no match","Offer":"","NAT":""}`)
 				So(w.Code, ShouldEqual, http.StatusOK)
 			})
 		})
@@ -279,7 +279,7 @@ func TestBroker(t *testing.T) {
 
 			<-polled
 			So(wP.Code, ShouldEqual, http.StatusOK)
-			So(wP.Body.String(), ShouldResemble, `{"Status":"client match","Offer":"fake offer"}`)
+			So(wP.Body.String(), ShouldResemble, `{"Status":"client match","Offer":"fake offer","NAT":"unknown"}`)
 			So(ctx.idToSnowflake["ymbcCMto7KHNGYlp"], ShouldNotBeNil)
 			// Follow up with the answer request afterwards
 			wA := httptest.NewRecorder()
@@ -543,7 +543,7 @@ func TestMetrics(t *testing.T) {
 				done <- true
 			}()
 			offer := <-snowflake.offerChannel
-			So(offer, ShouldResemble, []byte("test"))
+			So(offer.sdp, ShouldResemble, []byte("test"))
 			snowflake.answerChannel <- []byte("fake answer")
 			<-done
 
diff --git a/broker/snowflake-heap.go b/broker/snowflake-heap.go
index 19a64b2..12fe557 100644
--- a/broker/snowflake-heap.go
+++ b/broker/snowflake-heap.go
@@ -11,7 +11,7 @@ over the offer and answer channels.
 type Snowflake struct {
 	id            string
 	proxyType     string
-	offerChannel  chan []byte
+	offerChannel  chan *ClientOffer
 	answerChannel chan []byte
 	clients       int
 	index         int
diff --git a/common/messages/proxy.go b/common/messages/proxy.go
index 923189b..2d9e58d 100644
--- a/common/messages/proxy.go
+++ b/common/messages/proxy.go
@@ -29,7 +29,8 @@ HTTP 200 OK
   {
     type: offer,
     sdp: [WebRTC SDP]
-  }
+  },
+  NAT: ["unknown"|"restricted"|"unrestricted"]
 }
 
 2) If a client is not matched:
@@ -120,13 +121,15 @@ func DecodePollRequest(data []byte) (string, string, string, error) {
 type ProxyPollResponse struct {
 	Status string
 	Offer  string
+	NAT    string
 }
 
-func EncodePollResponse(offer string, success bool) ([]byte, error) {
+func EncodePollResponse(offer string, success bool, natType string) ([]byte, error) {
 	if success {
 		return json.Marshal(ProxyPollResponse{
 			Status: "client match",
 			Offer:  offer,
+			NAT:    natType,
 		})
 
 	}
@@ -135,28 +138,33 @@ func EncodePollResponse(offer string, success bool) ([]byte, error) {
 	})
 }
 
-// Decodes a poll response from the broker and returns an offer
+// Decodes a poll response from the broker and returns an offer and the client's NAT type
 // If there is a client match, the returned offer string will be non-empty
-func DecodePollResponse(data []byte) (string, error) {
+func DecodePollResponse(data []byte) (string, string, error) {
 	var message ProxyPollResponse
 
 	err := json.Unmarshal(data, &message)
 	if err != nil {
-		return "", err
+		return "", "", err
 	}
 	if message.Status == "" {
-		return "", fmt.Errorf("received invalid data")
+		return "", "", fmt.Errorf("received invalid data")
 	}
 
 	if message.Status == "client match" {
 		if message.Offer == "" {
-			return "", fmt.Errorf("no supplied offer")
+			return "", "", fmt.Errorf("no supplied offer")
 		}
 	} else {
 		message.Offer = ""
 	}
 
-	return message.Offer, nil
+	natType := message.NAT
+	if natType == "" {
+		natType = "unknown"
+	}
+
+	return message.Offer, natType, nil
 }
 
 type ProxyAnswerRequest struct {
diff --git a/common/messages/proxy_test.go b/common/messages/proxy_test.go
index 3aa67fb..f4191e1 100644
--- a/common/messages/proxy_test.go
+++ b/common/messages/proxy_test.go
@@ -109,7 +109,7 @@ func TestDecodeProxyPollResponse(t *testing.T) {
 		}{
 			{
 				"fake offer",
-				`{"Status":"client match","Offer":"fake offer"}`,
+				`{"Status":"client match","Offer":"fake offer","NAT":"unknown"}`,
 				nil,
 			},
 			{
@@ -128,9 +128,9 @@ func TestDecodeProxyPollResponse(t *testing.T) {
 				fmt.Errorf(""),
 			},
 		} {
-			offer, err := DecodePollResponse([]byte(test.data))
-			So(offer, ShouldResemble, test.offer)
+			offer, _, err := DecodePollResponse([]byte(test.data))
 			So(err, ShouldHaveSameTypeAs, test.err)
+			So(offer, ShouldResemble, test.offer)
 		}
 
 	})
@@ -138,16 +138,18 @@ func TestDecodeProxyPollResponse(t *testing.T) {
 
 func TestEncodeProxyPollResponse(t *testing.T) {
 	Convey("Context", t, func() {
-		b, err := EncodePollResponse("fake offer", true)
+		b, err := EncodePollResponse("fake offer", true, "restricted")
 		So(err, ShouldEqual, nil)
-		offer, err := DecodePollResponse(b)
+		offer, natType, err := DecodePollResponse(b)
 		So(offer, ShouldEqual, "fake offer")
+		So(natType, ShouldEqual, "restricted")
 		So(err, ShouldEqual, nil)
 
-		b, err = EncodePollResponse("", false)
+		b, err = EncodePollResponse("", false, "unknown")
 		So(err, ShouldEqual, nil)
-		offer, err = DecodePollResponse(b)
+		offer, natType, err = DecodePollResponse(b)
 		So(offer, ShouldEqual, "")
+		So(natType, ShouldEqual, "unknown")
 		So(err, ShouldEqual, nil)
 	})
 }
diff --git a/proxy/proxy-go_test.go b/proxy/proxy-go_test.go
index 03d7307..168ca25 100644
--- a/proxy/proxy-go_test.go
+++ b/proxy/proxy-go_test.go
@@ -249,7 +249,7 @@ func TestBrokerInteractions(t *testing.T) {
 		Convey("polls broker correctly", func() {
 			var err error
 
-			b, err := messages.EncodePollResponse(sampleOffer, true)
+			b, err := messages.EncodePollResponse(sampleOffer, true, "unknown")
 			So(err, ShouldEqual, nil)
 			broker.transport = &MockTransport{
 				http.StatusOK,
diff --git a/proxy/snowflake.go b/proxy/snowflake.go
index ac67748..a1886c2 100644
--- a/proxy/snowflake.go
+++ b/proxy/snowflake.go
@@ -201,7 +201,7 @@ func (b *Broker) pollOffer(sid string) *webrtc.SessionDescription {
 					log.Printf("error reading broker response: %s", err)
 				} else {
 
-					offer, err := messages.DecodePollResponse(body)
+					offer, _, err := messages.DecodePollResponse(body)
 					if err != nil {
 						log.Printf("error reading broker response: %s", err.Error())
 						log.Printf("body: %s", body)