feat: Add PyPI dependency cooldown support (#221)

* refactor: Extract shared cooldown helpers to package-level functions

* feat: Add PyPI cooldown handler with PEP 691 file parsing

* feat: Add PyPI cooldown file stripping logic

* feat: Implement PyPI cooldown HandleMetadataRequest with PEP 691 filtering

* feat: Wire PyPI cooldown into pypi_registry interceptor

* update headers for no cache

* fix: Strip conditional GET headers to prevent 304 bypass in cooldown handlers

pip and npm clients cache Simple API / registry responses with ETags. On
subsequent requests they send If-None-Match, which causes the server to
return 304 Not Modified with no body. The cooldown response modifier
received an empty body, failed to parse it, and failed-open — letting
the client use its stale cached (unfiltered) response.

Fix: delete If-None-Match and If-Modified-Since from the request before
forwarding, forcing a full 200 response so the modifier always has a
body to filter.

Also removes the Content-Type guard from the PyPI modifier (the empty
Content-Type on 304 responses was a symptom of the same root cause) and
replaces Cache-Control: no-cache with the more targeted header deletion.

* docs: Add PyPI cooldown limitation for pip < 22.3 to dependency-cooldown docs
This commit is contained in:
Sahil Bansal
2026-04-20 21:10:05 +05:30
committed by GitHub
parent 365deb1897
commit 544b38b278
8 changed files with 1013 additions and 168 deletions
+9 -1
View File
@@ -28,4 +28,12 @@ pmg --skip-dependency-cooldown npm install express
## Requirements
Dependency cooldown requires [proxy mode](proxy-mode.md) to be enabled. It is currently supported for npm packages.
Dependency cooldown requires [proxy mode](proxy-mode.md) to be enabled. It is supported for npm and PyPI packages.
## Limitations
### PyPI: requires pip 22.3+ or a PEP 691-capable client
PyPI cooldown is enforced by filtering the [PEP 691 JSON Simple API](https://peps.python.org/pep-0691/) response, which includes a per-file `upload-time` field needed to determine when each version was published. This JSON format is only supported by pip 22.3+ (released October 2022) and other modern tools such as uv, Poetry, and PDM.
Older pip versions request the HTML Simple API, which carries no publish timestamps. PMG cannot apply cooldown filtering to HTML responses and fails open; the request passes through unchanged and the client receives the full version list. Old pip gets no cooldown protection but does not break.
+48
View File
@@ -0,0 +1,48 @@
package interceptors
import "time"
// cooldownIsWithinWindow reports whether a version published at publishDate is still
// within the cooldown window of cooldownDays. Returns withinCooldown, daysSincePublish,
// and daysRemaining.
func cooldownIsWithinWindow(publishDate time.Time, cooldownDays int) (withinCooldown bool, daysSincePublish int, daysRemaining int) {
daysSincePublish = int(time.Since(publishDate).Hours() / 24)
if daysSincePublish < 0 {
daysSincePublish = 0
}
daysRemaining = cooldownDays - daysSincePublish
if daysRemaining < 0 {
daysRemaining = 0
}
return daysSincePublish < cooldownDays, daysSincePublish, daysRemaining
}
// cooldownOldestVersion returns the version with the earliest publish date.
// When all versions are in cooldown, this is the one closest to exiting the window.
func cooldownOldestVersion(dates map[string]time.Time) (string, time.Time) {
var oldest string
var oldestTime time.Time
for version, publishDate := range dates {
if oldestTime.IsZero() || publishDate.Before(oldestTime) {
oldest = version
oldestTime = publishDate
}
}
return oldest, oldestTime
}
// cooldownLatestEligibleVersion returns the most recently published version not in tooNew.
func cooldownLatestEligibleVersion(dates map[string]time.Time, tooNew map[string]bool) string {
var latest string
var latestTime time.Time
for version, publishDate := range dates {
if tooNew[version] {
continue
}
if publishDate.After(latestTime) {
latest = version
latestTime = publishDate
}
}
return latest
}
+164
View File
@@ -0,0 +1,164 @@
package interceptors
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestCooldownIsWithinWindow(t *testing.T) {
now := time.Now()
day := 24 * time.Hour
tests := []struct {
name string
publishDate time.Time
cooldownDays int
wantWithinCooldown bool
wantDaysSincePublish int
wantDaysRemaining int
}{
{
name: "published today with 30 day cooldown",
publishDate: now,
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 30,
},
{
name: "published exactly at cooldown boundary",
publishDate: now.Add(-30 * day),
cooldownDays: 30,
wantWithinCooldown: false,
wantDaysSincePublish: 30,
wantDaysRemaining: 0,
},
{
name: "published one day before cooldown expires",
publishDate: now.Add(-29 * day),
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 29,
wantDaysRemaining: 1,
},
{
name: "published well beyond cooldown",
publishDate: now.Add(-365 * day),
cooldownDays: 30,
wantWithinCooldown: false,
wantDaysSincePublish: 365,
wantDaysRemaining: 0,
},
{
name: "zero cooldown days",
publishDate: now,
cooldownDays: 0,
wantWithinCooldown: false,
wantDaysSincePublish: 0,
wantDaysRemaining: 0,
},
{
name: "future publish date clamped to zero days",
publishDate: now.Add(5 * day),
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 30,
},
{
name: "one day cooldown with publish today",
publishDate: now,
cooldownDays: 1,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 1,
},
{
name: "one day cooldown with publish yesterday",
publishDate: now.Add(-1 * day),
cooldownDays: 1,
wantWithinCooldown: false,
wantDaysSincePublish: 1,
wantDaysRemaining: 0,
},
{
name: "max int cooldown days does not overflow",
publishDate: now.Add(-1 * day),
cooldownDays: int(^uint(0) >> 1),
wantWithinCooldown: true,
wantDaysSincePublish: 1,
wantDaysRemaining: int(^uint(0)>>1) - 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
withinCooldown, daysSincePublish, daysRemaining := cooldownIsWithinWindow(tt.publishDate, tt.cooldownDays)
assert.Equal(t, tt.wantWithinCooldown, withinCooldown, "withinCooldown")
assert.Equal(t, tt.wantDaysSincePublish, daysSincePublish, "daysSincePublish")
assert.Equal(t, tt.wantDaysRemaining, daysRemaining, "daysRemaining")
})
}
}
func TestCooldownOldestVersion(t *testing.T) {
now := time.Now()
day := 24 * time.Hour
t.Run("returns version with earliest publish date", func(t *testing.T) {
dates := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-10 * day),
"3.0.0": now.Add(-1 * day),
}
ver, ts := cooldownOldestVersion(dates)
assert.Equal(t, "1.0.0", ver)
assert.False(t, ts.IsZero())
})
t.Run("single version", func(t *testing.T) {
dates := map[string]time.Time{"1.0.0": now.Add(-5 * day)}
ver, _ := cooldownOldestVersion(dates)
assert.Equal(t, "1.0.0", ver)
})
t.Run("empty map returns empty string and zero time", func(t *testing.T) {
ver, ts := cooldownOldestVersion(map[string]time.Time{})
assert.Empty(t, ver)
assert.True(t, ts.IsZero())
})
}
func TestCooldownLatestEligibleVersion(t *testing.T) {
now := time.Now()
day := 24 * time.Hour
t.Run("returns most recently published non-blocked version", func(t *testing.T) {
dates := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-10 * day),
"3.0.0": now.Add(-1 * day),
}
tooNew := map[string]bool{"3.0.0": true}
ver := cooldownLatestEligibleVersion(dates, tooNew)
assert.Equal(t, "2.0.0", ver)
})
t.Run("all versions blocked returns empty string", func(t *testing.T) {
dates := map[string]time.Time{"1.0.0": now, "2.0.0": now.Add(-1 * day)}
tooNew := map[string]bool{"1.0.0": true, "2.0.0": true}
ver := cooldownLatestEligibleVersion(dates, tooNew)
assert.Empty(t, ver)
})
t.Run("empty tooNew returns latest version", func(t *testing.T) {
dates := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-10 * day),
}
ver := cooldownLatestEligibleVersion(dates, map[string]bool{})
assert.Equal(t, "2.0.0", ver)
})
}
+11 -52
View File
@@ -45,6 +45,13 @@ func (h *npmCooldownHandler) HandleMetadataRequest(ctx *proxy.RequestContext, pa
// we'd get raw gzip bytes that fail JSON parsing.
ctx.Headers.Set("Accept-Encoding", "identity")
// Strip conditional-GET headers so the registry cannot return 304 Not Modified.
// A 304 has no body — the modifier would receive an empty body, fail to parse
// it as JSON, and fail-open, letting the client use its cached (unfiltered)
// response. Removing these forces a full 200 response on every request.
ctx.Headers.Del("If-None-Match")
ctx.Headers.Del("If-Modified-Since")
modifier := func(statusCode int, headers http.Header, body []byte) (int, http.Header, []byte, error) {
dates, err := h.parseMetadataTime(body)
if err != nil {
@@ -60,9 +67,9 @@ func (h *npmCooldownHandler) HandleMetadataRequest(ctx *proxy.RequestContext, pa
ctx.RequestID, stripped, packageName, cooldownDays, remaining)
if remaining == 0 && h.statsCollector != nil {
oldestVer, oldestDate := h.oldestVersion(dates)
oldestVer, oldestDate := cooldownOldestVersion(dates)
if oldestVer != "" {
_, daysAgo, daysLeft := h.isWithinCooldown(oldestDate, cooldownDays)
_, daysAgo, daysLeft := cooldownIsWithinWindow(oldestDate, cooldownDays)
h.statsCollector.RecordCooldownBlocked(packageName, oldestVer, oldestDate, daysAgo, daysLeft, cooldownDays)
}
}
@@ -124,7 +131,7 @@ func (h *npmCooldownHandler) parseMetadataTime(body []byte) (map[string]time.Tim
func (h *npmCooldownHandler) stripCooldownVersions(body []byte, dates map[string]time.Time, cooldownDays int) ([]byte, int, int) {
tooNew := make(map[string]bool)
for version, publishDate := range dates {
if withinCooldown, _, _ := h.isWithinCooldown(publishDate, cooldownDays); withinCooldown {
if withinCooldown, _, _ := cooldownIsWithinWindow(publishDate, cooldownDays); withinCooldown {
tooNew[version] = true
}
}
@@ -181,7 +188,7 @@ func (h *npmCooldownHandler) stripCooldownVersions(body []byte, dates map[string
changed := false
for tag, version := range distTags {
if tooNew[version] {
latest := h.latestNonCooldownVersion(dates, tooNew)
latest := cooldownLatestEligibleVersion(dates, tooNew)
if latest != "" {
distTags[tag] = latest
} else {
@@ -209,51 +216,3 @@ func (h *npmCooldownHandler) stripCooldownVersions(body []byte, dates map[string
return result, len(tooNew), remaining
}
// oldestVersion returns the version with the earliest publish date.
// When all versions are blocked by cooldown, this is the version closest
// to exiting the cooldown window (shortest wait for the user).
func (h *npmCooldownHandler) oldestVersion(dates map[string]time.Time) (string, time.Time) {
var oldest string
var oldestTime time.Time
for version, publishDate := range dates {
if oldestTime.IsZero() || publishDate.Before(oldestTime) {
oldest = version
oldestTime = publishDate
}
}
return oldest, oldestTime
}
// isWithinCooldown reports whether a version published at publishDate is still
// within the cooldown window of cooldownDays. It also returns the number of
// whole days since publication.
func (h *npmCooldownHandler) isWithinCooldown(publishDate time.Time, cooldownDays int) (withinCooldown bool, daysSincePublish int, daysRemaining int) {
daysSincePublish = int(time.Since(publishDate).Hours() / 24)
if daysSincePublish < 0 {
daysSincePublish = 0
}
daysRemaining = cooldownDays - daysSincePublish
if daysRemaining < 0 {
daysRemaining = 0
}
return daysSincePublish < cooldownDays, daysSincePublish, daysRemaining
}
func (h *npmCooldownHandler) latestNonCooldownVersion(dates map[string]time.Time, tooNew map[string]bool) string {
var latest string
var latestTime time.Time
for version, publishDate := range dates {
if tooNew[version] {
continue
}
if publishDate.After(latestTime) {
latest = version
latestTime = publishDate
}
}
return latest
}
+4 -113
View File
@@ -298,12 +298,16 @@ func TestNpmCooldown_HandleMetadataRequest_OverridesHeaders(t *testing.T) {
ctx := makeTestRequestContext("https://registry.npmjs.org/lodash")
ctx.Headers.Set("Accept", "application/vnd.npm.install-v1+json")
ctx.Headers.Set("Accept-Encoding", "gzip")
ctx.Headers.Set("If-None-Match", `"abc123"`)
ctx.Headers.Set("If-Modified-Since", "Wed, 01 Jan 2025 00:00:00 GMT")
resp, err := handler.HandleMetadataRequest(ctx, "lodash", 5)
require.NoError(t, err)
assert.Equal(t, proxy.ActionModifyResponse, resp.Action)
assert.Equal(t, "application/json", ctx.Headers.Get("Accept"))
assert.Equal(t, "identity", ctx.Headers.Get("Accept-Encoding"))
assert.Empty(t, ctx.Headers.Get("If-None-Match"))
assert.Empty(t, ctx.Headers.Get("If-Modified-Since"))
}
func TestNpmCooldown_HandleMetadataRequest_StripsRecentVersions(t *testing.T) {
@@ -485,116 +489,3 @@ func TestNpmCooldown_TarballRequestBypassesCooldown(t *testing.T) {
// Accept header should not be set to application/json for tarball requests
assert.NotEqual(t, proxy.ActionModifyResponse, resp.Action)
}
func TestIsWithinCooldown(t *testing.T) {
h := newNpmCooldownHandler(nil)
now := time.Now()
day := 24 * time.Hour
tests := []struct {
name string
publishDate time.Time
cooldownDays int
wantWithinCooldown bool
wantDaysSincePublish int
wantDaysRemaining int
}{
{
name: "published today with 30 day cooldown",
publishDate: now,
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 30,
},
{
name: "published exactly at cooldown boundary",
publishDate: now.Add(-30 * day),
cooldownDays: 30,
wantWithinCooldown: false,
wantDaysSincePublish: 30,
wantDaysRemaining: 0,
},
{
name: "published one day before cooldown expires",
publishDate: now.Add(-29 * day),
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 29,
wantDaysRemaining: 1,
},
{
name: "published well beyond cooldown",
publishDate: now.Add(-365 * day),
cooldownDays: 30,
wantWithinCooldown: false,
wantDaysSincePublish: 365,
wantDaysRemaining: 0,
},
{
name: "zero cooldown days",
publishDate: now,
cooldownDays: 0,
wantWithinCooldown: false,
wantDaysSincePublish: 0,
wantDaysRemaining: 0,
},
{
name: "future publish date is clamped to zero days",
publishDate: now.Add(5 * day),
cooldownDays: 30,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 30,
},
{
name: "large cooldown days that previously caused overflow",
publishDate: now.Add(-10 * day),
cooldownDays: 1000000,
wantWithinCooldown: true,
wantDaysSincePublish: 10,
wantDaysRemaining: 999990,
},
{
name: "max int cooldown days does not overflow",
publishDate: now.Add(-1 * day),
cooldownDays: int(^uint(0) >> 1), // math.MaxInt
wantWithinCooldown: true,
wantDaysSincePublish: 1,
wantDaysRemaining: int(^uint(0)>>1) - 1,
},
{
name: "very old publish date well past cooldown",
publishDate: now.Add(-3652 * day),
cooldownDays: 30,
wantWithinCooldown: false,
wantDaysSincePublish: 3652,
wantDaysRemaining: 0,
},
{
name: "one day cooldown with publish yesterday",
publishDate: now.Add(-1 * day),
cooldownDays: 1,
wantWithinCooldown: false,
wantDaysSincePublish: 1,
wantDaysRemaining: 0,
},
{
name: "one day cooldown with publish today",
publishDate: now,
cooldownDays: 1,
wantWithinCooldown: true,
wantDaysSincePublish: 0,
wantDaysRemaining: 1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
withinCooldown, daysSincePublish, daysRemaining := h.isWithinCooldown(tt.publishDate, tt.cooldownDays)
assert.Equal(t, tt.wantWithinCooldown, withinCooldown, "withinCooldown")
assert.Equal(t, tt.wantDaysSincePublish, daysSincePublish, "daysSincePublish")
assert.Equal(t, tt.wantDaysRemaining, daysRemaining, "daysRemaining")
})
}
}
+205
View File
@@ -0,0 +1,205 @@
package interceptors
import (
"encoding/json"
"fmt"
"net/http"
"time"
"github.com/safedep/dry/log"
"github.com/safedep/pmg/proxy"
)
const pypiSimpleAPIContentType = "application/vnd.pypi.simple.v1+json"
// pypiCooldownHandler handles dependency cooldown for PyPI packages.
// It strips recently-published file entries from PEP 691 Simple API responses
// so pip's resolver naturally falls back to the latest eligible version.
type pypiCooldownHandler struct {
statsCollector *AnalysisStatsCollector
}
func newPypiCooldownHandler(statsCollector *AnalysisStatsCollector) *pypiCooldownHandler {
return &pypiCooldownHandler{statsCollector: statsCollector}
}
// HandleMetadataRequest overrides the Accept header to force a PEP 691 JSON response,
// then registers a response modifier that strips files for versions within the cooldown window.
func (h *pypiCooldownHandler) HandleMetadataRequest(ctx *proxy.RequestContext, packageName string, cooldownDays int) (*proxy.InterceptorResponse, error) {
log.Debugf("[%s] Cooldown: registering metadata modifier for %s", ctx.RequestID, packageName)
// Force PEP 691 JSON so we receive upload-time per file entry.
ctx.Headers.Set("Accept", pypiSimpleAPIContentType)
// Prevent compression so the response body can be parsed as JSON directly.
ctx.Headers.Set("Accept-Encoding", "identity")
// Strip conditional-GET headers so PyPI cannot return 304 Not Modified.
// A 304 has no body — the modifier would receive an empty body, fail to parse
// it as JSON, and fail-open, letting the client use its cached (unfiltered)
// response. Removing these forces a full 200 response on every request.
ctx.Headers.Del("If-None-Match")
ctx.Headers.Del("If-Modified-Since")
modifier := func(statusCode int, headers http.Header, body []byte) (int, http.Header, []byte, error) {
dates, err := h.parsePEP691Files(body)
if err != nil {
log.Warnf("[%s] Cooldown: failed to parse PEP 691 metadata for %s: %v", ctx.RequestID, packageName, err)
return statusCode, headers, body, nil
}
log.Debugf("[%s] Cooldown: parsed %d versions for %s", ctx.RequestID, len(dates), packageName)
strippedBody, stripped, remaining := h.stripCooldownFiles(body, dates, cooldownDays)
if stripped > 0 {
log.Infof("[%s] Cooldown: stripped %d version(s) from %s metadata (%d days, %d eligible remain)",
ctx.RequestID, stripped, packageName, cooldownDays, remaining)
if remaining == 0 && h.statsCollector != nil {
oldestVer, oldestDate := cooldownOldestVersion(dates)
if oldestVer != "" {
_, daysAgo, daysLeft := cooldownIsWithinWindow(oldestDate, cooldownDays)
h.statsCollector.RecordCooldownBlocked(packageName, oldestVer, oldestDate, daysAgo, daysLeft, cooldownDays)
}
}
headers.Set("Cache-Control", "no-store")
return statusCode, headers, strippedBody, nil
}
return statusCode, headers, body, nil
}
return &proxy.InterceptorResponse{
Action: proxy.ActionModifyResponse,
ResponseModifier: modifier,
}, nil
}
// parsePEP691Files extracts the earliest upload-time per version from a PEP 691 JSON body.
// Files with missing or unparseable upload-time are skipped (treated as eligible).
// Multiple files for the same version (sdist + wheels) use the earliest upload-time.
func (h *pypiCooldownHandler) parsePEP691Files(body []byte) (map[string]time.Time, error) {
var resp struct {
Files []struct {
Filename string `json:"filename"`
UploadTime string `json:"upload-time"`
} `json:"files"`
}
if err := json.Unmarshal(body, &resp); err != nil {
return nil, fmt.Errorf("failed to unmarshal PEP 691 response: %w", err)
}
dates := make(map[string]time.Time)
for _, f := range resp.Files {
if f.UploadTime == "" {
log.Debugf("Cooldown: skipping file %s with missing upload-time", f.Filename)
continue
}
t, err := parsePEP691UploadTime(f.UploadTime)
if err != nil {
log.Debugf("Cooldown: skipping file %s with unparseable upload-time %q: %v", f.Filename, f.UploadTime, err)
continue
}
pkgInfo, err := parseFilename(f.Filename)
if err != nil {
log.Debugf("Cooldown: skipping file %s with unparseable filename: %v", f.Filename, err)
continue
}
version := pkgInfo.GetVersion()
if version == "" {
continue
}
// Use the earliest upload-time across all files for a given version
if existing, ok := dates[version]; !ok || t.Before(existing) {
dates[version] = t
}
}
return dates, nil
}
// stripCooldownFiles removes all file entries for versions within the cooldown window
// from a PEP 691 JSON body. Returns the modified body, number of versions stripped,
// and number of versions remaining.
func (h *pypiCooldownHandler) stripCooldownFiles(body []byte, dates map[string]time.Time, cooldownDays int) ([]byte, int, int) {
tooNew := make(map[string]bool)
for version, uploadDate := range dates {
if within, _, _ := cooldownIsWithinWindow(uploadDate, cooldownDays); within {
tooNew[version] = true
}
}
remaining := len(dates) - len(tooNew)
if len(tooNew) == 0 {
return body, 0, remaining
}
var resp map[string]json.RawMessage
if err := json.Unmarshal(body, &resp); err != nil {
log.Warnf("Cooldown: failed to unmarshal PEP 691 body for stripping: %v", err)
return body, 0, remaining
}
rawFiles, ok := resp["files"]
if !ok {
return body, 0, remaining
}
var files []json.RawMessage
if err := json.Unmarshal(rawFiles, &files); err != nil {
log.Warnf("Cooldown: failed to unmarshal files array: %v", err)
return body, 0, remaining
}
filtered := make([]json.RawMessage, 0, len(files))
for _, rawFile := range files {
var f struct {
Filename string `json:"filename"`
}
if err := json.Unmarshal(rawFile, &f); err != nil {
// Keep files we cannot parse to avoid accidentally dropping valid entries
filtered = append(filtered, rawFile)
continue
}
pkgInfo, err := parseFilename(f.Filename)
if err != nil {
// unparseable filename — keep it (fail-open)
filtered = append(filtered, rawFile)
continue
}
if tooNew[pkgInfo.GetVersion()] {
continue // strip
}
filtered = append(filtered, rawFile)
}
updatedFiles, err := json.Marshal(filtered)
if err != nil {
log.Warnf("Cooldown: failed to marshal filtered files array: %v", err)
return body, 0, remaining
}
resp["files"] = updatedFiles
result, err := json.Marshal(resp)
if err != nil {
log.Warnf("Cooldown: failed to marshal final PEP 691 response: %v", err)
return body, 0, remaining
}
return result, len(tooNew), remaining
}
// parsePEP691UploadTime parses the ISO 8601 upload-time field from PEP 691 responses.
// Example: "2023-05-22T15:12:44.000000+00:00"
func parsePEP691UploadTime(s string) (time.Time, error) {
if t, err := time.Parse(time.RFC3339Nano, s); err == nil {
return t, nil
}
return time.Parse(time.RFC3339, s)
}
+559
View File
@@ -0,0 +1,559 @@
package interceptors
import (
"encoding/json"
"fmt"
"net/http"
"testing"
"time"
"github.com/safedep/pmg/config"
"github.com/safedep/pmg/proxy"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// buildTestPEP691Response builds a PEP 691 JSON Simple API response for testing.
// versions maps version string to upload time for a single .tar.gz file per version.
func buildTestPEP691Response(versions map[string]time.Time) []byte {
type fileEntry struct {
Filename string `json:"filename"`
URL string `json:"url"`
UploadTime string `json:"upload-time"`
Hashes map[string]string `json:"hashes"`
}
files := make([]fileEntry, 0, len(versions))
for version, t := range versions {
files = append(files, fileEntry{
Filename: fmt.Sprintf("testpkg-%s.tar.gz", version),
URL: fmt.Sprintf("https://files.pythonhosted.org/packages/testpkg-%s.tar.gz", version),
UploadTime: t.UTC().Format(time.RFC3339Nano),
Hashes: map[string]string{"sha256": "abc123"},
})
}
resp := map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": files,
}
b, _ := json.Marshal(resp)
return b
}
func TestParsePEP691Files_ValidResponse(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now().UTC().Truncate(time.Second)
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-10 * day),
}
body := buildTestPEP691Response(versions)
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
assert.Len(t, dates, 2)
assert.Contains(t, dates, "1.0.0")
assert.Contains(t, dates, "2.0.0")
}
func TestParsePEP691Files_MultipleFilesPerVersion(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now().UTC().Truncate(time.Second)
day := 24 * time.Hour
// Two files for 1.0.0: sdist uploaded 5 days ago, wheel uploaded 3 days ago.
// parsePEP691Files must use the earliest (5 days ago).
sdistTime := now.Add(-5 * day)
wheelTime := now.Add(-3 * day)
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []map[string]any{
{
"filename": "testpkg-1.0.0.tar.gz",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0.tar.gz",
"upload-time": sdistTime.Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "abc"},
},
{
"filename": "testpkg-1.0.0-py3-none-any.whl",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0-py3-none-any.whl",
"upload-time": wheelTime.Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "def"},
},
},
})
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
require.Contains(t, dates, "1.0.0")
// Should use the earliest upload-time (sdist, 5 days ago)
assert.WithinDuration(t, sdistTime, dates["1.0.0"], time.Second)
}
func TestParsePEP691Files_MissingUploadTime(t *testing.T) {
handler := newPypiCooldownHandler(nil)
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []map[string]any{
{
"filename": "testpkg-1.0.0.tar.gz",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0.tar.gz",
"hashes": map[string]string{"sha256": "abc"},
},
{
"filename": "testpkg-2.0.0.tar.gz",
"url": "https://files.pythonhosted.org/packages/testpkg-2.0.0.tar.gz",
"upload-time": time.Now().Add(-10 * 24 * time.Hour).UTC().Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "def"},
},
},
})
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
assert.NotContains(t, dates, "1.0.0")
assert.Contains(t, dates, "2.0.0")
}
func TestParsePEP691Files_YankedFiles(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now().UTC()
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []map[string]any{
{
"filename": "testpkg-1.0.0.tar.gz",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0.tar.gz",
"upload-time": now.Add(-2 * 24 * time.Hour).Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "abc"},
"yanked": true,
},
},
})
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
// Yanked file is still parsed — cooldown applies, pip handles yanked behaviour
assert.Contains(t, dates, "1.0.0")
}
func TestParsePEP691Files_InvalidJSON(t *testing.T) {
handler := newPypiCooldownHandler(nil)
_, err := handler.parsePEP691Files([]byte(`not-json`))
assert.Error(t, err)
}
func TestParsePEP691Files_EmptyFiles(t *testing.T) {
handler := newPypiCooldownHandler(nil)
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []any{},
})
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
assert.Empty(t, dates)
}
func TestStripCooldownFiles_MixedVersions(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-30 * day), // old — eligible
"2.0.0": now.Add(-1 * day), // too new (5d cooldown)
}
body := buildTestPEP691Response(versions)
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
newBody, stripped, remaining := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 1, stripped)
assert.Equal(t, 1, remaining)
var result struct {
Files []struct {
Filename string `json:"filename"`
} `json:"files"`
}
require.NoError(t, json.Unmarshal(newBody, &result))
filenames := make([]string, 0, len(result.Files))
for _, f := range result.Files {
filenames = append(filenames, f.Filename)
}
assert.Contains(t, filenames, "testpkg-1.0.0.tar.gz")
assert.NotContains(t, filenames, "testpkg-2.0.0.tar.gz")
}
func TestStripCooldownFiles_AllVersionsTooNew(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-1 * day),
"2.0.0": now.Add(-2 * day),
}
body := buildTestPEP691Response(versions)
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
newBody, stripped, remaining := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 2, stripped)
assert.Equal(t, 0, remaining)
var result struct {
Files []json.RawMessage `json:"files"`
}
require.NoError(t, json.Unmarshal(newBody, &result))
assert.Empty(t, result.Files)
}
func TestStripCooldownFiles_NoVersionsTooNew(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-10 * day),
"2.0.0": now.Add(-20 * day),
}
body := buildTestPEP691Response(versions)
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
newBody, stripped, remaining := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 0, stripped)
assert.Equal(t, 2, remaining)
assert.Equal(t, body, newBody)
}
func TestStripCooldownFiles_SingleVersionInCooldown(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
versions := map[string]time.Time{
"1.0.0": now.Add(-1 * 24 * time.Hour),
}
body := buildTestPEP691Response(versions)
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
_, stripped, remaining := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 1, stripped)
assert.Equal(t, 0, remaining)
}
func TestStripCooldownFiles_MultipleFilesPerVersion_AllStripped(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []map[string]any{
{
"filename": "testpkg-1.0.0.tar.gz",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0.tar.gz",
"upload-time": now.Add(-1 * 24 * time.Hour).UTC().Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "abc"},
},
{
"filename": "testpkg-1.0.0-py3-none-any.whl",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0-py3-none-any.whl",
"upload-time": now.Add(-2 * 24 * time.Hour).UTC().Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "def"},
},
},
})
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
newBody, stripped, remaining := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 1, stripped) // 1 version stripped
assert.Equal(t, 0, remaining)
var result struct {
Files []json.RawMessage `json:"files"`
}
require.NoError(t, json.Unmarshal(newBody, &result))
assert.Empty(t, result.Files) // both files (sdist + wheel) removed
}
func TestStripCooldownFiles_MalformedJSON(t *testing.T) {
handler := newPypiCooldownHandler(nil)
body := []byte(`not-json`)
dates := map[string]time.Time{"1.0.0": time.Now().Add(-1 * time.Hour)}
newBody, stripped, _ := handler.stripCooldownFiles(body, dates, 5)
assert.Equal(t, 0, stripped)
assert.Equal(t, body, newBody)
}
func TestStripCooldownFiles_UnparseableFilename_KeepFile(t *testing.T) {
handler := newPypiCooldownHandler(nil)
now := time.Now()
// .egg is an unsupported extension — parseFilename will fail
// The file must be kept (fail-open), not stripped
body, _ := json.Marshal(map[string]any{
"meta": map[string]string{"api-version": "1.0"},
"name": "testpkg",
"files": []map[string]any{
{
"filename": "testpkg-1.0.0.egg",
"url": "https://files.pythonhosted.org/packages/testpkg-1.0.0.egg",
"upload-time": now.Add(-1 * 24 * time.Hour).UTC().Format(time.RFC3339Nano),
"hashes": map[string]string{"sha256": "abc"},
},
},
})
// parsePEP691Files will skip the .egg file (no version extracted),
// so dates will be empty — nothing to strip
dates, err := handler.parsePEP691Files(body)
require.NoError(t, err)
assert.Empty(t, dates)
// Force tooNew to include a version that matches nothing, to exercise the
// stripCooldownFiles path with a non-empty tooNew map
forcedDates := map[string]time.Time{
"1.0.0": now.Add(-1 * 24 * time.Hour),
}
newBody, stripped, _ := handler.stripCooldownFiles(body, forcedDates, 5)
assert.Equal(t, 1, stripped) // version is "stripped" from the date map perspective
var result struct {
Files []json.RawMessage `json:"files"`
}
require.NoError(t, json.Unmarshal(newBody, &result))
// The .egg file must still be present — unparseable filename means fail-open
assert.Len(t, result.Files, 1)
}
func TestPyPICooldown_HandleMetadataRequest_OverridesHeaders(t *testing.T) {
handler := newPypiCooldownHandler(nil)
ctx := makeTestRequestContext("https://pypi.org/simple/requests/")
ctx.Headers.Set("Accept", "text/html")
ctx.Headers.Set("Accept-Encoding", "gzip")
ctx.Headers.Set("If-None-Match", `"abc123"`)
ctx.Headers.Set("If-Modified-Since", "Wed, 01 Jan 2025 00:00:00 GMT")
resp, err := handler.HandleMetadataRequest(ctx, "requests", 5)
require.NoError(t, err)
assert.Equal(t, proxy.ActionModifyResponse, resp.Action)
assert.Equal(t, "application/vnd.pypi.simple.v1+json", ctx.Headers.Get("Accept"))
assert.Equal(t, "identity", ctx.Headers.Get("Accept-Encoding"))
assert.Empty(t, ctx.Headers.Get("If-None-Match"))
assert.Empty(t, ctx.Headers.Get("If-Modified-Since"))
}
func TestPyPICooldown_HandleMetadataRequest_NonJSONResponse_FailOpen(t *testing.T) {
handler := newPypiCooldownHandler(nil)
ctx := makeTestRequestContext("https://pypi.org/simple/requests/")
resp, err := handler.HandleMetadataRequest(ctx, "requests", 5)
require.NoError(t, err)
require.NotNil(t, resp.ResponseModifier)
htmlBody := []byte(`<!DOCTYPE html><html><body><a href="/packages/requests-2.31.0.tar.gz">requests-2.31.0.tar.gz</a></body></html>`)
headers := http.Header{}
headers.Set("Content-Type", "text/html")
_, retHeaders, retBody, err := resp.ResponseModifier(200, headers, htmlBody)
require.NoError(t, err)
assert.Equal(t, htmlBody, retBody)
assert.NotEqual(t, "no-store", retHeaders.Get("Cache-Control"))
}
func TestPyPICooldown_HandleMetadataRequest_StripsRecentVersions(t *testing.T) {
now := time.Now()
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-1 * day),
}
body := buildTestPEP691Response(versions)
handler := newPypiCooldownHandler(NewAnalysisStatsCollector())
ctx := makeTestRequestContext("https://pypi.org/simple/testpkg/")
resp, err := handler.HandleMetadataRequest(ctx, "testpkg", 5)
require.NoError(t, err)
require.NotNil(t, resp.ResponseModifier)
headers := http.Header{}
headers.Set("Content-Type", "application/vnd.pypi.simple.v1+json")
_, retHeaders, retBody, err := resp.ResponseModifier(200, headers, body)
require.NoError(t, err)
assert.Equal(t, "no-store", retHeaders.Get("Cache-Control"))
var result struct {
Files []struct {
Filename string `json:"filename"`
} `json:"files"`
}
require.NoError(t, json.Unmarshal(retBody, &result))
filenames := make([]string, 0, len(result.Files))
for _, f := range result.Files {
filenames = append(filenames, f.Filename)
}
assert.Contains(t, filenames, "testpkg-1.0.0.tar.gz")
assert.NotContains(t, filenames, "testpkg-2.0.0.tar.gz")
}
func TestPyPICooldown_HandleMetadataRequest_AllVersionsInCooldown_RecordsStats(t *testing.T) {
now := time.Now()
versions := map[string]time.Time{
"1.0.0": now.Add(-1 * 24 * time.Hour),
}
body := buildTestPEP691Response(versions)
collector := NewAnalysisStatsCollector()
handler := newPypiCooldownHandler(collector)
ctx := makeTestRequestContext("https://pypi.org/simple/newpkg/")
resp, err := handler.HandleMetadataRequest(ctx, "newpkg", 5)
require.NoError(t, err)
require.NotNil(t, resp.ResponseModifier)
headers := http.Header{}
headers.Set("Content-Type", "application/vnd.pypi.simple.v1+json")
_, _, _, err = resp.ResponseModifier(200, headers, body)
require.NoError(t, err)
blocks := collector.GetCooldownBlocks()
require.Len(t, blocks, 1)
assert.Equal(t, "newpkg", blocks[0].Name)
assert.Equal(t, "1.0.0", blocks[0].Version)
assert.Equal(t, 5, blocks[0].CooldownDays)
stats := collector.GetStats()
assert.Equal(t, 1, stats.CooldownBlockedCount)
assert.Equal(t, 1, stats.BlockedCount)
}
func TestPyPICooldown_HandleMetadataRequest_NoVersionsInCooldown_BodyUnchanged(t *testing.T) {
now := time.Now()
day := 24 * time.Hour
versions := map[string]time.Time{
"1.0.0": now.Add(-30 * day),
"2.0.0": now.Add(-20 * day),
}
body := buildTestPEP691Response(versions)
handler := newPypiCooldownHandler(nil)
ctx := makeTestRequestContext("https://pypi.org/simple/testpkg/")
resp, err := handler.HandleMetadataRequest(ctx, "testpkg", 5)
require.NoError(t, err)
require.NotNil(t, resp.ResponseModifier)
headers := http.Header{}
headers.Set("Content-Type", "application/vnd.pypi.simple.v1+json")
_, _, retBody, err := resp.ResponseModifier(200, headers, body)
require.NoError(t, err)
assert.Equal(t, body, retBody)
}
func TestPyPICooldown_HandleMetadataRequest_MalformedJSON_FailOpen(t *testing.T) {
handler := newPypiCooldownHandler(nil)
ctx := makeTestRequestContext("https://pypi.org/simple/badpkg/")
resp, err := handler.HandleMetadataRequest(ctx, "badpkg", 5)
require.NoError(t, err)
require.NotNil(t, resp.ResponseModifier)
body := []byte(`not-json`)
headers := http.Header{}
headers.Set("Content-Type", "application/vnd.pypi.simple.v1+json")
_, _, retBody, err := resp.ResponseModifier(200, headers, body)
require.NoError(t, err)
assert.Equal(t, body, retBody)
}
func TestPyPICooldown_InterceptorDelegation_CooldownEnabled(t *testing.T) {
setCooldownConfig(t, config.DependencyCooldownConfig{Enabled: true, Days: 5})
interceptor := NewPypiRegistryInterceptor(nil, NewInMemoryAnalysisCache(), NewAnalysisStatsCollector(), make(chan *ConfirmationRequest, 1))
ctx := makeTestRequestContext("https://pypi.org/simple/requests/")
ctx.Hostname = "pypi.org"
ctx.Headers.Set("Accept", "text/html")
resp, err := interceptor.HandleRequest(ctx)
require.NoError(t, err)
assert.Equal(t, proxy.ActionModifyResponse, resp.Action)
assert.Equal(t, "application/vnd.pypi.simple.v1+json", ctx.Headers.Get("Accept"))
}
func TestPyPICooldown_InterceptorDelegation_CooldownDisabled(t *testing.T) {
setCooldownConfig(t, config.DependencyCooldownConfig{Enabled: false, Days: 5})
interceptor := NewPypiRegistryInterceptor(nil, NewInMemoryAnalysisCache(), NewAnalysisStatsCollector(), make(chan *ConfirmationRequest, 1))
ctx := makeTestRequestContext("https://pypi.org/simple/requests/")
ctx.Hostname = "pypi.org"
ctx.Headers.Set("Accept", "text/html")
resp, err := interceptor.HandleRequest(ctx)
require.NoError(t, err)
assert.Equal(t, proxy.ActionAllow, resp.Action)
assert.Equal(t, "text/html", ctx.Headers.Get("Accept"))
}
func TestPyPICooldown_JSONAPIRequest_NotIntercepted(t *testing.T) {
setCooldownConfig(t, config.DependencyCooldownConfig{Enabled: true, Days: 5})
interceptor := NewPypiRegistryInterceptor(nil, NewInMemoryAnalysisCache(), NewAnalysisStatsCollector(), make(chan *ConfirmationRequest, 1))
ctx := makeTestRequestContext("https://pypi.org/pypi/requests/json")
ctx.Hostname = "pypi.org"
ctx.Headers.Set("Accept", "application/json")
resp, err := interceptor.HandleRequest(ctx)
require.NoError(t, err)
assert.Equal(t, proxy.ActionAllow, resp.Action)
assert.Equal(t, "application/json", ctx.Headers.Get("Accept"))
}
func TestPyPICooldown_FileDownloadBypassesCooldown(t *testing.T) {
setCooldownConfig(t, config.DependencyCooldownConfig{Enabled: true, Days: 5})
origInsecure := config.Get().InsecureInstallation
config.Get().InsecureInstallation = true
t.Cleanup(func() { config.Get().InsecureInstallation = origInsecure })
interceptor := NewPypiRegistryInterceptor(nil, NewInMemoryAnalysisCache(), NewAnalysisStatsCollector(), make(chan *ConfirmationRequest, 1))
ctx := makeTestRequestContext("https://files.pythonhosted.org/packages/ab/cd/ef/requests-2.31.0-py3-none-any.whl")
ctx.Hostname = "files.pythonhosted.org"
resp, err := interceptor.HandleRequest(ctx)
require.NoError(t, err)
assert.NotEqual(t, proxy.ActionModifyResponse, resp.Action)
assert.NotEqual(t, "application/vnd.pypi.simple.v1+json", ctx.Headers.Get("Accept"))
}
+13 -2
View File
@@ -1,9 +1,12 @@
package interceptors
import (
"strings"
packagev1 "buf.build/gen/go/safedep/api/protocolbuffers/go/safedep/messages/package/v1"
"github.com/safedep/dry/log"
"github.com/safedep/pmg/analyzer"
pmgconfig "github.com/safedep/pmg/config"
"github.com/safedep/pmg/proxy"
)
@@ -35,6 +38,7 @@ var pypiRegistryDomains = registryConfigMap{
// It embeds baseRegistryInterceptor to reuse ecosystem agnostic functionality
type PypiRegistryInterceptor struct {
baseRegistryInterceptor
cooldownHandler *pypiCooldownHandler
}
var _ proxy.Interceptor = (*PypiRegistryInterceptor)(nil)
@@ -55,6 +59,7 @@ func NewPypiRegistryInterceptor(
confirmationChan: confirmationChan,
circuitBreaker: newAnalyzerCircuitBreaker("malysis-analyzer-pypi"),
},
cooldownHandler: newPypiCooldownHandler(statsCollector),
}
}
@@ -105,9 +110,15 @@ func (i *PypiRegistryInterceptor) HandleRequest(ctx *proxy.RequestContext) (*pro
return &proxy.InterceptorResponse{Action: proxy.ActionAllow}, nil
}
// Only analyze actual file downloads (sdist or wheel)
// Metadata requests (Simple API or JSON API) are allowed through
if !pkgInfo.IsFileDownload() {
depCooldownConfig := pmgconfig.Get().Config.DependencyCooldown
// Only apply cooldown to Simple API requests (/simple/{pkg}/) — pip uses these
// for version resolution. JSON API requests (/pypi/{pkg}/json) are allowed through;
// they have a different response structure and pip does not use them for installs.
if depCooldownConfig.Enabled && strings.HasPrefix(ctx.URL.Path, "/simple/") {
return i.cooldownHandler.HandleMetadataRequest(ctx, pkgInfo.GetName(), depCooldownConfig.Days)
}
log.Debugf("[%s] Skipping analysis for metadata request: %s", ctx.RequestID, pkgInfo.GetName())
return &proxy.InterceptorResponse{Action: proxy.ActionAllow}, nil
}