<commit_before>/*
 * Copyright 2018 The Kythe Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package kzip implements the kzip compilation storage file format.
//
// The package exports two types of interest: A kzip.Reader can be used to read
// the contents of an existing kzip archive, and a kzip.Writer can be used to
// construct a new kzip archive.
//
// Reading an Archive:
//
//   r, err := kzip.NewReader(file, size)
//   ...
//
//   // Look up a compilation record by its digest.
//   unit, err := r.Lookup(unitDigest)
//   ...
//
//   // Scan all the compilation records stored.
//   err := r.Scan(func(unit *kzip.Unit) error {
//      if hasInterestingProperty(unit) {
//         doStuffWith(unit)
//      }
//      return nil
//   })
//
//   // Open a reader for a stored file.
//   rc, err := r.Open(fileDigest)
//   ...
//   defer rc.Close()
//
//   // Read the complete contents of a stored file.
//   bits, err := r.ReadAll(fileDigest)
//   ...
//
// Writing an Archive:
//
//   w, err := kzip.NewWriter(file)
//   ...
//
//   // Add a compilation record and (optional) index data.
//   udigest, err := w.AddUnit(unit, nil)
//   ...
//
//   // Add file contents.
//   fdigest, err := w.AddFile(file)
//   ...
//
package kzip // import "kythe.io/kythe/go/platform/kzip"

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"kythe.io/kythe/go/platform/kcd/kythe"

	"bitbucket.org/creachadair/stringset"
	"golang.org/x/sync/errgroup"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"

	apb "kythe.io/kythe/proto/analysis_go_proto"

	// These are common detail messages used by Kythe compilations, and
	// required for JSON (un)marshaling to work.
	_ "kythe.io/kythe/proto/buildinfo_go_proto"
	_ "kythe.io/kythe/proto/cxx_go_proto"
	_ "kythe.io/kythe/proto/filecontext_go_proto"
	_ "kythe.io/kythe/proto/go_go_proto"
	_ "kythe.io/kythe/proto/java_go_proto"
)

// Encoding describes how compilation units will be encoded when written to a kzip.
type Encoding int

const (
	// EncodingJSON specifies to use JSON encoding
	EncodingJSON Encoding = 1
	// EncodingProto specifies to use Proto encoding
	EncodingProto Encoding = 2
	// EncodingAll specifies to encode using all known encodings
	EncodingAll Encoding = EncodingJSON | EncodingProto

	prefixJSON  = "units"
	prefixProto = "pbunits"
)

var (
	// Use a constant file modification time in the kzip so file diffs only compare the contents,
	// not when the kzips were created.
	modifiedTime = time.Unix(0, 0)
)

// EncodingFor converts a string to an Encoding.
func EncodingFor(v string) (Encoding, error) {
	v = strings.ToUpper(v)
	switch {
	case v == "ALL":
		return EncodingAll, nil
	case v == "JSON":
		return EncodingJSON, nil
	case v == "PROTO":
		return EncodingProto, nil
	default:
		return EncodingProto, fmt.Errorf("unknown encoding %s", v)
	}
}

// String stringifies an Encoding
func (e Encoding) String() string {
	switch {
	case e == EncodingAll:
		return "All"
	case e == EncodingJSON:
		return "JSON"
	case e == EncodingProto:
		return "Proto"
	default:
		return "Encoding" + strconv.FormatInt(int64(e), 10)
	}
}

// DefaultEncoding returns the default kzip encoding
func DefaultEncoding() Encoding {
	if e := os.Getenv("KYTHE_KZIP_ENCODING"); e != "" {
		enc, err := EncodingFor(e)
		if err == nil {
			return enc
		}
		log.Printf("Unknown kzip encoding: %s", e)
	}
	return EncodingProto
}

// A Reader permits reading and scanning compilation records and file contents
// stored in a .kzip archive. The Lookup and Scan methods are mutually safe for
// concurrent use by multiple goroutines.
type Reader struct {
	zip *zip.Reader

	// The archives written by this library always use "root/" for the root
	// directory, but it's not required by the spec. Use whatever name the
	// archive actually specifies in the leading directory.
	root string

	// The prefix used for the compilation unit directory; one of
	// prefixJSON or prefixProto
	unitsPrefix string
}

// NewReader constructs a new Reader that consumes zip data from r, whose total
// size in bytes is given.
func NewReader(r io.ReaderAt, size int64) (*Reader, error) {
	archive, err := zip.NewReader(r, size)
	if err != nil {
		return nil, err
	}
	// Order the files in the archive by path, so we can binary search.
	sort.Slice(archive.File, func(i, j int) bool {
		return archive.File[i].Name < archive.File[j].Name
	})

	if len(archive.File) == 0 {
		return nil, errors.New("archive is empty")
	} else if fi := archive.File[0].FileInfo(); !fi.IsDir() {
		return nil, errors.New("archive root is not a directory")
	}
	root := archive.File[0].Name
	pref, err := unitPrefix(root, archive.File)
	if err != nil {
		return nil, err
	}
	return &Reader{
		zip:         archive,
		root:        root,
		unitsPrefix: pref,
	}, nil
}

func unitPrefix(root string, fs []*zip.File) (string, error) {
	jsonDir := root + prefixJSON + "/"
	protoDir := root + prefixProto + "/"
	j := sort.Search(len(fs), func(i int) bool {
		return fs[i].Name > jsonDir
	})
	hasJSON := j < len(fs) && strings.HasPrefix(fs[j].Name, jsonDir)
	p := sort.Search(len(fs), func(i int) bool {
		return fs[i].Name > protoDir
	})
	hasProto := p < len(fs) && strings.HasPrefix(fs[p].Name, protoDir)
	if hasJSON && hasProto {
		// validate that they have identical units based on hash
		for p < len(fs) && j < len(fs) {
			ispb := strings.HasPrefix(fs[p].Name, protoDir)
			isjson := strings.HasPrefix(fs[j].Name, jsonDir)
			if ispb != isjson {<commit_msg>chore: add context to kzip error message (#5168)<commit_after>/*
 * Copyright 2018 The Kythe Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Package kzip implements the kzip compilation storage file format.
//
// The package exports two types of interest: A kzip.Reader can be used to read
// the contents of an existing kzip archive, and a kzip.Writer can be used to
// construct a new kzip archive.
//
// Reading an Archive:
//
//   r, err := kzip.NewReader(file, size)
//   ...
//
//   // Look up a compilation record by its digest.
//   unit, err := r.Lookup(unitDigest)
//   ...
//
//   // Scan all the compilation records stored.
//   err := r.Scan(func(unit *kzip.Unit) error {
//      if hasInterestingProperty(unit) {
//         doStuffWith(unit)
//      }
//      return nil
//   })
//
//   // Open a reader for a stored file.
//   rc, err := r.Open(fileDigest)
//   ...
//   defer rc.Close()
//
//   // Read the complete contents of a stored file.
//   bits, err := r.ReadAll(fileDigest)
//   ...
//
// Writing an Archive:
//
//   w, err := kzip.NewWriter(file)
//   ...
//
//   // Add a compilation record and (optional) index data.
//   udigest, err := w.AddUnit(unit, nil)
//   ...
//
//   // Add file contents.
//   fdigest, err := w.AddFile(file)
//   ...
//
package kzip // import "kythe.io/kythe/go/platform/kzip"

import (
	"archive/zip"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"kythe.io/kythe/go/platform/kcd/kythe"

	"bitbucket.org/creachadair/stringset"
	"golang.org/x/sync/errgroup"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"

	apb "kythe.io/kythe/proto/analysis_go_proto"

	// These are common detail messages used by Kythe compilations, and
	// required for JSON (un)marshaling to work.
	_ "kythe.io/kythe/proto/buildinfo_go_proto"
	_ "kythe.io/kythe/proto/cxx_go_proto"
	_ "kythe.io/kythe/proto/filecontext_go_proto"
	_ "kythe.io/kythe/proto/go_go_proto"
	_ "kythe.io/kythe/proto/java_go_proto"
)

// Encoding describes how compilation units will be encoded when written to a kzip.
type Encoding int

const (
	// EncodingJSON specifies to use JSON encoding
	EncodingJSON Encoding = 1
	// EncodingProto specifies to use Proto encoding
	EncodingProto Encoding = 2
	// EncodingAll specifies to encode using all known encodings
	EncodingAll Encoding = EncodingJSON | EncodingProto

	prefixJSON  = "units"
	prefixProto = "pbunits"
)

var (
	// Use a constant file modification time in the kzip so file diffs only compare the contents,
	// not when the kzips were created.
	modifiedTime = time.Unix(0, 0)
)

// EncodingFor converts a string to an Encoding.
func EncodingFor(v string) (Encoding, error) {
	v = strings.ToUpper(v)
	switch {
	case v == "ALL":
		return EncodingAll, nil
	case v == "JSON":
		return EncodingJSON, nil
	case v == "PROTO":
		return EncodingProto, nil
	default:
		return EncodingProto, fmt.Errorf("unknown encoding %s", v)
	}
}

// String stringifies an Encoding
func (e Encoding) String() string {
	switch {
	case e == EncodingAll:
		return "All"
	case e == EncodingJSON:
		return "JSON"
	case e == EncodingProto:
		return "Proto"
	default:
		return "Encoding" + strconv.FormatInt(int64(e), 10)
	}
}

// DefaultEncoding returns the default kzip encoding
func DefaultEncoding() Encoding {
	if e := os.Getenv("KYTHE_KZIP_ENCODING"); e != "" {
		enc, err := EncodingFor(e)
		if err == nil {
			return enc
		}
		log.Printf("Unknown kzip encoding: %s", e)
	}
	return EncodingProto
}

// A Reader permits reading and scanning compilation records and file contents
// stored in a .kzip archive. The Lookup and Scan methods are mutually safe for
// concurrent use by multiple goroutines.
type Reader struct {
	zip *zip.Reader

	// The archives written by this library always use "root/" for the root
	// directory, but it's not required by the spec. Use whatever name the
	// archive actually specifies in the leading directory.
	root string

	// The prefix used for the compilation unit directory; one of
	// prefixJSON or prefixProto
	unitsPrefix string
}

// NewReader constructs a new Reader that consumes zip data from r, whose total
// size in bytes is given.
func NewReader(r io.ReaderAt, size int64) (*Reader, error) {
	archive, err := zip.NewReader(r, size)
	if err != nil {
		return nil, err
	}
	// Order the files in the archive by path, so we can binary search.
	sort.Slice(archive.File, func(i, j int) bool {
		return archive.File[i].Name < archive.File[j].Name
	})

	if len(archive.File) == 0 {
		return nil, errors.New("archive is empty")
	} else if fi := archive.File[0].FileInfo(); !fi.IsDir() {
		return nil, fmt.Errorf("archive root directory missing: expected a directory but got %v - see https://kythe.io/docs/kythe-kzip.html#_directory_and_file_layout", archive.File[0].Name)
	}
	root := archive.File[0].Name
	pref, err := unitPrefix(root, archive.File)
	if err != nil {
		return nil, err
	}
	return &Reader{
		zip:         archive,
		root:        root,
		unitsPrefix: pref,
	}, nil
}

func unitPrefix(root string, fs []*zip.File) (string, error) {
	jsonDir := root + prefixJSON + "/"
	protoDir := root + prefixProto + "/"
	j := sort.Search(len(fs), func(i int) bool {
		return fs[i].Name > jsonDir
	})
	hasJSON := j < len(fs) && strings.HasPrefix(fs[j].Name, jsonDir)
	p := sort.Search(len(fs), func(i int) bool {
		return fs[i].Name > protoDir
	})
	hasProto := p < len(fs) && strings.HasPrefix(fs[p].Name, protoDir)
	if hasJSON && hasProto {
		// validate that they have identical units based on hash
		for p < len(fs) && j < len(fs) {
			ispb := strings.HasPrefix(fs[p].Name, protoDir)
			isjson := strings.HasPrefix(fs[j].Name, jsonDir)
			if ispb != isjson {
{'fraction_non_alphanumeric': 0.12764084507042253, 'fraction_numerical': 0.004241357234314981, 'mean_word_length': 3.2134187457855696, 'pattern_counts': {'":': 6, '<': 13, '<?xml version=': 0, '>': 7, 'https://': 1, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4097514', 'n_tokens_mistral': 4306, 'n_tokens_neox': 4133, 'n_words': 1374}