snowflake/broker/snowflake-broker_test.go
Cecylia Bocovich c4ae64905b Redo protocol for proxy--broker messages
Switch to containing all communication between the proxy and the broker
in the HTTP response body. This will make things easier if we ever use
something other than HTTP communicate between different actors in the
snowflake system.

Other changes to the protocol are as follows:
- requests are accompanied by a version number so the broker can be
backwards compatable if desired in the future
- all responses are 200 OK unless the request was badly formatted
2019-11-13 10:54:48 -05:00

525 lines
15 KiB
Go

package main
import (
"bytes"
"container/heap"
"io/ioutil"
"log"
"net"
"net/http"
"net/http/httptest"
"os"
"testing"
"time"
. "github.com/smartystreets/goconvey/convey"
)
func NullLogger() *log.Logger {
logger := log.New(os.Stdout, "", 0)
logger.SetOutput(ioutil.Discard)
return logger
}
func TestBroker(t *testing.T) {
Convey("Context", t, func() {
ctx := NewBrokerContext(NullLogger())
Convey("Adds Snowflake", func() {
So(ctx.snowflakes.Len(), ShouldEqual, 0)
So(len(ctx.idToSnowflake), ShouldEqual, 0)
ctx.AddSnowflake("foo")
So(ctx.snowflakes.Len(), ShouldEqual, 1)
So(len(ctx.idToSnowflake), ShouldEqual, 1)
})
Convey("Broker goroutine matches clients with proxies", func() {
p := new(ProxyPoll)
p.id = "test"
p.offerChannel = make(chan []byte)
go func(ctx *BrokerContext) {
ctx.proxyPolls <- p
close(ctx.proxyPolls)
}(ctx)
ctx.Broker()
So(ctx.snowflakes.Len(), ShouldEqual, 1)
snowflake := heap.Pop(ctx.snowflakes).(*Snowflake)
snowflake.offerChannel <- []byte("test offer")
offer := <-p.offerChannel
So(ctx.idToSnowflake["test"], ShouldNotBeNil)
So(offer, ShouldResemble, []byte("test offer"))
So(ctx.snowflakes.Len(), ShouldEqual, 0)
})
Convey("Request an offer from the Snowflake Heap", func() {
done := make(chan []byte)
go func() {
offer := ctx.RequestOffer("test")
done <- offer
}()
request := <-ctx.proxyPolls
request.offerChannel <- []byte("test offer")
offer := <-done
So(offer, ShouldResemble, []byte("test offer"))
})
Convey("Responds to client offers...", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte("test"))
r, err := http.NewRequest("POST", "snowflake.broker/client", data)
So(err, ShouldBeNil)
Convey("with 503 when no snowflakes are available.", func() {
clientOffers(ctx, w, r)
So(w.Code, ShouldEqual, http.StatusServiceUnavailable)
So(w.Body.String(), ShouldEqual, "")
})
Convey("with a proxy answer if available.", func() {
done := make(chan bool)
// Prepare a fake proxy to respond with.
snowflake := ctx.AddSnowflake("fake")
go func() {
clientOffers(ctx, w, r)
done <- true
}()
offer := <-snowflake.offerChannel
So(offer, ShouldResemble, []byte("test"))
snowflake.answerChannel <- []byte("fake answer")
<-done
So(w.Body.String(), ShouldEqual, "fake answer")
So(w.Code, ShouldEqual, http.StatusOK)
})
Convey("Times out when no proxy responds.", func() {
if testing.Short() {
return
}
done := make(chan bool)
snowflake := ctx.AddSnowflake("fake")
go func() {
clientOffers(ctx, w, r)
// Takes a few seconds here...
done <- true
}()
offer := <-snowflake.offerChannel
So(offer, ShouldResemble, []byte("test"))
<-done
So(w.Code, ShouldEqual, http.StatusGatewayTimeout)
})
})
Convey("Responds to proxy polls...", func() {
done := make(chan bool)
w := httptest.NewRecorder()
data := bytes.NewReader([]byte(`{"Sid":"ymbcCMto7KHNGYlp","Version":"1.0"}`))
r, err := http.NewRequest("POST", "snowflake.broker/proxy", data)
So(err, ShouldBeNil)
Convey("with a client offer if available.", func() {
go func(ctx *BrokerContext) {
proxyPolls(ctx, w, r)
done <- true
}(ctx)
// Pass a fake client offer to this proxy
p := <-ctx.proxyPolls
So(p.id, ShouldEqual, "ymbcCMto7KHNGYlp")
p.offerChannel <- []byte("fake offer")
<-done
So(w.Code, ShouldEqual, http.StatusOK)
So(w.Body.String(), ShouldEqual, `{"Status":"client match","Offer":"fake offer"}`)
})
Convey("return empty 200 OK when no client offer is available.", func() {
go func(ctx *BrokerContext) {
proxyPolls(ctx, w, r)
done <- true
}(ctx)
p := <-ctx.proxyPolls
So(p.id, ShouldEqual, "ymbcCMto7KHNGYlp")
// nil means timeout
p.offerChannel <- nil
<-done
So(w.Body.String(), ShouldEqual, `{"Status":"no match","Offer":""}`)
So(w.Code, ShouldEqual, http.StatusOK)
})
})
Convey("Responds to proxy answers...", func() {
s := ctx.AddSnowflake("test")
w := httptest.NewRecorder()
data := bytes.NewReader([]byte(`{"Version":"1.0","Sid":"test","Answer":"test"}`))
Convey("by passing to the client if valid.", func() {
r, err := http.NewRequest("POST", "snowflake.broker/answer", data)
So(err, ShouldBeNil)
go func(ctx *BrokerContext) {
proxyAnswers(ctx, w, r)
}(ctx)
answer := <-s.answerChannel
So(w.Code, ShouldEqual, http.StatusOK)
So(answer, ShouldResemble, []byte("test"))
})
Convey("with client gone status if the proxy is not recognized", func() {
data = bytes.NewReader([]byte(`{"Version":"1.0","Sid":"invalid","Answer":"test"}`))
r, err := http.NewRequest("POST", "snowflake.broker/answer", data)
So(err, ShouldBeNil)
proxyAnswers(ctx, w, r)
So(w.Code, ShouldEqual, http.StatusOK)
b, err := ioutil.ReadAll(w.Body)
So(err, ShouldBeNil)
So(b, ShouldResemble, []byte(`{"Status":"client gone"}`))
})
Convey("with error if the proxy gives invalid answer", func() {
data := bytes.NewReader(nil)
r, err := http.NewRequest("POST", "snowflake.broker/answer", data)
So(err, ShouldBeNil)
proxyAnswers(ctx, w, r)
So(w.Code, ShouldEqual, http.StatusBadRequest)
})
Convey("with error if the proxy writes too much data", func() {
data := bytes.NewReader(make([]byte, 100001))
r, err := http.NewRequest("POST", "snowflake.broker/answer", data)
So(err, ShouldBeNil)
proxyAnswers(ctx, w, r)
So(w.Code, ShouldEqual, http.StatusBadRequest)
})
})
})
Convey("End-To-End", t, func() {
done := make(chan bool)
polled := make(chan bool)
ctx := NewBrokerContext(NullLogger())
// Proxy polls with its ID first...
dataP := bytes.NewReader([]byte(`{"Sid":"ymbcCMto7KHNGYlp","Version":"1.0"}`))
wP := httptest.NewRecorder()
rP, err := http.NewRequest("POST", "snowflake.broker/proxy", dataP)
So(err, ShouldBeNil)
go func() {
proxyPolls(ctx, wP, rP)
polled <- true
}()
// Manually do the Broker goroutine action here for full control.
p := <-ctx.proxyPolls
So(p.id, ShouldEqual, "ymbcCMto7KHNGYlp")
s := ctx.AddSnowflake(p.id)
go func() {
offer := <-s.offerChannel
p.offerChannel <- offer
}()
So(ctx.idToSnowflake["ymbcCMto7KHNGYlp"], ShouldNotBeNil)
// Client request blocks until proxy answer arrives.
dataC := bytes.NewReader([]byte("fake offer"))
wC := httptest.NewRecorder()
rC, err := http.NewRequest("POST", "snowflake.broker/client", dataC)
So(err, ShouldBeNil)
go func() {
clientOffers(ctx, wC, rC)
done <- true
}()
<-polled
So(wP.Code, ShouldEqual, http.StatusOK)
So(wP.Body.String(), ShouldResemble, `{"Status":"client match","Offer":"fake offer"}`)
So(ctx.idToSnowflake["ymbcCMto7KHNGYlp"], ShouldNotBeNil)
// Follow up with the answer request afterwards
wA := httptest.NewRecorder()
dataA := bytes.NewReader([]byte(`{"Version":"1.0","Sid":"ymbcCMto7KHNGYlp","Answer":"test"}`))
rA, err := http.NewRequest("POST", "snowflake.broker/answer", dataA)
So(err, ShouldBeNil)
proxyAnswers(ctx, wA, rA)
So(wA.Code, ShouldEqual, http.StatusOK)
<-done
So(wC.Code, ShouldEqual, http.StatusOK)
So(wC.Body.String(), ShouldEqual, "test")
})
}
func TestSnowflakeHeap(t *testing.T) {
Convey("SnowflakeHeap", t, func() {
h := new(SnowflakeHeap)
heap.Init(h)
So(h.Len(), ShouldEqual, 0)
s1 := new(Snowflake)
s2 := new(Snowflake)
s3 := new(Snowflake)
s4 := new(Snowflake)
s1.clients = 4
s2.clients = 5
s3.clients = 3
s4.clients = 1
heap.Push(h, s1)
So(h.Len(), ShouldEqual, 1)
heap.Push(h, s2)
So(h.Len(), ShouldEqual, 2)
heap.Push(h, s3)
So(h.Len(), ShouldEqual, 3)
heap.Push(h, s4)
So(h.Len(), ShouldEqual, 4)
heap.Remove(h, 0)
So(h.Len(), ShouldEqual, 3)
r := heap.Pop(h).(*Snowflake)
So(h.Len(), ShouldEqual, 2)
So(r.clients, ShouldEqual, 3)
So(r.index, ShouldEqual, -1)
r = heap.Pop(h).(*Snowflake)
So(h.Len(), ShouldEqual, 1)
So(r.clients, ShouldEqual, 4)
So(r.index, ShouldEqual, -1)
r = heap.Pop(h).(*Snowflake)
So(h.Len(), ShouldEqual, 0)
So(r.clients, ShouldEqual, 5)
So(r.index, ShouldEqual, -1)
})
}
func TestGeoip(t *testing.T) {
Convey("Geoip", t, func() {
tv4 := new(GeoIPv4Table)
err := GeoIPLoadFile(tv4, "test_geoip")
So(err, ShouldEqual, nil)
tv6 := new(GeoIPv6Table)
err = GeoIPLoadFile(tv6, "test_geoip6")
So(err, ShouldEqual, nil)
Convey("IPv4 Country Mapping Tests", func() {
for _, test := range []struct {
addr, cc string
ok bool
}{
{
"129.97.208.23", //uwaterloo
"CA",
true,
},
{
"127.0.0.1",
"",
false,
},
{
"255.255.255.255",
"",
false,
},
{
"0.0.0.0",
"",
false,
},
{
"223.252.127.255", //test high end of range
"JP",
true,
},
{
"223.252.127.255", //test low end of range
"JP",
true,
},
} {
country, ok := GetCountryByAddr(tv4, net.ParseIP(test.addr))
So(country, ShouldEqual, test.cc)
So(ok, ShouldResemble, test.ok)
}
})
Convey("IPv6 Country Mapping Tests", func() {
for _, test := range []struct {
addr, cc string
ok bool
}{
{
"2620:101:f000:0:250:56ff:fe80:168e", //uwaterloo
"CA",
true,
},
{
"fd00:0:0:0:0:0:0:1",
"",
false,
},
{
"0:0:0:0:0:0:0:0",
"",
false,
},
{
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff",
"",
false,
},
{
"2a07:2e47:ffff:ffff:ffff:ffff:ffff:ffff", //test high end of range
"FR",
true,
},
{
"2a07:2e40::", //test low end of range
"FR",
true,
},
} {
country, ok := GetCountryByAddr(tv6, net.ParseIP(test.addr))
So(country, ShouldEqual, test.cc)
So(ok, ShouldResemble, test.ok)
}
})
// Make sure things behave properly if geoip file fails to load
ctx := NewBrokerContext(NullLogger())
if err := ctx.metrics.LoadGeoipDatabases("invalid_filename", "invalid_filename6"); err != nil {
log.Printf("loading geo ip databases returned error: %v", err)
}
ctx.metrics.UpdateCountryStats("127.0.0.1")
So(ctx.metrics.tablev4, ShouldEqual, nil)
})
}
func TestMetrics(t *testing.T) {
Convey("Test metrics...", t, func() {
done := make(chan bool)
buf := new(bytes.Buffer)
ctx := NewBrokerContext(log.New(buf, "", 0))
err := ctx.metrics.LoadGeoipDatabases("test_geoip", "test_geoip6")
So(err, ShouldEqual, nil)
//Test addition of proxy polls
Convey("for proxy polls", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte("{\"Sid\":\"ymbcCMto7KHNGYlp\",\"Version\":\"1.0\"}"))
r, err := http.NewRequest("POST", "snowflake.broker/proxy", data)
r.Header.Set("X-Session-ID", "test")
r.RemoteAddr = "129.97.208.23:8888" //CA geoip
So(err, ShouldBeNil)
go func(ctx *BrokerContext) {
proxyPolls(ctx, w, r)
done <- true
}(ctx)
p := <-ctx.proxyPolls //manually unblock poll
p.offerChannel <- nil
<-done
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips CA=1\nsnowflake-ips-total 1\nsnowflake-idle-count 8\nclient-denied-count 0\nclient-snowflake-match-count 0\n")
})
//Test addition of client failures
Convey("for no proxies available", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte("test"))
r, err := http.NewRequest("POST", "snowflake.broker/client", data)
So(err, ShouldBeNil)
clientOffers(ctx, w, r)
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips \nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 8\nclient-snowflake-match-count 0\n")
// Test reset
buf.Reset()
ctx.metrics.zeroMetrics()
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips \nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 0\nclient-snowflake-match-count 0\n")
})
//Test addition of client matches
Convey("for client-proxy match", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte("test"))
r, err := http.NewRequest("POST", "snowflake.broker/client", data)
So(err, ShouldBeNil)
// Prepare a fake proxy to respond with.
snowflake := ctx.AddSnowflake("fake")
go func() {
clientOffers(ctx, w, r)
done <- true
}()
offer := <-snowflake.offerChannel
So(offer, ShouldResemble, []byte("test"))
snowflake.answerChannel <- []byte("fake answer")
<-done
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips \nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 0\nclient-snowflake-match-count 8\n")
})
//Test rounding boundary
Convey("binning boundary", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte("test"))
r, err := http.NewRequest("POST", "snowflake.broker/client", data)
So(err, ShouldBeNil)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
clientOffers(ctx, w, r)
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips \nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 8\nclient-snowflake-match-count 0\n")
clientOffers(ctx, w, r)
buf.Reset()
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips \nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 16\nclient-snowflake-match-count 0\n")
})
//Test unique ip
Convey("proxy counts by unique ip", func() {
w := httptest.NewRecorder()
data := bytes.NewReader([]byte(`{"Sid":"ymbcCMto7KHNGYlp","Version":"1.0"}`))
r, err := http.NewRequest("POST", "snowflake.broker/proxy", data)
r.Header.Set("X-Session-ID", "test")
r.RemoteAddr = "129.97.208.23:8888" //CA geoip
So(err, ShouldBeNil)
go func(ctx *BrokerContext) {
proxyPolls(ctx, w, r)
done <- true
}(ctx)
p := <-ctx.proxyPolls //manually unblock poll
p.offerChannel <- nil
<-done
data = bytes.NewReader([]byte(`{"Sid":"ymbcCMto7KHNGYlp","Version":"1.0"}`))
r, err = http.NewRequest("POST", "snowflake.broker/proxy", data)
if err != nil {
log.Printf("unable to get NewRequest with error: %v", err)
}
r.Header.Set("X-Session-ID", "test")
r.RemoteAddr = "129.97.208.23:8888" //CA geoip
go func(ctx *BrokerContext) {
proxyPolls(ctx, w, r)
done <- true
}(ctx)
p = <-ctx.proxyPolls //manually unblock poll
p.offerChannel <- nil
<-done
ctx.metrics.printMetrics()
So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips CA=1\nsnowflake-ips-total 1\nsnowflake-idle-count 8\nclient-denied-count 0\nclient-snowflake-match-count 0\n")
})
})
}