|
- package martian
-
- import (
- "bufio"
- "bytes"
- "io/ioutil"
- "net"
- "net/http"
- "net/http/httptest"
- "strings"
- "testing"
- "time"
-
- "github.com/google/martian/log"
- "github.com/google/martian/martiantest"
- "github.com/google/martian/trafficshape"
- )
-
- // Tests that sending data of length 600 bytes with max bandwidth of 100 bytes/s takes
- // atleast 4.9s. Uses the Close Connection action to immediately close the connection
- // upon the proxy writing 600 bytes. (4.9s ~ 5s = 600b /100b/s - 1s)
- func TestConstantThrottleAndClose(t *testing.T) {
- log.SetLevel(log.Info)
-
- l, err := net.Listen("tcp", "[::]:0")
- if err != nil {
- t.Fatalf("net.Listen(): got %v, want no error", err)
- }
-
- tsl := trafficshape.NewListener(l)
- tsh := trafficshape.NewHandler(tsl)
-
- // This is the data to be sent.
- testString := strings.Repeat("0", 600)
-
- // Traffic shaping config request.
- jsonString :=
- `{
- "trafficshape": {
- "shapes": [
- {
- "url_regex": "http://example/example",
- "throttles": [
- {
- "bytes": "0-",
- "bandwidth": 100
- }
- ],
- "close_connections": [
- {
- "byte": 600,
- "count": 1
- }
- ]
- }
- ]
- }
- }`
-
- tsReq, err := http.NewRequest("POST", "test", bytes.NewBufferString(jsonString))
- rw := httptest.NewRecorder()
- tsh.ServeHTTP(rw, tsReq)
- res := rw.Result()
-
- if got, want := res.StatusCode, 200; got != want {
- t.Fatalf("res.StatusCode: got %d, want %d", got, want)
- }
-
- p := NewProxy()
- defer p.Close()
-
- p.SetRequestModifier(nil)
- p.SetResponseModifier(nil)
-
- tr := martiantest.NewTransport()
- p.SetRoundTripper(tr)
- p.SetTimeout(15 * time.Second)
-
- tm := martiantest.NewModifier()
-
- tm.RequestFunc(func(req *http.Request) {
- ctx := NewContext(req)
- ctx.SkipRoundTrip()
- })
-
- tm.ResponseFunc(func(res *http.Response) {
- res.StatusCode = http.StatusOK
- res.Body = ioutil.NopCloser(bytes.NewBufferString(testString))
- })
-
- p.SetRequestModifier(tm)
- p.SetResponseModifier(tm)
-
- go p.Serve(tsl)
-
- c1 := make(chan string)
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
-
- go func() {
- req, err := http.NewRequest("GET", "http://example/example", nil)
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
- c1 <- bodystr
- }()
-
- var bodystr string
- select {
- case bodystringc := <-c1:
- t.Errorf("took < 4.9s, should take at least 4.9s")
- bodystr = bodystringc
- case <-time.After(4900 * time.Millisecond):
- bodystringc := <-c1
- bodystr = bodystringc
- }
-
- if bodystr != testString {
- t.Errorf("res.Body: got %s, want %s", bodystr, testString)
- }
- }
-
- // Tests that sleeping for 5s and then closing the connection
- // upon reading 200 bytes, with a bandwidth of 5000 bytes/s
- // takes at least 4.9s, and results in a correctly trimmed
- // response body. (200 0s instead of 500 0s)
- func TestSleepAndClose(t *testing.T) {
- log.SetLevel(log.Info)
-
- l, err := net.Listen("tcp", "[::]:0")
- if err != nil {
- t.Fatalf("net.Listen(): got %v, want no error", err)
- }
-
- tsl := trafficshape.NewListener(l)
- tsh := trafficshape.NewHandler(tsl)
-
- // This is the data to be sent.
- testString := strings.Repeat("0", 500)
-
- // Traffic shaping config request.
- jsonString :=
- `{
- "trafficshape": {
- "shapes": [
- {
- "url_regex": "http://example/example",
- "throttles": [
- {
- "bytes": "0-",
- "bandwidth": 5000
- }
- ],
- "halts": [
- {
- "byte": 100,
- "duration": 5000,
- "count": 1
- }
- ],
- "close_connections": [
- {
- "byte": 200,
- "count": 1
- }
- ]
- }
- ]
- }
- }`
-
- tsReq, err := http.NewRequest("POST", "test", bytes.NewBufferString(jsonString))
- rw := httptest.NewRecorder()
- tsh.ServeHTTP(rw, tsReq)
- res := rw.Result()
-
- if got, want := res.StatusCode, 200; got != want {
- t.Fatalf("res.StatusCode: got %d, want %d", got, want)
- }
-
- p := NewProxy()
- defer p.Close()
-
- p.SetRequestModifier(nil)
- p.SetResponseModifier(nil)
-
- tr := martiantest.NewTransport()
- p.SetRoundTripper(tr)
- p.SetTimeout(15 * time.Second)
-
- tm := martiantest.NewModifier()
-
- tm.RequestFunc(func(req *http.Request) {
- ctx := NewContext(req)
- ctx.SkipRoundTrip()
- })
-
- tm.ResponseFunc(func(res *http.Response) {
- res.StatusCode = http.StatusOK
- res.Body = ioutil.NopCloser(bytes.NewBufferString(testString))
- })
-
- p.SetRequestModifier(tm)
- p.SetResponseModifier(tm)
-
- go p.Serve(tsl)
-
- c1 := make(chan string)
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
-
- go func() {
- req, err := http.NewRequest("GET", "http://example/example", nil)
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
- c1 <- bodystr
- }()
-
- var bodystr string
- select {
- case bodystringc := <-c1:
- t.Errorf("took < 4.9s, should take at least 4.9s")
- bodystr = bodystringc
- case <-time.After(4900 * time.Millisecond):
- bodystringc := <-c1
- bodystr = bodystringc
- }
-
- if want := strings.Repeat("0", 200); bodystr != want {
- t.Errorf("res.Body: got %s, want %s", bodystr, want)
- }
- }
-
- // Similar to TestConstantThrottleAndClose, except that it applies
- // the throttle only in a specific byte range, and modifies the
- // the response to lie in the byte range.
- func TestConstantThrottleAndCloseByteRange(t *testing.T) {
- log.SetLevel(log.Info)
-
- l, err := net.Listen("tcp", "[::]:0")
- if err != nil {
- t.Fatalf("net.Listen(): got %v, want no error", err)
- }
-
- tsl := trafficshape.NewListener(l)
- tsh := trafficshape.NewHandler(tsl)
-
- // This is the data to be sent.
- testString := strings.Repeat("0", 600)
-
- // Traffic shaping config request.
- jsonString :=
- `{
- "trafficshape": {
- "shapes": [
- {
- "url_regex": "http://example/example",
- "throttles": [
- {
- "bytes": "500-",
- "bandwidth": 100
- }
- ],
- "close_connections": [
- {
- "byte": 1100,
- "count": 1
- }
- ]
- }
- ]
- }
- }`
-
- tsReq, err := http.NewRequest("POST", "test", bytes.NewBufferString(jsonString))
- rw := httptest.NewRecorder()
- tsh.ServeHTTP(rw, tsReq)
- res := rw.Result()
-
- if got, want := res.StatusCode, 200; got != want {
- t.Fatalf("res.StatusCode: got %d, want %d", got, want)
- }
-
- p := NewProxy()
- defer p.Close()
-
- p.SetRequestModifier(nil)
- p.SetResponseModifier(nil)
-
- tr := martiantest.NewTransport()
- p.SetRoundTripper(tr)
- p.SetTimeout(15 * time.Second)
-
- tm := martiantest.NewModifier()
-
- tm.RequestFunc(func(req *http.Request) {
- ctx := NewContext(req)
- ctx.SkipRoundTrip()
- })
-
- tm.ResponseFunc(func(res *http.Response) {
- res.StatusCode = http.StatusPartialContent
- res.Body = ioutil.NopCloser(bytes.NewBufferString(testString))
- res.Header.Set("Content-Range", "bytes 500-1100/1100")
- })
-
- p.SetRequestModifier(tm)
- p.SetResponseModifier(tm)
-
- go p.Serve(tsl)
-
- c1 := make(chan string)
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
-
- go func() {
- req, err := http.NewRequest("GET", "http://example/example", nil)
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
-
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
- c1 <- bodystr
- }()
-
- var bodystr string
- select {
- case bodystringc := <-c1:
- t.Errorf("took < 4.9s, should take at least 4.9s")
- bodystr = bodystringc
- case <-time.After(4900 * time.Millisecond):
- bodystringc := <-c1
- bodystr = bodystringc
- }
-
- if bodystr != testString {
- t.Errorf("res.Body: got %s, want %s", bodystr, testString)
- }
- }
-
- // Opens up 5 concurrent connections, and sets the
- // max global bandwidth for the url regex to be 250b/s.
- // Every connection tries to read 500b of data, but since
- // the global bandwidth for the particular regex is 250,
- // it should take at least 5 * 500b / 250b/s -1s = 9s to read
- // everything.
- func TestMaxBandwidth(t *testing.T) {
- log.SetLevel(log.Info)
-
- l, err := net.Listen("tcp", "[::]:0")
- if err != nil {
- t.Fatalf("net.Listen(): got %v, want no error", err)
- }
-
- tsl := trafficshape.NewListener(l)
- tsh := trafficshape.NewHandler(tsl)
-
- // This is the data to be sent.
- testString := strings.Repeat("0", 500)
-
- // Traffic shaping config request.
- jsonString :=
- `{
- "trafficshape": {
- "shapes": [
- {
- "url_regex": "http://example/example",
- "max_global_bandwidth": 250,
- "close_connections": [
- {
- "byte": 500,
- "count": 5
- }
- ]
- }
- ]
- }
- }`
-
- tsReq, err := http.NewRequest("POST", "test", bytes.NewBufferString(jsonString))
- rw := httptest.NewRecorder()
- tsh.ServeHTTP(rw, tsReq)
- res := rw.Result()
-
- if got, want := res.StatusCode, 200; got != want {
- t.Fatalf("res.StatusCode: got %d, want %d", got, want)
- }
-
- p := NewProxy()
- defer p.Close()
-
- p.SetRequestModifier(nil)
- p.SetResponseModifier(nil)
-
- tr := martiantest.NewTransport()
- p.SetRoundTripper(tr)
- p.SetTimeout(20 * time.Second)
-
- tm := martiantest.NewModifier()
-
- tm.RequestFunc(func(req *http.Request) {
- ctx := NewContext(req)
- ctx.SkipRoundTrip()
- })
-
- tm.ResponseFunc(func(res *http.Response) {
- res.StatusCode = http.StatusOK
- res.Body = ioutil.NopCloser(bytes.NewBufferString(testString))
- })
-
- p.SetRequestModifier(tm)
- p.SetResponseModifier(tm)
-
- go p.Serve(tsl)
-
- numChannels := 5
-
- channels := make([]chan string, numChannels)
-
- for i := 0; i < numChannels; i++ {
- channels[i] = make(chan string)
- }
-
- for i := 0; i < numChannels; i++ {
- go func(i int) {
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
- req, err := http.NewRequest("GET", "http://example/example", nil)
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
-
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
-
- if i != 0 {
- <-channels[i-1]
- }
-
- channels[i] <- bodystr
- }(i)
- }
-
- var bodystr string
- select {
- case bodystringc := <-channels[numChannels-1]:
- t.Errorf("took < 8.9s, should take at least 8.9s")
- bodystr = bodystringc
- case <-time.After(8900 * time.Millisecond):
- bodystringc := <-channels[numChannels-1]
- bodystr = bodystringc
- }
-
- if bodystr != testString {
- t.Errorf("res.Body: got %s, want %s", bodystr, testString)
- }
- }
-
- // Makes 2 requests, with the first one having a byte range starting
- // at byte 250, and adds a close connection action at byte 450.
- // The first request should hit the action sooner,
- // and delete it. The second request should read the whole
- // data (500b)
- func TestConcurrentResponseActions(t *testing.T) {
- log.SetLevel(log.Info)
-
- l, err := net.Listen("tcp", "[::]:0")
- if err != nil {
- t.Fatalf("net.Listen(): got %v, want no error", err)
- }
-
- tsl := trafficshape.NewListener(l)
- tsh := trafficshape.NewHandler(tsl)
-
- // This is the data to be sent.
- testString := strings.Repeat("0", 500)
-
- // Traffic shaping config request.
- jsonString :=
- `{
- "trafficshape": {
- "shapes": [
- {
- "url_regex": "http://example/example",
- "throttles": [
- {
- "bytes": "-",
- "bandwidth": 250
- }
- ],
- "close_connections": [
- {
- "byte": 450,
- "count": 1
- },
- {
- "byte": 500,
- "count": 1
- }
- ]
- }
- ]
- }
- }`
-
- tsReq, err := http.NewRequest("POST", "test", bytes.NewBufferString(jsonString))
- rw := httptest.NewRecorder()
- tsh.ServeHTTP(rw, tsReq)
- res := rw.Result()
-
- if got, want := res.StatusCode, 200; got != want {
- t.Fatalf("res.StatusCode: got %d, want %d", got, want)
- }
-
- p := NewProxy()
- defer p.Close()
-
- p.SetRequestModifier(nil)
- p.SetResponseModifier(nil)
-
- tr := martiantest.NewTransport()
- p.SetRoundTripper(tr)
- p.SetTimeout(20 * time.Second)
-
- tm := martiantest.NewModifier()
-
- tm.RequestFunc(func(req *http.Request) {
- ctx := NewContext(req)
- ctx.SkipRoundTrip()
- })
-
- tm.ResponseFunc(func(res *http.Response) {
- cr := res.Request.Header.Get("ContentRange")
- res.StatusCode = http.StatusOK
- res.Body = ioutil.NopCloser(bytes.NewBufferString(testString))
- if cr != "" {
- res.StatusCode = http.StatusPartialContent
- res.Header.Set("Content-Range", cr)
- }
- })
-
- p.SetRequestModifier(tm)
- p.SetResponseModifier(tm)
-
- go p.Serve(tsl)
-
- c1 := make(chan string)
- c2 := make(chan string)
-
- go func() {
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
- req, err := http.NewRequest("GET", "http://example/example", nil)
- req.Header.Set("ContentRange", "bytes 250-1000/1000")
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
-
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
- c1 <- bodystr
- }()
-
- go func() {
- conn, err := net.Dial("tcp", l.Addr().String())
- defer conn.Close()
- if err != nil {
- t.Fatalf("net.Dial(): got %v, want no error", err)
- }
- req, err := http.NewRequest("GET", "http://example/example", nil)
- if err != nil {
- t.Fatalf("http.NewRequest(): got %v, want no error", err)
- }
-
- if err := req.WriteProxy(conn); err != nil {
- t.Fatalf("req.WriteProxy(): got %v, want no error", err)
- }
-
- res, err := http.ReadResponse(bufio.NewReader(conn), req)
- if err != nil {
- t.Fatalf("http.ReadResponse(): got %v, want no error", err)
- }
-
- body, _ := ioutil.ReadAll(res.Body)
- bodystr := string(body)
- c2 <- bodystr
- }()
-
- bodystr1 := <-c1
- bodystr2 := <-c2
-
- if want1 := strings.Repeat("0", 200); bodystr1 != want1 {
- t.Errorf("res.Body: got %s, want %s", bodystr1, want1)
- }
- if want2 := strings.Repeat("0", 500); bodystr2 != want2 {
- t.Errorf("res.Body: got %s, want %s", bodystr2, want2)
- }
- }
|