|
- // 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
- }
|