diff --git a/internal/handler/conda.go b/internal/handler/conda.go index f593820c..c9db7a7b 100644 --- a/internal/handler/conda.go +++ b/internal/handler/conda.go @@ -461,7 +461,7 @@ func (h *CondaHandler) filterRepodata(r io.Reader, w io.Writer) error { // filterRepodataWith is filterRepodata deciding by filter. func (h *CondaHandler) filterRepodataWith(r io.Reader, w io.Writer, filter *repodataFilter) error { - return jsonstream.CopyObject(json.NewDecoder(r), w, func(m *jsonstream.Member) error { + return jsonstream.CopyObject(r, w, func(m *jsonstream.Member) error { if m.Key != "packages" && m.Key != "packages.conda" { return m.Copy() } @@ -619,7 +619,7 @@ func (h *CondaHandler) streamRepodataIndex(r *http.Request, upstreamURL string) // is not limited by metadata_max_size. func readRepodataIndex(r io.Reader) (map[string]time.Time, error) { index := make(map[string]time.Time) - err := jsonstream.CopyObject(json.NewDecoder(io.LimitReader(r, maxRepodataSize)), io.Discard, func(m *jsonstream.Member) error { + err := jsonstream.CopyObject(io.LimitReader(r, maxRepodataSize), io.Discard, func(m *jsonstream.Member) error { if m.Key != "packages" && m.Key != "packages.conda" { return m.Copy() } diff --git a/internal/handler/conda_repodata_bench_test.go b/internal/handler/conda_repodata_bench_test.go new file mode 100644 index 00000000..61e68bda --- /dev/null +++ b/internal/handler/conda_repodata_bench_test.go @@ -0,0 +1,59 @@ +package handler + +import ( + "bytes" + "io" + "log/slog" + "os" + "testing" + "time" + + "github.com/git-pkgs/cooldown" +) + +// repodataBenchBody reads the repodata.json named by CONDA_REPODATA_BENCH, +// such as conda-forge's noarch one, and skips the benchmark without it. +func repodataBenchBody(b *testing.B) []byte { + b.Helper() + path := os.Getenv("CONDA_REPODATA_BENCH") + if path == "" { + b.Skip("CONDA_REPODATA_BENCH names no repodata.json") + } + body, err := os.ReadFile(path) + if err != nil { + b.Fatal(err) + } + return body +} + +func BenchmarkCondaFilterRepodata(b *testing.B) { + body := repodataBenchBody(b) + policy, err := cooldown.NewPolicy(&cooldown.Config{Default: "3d"}, nil) + if err != nil { + b.Fatal(err) + } + proxy := testProxy() + proxy.Logger = slog.New(slog.DiscardHandler) + proxy.Cooldown = policy + h := NewCondaHandler(proxy, "http://proxy.example") + now := time.Now() + b.SetBytes(int64(len(body))) + b.ReportAllocs() + for b.Loop() { + if err := h.filterRepodataWith(bytes.NewReader(body), io.Discard, h.newRepodataFilter(now)); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkCondaReadRepodataIndex(b *testing.B) { + body := repodataBenchBody(b) + b.SetBytes(int64(len(body))) + b.ReportAllocs() + for b.Loop() { + index, err := readRepodataIndex(bytes.NewReader(body)) + if err != nil || len(index) == 0 { + b.Fatalf("index of %d files, %v", len(index), err) + } + } +} diff --git a/internal/jsonstream/jsonstream.go b/internal/jsonstream/jsonstream.go index 2d3d1200..cf213ce9 100644 --- a/internal/jsonstream/jsonstream.go +++ b/internal/jsonstream/jsonstream.go @@ -5,7 +5,7 @@ package jsonstream import ( "bytes" - "encoding/json" + "encoding/json/jsontext" "errors" "fmt" "io" @@ -20,26 +20,32 @@ var ( ) // Member is one member of an object being copied. Its value must be consumed -// exactly once, with Raw, Copy or Object. +// exactly once, with Raw, Skip, Copy or Object. type Member struct { Key string - dec *json.Decoder + rawKey []byte + dec *jsontext.Decoder out *objectWriter consumed bool } // Raw reads the member's value without writing it. Follow it with Write to -// keep the member. -func (m *Member) Raw() (json.RawMessage, error) { +// keep the member. The bytes are only valid until the next member is read. +func (m *Member) Raw() ([]byte, error) { if m.consumed { return nil, errConsumed } m.consumed = true - var raw json.RawMessage - if err := m.dec.Decode(&raw); err != nil { - return nil, err + return m.dec.ReadValue() +} + +// Skip reads past the member's value, leaving the member out. +func (m *Member) Skip() error { + if m.consumed { + return errConsumed } - return raw, nil + m.consumed = true + return m.dec.SkipValue() } // Write writes the member with value as its raw JSON value. @@ -47,7 +53,7 @@ func (m *Member) Write(value []byte) error { if !m.consumed { return errUnconsumed } - if err := m.out.key(m.Key); err != nil { + if err := m.out.key(m.rawKey); err != nil { return err } _, err := m.out.w.Write(value) @@ -70,54 +76,66 @@ func (m *Member) Object(fn func(*Member) error) error { return errConsumed } m.consumed = true - if err := m.out.key(m.Key); err != nil { + if err := m.out.key(m.rawKey); err != nil { return err } return copyObject(m.dec, m.out.w, fn, true) } -// CopyObject copies the object read from dec to w, calling fn for each -// member. fn decides what is written by how it consumes the member's value; a -// member whose value fn leaves unconsumed is copied unchanged. -func CopyObject(dec *json.Decoder, w io.Writer, fn func(*Member) error) error { +// CopyObject copies the object read from r to w, calling fn for each member. +// fn decides what is written by how it consumes the member's value; a member +// whose value fn leaves unconsumed is copied unchanged. As with +// encoding/json, of repeated names each is passed to fn, and invalid UTF-8 +// is accepted. +func CopyObject(r io.Reader, w io.Writer, fn func(*Member) error) error { + dec := jsontext.NewDecoder(r, jsontext.AllowDuplicateNames(true), jsontext.AllowInvalidUTF8(true)) if err := copyObject(dec, w, fn, false); err != nil { return err } - if _, err := dec.Token(); err != io.EOF { - if err == nil { - err = errors.New("unexpected data after the top-level object") - } - return err + _, err := dec.ReadToken() + switch err { + case io.EOF: + return nil + case nil: + return errors.New("unexpected data after the top-level object") } - return nil + return err } -func copyObject(dec *json.Decoder, w io.Writer, fn func(*Member) error, allowNull bool) error { - tok, err := dec.Token() +func copyObject(dec *jsontext.Decoder, w io.Writer, fn func(*Member) error, allowNull bool) error { + tok, err := dec.ReadToken() if err != nil { return err } - if tok == nil && allowNull { + if tok.Kind() == 'n' && allowNull { _, err := io.WriteString(w, "null") return err } - if tok != json.Delim('{') { + if tok.Kind() != '{' { return ErrNotObject } if _, err := io.WriteString(w, "{"); err != nil { return err } out := &objectWriter{w: w} - for dec.More() { - tok, err := dec.Token() + m := &Member{dec: dec, out: out} + // keyBuf holds the current key; ReadValue's bytes don't outlast fn. + var keyBuf []byte + for dec.PeekKind() != '}' { + // The key is read raw, so it is written with its own bytes. + rawKey, err := dec.ReadValue() if err != nil { return err } - key, ok := tok.(string) - if !ok { - return fmt.Errorf("object key is %T, not a string", tok) + if rawKey.Kind() != '"' { + return fmt.Errorf("object key is %v, not a string", rawKey.Kind()) } - m := &Member{Key: key, dec: dec, out: out} + keyBuf = append(keyBuf[:0], rawKey...) + // The decoder has checked the key, so the only error left is invalid + // UTF-8, which AppendUnquote reports after replacing it with U+FFFD, + // as encoding/json does. + key, _ := jsontext.AppendUnquote(nil, keyBuf) + *m = Member{Key: string(key), rawKey: keyBuf, dec: dec, out: out} if err := fn(m); err != nil { return err } @@ -127,7 +145,7 @@ func copyObject(dec *json.Decoder, w io.Writer, fn func(*Member) error, allowNul } } } - if _, err := dec.Token(); err != nil { + if _, err := dec.ReadToken(); err != nil { return err } _, err = io.WriteString(w, "}") @@ -137,23 +155,18 @@ func copyObject(dec *json.Decoder, w io.Writer, fn func(*Member) error, allowNul type objectWriter struct { w io.Writer wrote bool + buf bytes.Buffer } -// key writes the separator and key for the next member. Keys are re-encoded -// without HTML escaping, so plain keys keep their original bytes. -func (o *objectWriter) key(k string) error { - var buf bytes.Buffer +// key writes the separator and the key's bytes for the next member. +func (o *objectWriter) key(rawKey []byte) error { + o.buf.Reset() if o.wrote { - buf.WriteByte(',') - } - enc := json.NewEncoder(&buf) - enc.SetEscapeHTML(false) - if err := enc.Encode(k); err != nil { - return err + o.buf.WriteByte(',') } - buf.Truncate(buf.Len() - 1) // Encode appends a newline - buf.WriteByte(':') + o.buf.Write(rawKey) + o.buf.WriteByte(':') o.wrote = true - _, err := o.w.Write(buf.Bytes()) + _, err := o.w.Write(o.buf.Bytes()) return err } diff --git a/internal/jsonstream/jsonstream_test.go b/internal/jsonstream/jsonstream_test.go index 44a877e6..cd201ae2 100644 --- a/internal/jsonstream/jsonstream_test.go +++ b/internal/jsonstream/jsonstream_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "errors" + "slices" "strings" "testing" ) @@ -11,7 +12,7 @@ import ( func copyString(t *testing.T, in string, fn func(*Member) error) (string, error) { t.Helper() var out bytes.Buffer - err := CopyObject(json.NewDecoder(strings.NewReader(in)), &out, fn) + err := CopyObject(strings.NewReader(in), &out, fn) return out.String(), err } @@ -140,3 +141,62 @@ func TestCopyObjectErrors(t *testing.T) { t.Errorf("err = %v, want ErrNotObject", err) } } + +// TestCopyObjectReadsJSONAsEncodingJSON checks CopyObject accepts what +// encoding/json does: repeated names, each passed to fn, and invalid UTF-8. +// Keys keep their bytes, escapes included, while Key is unquoted. +func TestCopyObjectReadsJSONAsEncodingJSON(t *testing.T) { + var keys []string + in := "{\"a\":1,\"a\":2,\"caf\\u00e9\":\"bad \xff byte\",\"skip\":{\"x\":[1,2]}}" + got, err := copyString(t, in, func(m *Member) error { + keys = append(keys, m.Key) + if m.Key == "skip" { + return m.Skip() + } + return nil + }) + if err != nil { + t.Fatal(err) + } + if want := "{\"a\":1,\"a\":2,\"caf\\u00e9\":\"bad \xff byte\"}"; got != want { + t.Errorf("got %q, want %q", got, want) + } + if want := []string{"a", "a", "café", "skip"}; !slices.Equal(keys, want) { + t.Errorf("keys %q, want %q", keys, want) + } +} + +func TestSkipThenWriteReplaces(t *testing.T) { + got, err := copyString(t, `{"a":{"big":[1,2,3]},"b":2}`, func(m *Member) error { + if m.Key != "a" { + return nil + } + if err := m.Skip(); err != nil { + return err + } + return m.Write([]byte(`"small"`)) + }) + if err != nil || got != `{"a":"small","b":2}` { + t.Errorf("got %s, %v", got, err) + } +} + +// TestCopyObjectInvalidUTF8Key checks a key with invalid UTF-8 is accepted: +// Key has U+FFFD in its place, as with encoding/json, and the key is written +// with its original bytes. +func TestCopyObjectInvalidUTF8Key(t *testing.T) { + var keys []string + got, err := copyString(t, "{\"k\xff\":1}", func(m *Member) error { + keys = append(keys, m.Key) + return nil + }) + if err != nil { + t.Fatal(err) + } + if want := "{\"k\xff\":1}"; got != want { + t.Errorf("got %q, want %q", got, want) + } + if want := []string{"k�"}; !slices.Equal(keys, want) { + t.Errorf("keys %q, want %q", keys, want) + } +} diff --git a/internal/metadata/npm/stream.go b/internal/metadata/npm/stream.go index b00f6a23..8788de1b 100644 --- a/internal/metadata/npm/stream.go +++ b/internal/metadata/npm/stream.go @@ -102,7 +102,7 @@ func RewriteStream(packageName string, scan *Scan, r io.Reader, w io.Writer, opt filtered: f, seen: map[string]int{}, } - if err := jsonstream.CopyObject(json.NewDecoder(r), w, sw.member); err != nil { + if err := jsonstream.CopyObject(r, w, sw.member); err != nil { return nil, err } h.tarballsRewritten(sw.rewritten) @@ -162,8 +162,7 @@ func (s *streamWriter) member(m *jsonstream.Member) error { } s.seen[m.Key]++ if s.seen[m.Key] < total { - _, err := m.Raw() // an earlier duplicate, dropped - return err + return m.Skip() // an earlier duplicate, dropped } switch { case m.Key == "versions": @@ -171,7 +170,7 @@ func (s *streamWriter) member(m *jsonstream.Member) error { case m.Key == "time" && s.timeChanged: return m.Object(s.time) case m.Key == "dist-tags" && s.tagsChanged: - if _, err := m.Raw(); err != nil { + if err := m.Skip(); err != nil { return err } encoded, err := json.Marshal(s.distTags) @@ -185,13 +184,13 @@ func (s *streamWriter) member(m *jsonstream.Member) error { // version writes a kept version with its tarball pointing at the proxy. func (s *streamWriter) version(v *jsonstream.Member) error { + if _, keep := s.versions[v.Key]; !keep { + return v.Skip() + } raw, err := v.Raw() if err != nil { return err } - if _, keep := s.versions[v.Key]; !keep { - return nil - } s.entry.Reset() if s.writeNPMVersion(&s.entry, s.packageName, v.Key, raw) { s.rewritten++ @@ -204,8 +203,7 @@ func (s *streamWriter) time(t *jsonstream.Member) error { if _, keep := s.timeMap[t.Key]; keep { return t.Copy() } - _, err := t.Raw() - return err + return t.Skip() } // objectValue records value under name in metadata, as Rewrite's diff --git a/internal/metadata/npm/stream_bench_test.go b/internal/metadata/npm/stream_bench_test.go index 827cf94a..12d04d84 100644 --- a/internal/metadata/npm/stream_bench_test.go +++ b/internal/metadata/npm/stream_bench_test.go @@ -2,6 +2,8 @@ package npm import ( "bytes" + "io" + "os" "testing" "github.com/git-pkgs/proxy/internal/metadata/npm/npmbench" @@ -29,3 +31,29 @@ func BenchmarkLookupVersion(b *testing.B) { } } } + +// BenchmarkRewriteStream times RewriteStream's second pass, the copy, over +// the 5,000-version packument, or the one NPM_PACKUMENT_BENCH names. +func BenchmarkRewriteStream(b *testing.B) { + body := npmbench.Packument(5000) + name := "big" + if path := os.Getenv("NPM_PACKUMENT_BENCH"); path != "" { + var err error + if body, err = os.ReadFile(path); err != nil { + b.Fatal(err) + } + name = os.Getenv("NPM_PACKUMENT_BENCH_NAME") + } + scan, err := ScanPackument(bytes.NewReader(body)) + if err != nil { + b.Fatal(err) + } + options := Options{ProxyURL: "http://proxy.example"} + b.SetBytes(int64(len(body))) + b.ReportAllocs() + for b.Loop() { + if _, err := RewriteStream(name, scan, bytes.NewReader(body), io.Discard, options); err != nil { + b.Fatal(err) + } + } +}