mirror of
https://github.com/dutchcoders/transfer.sh.git
synced 2020-11-18 19:53:40 -08:00
665 lines
15 KiB
Go
665 lines
15 KiB
Go
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)
|
|
}
|
|
}
|