Skip to content

Commit aa51557

Browse files
wk8wk8
committed
Making it possible to specify connect and header timeouts on registry backends
With unit tests. Signed-off-by: Jean Rouge <rougej+github@gmail.com>
1 parent e435b83 commit aa51557

File tree

7 files changed

+139
-10
lines changed

7 files changed

+139
-10
lines changed

‎Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ bins: $(LINUX_BINS)
8282
# ==== TEST ====
8383
.PHONY: unit-test
8484
unit-test:
85-
-rm coverage.txt
85+
-rm -f coverage.txt
8686
$(GO) test -timeout=30s -race -coverprofile=coverage.txt $(ALL_PKGS) --tags "unit"
8787

8888
.PHONY: docker_stop

‎lib/backend/registrybackend/blobclient.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ type BlobClient struct {
6363
// NewBlobClient creates a new BlobClient.
6464
func NewBlobClient(config Config) (*BlobClient, error) {
6565
config = config.applyDefaults()
66-
authenticator, err := security.NewAuthenticator(config.Address, config.Security)
66+
authenticator, err := config.Authenticator()
6767
if err != nil {
6868
return nil, fmt.Errorf("cannot create tag client authenticator: %s", err)
6969
}

‎lib/backend/registrybackend/blobclient_test.go

+38
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@ import (
1919
"io"
2020
"net/http"
2121
"testing"
22+
"time"
2223

2324
"github.com/pressly/chi"
25+
"github.com/stretchr/testify/assert"
2426
"github.com/stretchr/testify/require"
2527
"github.com/uber/kraken/core"
2628
"github.com/uber/kraken/lib/backend/backenderrors"
@@ -125,3 +127,39 @@ func TestBlobDownloadFileNotFound(t *testing.T) {
125127
var b bytes.Buffer
126128
require.Equal(backenderrors.ErrBlobNotFound, client.Download(namespace, "data", &b))
127129
}
130+
131+
func TestBlobDownloadHeaderTimeout(t *testing.T) {
132+
require := require.New(t)
133+
134+
blob := randutil.Blob(32 * memsize.KB)
135+
namespace := core.NamespaceFixture()
136+
137+
r := chi.NewRouter()
138+
r.Get(fmt.Sprintf("/v2/%s/blobs/{blob}", namespace), func(w http.ResponseWriter, req *http.Request) {
139+
time.Sleep(time.Second)
140+
// ignoring errors here, as this will fail after we timeout below
141+
_, _ = io.Copy(w, bytes.NewReader(blob))
142+
})
143+
r.Head(fmt.Sprintf("/v2/%s/blobs/{blob}", namespace), func(w http.ResponseWriter, req *http.Request) {
144+
time.Sleep(time.Second)
145+
w.Header().Set("Content-Length", fmt.Sprintf("%d", len(blob)))
146+
})
147+
addr, stop := testutil.StartServer(r)
148+
defer stop()
149+
150+
config := newTestConfig(addr)
151+
config.ResponseHeaderTimeout = 100 * time.Millisecond
152+
client, err := NewBlobClient(config)
153+
require.NoError(err)
154+
155+
_, err = client.Stat(namespace, "data")
156+
if assert.NotNil(t, err) {
157+
assert.Contains(t, err.Error(), "timeout awaiting response headers")
158+
}
159+
160+
var b bytes.Buffer
161+
err = client.Download(namespace, "data", &b)
162+
if assert.NotNil(t, err) {
163+
assert.Contains(t, err.Error(), "timeout awaiting response headers")
164+
}
165+
}

‎lib/backend/registrybackend/config.go

+27-3
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,22 @@
1414
package registrybackend
1515

1616
import (
17+
"net"
18+
"net/http"
1719
"time"
1820

1921
"github.com/uber/kraken/lib/backend/registrybackend/security"
2022
)
2123

2224
// Config defines the registry address, timeout and security options.
2325
type Config struct {
24-
Address string `yaml:"address"`
25-
Timeout time.Duration `yaml:"timeout"`
26-
Security security.Config `yaml:"security"`
26+
Address string `yaml:"address"`
27+
Timeout time.Duration `yaml:"timeout"`
28+
// ConnectTimeout limits the time spent establishing the TCP connection (if a new one is needed).
29+
ConnectTimeout time.Duration `yaml:"connect_timeout"`
30+
// ResponseHeaderTimeout limits the time spent reading the headers of the response.
31+
ResponseHeaderTimeout time.Duration `yaml:"response_header_timeout"`
32+
Security security.Config `yaml:"security"`
2733
}
2834

2935
// Set default configuration
@@ -33,3 +39,21 @@ func (c Config) applyDefaults() Config {
3339
}
3440
return c
3541
}
42+
43+
func (c Config) Authenticator() (security.Authenticator, error) {
44+
transport := http.DefaultTransport.(*http.Transport).Clone()
45+
46+
if c.ConnectTimeout != 0 {
47+
dialer := &net.Dialer{
48+
Timeout: c.ConnectTimeout,
49+
KeepAlive: 30 * time.Second,
50+
}
51+
transport.DialContext = dialer.DialContext
52+
}
53+
54+
if c.ResponseHeaderTimeout != 0 {
55+
transport.ResponseHeaderTimeout = c.ResponseHeaderTimeout
56+
}
57+
58+
return security.NewAuthenticator(c.Address, c.Security, transport)
59+
}

‎lib/backend/registrybackend/security/security.go

+3-4
Original file line numberDiff line numberDiff line change
@@ -72,17 +72,16 @@ type authenticator struct {
7272
// address, TLS, and credentials configuration. It supports both basic auth and
7373
// token based authentication challenges. If TLS is disabled, no authentication
7474
// is attempted.
75-
func NewAuthenticator(address string, config Config) (Authenticator, error) {
76-
rt := http.DefaultTransport.(*http.Transport).Clone()
75+
func NewAuthenticator(address string, config Config, transport *http.Transport) (Authenticator, error) {
7776
tlsClientConfig, err := config.TLS.BuildClient()
7877
if err != nil {
7978
return nil, fmt.Errorf("build tls config for %q: %s", address, err)
8079
}
81-
rt.TLSClientConfig = tlsClientConfig
80+
transport.TLSClientConfig = tlsClientConfig
8281
return &authenticator{
8382
address: address,
8483
config: config,
85-
roundTripper: rt,
84+
roundTripper: transport,
8685
credentialStore: newCredentialStore(address, config),
8786
challengeManager: challenge.NewSimpleManager(),
8887
}, nil

‎lib/backend/registrybackend/tagclient.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ type TagClient struct {
6464
// NewTagClient creates a new TagClient.
6565
func NewTagClient(config Config) (*TagClient, error) {
6666
config = config.applyDefaults()
67-
authenticator, err := security.NewAuthenticator(config.Address, config.Security)
67+
authenticator, err := config.Authenticator()
6868
if err != nil {
6969
return nil, fmt.Errorf("cannot create tag client authenticator: %s", err)
7070
}

‎lib/backend/registrybackend/tagclient_test.go

+68
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@ import (
2020
"net/http"
2121
"strings"
2222
"testing"
23+
"time"
2324

2425
"github.com/pressly/chi"
26+
"github.com/stretchr/testify/assert"
2527
"github.com/stretchr/testify/require"
2628
"github.com/uber/kraken/core"
2729
"github.com/uber/kraken/lib/backend/backenderrors"
@@ -97,3 +99,69 @@ func TestTagDownloadFileNotFound(t *testing.T) {
9799
var b bytes.Buffer
98100
require.Equal(backenderrors.ErrBlobNotFound, client.Download(tag, tag, &b))
99101
}
102+
103+
func TestTagDownloadHeaderTimeout(t *testing.T) {
104+
require := require.New(t)
105+
106+
imageConfig := core.NewBlobFixture()
107+
layer1 := core.NewBlobFixture()
108+
layer2 := core.NewBlobFixture()
109+
digest, manifest := dockerutil.ManifestFixture(
110+
imageConfig.Digest, layer1.Digest, layer2.Digest)
111+
112+
tag := core.TagFixture()
113+
namespace := strings.Split(tag, ":")[0]
114+
115+
r := chi.NewRouter()
116+
r.Get(fmt.Sprintf("/v2/%s/manifests/{tag}", namespace), func(w http.ResponseWriter, req *http.Request) {
117+
time.Sleep(time.Second)
118+
w.Header().Set("Content-Length", fmt.Sprintf("%d", len(manifest)))
119+
w.Header().Set("Docker-Content-Digest", digest.String())
120+
_, _ = io.Copy(w, bytes.NewReader(manifest))
121+
})
122+
r.Head(fmt.Sprintf("/v2/%s/manifests/{tag}", namespace), func(w http.ResponseWriter, req *http.Request) {
123+
time.Sleep(time.Second)
124+
w.Header().Set("Content-Length", fmt.Sprintf("%d", len(manifest)))
125+
w.Header().Set("Docker-Content-Digest", digest.String())
126+
_, _ = io.Copy(w, bytes.NewReader(manifest))
127+
})
128+
addr, stop := testutil.StartServer(r)
129+
defer stop()
130+
131+
config := newTestConfig(addr)
132+
config.ResponseHeaderTimeout = 100 * time.Millisecond
133+
client, err := NewTagClient(config)
134+
require.NoError(err)
135+
136+
_, err = client.Stat(tag, tag)
137+
if assert.NotNil(t, err) {
138+
assert.Contains(t, err.Error(), "timeout awaiting response headers")
139+
}
140+
141+
var b bytes.Buffer
142+
err = client.Download(tag, tag, &b)
143+
if assert.NotNil(t, err) {
144+
assert.Contains(t, err.Error(), "timeout awaiting response headers")
145+
}
146+
}
147+
148+
func TestTagDownloadConnectTimeout(t *testing.T) {
149+
require := require.New(t)
150+
151+
// unroutable address, courtesy of https://stackoverflow.com/a/904609/4867444
152+
config := newTestConfig("10.255.255.1")
153+
config.ConnectTimeout = 100 * time.Millisecond
154+
client, err := NewTagClient(config)
155+
require.NoError(err)
156+
157+
_, err = client.Stat("dummynamespace", "image:tag")
158+
if assert.NotNil(t, err) {
159+
assert.Contains(t, err.Error(), "i/o timeout")
160+
}
161+
162+
var b bytes.Buffer
163+
err = client.Download("dummynamespace", "image:tag", &b)
164+
if assert.NotNil(t, err) {
165+
assert.Contains(t, err.Error(), "i/o timeout")
166+
}
167+
}

0 commit comments

Comments
 (0)