1// Copyright 2015 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package proxyauth
16
17import (
18	"encoding/base64"
19	"errors"
20	"net/http"
21	"testing"
22
23	"github.com/google/martian/v3"
24	"github.com/google/martian/v3/auth"
25	"github.com/google/martian/v3/martiantest"
26	"github.com/google/martian/v3/proxyutil"
27)
28
29func encode(v string) string {
30	return base64.StdEncoding.EncodeToString([]byte(v))
31}
32
33func TestNoModifiers(t *testing.T) {
34	m := NewModifier()
35	m.SetRequestModifier(nil)
36	m.SetResponseModifier(nil)
37
38	req, err := http.NewRequest("GET", "http://example.com", nil)
39	if err != nil {
40		t.Fatalf("http.NewRequest(): got %v, want no error", err)
41	}
42
43	_, remove, err := martian.TestContext(req, nil, nil)
44	if err != nil {
45		t.Fatalf("martian.TestContext(): got %v, want no error", err)
46	}
47	defer remove()
48
49	if err := m.ModifyRequest(req); err != nil {
50		t.Errorf("ModifyRequest(): got %v, want no error", err)
51	}
52
53	res := proxyutil.NewResponse(200, nil, req)
54	if err := m.ModifyResponse(res); err != nil {
55		t.Errorf("ModifyResponse(): got %v, want no error", err)
56	}
57}
58
59func TestProxyAuth(t *testing.T) {
60	m := NewModifier()
61	tm := martiantest.NewModifier()
62	m.SetRequestModifier(tm)
63	m.SetResponseModifier(tm)
64
65	req, err := http.NewRequest("GET", "http://example.com", nil)
66	if err != nil {
67		t.Fatalf("http.NewRequest(): got %v, want no error", err)
68	}
69	req.Header.Set("Proxy-Authorization", "Basic "+encode("user:pass"))
70
71	ctx, remove, err := martian.TestContext(req, nil, nil)
72	if err != nil {
73		t.Fatalf("martian.TestContext(): got %v, want no error", err)
74	}
75	defer remove()
76
77	if err := m.ModifyRequest(req); err != nil {
78		t.Fatalf("ModifyRequest(): got %v, want no error", err)
79	}
80
81	actx := auth.FromContext(ctx)
82	if got, want := actx.ID(), "user:pass"; got != want {
83		t.Fatalf("actx.ID(): got %q, want %q", got, want)
84	}
85
86	if err := m.ModifyRequest(req); err != nil {
87		t.Fatalf("ModifyRequest(): got %v, want no error", err)
88	}
89	if !tm.RequestModified() {
90		t.Error("tm.RequestModified(): got false, want true")
91	}
92
93	res := proxyutil.NewResponse(200, nil, req)
94
95	if err := m.ModifyResponse(res); err != nil {
96		t.Fatalf("ModifyResponse(): got %v, want no error", err)
97	}
98
99	if err := m.ModifyResponse(res); err != nil {
100		t.Fatalf("ModifyResponse(): got %v, want no error", err)
101	}
102
103	if !tm.ResponseModified() {
104		t.Error("tm.ResponseModified(): got false, want true")
105	}
106
107	if got, want := res.StatusCode, 200; got != want {
108		t.Errorf("res.StatusCode: got %d, want %d", got, want)
109	}
110	if got, want := res.Header.Get("Proxy-Authenticate"), ""; got != want {
111		t.Errorf("res.Header.Get(%q): got %q, want %q", "Proxy-Authenticate", got, want)
112	}
113}
114
115func TestProxyAuthInvalidCredentials(t *testing.T) {
116	m := NewModifier()
117	authErr := errors.New("auth error")
118
119	tm := martiantest.NewModifier()
120	tm.RequestFunc(func(req *http.Request) {
121		ctx := martian.NewContext(req)
122		actx := auth.FromContext(ctx)
123
124		actx.SetError(authErr)
125	})
126	tm.ResponseFunc(func(res *http.Response) {
127		ctx := martian.NewContext(res.Request)
128		actx := auth.FromContext(ctx)
129
130		actx.SetError(authErr)
131	})
132
133	m.SetRequestModifier(tm)
134	m.SetResponseModifier(tm)
135
136	req, err := http.NewRequest("GET", "http://example.com", nil)
137	if err != nil {
138		t.Fatalf("http.NewRequest(): got %v, want no error", err)
139	}
140	req.Header.Set("Proxy-Authorization", "Basic "+encode("user:pass"))
141
142	ctx, remove, err := martian.TestContext(req, nil, nil)
143	if err != nil {
144		t.Fatalf("martian.TestContext(): got %v, want no error", err)
145	}
146	defer remove()
147
148	if err := m.ModifyRequest(req); err != nil {
149		t.Fatalf("ModifyRequest(): got %v, want no error", err)
150	}
151
152	if !tm.RequestModified() {
153		t.Error("tm.RequestModified(): got false, want true")
154	}
155
156	actx := auth.FromContext(ctx)
157	if actx.Error() != authErr {
158		t.Fatalf("auth.Error(): got %v, want %v", actx.Error(), authErr)
159	}
160	actx.SetError(nil)
161
162	res := proxyutil.NewResponse(200, nil, req)
163
164	if err := m.ModifyResponse(res); err != nil {
165		t.Fatalf("ModifyResponse(): got %v, want no error", err)
166	}
167
168	if !tm.ResponseModified() {
169		t.Error("tm.ResponseModified(): got false, want true")
170	}
171
172	if actx.Error() != authErr {
173		t.Fatalf("auth.Error(): got %v, want %v", actx.Error(), authErr)
174	}
175
176	if got, want := res.StatusCode, http.StatusProxyAuthRequired; got != want {
177		t.Errorf("res.StatusCode: got %d, want %d", got, want)
178	}
179	if got, want := res.Header.Get("Proxy-Authenticate"), "Basic"; got != want {
180		t.Errorf("res.Header.Get(%q): got %q, want %q", "Proxy-Authenticate", got, want)
181	}
182}
183