Cache issued tgz in memory for one-time download
Instead of streaming the tgz directly to the response (which was
fragile under server write timeouts), handleIssueCert now:
- Builds the tgz into a bytes.Buffer
- Stores it in a sync.Map (tgzCache) under a random 16-byte hex token
- Redirects the browser to /pki/download/{token}
handleTGZDownload serves the cached bytes via LoadAndDelete, so the
archive is removed from memory after the first (and only) download.
An unknown or already-used token returns 404.
Also adds TestHandleTGZDownload covering the one-time-use and
not-found cases, and wires issueCertFn into mockVault.
Co-authored-by: Junie <junie@jetbrains.com>
This commit is contained in:
@@ -24,6 +24,7 @@ type mockVault struct {
|
||||
validateTokenFn func(ctx context.Context, token string) (*TokenInfo, error)
|
||||
listMountsFn func(ctx context.Context, token string) ([]MountInfo, error)
|
||||
getCertFn func(ctx context.Context, token, mount, serial string) (*CertDetail, error)
|
||||
issueCertFn func(ctx context.Context, token string, req IssueCertRequest) (*IssuedCert, error)
|
||||
}
|
||||
|
||||
func (m *mockVault) Status(ctx context.Context) (string, error) {
|
||||
@@ -80,6 +81,9 @@ func (m *mockVault) ListIssuers(ctx context.Context, token, mount string) ([]str
|
||||
}
|
||||
|
||||
func (m *mockVault) IssueCert(ctx context.Context, token string, req IssueCertRequest) (*IssuedCert, error) {
|
||||
if m.issueCertFn != nil {
|
||||
return m.issueCertFn(ctx, token, req)
|
||||
}
|
||||
return nil, fmt.Errorf("not implemented")
|
||||
}
|
||||
|
||||
@@ -108,6 +112,62 @@ func (m *mockVault) DeleteCert(ctx context.Context, token, mount, serial string)
|
||||
|
||||
func (m *mockVault) Close() error { return nil }
|
||||
|
||||
// ---- handleTGZDownload tests ----
|
||||
|
||||
func TestHandleTGZDownload(t *testing.T) {
|
||||
t.Run("valid token serves archive and removes it from cache", func(t *testing.T) {
|
||||
ws := newTestWebServer(t, &mockVault{})
|
||||
|
||||
const dlToken = "abc123"
|
||||
ws.tgzCache.Store(dlToken, &tgzEntry{
|
||||
filename: "test.tgz",
|
||||
data: []byte("fake-tgz-data"),
|
||||
})
|
||||
|
||||
r := newChiRequest(http.MethodGet, "/pki/download/"+dlToken, map[string]string{"token": dlToken})
|
||||
r = addAuthCookie(r, &TokenInfo{Username: "testuser"})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
ws.handleTGZDownload(w, r)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("status = %d, want %d", w.Code, http.StatusOK)
|
||||
}
|
||||
if ct := w.Header().Get("Content-Type"); ct != "application/gzip" {
|
||||
t.Errorf("Content-Type = %q, want application/gzip", ct)
|
||||
}
|
||||
if body := w.Body.String(); body != "fake-tgz-data" {
|
||||
t.Errorf("body = %q, want fake-tgz-data", body)
|
||||
}
|
||||
|
||||
// Entry must have been removed — a second request should 404.
|
||||
w2 := httptest.NewRecorder()
|
||||
r2 := newChiRequest(http.MethodGet, "/pki/download/"+dlToken, map[string]string{"token": dlToken})
|
||||
r2 = addAuthCookie(r2, &TokenInfo{Username: "testuser"})
|
||||
ws.handleTGZDownload(w2, r2)
|
||||
if w2.Code != http.StatusNotFound {
|
||||
t.Errorf("second request status = %d, want %d", w2.Code, http.StatusNotFound)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unknown token returns 404", func(t *testing.T) {
|
||||
ws := newTestWebServer(t, &mockVault{})
|
||||
|
||||
r := newChiRequest(http.MethodGet, "/pki/download/nosuchtoken", map[string]string{"token": "nosuchtoken"})
|
||||
r = addAuthCookie(r, &TokenInfo{Username: "testuser"})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
ws.handleTGZDownload(w, r)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("status = %d, want %d", w.Code, http.StatusNotFound)
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "download not found") {
|
||||
t.Errorf("body %q does not contain 'download not found'", w.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// newTestWebServer builds a WebServer wired to the given mock, suitable for unit tests.
|
||||
func newTestWebServer(t *testing.T, vault vaultBackend) *WebServer {
|
||||
t.Helper()
|
||||
|
||||
Reference in New Issue
Block a user