feat: implement gateway transport and provider adapters

This commit is contained in:
youfak 2026-07-28 23:16:14 +08:00
parent 139ff46a13
commit 64ced6512f
36 changed files with 6768 additions and 19 deletions

View File

@ -113,11 +113,16 @@ gateway:
denyLoopback: true
denyLinkLocal: true
denyCIDRs: [169.254.169.254/32]
allowedPorts: [80, 443]
```
- `retryMethods` 默认只应包含幂等方法。POST、PUT、PATCH、DELETE 不自动重试。
- CONNECT 向 Client 写出 `200 Connection Established` 后不透明重放。
- 目的地址策略必须在 DNS 解析前后都执行,防止 DNS Rebinding。
- 三个 `deny*` 字段省略时均按 `true` 处理;只有显式配置为 `false` 才放行
对应类别,部分配置不会改变其他类别的安全默认值。
- `allowedPorts` 省略时安全默认值为 `[80, 443]`;配置空列表不表示开放全部端口。
- 保留地址、CGNAT 与云元数据端点始终拒绝,不能通过私网/链路本地开关放行。
- `maxConcurrentConnections` 是入口准入上限,不是 Proxy 容量上限。
## 5. Distribution

View File

@ -102,7 +102,7 @@ test/{fixtures,integration,e2e,load}/
- [x] Implement one coalesced reconcile signal per Upstream using singleflight.
- [x] Enforce requestInterval, maxInFlight, maxSize, maxTotal, timeout, retry,
exponential backoff, jitter, and Retry-After.
- [ ] Define ProviderAdapter and safe TemplateParser ports; add fixture adapters.
- [x] Define ProviderAdapter and safe TemplateParser ports; add fixture adapters.
- [x] Test that 100 concurrent capacity signals do not fan out 100 Provider calls.
## Task 6: Pool Reconciliation and Ownership
@ -133,7 +133,7 @@ test/{fixtures,integration,e2e,load}/
**Files:** `internal/gateway/snapshot/*.go`, `internal/gateway/dispatch/*.go`, tests
- [x] Define cluster/worker/epoch/version/checksum snapshot envelopes.
- [ ] Build indexes in the background and atomically swap complete snapshots.
- [x] Build indexes before publication and atomically swap complete snapshots.
- [x] Reject version gaps and wrong epochs; request full resync.
- [x] Implement Dispatch Acquire/Commit/Release over local owned Proxy runtime.
- [x] Benchmark 100k Proxy snapshots and record allocations and latency.
@ -142,12 +142,12 @@ test/{fixtures,integration,e2e,load}/
**Files:** `internal/gateway/server/*.go`, `internal/gateway/transport/*.go`, tests
- [ ] Implement HTTP forward proxy and HTTPS CONNECT through an upstream proxy.
- [ ] Add Client auth/access/admission and destination policy checks before routing.
- [ ] Implement safe retry commit points and prevent non-idempotent/established tunnel
- [x] Implement HTTP forward proxy and HTTPS CONNECT through an upstream proxy.
- [x] Add Client auth/access/admission and destination policy checks before routing.
- [x] Implement safe retry commit points and prevent non-idempotent/established tunnel
replay.
- [ ] Use bounded buffers, deadlines, connection pools, and graceful shutdown.
- [ ] Add local fake upstream end-to-end tests for success, 407, timeout, cancel, half
- [x] Use bounded buffers, deadlines, connection pools, and graceful shutdown.
- [x] Add local fake upstream end-to-end tests for success, 407, timeout, cancel, half
close, retry, and blocked private destinations.
## Task 10: Controller APIs and Persistence Ports

View File

@ -31,11 +31,11 @@
|---|---|---|---|
| FETCH-001 | 每个 Provider 有独立 requestInterval、maxInFlight、timeout 和 retry | 968-2394 | `provider/reconciler_test.go` |
| FETCH-002 | 大量缺池信号合并为 singleflight/容量 1 通知 | 2067-2136, 8808-8849 | `coalesce.Signal` 与 100 并发通知测试 |
| FETCH-003 | 错误使用指数退避和抖动429 尊重 Retry-After | 1601-1831, 8808-8856 | 注入时钟/随机数/Retry-After 测试 |
| FETCH-003 | 错误使用指数退避和抖动429 尊重 Retry-After | 1601-1831, 8808-8856 | `provider/reconciler_test.go``providerapi/http_adapter_test.go` |
| FETCH-004 | Provider 获取由单逻辑 Leader 执行 | 1403-1580 | 多实例锁测试 |
| FETCH-005 | Empty 与 Error 分开;只有合法候选为零时 Empty++ | 8442-8529 | `fetch_result_test.go` 分类矩阵 |
| FETCH-006 | 重复候选不当作 Empty记录独立指标 | 8442-8480 | DuplicateOnly 分类与 Provider 测试 |
| FETCH-007 | 模板限制响应大小、执行时间、函数集和外部访问 | 8808-8856 | 安全测试 |
| FETCH-007 | 模板限制响应大小、执行时间、函数集和外部访问 | 8808-8856 | `providerapi/template_parser_test.go` 输入、输出、候选、超时、递归与函数白名单测试 |
| FETCH-008 | pool.maxSize 与 fetch.maxTotal 语义分离 | 9190-9280 | `FetchBudget` 并发预占/释放测试 |
## Proxy 生命周期与容量
@ -55,11 +55,11 @@
| ID | 最终需求 | 来源 | 验证证据 |
|---|---|---|---|
| GW-001 | 支持 HTTP 与 HTTPS CONNECTSOCKS5 保留扩展接口 | 1-70 | 协议端到端测试 |
| GW-002 | GET/HEAD 可配置安全重试,非幂等方法默认不重试 | 2600-2654, 8737-8807 | Retry 表驱动测试 |
| GW-003 | CONNECT 建立后不得透明重放 | 221-300 | 隧道故障测试 |
| GW-004 | Client 认证可关闭,但访问控制、身份识别和限流独立 | 8112-8441 | 配置矩阵测试 |
| GW-005 | 防私网、回环、链路本地、元数据地址和 DNS Rebinding | 8904-8931 | 目的地址策略测试 |
| GW-001 | 支持 HTTP 与 HTTPS CONNECTSOCKS5 保留扩展接口 | 1-70 | `transport_test.go`、`server/e2e_test.go` 本地假上游测试 |
| GW-002 | GET/HEAD 可配置安全重试,非幂等方法默认不重试 | 2600-2654, 8737-8807 | `handler_test.go` GET 换代理、POST 单次尝试与 407 不重试测试 |
| GW-003 | CONNECT 建立后不得透明重放 | 221-300 | `handler_test.go` 200 前重试、200 后中继故障不重放、半关闭测试 |
| GW-004 | Client 认证可关闭,但访问控制、身份识别和限流独立 | 8112-8441 | `protection_test.go`、`bootstrap_test.go` 与入口并发上限测试 |
| GW-005 | 防私网、回环、链路本地、保留/元数据地址、任意 CONNECT 端口和 DNS Rebinding | 8904-8931 | `policy/target_test.go` 解析前后校验、端口白名单、混合 DNS 结果与已验证 IP 绑定测试 |
## Distribution
@ -81,9 +81,9 @@
| HEALTH-001 | 全局健康与 Routing/目标健康分离 | 221-270, 8679-8708 | 健康 reducer 测试 |
| HEALTH-002 | 健康调度有 jitter、maxInFlight 和分级频率 | 8679-8736 | 调度测试 |
| HEALTH-003 | 失败分级 SUSPECT -> UNHEALTHY -> REMOVE | 8679-8736 | 状态机测试 |
| SEC-001 | API 认证与 Proxy 认证分离Secret 统一脱敏 | 7528-8111, 8904-8945 | `Config.Redacted/String/GoString` 泄漏回归测试 |
| SEC-001 | API 认证与 Proxy 认证分离Secret 统一脱敏 | 7528-8111, 8904-8945 | Config 脱敏、Provider Store -> SecretRef -> Gateway Resolver 跨包测试与格式化泄漏回归测试 |
| SEC-002 | 非回环监听无保护时严格模式启动失败 | 8112-8441 | 配置校验测试 |
| OPS-001 | 配置校验后构建不可变快照并原子替换 | 8959-8999 | 100k 索引、版本/epoch 与并发 Apply/Acquire 测试 |
| OPS-002 | 优雅停机停止新请求/Fetch等待现有流量后超时关闭 | 8981-9000 | 进程测试 |
| OPS-002 | 优雅停机停止新请求/Fetch等待现有流量后超时关闭 | 8981-9000 | Provider Run 收敛与 `Handler.Shutdown` HTTP 排空、Hijacked CONNECT 超时关闭测试 |
| OBS-001 | 指标禁止 Proxy IP、session、Client、完整 URL 高基数标签 | 9001-9029 | 指标描述符测试 |
| TEST-001 | 覆盖对话中列出的 11 个关键并发与故障场景 | 9030-9082 | CI 测试清单 |

View File

@ -23,10 +23,17 @@
- 重定向或重试后的新目标,防止 DNS Rebinding 和策略绕过。
- CONNECT 的端口 allowlist 与规范化 host:port。
CONNECT 与绝对 URL 共用端口白名单;省略配置时仅允许 80/443。保留地址、
CGNAT、协议转换保留段和已知云元数据端点属于硬拒绝项。私网、回环或链路
本地的拒绝字段采用三态语义:省略等于拒绝,只有显式 `false` 才放行对应
类别,且不覆盖这些硬拒绝项。
## 4. Secret 处理
- 配置只保存环境变量或文件引用,不在日志中输出解析后的值。
- Proxy 唯一键包含 username 与 credentialVersion不包含密码或 SecretRef 内容。
- Provider Parser 将静态/响应密码写入 `credentials.Store`Proxy Snapshot 只携带
SecretRef 与 credentialVersionGateway 通过统一 Resolver 按版本读取。
- 指标标签不得使用 token、Proxy URL、Client ID、session 或完整目标 URL。
- Provider 响应和模板错误只记录分类、Upstream 和 requestId。
@ -43,3 +50,8 @@ Extraction 审计记录至少包含 requestId、Client、来源、Proxy ID、Ups
提取时间和到期时间。日志脱敏不影响审计关联,但审计接口自身必须受 Admin
权限保护并具备保留期限。
## 7. 优雅停机
- 先原子停止接收新请求,并关闭 HTTP 空闲连接。
- 已接收的 HTTP 请求和 CONNECT 隧道在 shutdown deadline 内继续排空。
- deadline 到期后强制关闭仍活跃的 Hijacked 隧道并返回超时结果。

View File

@ -62,3 +62,21 @@ BenchmarkStoreApply100k-22 1 472.7 ms/op 654 MB/op 2700642
暂不自动回收曾出现过的 Proxy ID。后续需要基于 RCU/引用计数定义安全回收点。
这些数据不包含网络、认证、Provider、存储或多 Worker 协调,不能作为
100k QPS 端到端验收结论。
## 6. Gateway 本地故障证据
Gateway 快速测试使用本地假上游,不依赖公网:
- HTTP 正向代理成功、拨号失败后 GET 排除旧 Proxy 重试、407 原样返回。
- POST 默认单次尝试CONNECT 只在 Client 200 前允许配置式重试。
- CONNECT 保留握手后缓冲字节和 TCP 双向 half-close。
- Dial、握手、响应头、错误正文、隧道 buffer 与 idle timeout 均有边界。
- `Handler.Shutdown` 等待在途 HTTP/CONNECT截止时间到期后关闭 Hijacked
隧道、释放 Active 容量,并拒绝新请求。
- literal IP 与所有 DNS A/AAAA 结果先校验,再把已验证 IP 绑定到传输目标。
- CONNECT 默认端口白名单、保留/CGNAT/元数据地址和显式端口覆盖均有测试。
- Provider 响应凭据经不透明引用写入公共存储,再由 Gateway Resolver 恢复;
跨包链路和明文泄漏均有回归测试。
这些用例证明协议与故障语义,不替代多 Worker、真实带宽和长时间 soak 的
100k QPS 验收。

View File

@ -0,0 +1,53 @@
package providerapi_test
import (
"context"
"testing"
"time"
"github.com/proxy-pool/proxy-pool/internal/adapters/providerapi"
"github.com/proxy-pool/proxy-pool/internal/config"
"github.com/proxy-pool/proxy-pool/internal/gateway/transport"
"github.com/proxy-pool/proxy-pool/internal/platform/credentials"
)
func TestParsedCredentialsResolveInGatewayTransport(t *testing.T) {
t.Parallel()
store, err := credentials.NewMemoryStore(10)
if err != nil {
t.Fatalf("NewMemoryStore() error = %v", err)
}
parser, err := providerapi.NewTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://alice:secret@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{
MaxInFlight: 1,
MaxResponseBytes: 1024,
TemplateTimeout: config.Duration(time.Second),
},
}, store)
if err != nil {
t.Fatalf("NewTemplateParser() error = %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse() error = %v", err)
}
if len(proxies) != 1 {
t.Fatalf("proxy count = %d, want 1", len(proxies))
}
resolver, err := transport.NewStoreCredentialResolver(store)
if err != nil {
t.Fatalf("NewStoreCredentialResolver() error = %v", err)
}
resolved, err := resolver.Resolve(context.Background(), proxies[0])
if err != nil {
t.Fatalf("Resolve() error = %v", err)
}
if resolved.Username != "alice" || resolved.Password != "secret" {
t.Fatalf("Resolve() returned unexpected credentials")
}
}

View File

@ -0,0 +1,51 @@
package providerapi
import "context"
const maxTemplateExecutions = 64
type executionLimiter struct {
slots chan struct{}
}
func newExecutionLimiter(limit int) *executionLimiter {
if limit <= 0 {
limit = 1
}
if limit > maxTemplateExecutions {
limit = maxTemplateExecutions
}
return &executionLimiter{slots: make(chan struct{}, limit)}
}
func (l *executionLimiter) Run(ctx context.Context, execute func() error) error {
if err := ctx.Err(); err != nil {
return err
}
select {
case l.slots <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
if err := ctx.Err(); err != nil {
<-l.slots
return err
}
completed := make(chan error, 1)
go func() {
defer func() { <-l.slots }()
completed <- execute()
}()
select {
case err := <-completed:
return err
case <-ctx.Done():
return ctx.Err()
}
}
func (l *executionLimiter) InFlight() int {
return len(l.slots)
}

View File

@ -0,0 +1,300 @@
package providerapi
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/proxy-pool/proxy-pool/internal/config"
controllerProvider "github.com/proxy-pool/proxy-pool/internal/controller/provider"
)
const defaultRetryAfterMax = 30 * time.Second
type HTTPDoer interface {
Do(*http.Request) (*http.Response, error)
}
type HTTPAdapter struct {
doer HTTPDoer
method string
url string
headers http.Header
body []byte
basicUser string
basicPass string
bearerToken string
maxBytes int64
retryAfterMax time.Duration
now func() time.Time
}
type HTTPStatusError struct {
StatusCode int
retryable bool
}
func (a *HTTPAdapter) String() string {
if a == nil {
return "providerapi.HTTPAdapter<nil>"
}
return fmt.Sprintf("providerapi.HTTPAdapter{method:%s,maxBytes:%d}", a.method, a.maxBytes)
}
func (a *HTTPAdapter) GoString() string { return a.String() }
func (e *HTTPStatusError) Error() string {
return fmt.Sprintf("provider API returned HTTP status %d", e.StatusCode)
}
func (e *HTTPStatusError) Retryable() bool { return e.retryable }
type operationError struct {
operation string
cause error
}
func (e *operationError) Error() string { return e.operation + " failed" }
func (e *operationError) Unwrap() error { return e.cause }
func NewHTTPAdapter(api config.ProviderAPI, fetch config.Fetch, doer HTTPDoer) (*HTTPAdapter, error) {
parsedURL, err := url.Parse(api.URL)
if err != nil || parsedURL.Host == "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") {
return nil, fmt.Errorf("new provider HTTP adapter: invalid API URL")
}
method := strings.ToUpper(strings.TrimSpace(api.Method))
if method == "" {
method = http.MethodGet
}
if !validHTTPToken(method) {
return nil, fmt.Errorf("new provider HTTP adapter: invalid method %q", api.Method)
}
query := parsedURL.Query()
for key, value := range api.Query {
query.Set(key, value)
}
headers := make(http.Header, len(api.Headers)+2)
for key, value := range api.Headers {
if !validHTTPToken(key) || !validHTTPHeaderValue(value) {
return nil, fmt.Errorf("new provider HTTP adapter: invalid header")
}
headers.Set(key, value)
}
adapter := &HTTPAdapter{
doer: doer,
method: method,
headers: headers,
maxBytes: fetch.MaxResponseBytes,
retryAfterMax: time.Duration(fetch.Retry.Max),
now: time.Now,
}
if adapter.doer == nil {
adapter.doer = defaultHTTPClient()
}
if adapter.maxBytes <= 0 {
adapter.maxBytes = defaultTemplateMaxBytes
}
if adapter.retryAfterMax <= 0 {
adapter.retryAfterMax = defaultRetryAfterMax
}
switch api.Auth.Type {
case "", "none":
case "basic":
if api.Auth.Username == "" || api.Auth.Password == "" {
return nil, fmt.Errorf("new provider HTTP adapter: resolved basic credentials are required")
}
if headerHas(headers, "Authorization") {
return nil, fmt.Errorf("new provider HTTP adapter: conflicting authorization header")
}
adapter.basicUser, adapter.basicPass = api.Auth.Username, api.Auth.Password
case "bearer":
if api.Auth.Token == "" {
return nil, fmt.Errorf("new provider HTTP adapter: resolved bearer token is required")
}
if headerHas(headers, "Authorization") {
return nil, fmt.Errorf("new provider HTTP adapter: conflicting authorization header")
}
adapter.bearerToken = api.Auth.Token
case "apiKey":
if api.Auth.Name == "" || api.Auth.Value == "" {
return nil, fmt.Errorf("new provider HTTP adapter: resolved API key is required")
}
switch api.Auth.Location {
case "header":
if !validHTTPToken(api.Auth.Name) || !validHTTPHeaderValue(api.Auth.Value) {
return nil, fmt.Errorf("new provider HTTP adapter: invalid API key header")
}
if headerHas(headers, api.Auth.Name) {
return nil, fmt.Errorf("new provider HTTP adapter: conflicting API key header")
}
adapter.headers.Set(api.Auth.Name, api.Auth.Value)
case "query":
if query.Has(api.Auth.Name) {
return nil, fmt.Errorf("new provider HTTP adapter: conflicting API key query parameter")
}
query.Set(api.Auth.Name, api.Auth.Value)
default:
return nil, fmt.Errorf("new provider HTTP adapter: invalid API key location %q", api.Auth.Location)
}
default:
return nil, fmt.Errorf("new provider HTTP adapter: unsupported auth type %q", api.Auth.Type)
}
parsedURL.RawQuery = query.Encode()
adapter.url = parsedURL.String()
switch api.Body.Type {
case "":
case "json":
adapter.body, err = json.Marshal(api.Body.Value)
if err != nil {
return nil, fmt.Errorf("new provider HTTP adapter: encode JSON body: %w", err)
}
if adapter.headers.Get("Content-Type") == "" {
adapter.headers.Set("Content-Type", "application/json")
}
case "form":
values := make(url.Values, len(api.Body.Value))
for key, value := range api.Body.Value {
values.Set(key, value)
}
adapter.body = []byte(values.Encode())
if adapter.headers.Get("Content-Type") == "" {
adapter.headers.Set("Content-Type", "application/x-www-form-urlencoded")
}
default:
return nil, fmt.Errorf("new provider HTTP adapter: unsupported body type %q", api.Body.Type)
}
return adapter, nil
}
func (a *HTTPAdapter) Fetch(ctx context.Context) (controllerProvider.FetchResponse, error) {
request, err := http.NewRequestWithContext(ctx, a.method, a.url, bytes.NewReader(a.body))
if err != nil {
return controllerProvider.FetchResponse{}, &operationError{operation: "build provider request", cause: err}
}
request.Header = a.headers.Clone()
if a.basicUser != "" {
request.SetBasicAuth(a.basicUser, a.basicPass)
}
if a.bearerToken != "" {
request.Header.Set("Authorization", "Bearer "+a.bearerToken)
}
response, err := a.doer.Do(request)
if err != nil {
if response != nil && response.Body != nil {
_ = response.Body.Close()
}
return controllerProvider.FetchResponse{}, &operationError{operation: "call provider API", cause: err}
}
if response == nil || response.Body == nil {
return controllerProvider.FetchResponse{}, ErrInvalidHTTPResponse
}
defer response.Body.Close()
body, readErr := readAllLimited(response.Body, a.maxBytes, ErrResponseTooLarge)
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
statusErr := &HTTPStatusError{
StatusCode: response.StatusCode,
retryable: retryableHTTPStatus(response.StatusCode),
}
result := controllerProvider.FetchResponse{}
if statusErr.retryable {
result.RetryAfter = parseRetryAfter(response.Header.Get("Retry-After"), a.now(), a.retryAfterMax)
}
if readErr != nil {
if !errors.Is(readErr, ErrResponseTooLarge) {
readErr = &operationError{operation: "read provider response", cause: readErr}
}
return result, errors.Join(statusErr, readErr)
}
return result, statusErr
}
if readErr != nil {
if errors.Is(readErr, ErrResponseTooLarge) {
return controllerProvider.FetchResponse{}, readErr
}
return controllerProvider.FetchResponse{}, &operationError{operation: "read provider response", cause: readErr}
}
return controllerProvider.FetchResponse{Body: body}, nil
}
func retryableHTTPStatus(status int) bool {
return status == http.StatusRequestTimeout || status == http.StatusTooEarly ||
status == http.StatusTooManyRequests ||
(status >= http.StatusInternalServerError && status <= 599)
}
func parseRetryAfter(value string, now time.Time, maximum time.Duration) time.Duration {
value = strings.TrimSpace(value)
if value == "" {
return 0
}
if seconds, err := strconv.ParseInt(value, 10, 64); err == nil {
if seconds <= 0 {
return 0
}
if seconds > int64(maximum/time.Second) {
return maximum
}
return time.Duration(seconds) * time.Second
}
when, err := http.ParseTime(value)
if err != nil {
return 0
}
delay := when.Sub(now)
if delay <= 0 {
return 0
}
if delay > maximum {
return maximum
}
return delay
}
func defaultHTTPClient() *http.Client {
return &http.Client{
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
func validHTTPToken(value string) bool {
if value == "" {
return false
}
for _, character := range value {
if character <= 0x20 || character >= 0x7f || strings.ContainsRune("()<>@,;:\\\"/[]?={}", character) {
return false
}
}
return true
}
func validHTTPHeaderValue(value string) bool {
for _, character := range value {
if character == '\t' {
continue
}
if character < 0x20 || character == 0x7f {
return false
}
}
return true
}
func headerHas(headers http.Header, key string) bool {
_, exists := headers[http.CanonicalHeaderKey(key)]
return exists
}

View File

@ -0,0 +1,389 @@
package providerapi
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/proxy-pool/proxy-pool/internal/config"
controllerProvider "github.com/proxy-pool/proxy-pool/internal/controller/provider"
)
var _ controllerProvider.ProviderAdapter = (*HTTPAdapter)(nil)
func TestHTTPAdapterBuildsStructuredRequestAndReturnsBody(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if request.Method != http.MethodPost {
t.Errorf("method = %q, want POST", request.Method)
}
if got := request.URL.Query().Get("count"); got != "100" {
t.Errorf("count query = %q, want 100", got)
}
if got := request.Header.Get("X-Provider-Key"); got != "secret" {
t.Errorf("X-Provider-Key = %q, want secret", got)
}
if got := request.Header.Get("X-Request-Mode"); got != "batch" {
t.Errorf("X-Request-Mode = %q, want batch", got)
}
var body map[string]string
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
t.Errorf("decode request body: %v", err)
}
if body["protocol"] != "http" {
t.Errorf("protocol body = %q, want http", body["protocol"])
}
_, _ = writer.Write([]byte("http://192.0.2.20:8080"))
}))
defer server.Close()
adapter, err := NewHTTPAdapter(config.ProviderAPI{
URL: server.URL + "/proxies",
Method: http.MethodPost,
Headers: map[string]string{"X-Request-Mode": "batch"},
Query: map[string]string{"count": "100"},
Auth: config.ProviderAuth{
Type: "apiKey",
Location: "header",
Name: "X-Provider-Key",
Value: "secret",
},
Body: config.APIBody{Type: "json", Value: map[string]string{"protocol": "http"}},
}, config.Fetch{MaxResponseBytes: 1024}, server.Client())
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
response, err := adapter.Fetch(context.Background())
if err != nil {
t.Fatalf("Fetch(): %v", err)
}
if got := string(response.Body); got != "http://192.0.2.20:8080" {
t.Fatalf("response body = %q", got)
}
}
func TestHTTPAdapterFormattingDoesNotExposeProviderSecrets(t *testing.T) {
t.Parallel()
adapter, err := NewHTTPAdapter(config.ProviderAPI{
URL: "https://provider.test/fetch",
Method: http.MethodGet,
Auth: config.ProviderAuth{
Type: "apiKey", Location: "query", Name: "token", Value: "provider-secret",
},
}, config.Fetch{MaxResponseBytes: 1024}, doerFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("unused")
}))
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
formatted := fmt.Sprintf("%v %+v %#v", adapter, adapter, adapter)
for _, secret := range []string{"provider-secret", "token=", "provider.test/fetch"} {
if strings.Contains(formatted, secret) {
t.Fatalf("formatted adapter contains %q: %s", secret, formatted)
}
}
}
func TestHTTPAdapterRedactsTransportErrorURLAndQuerySecret(t *testing.T) {
transportErr := errors.New("dial failed")
adapter, err := NewHTTPAdapter(config.ProviderAPI{
URL: "https://provider.example/proxies",
Method: http.MethodGet,
Auth: config.ProviderAuth{
Type: "apiKey",
Location: "query",
Name: "api_key",
Value: "top-secret",
},
}, config.Fetch{MaxResponseBytes: 1024}, doerFunc(func(request *http.Request) (*http.Response, error) {
return nil, fmt.Errorf("request %s: %w", request.URL.String(), transportErr)
}))
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
_, err = adapter.Fetch(context.Background())
if !errors.Is(err, transportErr) {
t.Fatalf("Fetch() error = %v, want wrapped transport cause", err)
}
if strings.Contains(err.Error(), "top-secret") || strings.Contains(err.Error(), "provider.example") {
t.Fatalf("transport error leaked Provider URL or API key: %v", err)
}
}
func TestHTTPAdapterClosesResponseBodyReturnedWithTransportError(t *testing.T) {
body := &trackingBody{Reader: strings.NewReader("unused")}
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: "https://provider.example/proxies", Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 1024},
doerFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusBadGateway, Body: body}, errors.New("transport failed")
}),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
_, _ = adapter.Fetch(context.Background())
if !body.closed {
t.Fatal("response body returned with transport error was not closed")
}
}
func TestHTTPAdapterRejectsConflictingAuthenticationFields(t *testing.T) {
tests := []struct {
name string
api config.ProviderAPI
}{
{
name: "bearer conflicts with authorization header",
api: config.ProviderAPI{
URL: "https://provider.example/proxies",
Headers: map[string]string{"Authorization": "configured"},
Auth: config.ProviderAuth{Type: "bearer", Token: "resolved-token"},
},
},
{
name: "API key conflicts with configured header",
api: config.ProviderAPI{
URL: "https://provider.example/proxies",
Headers: map[string]string{"X-API-Key": "configured"},
Auth: config.ProviderAuth{
Type: "apiKey", Location: "header", Name: "x-api-key", Value: "resolved-key",
},
},
},
{
name: "API key conflicts with URL query",
api: config.ProviderAPI{
URL: "https://provider.example/proxies?api_key=configured",
Auth: config.ProviderAuth{
Type: "apiKey", Location: "query", Name: "api_key", Value: "resolved-key",
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if _, err := NewHTTPAdapter(tt.api, config.Fetch{MaxResponseBytes: 1024}, nil); err == nil {
t.Fatal("NewHTTPAdapter() error = nil, want authentication conflict")
}
})
}
}
type doerFunc func(*http.Request) (*http.Response, error)
func (f doerFunc) Do(request *http.Request) (*http.Response, error) { return f(request) }
func TestHTTPAdapterClassifiesStatusAndRetryAfter(t *testing.T) {
tests := []struct {
name string
status int
retryAfter string
retryMax config.Duration
wantRetry bool
wantDelay time.Duration
}{
{name: "bad request is permanent", status: http.StatusBadRequest},
{name: "authentication is permanent", status: http.StatusUnauthorized},
{name: "request timeout is transient", status: http.StatusRequestTimeout, wantRetry: true},
{name: "rate limit honors bounded retry after", status: http.StatusTooManyRequests, retryAfter: "60", retryMax: config.Duration(2 * time.Second), wantRetry: true, wantDelay: 2 * time.Second},
{name: "server error is transient", status: http.StatusServiceUnavailable, retryAfter: "1", retryMax: config.Duration(5 * time.Second), wantRetry: true, wantDelay: time.Second},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
writer.Header().Set("Retry-After", tt.retryAfter)
writer.WriteHeader(tt.status)
_, _ = writer.Write([]byte("sensitive provider response"))
}))
defer server.Close()
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: server.URL, Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 1024, Retry: config.Backoff{Max: tt.retryMax}},
server.Client(),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
response, err := adapter.Fetch(context.Background())
if err == nil {
t.Fatal("Fetch() error = nil, want HTTP status error")
}
var statusErr *HTTPStatusError
if !errors.As(err, &statusErr) || statusErr.StatusCode != tt.status {
t.Fatalf("Fetch() error = %v, want HTTPStatusError{%d}", err, tt.status)
}
var retryable controllerProvider.RetryableError
if !errors.As(err, &retryable) || retryable.Retryable() != tt.wantRetry {
t.Fatalf("Retryable() = %v, want %v", retryable != nil && retryable.Retryable(), tt.wantRetry)
}
if response.RetryAfter != tt.wantDelay {
t.Fatalf("RetryAfter = %s, want %s", response.RetryAfter, tt.wantDelay)
}
if strings.Contains(err.Error(), "sensitive") || strings.Contains(err.Error(), server.URL) {
t.Fatalf("status error leaked response or URL: %v", err)
}
})
}
}
func TestHTTPAdapterRejectsResponseOverMaxBytes(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
_, _ = writer.Write([]byte("12345"))
}))
defer server.Close()
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: server.URL, Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 4},
server.Client(),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
response, err := adapter.Fetch(context.Background())
if !errors.Is(err, ErrResponseTooLarge) {
t.Fatalf("Fetch() error = %v, want ErrResponseTooLarge", err)
}
if response.Body != nil {
t.Fatalf("response body = %q, want nil on overflow", response.Body)
}
}
func TestHTTPAdapterPreservesPermanentStatusWhenErrorBodyIsOversized(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
writer.WriteHeader(http.StatusUnauthorized)
_, _ = writer.Write([]byte("12345"))
}))
defer server.Close()
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: server.URL, Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 4},
server.Client(),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
_, err = adapter.Fetch(context.Background())
var statusErr *HTTPStatusError
if !errors.As(err, &statusErr) || statusErr.StatusCode != http.StatusUnauthorized {
t.Fatalf("Fetch() error = %v, want HTTPStatusError{401}", err)
}
if !errors.Is(err, ErrResponseTooLarge) {
t.Fatalf("Fetch() error = %v, want joined ErrResponseTooLarge", err)
}
if statusErr.Retryable() {
t.Fatal("HTTPStatusError{401}.Retryable() = true, want false")
}
}
func TestHTTPAdapterAcceptsResponseAtExactByteLimitAndClosesBody(t *testing.T) {
body := &trackingBody{Reader: strings.NewReader("1234")}
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: "https://provider.example/proxies", Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 4},
doerFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil
}),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
response, err := adapter.Fetch(context.Background())
if err != nil {
t.Fatalf("Fetch(): %v", err)
}
if string(response.Body) != "1234" {
t.Fatalf("response body = %q, want exact-limit body", response.Body)
}
if !body.closed {
t.Fatal("response body was not closed")
}
}
func TestHTTPAdapterParsesHTTPDateAndUsesSafeDefaultRetryAfterCap(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
writer.Header().Set("Retry-After", time.Now().Add(time.Hour).UTC().Format(http.TimeFormat))
writer.WriteHeader(http.StatusTooManyRequests)
}))
defer server.Close()
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: server.URL, Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 1024},
server.Client(),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
response, err := adapter.Fetch(context.Background())
if err == nil {
t.Fatal("Fetch() error = nil, want rate limit status")
}
if response.RetryAfter != defaultRetryAfterMax {
t.Fatalf("RetryAfter = %s, want safe default cap %s", response.RetryAfter, defaultRetryAfterMax)
}
}
func TestHTTPAdapterRejectsMalformedDoerResponse(t *testing.T) {
tests := []struct {
name string
response *http.Response
}{
{name: "nil response"},
{name: "nil body", response: &http.Response{StatusCode: http.StatusOK}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
adapter, err := NewHTTPAdapter(
config.ProviderAPI{URL: "https://provider.example/proxies", Method: http.MethodGet},
config.Fetch{MaxResponseBytes: 1024},
doerFunc(func(*http.Request) (*http.Response, error) { return tt.response, nil }),
)
if err != nil {
t.Fatalf("NewHTTPAdapter(): %v", err)
}
_, err = adapter.Fetch(context.Background())
if !errors.Is(err, ErrInvalidHTTPResponse) {
t.Fatalf("Fetch() error = %v, want ErrInvalidHTTPResponse", err)
}
})
}
}
func TestRetryableHTTPStatusRejectsOutOfRangeStatus(t *testing.T) {
if retryableHTTPStatus(600) {
t.Fatal("retryableHTTPStatus(600) = true, want false")
}
}
type trackingBody struct {
io.Reader
closed bool
}
func (b *trackingBody) Close() error {
b.closed = true
return nil
}

View File

@ -0,0 +1,77 @@
package providerapi
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"math"
)
var (
ErrResponseTooLarge = errors.New("provider response exceeds byte limit")
ErrTemplateInputTooLarge = errors.New("provider template input exceeds byte limit")
ErrTemplateOutputTooLarge = errors.New("provider template output exceeds byte limit")
ErrTooManyCandidates = errors.New("provider template output exceeds candidate limit")
ErrInvalidProxyOutput = errors.New("provider template output contains no valid proxy")
ErrRecursiveTemplate = errors.New("provider template contains recursive calls")
ErrTemplateTooComplex = errors.New("provider template exceeds static complexity limit")
ErrCredentialStoreRequired = errors.New("provider credentials require a credential store")
ErrInvalidHTTPResponse = errors.New("provider HTTP client returned an invalid response")
)
type limitError struct {
kind error
size int64
limit int64
}
func (e *limitError) Error() string {
return fmt.Sprintf("%v: size %d, limit %d", e.kind, e.size, e.limit)
}
func (e *limitError) Unwrap() error { return e.kind }
func enforceByteLimit(kind error, size, limit int64) error {
if size <= limit {
return nil
}
return &limitError{kind: kind, size: size, limit: limit}
}
func readAllLimited(reader io.Reader, limit int64, kind error) ([]byte, error) {
readLimit := limit
if readLimit < math.MaxInt64 {
readLimit++
}
limited := &io.LimitedReader{R: reader, N: readLimit}
data, err := io.ReadAll(limited)
if err != nil {
return nil, err
}
if err := enforceByteLimit(kind, int64(len(data)), limit); err != nil {
return nil, err
}
return data, nil
}
type limitedBuffer struct {
ctx context.Context
kind error
limit int64
buffer bytes.Buffer
}
func (w *limitedBuffer) Write(data []byte) (int, error) {
if err := w.ctx.Err(); err != nil {
return 0, err
}
current := int64(w.buffer.Len())
if err := enforceByteLimit(w.kind, current+int64(len(data)), w.limit); err != nil {
return 0, err
}
return w.buffer.Write(data)
}
func (w *limitedBuffer) String() string { return w.buffer.String() }

View File

@ -0,0 +1,23 @@
package providerapi
import (
"errors"
"testing"
"github.com/proxy-pool/proxy-pool/internal/config"
)
func TestTemplateParserRejectsNestedRanges(t *testing.T) {
t.Parallel()
_, err := NewTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{
Template: `{{range .}}{{range .}}{{end}}{{end}}`,
},
Pool: config.Pool{MaxSize: 10},
}, nil)
if !errors.Is(err, ErrTemplateTooComplex) {
t.Fatalf("NewTemplateParser() error = %v, want ErrTemplateTooComplex", err)
}
}

View File

@ -0,0 +1,340 @@
package providerapi
import (
"context"
"fmt"
"net/url"
"regexp"
"strconv"
"strings"
"sync"
"text/template"
"time"
"github.com/proxy-pool/proxy-pool/internal/config"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
"github.com/proxy-pool/proxy-pool/internal/platform/credentials"
)
const (
defaultTemplateTimeout = 100 * time.Millisecond
defaultTemplateMaxBytes = int64(1 << 20)
defaultMaxCandidates = 10_000
maxRegexPatterns = 64
maxRegexPatternBytes = 1024
)
type TemplateParser struct {
upstreamID string
template *template.Template
defaultScheme proxyDomain.Scheme
allowed map[proxyDomain.Scheme]struct{}
proxyAuthType string
username string
password string
maxConcurrency int64
timeout time.Duration
maxInputBytes int64
maxOutputBytes int64
maxCandidates int
credentialStore credentials.Store
executions *executionLimiter
regexMu sync.Mutex
regexes map[string]*regexp.Regexp
}
func (p *TemplateParser) String() string {
if p == nil {
return "providerapi.TemplateParser<nil>"
}
return fmt.Sprintf(
"providerapi.TemplateParser{upstream:%q,maxInputBytes:%d,maxOutputBytes:%d,maxCandidates:%d}",
p.upstreamID,
p.maxInputBytes,
p.maxOutputBytes,
p.maxCandidates,
)
}
func (p *TemplateParser) GoString() string { return p.String() }
func NewTemplateParser(
upstreamID string,
upstream config.Upstream,
credentialStore credentials.Store,
) (*TemplateParser, error) {
if strings.TrimSpace(upstreamID) == "" {
return nil, fmt.Errorf("new provider template parser: upstream ID is required")
}
parser := &TemplateParser{
upstreamID: upstreamID,
allowed: make(map[proxyDomain.Scheme]struct{}),
proxyAuthType: upstream.ProxyAuth.Type,
username: upstream.ProxyAuth.Username,
password: upstream.ProxyAuth.Password,
maxConcurrency: int64(upstream.Capacity.MaxConcurrencyPerProxy),
timeout: time.Duration(upstream.Fetch.TemplateTimeout),
maxInputBytes: upstream.Fetch.MaxResponseBytes,
maxOutputBytes: upstream.Fetch.MaxResponseBytes,
maxCandidates: upstream.Pool.MaxSize,
credentialStore: credentialStore,
executions: newExecutionLimiter(upstream.Fetch.MaxInFlight),
regexes: make(map[string]*regexp.Regexp),
}
if parser.timeout <= 0 {
parser.timeout = defaultTemplateTimeout
}
if parser.maxInputBytes <= 0 {
parser.maxInputBytes = defaultTemplateMaxBytes
}
if parser.maxOutputBytes <= 0 {
parser.maxOutputBytes = defaultTemplateMaxBytes
}
if parser.maxCandidates <= 0 {
parser.maxCandidates = defaultMaxCandidates
}
if parser.maxConcurrency <= 0 {
parser.maxConcurrency = 1
}
if parser.proxyAuthType == "" {
parser.proxyAuthType = "response"
}
switch parser.proxyAuthType {
case "response", "static", "ipWhitelist":
default:
return nil, fmt.Errorf("new provider template parser: unsupported proxy auth type %q", parser.proxyAuthType)
}
for _, protocol := range upstream.Provider.Protocols {
scheme := proxyDomain.Scheme(strings.ToLower(strings.TrimSpace(protocol)))
if !supportedScheme(scheme) {
return nil, fmt.Errorf("new provider template parser: unsupported protocol %q", protocol)
}
if parser.defaultScheme == "" {
parser.defaultScheme = scheme
}
parser.allowed[scheme] = struct{}{}
}
if parser.defaultScheme == "" {
parser.defaultScheme = proxyDomain.SchemeHTTP
parser.allowed[parser.defaultScheme] = struct{}{}
}
parsed, err := template.New("provider-response").Option("missingkey=error").Funcs(template.FuncMap{
"regexFind": parser.regexFind,
"regexFindAll": parser.regexFindAll,
}).Parse(upstream.API.Template)
if err != nil {
return nil, fmt.Errorf("new provider template parser: parse template: %w", err)
}
if err := validateTemplateComplexity(parsed); err != nil {
return nil, err
}
if err := rejectRecursiveTemplates(parsed); err != nil {
return nil, err
}
parser.template = parsed
return parser, nil
}
func (p *TemplateParser) Parse(ctx context.Context, body []byte) ([]proxyDomain.Proxy, error) {
if err := ctx.Err(); err != nil {
return nil, err
}
if err := enforceByteLimit(ErrTemplateInputTooLarge, int64(len(body)), p.maxInputBytes); err != nil {
return nil, err
}
execCtx, cancel := context.WithTimeout(ctx, p.timeout)
defer cancel()
output := limitedBuffer{ctx: execCtx, kind: ErrTemplateOutputTooLarge, limit: p.maxOutputBytes}
err := p.executions.Run(execCtx, func() error {
return p.template.Execute(&output, string(body))
})
if ctxErr := execCtx.Err(); ctxErr != nil {
return nil, ctxErr
}
if err != nil {
return nil, fmt.Errorf("execute provider template: %w", err)
}
tokens := strings.Fields(output.String())
if len(tokens) > p.maxCandidates {
return nil, &limitError{kind: ErrTooManyCandidates, size: int64(len(tokens)), limit: int64(p.maxCandidates)}
}
proxies := make([]proxyDomain.Proxy, 0, len(tokens))
credentialIndexes := make(map[string]int)
for _, token := range tokens {
candidate, credential, ok := p.parseCandidate(token)
if ok {
credentialKey := ""
if credential != nil {
if p.credentialStore == nil {
return nil, ErrCredentialStoreRequired
}
reference, err := p.credentialStore.Put(ctx, p.credentialScope(candidate), *credential)
if err != nil {
return nil, &operationError{operation: "store provider credentials", cause: err}
}
if reference.SecretRef == "" || reference.CredentialVersion == "" {
return nil, &operationError{
operation: "store provider credentials",
cause: credentials.ErrInvalidReference,
}
}
candidate.SecretRef = reference.SecretRef
candidate.CredentialVersion = reference.CredentialVersion
credentialKey = candidateCredentialKey(candidate)
if index, exists := credentialIndexes[credentialKey]; exists {
proxies[index] = candidate
continue
}
}
if len(proxies) >= p.maxCandidates {
return nil, &limitError{kind: ErrTooManyCandidates, size: int64(len(proxies) + 1), limit: int64(p.maxCandidates)}
}
proxies = append(proxies, candidate)
if credentialKey != "" {
credentialIndexes[credentialKey] = len(proxies) - 1
}
}
}
if len(tokens) > 0 && len(proxies) == 0 {
return nil, ErrInvalidProxyOutput
}
return proxies, nil
}
func (p *TemplateParser) regexFind(pattern, value string) (string, error) {
compiled, err := p.compileRegex(pattern)
if err != nil {
return "", err
}
return compiled.FindString(value), nil
}
func (p *TemplateParser) regexFindAll(pattern, value string, count int) ([]string, error) {
compiled, err := p.compileRegex(pattern)
if err != nil {
return nil, err
}
matchLimit := p.maxCandidates
if matchLimit < int(^uint(0)>>1) {
matchLimit++
}
effectiveCount := count
if effectiveCount < 0 || effectiveCount > matchLimit {
effectiveCount = matchLimit
}
matches := compiled.FindAllString(value, effectiveCount)
if len(matches) > p.maxCandidates {
return nil, &limitError{kind: ErrTooManyCandidates, size: int64(len(matches)), limit: int64(p.maxCandidates)}
}
return matches, nil
}
func (p *TemplateParser) compileRegex(pattern string) (*regexp.Regexp, error) {
if len(pattern) > maxRegexPatternBytes {
return nil, fmt.Errorf("regex pattern exceeds %d bytes", maxRegexPatternBytes)
}
p.regexMu.Lock()
defer p.regexMu.Unlock()
if compiled := p.regexes[pattern]; compiled != nil {
return compiled, nil
}
if len(p.regexes) >= maxRegexPatterns {
return nil, fmt.Errorf("template exceeds %d distinct regex patterns", maxRegexPatterns)
}
compiled, err := regexp.Compile(pattern)
if err != nil {
return nil, fmt.Errorf("compile template regex: invalid pattern")
}
p.regexes[pattern] = compiled
return compiled, nil
}
func (p *TemplateParser) parseCandidate(raw string) (proxyDomain.Proxy, *credentials.Value, bool) {
if !strings.Contains(raw, "://") {
raw = string(p.defaultScheme) + "://" + raw
}
parsed, err := url.Parse(raw)
if err != nil || parsed.Opaque != "" || parsed.Hostname() == "" || parsed.Port() == "" ||
parsed.Path != "" || parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" {
return proxyDomain.Proxy{}, nil, false
}
scheme := proxyDomain.Scheme(strings.ToLower(parsed.Scheme))
if _, ok := p.allowed[scheme]; !ok || !supportedScheme(scheme) {
return proxyDomain.Proxy{}, nil, false
}
port, err := strconv.ParseUint(parsed.Port(), 10, 16)
if err != nil || port == 0 {
return proxyDomain.Proxy{}, nil, false
}
username := ""
var credential *credentials.Value
switch p.proxyAuthType {
case "static":
username = p.username
if p.username != "" || p.password != "" {
credential = &credentials.Value{Username: p.username, Password: p.password}
}
case "response":
if parsed.User != nil {
username = parsed.User.Username()
password, hasPassword := parsed.User.Password()
if username != "" || hasPassword {
credential = &credentials.Value{Username: username, Password: password}
}
}
}
return proxyDomain.Proxy{
Scheme: scheme,
Host: parsed.Hostname(),
Port: uint16(port),
Username: username,
SourceUpstream: p.upstreamID,
MaxConcurrency: p.maxConcurrency,
State: proxyDomain.StateFetched,
}, credential, true
}
func (p *TemplateParser) credentialScope(candidate proxyDomain.Proxy) string {
if p.proxyAuthType == "static" {
return lengthPrefixedScope("provider", p.upstreamID, "static")
}
return lengthPrefixedScope(
"provider",
p.upstreamID,
string(candidate.Scheme),
strings.ToLower(candidate.Host),
strconv.FormatUint(uint64(candidate.Port), 10),
candidate.Username,
)
}
func lengthPrefixedScope(parts ...string) string {
var scope strings.Builder
for _, part := range parts {
scope.WriteString(strconv.Itoa(len(part)))
scope.WriteByte(':')
scope.WriteString(part)
}
return scope.String()
}
func candidateCredentialKey(candidate proxyDomain.Proxy) string {
return lengthPrefixedScope(
string(candidate.Scheme),
strings.ToLower(candidate.Host),
strconv.FormatUint(uint64(candidate.Port), 10),
candidate.Username,
)
}
func supportedScheme(scheme proxyDomain.Scheme) bool {
switch scheme {
case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5:
return true
default:
return false
}
}

View File

@ -0,0 +1,681 @@
package providerapi
import (
"context"
"errors"
"fmt"
"runtime"
"strings"
"sync"
"testing"
"time"
"github.com/proxy-pool/proxy-pool/internal/config"
controllerProvider "github.com/proxy-pool/proxy-pool/internal/controller/provider"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
"github.com/proxy-pool/proxy-pool/internal/platform/credentials"
)
var _ controllerProvider.Parser = (*TemplateParser)(nil)
func newTemplateParser(
upstreamID string,
upstream config.Upstream,
stores ...credentials.Store,
) (*TemplateParser, error) {
var store credentials.Store
if len(stores) > 0 {
store = stores[0]
}
return NewTemplateParser(upstreamID, upstream, store)
}
func TestTemplateParserFormattingDoesNotExposeTemplateCredentials(t *testing.T) {
t.Parallel()
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://response-user:response-secret@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "static", Username: "static-user", Password: "static-secret"},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
formatted := fmt.Sprintf("%v %+v %#v", parser, parser, parser)
for _, secret := range []string{"response-secret", "static-secret", "response-user"} {
if strings.Contains(formatted, secret) {
t.Fatalf("formatted parser contains %q: %s", secret, formatted)
}
}
}
func TestTemplateParserExecutesWhitelistedRegexAndBuildsProxy(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{$value := regexFind "[0-9.]+:[0-9]+" .}}{{if $value}}{{printf "http://%s\n" $value}}{{end}}`},
Pool: config.Pool{MaxSize: 10},
Capacity: config.Capacity{MaxConcurrencyPerProxy: 7},
Fetch: config.Fetch{
MaxResponseBytes: 1024,
TemplateTimeout: config.Duration(time.Second),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), []byte("address=192.0.2.10:8080"))
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 {
t.Fatalf("proxy count = %d, want 1", len(proxies))
}
got := proxies[0]
if got.Scheme != proxyDomain.SchemeHTTP || got.Host != "192.0.2.10" || got.Port != 8080 {
t.Fatalf("proxy endpoint = %s://%s:%d, want http://192.0.2.10:8080", got.Scheme, got.Host, got.Port)
}
if got.SourceUpstream != "provider-a" || got.MaxConcurrency != 7 || got.State != proxyDomain.StateFetched {
t.Fatalf("proxy metadata = %+v", got)
}
}
func TestTemplateParserRejectsOversizedInput(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{.}}`},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{
MaxResponseBytes: 4,
TemplateTimeout: config.Duration(time.Second),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), []byte("12345"))
if !errors.Is(err, ErrTemplateInputTooLarge) {
t.Fatalf("Parse() error = %v, want ErrTemplateInputTooLarge", err)
}
}
func TestTemplateParserStopsOversizedOutputDuringExecution(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `0123456789`},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{
MaxResponseBytes: 8,
TemplateTimeout: config.Duration(time.Second),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), []byte("x"))
if !errors.Is(err, ErrTemplateOutputTooLarge) {
t.Fatalf("Parse() error = %v, want ErrTemplateOutputTooLarge", err)
}
}
func TestTemplateParserRejectsTooManyCandidates(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: "http://192.0.2.1:8001\nhttp://192.0.2.2:8002\n"},
Pool: config.Pool{MaxSize: 1},
Fetch: config.Fetch{
MaxResponseBytes: 1024,
TemplateTimeout: config.Duration(time.Second),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), nil)
if !errors.Is(err, ErrTooManyCandidates) {
t.Fatalf("Parse() error = %v, want ErrTooManyCandidates", err)
}
}
func TestTemplateParserRegexFindAllCountSurvivesMaxCandidateBoundary(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{range regexFindAll "a" . 1}}http://192.0.2.10:8080
{{end}}`},
Pool: config.Pool{MaxSize: int(^uint(0) >> 1)},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), []byte("aaa"))
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 {
t.Fatalf("proxy count = %d, want regex count limit 1", len(proxies))
}
}
func TestTemplateParserEnforcesExecutionTimeout(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{$matches := regexFindAll "a" . -1}}{{range $matches}}http://192.0.2.10:8080
{{end}}`},
Pool: config.Pool{MaxSize: 20_000},
Fetch: config.Fetch{
MaxResponseBytes: 1 << 20,
TemplateTimeout: config.Duration(time.Nanosecond),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), []byte(strings.Repeat("a", 10_000)))
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Parse() error = %v, want context.DeadlineExceeded", err)
}
}
func TestTemplateParserTimedOutExecutionsRemainBounded(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{$matches := regexFindAll "a" . -1}}{{range $matches}}{{printf "%s" .}}{{end}}`},
Pool: config.Pool{MaxSize: 20_000},
Fetch: config.Fetch{
MaxResponseBytes: 1 << 20,
TemplateTimeout: config.Duration(time.Nanosecond),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
baseline := runtime.NumGoroutine()
for range 20 {
_, err := parser.Parse(context.Background(), []byte(strings.Repeat("a", 10_000)))
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Parse() error = %v, want context.DeadlineExceeded", err)
}
}
deadline := time.Now().Add(2 * time.Second)
for runtime.NumGoroutine() > baseline+4 && time.Now().Before(deadline) {
runtime.Gosched()
time.Sleep(time.Millisecond)
}
if got := runtime.NumGoroutine(); got > baseline+4 {
t.Fatalf("goroutines after timed-out executions = %d, baseline = %d", got, baseline)
}
}
func TestTemplateParserRejectsRecursiveTemplates(t *testing.T) {
_, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `
{{define "first"}}{{template "second"}}{{end}}
{{define "second"}}{{template "first"}}{{end}}
{{template "first"}}`},
Pool: config.Pool{MaxSize: 10},
})
if !errors.Is(err, ErrRecursiveTemplate) {
t.Fatalf("newTemplateParser() error = %v, want ErrRecursiveTemplate", err)
}
}
func TestTemplateParserRejectsNonEmptyAllInvalidOutput(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `not-a-proxy`},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{
MaxResponseBytes: 1024,
TemplateTimeout: config.Duration(time.Second),
},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), nil)
if !errors.Is(err, ErrInvalidProxyOutput) {
t.Fatalf("Parse() error = %v, want ErrInvalidProxyOutput", err)
}
}
func TestTemplateParserAcceptsIPv6AndFiltersUnsafeEndpoints(t *testing.T) {
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: strings.Join([]string{
"http://[2001:db8::1]:8080",
"http://192.0.2.1:8080/path",
"http://192.0.2.2:8080?token=secret",
"http://192.0.2.3:0",
"http://192.0.2.4:65536",
"socks5://192.0.2.5:1080",
}, "\n")},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 || proxies[0].Host != "2001:db8::1" || proxies[0].Port != 8080 {
t.Fatalf("proxies = %+v, want only IPv6 endpoint", proxies)
}
}
func TestTemplateParserDoesNotOverrideStaticProxyAuthFromResponse(t *testing.T) {
store, err := credentials.NewMemoryStore(10)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://response-user:response-pass@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "static", Username: "configured-user", Password: "configured-pass"},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 || proxies[0].Username != "configured-user" {
t.Fatalf("proxies = %+v, want configured static username", proxies)
}
if strings.Contains(proxies[0].SecretRef, "configured-pass") || strings.Contains(proxies[0].SecretRef, "response-pass") {
t.Fatalf("SecretRef contains plaintext password: %q", proxies[0].SecretRef)
}
if proxies[0].SecretRef == "" || proxies[0].CredentialVersion == "" {
t.Fatalf("proxy credential reference is incomplete: %+v", proxies[0])
}
value, err := store.Resolve(context.Background(), credentials.Reference{
SecretRef: proxies[0].SecretRef,
CredentialVersion: proxies[0].CredentialVersion,
})
if err != nil {
t.Fatalf("Resolve(): %v", err)
}
if value.Username != "configured-user" || value.Password != "configured-pass" {
t.Fatalf("resolved static credentials = %v, want configured credentials", value)
}
}
func TestTemplateParserRetainsDistinctEndpointsSharingStaticCredentials(t *testing.T) {
store, err := credentials.NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: strings.Join([]string{
"http://192.0.2.10:8080",
"http://192.0.2.11:8080",
}, "\n")},
ProxyAuth: config.ProxyAuth{Type: "static", Username: "configured-user", Password: "configured-pass"},
Pool: config.Pool{MaxSize: 2},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("NewTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 2 {
t.Fatalf("proxy count = %d, want both static-auth endpoints", len(proxies))
}
if proxies[0].SecretRef == "" || proxies[0].SecretRef != proxies[1].SecretRef {
t.Fatalf("static credential references = %q and %q, want same opaque reference", proxies[0].SecretRef, proxies[1].SecretRef)
}
}
func TestTemplateParserStoresResponseCredentialsByOpaqueReference(t *testing.T) {
store, err := credentials.NewMemoryStore(10)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://response-user:response-secret@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 {
t.Fatalf("proxy count = %d, want 1", len(proxies))
}
got := proxies[0]
if got.Username != "response-user" || got.SecretRef == "" || got.CredentialVersion == "" {
t.Fatalf("proxy credentials metadata = %+v", got)
}
formatted := fmt.Sprintf("%v %+v %#v", got, got, got)
if strings.Contains(formatted, "response-secret") || strings.Contains(got.SecretRef, "response-secret") {
t.Fatalf("proxy formatting or reference contains password: %s", formatted)
}
value, err := store.Resolve(context.Background(), credentials.Reference{
SecretRef: got.SecretRef,
CredentialVersion: got.CredentialVersion,
})
if err != nil {
t.Fatalf("Resolve(): %v", err)
}
if value.Username != "response-user" || value.Password != "response-secret" {
t.Fatalf("resolved response credentials = %v, want response credentials", value)
}
}
func TestTemplateParserKeepsDistinctAccountsForSameEndpointResolvable(t *testing.T) {
store, err := credentials.NewMemoryStore(2)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: strings.Join([]string{
"http://alice:alice-secret@192.0.2.10:8080",
"http://bob:bob-secret@192.0.2.10:8080",
}, "\n")},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 2},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("NewTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 2 {
t.Fatalf("proxy count = %d, want 2", len(proxies))
}
for _, candidate := range proxies {
value, err := store.Resolve(context.Background(), credentials.Reference{
SecretRef: candidate.SecretRef,
CredentialVersion: candidate.CredentialVersion,
})
if err != nil {
t.Fatalf("Resolve(%s): %v", candidate.Username, err)
}
if value.Username != candidate.Username {
t.Fatalf("resolved username = %q, want %q", value.Username, candidate.Username)
}
}
}
func TestTemplateParserKeepsLatestCredentialVersionWithinOneResponse(t *testing.T) {
store, err := credentials.NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: strings.Join([]string{
"http://alice:old-secret@192.0.2.10:8080",
"http://alice:new-secret@192.0.2.10:8080",
}, "\n")},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 2},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("NewTemplateParser(): %v", err)
}
proxies, err := parser.Parse(context.Background(), nil)
if err != nil {
t.Fatalf("Parse(): %v", err)
}
if len(proxies) != 1 {
t.Fatalf("proxy count = %d, want latest credential only", len(proxies))
}
value, err := store.Resolve(context.Background(), credentials.Reference{
SecretRef: proxies[0].SecretRef,
CredentialVersion: proxies[0].CredentialVersion,
})
if err != nil {
t.Fatalf("Resolve(): %v", err)
}
if value.Password != "new-secret" {
t.Fatalf("resolved latest password mismatch")
}
}
func TestTemplateParserRedactsCredentialStoreErrors(t *testing.T) {
const password = "provider-password"
storeErr := errors.New("store failed for " + password)
store := &failingCredentialStore{err: storeErr}
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://user:` + password + `@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}, store)
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), nil)
if !errors.Is(err, storeErr) {
t.Fatalf("Parse() error = %v, want wrapped store error", err)
}
if strings.Contains(err.Error(), password) {
t.Fatalf("Parse() error leaked password: %v", err)
}
}
func TestTemplateParserRejectsMissingStoreAndEmptyCredentialReference(t *testing.T) {
upstream := config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `http://user:password@192.0.2.10:8080`},
ProxyAuth: config.ProxyAuth{Type: "response"},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
}
parser, err := newTemplateParser("provider-a", upstream)
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
if _, err := parser.Parse(context.Background(), nil); !errors.Is(err, ErrCredentialStoreRequired) {
t.Fatalf("Parse() error = %v, want ErrCredentialStoreRequired", err)
}
parser, err = newTemplateParser("provider-a", upstream, emptyReferenceCredentialStore{})
if err != nil {
t.Fatalf("newTemplateParser(with store): %v", err)
}
if _, err := parser.Parse(context.Background(), nil); !errors.Is(err, credentials.ErrInvalidReference) {
t.Fatalf("Parse() error = %v, want credentials.ErrInvalidReference", err)
}
}
type failingCredentialStore struct{ err error }
func (s *failingCredentialStore) Put(context.Context, string, credentials.Value) (credentials.Reference, error) {
return credentials.Reference{}, s.err
}
func (s *failingCredentialStore) Resolve(context.Context, credentials.Reference) (credentials.Value, error) {
return credentials.Value{}, s.err
}
type emptyReferenceCredentialStore struct{}
func (emptyReferenceCredentialStore) Put(context.Context, string, credentials.Value) (credentials.Reference, error) {
return credentials.Reference{}, nil
}
func (emptyReferenceCredentialStore) Resolve(context.Context, credentials.Reference) (credentials.Value, error) {
return credentials.Value{}, credentials.ErrCredentialMissing
}
func TestTemplateParserRejectsFunctionsOutsideWhitelist(t *testing.T) {
tests := []string{
`{{env "SECRET_TOKEN"}}`,
`{{readFile "credentials.txt"}}`,
`{{httpGet "https://example.invalid"}}`,
`{{exec "command"}}`,
}
for _, source := range tests {
if _, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: source},
Pool: config.Pool{MaxSize: 10},
}); err == nil {
t.Fatalf("newTemplateParser(%q) error = nil, want unknown function error", source)
}
}
}
func TestTemplateParserRedactsInvalidRegexPattern(t *testing.T) {
const secretPattern = "(?P<secret-token>"
parser, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: `{{regexFind . "value"}}`},
Pool: config.Pool{MaxSize: 10},
Fetch: config.Fetch{MaxResponseBytes: 1024, TemplateTimeout: config.Duration(time.Second)},
})
if err != nil {
t.Fatalf("newTemplateParser(): %v", err)
}
_, err = parser.Parse(context.Background(), []byte(secretPattern))
if err == nil {
t.Fatal("Parse() error = nil, want invalid regex error")
}
if strings.Contains(err.Error(), secretPattern) || strings.Contains(err.Error(), "secret-token") {
t.Fatalf("Parse() error leaked regex pattern: %v", err)
}
}
func TestTemplateParserRejectsExcessiveASTNesting(t *testing.T) {
const nesting = 40
source := strings.Repeat("{{range .}}", nesting) + strings.Repeat("{{end}}", nesting)
_, err := newTemplateParser("provider-a", config.Upstream{
Provider: config.Provider{Protocols: []string{"http"}},
API: config.ProviderAPI{Template: source},
Pool: config.Pool{MaxSize: 10},
})
if !errors.Is(err, ErrTemplateTooComplex) {
t.Fatalf("newTemplateParser() error = %v, want ErrTemplateTooComplex", err)
}
}
func TestExecutionLimiterRetainsSlotUntilTimedOutExecutionExits(t *testing.T) {
limiter := newExecutionLimiter(1)
started := make(chan struct{})
release := make(chan struct{})
firstDone := make(chan error, 1)
firstCtx, cancelFirst := context.WithCancel(context.Background())
go func() {
firstDone <- limiter.Run(firstCtx, func() error {
close(started)
<-release
return nil
})
}()
<-started
cancelFirst()
if err := <-firstDone; !errors.Is(err, context.Canceled) {
t.Fatalf("first Run() error = %v, want context.Canceled", err)
}
secondStarted := make(chan struct{})
secondCtx, cancelSecond := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancelSecond()
err := limiter.Run(secondCtx, func() error {
close(secondStarted)
return nil
})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("second Run() error = %v, want context.DeadlineExceeded", err)
}
select {
case <-secondStarted:
t.Fatal("second execution started while timed-out execution retained the only slot")
default:
}
close(release)
deadline := time.Now().Add(time.Second)
for limiter.InFlight() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := limiter.InFlight(); got != 0 {
t.Fatalf("in-flight executions = %d, want 0 after execution exit", got)
}
}
func TestExecutionLimiterIsSafeForConcurrentInspection(t *testing.T) {
limiter := newExecutionLimiter(2)
var wait sync.WaitGroup
for range 10 {
wait.Add(1)
go func() {
defer wait.Done()
_ = limiter.InFlight()
}()
}
wait.Wait()
}
func TestExecutionLimiterDoesNotStartWithCanceledContext(t *testing.T) {
limiter := newExecutionLimiter(1)
ctx, cancel := context.WithCancel(context.Background())
cancel()
for range 20 {
started := make(chan struct{}, 1)
release := make(chan struct{})
err := limiter.Run(ctx, func() error {
started <- struct{}{}
<-release
return nil
})
if !errors.Is(err, context.Canceled) {
t.Fatalf("Run() error = %v, want context.Canceled", err)
}
select {
case <-started:
close(release)
t.Fatal("execution started with an already canceled context")
case <-time.After(5 * time.Millisecond):
close(release)
}
}
}
func TestExecutionLimiterCapsConfiguredConcurrency(t *testing.T) {
limiter := newExecutionLimiter(maxTemplateExecutions + 1)
if got := cap(limiter.slots); got != maxTemplateExecutions {
t.Fatalf("execution slot capacity = %d, want hard cap %d", got, maxTemplateExecutions)
}
}

View File

@ -0,0 +1,186 @@
package providerapi
import (
"fmt"
"reflect"
"text/template"
"text/template/parse"
)
const (
maxTemplateASTDepth = 32
maxTemplateASTNodes = 4096
maxTemplateRangeDepth = 1
)
func validateTemplateComplexity(parsed *template.Template) error {
nodes := 0
for _, definition := range parsed.Templates() {
if definition.Tree == nil || definition.Tree.Root == nil {
continue
}
if err := walkTemplateNode(definition.Tree.Root, 0, 0, &nodes); err != nil {
return err
}
}
return nil
}
func walkTemplateNode(node parse.Node, depth, rangeDepth int, nodes *int) error {
if node == nil {
return nil
}
value := reflect.ValueOf(node)
if value.Kind() == reflect.Pointer && value.IsNil() {
return nil
}
*nodes++
if *nodes > maxTemplateASTNodes {
return fmt.Errorf("%w: node count exceeds %d", ErrTemplateTooComplex, maxTemplateASTNodes)
}
if depth > maxTemplateASTDepth {
return fmt.Errorf("%w: nesting depth exceeds %d", ErrTemplateTooComplex, maxTemplateASTDepth)
}
walk := func(child parse.Node, childDepth, childRangeDepth int) error {
return walkTemplateNode(child, childDepth, childRangeDepth, nodes)
}
switch current := node.(type) {
case *parse.ListNode:
for _, child := range current.Nodes {
if err := walk(child, depth, rangeDepth); err != nil {
return err
}
}
case *parse.ActionNode:
return walk(current.Pipe, depth, rangeDepth)
case *parse.CommandNode:
for _, argument := range current.Args {
if err := walk(argument, depth, rangeDepth); err != nil {
return err
}
}
case *parse.PipeNode:
for _, declaration := range current.Decl {
if err := walk(declaration, depth, rangeDepth); err != nil {
return err
}
}
for _, command := range current.Cmds {
if err := walk(command, depth, rangeDepth); err != nil {
return err
}
}
case *parse.ChainNode:
return walk(current.Node, depth, rangeDepth)
case *parse.TemplateNode:
return walk(current.Pipe, depth, rangeDepth)
case *parse.IfNode:
if err := walk(current.Pipe, depth, rangeDepth); err != nil {
return err
}
if err := walk(current.List, depth+1, rangeDepth); err != nil {
return err
}
return walk(current.ElseList, depth+1, rangeDepth)
case *parse.RangeNode:
if rangeDepth >= maxTemplateRangeDepth {
return fmt.Errorf("%w: nested range exceeds depth %d", ErrTemplateTooComplex, maxTemplateRangeDepth)
}
if err := walk(current.Pipe, depth, rangeDepth); err != nil {
return err
}
if err := walk(current.List, depth+1, rangeDepth+1); err != nil {
return err
}
return walk(current.ElseList, depth+1, rangeDepth+1)
case *parse.WithNode:
if err := walk(current.Pipe, depth, rangeDepth); err != nil {
return err
}
if err := walk(current.List, depth+1, rangeDepth); err != nil {
return err
}
return walk(current.ElseList, depth+1, rangeDepth)
}
return nil
}
func rejectRecursiveTemplates(parsed *template.Template) error {
graph := make(map[string][]string)
for _, definition := range parsed.Templates() {
if definition.Tree == nil || definition.Tree.Root == nil {
continue
}
var calls []string
collectTemplateCalls(definition.Tree.Root, &calls)
graph[definition.Name()] = calls
}
const (
visiting = 1
visited = 2
)
state := make(map[string]int, len(graph))
var visit func(string) bool
visit = func(name string) bool {
switch state[name] {
case visiting:
return true
case visited:
return false
}
state[name] = visiting
for _, called := range graph[name] {
if visit(called) {
return true
}
}
state[name] = visited
return false
}
for name := range graph {
if visit(name) {
return ErrRecursiveTemplate
}
}
return nil
}
func collectTemplateCalls(node parse.Node, calls *[]string) {
if node == nil {
return
}
switch current := node.(type) {
case *parse.ListNode:
if current == nil {
return
}
for _, child := range current.Nodes {
collectTemplateCalls(child, calls)
}
case *parse.TemplateNode:
if current == nil {
return
}
*calls = append(*calls, current.Name)
case *parse.IfNode:
if current == nil {
return
}
collectTemplateCalls(current.List, calls)
collectTemplateCalls(current.ElseList, calls)
case *parse.RangeNode:
if current == nil {
return
}
collectTemplateCalls(current.List, calls)
collectTemplateCalls(current.ElseList, calls)
case *parse.WithNode:
if current == nil {
return
}
collectTemplateCalls(current.List, calls)
collectTemplateCalls(current.ElseList, calls)
}
}

View File

@ -96,10 +96,11 @@ type Retry struct {
}
type DestinationPolicy struct {
DenyPrivateNetworks bool `yaml:"denyPrivateNetworks"`
DenyLoopback bool `yaml:"denyLoopback"`
DenyLinkLocal bool `yaml:"denyLinkLocal"`
DenyPrivateNetworks *bool `yaml:"denyPrivateNetworks"`
DenyLoopback *bool `yaml:"denyLoopback"`
DenyLinkLocal *bool `yaml:"denyLinkLocal"`
DenyCIDRs []string `yaml:"denyCIDRs"`
AllowedPorts []uint16 `yaml:"allowedPorts"`
}
type ClientIdentification struct {

View File

@ -87,6 +87,30 @@ func TestLoadStrictValidConfiguration(t *testing.T) {
}
}
func TestLoadPreservesOmittedAndExplicitDestinationPolicyBooleans(t *testing.T) {
source := strings.Replace(validConfig,
" auth:\n mode: none\n",
" auth:\n mode: none\n destinationPolicy:\n denyLoopback: false\n",
1,
)
cfg, err := Load(strings.NewReader(source))
if err != nil {
t.Fatalf("Load() error = %v", err)
}
policy := cfg.Gateway.DestinationPolicy
if policy.DenyLoopback == nil || *policy.DenyLoopback {
t.Fatalf("denyLoopback = %v, want explicit false", policy.DenyLoopback)
}
if policy.DenyPrivateNetworks != nil || policy.DenyLinkLocal != nil {
t.Fatalf("omitted policy fields = (%v, %v), want nil", policy.DenyPrivateNetworks, policy.DenyLinkLocal)
}
redacted := cfg.Redacted()
*redacted.Gateway.DestinationPolicy.DenyLoopback = true
if *cfg.Gateway.DestinationPolicy.DenyLoopback {
t.Fatal("Redacted() destination policy aliases source configuration")
}
}
func TestLoadRejectsUnknownFields(t *testing.T) {
_, err := Load(strings.NewReader(validConfig + "unknownField: true\n"))
if err == nil || !strings.Contains(err.Error(), "unknownField") {
@ -266,6 +290,20 @@ func TestValidateRejectsInvalidConfigurationMatrix(t *testing.T) {
},
want: "denyCIDRs",
},
{
name: "invalid destination allowed port",
mutate: func(cfg *Config) {
cfg.Gateway.DestinationPolicy.AllowedPorts = []uint16{443, 0}
},
want: "allowedPorts",
},
{
name: "duplicate destination allowed port",
mutate: func(cfg *Config) {
cfg.Gateway.DestinationPolicy.AllowedPorts = []uint16{443, 443}
},
want: "allowedPorts",
},
{
name: "unsupported routing purpose",
mutate: func(cfg *Config) {

View File

@ -85,10 +85,22 @@ func cloneListener(source Listener) Listener {
cloned.Auth.Methods[index].CIDRs = cloneStrings(source.Auth.Methods[index].CIDRs)
}
cloned.Retry.RetryMethods = cloneStrings(source.Retry.RetryMethods)
cloned.DestinationPolicy.DenyPrivateNetworks = cloneBool(source.DestinationPolicy.DenyPrivateNetworks)
cloned.DestinationPolicy.DenyLoopback = cloneBool(source.DestinationPolicy.DenyLoopback)
cloned.DestinationPolicy.DenyLinkLocal = cloneBool(source.DestinationPolicy.DenyLinkLocal)
cloned.DestinationPolicy.DenyCIDRs = cloneStrings(source.DestinationPolicy.DenyCIDRs)
cloned.DestinationPolicy.AllowedPorts = append([]uint16(nil), source.DestinationPolicy.AllowedPorts...)
return cloned
}
func cloneBool(source *bool) *bool {
if source == nil {
return nil
}
cloned := *source
return &cloned
}
func cloneRouting(source Routing) Routing {
cloned := source
cloned.Match.Methods = cloneStrings(source.Match.Methods)

View File

@ -100,6 +100,16 @@ func validateListener(name string, listener Listener, security Security) error {
if err := validateCIDRs(name+" destinationPolicy.denyCIDRs", listener.DestinationPolicy.DenyCIDRs); err != nil {
return err
}
seenPorts := make(map[uint16]struct{}, len(listener.DestinationPolicy.AllowedPorts))
for _, port := range listener.DestinationPolicy.AllowedPorts {
if port == 0 {
return fmt.Errorf("validate %s destinationPolicy.allowedPorts: port must be greater than zero", name)
}
if _, exists := seenPorts[port]; exists {
return fmt.Errorf("validate %s destinationPolicy.allowedPorts: duplicate port %d", name, port)
}
seenPorts[port] = struct{}{}
}
return nil
}

View File

@ -0,0 +1,301 @@
package policy
import (
"context"
"errors"
"fmt"
"net"
"net/netip"
"net/url"
"strconv"
"strings"
)
var (
ErrInvalidAuthority = errors.New("invalid target authority")
ErrTargetDenied = errors.New("target address is denied by policy")
ErrDNSResolution = errors.New("target dns resolution failed")
)
type Resolver interface {
LookupNetIP(ctx context.Context, host string) ([]netip.Addr, error)
}
type Config struct {
Resolver Resolver
DenyCIDRs []string
AllowedPorts []uint16
AllowPrivateNetworks bool
AllowLoopback bool
AllowLinkLocal bool
}
type TargetPolicy struct {
resolver Resolver
deny CIDRMatcher
allowPrivateNetworks bool
allowLoopback bool
allowLinkLocal bool
allowedPorts map[uint16]struct{}
}
type Authority struct {
Host string
Port uint16
LiteralIP netip.Addr
ResolvedIP netip.Addr
}
func (authority Authority) DialAddress() string {
host := authority.Host
if authority.ResolvedIP.IsValid() {
host = authority.ResolvedIP.Unmap().String()
} else if authority.LiteralIP.IsValid() {
host = authority.LiteralIP.Unmap().String()
}
return net.JoinHostPort(host, strconv.FormatUint(uint64(authority.Port), 10))
}
type CIDRMatcher struct {
prefixes []netip.Prefix
}
func NewTargetPolicy(config Config) (*TargetPolicy, error) {
matcher, err := NewCIDRMatcher(config.DenyCIDRs)
if err != nil {
return nil, err
}
resolver := config.Resolver
if resolver == nil {
resolver = defaultResolver{}
}
ports := append([]uint16(nil), config.AllowedPorts...)
if len(ports) == 0 {
ports = []uint16{80, 443}
}
allowedPorts := make(map[uint16]struct{}, len(ports))
for _, port := range ports {
if port == 0 {
return nil, fmt.Errorf("create target policy: allowed port must be positive")
}
allowedPorts[port] = struct{}{}
}
return &TargetPolicy{
resolver: resolver,
deny: matcher,
allowPrivateNetworks: config.AllowPrivateNetworks,
allowLoopback: config.AllowLoopback,
allowLinkLocal: config.AllowLinkLocal,
allowedPorts: allowedPorts,
}, nil
}
func NewCIDRMatcher(cidrs []string) (CIDRMatcher, error) {
prefixes := make([]netip.Prefix, 0, len(cidrs))
for _, raw := range cidrs {
raw = strings.TrimSpace(raw)
if raw == "" {
continue
}
prefix, err := netip.ParsePrefix(raw)
if err != nil {
return CIDRMatcher{}, fmt.Errorf("parse deny cidr %q: %w", raw, err)
}
prefixes = append(prefixes, prefix.Masked())
}
return CIDRMatcher{prefixes: prefixes}, nil
}
func (m CIDRMatcher) Match(addr netip.Addr) bool {
if !addr.IsValid() {
return false
}
addr = addr.Unmap()
for _, prefix := range m.prefixes {
if prefix.Contains(addr) {
return true
}
}
return false
}
func ParseURLAuthority(raw string) (Authority, error) {
parsed, err := url.Parse(raw)
if err != nil {
return Authority{}, fmt.Errorf("%w: parse url: %v", ErrInvalidAuthority, err)
}
if parsed.Host == "" {
return Authority{}, fmt.Errorf("%w: missing host", ErrInvalidAuthority)
}
port := parsed.Port()
if port == "" {
switch strings.ToLower(parsed.Scheme) {
case "http":
port = "80"
case "https":
port = "443"
default:
return Authority{}, fmt.Errorf("%w: unsupported url scheme %q", ErrInvalidAuthority, parsed.Scheme)
}
}
return parseHostPort(parsed.Hostname(), port)
}
func ParseConnectAuthority(raw string) (Authority, error) {
host, port, err := net.SplitHostPort(raw)
if err != nil {
return Authority{}, fmt.Errorf("%w: %v", ErrInvalidAuthority, err)
}
return parseHostPort(host, port)
}
func (p *TargetPolicy) EvaluateURL(ctx context.Context, raw string) (Authority, error) {
authority, err := ParseURLAuthority(raw)
if err != nil {
return Authority{}, err
}
resolved, err := p.evaluateAuthority(ctx, authority)
authority.ResolvedIP = resolved
return authority, err
}
func (p *TargetPolicy) EvaluateConnectAuthority(ctx context.Context, raw string) (Authority, error) {
authority, err := ParseConnectAuthority(raw)
if err != nil {
return Authority{}, err
}
resolved, err := p.evaluateAuthority(ctx, authority)
authority.ResolvedIP = resolved
return authority, err
}
func (p *TargetPolicy) EvaluateAuthority(ctx context.Context, authority Authority) error {
_, err := p.evaluateAuthority(ctx, authority)
return err
}
func (p *TargetPolicy) evaluateAuthority(ctx context.Context, authority Authority) (netip.Addr, error) {
if p == nil {
return netip.Addr{}, fmt.Errorf("evaluate target: nil policy")
}
if err := ctx.Err(); err != nil {
return netip.Addr{}, err
}
if _, allowed := p.allowedPorts[authority.Port]; !allowed {
return netip.Addr{}, fmt.Errorf("%w: destination port %d is not allowed", ErrTargetDenied, authority.Port)
}
if authority.LiteralIP.IsValid() {
return authority.LiteralIP.Unmap(), p.validateAddress(authority.Host, authority.LiteralIP)
}
addrs, err := p.resolver.LookupNetIP(ctx, authority.Host)
if err != nil {
return netip.Addr{}, err
}
if len(addrs) == 0 {
return netip.Addr{}, fmt.Errorf("%w: no addresses returned for %q", ErrDNSResolution, authority.Host)
}
for _, addr := range addrs {
if err := p.validateAddress(authority.Host, addr); err != nil {
return netip.Addr{}, err
}
}
return addrs[0].Unmap(), nil
}
func (p *TargetPolicy) validateAddress(host string, addr netip.Addr) error {
addr = addr.Unmap()
if p.deny.Match(addr) {
return fmt.Errorf("%w: %s matched deny cidr", ErrTargetDenied, addr)
}
if isSpecialUse(addr) {
return fmt.Errorf("%w: %s is a special-use address", ErrTargetDenied, addr)
}
if p.allowLoopback && addr.IsLoopback() {
return nil
}
if p.allowLinkLocal && addr.IsLinkLocalUnicast() {
return nil
}
if p.allowPrivateNetworks && addr.IsPrivate() {
return nil
}
if isDefaultDenied(addr) {
return fmt.Errorf("%w: %s for host %q", ErrTargetDenied, addr, host)
}
return nil
}
func isSpecialUse(addr netip.Addr) bool {
addr = addr.Unmap()
for _, prefix := range specialUsePrefixes {
if prefix.Contains(addr) {
return true
}
}
return false
}
var specialUsePrefixes = []netip.Prefix{
netip.MustParsePrefix("0.0.0.0/8"),
netip.MustParsePrefix("100.64.0.0/10"),
netip.MustParsePrefix("168.63.129.16/32"),
netip.MustParsePrefix("169.254.169.254/32"),
netip.MustParsePrefix("169.254.170.2/32"),
netip.MustParsePrefix("192.0.0.0/24"),
netip.MustParsePrefix("192.0.2.0/24"),
netip.MustParsePrefix("192.88.99.0/24"),
netip.MustParsePrefix("198.18.0.0/15"),
netip.MustParsePrefix("198.51.100.0/24"),
netip.MustParsePrefix("203.0.113.0/24"),
netip.MustParsePrefix("240.0.0.0/4"),
netip.MustParsePrefix("64:ff9b::/96"),
netip.MustParsePrefix("64:ff9b:1::/48"),
netip.MustParsePrefix("100::/64"),
netip.MustParsePrefix("2001::/23"),
netip.MustParsePrefix("2001:db8::/32"),
netip.MustParsePrefix("2002::/16"),
netip.MustParsePrefix("3fff::/20"),
netip.MustParsePrefix("fd00:ec2::254/128"),
}
func parseHostPort(host, port string) (Authority, error) {
host = strings.TrimSpace(host)
port = strings.TrimSpace(port)
if host == "" || port == "" {
return Authority{}, fmt.Errorf("%w: missing host or port", ErrInvalidAuthority)
}
numericPort, err := strconv.ParseUint(port, 10, 16)
if err != nil || numericPort == 0 {
return Authority{}, fmt.Errorf("%w: invalid port %q", ErrInvalidAuthority, port)
}
authority := Authority{
Host: strings.TrimSuffix(host, "."),
Port: uint16(numericPort),
}
if ip, err := netip.ParseAddr(authority.Host); err == nil {
authority.Host = ip.String()
authority.LiteralIP = ip.Unmap()
}
return authority, nil
}
func isDefaultDenied(addr netip.Addr) bool {
return addr.IsLoopback() ||
addr.IsPrivate() ||
addr.IsLinkLocalUnicast() ||
addr.IsLinkLocalMulticast() ||
addr.IsUnspecified() ||
addr.IsMulticast()
}
type defaultResolver struct{}
func (defaultResolver) LookupNetIP(ctx context.Context, host string) ([]netip.Addr, error) {
addrs, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
if err != nil {
return nil, fmt.Errorf("%w: %v", ErrDNSResolution, err)
}
return addrs, nil
}

View File

@ -0,0 +1,262 @@
package policy
import (
"context"
"errors"
"net/netip"
"reflect"
"testing"
)
func TestParseURLAuthorityUsesDefaultPortAndPreservesHost(t *testing.T) {
authority, err := ParseURLAuthority("https://example.com/path?q=1")
if err != nil {
t.Fatalf("ParseURLAuthority(): %v", err)
}
if authority.Host != "example.com" {
t.Fatalf("host = %q, want example.com", authority.Host)
}
if authority.Port != 443 {
t.Fatalf("port = %d, want 443", authority.Port)
}
if authority.LiteralIP.IsValid() {
t.Fatalf("literal ip = %v, want invalid", authority.LiteralIP)
}
}
func TestParseConnectAuthoritySupportsIPv6Literal(t *testing.T) {
authority, err := ParseConnectAuthority("[2001:db8::1]:8443")
if err != nil {
t.Fatalf("ParseConnectAuthority(): %v", err)
}
if authority.Host != "2001:db8::1" {
t.Fatalf("host = %q, want 2001:db8::1", authority.Host)
}
if authority.Port != 8443 {
t.Fatalf("port = %d, want 8443", authority.Port)
}
if got := authority.LiteralIP.String(); got != "2001:db8::1" {
t.Fatalf("literal ip = %q, want 2001:db8::1", got)
}
}
func TestEvaluateURLRejectsLiteralMetadataAddressWithoutDNSLookup(t *testing.T) {
resolver := &stubResolver{}
policy, err := NewTargetPolicy(Config{Resolver: resolver})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
_, err = policy.EvaluateURL(context.Background(), "http://169.254.169.254/latest/meta-data")
if !errors.Is(err, ErrTargetDenied) {
t.Fatalf("EvaluateURL() error = %v, want ErrTargetDenied", err)
}
if resolver.calls != nil {
t.Fatalf("resolver calls = %v, want nil", resolver.calls)
}
}
func TestEvaluateConnectAuthorityRejectsResolvedMixedResults(t *testing.T) {
resolver := &stubResolver{
results: map[string][]netip.Addr{
"example.com": {
netip.MustParseAddr("93.184.216.34"),
netip.MustParseAddr("10.0.0.8"),
},
},
}
policy, err := NewTargetPolicy(Config{Resolver: resolver})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
_, err = policy.EvaluateConnectAuthority(context.Background(), "example.com:443")
if !errors.Is(err, ErrTargetDenied) {
t.Fatalf("EvaluateConnectAuthority() error = %v, want ErrTargetDenied", err)
}
if got, want := resolver.calls, []string{"example.com"}; !reflect.DeepEqual(got, want) {
t.Fatalf("resolver calls = %v, want %v", got, want)
}
}
func TestEvaluateURLPinsTheValidatedResolutionForTransport(t *testing.T) {
t.Parallel()
policy, err := NewTargetPolicy(Config{Resolver: &stubResolver{results: map[string][]netip.Addr{
"example.test": {
netip.MustParseAddr("8.8.8.8"),
netip.MustParseAddr("2001:4860:4860::8888"),
},
}}})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
authority, err := policy.EvaluateURL(context.Background(), "http://example.test/path")
if err != nil {
t.Fatalf("EvaluateURL() error = %v", err)
}
if got := authority.DialAddress(); got != "8.8.8.8:80" {
t.Fatalf("DialAddress() = %q, want 8.8.8.8:80", got)
}
}
func TestEvaluateAuthorityRejectsReservedAndDisallowedPortsBeforeDNS(t *testing.T) {
t.Parallel()
resolver := &stubResolver{}
targets, err := NewTargetPolicy(Config{Resolver: resolver})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
for _, target := range []string{
"http://100.64.0.1/",
"http://198.51.100.10/",
"http://[2001:db8::1]/",
"http://[3fff::1]/",
"http://168.63.129.16/",
"http://example.test:25/",
} {
if _, err := targets.EvaluateURL(context.Background(), target); !errors.Is(err, ErrTargetDenied) {
t.Fatalf("EvaluateURL(%q) error = %v, want ErrTargetDenied", target, err)
}
}
if resolver.calls != nil {
t.Fatalf("resolver calls = %v, want nil for disallowed port", resolver.calls)
}
}
func TestEvaluateAuthorityAllowsExplicitPort(t *testing.T) {
t.Parallel()
targets, err := NewTargetPolicy(Config{
AllowedPorts: []uint16{8443},
Resolver: &stubResolver{results: map[string][]netip.Addr{
"example.test": {netip.MustParseAddr("8.8.8.8")},
}},
})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
if _, err := targets.EvaluateConnectAuthority(context.Background(), "example.test:8443"); err != nil {
t.Fatalf("EvaluateConnectAuthority() error = %v", err)
}
}
func TestEvaluateConnectAuthorityRejectsIPv6LoopbackLiteral(t *testing.T) {
policy, err := NewTargetPolicy(Config{Resolver: &stubResolver{}})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
_, err = policy.EvaluateConnectAuthority(context.Background(), "[::1]:443")
if !errors.Is(err, ErrTargetDenied) {
t.Fatalf("EvaluateConnectAuthority() error = %v, want ErrTargetDenied", err)
}
}
func TestEvaluateURLAllowsPrivateNetworksWhenExplicitlyEnabled(t *testing.T) {
resolver := &stubResolver{
results: map[string][]netip.Addr{
"internal.example": {
netip.MustParseAddr("10.1.2.3"),
},
},
}
policy, err := NewTargetPolicy(Config{
Resolver: resolver,
AllowPrivateNetworks: true,
})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
authority, err := policy.EvaluateURL(context.Background(), "http://internal.example")
if err != nil {
t.Fatalf("EvaluateURL(): %v", err)
}
if authority.Host != "internal.example" || authority.Port != 80 {
t.Fatalf("authority = %+v, want internal.example:80", authority)
}
}
func TestEvaluateURLCanConfigureLoopbackAndLinkLocalIndependently(t *testing.T) {
t.Parallel()
loopback, err := NewTargetPolicy(Config{AllowLoopback: true})
if err != nil {
t.Fatalf("NewTargetPolicy(loopback): %v", err)
}
if _, err := loopback.EvaluateURL(context.Background(), "http://127.0.0.1/"); err != nil {
t.Fatalf("EvaluateURL(loopback) error = %v", err)
}
linkLocal, err := NewTargetPolicy(Config{AllowLinkLocal: true})
if err != nil {
t.Fatalf("NewTargetPolicy(link-local): %v", err)
}
if _, err := linkLocal.EvaluateURL(context.Background(), "http://169.254.10.20/"); err != nil {
t.Fatalf("EvaluateURL(link-local) error = %v", err)
}
if _, err := linkLocal.EvaluateURL(context.Background(), "http://127.0.0.1/"); !errors.Is(err, ErrTargetDenied) {
t.Fatalf("link-local policy loopback error = %v, want ErrTargetDenied", err)
}
if _, err := linkLocal.EvaluateURL(context.Background(), "http://169.254.169.254/"); !errors.Is(err, ErrTargetDenied) {
t.Fatalf("link-local policy metadata error = %v, want ErrTargetDenied", err)
}
}
func TestEvaluateURLAllowsPrivateWhenEnabledButStillHonorsDenyCIDRs(t *testing.T) {
resolver := &stubResolver{
results: map[string][]netip.Addr{
"internal.example": {
netip.MustParseAddr("10.1.2.3"),
netip.MustParseAddr("169.254.169.254"),
},
},
}
policy, err := NewTargetPolicy(Config{
Resolver: resolver,
AllowPrivateNetworks: true,
DenyCIDRs: []string{"169.254.169.254/32"},
})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
_, err = policy.EvaluateURL(context.Background(), "http://internal.example")
if !errors.Is(err, ErrTargetDenied) {
t.Fatalf("EvaluateURL() error = %v, want ErrTargetDenied", err)
}
}
func TestEvaluateConnectAuthorityPropagatesContextCancellation(t *testing.T) {
resolver := &stubResolver{
err: context.Canceled,
}
policy, err := NewTargetPolicy(Config{Resolver: resolver})
if err != nil {
t.Fatalf("NewTargetPolicy(): %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err = policy.EvaluateConnectAuthority(ctx, "example.com:443")
if !errors.Is(err, context.Canceled) {
t.Fatalf("EvaluateConnectAuthority() error = %v, want context.Canceled", err)
}
}
type stubResolver struct {
results map[string][]netip.Addr
err error
calls []string
}
func (s *stubResolver) LookupNetIP(ctx context.Context, host string) ([]netip.Addr, error) {
s.calls = append(s.calls, host)
if s.err != nil {
return nil, s.err
}
return append([]netip.Addr(nil), s.results[host]...), nil
}

View File

@ -0,0 +1,108 @@
package server
import (
"errors"
"fmt"
"time"
"github.com/proxy-pool/proxy-pool/internal/config"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
platformAdmission "github.com/proxy-pool/proxy-pool/internal/platform/admission"
)
type Protection struct {
Auth Guard
Access Guard
Admission Guard
ClientIPs *ClientIPResolver
}
func BuildProtection(listener config.Listener) (Protection, error) {
clientIPs, err := NewClientIPResolver(listener.Access.TrustedProxies)
if err != nil {
return Protection{}, err
}
access, err := NewAccessGuard(clientIPs, listener.Access.AllowCIDRs)
if err != nil {
return Protection{}, err
}
auth, err := buildConfiguredAuth(listener.Auth, clientIPs)
if err != nil {
return Protection{}, err
}
var admission Guard
if listener.Limits.RequestsPerMinute > 0 || listener.Limits.RequestsPerMinutePerClient > 0 {
limiter, limiterErr := platformAdmission.NewFixedWindow(platformAdmission.FixedWindowConfig{
Window: time.Minute,
Global: listener.Limits.RequestsPerMinute,
PerKey: listener.Limits.RequestsPerMinutePerClient,
})
if limiterErr != nil {
return Protection{}, fmt.Errorf("build gateway admission: %w", limiterErr)
}
admission = NewAdmissionGuard(clientIPs, limiter)
}
return Protection{Auth: auth, Access: access, Admission: admission, ClientIPs: clientIPs}, nil
}
func ConfigFromListener(listener config.Listener) Config {
return Config{
MaxAttempts: listener.Retry.MaxAttempts,
RetryMethods: append([]string(nil), listener.Retry.RetryMethods...),
MaxConcurrentRequests: listener.Limits.MaxConcurrentConnections,
}
}
func TargetPolicyFromListener(listener config.Listener) (*policy.TargetPolicy, error) {
destination := listener.DestinationPolicy
return policy.NewTargetPolicy(policy.Config{
DenyCIDRs: append([]string(nil), destination.DenyCIDRs...),
AllowedPorts: append([]uint16(nil), destination.AllowedPorts...),
AllowPrivateNetworks: explicitlyAllowed(destination.DenyPrivateNetworks),
AllowLoopback: explicitlyAllowed(destination.DenyLoopback),
AllowLinkLocal: explicitlyAllowed(destination.DenyLinkLocal),
})
}
func explicitlyAllowed(deny *bool) bool {
return deny != nil && !*deny
}
func buildConfiguredAuth(auth config.Auth, clientIPs *ClientIPResolver) (Guard, error) {
switch auth.Mode {
case "", "none":
return nil, nil
case "usernamePassword":
return NewBasicAuthGuard(auth.Username, auth.Password), nil
case "apiKey":
return NewAPIKeyGuard(auth.Header, auth.Token), nil
case "ipWhitelist":
return NewAccessGuard(clientIPs, auth.CIDRs)
case "any":
methods := make([]Guard, 0, len(auth.Methods))
for index, method := range auth.Methods {
guard, err := buildConfiguredMethod(method, clientIPs)
if err != nil {
return nil, fmt.Errorf("build gateway auth method %d: %w", index, err)
}
methods = append(methods, guard)
}
return NewAnyGuard(methods...), nil
default:
return nil, fmt.Errorf("build gateway auth: unsupported mode %q", auth.Mode)
}
}
func buildConfiguredMethod(method config.AuthMethod, clientIPs *ClientIPResolver) (Guard, error) {
switch method.Mode {
case "usernamePassword":
return NewBasicAuthGuard(method.Username, method.Password), nil
case "apiKey":
return NewAPIKeyGuard(method.Header, method.Value), nil
case "ipWhitelist":
return NewAccessGuard(clientIPs, method.CIDRs)
default:
return nil, errors.New("unsupported authentication method: " + method.Mode)
}
}

View File

@ -0,0 +1,120 @@
package server
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/proxy-pool/proxy-pool/internal/config"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
)
func TestBuildProtectionFromListenerConfig(t *testing.T) {
t.Parallel()
protection, err := BuildProtection(config.Listener{
Access: config.Access{AllowCIDRs: []string{"198.51.100.0/24"}},
Auth: config.Auth{
Mode: "usernamePassword",
Username: "client",
Password: "secret",
},
Limits: config.Limits{RequestsPerMinute: 10, RequestsPerMinutePerClient: 2},
})
if err != nil {
t.Fatalf("BuildProtection() error = %v", err)
}
request := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
request.RemoteAddr = "198.51.100.8:1234"
request.Header.Set("Proxy-Authorization", "Basic Y2xpZW50OnNlY3JldA==")
for name, guard := range map[string]Guard{
"auth": protection.Auth, "access": protection.Access, "admission": protection.Admission,
} {
if err := guard.Check(context.Background(), request); err != nil {
t.Fatalf("%s guard error = %v", name, err)
}
}
}
func TestTargetPolicyFromOmittedConfigDefaultsToDeny(t *testing.T) {
t.Parallel()
targets, err := TargetPolicyFromListener(config.Listener{})
if err != nil {
t.Fatalf("TargetPolicyFromListener() error = %v", err)
}
if _, err := targets.EvaluateURL(context.Background(), "http://127.0.0.1/"); !errors.Is(err, policy.ErrTargetDenied) {
t.Fatalf("EvaluateURL(loopback) error = %v, want ErrTargetDenied", err)
}
if _, err := targets.EvaluateConnectAuthority(context.Background(), "8.8.8.8:22"); !errors.Is(err, policy.ErrTargetDenied) {
t.Fatalf("EvaluateConnectAuthority(port 22) error = %v, want ErrTargetDenied", err)
}
}
func TestTargetPolicyFromListenerMapsAllowedPorts(t *testing.T) {
t.Parallel()
targets, err := TargetPolicyFromListener(config.Listener{
DestinationPolicy: config.DestinationPolicy{AllowedPorts: []uint16{8443}},
})
if err != nil {
t.Fatalf("TargetPolicyFromListener() error = %v", err)
}
if _, err := targets.EvaluateConnectAuthority(context.Background(), "8.8.8.8:8443"); err != nil {
t.Fatalf("EvaluateConnectAuthority(port 8443) error = %v", err)
}
if _, err := targets.EvaluateConnectAuthority(context.Background(), "8.8.8.8:443"); !errors.Is(err, policy.ErrTargetDenied) {
t.Fatalf("EvaluateConnectAuthority(port 443) error = %v, want ErrTargetDenied", err)
}
}
func TestTargetPolicyFromPartialConfigKeepsOmittedCategoriesDenied(t *testing.T) {
t.Parallel()
deny := true
targets, err := TargetPolicyFromListener(config.Listener{
DestinationPolicy: config.DestinationPolicy{DenyPrivateNetworks: &deny},
})
if err != nil {
t.Fatalf("TargetPolicyFromListener() error = %v", err)
}
for _, target := range []string{"http://127.0.0.1/", "http://169.254.10.20/"} {
if _, err := targets.EvaluateURL(context.Background(), target); !errors.Is(err, policy.ErrTargetDenied) {
t.Fatalf("EvaluateURL(%q) error = %v, want ErrTargetDenied", target, err)
}
}
}
func TestTargetPolicyFromListenerAllowsOnlyExplicitlyFalseCategory(t *testing.T) {
t.Parallel()
allow := false
targets, err := TargetPolicyFromListener(config.Listener{
DestinationPolicy: config.DestinationPolicy{DenyLoopback: &allow},
})
if err != nil {
t.Fatalf("TargetPolicyFromListener() error = %v", err)
}
if _, err := targets.EvaluateURL(context.Background(), "http://127.0.0.1/"); err != nil {
t.Fatalf("EvaluateURL(loopback) error = %v", err)
}
for _, target := range []string{"http://10.0.0.1/", "http://169.254.10.20/"} {
if _, err := targets.EvaluateURL(context.Background(), target); !errors.Is(err, policy.ErrTargetDenied) {
t.Fatalf("EvaluateURL(%q) error = %v, want ErrTargetDenied", target, err)
}
}
}
func TestConfigFromListenerMapsRetryAndConcurrency(t *testing.T) {
t.Parallel()
result := ConfigFromListener(config.Listener{
Limits: config.Limits{MaxConcurrentConnections: 123},
Retry: config.Retry{MaxAttempts: 2, RetryMethods: []string{"GET", "HEAD"}},
})
if result.MaxConcurrentRequests != 123 || result.MaxAttempts != 2 || len(result.RetryMethods) != 2 {
t.Fatalf("handler config = %+v", result)
}
}

View File

@ -0,0 +1,304 @@
package server
import (
"bufio"
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"net/netip"
"net/url"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
"github.com/proxy-pool/proxy-pool/internal/gateway/dispatch"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
"github.com/proxy-pool/proxy-pool/internal/gateway/snapshot"
transportDomain "github.com/proxy-pool/proxy-pool/internal/gateway/transport"
)
func TestHTTPProxyEndToEndRetriesDialFailure(t *testing.T) {
t.Parallel()
requestURI := make(chan string, 1)
goodProxy := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
requestURI <- request.RequestURI
writer.WriteHeader(http.StatusOK)
_, _ = writer.Write([]byte("through-good-proxy"))
}))
defer goodProxy.Close()
closedAddress := reserveClosedAddress(t)
dispatcher, view := dispatcherWithDescriptors(t,
proxyDescriptor(t, "proxy-a", "http://"+closedAddress),
proxyDescriptor(t, "proxy-b", goodProxy.URL),
)
proxyTransport := transportDomain.New(transportDomain.Config{
DialTimeout: 100 * time.Millisecond,
ResponseHeaderTimeout: time.Second,
}, nil)
defer proxyTransport.CloseIdleConnections()
handler := newTestHandler(t, Config{
MaxAttempts: 2,
RetryMethods: []string{http.MethodGet},
}, dispatcher, proxyTransport)
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://TARGET/resource?q=1", nil))
if response.Code != http.StatusOK || response.Body.String() != "through-good-proxy" {
t.Fatalf("response = (%d, %q)", response.Code, response.Body.String())
}
if got := <-requestURI; got != "http://TARGET/resource?q=1" {
t.Fatalf("upstream request URI = %q", got)
}
assertNoLeakedCapacity(t, view)
}
func TestHTTPProxyEndToEndDoesNotRetry407(t *testing.T) {
t.Parallel()
var deniedCalls atomic.Int64
deniedProxy := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
deniedCalls.Add(1)
writer.Header().Set("Proxy-Authenticate", "Basic")
writer.WriteHeader(http.StatusProxyAuthRequired)
_, _ = writer.Write([]byte("denied"))
}))
defer deniedProxy.Close()
var fallbackCalls atomic.Int64
fallbackProxy := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
fallbackCalls.Add(1)
writer.WriteHeader(http.StatusOK)
}))
defer fallbackProxy.Close()
dispatcher, view := dispatcherWithDescriptors(t,
proxyDescriptor(t, "proxy-a", deniedProxy.URL),
proxyDescriptor(t, "proxy-b", fallbackProxy.URL),
)
proxyTransport := transportDomain.New(transportDomain.Config{}, nil)
defer proxyTransport.CloseIdleConnections()
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodGet}}, dispatcher, proxyTransport)
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://TARGET/resource", nil))
if response.Code != http.StatusProxyAuthRequired || response.Body.String() != "denied" {
t.Fatalf("response = (%d, %q), want (407, denied)", response.Code, response.Body.String())
}
if deniedCalls.Load() != 1 || fallbackCalls.Load() != 0 {
t.Fatalf("proxy calls = denied:%d fallback:%d", deniedCalls.Load(), fallbackCalls.Load())
}
assertNoLeakedCapacity(t, view)
}
func TestHTTPProxyEndToEndBlocksPrivateTargetBeforeDispatch(t *testing.T) {
t.Parallel()
targets, err := policy.NewTargetPolicy(policy.Config{Resolver: staticResolver{
addresses: []netip.Addr{netip.MustParseAddr("127.0.0.1")},
}})
if err != nil {
t.Fatalf("NewTargetPolicy() error = %v", err)
}
var dispatchCalls atomic.Int64
handler, err := New(Config{}, Dependencies{
Targets: targets,
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
t.Fatal("routing must not run for a blocked target")
return dispatch.Request{}, nil
}),
Dispatcher: DispatcherFunc(func(dispatch.Request) (*dispatch.Lease, error) {
dispatchCalls.Add(1)
return nil, dispatch.ErrNoCandidate
}),
Transport: &fakeTransport{},
})
if err != nil {
t.Fatalf("New() error = %v", err)
}
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://127.0.0.1/admin", nil))
if response.Code != http.StatusForbidden {
t.Fatalf("status = %d, want 403", response.Code)
}
if dispatchCalls.Load() != 0 {
t.Fatalf("dispatch calls = %d, want 0", dispatchCalls.Load())
}
}
func TestCONNECTEndToEndRelaysDataAndHalfClose(t *testing.T) {
t.Parallel()
upstream, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen fake CONNECT proxy: %v", err)
}
defer upstream.Close()
upstreamDone := make(chan error, 1)
go func() {
connection, acceptErr := upstream.Accept()
if acceptErr != nil {
upstreamDone <- acceptErr
return
}
defer connection.Close()
request, readErr := http.ReadRequest(bufio.NewReader(connection))
if readErr != nil {
upstreamDone <- readErr
return
}
if request.Method != http.MethodConnect || request.Host != "example.test:443" {
upstreamDone <- errors.New("unexpected upstream CONNECT target: " + request.Host)
return
}
if _, writeErr := io.WriteString(connection, "HTTP/1.1 200 Connection Established\r\n\r\n"); writeErr != nil {
upstreamDone <- writeErr
return
}
buffer := make([]byte, 1024)
for {
count, readErr := connection.Read(buffer)
if count > 0 {
if _, writeErr := connection.Write(buffer[:count]); writeErr != nil {
upstreamDone <- writeErr
return
}
}
if readErr != nil {
if readErr == io.EOF {
upstreamDone <- nil
} else {
upstreamDone <- readErr
}
return
}
}
}()
selected := proxyDescriptor(t, "proxy-a", "http://"+upstream.Addr().String())
dispatcher, view := dispatcherWithDescriptors(t, selected)
proxyTransport := transportDomain.New(transportDomain.Config{TunnelIdleTimeout: time.Second}, nil)
handler := newTestHandler(t, Config{MaxAttempts: 1}, dispatcher, proxyTransport)
gateway := httptest.NewServer(handler)
defer gateway.Close()
parsedGateway, _ := url.Parse(gateway.URL)
client, err := net.DialTCP("tcp", nil, tcpAddress(t, parsedGateway.Host))
if err != nil {
t.Fatalf("dial gateway: %v", err)
}
defer client.Close()
if _, err := io.WriteString(client, "CONNECT example.test:443 HTTP/1.1\r\nHost: example.test:443\r\n\r\n"); err != nil {
t.Fatalf("write client CONNECT: %v", err)
}
response, err := http.ReadResponse(bufio.NewReader(client), &http.Request{Method: http.MethodConnect})
if err != nil {
t.Fatalf("read gateway CONNECT response: %v", err)
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("CONNECT status = %d, want 200", response.StatusCode)
}
if _, err := io.WriteString(client, "ping"); err != nil {
t.Fatalf("write tunnel data: %v", err)
}
echo := make([]byte, 4)
if _, err := io.ReadFull(response.Body, echo); err != nil {
t.Fatalf("read tunnel echo: %v", err)
}
if string(echo) != "ping" {
t.Fatalf("tunnel echo = %q", echo)
}
if err := client.CloseWrite(); err != nil {
t.Fatalf("half-close client tunnel: %v", err)
}
if _, err := io.Copy(io.Discard, response.Body); err != nil {
t.Fatalf("drain tunnel after half-close: %v", err)
}
if err := <-upstreamDone; err != nil {
t.Fatalf("fake upstream: %v", err)
}
assertNoLeakedCapacity(t, view)
}
type staticResolver struct{ addresses []netip.Addr }
func (resolver staticResolver) LookupNetIP(context.Context, string) ([]netip.Addr, error) {
return append([]netip.Addr(nil), resolver.addresses...), nil
}
func dispatcherWithDescriptors(t *testing.T, proxies ...proxyDomain.Proxy) (*dispatch.Dispatcher, *snapshot.View) {
t.Helper()
store := snapshot.NewStore("cluster-e2e", "worker-e2e")
envelope := snapshot.Envelope{
ClusterID: "cluster-e2e",
WorkerID: "worker-e2e",
Epoch: 1,
Version: 1,
Full: true,
Proxies: proxies,
}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("apply snapshot: %v", err)
}
return dispatch.New(store), store.Current()
}
func proxyDescriptor(t *testing.T, id, rawURL string) proxyDomain.Proxy {
t.Helper()
parsed, err := url.Parse(rawURL)
if err != nil {
t.Fatalf("parse proxy URL: %v", err)
}
host, portText, err := net.SplitHostPort(parsed.Host)
if err != nil {
t.Fatalf("split proxy address: %v", err)
}
port, err := strconv.ParseUint(portText, 10, 16)
if err != nil {
t.Fatalf("parse proxy port: %v", err)
}
return proxyDomain.Proxy{
ID: id,
Scheme: proxyDomain.Scheme(strings.ToLower(parsed.Scheme)),
Host: host,
Port: uint16(port),
SourceUpstream: "provider-a",
MaxConcurrency: 2,
State: proxyDomain.StateAvailable,
CredentialVersion: "v1",
}
}
func reserveClosedAddress(t *testing.T) string {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("reserve address: %v", err)
}
address := listener.Addr().String()
if err := listener.Close(); err != nil {
t.Fatalf("close reserved address: %v", err)
}
return address
}
func tcpAddress(t *testing.T, address string) *net.TCPAddr {
t.Helper()
resolved, err := net.ResolveTCPAddr("tcp", address)
if err != nil {
t.Fatalf("resolve TCP address: %v", err)
}
return resolved
}

View File

@ -0,0 +1,569 @@
package server
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"slices"
"strings"
"sync"
"sync/atomic"
"time"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
"github.com/proxy-pool/proxy-pool/internal/gateway/dispatch"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
transportDomain "github.com/proxy-pool/proxy-pool/internal/gateway/transport"
)
type Config struct {
MaxAttempts int
RetryMethods []string
SafetyMargin time.Duration
CopyBufferSize int
MaxConcurrentRequests int
}
const requestClosingBit uint64 = 1 << 63
type Guard interface {
Check(context.Context, *http.Request) error
}
type GuardFunc func(context.Context, *http.Request) error
func (guard GuardFunc) Check(ctx context.Context, request *http.Request) error {
return guard(ctx, request)
}
type TargetPolicy interface {
EvaluateURL(context.Context, string) (policy.Authority, error)
EvaluateConnectAuthority(context.Context, string) (policy.Authority, error)
}
type Router interface {
Route(*http.Request) (dispatch.Request, error)
}
type RouteFunc func(*http.Request) (dispatch.Request, error)
func (route RouteFunc) Route(request *http.Request) (dispatch.Request, error) {
return route(request)
}
type Dispatcher interface {
Acquire(dispatch.Request) (*dispatch.Lease, error)
}
type DispatcherFunc func(dispatch.Request) (*dispatch.Lease, error)
func (acquire DispatcherFunc) Acquire(request dispatch.Request) (*dispatch.Lease, error) {
return acquire(request)
}
type ProxyTransport interface {
RoundTrip(context.Context, proxyDomain.Proxy, *http.Request, ...func() error) (*http.Response, error)
OpenTunnel(context.Context, proxyDomain.Proxy, string) (net.Conn, error)
Relay(context.Context, net.Conn, net.Conn) error
}
type Dependencies struct {
Auth Guard
Access Guard
Admission Guard
Targets TargetPolicy
Router Router
Dispatcher Dispatcher
Transport ProxyTransport
}
type Handler struct {
config Config
guards [3]Guard
targets TargetPolicy
router Router
dispatcher Dispatcher
transport ProxyTransport
buffers sync.Pool
inFlight chan struct{}
forceClose atomic.Bool
requestState atomic.Uint64
shutdownStart sync.Once
drainOnce sync.Once
shutdownDone chan struct{}
tunnelMu sync.Mutex
tunnels map[*activeTunnel]struct{}
}
func New(config Config, dependencies Dependencies) (*Handler, error) {
if dependencies.Targets == nil {
return nil, errors.New("create gateway handler: target policy is required")
}
if dependencies.Router == nil {
return nil, errors.New("create gateway handler: router is required")
}
if dependencies.Dispatcher == nil {
return nil, errors.New("create gateway handler: dispatcher is required")
}
if dependencies.Transport == nil {
return nil, errors.New("create gateway handler: transport is required")
}
if config.MaxAttempts <= 0 {
config.MaxAttempts = 1
}
if config.CopyBufferSize <= 0 {
config.CopyBufferSize = 32 << 10
}
handler := &Handler{
config: config,
guards: [3]Guard{dependencies.Auth, dependencies.Access, dependencies.Admission},
targets: dependencies.Targets,
router: dependencies.Router,
dispatcher: dependencies.Dispatcher,
transport: dependencies.Transport,
tunnels: make(map[*activeTunnel]struct{}),
shutdownDone: make(chan struct{}),
}
if config.MaxConcurrentRequests > 0 {
handler.inFlight = make(chan struct{}, config.MaxConcurrentRequests)
}
handler.buffers.New = func() any { return make([]byte, config.CopyBufferSize) }
return handler, nil
}
func (handler *Handler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
if !handler.beginRequest() {
http.Error(writer, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
return
}
defer handler.finishRequest()
if handler.inFlight != nil {
select {
case handler.inFlight <- struct{}{}:
defer func() { <-handler.inFlight }()
default:
http.Error(writer, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
return
}
}
for _, guard := range handler.guards {
if guard == nil {
continue
}
if err := guard.Check(request.Context(), request); err != nil {
writeGatewayError(writer, err)
return
}
}
if request.Method == http.MethodConnect {
target, err := handler.targets.EvaluateConnectAuthority(request.Context(), request.Host)
if err != nil {
writeGatewayError(writer, err)
return
}
route, err := handler.router.Route(request)
if err != nil {
writeGatewayError(writer, fmt.Errorf("route gateway CONNECT: %w", err))
return
}
handler.connect(writer, request, target, route)
return
}
if request.URL == nil || !request.URL.IsAbs() {
http.Error(writer, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)
return
}
target, err := handler.targets.EvaluateURL(request.Context(), request.URL.String())
if err != nil {
writeGatewayError(writer, err)
return
}
route, err := handler.router.Route(request)
if err != nil {
writeGatewayError(writer, fmt.Errorf("route gateway request: %w", err))
return
}
handler.forwardHTTP(writer, request, target, route)
}
func (handler *Handler) connect(
writer http.ResponseWriter,
request *http.Request,
target policy.Authority,
route dispatch.Request,
) {
attempts := handler.attemptLimit(request)
excluded := cloneSet(route.Exclude)
var lastErr error
for attempt := 0; attempt < attempts; attempt++ {
route.Now = time.Now().UTC()
route.Exclude = excluded
route.SafetyMargin = handler.config.SafetyMargin
lease, err := handler.dispatcher.Acquire(route)
if err != nil {
lastErr = err
break
}
upstream, err := handler.transport.OpenTunnel(request.Context(), lease.Proxy, target.DialAddress())
if err != nil {
finishLease(lease, false)
excluded[lease.Proxy.ID] = struct{}{}
var responseError *transportDomain.ProxyResponseError
if errors.As(err, &responseError) {
lastErr = responseError
if responseError.Retryable() && attempt+1 < attempts {
continue
}
handler.writeConnectError(writer, responseError)
return
}
lastErr = err
continue
}
if err := lease.Commit(); err != nil {
_ = upstream.Close()
finishLease(lease, false)
lastErr = err
break
}
handler.serveTunnel(writer, request, lease, upstream)
return
}
writeGatewayError(writer, fmt.Errorf("establish gateway CONNECT: %w", lastErr))
}
func (handler *Handler) writeConnectError(writer http.ResponseWriter, responseError *transportDomain.ProxyResponseError) {
header := responseError.Header.Clone()
removeHopByHop(header)
copyHeaders(writer.Header(), header)
writer.WriteHeader(responseError.StatusCode)
_, _ = writer.Write(responseError.Body)
}
func (handler *Handler) serveTunnel(
writer http.ResponseWriter,
request *http.Request,
lease *dispatch.Lease,
upstream net.Conn,
) {
defer finishLease(lease, true)
defer upstream.Close()
hijacker, ok := writer.(http.Hijacker)
if !ok {
http.Error(writer, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
return
}
client, readWriter, err := hijacker.Hijack()
if err != nil {
http.Error(writer, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
return
}
defer client.Close()
tunnel := &activeTunnel{client: client, upstream: upstream}
if !handler.registerTunnel(tunnel) {
return
}
defer handler.unregisterTunnel(tunnel)
if _, err := readWriter.WriteString("HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil {
return
}
if err := readWriter.Flush(); err != nil {
return
}
bufferedClient := &bufferedClientConn{Conn: client, reader: readWriter.Reader}
_ = handler.transport.Relay(request.Context(), bufferedClient, upstream)
}
func (handler *Handler) Shutdown(ctx context.Context) error {
if handler == nil {
return nil
}
handler.shutdownStart.Do(func() {
for {
state := handler.requestState.Load()
if handler.requestState.CompareAndSwap(state, state|requestClosingBit) {
if state == 0 {
handler.signalDrained()
}
break
}
}
if closer, ok := handler.transport.(interface{ CloseIdleConnections() }); ok {
closer.CloseIdleConnections()
}
})
select {
case <-handler.shutdownDone:
return nil
case <-ctx.Done():
handler.forceClose.Store(true)
handler.closeActiveTunnels()
return ctx.Err()
}
}
func (handler *Handler) beginRequest() bool {
for {
state := handler.requestState.Load()
if state&requestClosingBit != 0 || state == requestClosingBit-1 {
return false
}
if handler.requestState.CompareAndSwap(state, state+1) {
return true
}
}
}
func (handler *Handler) finishRequest() {
state := handler.requestState.Add(^uint64(0))
if state == requestClosingBit {
handler.signalDrained()
}
}
func (handler *Handler) signalDrained() {
handler.drainOnce.Do(func() { close(handler.shutdownDone) })
}
func (handler *Handler) closeActiveTunnels() {
handler.tunnelMu.Lock()
tunnels := make([]*activeTunnel, 0, len(handler.tunnels))
for tunnel := range handler.tunnels {
tunnels = append(tunnels, tunnel)
}
handler.tunnelMu.Unlock()
for _, tunnel := range tunnels {
_ = tunnel.client.Close()
_ = tunnel.upstream.Close()
}
}
func (handler *Handler) registerTunnel(tunnel *activeTunnel) bool {
handler.tunnelMu.Lock()
defer handler.tunnelMu.Unlock()
if handler.forceClose.Load() {
_ = tunnel.client.Close()
_ = tunnel.upstream.Close()
return false
}
handler.tunnels[tunnel] = struct{}{}
return true
}
func (handler *Handler) unregisterTunnel(tunnel *activeTunnel) {
handler.tunnelMu.Lock()
delete(handler.tunnels, tunnel)
handler.tunnelMu.Unlock()
}
func (handler *Handler) forwardHTTP(
writer http.ResponseWriter,
request *http.Request,
target policy.Authority,
route dispatch.Request,
) {
attempts := handler.attemptLimit(request)
excluded := cloneSet(route.Exclude)
var lastErr error
for attempt := 0; attempt < attempts; attempt++ {
attemptRequest, err := requestForAttempt(request, attempt)
if err != nil {
lastErr = err
break
}
pinHTTPDestination(attemptRequest, target)
removeHopByHop(attemptRequest.Header)
route.Now = time.Now().UTC()
route.Exclude = excluded
route.SafetyMargin = handler.config.SafetyMargin
lease, err := handler.dispatcher.Acquire(route)
if err != nil {
lastErr = err
break
}
var committed atomic.Bool
commit := func() error {
if err := lease.Commit(); err != nil {
return err
}
committed.Store(true)
return nil
}
response, err := handler.transport.RoundTrip(request.Context(), lease.Proxy, attemptRequest, commit)
if err != nil {
finishLease(lease, committed.Load())
excluded[lease.Proxy.ID] = struct{}{}
lastErr = err
continue
}
if !committed.Load() {
if err := commit(); err != nil {
_ = response.Body.Close()
finishLease(lease, false)
lastErr = err
break
}
}
handler.writeResponse(writer, response)
finishLease(lease, true)
return
}
writeGatewayError(writer, fmt.Errorf("forward gateway request: %w", lastErr))
}
func pinHTTPDestination(request *http.Request, target policy.Authority) {
if request == nil || request.URL == nil || (!target.ResolvedIP.IsValid() && !target.LiteralIP.IsValid()) {
return
}
request.URL.Host = target.DialAddress()
}
func (handler *Handler) writeResponse(writer http.ResponseWriter, response *http.Response) {
defer response.Body.Close()
removeHopByHop(response.Header)
copyHeaders(writer.Header(), response.Header)
writer.WriteHeader(response.StatusCode)
buffer := handler.buffers.Get().([]byte)
defer handler.buffers.Put(buffer)
_, _ = io.CopyBuffer(writer, response.Body, buffer)
}
func (handler *Handler) attemptLimit(request *http.Request) int {
if handler.config.MaxAttempts <= 1 || !containsMethod(handler.config.RetryMethods, request.Method) {
return 1
}
if request.Body != nil && request.Body != http.NoBody && request.GetBody == nil {
return 1
}
return handler.config.MaxAttempts
}
func requestForAttempt(request *http.Request, attempt int) (*http.Request, error) {
clone := request.Clone(request.Context())
clone.RequestURI = ""
clone.Header = request.Header.Clone()
if attempt == 0 || request.Body == nil || request.Body == http.NoBody {
return clone, nil
}
body, err := request.GetBody()
if err != nil {
return nil, fmt.Errorf("replay gateway request body: %w", err)
}
clone.Body = body
return clone, nil
}
func finishLease(lease *dispatch.Lease, committed bool) {
if committed {
_ = lease.Release()
return
}
_ = lease.Cancel()
}
func containsMethod(methods []string, method string) bool {
return slices.ContainsFunc(methods, func(candidate string) bool {
return strings.EqualFold(candidate, method)
})
}
func cloneSet(source map[string]struct{}) map[string]struct{} {
cloned := make(map[string]struct{}, len(source)+1)
for key := range source {
cloned[key] = struct{}{}
}
return cloned
}
func copyHeaders(destination, source http.Header) {
for name, values := range source {
for _, value := range values {
destination.Add(name, value)
}
}
}
func removeHopByHop(header http.Header) {
for _, value := range header.Values("Connection") {
for name := range strings.SplitSeq(value, ",") {
header.Del(strings.TrimSpace(name))
}
}
for _, name := range []string{
"Connection",
"Keep-Alive",
"Proxy-Connection",
"TE",
"Trailer",
"Transfer-Encoding",
"Upgrade",
} {
header.Del(name)
}
}
func writeGatewayError(writer http.ResponseWriter, err error) {
status := http.StatusBadGateway
var httpError *HTTPError
if errors.As(err, &httpError) {
status = httpError.StatusCode
copyHeaders(writer.Header(), httpError.Header)
} else {
switch {
case errors.Is(err, policy.ErrInvalidAuthority):
status = http.StatusBadRequest
case errors.Is(err, policy.ErrTargetDenied):
status = http.StatusForbidden
case errors.Is(err, ErrRouteRejected):
status = http.StatusForbidden
case errors.Is(err, dispatch.ErrNoCandidate), errors.Is(err, ErrRouteNotFound):
status = http.StatusServiceUnavailable
case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded):
status = http.StatusGatewayTimeout
}
}
http.Error(writer, http.StatusText(status), status)
}
type bufferedClientConn struct {
net.Conn
reader *bufio.Reader
}
type activeTunnel struct {
client net.Conn
upstream net.Conn
}
func (connection *bufferedClientConn) Read(buffer []byte) (int, error) {
return connection.reader.Read(buffer)
}
func (connection *bufferedClientConn) CloseWrite() error {
if halfCloser, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
return halfCloser.CloseWrite()
}
return connection.Close()
}
func (connection *bufferedClientConn) CloseRead() error {
if halfCloser, ok := connection.Conn.(interface{ CloseRead() error }); ok {
return halfCloser.CloseRead()
}
return nil
}

View File

@ -0,0 +1,593 @@
package server
import (
"bufio"
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"net/netip"
"net/url"
"strings"
"sync"
"testing"
"time"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
"github.com/proxy-pool/proxy-pool/internal/gateway/dispatch"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
"github.com/proxy-pool/proxy-pool/internal/gateway/snapshot"
transportDomain "github.com/proxy-pool/proxy-pool/internal/gateway/transport"
)
func TestHandlerRunsProtectionAndTargetPolicyBeforeRouting(t *testing.T) {
t.Parallel()
var calls []string
var mu sync.Mutex
record := func(name string) {
mu.Lock()
calls = append(calls, name)
mu.Unlock()
}
handler, err := New(Config{MaxAttempts: 1}, Dependencies{
Auth: GuardFunc(func(context.Context, *http.Request) error { record("auth"); return nil }),
Access: GuardFunc(func(context.Context, *http.Request) error { record("access"); return nil }),
Admission: GuardFunc(func(context.Context, *http.Request) error { record("admission"); return nil }),
Targets: fakeTargets{evaluateURL: func(context.Context, string) (policy.Authority, error) {
record("target")
return policy.Authority{Host: "example.test", Port: 80}, nil
}},
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
record("route")
return dispatch.Request{}, errors.New("route stopped")
}),
Dispatcher: DispatcherFunc(func(dispatch.Request) (*dispatch.Lease, error) {
t.Fatal("dispatcher must not run after route error")
return nil, nil
}),
Transport: &fakeTransport{},
})
if err != nil {
t.Fatalf("New() error = %v", err)
}
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://example.test/resource", nil))
if got, want := strings.Join(calls, ","), "auth,access,admission,target,route"; got != want {
t.Fatalf("pipeline order = %q, want %q", got, want)
}
if response.Code != http.StatusBadGateway {
t.Fatalf("status = %d, want 502", response.Code)
}
}
func TestHandlerRetriesGETWithAnotherProxyBeforeResponseCommit(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{}
transport.roundTrip = func(
_ context.Context,
selected proxyDomain.Proxy,
_ *http.Request,
commit ...func() error,
) (*http.Response, error) {
if selected.ID == "proxy-a" {
return nil, errors.New("dial failed")
}
if err := commit[0](); err != nil {
return nil, err
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader("ok")),
}, nil
}
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodGet}}, dispatcher, transport)
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://example.test/resource", nil))
if response.Code != http.StatusOK || response.Body.String() != "ok" {
t.Fatalf("response = (%d, %q), want (200, ok)", response.Code, response.Body.String())
}
if got := strings.Join(transport.attempts(), ","); got != "proxy-a,proxy-b" {
t.Fatalf("proxy attempts = %q", got)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerDoesNotRetryPOST(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{roundTrip: func(
_ context.Context,
selected proxyDomain.Proxy,
_ *http.Request,
_ ...func() error,
) (*http.Response, error) {
return nil, errors.New("dial failed through " + selected.ID)
}}
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodGet}}, dispatcher, transport)
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodPost, "http://example.test/resource", strings.NewReader("body")))
if response.Code != http.StatusBadGateway {
t.Fatalf("status = %d, want 502", response.Code)
}
if len(transport.attempts()) != 1 {
t.Fatalf("attempts = %v, want one POST attempt", transport.attempts())
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerReturns407WithoutRetry(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
_ *http.Request,
commit ...func() error,
) (*http.Response, error) {
if err := commit[0](); err != nil {
return nil, err
}
return &http.Response{
StatusCode: http.StatusProxyAuthRequired,
Header: http.Header{"Proxy-Authenticate": []string{"Basic"}},
Body: io.NopCloser(strings.NewReader("denied")),
}, nil
}}
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodGet}}, dispatcher, transport)
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://example.test/resource", nil))
if response.Code != http.StatusProxyAuthRequired || response.Body.String() != "denied" {
t.Fatalf("response = (%d, %q), want (407, denied)", response.Code, response.Body.String())
}
if len(transport.attempts()) != 1 {
t.Fatalf("attempts = %v, want no retry for 407", transport.attempts())
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerPinsValidatedHTTPDestinationWithoutChangingHost(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a")
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
request *http.Request,
commit ...func() error,
) (*http.Response, error) {
if request.URL.Host != "198.51.100.10:80" {
t.Fatalf("pinned URL host = %q", request.URL.Host)
}
if request.Host != "example.test" {
t.Fatalf("original Host = %q", request.Host)
}
if err := commit[0](); err != nil {
return nil, err
}
return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil
}}
handler, err := New(Config{}, Dependencies{
Targets: fakeTargets{evaluateURL: func(context.Context, string) (policy.Authority, error) {
return policy.Authority{
Host: "example.test", Port: 80,
ResolvedIP: netip.MustParseAddr("198.51.100.10"),
}, nil
}},
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
return dispatch.Request{Upstreams: []string{"provider-a"}}, nil
}),
Dispatcher: dispatcher,
Transport: transport,
})
if err != nil {
t.Fatalf("New() error = %v", err)
}
response := httptest.NewRecorder()
handler.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "http://example.test/resource", nil))
if response.Code != http.StatusNoContent {
t.Fatalf("status = %d, want 204", response.Code)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerEnforcesConcurrentRequestLimit(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a")
started := make(chan struct{})
release := make(chan struct{})
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
_ *http.Request,
commit ...func() error,
) (*http.Response, error) {
if err := commit[0](); err != nil {
return nil, err
}
close(started)
<-release
return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil
}}
handler := newTestHandler(t, Config{MaxAttempts: 1, MaxConcurrentRequests: 1}, dispatcher, transport)
first := httptest.NewRecorder()
firstDone := make(chan struct{})
go func() {
handler.ServeHTTP(first, httptest.NewRequest(http.MethodGet, "http://example.test/first", nil))
close(firstDone)
}()
<-started
second := httptest.NewRecorder()
handler.ServeHTTP(second, httptest.NewRequest(http.MethodGet, "http://example.test/second", nil))
if second.Code != http.StatusServiceUnavailable {
t.Fatalf("second status = %d, want 503", second.Code)
}
close(release)
<-firstDone
assertNoLeakedCapacity(t, view)
}
func TestHandlerRetriesCONNECTOnlyBeforeClientSuccess(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{}
transport.openTunnel = func(_ context.Context, selected proxyDomain.Proxy, target string) (net.Conn, error) {
if target != "example.test:443" {
t.Fatalf("CONNECT target = %q", target)
}
if selected.ID == "proxy-a" {
return nil, errors.New("handshake failed")
}
gateway, peer := net.Pipe()
_ = peer.Close()
return gateway, nil
}
transport.relay = func(context.Context, net.Conn, net.Conn) error {
return errors.New("tunnel broke after client commit")
}
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodConnect}}, dispatcher, transport)
gateway := httptest.NewServer(handler)
defer gateway.Close()
response := sendConnect(t, gateway.URL, "example.test:443")
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("CONNECT status = %d, want 200", response.StatusCode)
}
if got := strings.Join(transport.attempts(), ","); got != "proxy-a,proxy-b" {
t.Fatalf("CONNECT proxy attempts = %q", got)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerReturnsCONNECT407WithoutRetry(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{openTunnel: func(context.Context, proxyDomain.Proxy, string) (net.Conn, error) {
return nil, &transportDomain.ProxyResponseError{
StatusCode: http.StatusProxyAuthRequired,
Status: "407 Proxy Authentication Required",
Header: http.Header{"Proxy-Authenticate": []string{"Basic"}},
Body: []byte("denied"),
}
}}
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodConnect}}, dispatcher, transport)
response := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodConnect, "http://example.test:443", nil)
request.Host = "example.test:443"
handler.ServeHTTP(response, request)
if response.Code != http.StatusProxyAuthRequired || strings.TrimSpace(response.Body.String()) != "denied" {
t.Fatalf("CONNECT response = (%d, %q), want (407, denied)", response.Code, response.Body.String())
}
if len(transport.attempts()) != 1 {
t.Fatalf("CONNECT attempts = %v, want no retry for 407", transport.attempts())
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerRetriesRetryableCONNECTResponseBeforeClientSuccess(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a", "proxy-b")
transport := &fakeTransport{openTunnel: func(_ context.Context, selected proxyDomain.Proxy, _ string) (net.Conn, error) {
if selected.ID == "proxy-a" {
return nil, &transportDomain.ProxyResponseError{
StatusCode: http.StatusServiceUnavailable,
Status: "503 Service Unavailable",
Header: make(http.Header),
Body: []byte("retry"),
}
}
gateway, peer := net.Pipe()
_ = peer.Close()
return gateway, nil
}}
transport.relay = func(context.Context, net.Conn, net.Conn) error { return nil }
handler := newTestHandler(t, Config{MaxAttempts: 2, RetryMethods: []string{http.MethodConnect}}, dispatcher, transport)
gateway := httptest.NewServer(handler)
defer gateway.Close()
response := sendConnect(t, gateway.URL, "example.test:443")
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
t.Fatalf("CONNECT status = %d, want 200", response.StatusCode)
}
if got := strings.Join(transport.attempts(), ","); got != "proxy-a,proxy-b" {
t.Fatalf("CONNECT attempts = %q, want proxy-a,proxy-b", got)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerShutdownForcesHijackedTunnelsAtDeadlineAndRejectsNewRequests(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a")
peerConnections := make(chan net.Conn, 1)
relayStarted := make(chan struct{})
transport := &fakeTransport{}
transport.openTunnel = func(context.Context, proxyDomain.Proxy, string) (net.Conn, error) {
gateway, peer := net.Pipe()
peerConnections <- peer
return gateway, nil
}
transport.relay = func(_ context.Context, client, _ net.Conn) error {
close(relayStarted)
_, err := io.Copy(io.Discard, client)
return err
}
handler := newTestHandler(t, Config{MaxAttempts: 1}, dispatcher, transport)
gateway := httptest.NewServer(handler)
defer gateway.Close()
response := sendConnect(t, gateway.URL, "example.test:443")
defer response.Body.Close()
peer := <-peerConnections
defer peer.Close()
<-relayStarted
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
if err := handler.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Shutdown() error = %v, want context deadline exceeded", err)
}
rejected := httptest.NewRecorder()
handler.ServeHTTP(rejected, httptest.NewRequest(http.MethodGet, "http://example.test/new", nil))
if rejected.Code != http.StatusServiceUnavailable {
t.Fatalf("request after shutdown status = %d, want 503", rejected.Code)
}
assertNoLeakedCapacity(t, view)
}
func TestHandlerShutdownWaitsForInFlightHTTP(t *testing.T) {
t.Parallel()
dispatcher, view := dispatcherWithProxies(t, "proxy-a")
started := make(chan struct{})
release := make(chan struct{})
transport := &fakeTransport{roundTrip: func(
_ context.Context,
_ proxyDomain.Proxy,
_ *http.Request,
commit ...func() error,
) (*http.Response, error) {
if err := commit[0](); err != nil {
return nil, err
}
close(started)
<-release
return &http.Response{StatusCode: http.StatusNoContent, Header: make(http.Header), Body: http.NoBody}, nil
}}
handler := newTestHandler(t, Config{}, dispatcher, transport)
requestDone := make(chan struct{})
go func() {
handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "http://example.test/resource", nil))
close(requestDone)
}()
<-started
shutdownDone := make(chan error, 1)
go func() { shutdownDone <- handler.Shutdown(context.Background()) }()
select {
case err := <-shutdownDone:
t.Fatalf("Shutdown() returned before HTTP completed: %v", err)
case <-time.After(20 * time.Millisecond):
}
close(release)
<-requestDone
if err := <-shutdownDone; err != nil {
t.Fatalf("Shutdown() error = %v", err)
}
assertNoLeakedCapacity(t, view)
}
type fakeTargets struct {
evaluateURL func(context.Context, string) (policy.Authority, error)
evaluateConnect func(context.Context, string) (policy.Authority, error)
}
func (targets fakeTargets) EvaluateURL(ctx context.Context, raw string) (policy.Authority, error) {
if targets.evaluateURL != nil {
return targets.evaluateURL(ctx, raw)
}
return policy.Authority{Host: "example.test", Port: 80}, nil
}
func (targets fakeTargets) EvaluateConnectAuthority(ctx context.Context, raw string) (policy.Authority, error) {
if targets.evaluateConnect != nil {
return targets.evaluateConnect(ctx, raw)
}
return policy.Authority{Host: "example.test", Port: 443}, nil
}
type fakeTransport struct {
mu sync.Mutex
seen []string
roundTrip func(context.Context, proxyDomain.Proxy, *http.Request, ...func() error) (*http.Response, error)
openTunnel func(context.Context, proxyDomain.Proxy, string) (net.Conn, error)
relay func(context.Context, net.Conn, net.Conn) error
}
func (transport *fakeTransport) RoundTrip(
ctx context.Context,
selected proxyDomain.Proxy,
request *http.Request,
commit ...func() error,
) (*http.Response, error) {
transport.record(selected.ID)
if transport.roundTrip == nil {
return nil, errors.New("round trip not configured")
}
return transport.roundTrip(ctx, selected, request, commit...)
}
func (transport *fakeTransport) OpenTunnel(ctx context.Context, selected proxyDomain.Proxy, target string) (net.Conn, error) {
transport.record(selected.ID)
if transport.openTunnel == nil {
return nil, errors.New("tunnel not configured")
}
return transport.openTunnel(ctx, selected, target)
}
func (transport *fakeTransport) Relay(ctx context.Context, left, right net.Conn) error {
if transport.relay == nil {
return errors.New("relay not configured")
}
return transport.relay(ctx, left, right)
}
func (transport *fakeTransport) record(id string) {
transport.mu.Lock()
defer transport.mu.Unlock()
transport.seen = append(transport.seen, id)
}
func (transport *fakeTransport) attempts() []string {
transport.mu.Lock()
defer transport.mu.Unlock()
return append([]string(nil), transport.seen...)
}
func newTestHandler(
t *testing.T,
config Config,
dispatcher Dispatcher,
transport ProxyTransport,
) *Handler {
t.Helper()
handler, err := New(config, Dependencies{
Targets: fakeTargets{},
Router: RouteFunc(func(*http.Request) (dispatch.Request, error) {
return dispatch.Request{Upstreams: []string{"provider-a"}}, nil
}),
Dispatcher: dispatcher,
Transport: transport,
})
if err != nil {
t.Fatalf("New() error = %v", err)
}
return handler
}
func dispatcherWithProxies(t *testing.T, ids ...string) (*dispatch.Dispatcher, *snapshot.View) {
t.Helper()
store := snapshot.NewStore("cluster-a", "worker-a")
proxies := make([]proxyDomain.Proxy, 0, len(ids))
for index, id := range ids {
proxies = append(proxies, proxyDomain.Proxy{
ID: id,
Scheme: proxyDomain.SchemeHTTP,
Host: "127.0.0.1",
Port: uint16(20000 + index),
SourceUpstream: "provider-a",
MaxConcurrency: 2,
State: proxyDomain.StateAvailable,
CredentialVersion: "v1",
})
}
envelope := snapshot.Envelope{
ClusterID: "cluster-a",
WorkerID: "worker-a",
Epoch: 1,
Version: 1,
Full: true,
Proxies: proxies,
}
envelope.Checksum = snapshot.Checksum(proxies)
if err := store.Apply(envelope); err != nil {
t.Fatalf("apply snapshot: %v", err)
}
return dispatch.New(store), store.Current()
}
func assertNoLeakedCapacity(t *testing.T, view *snapshot.View) {
t.Helper()
deadline := time.Now().Add(500 * time.Millisecond)
for {
allReleased := true
for _, entry := range view.Entries {
if entry.Runtime.Active() != 0 || entry.Runtime.Reserved() != 0 {
allReleased = false
break
}
}
if allReleased {
return
}
if time.Now().After(deadline) {
for _, entry := range view.Entries {
if active, reserved := entry.Runtime.Active(), entry.Runtime.Reserved(); active != 0 || reserved != 0 {
t.Fatalf("proxy %s capacity leaked: active=%d reserved=%d", entry.Proxy.ID, active, reserved)
}
}
}
time.Sleep(time.Millisecond)
}
}
func sendConnect(t *testing.T, gatewayURL, target string) *http.Response {
t.Helper()
parsed, err := url.Parse(gatewayURL)
if err != nil {
t.Fatalf("parse gateway URL: %v", err)
}
connection, err := net.Dial("tcp", parsed.Host)
if err != nil {
t.Fatalf("dial gateway: %v", err)
}
t.Cleanup(func() { _ = connection.Close() })
if _, err := io.WriteString(connection, "CONNECT "+target+" HTTP/1.1\r\nHost: "+target+"\r\n\r\n"); err != nil {
t.Fatalf("write CONNECT request: %v", err)
}
response, err := http.ReadResponse(bufio.NewReader(connection), &http.Request{Method: http.MethodConnect})
if err != nil {
t.Fatalf("read CONNECT response: %v", err)
}
return response
}
var _ ProxyTransport = (*transportDomain.Transport)(nil)

View File

@ -0,0 +1,300 @@
package server
import (
"context"
"crypto/subtle"
"encoding/base64"
"errors"
"fmt"
"net"
"net/http"
"net/netip"
"strconv"
"strings"
"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
)
type HTTPError struct {
StatusCode int
Header http.Header
Cause error
}
func (err *HTTPError) Error() string {
if err.Cause != nil {
return err.Cause.Error()
}
return http.StatusText(err.StatusCode)
}
func (err *HTTPError) Unwrap() error { return err.Cause }
type BasicAuthGuard struct {
username string
password string
}
func NewBasicAuthGuard(username, password string) *BasicAuthGuard {
return &BasicAuthGuard{username: username, password: password}
}
func (guard *BasicAuthGuard) Check(_ context.Context, request *http.Request) error {
username, password, ok := parseBasicCredentials(request.Header.Get("Proxy-Authorization"))
if ok && constantTimeEqual(username, guard.username) && constantTimeEqual(password, guard.password) {
return nil
}
return &HTTPError{
StatusCode: http.StatusProxyAuthRequired,
Header: http.Header{"Proxy-Authenticate": []string{`Basic realm="proxy"`}},
Cause: errors.New("proxy authentication failed"),
}
}
type APIKeyGuard struct {
header string
value string
}
func NewAPIKeyGuard(header, value string) *APIKeyGuard {
if strings.TrimSpace(header) == "" {
header = "X-API-Key"
}
return &APIKeyGuard{header: header, value: value}
}
func (guard *APIKeyGuard) Check(_ context.Context, request *http.Request) error {
if constantTimeEqual(request.Header.Get(guard.header), guard.value) {
return nil
}
return &HTTPError{StatusCode: http.StatusProxyAuthRequired, Cause: errors.New("proxy API key authentication failed")}
}
type AnyGuard struct {
guards []Guard
}
func NewAnyGuard(guards ...Guard) *AnyGuard {
return &AnyGuard{guards: append([]Guard(nil), guards...)}
}
func (guard *AnyGuard) Check(ctx context.Context, request *http.Request) error {
var lastErr error
for _, candidate := range guard.guards {
if candidate == nil {
continue
}
if err := candidate.Check(ctx, request); err == nil {
return nil
} else {
lastErr = err
}
}
if lastErr == nil {
lastErr = errors.New("no authentication method is configured")
}
return lastErr
}
type ClientIPResolver struct {
trusted policy.CIDRMatcher
}
func NewClientIPResolver(trustedCIDRs []string) (*ClientIPResolver, error) {
trusted, err := policy.NewCIDRMatcher(trustedCIDRs)
if err != nil {
return nil, fmt.Errorf("create client IP resolver: %w", err)
}
return &ClientIPResolver{trusted: trusted}, nil
}
func (resolver *ClientIPResolver) Resolve(request *http.Request) (netip.Addr, error) {
if resolver == nil || request == nil {
return netip.Addr{}, errors.New("resolve client IP: resolver and request are required")
}
peer, err := parseRemoteAddress(request.RemoteAddr)
if err != nil {
return netip.Addr{}, err
}
if !resolver.trusted.Match(peer) {
return peer, nil
}
chain, present, err := parseForwardedChain(request.Header.Values("Forwarded"))
if err != nil {
return netip.Addr{}, err
}
if !present {
chain, err = parseXForwardedFor(request.Header.Values("X-Forwarded-For"))
if err != nil {
return netip.Addr{}, err
}
}
if len(chain) == 0 {
return peer, nil
}
for index := len(chain) - 1; index >= 0; index-- {
if !resolver.trusted.Match(chain[index]) {
return chain[index], nil
}
}
if len(chain) > 0 {
return chain[0], nil
}
return peer, nil
}
func parseXForwardedFor(fields []string) ([]netip.Addr, error) {
chain := make([]netip.Addr, 0, len(fields)+1)
for _, field := range fields {
for value := range strings.SplitSeq(field, ",") {
address, err := netip.ParseAddr(strings.TrimSpace(value))
if err != nil {
return nil, fmt.Errorf("resolve client IP: invalid X-Forwarded-For address %q", value)
}
chain = append(chain, address.Unmap())
}
}
return chain, nil
}
func parseForwardedChain(fields []string) ([]netip.Addr, bool, error) {
if len(fields) == 0 {
return nil, false, nil
}
chain := make([]netip.Addr, 0, len(fields)+1)
for _, field := range fields {
for element := range strings.SplitSeq(field, ",") {
found := false
for parameter := range strings.SplitSeq(element, ";") {
name, value, ok := strings.Cut(strings.TrimSpace(parameter), "=")
if !ok || !strings.EqualFold(name, "for") {
continue
}
address, err := parseForwardedIdentifier(value)
if err != nil {
return nil, true, err
}
chain = append(chain, address)
found = true
break
}
if !found {
return nil, true, errors.New("resolve client IP: Forwarded element is missing for parameter")
}
}
}
return chain, true, nil
}
func parseForwardedIdentifier(raw string) (netip.Addr, error) {
value := strings.TrimSpace(raw)
if strings.HasPrefix(value, `"`) {
unquoted, err := strconv.Unquote(value)
if err != nil {
return netip.Addr{}, fmt.Errorf("resolve client IP: invalid quoted Forwarded identifier")
}
value = unquoted
}
if strings.EqualFold(value, "unknown") || strings.HasPrefix(value, "_") {
return netip.Addr{}, fmt.Errorf("resolve client IP: non-IP Forwarded identifier")
}
if strings.HasPrefix(value, "[") {
closing := strings.IndexByte(value, ']')
if closing < 0 {
return netip.Addr{}, fmt.Errorf("resolve client IP: invalid Forwarded IPv6 identifier")
}
value = value[1:closing]
} else if host, _, err := net.SplitHostPort(value); err == nil {
value = host
}
address, err := netip.ParseAddr(value)
if err != nil {
return netip.Addr{}, fmt.Errorf("resolve client IP: invalid Forwarded address")
}
return address.Unmap(), nil
}
type AccessGuard struct {
resolver *ClientIPResolver
allow policy.CIDRMatcher
unrestricted bool
}
func NewAccessGuard(resolver *ClientIPResolver, allowCIDRs []string) (*AccessGuard, error) {
if resolver == nil {
return nil, errors.New("create access guard: client IP resolver is required")
}
allow, err := policy.NewCIDRMatcher(allowCIDRs)
if err != nil {
return nil, fmt.Errorf("create access guard: %w", err)
}
return &AccessGuard{resolver: resolver, allow: allow, unrestricted: len(allowCIDRs) == 0}, nil
}
func (guard *AccessGuard) Check(_ context.Context, request *http.Request) error {
address, err := guard.resolver.Resolve(request)
if err != nil {
return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
}
if guard.unrestricted || guard.allow.Match(address) {
return nil
}
return &HTTPError{StatusCode: http.StatusForbidden, Cause: fmt.Errorf("client address %s is not allowed", address)}
}
type Admitter interface {
Admit(context.Context, string) error
}
type AdmissionGuard struct {
resolver *ClientIPResolver
admitter Admitter
}
func NewAdmissionGuard(resolver *ClientIPResolver, admitter Admitter) *AdmissionGuard {
return &AdmissionGuard{resolver: resolver, admitter: admitter}
}
func (guard *AdmissionGuard) Check(ctx context.Context, request *http.Request) error {
if guard == nil || guard.resolver == nil || guard.admitter == nil {
return &HTTPError{StatusCode: http.StatusInternalServerError, Cause: errors.New("gateway admission is not configured")}
}
address, err := guard.resolver.Resolve(request)
if err != nil {
return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
}
if err := guard.admitter.Admit(ctx, address.String()); err != nil {
return &HTTPError{StatusCode: http.StatusTooManyRequests, Cause: err}
}
return nil
}
func parseBasicCredentials(value string) (string, string, bool) {
scheme, encoded, ok := strings.Cut(strings.TrimSpace(value), " ")
if !ok || !strings.EqualFold(scheme, "Basic") {
return "", "", false
}
decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(encoded))
if err != nil {
return "", "", false
}
username, password, ok := strings.Cut(string(decoded), ":")
return username, password, ok
}
func constantTimeEqual(actual, expected string) bool {
return subtle.ConstantTimeCompare([]byte(actual), []byte(expected)) == 1
}
func parseRemoteAddress(remote string) (netip.Addr, error) {
host, _, err := net.SplitHostPort(strings.TrimSpace(remote))
if err != nil {
host = strings.TrimSpace(remote)
}
address, err := netip.ParseAddr(host)
if err != nil {
return netip.Addr{}, fmt.Errorf("resolve client IP: invalid remote address %q", remote)
}
return address.Unmap(), nil
}

View File

@ -0,0 +1,139 @@
package server
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"net/netip"
"testing"
)
func TestBasicAuthGuardUsesProxyAuthorizationAndReturnsChallenge(t *testing.T) {
t.Parallel()
guard := NewBasicAuthGuard("client", "secret")
allowed := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
allowed.SetBasicAuth("ignored", "ignored")
allowed.Header.Set("Proxy-Authorization", "Basic Y2xpZW50OnNlY3JldA==")
if err := guard.Check(context.Background(), allowed); err != nil {
t.Fatalf("Check(valid) error = %v", err)
}
denied := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
denied.Header.Set("Proxy-Authorization", "Basic Y2xpZW50Ondyb25n")
err := guard.Check(context.Background(), denied)
var httpError *HTTPError
if !errors.As(err, &httpError) {
t.Fatalf("Check(invalid) error = %T %v, want *HTTPError", err, err)
}
if httpError.StatusCode != http.StatusProxyAuthRequired {
t.Fatalf("status = %d, want 407", httpError.StatusCode)
}
if got := httpError.Header.Get("Proxy-Authenticate"); got != `Basic realm="proxy"` {
t.Fatalf("challenge = %q", got)
}
}
func TestClientIPResolverOnlyTrustsForwardedChainFromTrustedPeer(t *testing.T) {
t.Parallel()
resolver, err := NewClientIPResolver([]string{"10.0.0.0/8", "192.0.2.0/24"})
if err != nil {
t.Fatalf("NewClientIPResolver() error = %v", err)
}
request := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
request.RemoteAddr = "10.0.0.9:1234"
request.Header.Set("X-Forwarded-For", "198.51.100.7, 192.0.2.5")
address, err := resolver.Resolve(request)
if err != nil {
t.Fatalf("Resolve(trusted) error = %v", err)
}
if address != netip.MustParseAddr("198.51.100.7") {
t.Fatalf("trusted forwarded address = %s", address)
}
request.RemoteAddr = "203.0.113.9:1234"
address, err = resolver.Resolve(request)
if err != nil {
t.Fatalf("Resolve(untrusted) error = %v", err)
}
if address != netip.MustParseAddr("203.0.113.9") {
t.Fatalf("untrusted forwarded address = %s", address)
}
}
func TestClientIPResolverSupportsRFCForwardedIPv4AndIPv6(t *testing.T) {
t.Parallel()
resolver, err := NewClientIPResolver([]string{"10.0.0.0/8", "192.0.2.0/24"})
if err != nil {
t.Fatalf("NewClientIPResolver() error = %v", err)
}
request := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
request.RemoteAddr = "10.0.0.9:1234"
request.Header.Set("Forwarded", `for="[2001:db8::7]:4711";proto=https, for=192.0.2.5`)
address, err := resolver.Resolve(request)
if err != nil {
t.Fatalf("Resolve() error = %v", err)
}
if address != netip.MustParseAddr("2001:db8::7") {
t.Fatalf("Forwarded address = %s", address)
}
}
func TestAccessAndAdmissionGuardsShareResolvedClientIdentity(t *testing.T) {
t.Parallel()
resolver, err := NewClientIPResolver(nil)
if err != nil {
t.Fatalf("NewClientIPResolver() error = %v", err)
}
access, err := NewAccessGuard(resolver, []string{"198.51.100.0/24"})
if err != nil {
t.Fatalf("NewAccessGuard() error = %v", err)
}
admitter := &recordingAdmitter{}
admission := NewAdmissionGuard(resolver, admitter)
request := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
request.RemoteAddr = "198.51.100.8:1234"
if err := access.Check(context.Background(), request); err != nil {
t.Fatalf("access.Check() error = %v", err)
}
if err := admission.Check(context.Background(), request); err != nil {
t.Fatalf("admission.Check() error = %v", err)
}
if admitter.key != "198.51.100.8" {
t.Fatalf("admission key = %q", admitter.key)
}
}
func TestAdmissionGuardMapsLimiterRejectionTo429(t *testing.T) {
t.Parallel()
resolver, _ := NewClientIPResolver(nil)
guard := NewAdmissionGuard(resolver, rejectingAdmitter{})
request := httptest.NewRequest(http.MethodGet, "http://example.test", nil)
request.RemoteAddr = "198.51.100.8:1234"
err := guard.Check(context.Background(), request)
var httpError *HTTPError
if !errors.As(err, &httpError) || httpError.StatusCode != http.StatusTooManyRequests {
t.Fatalf("Check() error = %T %v, want HTTP 429", err, err)
}
}
type recordingAdmitter struct{ key string }
func (admitter *recordingAdmitter) Admit(_ context.Context, key string) error {
admitter.key = key
return nil
}
type rejectingAdmitter struct{}
func (rejectingAdmitter) Admit(context.Context, string) error {
return errors.New("rate limited")
}

View File

@ -0,0 +1,65 @@
package server
import (
"errors"
"fmt"
"net"
"net/http"
"strings"
"github.com/proxy-pool/proxy-pool/internal/domain/routing"
"github.com/proxy-pool/proxy-pool/internal/gateway/dispatch"
)
var (
ErrRouteNotFound = errors.New("no gateway routing rule matched")
ErrRouteRejected = errors.New("gateway routing rule rejected the request")
ErrDirectRouteUnsupported = errors.New("direct gateway routing is not implemented")
)
type RulesRouter struct {
rules *routing.RuleSet
}
func NewRulesRouter(rules *routing.RuleSet) *RulesRouter {
return &RulesRouter{rules: rules}
}
func (router *RulesRouter) Route(request *http.Request) (dispatch.Request, error) {
if router == nil || router.rules == nil || request == nil {
return dispatch.Request{}, ErrRouteNotFound
}
host := request.Host
if request.URL != nil && request.URL.Hostname() != "" {
host = request.URL.Hostname()
} else if parsed, _, err := net.SplitHostPort(host); err == nil {
host = parsed
}
path := "/"
if request.URL != nil && request.URL.Path != "" {
path = request.URL.Path
}
headers := make(map[string]string, len(request.Header))
for name := range request.Header {
headers[name] = request.Header.Get(name)
}
matched, ok := router.rules.Match(routing.Request{
Host: strings.ToLower(strings.TrimSuffix(host, ".")),
Method: request.Method,
Path: path,
Headers: headers,
})
if !ok {
return dispatch.Request{}, ErrRouteNotFound
}
switch matched.Action {
case routing.ActionProxy:
return dispatch.Request{Upstreams: append([]string(nil), matched.Upstreams...)}, nil
case routing.ActionReject:
return dispatch.Request{}, fmt.Errorf("%w: %s", ErrRouteRejected, matched.Name)
case routing.ActionDirect:
return dispatch.Request{}, fmt.Errorf("%w: %s", ErrDirectRouteUnsupported, matched.Name)
default:
return dispatch.Request{}, fmt.Errorf("%w: %s has action %q", ErrRouteRejected, matched.Name, matched.Action)
}
}

View File

@ -0,0 +1,56 @@
package server
import (
"errors"
"net/http"
"net/http/httptest"
"reflect"
"testing"
"github.com/proxy-pool/proxy-pool/internal/domain/routing"
)
func TestRulesRouterReturnsMatchedUpstreams(t *testing.T) {
t.Parallel()
rules, err := routing.Compile([]routing.Rule{{
Name: "api",
Match: routing.Match{HostRegex: `^example\.test$`, Methods: []string{http.MethodGet}, PathRegex: `^/v1/`},
Upstreams: []string{"provider-a", "provider-b"},
Action: routing.ActionProxy,
}})
if err != nil {
t.Fatalf("routing.Compile() error = %v", err)
}
router := NewRulesRouter(rules)
request := httptest.NewRequest(http.MethodGet, "http://example.test/v1/items", nil)
result, err := router.Route(request)
if err != nil {
t.Fatalf("Route() error = %v", err)
}
if !reflect.DeepEqual(result.Upstreams, []string{"provider-a", "provider-b"}) {
t.Fatalf("upstreams = %v", result.Upstreams)
}
}
func TestRulesRouterRejectsExplicitRejectAndMissingRoute(t *testing.T) {
t.Parallel()
rules, err := routing.Compile([]routing.Rule{{
Name: "blocked",
Match: routing.Match{HostRegex: `^blocked\.test$`},
Action: routing.ActionReject,
}})
if err != nil {
t.Fatalf("routing.Compile() error = %v", err)
}
router := NewRulesRouter(rules)
if _, err := router.Route(httptest.NewRequest(http.MethodGet, "http://blocked.test/", nil)); !errors.Is(err, ErrRouteRejected) {
t.Fatalf("blocked Route() error = %v", err)
}
if _, err := router.Route(httptest.NewRequest(http.MethodGet, "http://missing.test/", nil)); !errors.Is(err, ErrRouteNotFound) {
t.Fatalf("missing Route() error = %v", err)
}
}

View File

@ -0,0 +1,58 @@
package transport
import (
"context"
"errors"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
platformCredentials "github.com/proxy-pool/proxy-pool/internal/platform/credentials"
)
var (
ErrCredentialStoreRequired = errors.New("credential store is required")
ErrIncompleteCredentialReference = errors.New("proxy credential reference is incomplete")
ErrCredentialUsernameMismatch = errors.New("resolved credential username does not match proxy metadata")
)
// StoreCredentialResolver adapts the shared credential store to the data-plane
// transport without placing plaintext secrets in proxy snapshots.
type StoreCredentialResolver struct {
store platformCredentials.Store
}
var _ CredentialResolver = (*StoreCredentialResolver)(nil)
func NewStoreCredentialResolver(store platformCredentials.Store) (*StoreCredentialResolver, error) {
if store == nil {
return nil, ErrCredentialStoreRequired
}
return &StoreCredentialResolver{store: store}, nil
}
func (resolver *StoreCredentialResolver) Resolve(
ctx context.Context,
selected proxyDomain.Proxy,
) (Credentials, error) {
if selected.SecretRef == "" && selected.CredentialVersion == "" {
return Credentials{Username: selected.Username}, nil
}
if selected.SecretRef == "" || selected.CredentialVersion == "" {
return Credentials{}, ErrIncompleteCredentialReference
}
if resolver == nil || resolver.store == nil {
return Credentials{}, ErrCredentialStoreRequired
}
value, err := resolver.store.Resolve(ctx, platformCredentials.Reference{
SecretRef: selected.SecretRef,
CredentialVersion: selected.CredentialVersion,
})
if err != nil {
return Credentials{}, err
}
if value.Username == "" {
value.Username = selected.Username
} else if selected.Username != "" && value.Username != selected.Username {
return Credentials{}, ErrCredentialUsernameMismatch
}
return Credentials{Username: value.Username, Password: value.Password}, nil
}

View File

@ -0,0 +1,112 @@
package transport
import (
"context"
"errors"
"fmt"
"strings"
"testing"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
platformCredentials "github.com/proxy-pool/proxy-pool/internal/platform/credentials"
)
func TestCredentialsFormattingRedactsPassword(t *testing.T) {
t.Parallel()
value := Credentials{Username: "alice", Password: "secret-password"}
for _, formatted := range []string{
fmt.Sprintf("%v", value),
fmt.Sprintf("%+v", value),
fmt.Sprintf("%#v", value),
} {
if strings.Contains(formatted, value.Username) || strings.Contains(formatted, value.Password) {
t.Fatalf("formatted Credentials exposes credential material: %s", formatted)
}
}
}
func TestStoreCredentialResolverResolvesProxyReference(t *testing.T) {
t.Parallel()
store, err := platformCredentials.NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore() error = %v", err)
}
reference, err := store.Put(context.Background(), "provider-a/proxy-a", platformCredentials.Value{
Username: "alice",
Password: "secret",
})
if err != nil {
t.Fatalf("Put() error = %v", err)
}
resolver, err := NewStoreCredentialResolver(store)
if err != nil {
t.Fatalf("NewStoreCredentialResolver() error = %v", err)
}
got, err := resolver.Resolve(context.Background(), proxyDomain.Proxy{
Username: "alice",
SecretRef: reference.SecretRef,
CredentialVersion: reference.CredentialVersion,
})
if err != nil {
t.Fatalf("Resolve() error = %v", err)
}
if got != (Credentials{Username: "alice", Password: "secret"}) {
t.Fatalf("Resolve() returned unexpected credentials")
}
}
func TestStoreCredentialResolverRejectsIncompleteOrMismatchedReference(t *testing.T) {
t.Parallel()
store, err := platformCredentials.NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore() error = %v", err)
}
reference, err := store.Put(context.Background(), "provider-a/proxy-a", platformCredentials.Value{
Username: "alice",
Password: "secret",
})
if err != nil {
t.Fatalf("Put() error = %v", err)
}
resolver, err := NewStoreCredentialResolver(store)
if err != nil {
t.Fatalf("NewStoreCredentialResolver() error = %v", err)
}
for name, selected := range map[string]proxyDomain.Proxy{
"missing version": {SecretRef: reference.SecretRef},
"missing ref": {CredentialVersion: reference.CredentialVersion},
} {
if _, err := resolver.Resolve(context.Background(), selected); !errors.Is(err, ErrIncompleteCredentialReference) {
t.Fatalf("%s Resolve() error = %v, want ErrIncompleteCredentialReference", name, err)
}
}
if _, err := resolver.Resolve(context.Background(), proxyDomain.Proxy{
Username: "bob",
SecretRef: reference.SecretRef,
CredentialVersion: reference.CredentialVersion,
}); !errors.Is(err, ErrCredentialUsernameMismatch) {
t.Fatalf("mismatched Resolve() error = %v, want ErrCredentialUsernameMismatch", err)
}
}
func TestStoreCredentialResolverAllowsCredentiallessProxy(t *testing.T) {
t.Parallel()
store, err := platformCredentials.NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore() error = %v", err)
}
resolver, err := NewStoreCredentialResolver(store)
if err != nil {
t.Fatalf("NewStoreCredentialResolver() error = %v", err)
}
got, err := resolver.Resolve(context.Background(), proxyDomain.Proxy{})
if err != nil || got != (Credentials{}) {
t.Fatalf("Resolve(credentialless) = (%+v, %v), want zero credentials", got, err)
}
}

View File

@ -0,0 +1,521 @@
package transport
import (
"bufio"
"bytes"
"context"
"crypto/tls"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptrace"
"net/url"
"sync"
"sync/atomic"
"time"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
)
const (
defaultDialTimeout = 10 * time.Second
defaultHandshakeTimeout = 15 * time.Second
defaultResponseHeaderTimeout = 30 * time.Second
defaultIdleConnTimeout = 90 * time.Second
defaultMaxErrorResponseBytes = int64(64 << 10)
defaultMaxResponseHeaderBytes = int64(64 << 10)
)
var (
ErrUnsupportedProxyScheme = errors.New("unsupported upstream proxy scheme")
ErrProxyResponseTooLarge = errors.New("upstream proxy response headers exceed the configured limit")
)
type Config struct {
DialTimeout time.Duration
HandshakeTimeout time.Duration
ResponseHeaderTimeout time.Duration
IdleConnTimeout time.Duration
MaxIdleConns int
MaxIdleConnsPerHost int
MaxErrorResponseBytes int64
MaxResponseHeaderBytes int64
TunnelBufferBytes int
TunnelIdleTimeout time.Duration
TLSClientConfig *tls.Config
}
type Credentials struct {
Username string
Password string
}
func (Credentials) Format(state fmt.State, _ rune) {
_, _ = state.Write([]byte("transport.Credentials{Username:<redacted>, Password:<redacted>}"))
}
type CredentialResolver interface {
Resolve(context.Context, proxyDomain.Proxy) (Credentials, error)
}
type CredentialResolverFunc func(context.Context, proxyDomain.Proxy) (Credentials, error)
func (resolve CredentialResolverFunc) Resolve(ctx context.Context, selected proxyDomain.Proxy) (Credentials, error) {
return resolve(ctx, selected)
}
type ProxyResponseError struct {
StatusCode int
Status string
Header http.Header
Body []byte
}
func (err *ProxyResponseError) Error() string {
return fmt.Sprintf("upstream proxy CONNECT response: %s", err.Status)
}
func (err *ProxyResponseError) Retryable() bool {
return err.StatusCode == http.StatusRequestTimeout ||
err.StatusCode == http.StatusTooEarly ||
err.StatusCode == http.StatusTooManyRequests ||
(err.StatusCode >= http.StatusInternalServerError && err.StatusCode <= 599)
}
type Transport struct {
config Config
resolver CredentialResolver
client *http.Transport
buffers sync.Pool
}
func New(config Config, resolver CredentialResolver) *Transport {
applyDefaults(&config)
transport := &Transport{config: config, resolver: resolver}
tlsConfig := config.TLSClientConfig
if tlsConfig != nil {
tlsConfig = tlsConfig.Clone()
}
transport.client = &http.Transport{
Proxy: func(request *http.Request) (*url.URL, error) {
proxyURL, ok := request.Context().Value(proxyURLContextKey{}).(*url.URL)
if !ok || proxyURL == nil {
return nil, errors.New("upstream proxy URL is missing from request context")
}
return proxyURL, nil
},
DialContext: (&net.Dialer{Timeout: config.DialTimeout, KeepAlive: 30 * time.Second}).DialContext,
ForceAttemptHTTP2: true,
MaxIdleConns: config.MaxIdleConns,
MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
IdleConnTimeout: config.IdleConnTimeout,
TLSHandshakeTimeout: config.HandshakeTimeout,
ResponseHeaderTimeout: config.ResponseHeaderTimeout,
TLSClientConfig: tlsConfig,
MaxResponseHeaderBytes: config.MaxResponseHeaderBytes,
}
transport.buffers.New = func() any {
return make([]byte, config.TunnelBufferBytes)
}
return transport
}
func (transport *Transport) RoundTrip(
ctx context.Context,
selected proxyDomain.Proxy,
request *http.Request,
commit ...func() error,
) (*http.Response, error) {
if request == nil {
return nil, errors.New("round trip through proxy: nil request")
}
credentials, err := transport.resolveCredentials(ctx, selected)
if err != nil {
return nil, fmt.Errorf("resolve upstream proxy credentials: %w", err)
}
proxyURL, err := upstreamURL(selected, credentials)
if err != nil {
return nil, err
}
var commitOnce sync.Once
var commitMu sync.Mutex
var commitErr error
if len(commit) > 0 && commit[0] != nil {
trace := &httptrace.ClientTrace{GotConn: func(httptrace.GotConnInfo) {
commitOnce.Do(func() {
commitMu.Lock()
commitErr = commit[0]()
commitMu.Unlock()
})
}}
ctx = httptrace.WithClientTrace(ctx, trace)
}
ctx = context.WithValue(ctx, proxyURLContextKey{}, proxyURL)
clone := request.Clone(ctx)
clone.RequestURI = ""
clone.Header = request.Header.Clone()
clone.Header.Del("Proxy-Authorization")
response, err := transport.client.RoundTrip(clone)
commitMu.Lock()
deferredCommitErr := commitErr
commitMu.Unlock()
if deferredCommitErr != nil {
if response != nil {
_ = response.Body.Close()
}
return nil, fmt.Errorf("commit upstream proxy allocation: %w", deferredCommitErr)
}
return response, err
}
func (transport *Transport) OpenTunnel(
ctx context.Context,
selected proxyDomain.Proxy,
target string,
) (net.Conn, error) {
credentials, err := transport.resolveCredentials(ctx, selected)
if err != nil {
return nil, fmt.Errorf("resolve upstream proxy credentials: %w", err)
}
handshakeContext, cancelHandshake := context.WithTimeout(ctx, transport.config.HandshakeTimeout)
defer cancelHandshake()
connection, err := transport.dialProxy(handshakeContext, selected)
if err != nil {
return nil, err
}
succeeded := false
defer func() {
if !succeeded {
_ = connection.Close()
}
}()
deadline := time.Now().Add(transport.config.HandshakeTimeout)
if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
deadline = contextDeadline
}
if err := connection.SetDeadline(deadline); err != nil {
return nil, fmt.Errorf("set upstream proxy handshake deadline: %w", err)
}
stopCancellation := context.AfterFunc(ctx, func() {
_ = connection.SetDeadline(time.Now())
})
defer stopCancellation()
request := &http.Request{
Method: http.MethodConnect,
URL: &url.URL{Opaque: target},
Host: target,
Header: make(http.Header),
}
if credentials.Username != "" || credentials.Password != "" {
request.Header.Set("Proxy-Authorization", basicAuth(credentials))
}
if err := request.Write(connection); err != nil {
return nil, contextError(ctx, fmt.Errorf("write upstream CONNECT: %w", err))
}
limited := &io.LimitedReader{R: connection, N: transport.config.MaxResponseHeaderBytes + 1}
reader := bufio.NewReader(limited)
response, err := http.ReadResponse(reader, request)
if err != nil {
if limited.N == 0 {
return nil, ErrProxyResponseTooLarge
}
return nil, contextError(ctx, fmt.Errorf("read upstream CONNECT response: %w", err))
}
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
body, readErr := io.ReadAll(io.LimitReader(response.Body, transport.config.MaxErrorResponseBytes))
_ = response.Body.Close()
if readErr != nil {
return nil, fmt.Errorf("read upstream CONNECT error response: %w", readErr)
}
return nil, &ProxyResponseError{
StatusCode: response.StatusCode,
Status: response.Status,
Header: response.Header.Clone(),
Body: body,
}
}
_ = response.Body.Close()
if err := connection.SetDeadline(time.Time{}); err != nil {
return nil, fmt.Errorf("clear upstream proxy handshake deadline: %w", err)
}
buffered := make([]byte, reader.Buffered())
if _, err := io.ReadFull(reader, buffered); err != nil {
return nil, fmt.Errorf("preserve buffered CONNECT bytes: %w", err)
}
succeeded = true
return &bufferedConn{Conn: connection, reader: io.MultiReader(bytes.NewReader(buffered), connection)}, nil
}
func (transport *Transport) CloseIdleConnections() {
if transport == nil {
return
}
transport.client.CloseIdleConnections()
}
func (transport *Transport) Relay(ctx context.Context, left, right net.Conn) error {
if left == nil || right == nil {
return errors.New("relay tunnel: nil connection")
}
stopCancellation := context.AfterFunc(ctx, func() {
deadline := time.Now()
_ = left.SetDeadline(deadline)
_ = right.SetDeadline(deadline)
})
defer stopCancellation()
activity := newTunnelActivity()
idleStopped := make(chan struct{})
go transport.enforceTunnelIdle(ctx, idleStopped, left, right, activity)
defer close(idleStopped)
activeLeft := &activityConn{Conn: left, activity: activity}
activeRight := &activityConn{Conn: right, activity: activity}
errorsByDirection := make(chan error, 2)
go func() { errorsByDirection <- transport.copyTunnel(activeRight, activeLeft) }()
go func() { errorsByDirection <- transport.copyTunnel(activeLeft, activeRight) }()
first := <-errorsByDirection
if first != nil {
_ = left.Close()
_ = right.Close()
}
second := <-errorsByDirection
if err := ctx.Err(); err != nil {
return err
}
return errors.Join(first, second)
}
func (transport *Transport) enforceTunnelIdle(
ctx context.Context,
stopped <-chan struct{},
left, right net.Conn,
activity *tunnelActivity,
) {
timer := time.NewTimer(transport.config.TunnelIdleTimeout)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return
case <-stopped:
return
case <-timer.C:
remaining := activity.remaining(transport.config.TunnelIdleTimeout)
if remaining > 0 {
timer.Reset(remaining)
continue
}
deadline := time.Now()
_ = left.SetDeadline(deadline)
_ = right.SetDeadline(deadline)
return
}
}
}
func (transport *Transport) copyTunnel(destination, source net.Conn) error {
buffer := transport.buffers.Get().([]byte)
defer transport.buffers.Put(buffer)
_, err := io.CopyBuffer(destination, source, buffer)
if halfCloser, ok := destination.(interface{ CloseWrite() error }); ok {
_ = halfCloser.CloseWrite()
}
if halfCloser, ok := source.(interface{ CloseRead() error }); ok {
_ = halfCloser.CloseRead()
}
if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
return nil
}
return err
}
func (transport *Transport) dialProxy(ctx context.Context, selected proxyDomain.Proxy) (net.Conn, error) {
if selected.Scheme != proxyDomain.SchemeHTTP && selected.Scheme != proxyDomain.SchemeHTTPS {
return nil, fmt.Errorf("%w: %s", ErrUnsupportedProxyScheme, selected.Scheme)
}
dialer := &net.Dialer{Timeout: transport.config.DialTimeout, KeepAlive: 30 * time.Second}
connection, err := dialer.DialContext(ctx, "tcp", selected.Address())
if err != nil {
return nil, contextError(ctx, fmt.Errorf("dial upstream proxy %s: %w", selected.ID, err))
}
if selected.Scheme == proxyDomain.SchemeHTTP {
return connection, nil
}
tlsConfig := &tls.Config{MinVersion: tls.VersionTLS12, ServerName: selected.Host}
if transport.config.TLSClientConfig != nil {
tlsConfig = transport.config.TLSClientConfig.Clone()
if tlsConfig.ServerName == "" {
tlsConfig.ServerName = selected.Host
}
}
tlsConnection := tls.Client(connection, tlsConfig)
if err := tlsConnection.HandshakeContext(ctx); err != nil {
_ = connection.Close()
return nil, contextError(ctx, fmt.Errorf("TLS handshake with upstream proxy %s: %w", selected.ID, err))
}
return tlsConnection, nil
}
func (transport *Transport) resolveCredentials(
ctx context.Context,
selected proxyDomain.Proxy,
) (Credentials, error) {
if transport.resolver == nil {
return Credentials{Username: selected.Username}, nil
}
credentials, err := transport.resolver.Resolve(ctx, selected)
if credentials.Username == "" {
credentials.Username = selected.Username
}
return credentials, err
}
func upstreamURL(selected proxyDomain.Proxy, credentials Credentials) (*url.URL, error) {
if selected.Scheme != proxyDomain.SchemeHTTP && selected.Scheme != proxyDomain.SchemeHTTPS {
return nil, fmt.Errorf("%w: %s", ErrUnsupportedProxyScheme, selected.Scheme)
}
proxyURL := &url.URL{Scheme: string(selected.Scheme), Host: selected.Address()}
if credentials.Username != "" || credentials.Password != "" {
proxyURL.User = url.UserPassword(credentials.Username, credentials.Password)
}
return proxyURL, nil
}
func applyDefaults(config *Config) {
if config.DialTimeout <= 0 {
config.DialTimeout = defaultDialTimeout
}
if config.HandshakeTimeout <= 0 {
config.HandshakeTimeout = defaultHandshakeTimeout
}
if config.ResponseHeaderTimeout <= 0 {
config.ResponseHeaderTimeout = defaultResponseHeaderTimeout
}
if config.IdleConnTimeout <= 0 {
config.IdleConnTimeout = defaultIdleConnTimeout
}
if config.MaxIdleConns <= 0 {
config.MaxIdleConns = 1024
}
if config.MaxIdleConnsPerHost <= 0 {
config.MaxIdleConnsPerHost = 64
}
if config.MaxErrorResponseBytes <= 0 {
config.MaxErrorResponseBytes = defaultMaxErrorResponseBytes
}
if config.MaxResponseHeaderBytes <= 0 {
config.MaxResponseHeaderBytes = defaultMaxResponseHeaderBytes
}
if config.TunnelBufferBytes <= 0 {
config.TunnelBufferBytes = 32 << 10
}
if config.TunnelIdleTimeout <= 0 {
config.TunnelIdleTimeout = 5 * time.Minute
}
}
func basicAuth(credentials Credentials) string {
request := &http.Request{Header: make(http.Header)}
request.SetBasicAuth(credentials.Username, credentials.Password)
return request.Header.Get("Authorization")
}
func contextError(ctx context.Context, fallback error) error {
if err := ctx.Err(); err != nil {
return err
}
var networkError net.Error
if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) &&
errors.As(fallback, &networkError) && networkError.Timeout() {
return context.DeadlineExceeded
}
return fallback
}
type bufferedConn struct {
net.Conn
reader io.Reader
}
type proxyURLContextKey struct{}
type tunnelActivity struct {
started time.Time
lastElapsed atomic.Int64
}
func newTunnelActivity() *tunnelActivity {
activity := &tunnelActivity{started: time.Now()}
activity.touch()
return activity
}
func (activity *tunnelActivity) touch() {
activity.lastElapsed.Store(int64(time.Since(activity.started)))
}
func (activity *tunnelActivity) remaining(timeout time.Duration) time.Duration {
elapsedSinceActivity := time.Since(activity.started) - time.Duration(activity.lastElapsed.Load())
return timeout - elapsedSinceActivity
}
type activityConn struct {
net.Conn
activity *tunnelActivity
}
func (connection *activityConn) Read(buffer []byte) (int, error) {
read, err := connection.Conn.Read(buffer)
if read > 0 {
connection.activity.touch()
}
return read, err
}
func (connection *activityConn) Write(buffer []byte) (int, error) {
written, err := connection.Conn.Write(buffer)
if written > 0 {
connection.activity.touch()
}
return written, err
}
func (connection *activityConn) CloseWrite() error {
if halfCloser, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
return halfCloser.CloseWrite()
}
return connection.Close()
}
func (connection *activityConn) CloseRead() error {
if halfCloser, ok := connection.Conn.(interface{ CloseRead() error }); ok {
return halfCloser.CloseRead()
}
return nil
}
func (connection *bufferedConn) Read(buffer []byte) (int, error) {
return connection.reader.Read(buffer)
}
func (connection *bufferedConn) CloseWrite() error {
if halfCloser, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
return halfCloser.CloseWrite()
}
return connection.Close()
}
func (connection *bufferedConn) CloseRead() error {
if halfCloser, ok := connection.Conn.(interface{ CloseRead() error }); ok {
return halfCloser.CloseRead()
}
return nil
}

View File

@ -0,0 +1,437 @@
package transport
import (
"bufio"
"context"
"encoding/base64"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
proxyDomain "github.com/proxy-pool/proxy-pool/internal/domain/proxy"
)
func TestRoundTripForwardsHTTPViaSelectedProxy(t *testing.T) {
t.Parallel()
requestSeen := make(chan *http.Request, 1)
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
requestSeen <- request.Clone(request.Context())
writer.Header().Set("X-Upstream", "selected")
writer.WriteHeader(http.StatusCreated)
_, _ = writer.Write([]byte("forwarded"))
}))
defer upstream.Close()
selected := proxyFromURL(t, upstream.URL)
selected.ID = "proxy-a"
selected.Username = "alice"
selected.CredentialVersion = "v1"
selected.SecretRef = "secret://proxy-a"
client := New(Config{}, CredentialResolverFunc(func(context.Context, proxyDomain.Proxy) (Credentials, error) {
return Credentials{Username: "alice", Password: "s3cret"}, nil
}))
request := httptest.NewRequest(http.MethodGet, "http://TARGET/resource?q=1", nil)
response, err := client.RoundTrip(request.Context(), selected, request)
if err != nil {
t.Fatalf("RoundTrip() error = %v", err)
}
defer response.Body.Close()
body, err := io.ReadAll(response.Body)
if err != nil {
t.Fatalf("read response body: %v", err)
}
if response.StatusCode != http.StatusCreated || string(body) != "forwarded" {
t.Fatalf("response = (%d, %q), want (201, forwarded)", response.StatusCode, body)
}
seen := <-requestSeen
if seen.RequestURI != "http://TARGET/resource?q=1" {
t.Fatalf("proxy request URI = %q", seen.RequestURI)
}
wantAuth := "Basic " + base64.StdEncoding.EncodeToString([]byte("alice:s3cret"))
if got := seen.Header.Get("Proxy-Authorization"); got != wantAuth {
t.Fatalf("Proxy-Authorization = %q, want %q", got, wantAuth)
}
}
func TestRoundTripCommitsReservationAfterConnectionAcquisition(t *testing.T) {
t.Parallel()
allowResponse := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
<-allowResponse
writer.WriteHeader(http.StatusNoContent)
}))
defer upstream.Close()
client := New(Config{}, nil)
request := httptest.NewRequest(http.MethodGet, "http://TARGET/resource", nil)
committed := make(chan struct{}, 1)
done := make(chan error, 1)
go func() {
response, err := client.RoundTrip(request.Context(), proxyFromURL(t, upstream.URL), request, func() error {
committed <- struct{}{}
return nil
})
if response != nil {
_ = response.Body.Close()
}
done <- err
}()
select {
case <-committed:
case <-time.After(time.Second):
t.Fatal("commit hook was not called after acquiring the proxy connection")
}
close(allowResponse)
if err := <-done; err != nil {
t.Fatalf("RoundTrip() error = %v", err)
}
}
func TestOpenTunnelPreservesBytesBufferedAfterSuccessfulHandshake(t *testing.T) {
t.Parallel()
listener, requests := startConnectProxy(t, func(connection net.Conn) {
_, _ = io.WriteString(connection, "HTTP/1.1 200 Connection Established\r\n\r\nREADY")
})
selected := proxyFromAddress(listener.Addr().String())
selected.Username = "alice"
client := New(Config{}, CredentialResolverFunc(func(context.Context, proxyDomain.Proxy) (Credentials, error) {
return Credentials{Username: "alice", Password: "s3cret"}, nil
}))
connection, err := client.OpenTunnel(context.Background(), selected, "example.test:443")
if err != nil {
t.Fatalf("OpenTunnel() error = %v", err)
}
defer connection.Close()
preface := make([]byte, len("READY"))
if _, err := io.ReadFull(connection, preface); err != nil {
t.Fatalf("read buffered tunnel bytes: %v", err)
}
if string(preface) != "READY" {
t.Fatalf("tunnel preface = %q", preface)
}
request := <-requests
if request.Method != http.MethodConnect || request.Host != "example.test:443" {
t.Fatalf("CONNECT request = %s %s", request.Method, request.Host)
}
wantAuth := "Basic " + base64.StdEncoding.EncodeToString([]byte("alice:s3cret"))
if got := request.Header.Get("Proxy-Authorization"); got != wantAuth {
t.Fatalf("Proxy-Authorization = %q, want %q", got, wantAuth)
}
}
func TestOpenTunnelReturnsBoundedProxyResponseError(t *testing.T) {
t.Parallel()
listener, _ := startConnectProxy(t, func(connection net.Conn) {
_, _ = io.WriteString(connection, "HTTP/1.1 407 Proxy Authentication Required\r\nContent-Length: 8\r\nProxy-Authenticate: Basic\r\n\r\ndenied!!")
})
client := New(Config{MaxErrorResponseBytes: 4}, nil)
connection, err := client.OpenTunnel(context.Background(), proxyFromAddress(listener.Addr().String()), "example.test:443")
if connection != nil {
_ = connection.Close()
t.Fatal("OpenTunnel() returned a connection for 407")
}
var responseError *ProxyResponseError
if !errors.As(err, &responseError) {
t.Fatalf("OpenTunnel() error = %T %v, want *ProxyResponseError", err, err)
}
if responseError.StatusCode != http.StatusProxyAuthRequired {
t.Fatalf("status = %d, want 407", responseError.StatusCode)
}
if string(responseError.Body) != "deni" {
t.Fatalf("bounded body = %q, want deni", responseError.Body)
}
}
func TestProxyResponseErrorRetryable(t *testing.T) {
t.Parallel()
tests := []struct {
statusCode int
want bool
}{
{statusCode: http.StatusProxyAuthRequired, want: false},
{statusCode: http.StatusRequestTimeout, want: true},
{statusCode: http.StatusTooEarly, want: true},
{statusCode: http.StatusTooManyRequests, want: true},
{statusCode: http.StatusInternalServerError, want: true},
{statusCode: http.StatusBadGateway, want: true},
{statusCode: http.StatusServiceUnavailable, want: true},
{statusCode: http.StatusGatewayTimeout, want: true},
{statusCode: http.StatusBadRequest, want: false},
}
for _, tt := range tests {
t.Run(fmt.Sprint(tt.statusCode), func(t *testing.T) {
err := &ProxyResponseError{StatusCode: tt.statusCode}
if got := err.Retryable(); got != tt.want {
t.Fatalf("Retryable() = %t, want %t", got, tt.want)
}
})
}
}
func TestOpenTunnelRejectsOversizedHandshakeResponse(t *testing.T) {
t.Parallel()
listener, _ := startConnectProxy(t, func(connection net.Conn) {
_, _ = io.WriteString(connection, "HTTP/1.1 200 Connection Established\r\nX-Large: "+strings.Repeat("x", 4096)+"\r\n\r\n")
})
client := New(Config{MaxResponseHeaderBytes: 256}, nil)
connection, err := client.OpenTunnel(context.Background(), proxyFromAddress(listener.Addr().String()), "example.test:443")
if connection != nil {
_ = connection.Close()
t.Fatal("OpenTunnel() returned a connection for an oversized handshake")
}
if !errors.Is(err, ErrProxyResponseTooLarge) {
t.Fatalf("OpenTunnel() error = %v, want ErrProxyResponseTooLarge", err)
}
}
func TestOpenTunnelHonorsContextCancellationDuringHandshake(t *testing.T) {
t.Parallel()
listener, _ := startConnectProxy(t, func(connection net.Conn) {
<-time.After(time.Second)
})
client := New(Config{HandshakeTimeout: time.Second}, nil)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
started := time.Now()
connection, err := client.OpenTunnel(ctx, proxyFromAddress(listener.Addr().String()), "example.test:443")
if connection != nil {
_ = connection.Close()
}
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("OpenTunnel() error = %v, want context deadline exceeded", err)
}
if elapsed := time.Since(started); elapsed > 300*time.Millisecond {
t.Fatalf("cancellation took %s", elapsed)
}
}
func TestRelayPreservesTCPHalfCloseInBothDirections(t *testing.T) {
t.Parallel()
client, gatewayClient := tcpPair(t)
gatewayUpstream, upstream := tcpPair(t)
defer client.Close()
defer gatewayClient.Close()
defer gatewayUpstream.Close()
defer upstream.Close()
relay := New(Config{TunnelBufferBytes: 1024}, nil)
relayDone := make(chan error, 1)
go func() {
relayDone <- relay.Relay(context.Background(), gatewayClient, gatewayUpstream)
}()
if _, err := io.WriteString(client, "request"); err != nil {
t.Fatalf("write client request: %v", err)
}
if err := client.CloseWrite(); err != nil {
t.Fatalf("half-close client: %v", err)
}
request, err := io.ReadAll(upstream)
if err != nil {
t.Fatalf("read upstream request: %v", err)
}
if string(request) != "request" {
t.Fatalf("upstream request = %q", request)
}
if _, err := io.WriteString(upstream, "response"); err != nil {
t.Fatalf("write upstream response: %v", err)
}
if err := upstream.CloseWrite(); err != nil {
t.Fatalf("half-close upstream: %v", err)
}
response, err := io.ReadAll(client)
if err != nil {
t.Fatalf("read client response: %v", err)
}
if string(response) != "response" {
t.Fatalf("client response = %q", response)
}
select {
case err := <-relayDone:
if err != nil {
t.Fatalf("Relay() error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("Relay() did not finish after both half-closes")
}
}
func TestRelayKeepsTunnelAliveWhileTrafficFlowsInOneDirection(t *testing.T) {
t.Parallel()
client, gatewayClient := tcpPair(t)
gatewayUpstream, upstream := tcpPair(t)
defer client.Close()
defer gatewayClient.Close()
defer gatewayUpstream.Close()
defer upstream.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
relay := New(Config{TunnelIdleTimeout: 100 * time.Millisecond}, nil)
done := make(chan error, 1)
go func() { done <- relay.Relay(ctx, gatewayClient, gatewayUpstream) }()
if err := client.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set client read deadline: %v", err)
}
for range 6 {
time.Sleep(30 * time.Millisecond)
if _, err := upstream.Write([]byte("x")); err != nil {
t.Fatalf("write one-way tunnel traffic: %v", err)
}
buffer := make([]byte, 1)
if _, err := io.ReadFull(client, buffer); err != nil {
t.Fatalf("read one-way tunnel traffic: %v", err)
}
}
select {
case err := <-done:
t.Fatalf("Relay() ended during one-way activity: %v", err)
default:
}
cancel()
select {
case err := <-done:
if !errors.Is(err, context.Canceled) {
t.Fatalf("Relay() error = %v, want context canceled", err)
}
case <-time.After(time.Second):
t.Fatal("Relay() did not stop after cancellation")
}
}
func TestRelayStopsAnIdleTunnelAtConfiguredDeadline(t *testing.T) {
t.Parallel()
left, leftPeer := net.Pipe()
right, rightPeer := net.Pipe()
defer leftPeer.Close()
defer rightPeer.Close()
relay := New(Config{TunnelIdleTimeout: 30 * time.Millisecond}, nil)
done := make(chan error, 1)
go func() { done <- relay.Relay(context.Background(), left, right) }()
select {
case err := <-done:
if err == nil {
t.Fatal("Relay() error = nil, want idle timeout")
}
case <-time.After(300 * time.Millisecond):
t.Fatal("Relay() did not enforce tunnel idle timeout")
}
}
func startConnectProxy(t *testing.T, respond func(net.Conn)) (net.Listener, <-chan *http.Request) {
t.Helper()
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
t.Cleanup(func() { _ = listener.Close() })
requests := make(chan *http.Request, 1)
go func() {
connection, acceptErr := listener.Accept()
if acceptErr != nil {
return
}
defer connection.Close()
request, readErr := http.ReadRequest(bufio.NewReader(connection))
if readErr != nil {
return
}
requests <- request
respond(connection)
}()
return listener, requests
}
func proxyFromURL(t *testing.T, rawURL string) proxyDomain.Proxy {
t.Helper()
parsed, err := url.Parse(rawURL)
if err != nil {
t.Fatalf("parse proxy URL: %v", err)
}
return proxyFromAddress(parsed.Host)
}
func proxyFromAddress(address string) proxyDomain.Proxy {
host, portText, err := net.SplitHostPort(address)
if err != nil {
panic(fmt.Sprintf("split proxy address %q: %v", address, err))
}
var port uint16
if _, err := fmt.Sscanf(portText, "%d", &port); err != nil {
panic(fmt.Sprintf("parse proxy port %q: %v", portText, err))
}
return proxyDomain.Proxy{
ID: strings.ReplaceAll(address, ":", "-"),
Scheme: proxyDomain.SchemeHTTP,
Host: host,
Port: port,
MaxConcurrency: 1,
}
}
func tcpPair(t *testing.T) (*net.TCPConn, *net.TCPConn) {
t.Helper()
listener, err := net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP("127.0.0.1")})
if err != nil {
t.Fatalf("listen TCP pair: %v", err)
}
defer listener.Close()
accepted := make(chan *net.TCPConn, 1)
acceptErrors := make(chan error, 1)
go func() {
connection, acceptErr := listener.AcceptTCP()
if acceptErr != nil {
acceptErrors <- acceptErr
return
}
accepted <- connection
}()
client, err := net.DialTCP("tcp", nil, listener.Addr().(*net.TCPAddr))
if err != nil {
t.Fatalf("dial TCP pair: %v", err)
}
select {
case server := <-accepted:
return client, server
case acceptErr := <-acceptErrors:
_ = client.Close()
t.Fatalf("accept TCP pair: %v", acceptErr)
}
return nil, nil
}

View File

@ -0,0 +1,28 @@
package credentials
import (
"fmt"
"strings"
"testing"
)
func TestReferenceFormattingRedactsSecretRef(t *testing.T) {
t.Parallel()
reference := Reference{
SecretRef: "cred_sensitive_reference",
CredentialVersion: "v7",
}
for _, formatted := range []string{
fmt.Sprintf("%v", reference),
fmt.Sprintf("%+v", reference),
fmt.Sprintf("%#v", reference),
} {
if strings.Contains(formatted, reference.SecretRef) {
t.Fatalf("formatted Reference exposes SecretRef: %s", formatted)
}
if !strings.Contains(formatted, reference.CredentialVersion) {
t.Fatalf("formatted Reference omits non-secret version: %s", formatted)
}
}
}

View File

@ -0,0 +1,222 @@
package credentials
import (
"context"
"crypto/rand"
"encoding/hex"
"errors"
"fmt"
"strconv"
"strings"
)
var (
ErrInvalidCapacity = errors.New("invalid credential store capacity")
ErrInvalidStore = errors.New("invalid credential store")
ErrInvalidScope = errors.New("invalid credential scope")
ErrInvalidReference = errors.New("invalid credential reference")
ErrCapacityExceeded = errors.New("credential store capacity exceeded")
ErrCredentialMissing = errors.New("credential not found")
ErrCredentialVersionMismatch = errors.New("credential version mismatch")
ErrReferenceCreation = errors.New("credential reference creation failed")
)
// Value is resolved credential material. Callers must not log this value.
type Value struct {
Username string
Password string
}
func (Value) Format(state fmt.State, _ rune) {
_, _ = state.Write([]byte("credentials.Value{Username:<redacted>, Password:<redacted>}"))
}
// Reference identifies one exact version of stored credential material.
type Reference struct {
SecretRef string
CredentialVersion string
}
func (reference Reference) Format(state fmt.State, _ rune) {
formatted := "credentials.Reference{SecretRef:<redacted>, CredentialVersion:" +
strconv.Quote(reference.CredentialVersion) + "}"
_, _ = state.Write([]byte(formatted))
}
type Store interface {
Put(context.Context, string, Value) (Reference, error)
Resolve(context.Context, Reference) (Value, error)
}
type entry struct {
value Value
version uint64
reference Reference
}
// MemoryStore keeps credentials in process memory and serializes access with a
// context-aware lock.
type MemoryStore struct {
lock chan struct{}
capacity int
byScope map[string]*entry
byRef map[string]*entry
}
func (s *MemoryStore) Format(state fmt.State, _ rune) {
if s == nil {
_, _ = state.Write([]byte("credentials.MemoryStore<nil>"))
return
}
_, _ = state.Write([]byte("credentials.MemoryStore{capacity:" + strconv.Itoa(s.capacity) + "}"))
}
var _ Store = (*MemoryStore)(nil)
func NewMemoryStore(capacity int) (*MemoryStore, error) {
if capacity <= 0 {
return nil, ErrInvalidCapacity
}
lock := make(chan struct{}, 1)
lock <- struct{}{}
return &MemoryStore{
lock: lock,
capacity: capacity,
byScope: make(map[string]*entry),
byRef: make(map[string]*entry),
}, nil
}
func (s *MemoryStore) Put(ctx context.Context, scope string, value Value) (Reference, error) {
if err := contextError(ctx); err != nil {
return Reference{}, err
}
if !s.valid() {
return Reference{}, ErrInvalidStore
}
if strings.TrimSpace(scope) == "" {
return Reference{}, ErrInvalidScope
}
if err := s.acquire(ctx); err != nil {
return Reference{}, err
}
defer s.release()
if err := ctx.Err(); err != nil {
return Reference{}, err
}
if current, ok := s.byScope[scope]; ok {
if current.value == value {
return current.reference, nil
}
current.value = value
current.version++
current.reference.CredentialVersion = versionString(current.version)
return current.reference, nil
}
if len(s.byScope) >= s.capacity {
return Reference{}, ErrCapacityExceeded
}
secretRef, err := s.newUniqueSecretRef()
if err != nil {
return Reference{}, err
}
created := &entry{
value: value,
version: 1,
reference: Reference{
SecretRef: secretRef,
CredentialVersion: versionString(1),
},
}
s.byScope[scope] = created
s.byRef[secretRef] = created
return created.reference, nil
}
func (s *MemoryStore) Resolve(ctx context.Context, reference Reference) (Value, error) {
if err := contextError(ctx); err != nil {
return Value{}, err
}
if !s.valid() {
return Value{}, ErrInvalidStore
}
if reference.SecretRef == "" || !validVersion(reference.CredentialVersion) {
return Value{}, ErrInvalidReference
}
if err := s.acquire(ctx); err != nil {
return Value{}, err
}
defer s.release()
if err := ctx.Err(); err != nil {
return Value{}, err
}
current, ok := s.byRef[reference.SecretRef]
if !ok {
return Value{}, ErrCredentialMissing
}
if current.reference.CredentialVersion != reference.CredentialVersion {
return Value{}, ErrCredentialVersionMismatch
}
return current.value, nil
}
func (s *MemoryStore) acquire(ctx context.Context) error {
if err := contextError(ctx); err != nil {
return err
}
select {
case <-ctx.Done():
return ctx.Err()
case <-s.lock:
return nil
}
}
func (s *MemoryStore) release() {
s.lock <- struct{}{}
}
func (s *MemoryStore) valid() bool {
return s != nil && s.lock != nil && s.capacity > 0 && s.byScope != nil && s.byRef != nil
}
func (s *MemoryStore) newUniqueSecretRef() (string, error) {
for {
secretRef, err := newSecretRef()
if err != nil {
return "", err
}
if _, exists := s.byRef[secretRef]; !exists {
return secretRef, nil
}
}
}
func newSecretRef() (string, error) {
var random [24]byte
if _, err := rand.Read(random[:]); err != nil {
return "", ErrReferenceCreation
}
return "cred_" + hex.EncodeToString(random[:]), nil
}
func versionString(version uint64) string {
return "v" + strconv.FormatUint(version, 10)
}
func validVersion(version string) bool {
if len(version) < 2 || version[0] != 'v' || version[1] == '0' {
return false
}
_, err := strconv.ParseUint(version[1:], 10, 64)
return err == nil
}
func contextError(ctx context.Context) error {
if ctx == nil {
return context.Canceled
}
return ctx.Err()
}

View File

@ -0,0 +1,358 @@
package credentials
import (
"context"
"errors"
"fmt"
"strconv"
"strings"
"sync"
"testing"
"time"
)
func TestMemoryStorePutIsIdempotentForUnchangedScope(t *testing.T) {
store, err := NewMemoryStore(2)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
value := Value{Username: "alice", Password: "secret-password"}
first, err := store.Put(context.Background(), "provider-a", value)
if err != nil {
t.Fatalf("Put(first): %v", err)
}
second, err := store.Put(context.Background(), "provider-a", value)
if err != nil {
t.Fatalf("Put(second): %v", err)
}
if first != second {
t.Fatalf("second reference = %#v, want %#v", second, first)
}
if first.SecretRef == "" || first.CredentialVersion != "v1" {
t.Fatalf("first reference = %#v, want opaque ref at v1", first)
}
for _, plaintext := range []string{"provider-a", value.Username, value.Password} {
if strings.Contains(first.SecretRef, plaintext) {
t.Fatalf("SecretRef %q contains plaintext %q", first.SecretRef, plaintext)
}
}
got, err := store.Resolve(context.Background(), first)
if err != nil {
t.Fatalf("Resolve(): %v", err)
}
if got != value {
t.Fatalf("Resolve() = %#v, want %#v", got, value)
}
}
func TestCredentialFormattingRedactsSensitiveMaterial(t *testing.T) {
value := Value{Username: "alice", Password: "secret-password"}
for _, formatted := range []string{
fmt.Sprintf("%v", value),
fmt.Sprintf("%+v", value),
fmt.Sprintf("%#v", value),
fmt.Sprintf("%s", value),
fmt.Sprintf("%q", value),
} {
if strings.Contains(formatted, value.Password) {
t.Fatalf("formatted Value exposes password: %s", formatted)
}
}
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
reference, err := store.Put(context.Background(), "provider-secret-scope", value)
if err != nil {
t.Fatalf("Put(): %v", err)
}
for _, formatted := range []string{
fmt.Sprintf("%v", store),
fmt.Sprintf("%+v", store),
fmt.Sprintf("%#v", store),
fmt.Sprintf("%s", store),
fmt.Sprintf("%q", store),
} {
for _, secret := range []string{
"provider-secret-scope",
value.Username,
value.Password,
reference.SecretRef,
} {
if strings.Contains(formatted, secret) {
t.Fatalf("formatted MemoryStore exposes sensitive material: %s", formatted)
}
}
}
}
func TestMemoryStorePutIncrementsVersionAndRejectsStaleReference(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
oldReference, err := store.Put(context.Background(), "provider-a", Value{
Username: "alice",
Password: "old-password",
})
if err != nil {
t.Fatalf("Put(old): %v", err)
}
want := Value{Username: "alice", Password: "new-password"}
newReference, err := store.Put(context.Background(), "provider-a", want)
if err != nil {
t.Fatalf("Put(new): %v", err)
}
if newReference.SecretRef != oldReference.SecretRef {
t.Fatalf("new SecretRef changed across versions")
}
if newReference.CredentialVersion != "v2" {
t.Fatalf("new CredentialVersion = %q, want v2", newReference.CredentialVersion)
}
if _, err := store.Resolve(context.Background(), oldReference); !errors.Is(err, ErrCredentialVersionMismatch) {
t.Fatalf("Resolve(stale) error = %v, want ErrCredentialVersionMismatch", err)
}
got, err := store.Resolve(context.Background(), newReference)
if err != nil {
t.Fatalf("Resolve(new): %v", err)
}
if got != want {
t.Fatalf("Resolve(new) returned unexpected credentials")
}
}
func TestMemoryStoreEnforcesCapacityWithoutChangingExistingCredentials(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
want := Value{Username: "alice", Password: "first-password"}
reference, err := store.Put(context.Background(), "provider-a", want)
if err != nil {
t.Fatalf("Put(first scope): %v", err)
}
if _, err := store.Put(context.Background(), "provider-b", Value{
Username: "bob",
Password: "second-password",
}); !errors.Is(err, ErrCapacityExceeded) {
t.Fatalf("Put(over capacity) error = %v, want ErrCapacityExceeded", err)
}
got, err := store.Resolve(context.Background(), reference)
if err != nil {
t.Fatalf("Resolve(existing): %v", err)
}
if got != want {
t.Fatalf("Resolve(existing) returned changed credentials")
}
if _, err := store.Put(context.Background(), "provider-a", want); err != nil {
t.Fatalf("Put(idempotent at capacity): %v", err)
}
}
func TestMemoryStoreStrictlyValidatesCredentialVersion(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
reference, err := store.Put(context.Background(), "provider-a", Value{Password: "password"})
if err != nil {
t.Fatalf("Put(): %v", err)
}
for _, version := range []string{"", "v0", "v01", "1", "latest", "v-1"} {
invalid := reference
invalid.CredentialVersion = version
if _, err := store.Resolve(context.Background(), invalid); !errors.Is(err, ErrInvalidReference) {
t.Fatalf("Resolve(version %q) error = %v, want ErrInvalidReference", version, err)
}
}
if _, err := store.Put(context.Background(), " \t\n", Value{}); !errors.Is(err, ErrInvalidScope) {
t.Fatalf("Put(blank scope) error = %v, want ErrInvalidScope", err)
}
}
func TestMemoryStoreValidatesCapacityAndHonorsCancellation(t *testing.T) {
if _, err := NewMemoryStore(0); !errors.Is(err, ErrInvalidCapacity) {
t.Fatalf("NewMemoryStore(0) error = %v, want ErrInvalidCapacity", err)
}
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(1): %v", err)
}
canceled, cancel := context.WithCancel(context.Background())
cancel()
if _, err := store.Put(canceled, "provider-a", Value{}); !errors.Is(err, context.Canceled) {
t.Fatalf("Put(canceled) error = %v, want context.Canceled", err)
}
if _, err := store.Resolve(canceled, Reference{
SecretRef: "cred_000000000000000000000000000000000000000000000000",
CredentialVersion: "v1",
}); !errors.Is(err, context.Canceled) {
t.Fatalf("Resolve(canceled) error = %v, want context.Canceled", err)
}
}
func TestMemoryStoreRejectsNilAndZeroValueStores(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
for _, store := range []*MemoryStore{nil, {}} {
if _, err := store.Put(ctx, "provider-a", Value{}); !errors.Is(err, ErrInvalidStore) {
t.Fatalf("Put() error = %v, want ErrInvalidStore", err)
}
if _, err := store.Resolve(ctx, Reference{
SecretRef: "cred_000000000000000000000000000000000000000000000000",
CredentialVersion: "v1",
}); !errors.Is(err, ErrInvalidStore) {
t.Fatalf("Resolve() error = %v, want ErrInvalidStore", err)
}
}
}
func TestMemoryStoreIsConcurrencySafeAndIdempotent(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
want := Value{Username: "alice", Password: "shared-password"}
const workers = 100
references := make(chan Reference, workers)
errorsSeen := make(chan error, workers)
var wait sync.WaitGroup
for range workers {
wait.Add(1)
go func() {
defer wait.Done()
reference, err := store.Put(context.Background(), "provider-a", want)
if err != nil {
errorsSeen <- err
return
}
references <- reference
}()
}
wait.Wait()
close(references)
close(errorsSeen)
for err := range errorsSeen {
t.Errorf("concurrent Put(): %v", err)
}
var first Reference
for reference := range references {
if first == (Reference{}) {
first = reference
}
if reference != first {
t.Errorf("concurrent Put() reference differs from first")
}
}
if first.CredentialVersion != "v1" {
t.Fatalf("concurrent version = %q, want v1", first.CredentialVersion)
}
got, err := store.Resolve(context.Background(), first)
if err != nil {
t.Fatalf("Resolve(): %v", err)
}
if got != want {
t.Fatalf("Resolve() returned unexpected credentials")
}
}
func TestMemoryStoreSerializesConcurrentCredentialChanges(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
type result struct {
value Value
reference Reference
err error
}
const workers = 100
results := make(chan result, workers)
var wait sync.WaitGroup
for index := range workers {
wait.Add(1)
go func() {
defer wait.Done()
value := Value{Username: "alice", Password: "password-" + strconv.Itoa(index)}
reference, err := store.Put(context.Background(), "provider-a", value)
results <- result{value: value, reference: reference, err: err}
}()
}
wait.Wait()
close(results)
versions := make(map[string]struct{}, workers)
var secretRef string
var latest result
for current := range results {
if current.err != nil {
t.Fatalf("concurrent Put(): %v", current.err)
}
if secretRef == "" {
secretRef = current.reference.SecretRef
}
if current.reference.SecretRef != secretRef {
t.Fatal("SecretRef changed across concurrent updates")
}
versions[current.reference.CredentialVersion] = struct{}{}
if current.reference.CredentialVersion == "v100" {
latest = current
}
}
if len(versions) != workers {
t.Fatalf("unique versions = %d, want %d", len(versions), workers)
}
if latest.reference == (Reference{}) {
t.Fatal("highest version v100 was not returned")
}
got, err := store.Resolve(context.Background(), latest.reference)
if err != nil {
t.Fatalf("Resolve(latest): %v", err)
}
if got != latest.value {
t.Fatal("Resolve(latest) returned a different concurrent write")
}
}
func TestMemoryStoreErrorsDoNotExposeInputs(t *testing.T) {
store, err := NewMemoryStore(1)
if err != nil {
t.Fatalf("NewMemoryStore(): %v", err)
}
value := Value{Username: "sensitive-user", Password: "sensitive-password"}
reference, err := store.Put(context.Background(), "sensitive-scope", value)
if err != nil {
t.Fatalf("Put(): %v", err)
}
_, capacityErr := store.Put(context.Background(), "other-sensitive-scope", value)
unknown := Reference{
SecretRef: "sensitive-secret-ref",
CredentialVersion: "v1",
}
_, missingErr := store.Resolve(context.Background(), unknown)
stale := reference
stale.CredentialVersion = "v2"
_, versionErr := store.Resolve(context.Background(), stale)
for _, operationErr := range []error{capacityErr, missingErr, versionErr} {
if operationErr == nil {
t.Fatal("operation unexpectedly succeeded")
}
for _, sensitive := range []string{
"sensitive-scope",
"other-sensitive-scope",
value.Username,
value.Password,
unknown.SecretRef,
reference.SecretRef,
} {
if strings.Contains(operationErr.Error(), sensitive) {
t.Fatalf("error exposes sensitive input: %v", operationErr)
}
}
}
}