Golang proxyutil.NewResponse函数代码示例

func TestModifyResponseResetAuth(t *testing.T) {
	auth := NewModifier()

	auth.SetRequestModifier(martian.RequestModifierFunc(
		func(ctx *martian.Context, req *http.Request) error {
			if ctx.Auth.ID != "secret:pass" {
				ctx.Auth.Error = errors.New("invalid auth")
			}
			return nil
		}))

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}
	req.Header.Set("Proxy-Authorization", "Basic "+encode("wrong:pass"))

	ctx := martian.NewContext()
	// This will set ctx.Auth.Error since the ID isn't "secret:pass".
	if err := auth.ModifyRequest(ctx, req); err != nil {
		t.Fatalf("ModifyRequest(): got %v, want no error", err)
	}

	res := proxyutil.NewResponse(200, nil, req)
	if err := auth.ModifyResponse(ctx, res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if got, want := res.StatusCode, 407; got != want {
		t.Errorf("res.StatusCode: got %d, want %d", got, want)
	}

	if got, want := ctx.Auth.ID, ""; got != want {
		t.Errorf("ctx.Auth.ID: got %q, want %q", got, want)
	}
	if err := ctx.Auth.Error; err != nil {
		t.Errorf("ctx.Auth.Error: got %v, want no error", err)
	}

	// This will be successful because the ID is "secret:pass".
	req.Header.Set("Proxy-Authorization", "Basic "+encode("secret:pass"))
	if err := auth.ModifyRequest(ctx, req); err != nil {
		t.Fatalf("ModifyRequest(): got %v, want no error", err)
	}

	// Reset the response.
	res = proxyutil.NewResponse(200, nil, req)
	if err := auth.ModifyResponse(ctx, res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if got, want := res.StatusCode, 200; got != want {
		t.Errorf("res.StatusCode: got %d, want %d", got, want)
	}
}

func TestModifyResponse(t *testing.T) {
	m := NewModifier()
	ctx := martian.NewContext()
	res := proxyutil.NewResponse(200, nil, nil)

	if err := m.ModifyResponse(ctx, res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}

	m.SetResponseModifier(martian.ResponseModifierFunc(
		func(*martian.Context, *http.Response) error {
			ctx.Auth.Error = errors.New("auth is required")
			return nil
		}))

	if err := m.ModifyResponse(ctx, res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if got, want := res.StatusCode, http.StatusProxyAuthRequired; got != want {
		t.Errorf("res.StatusCode: got %d, want %d", got, want)
	}
	if got, want := res.Header.Get("Proxy-Authenticate"), "Basic"; got != want {
		t.Errorf("res.Header.Get(%q): got %q, want %q", "Proxy-Authenticate", got, want)
	}
}

func TestCopyModifier(t *testing.T) {
	m := NewCopyModifier("Original", "Copy")

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}
	req.Header.Set("Original", "test")

	if err := m.ModifyRequest(req); err != nil {
		t.Fatalf("ModifyRequest(): got %v, want no error", err)
	}

	if got, want := req.Header.Get("Copy"), "test"; got != want {
		t.Errorf("req.Header.Get(%q): got %q, want %q", "Copy", got, want)
	}

	res := proxyutil.NewResponse(200, nil, req)
	res.Header.Set("Original", "test")

	if err := m.ModifyResponse(res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}

	if got, want := res.Header.Get("Copy"), "test"; got != want {
		t.Errorf("res.Header.Get(%q): got %q, want %q", "Copy", got, want)
	}
}

func TestModifyResponse(t *testing.T) {
	p := NewProxy(mitm)

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(%q, ..., nil): got %v, want no error", "GET", err)
	}
	res := proxyutil.NewResponse(200, nil, req)

	if err := p.ModifyResponse(NewContext(), res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}

	modifierRun := false
	f := func(*Context, *http.Response) error {
		modifierRun = true
		return nil
	}
	p.SetResponseModifier(ResponseModifierFunc(f))

	if err := p.ModifyResponse(NewContext(), res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if !modifierRun {
		t.Error("modifierRun: got false, want true")
	}
}

func TestModifierFromJSON(t *testing.T) {
	msg := []byte(`{
    "cookie.Modifier": {
      "scope": ["request", "response"],
		  "name": "martian",
			"value": "value"
		}
	}`)

	r, err := parse.FromJSON(msg)
	if err != nil {
		t.Fatalf("parse.FromJSON(): got %v, want no error", err)
	}

	req, err := http.NewRequest("GET", "http://example.com/path/", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}

	reqmod := r.RequestModifier()

	if reqmod == nil {
		t.Fatal("reqmod: got nil, want not nil")
	}

	if err := reqmod.ModifyRequest(martian.NewContext(), req); err != nil {
		t.Fatalf("reqmod.ModifyRequest(): got %v, want no error", err)
	}

	if got, want := len(req.Cookies()), 1; got != want {
		t.Fatalf("len(req.Cookies): got %v, want %v", got, want)
	}
	if got, want := req.Cookies()[0].Name, "martian"; got != want {
		t.Errorf("req.Cookies()[0].Name: got %v, want %v", got, want)
	}
	if got, want := req.Cookies()[0].Value, "value"; got != want {
		t.Errorf("req.Cookies()[0].Value: got %v, want %v", got, want)
	}

	resmod := r.ResponseModifier()

	if resmod == nil {
		t.Fatal("resmod: got nil, want not nil")
	}

	res := proxyutil.NewResponse(200, nil, req)
	if err := resmod.ModifyResponse(martian.NewContext(), res); err != nil {
		t.Fatalf("resmod.ModifyResponse(): got %v, want no error", err)
	}

	if got, want := len(res.Cookies()), 1; got != want {
		t.Fatalf("len(res.Cookies): got %v, want %v", got, want)
	}
	if got, want := res.Cookies()[0].Name, "martian"; got != want {
		t.Errorf("res.Cookies()[0].Name: got %v, want %v", got, want)
	}
	if got, want := res.Cookies()[0].Value, "value"; got != want {
		t.Errorf("res.Cookies()[0].Value: got %v, want %v", got, want)
	}
}

func TestFromJSON(t *testing.T) {
	msg := []byte(`{
    "status.Modifier": {
      "scope": ["response"],
      "statusCode": 400
    }
  }`)

	r, err := parse.FromJSON(msg)
	if err != nil {
		t.Fatalf("parse.FromJSON(): got %v, want no error", err)
	}
	resmod := r.ResponseModifier()
	if resmod == nil {
		t.Fatal("resmod: got nil, want not nil")
	}

	res := proxyutil.NewResponse(200, nil, nil)
	if err := resmod.ModifyResponse(res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if got, want := res.StatusCode, 400; got != want {
		t.Errorf("res.StatusCode: got %d, want %d", got, want)
	}
}

func TestModifierFuncs(t *testing.T) {
	reqmod := RequestModifierFunc(
		func(req *http.Request) error {
			req.Header.Set("Request-Modified", "true")
			return nil
		})

	resmod := ResponseModifierFunc(
		func(res *http.Response) error {
			res.Header.Set("Response-Modified", "true")
			return nil
		})

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}

	if err := reqmod.ModifyRequest(req); err != nil {
		t.Fatalf("reqmod.ModifyRequest(): got %v, want no error", err)
	}
	if got, want := req.Header.Get("Request-Modified"), "true"; got != want {
		t.Errorf("req.Header.Get(%q): got %q, want %q", "Request-Modified", got, want)
	}

	res := proxyutil.NewResponse(200, nil, req)

	if err := resmod.ModifyResponse(res); err != nil {
		t.Fatalf("resmod.ModifyResponse(): got %v, want no error", err)
	}
	if got, want := res.Header.Get("Response-Modified"), "true"; got != want {
		t.Errorf("res.Header.Get(%q): got %q, want %q", "Response-Modified", got, want)
	}
}

func TestStatusModifierOnResponse(t *testing.T) {
	for i, status := range []int{
		http.StatusForbidden,
		http.StatusOK,
		http.StatusTemporaryRedirect,
	} {
		req, err := http.NewRequest("GET", "/", nil)
		if err != nil {
			t.Fatalf("NewRequest(): got %v, want no error", err)
		}

		res := proxyutil.NewResponse(200, nil, req)

		mod := NewModifier(status)

		if err := mod.ModifyResponse(res); err != nil {
			t.Fatalf("%d. ModifyResponse(): got %v, want no error", i, err)
		}

		if got, want := res.StatusCode, status; got != want {
			t.Errorf("%d. res.StatusCode: got %v, want %v", i, got, want)
		}
		if got, want := res.Status, fmt.Sprintf("%d %s", res.StatusCode, http.StatusText(status)); got != want {
			t.Errorf("%d. res.Status: got %q, want %q", i, got, want)
		}
	}
}

func TestExportIgnoresOrphanedResponse(t *testing.T) {
	logger := NewLogger("martian", "2.0.0")

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}

	_, remove, err := martian.TestContext(req)
	if err != nil {
		t.Fatalf("martian.TestContext(): got %v, want no error", err)
	}
	defer remove()

	if err := logger.ModifyRequest(req); err != nil {
		t.Fatalf("ModifyRequest(): got %v, want no error", err)
	}

	// Reset before the response comes back.
	logger.Reset()

	res := proxyutil.NewResponse(200, nil, req)
	if err := logger.ModifyResponse(res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}

	log := logger.Export().Log
	if got, want := len(log.Entries), 0; got != want {
		t.Errorf("len(log.Entries): got %d, want %d", got, want)
	}
}

func TestVerifierFromJSON(t *testing.T) {
	msg := []byte(`{
    "status.Verifier": {
      "scope": ["response"],
      "statusCode": 400
    }
  }`)

	r, err := parse.FromJSON(msg)
	if err != nil {
		t.Fatalf("parse.FromJSON(): got %v, want no error", err)
	}
	resmod := r.ResponseModifier()
	if resmod == nil {
		t.Fatal("resmod: got nil, want not nil")
	}
	resv, ok := resmod.(verify.ResponseVerifier)
	if !ok {
		t.Fatal("reqmod.(verify.RequestVerifier): got !ok, want ok")
	}

	req, err := http.NewRequest("GET", "http://www.example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}

	res := proxyutil.NewResponse(200, nil, req)
	if err := resv.ModifyResponse(martian.NewContext(), res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}
	if err := resv.VerifyResponses(); err == nil {
		t.Error("VerifyResponses(): got nil, want not nil")
	}
}

func (p *Proxy) connect(req *http.Request) (*http.Response, net.Conn, error) {
	if p.proxyURL != nil {
		log.Debugf("martian: CONNECT with downstream proxy: %s", p.proxyURL.Host)

		conn, err := net.Dial("tcp", p.proxyURL.Host)
		if err != nil {
			return nil, nil, err
		}
		pbw := bufio.NewWriter(conn)
		pbr := bufio.NewReader(conn)

		req.Write(pbw)
		pbw.Flush()

		res, err := http.ReadResponse(pbr, req)
		if err != nil {
			return nil, nil, err
		}

		return res, conn, nil
	}

	log.Debugf("martian: CONNECT to host directly: %s", req.URL.Host)

	conn, err := net.Dial("tcp", req.URL.Host)
	if err != nil {
		return nil, nil, err
	}

	return proxyutil.NewResponse(200, nil, req), conn, nil
}

func TestBodyModifier(t *testing.T) {
	mod := NewModifier([]byte("text"), "text/plain")

	req, err := http.NewRequest("GET", "/", strings.NewReader(""))
	if err != nil {
		t.Fatalf("NewRequest(): got %v, want no error", err)
	}
	req.Header.Set("Content-Encoding", "gzip")

	if err := mod.ModifyRequest(req); err != nil {
		t.Fatalf("ModifyRequest(): got %v, want no error", err)
	}

	if got, want := req.Header.Get("Content-Type"), "text/plain"; got != want {
		t.Errorf("req.Header.Get(%q): got %v, want %v", "Content-Type", got, want)
	}
	if got, want := req.ContentLength, int64(len([]byte("text"))); got != want {
		t.Errorf("req.ContentLength: got %d, want %d", got, want)
	}
	if got, want := req.Header.Get("Content-Encoding"), ""; got != want {
		t.Errorf("req.Header.Get(%q): got %q, want %q", "Content-Encoding", got, want)
	}

	got, err := ioutil.ReadAll(req.Body)
	if err != nil {
		t.Fatalf("ioutil.ReadAll(): got %v, want no error", err)
	}
	req.Body.Close()

	if want := []byte("text"); !bytes.Equal(got, want) {
		t.Errorf("res.Body: got %q, want %q", got, want)
	}

	res := proxyutil.NewResponse(200, nil, req)
	res.Header.Set("Content-Encoding", "gzip")

	if err := mod.ModifyResponse(res); err != nil {
		t.Fatalf("ModifyResponse(): got %v, want no error", err)
	}

	if got, want := res.Header.Get("Content-Type"), "text/plain"; got != want {
		t.Errorf("res.Header.Get(%q): got %v, want %v", "Content-Type", got, want)
	}
	if got, want := res.ContentLength, int64(len([]byte("text"))); got != want {
		t.Errorf("res.ContentLength: got %d, want %d", got, want)
	}
	if got, want := res.Header.Get("Content-Encoding"), ""; got != want {
		t.Errorf("res.Header.Get(%q): got %q, want %q", "Content-Encoding", got, want)
	}

	got, err = ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatalf("ioutil.ReadAll(): got %v, want no error", err)
	}
	res.Body.Close()

	if want := []byte("text"); !bytes.Equal(got, want) {
		t.Errorf("res.Body: got %q, want %q", got, want)
	}
}

func TestNoModifiers(t *testing.T) {
	m := NewModifier()
	m.SetRequestModifier(nil)
	m.SetResponseModifier(nil)

	req, err := http.NewRequest("GET", "http://example.com", nil)
	if err != nil {
		t.Fatalf("http.NewRequest(): got %v, want no error", err)
	}

	ctx, err := session.FromContext(nil)
	if err != nil {
		t.Fatalf("session.FromContext(): got %v, want no error", err)
	}

	martian.SetContext(req, ctx)
	defer martian.RemoveContext(req)

	if err := m.ModifyRequest(req); err != nil {
		t.Errorf("ModifyRequest(): got %v, want no error", err)
	}

	res := proxyutil.NewResponse(200, nil, req)
	if err := m.ModifyResponse(res); err != nil {
		t.Errorf("ModifyResponse(): got %v, want no error", err)
	}
}

func ExampleLogger() {
	l := NewLogger()
	l.SetLogFunc(func(line string) {
		// Remove \r to make it easier to test with examples.
		fmt.Print(strings.Replace(line, "\r", "", -1))
	})
	l.SetDecode(true)

	buf := new(bytes.Buffer)
	gw := gzip.NewWriter(buf)
	gw.Write([]byte("request content"))
	gw.Close()

	req, err := http.NewRequest("GET", "http://example.com/path?querystring", buf)
	if err != nil {
		fmt.Println(err)
		return
	}
	req.TransferEncoding = []string{"chunked"}
	req.Header.Set("Content-Encoding", "gzip")

	if err := l.ModifyRequest(req); err != nil {
		fmt.Println(err)
		return
	}

	res := proxyutil.NewResponse(200, strings.NewReader("response content"), req)
	res.ContentLength = 16
	res.Header.Set("Date", "Tue, 15 Nov 1994 08:12:31 GMT")
	res.Header.Set("Other-Header", "values")

	if err := l.ModifyResponse(res); err != nil {
		fmt.Println(err)
		return
	}
	// Output:
	// --------------------------------------------------------------------------------
	// Request to http://example.com/path?querystring
	// --------------------------------------------------------------------------------
	// GET http://example.com/path?querystring HTTP/1.1
	// Host: example.com
	// Transfer-Encoding: chunked
	// Content-Encoding: gzip
	//
	// request content
	//
	// --------------------------------------------------------------------------------
	//
	// --------------------------------------------------------------------------------
	// Response from http://example.com/path?querystring
	// --------------------------------------------------------------------------------
	// HTTP/1.1 200 OK
	// Content-Length: 16
	// Date: Tue, 15 Nov 1994 08:12:31 GMT
	// Other-Header: values
	//
	// response content
	// --------------------------------------------------------------------------------
}

func TestModifyResponseNoModifier(t *testing.T) {
	m := NewModifier()
	res := proxyutil.NewResponse(200, nil, nil)

	if err := m.ModifyResponse(martian.NewContext(), res); err != nil {
		t.Errorf("ModifyResponse(): got %v, want no error", err)
	}
}