1
0
mirror of https://github.com/dutchcoders/transfer.sh.git synced 2020-11-18 19:53:40 -08:00
2019-03-17 20:19:56 +01:00

406 lines
11 KiB
Go

// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package marbl
import (
"bytes"
"io"
"net/http"
"strconv"
"strings"
"testing"
"time"
"github.com/google/martian"
"github.com/google/martian/proxyutil"
)
func TestMarkAPIRequestsWithHeader(t *testing.T) {
areq, err := http.NewRequest("POST", "http://localhost:8080/configure", nil)
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
ctx, remove, err := martian.TestContext(areq, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
ctx.APIRequest()
req, err := http.NewRequest("GET", "http://example.com", nil)
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, removereq, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer removereq()
var b bytes.Buffer
s := NewStream(&b)
s.LogRequest("00000000", areq)
s.LogRequest("00000001", req)
s.Close()
headers := make(map[string]string)
reader := NewReader(&b)
for {
frame, err := reader.ReadFrame()
if frame == nil {
break
}
if err != nil && err != io.EOF {
t.Fatalf("reader.ReadFrame(): got %v, want no error or io.EOF", err)
}
headerFrame, ok := frame.(Header)
if !ok {
t.Fatalf("frame.(Header): couldn't convert frame '%v' to a headerFrame", frame)
}
headers[headerFrame.ID+headerFrame.Name] = headerFrame.Value
}
apih, ok := headers["00000000:api"]
if !ok {
t.Errorf("headers[00000000:api]: got no such header, want :api (headers were: %v)", headers)
}
if got, want := apih, "true"; got != want {
t.Errorf("headers[%q]: got %v, want %q", "00000000:api", got, want)
}
_, ok = headers["00000001:api"]
if got, want := ok, false; got != want {
t.Error("headers[00000001:api]: got :api header, want no header for non-api requests")
}
}
func TestSendTimestampWithLogRequest(t *testing.T) {
req, err := http.NewRequest("POST", "http://example.com", nil)
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, remove, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
var b bytes.Buffer
s := NewStream(&b)
before := time.Now().UnixNano() / 1000 / 1000
s.LogRequest("Fake_Id0", req)
s.Close()
after := time.Now().UnixNano() / 1000 / 1000
headers := make(map[string]string)
reader := NewReader(&b)
for {
frame, err := reader.ReadFrame()
if frame == nil {
break
}
if err != nil && err != io.EOF {
t.Fatalf("reader.ReadFrame(): got %v, want no error or io.EOF", err)
}
headerFrame, ok := frame.(Header)
if !ok {
t.Fatalf("frame.(Header): couldn't convert frame '%v' to a headerFrame", frame)
}
headers[headerFrame.Name] = headerFrame.Value
}
timestr, ok := headers[":timestamp"]
if !ok {
t.Fatalf("headers[:timestamp]: got no such header, want :timestamp (headers were: %v)", headers)
}
ts, err := strconv.ParseInt(timestr, 10, 64)
if err != nil {
t.Fatalf("strconv.ParseInt: got %s, want no error. Invalidly formatted timestamp ('%s')", err, timestr)
}
if ts < before || ts > after {
t.Fatalf("headers[:timestamp]: got %d, want timestamp between %d and %d", ts, before, after)
}
}
func TestSendTimestampWithLogResponse(t *testing.T) {
req, err := http.NewRequest("POST", "http://example.com", nil)
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, remove, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
res := proxyutil.NewResponse(200, nil, req)
var b bytes.Buffer
s := NewStream(&b)
before := time.Now().UnixNano() / 1000 / 1000
s.LogResponse("Fake_Id1", res)
s.Close()
after := time.Now().UnixNano() / 1000 / 1000
headers := make(map[string]string)
reader := NewReader(&b)
for {
frame, err := reader.ReadFrame()
if frame == nil {
break
}
if err != nil && err != io.EOF {
t.Fatalf("reader.ReadFrame(): got %v, want no error or io.EOF", err)
}
headerFrame, ok := frame.(Header)
if !ok {
t.Fatalf("frame.(Header): couldn't convert frame '%v' to a headerFrame", frame)
}
headers[headerFrame.Name] = headerFrame.Value
}
timestr, ok := headers[":timestamp"]
if !ok {
t.Fatalf("headers[:timestamp]: got no such header, want :timestamp (headers were: %v)", headers)
}
ts, err := strconv.ParseInt(timestr, 10, 64)
if err != nil {
t.Fatalf("strconv.ParseInt: got %s, want no error. Invalidly formatted timestamp ('%s')", err, timestr)
}
if ts < before || ts > after {
t.Fatalf("headers[:timestamp]: got %d, want timestamp between %d and %d (headers were: %v)", ts, before, after, headers)
}
}
func TestBodyLoggingWithOneRead(t *testing.T) {
// Test scenario:
// 1. Prepare HTTP request with body containing a string.
// 2. Initialize marbl logging on this request.
// 3. Read body of the request in single Read() and verity that it matches
// original string.
// 4. Parse marbl data, extract DataFrames and verify that they match
// . original string.
body := "hello, world"
req, err := http.NewRequest("POST", "http://example.com", strings.NewReader(body))
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, remove, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
var b bytes.Buffer
s := NewStream(&b)
s.LogRequest("Fake_Id0", req)
// Read request body into big slice.
bodybytes := make([]byte, 100)
// First read. Due to implementation details of strings.Read
// it reads all bytes but doesn't return EOF.
n, err := req.Body.Read(bodybytes)
if n != len(body) {
t.Fatalf("req.Body.Read(): expected to read %v bytes but read %v", len(body), n)
}
if body != string(bodybytes[:n]) {
t.Fatalf("req.Body.Read(): expected to read %v but read %v", body, string(bodybytes[:n]))
}
if err != nil {
t.Fatalf("req.Body.Read(): first read expected to be successful but got error %v", err)
}
// second read. We already consumed the whole string on the first read
// so now it should be 0 bytes and EOF.
n, err = req.Body.Read(bodybytes)
if n != 0 {
t.Fatalf("req.Body.Read(): expected to read 0 bytes but read %v", n)
}
if err != io.EOF {
t.Fatalf("req.Body.Read(): expected EOF but got %v", err)
}
s.Close()
reader := NewReader(&b)
bodybytes = readAllDataFrames(reader, "Fake_Id0", t)
if len(bodybytes) != len(body) {
t.Fatalf("readAllDataFrames(): expected .marbl data to have %v bytes, but got %v", len(body), len(bodybytes))
}
if body != string(bodybytes) {
t.Fatalf("readAllDataFrames(): expected .marbl data to have string %v but got %v", body, string(bodybytes))
}
}
func TestBodyLogging_ManyReads(t *testing.T) {
// Test scenario:
// 1. Prepare HTTP request with body containing a string.
// 2. Initialize marbl logging on this request.
// 3. Read body of the request in many reads, 1 byte per read and
// . verify that it matches original string.
// 4. Parse marbl data, extract DataFrames and verify that they match
// . original string.
body := "hello, world"
req, err := http.NewRequest("POST", "http://example.com", strings.NewReader(body))
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, remove, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
var b bytes.Buffer
s := NewStream(&b)
s.LogRequest("Fake_Id0", req)
// Read request body into single byte slice.
bodybytes := make([]byte, 1)
for i := 0; i < len(body); i++ {
// first read
n, err := req.Body.Read(bodybytes)
if n != 1 {
t.Fatalf("req.Body.Read(): expected to read 1 byte but read %v", n)
}
if body[i] != bodybytes[0] {
t.Fatalf("req.Body.Read(): expected to read %v but read %v", body[i], bodybytes[0])
}
if err != nil {
t.Fatalf("req.Body.Read(): read expected to be successfully but got error %v", err)
}
}
// last read. We already consumed the whole string on the previous reads
// so now it should be 0 bytes and EOF.
n, err := req.Body.Read(bodybytes)
if n != 0 {
t.Fatalf("req.Body.Read(): expected to read 0 bytes but read %v", n)
}
if err != io.EOF {
t.Fatalf("req.Body.Read(): expected EOF but got %v", err)
}
s.Close()
reader := NewReader(&b)
bodybytes = readAllDataFrames(reader, "Fake_Id0", t)
if len(bodybytes) != len(body) {
t.Fatalf("readAllDataFrames(): expected .marbl data to have %v bytes, but got %v", len(body), len(bodybytes))
}
if body != string(bodybytes) {
t.Fatalf("readAllDataFrames(): expected .marbl data to have string %v but got %v", body, string(bodybytes))
}
}
func TestReturnOriginalRequestPathAndQuery(t *testing.T) {
req, err := http.NewRequest("GET", "http://example.com/foo%20bar?baz%20qux", nil)
if err != nil {
t.Fatalf("http.NewRequest(): got %v, want no error", err)
}
_, remove, err := martian.TestContext(req, nil, nil)
if err != nil {
t.Fatalf("TestContext(): got %v, want no error", err)
}
defer remove()
var b bytes.Buffer
s := NewStream(&b)
s.LogRequest("Fake_Id0", req)
s.Close()
headers := make(map[string]string)
reader := NewReader(&b)
for {
frame, err := reader.ReadFrame()
if frame == nil {
break
}
if err != nil && err != io.EOF {
t.Fatalf("reader.ReadFrame(): got %v, want no error or io.EOF", err)
}
headerFrame, ok := frame.(Header)
if !ok {
t.Fatalf("frame.(Header): couldn't convert frame '%v' to a headerFrame", frame)
}
headers[headerFrame.Name] = headerFrame.Value
}
path := headers[":path"]
if path != "/foo%20bar" {
t.Fatalf("headers[:path]: expected /foo%%20bar but got %s", path)
}
query := headers[":query"]
if query != "baz%20qux" {
t.Fatalf("headers[:query]: expected baz%%20qux but got %s", query)
}
}
// readAllDataFrames reads all DataFrames with reader, filters the one that match provided
// id and assembles data from all frames into single slice. It expects that
// there is only one slice of DataFrames with provided id.
func readAllDataFrames(reader *Reader, id string, t *testing.T) []byte {
res := make([]byte, 0)
term := false
var i uint32
for {
frame, _ := reader.ReadFrame()
if frame == nil {
break
}
if frame.FrameType() == DataFrame {
df := frame.(Data)
if df.ID != id {
continue
}
if term {
t.Fatal("DataFrame after terminal frame are not allowed.")
}
if df.Index != i {
t.Fatalf("expected DataFrame index %v but got %v", i, df.Index)
}
term = df.Terminal
res = append(res, df.Data...)
i++
}
}
if !term {
t.Fatal("didn't see terminal DataFrame")
}
return res
}