Cut memory use of npm metadata rewriting

internal/handler/composer.go

@@ -98,6 +98,12 @@ func (h *ComposerHandler) handlePackageMetadata(w http.ResponseWriter, r *http.R
 
 	upstreamURL := fmt.Sprintf("%s/p2/%s/%s.json", h.repoURL, vendor, pkg)
 
+	if rewritten, ok := h.proxy.storedRewrite("composer", packageName, h.proxyURL, packageName); ok {
+		w.Header().Set(headerContentType, "application/json")
+		_, _ = w.Write(rewritten)
+		return
+	}
+
 	body, _, err := h.proxy.FetchOrCacheMetadata(r.Context(), "composer", packageName, upstreamURL)
 	if err != nil {
 		if errors.Is(err, ErrUpstreamNotFound) {

internal/handler/handler.go

@@ -1383,7 +1383,7 @@ func (p *Proxy) cacheMetadataBlob(ctx context.Context, ecosystem, cacheKey, stor
 		return
 	}
 
-	size, _, err := p.Storage.Store(ctx, storagePath, bytes.NewReader(meta.body))
+	size, hash, err := p.Storage.Store(ctx, storagePath, bytes.NewReader(meta.body))
 	if err != nil {
 		p.Logger.Warn("failed to cache metadata", "ecosystem", ecosystem, "key", cacheKey, "error", err)
 		return
@@ -1396,9 +1396,12 @@ func (p *Proxy) cacheMetadataBlob(ctx context.Context, ecosystem, cacheKey, stor
 		ETag:            sql.NullString{String: meta.etag, Valid: meta.etag != ""},
 		ContentType:     sql.NullString{String: meta.contentType, Valid: meta.contentType != ""},
 		ContentEncoding: sql.NullString{String: meta.contentEncoding, Valid: meta.contentEncoding != ""},
-		Size:            sql.NullInt64{Int64: size, Valid: true},
-		LastModified:    sql.NullTime{Time: meta.lastModified, Valid: !meta.lastModified.IsZero()},
-		FetchedAt:       sql.NullTime{Time: time.Now(), Valid: true},
+		// The digest identifies the stored bytes, so a rewrite cached for them
+		// can be found without reading them back (see storedRewrite).
+		ContentDigest: sql.NullString{String: "sha256:" + hash, Valid: hash != ""},
+		Size:          sql.NullInt64{Int64: size, Valid: true},
+		LastModified:  sql.NullTime{Time: meta.lastModified, Valid: !meta.lastModified.IsZero()},
+		FetchedAt:     sql.NullTime{Time: time.Now(), Valid: true},
 	})
 	if err != nil {
 		// The blob is written but the row describing it is not, so a later

internal/handler/jsonscan.go

@@ -0,0 +1,213 @@
+package handler
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+)
+
+// The helpers in this file walk a JSON document in place, reporting where
+// object members sit as offsets into the original bytes. They let a handler
+// read or replace one value in a large metadata document without decoding the
+// rest of it into Go values, which for an npm packument costs many times the
+// document's size.
+//
+// They check structure only as far as they need to find their way through
+// the document. Callers that copy unvisited bytes into a response validate the
+// whole document first with json.Valid.
+
+var errMalformedJSON = errors.New("malformed JSON")
+
+// errNotJSONObject is returned when a value expected to be an object is not.
+var errNotJSONObject = errors.New("JSON value is not an object")
+
+// jsonMember is one member of a JSON object as offsets into the bytes that
+// were scanned. The key span includes its quotes.
+type jsonMember struct {
+	keyStart, keyEnd int
+	valStart, valEnd int
+}
+
+func (m jsonMember) key(data []byte) []byte   { return data[m.keyStart:m.keyEnd] }
+func (m jsonMember) value(data []byte) []byte { return data[m.valStart:m.valEnd] }
+
+func skipJSONSpace(data []byte, i int) int {
+	for i < len(data) {
+		switch data[i] {
+		case ' ', '\t', '\n', '\r':
+			i++
+		default:
+			return i
+		}
+	}
+	return i
+}
+
+// skipJSONString returns the index just past the string starting at data[i].
+func skipJSONString(data []byte, i int) (int, error) {
+	for j := i + 1; j < len(data); j++ {
+		switch data[j] {
+		case '\\':
+			j++
+		case '"':
+			return j + 1, nil
+		}
+	}
+	return 0, errMalformedJSON
+}
+
+// skipJSONValue returns the index just past the value starting at data[i].
+func skipJSONValue(data []byte, i int) (int, error) {
+	if i >= len(data) {
+		return 0, errMalformedJSON
+	}
+	switch data[i] {
+	case '"':
+		return skipJSONString(data, i)
+	case '{', '[':
+		depth := 0
+		for j := i; j < len(data); j++ {
+			switch data[j] {
+			case '"':
+				end, err := skipJSONString(data, j)
+				if err != nil {
+					return 0, err
+				}
+				j = end - 1
+			case '{', '[':
+				depth++
+			case '}', ']':
+				depth--
+				if depth == 0 {
+					return j + 1, nil
+				}
+			}
+		}
+		return 0, errMalformedJSON
+	case '}', ']', ',', ':':
+		return 0, errMalformedJSON
+	default:
+		// A number, true, false or null runs to the next delimiter.
+		j := i
+		for j < len(data) {
+			switch data[j] {
+			case ',', '}', ']', ' ', '\t', '\n', '\r':
+				return j, nil
+			}
+			j++
+		}
+		return j, nil
+	}
+}
+
+// forEachJSONMember calls fn for each member of the object obj holds, in
+// document order. It returns errNotJSONObject if obj is not an object, and
+// stops early with fn's error if fn returns one.
+func forEachJSONMember(obj []byte, fn func(jsonMember) error) error {
+	i := skipJSONSpace(obj, 0)
+	if i >= len(obj) || obj[i] != '{' {
+		return errNotJSONObject
+	}
+	i = skipJSONSpace(obj, i+1)
+	if i < len(obj) && obj[i] == '}' {
+		return nil
+	}
+	for {
+		if i >= len(obj) || obj[i] != '"' {
+			return errMalformedJSON
+		}
+		keyEnd, err := skipJSONString(obj, i)
+		if err != nil {
+			return err
+		}
+		colon := skipJSONSpace(obj, keyEnd)
+		if colon >= len(obj) || obj[colon] != ':' {
+			return errMalformedJSON
+		}
+		valStart := skipJSONSpace(obj, colon+1)
+		valEnd, err := skipJSONValue(obj, valStart)
+		if err != nil {
+			return err
+		}
+		if err := fn(jsonMember{keyStart: i, keyEnd: keyEnd, valStart: valStart, valEnd: valEnd}); err != nil {
+			return err
+		}
+		next := skipJSONSpace(obj, valEnd)
+		if next >= len(obj) {
+			return errMalformedJSON
+		}
+		switch obj[next] {
+		case ',':
+			i = skipJSONSpace(obj, next+1)
+		case '}':
+			return nil
+		default:
+			return errMalformedJSON
+		}
+	}
+}
+
+// jsonKey decodes a quoted key as forEachJSONMember reports it.
+func jsonKey(raw []byte) (string, error) {
+	if bytes.IndexByte(raw, '\\') < 0 {
+		return string(raw[1 : len(raw)-1]), nil
+	}
+	var key string
+	err := json.Unmarshal(raw, &key)
+	return key, err
+}
+
+// jsonKeyIs reports whether a quoted key decodes to name.
+func jsonKeyIs(raw []byte, name string) bool {
+	if bytes.IndexByte(raw, '\\') < 0 {
+		return string(raw[1:len(raw)-1]) == name
+	}
+	key, err := jsonKey(raw)
+	return err == nil && key == name
+}
+
+// findJSONMember returns the member of obj named name. When the key repeats,
+// the last one wins, as it does for JSON.parse and encoding/json.
+func findJSONMember(obj []byte, name string) (jsonMember, bool, error) {
+	var found jsonMember
+	ok := false
+	err := forEachJSONMember(obj, func(m jsonMember) error {
+		if jsonKeyIs(m.key(obj), name) {
+			found, ok = m, true
+		}
+		return nil
+	})
+	return found, ok, err
+}
+
+// lookupJSON follows path through nested objects in doc and returns the
+// value at its end. It returns nil and no error when a key along the path is
+// missing or names something other than an object.
+func lookupJSON(doc []byte, path ...string) ([]byte, error) {
+	value := doc
+	for _, name := range path {
+		m, ok, err := findJSONMember(value, name)
+		if errors.Is(err, errNotJSONObject) {
+			return nil, nil
+		}
+		if err != nil || !ok {
+			return nil, err
+		}
+		value = m.value(value)
+	}
+	return value, nil
+}
+
+// lookupJSONString is lookupJSON for a string value. ok is false when the
+// value is missing or is not a string.
+func lookupJSONString(doc []byte, path ...string) (string, bool, error) {
+	raw, err := lookupJSON(doc, path...)
+	if err != nil || len(raw) == 0 || raw[0] != '"' {
+		return "", false, err
+	}
+	var s string
+	if err := json.Unmarshal(raw, &s); err != nil {
+		return "", false, err
+	}
+	return s, true, nil
+}

internal/handler/jsonscan_test.go

@@ -0,0 +1,98 @@
+package handler
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+)
+
+func TestLookupJSONString(t *testing.T) {
+	doc := []byte(`{
+		"name": "demo",
+		"versions": {
+			"1.0.0": {"dist": {"tarball": "https://example.com/a.tgz", "shasum": "x"}},
+			"2.0.0": {"readme": "a } tricky \" string ] with {brackets}", "dist": {"tarball": "https://example.com/b.tgz"}},
+			"3.0.0": "not an object",
+			"4.0.0": {"dist": {"tarball": 42}}
+		},
+		"time": {"1.0.0": "2020-01-01T00:00:00Z", "1.0.0": "2021-01-01T00:00:00Z"},
+		"n": -1.5e3, "t": true, "z": null, "list": [1, {"a": []}, "]"]
+	}`)
+
+	cases := []struct {
+		path   []string
+		want   string
+		wantOK bool
+	}{
+		{[]string{"versions", "1.0.0", "dist", "tarball"}, "https://example.com/a.tgz", true},
+		{[]string{"versions", "2.0.0", "dist", "tarball"}, "https://example.com/b.tgz", true},
+		{[]string{"versions", "3.0.0", "dist", "tarball"}, "", false},
+		{[]string{"versions", "4.0.0", "dist", "tarball"}, "", false},
+		{[]string{"versions", "9.9.9", "dist", "tarball"}, "", false},
+		{[]string{"time", "1.0.0"}, "2021-01-01T00:00:00Z", true}, // last duplicate wins
+		{[]string{"name"}, "demo", true},
+		{[]string{"name", "x"}, "", false},
+	}
+	for _, c := range cases {
+		got, ok, err := lookupJSONString(doc, c.path...)
+		if err != nil || ok != c.wantOK || got != c.want {
+			t.Errorf("lookup %v = %q, %v, %v; want %q, %v", c.path, got, ok, err, c.want, c.wantOK)
+		}
+	}
+}
+
+func TestLookupJSONEscapedKey(t *testing.T) {
+	doc := []byte(`{"versions": {"1.0.0": {"dist": {"tarball": "u"}}}}`)
+	got, ok, err := lookupJSONString(doc, "versions", "1.0.0", "dist", "tarball")
+	if err != nil || !ok || got != "u" {
+		t.Errorf("lookup = %q, %v, %v", got, ok, err)
+	}
+}
+
+func TestLookupJSONMalformed(t *testing.T) {
+	for _, doc := range []string{
+		`{"versions": {"1.0.0": {"dist": `,
+		`{"versions" {}}`,
+		`{"versions": {"a": 1 "b": 2}}`,
+		`{"a": "unterminated}`,
+		`{"a": [1, 2}`,
+	} {
+		_, _, err := lookupJSONString([]byte(doc), "versions", "1.0.0", "dist", "tarball")
+		if !errors.Is(err, errMalformedJSON) {
+			t.Errorf("%s: err = %v, want errMalformedJSON", doc, err)
+		}
+	}
+}
+
+func TestForEachJSONMemberOffsets(t *testing.T) {
+	doc := []byte(` { "a" : 1 , "b":{"c":[true,null]} ,"d":"x\"y" } `)
+	var got [][2]string
+	err := forEachJSONMember(doc, func(m jsonMember) error {
+		got = append(got, [2]string{string(m.key(doc)), string(m.value(doc))})
+		return nil
+	})
+	want := [][2]string{{`"a"`, `1`}, {`"b"`, `{"c":[true,null]}`}, {`"d"`, `"x\"y"`}}
+	if err != nil || len(got) != len(want) {
+		t.Fatalf("members = %v, %v", got, err)
+	}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("member %d = %v, want %v", i, got[i], want[i])
+		}
+	}
+
+	if err := forEachJSONMember([]byte(`[1]`), func(jsonMember) error { return nil }); !errors.Is(err, errNotJSONObject) {
+		t.Errorf("array: err = %v, want errNotJSONObject", err)
+	}
+	if err := forEachJSONMember([]byte(`{}`), func(jsonMember) error { t.Error("called for empty object"); return nil }); err != nil {
+		t.Errorf("empty object: err = %v", err)
+	}
+}
+
+func TestWriteFilteredJSONObject(t *testing.T) {
+	var out bytes.Buffer
+	err := writeFilteredJSONObject(&out, []byte(`{"a": 1, "b": {"x": 2}, "c": 3}`), func(k string) bool { return k != "b" })
+	if err != nil || out.String() != `{"a": 1,"c": 3}` {
+		t.Errorf("filtered = %s, %v", out.String(), err)
+	}
+}