<commit_before>/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is murmurhash,
a non-cryptographic hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
Bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"

	"github.com/willf/bitset"
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

func max(x, y uint) uint {
	if x > y {
		return x
	}
	return y
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
// We force _m_ and _k_ to be at least one to avoid panics.
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{max(1, m), max(1, k), bitset.New(m)}
}

// From creates a new Bloom filter with len(_data_) * 64 bits and _k_ hashing
// functions. The data slice is not going to be reset.
func From(data []uint64, k uint) *BloomFilter {
	m := uint(len(data) * 64)
	return &BloomFilter{m, k, bitset.From(data)}
}

// baseHashes returns the four hash values of data that are used to create k
// hashes
func baseHashes(data []byte) [4]uint64 {
	var d digest128;
	hash1, hash2, hash3, hash4 := d.sum256(data)
	return [4]uint64{
		hash1, hash2, hash3, hash4,
	}
}

// location returns the ith hashed location using the four base hash values
func location(h [4]uint64, i uint) uint64 {
	ii := uint64(i)
	return h[ii%2] + ii*h[2+(((ii+(ii%2))%4)/2)]
}

// location returns the ith hashed location using the four base hash values
func (f *BloomFilter) location(h [4]uint64, i uint) uint {
	return uint(location(h, i) % uint64(f.m))
}

// EstimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func EstimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := EstimateParameters(n, fp)
	return New(m, k)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h, i))
	}
	return f
}

// Merge the data from two Bloom Filters.
func (f *BloomFilter) Merge(g *BloomFilter) error {
	// Make sure the m's and k's are the same, otherwise merging has no real use.
	if f.m != g.m {
		return fmt.Errorf("m's don't match: %d != %d", f.m, g.m)
	}

	if f.k != g.k {
		return fmt.Errorf("k's don't match: %d != %d", f.m, g.m)
	}

	f.b.InPlaceUnion(g.b)
	return nil
}

// Copy creates a copy of a Bloom filter.
func (f *BloomFilter) Copy() *BloomFilter {
	fc := New(f.m, f.k)
	fc.Merge(f) // #nosec
	return fc
}

// AddString to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) AddString(data string) *BloomFilter {
	return f.Add([]byte(data))
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) bool {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h, i)) {
			return false
		}
	}
	return true
}

// TestString returns true if the string is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) TestString(data string) bool {
	return f.Test([]byte(data))
}

// TestLocations returns true if all locations are set in the BloomFilter, false
// otherwise.
func (f *BloomFilter) TestLocations(locs []uint64) bool {
	for i := 0; i < len(locs); i++ {
		if !f.b.Test(uint(locs[i] % uint64(f.m))) {
			return false
		}
	}
	return true
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// TestAndAddString is the equivalent to calling Test(string) then Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestAndAddString(data string) bool {
	return f.TestAndAdd([]byte(data))
}

// ClearAll clears all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// EstimateFalsePositiveRate returns, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// while storing n entries; runs 100,000 tests. This is an empirical
// test using integers as keys. As a side-effect, it clears the BloomFilter.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}

// Equal tests for the equality of two Bloom filters
func (f *BloomFilter) Equal(g *BloomFilter) bool {
	return f.m == g.m && f.k == g.k && f.b.Equal(g.b)
}

// Locations returns a list of hash locations representing a data item.
func Locations(data []byte, k uint) []uint64 {
	locs := make([]uint64, k)

	// calculate locations
	h := baseHashes(data)
	for i := uint(0); i < k; i++ {
		locs[i] = location(h, i)
	}

	return locs
}
<commit_msg>Removing ';', adding comment.<commit_after>/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is murmurhash,
a non-cryptographic hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
Bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"

	"github.com/willf/bitset"
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

func max(x, y uint) uint {
	if x > y {
		return x
	}
	return y
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
// We force _m_ and _k_ to be at least one to avoid panics.
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{max(1, m), max(1, k), bitset.New(m)}
}

// From creates a new Bloom filter with len(_data_) * 64 bits and _k_ hashing
// functions. The data slice is not going to be reset.
func From(data []uint64, k uint) *BloomFilter {
	m := uint(len(data) * 64)
	return &BloomFilter{m, k, bitset.From(data)}
}

// baseHashes returns the four hash values of data that are used to create k
// hashes
func baseHashes(data []byte) [4]uint64 {
	var d digest128 // murmur hashing
	hash1, hash2, hash3, hash4 := d.sum256(data)
	return [4]uint64{
		hash1, hash2, hash3, hash4,
	}
}

// location returns the ith hashed location using the four base hash values
func location(h [4]uint64, i uint) uint64 {
	ii := uint64(i)
	return h[ii%2] + ii*h[2+(((ii+(ii%2))%4)/2)]
}

// location returns the ith hashed location using the four base hash values
func (f *BloomFilter) location(h [4]uint64, i uint) uint {
	return uint(location(h, i) % uint64(f.m))
}

// EstimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func EstimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := EstimateParameters(n, fp)
	return New(m, k)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h, i))
	}
	return f
}

// Merge the data from two Bloom Filters.
func (f *BloomFilter) Merge(g *BloomFilter) error {
	// Make sure the m's and k's are the same, otherwise merging has no real use.
	if f.m != g.m {
		return fmt.Errorf("m's don't match: %d != %d", f.m, g.m)
	}

	if f.k != g.k {
		return fmt.Errorf("k's don't match: %d != %d", f.m, g.m)
	}

	f.b.InPlaceUnion(g.b)
	return nil
}

// Copy creates a copy of a Bloom filter.
func (f *BloomFilter) Copy() *BloomFilter {
	fc := New(f.m, f.k)
	fc.Merge(f) // #nosec
	return fc
}

// AddString to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) AddString(data string) *BloomFilter {
	return f.Add([]byte(data))
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) bool {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h, i)) {
			return false
		}
	}
	return true
}

// TestString returns true if the string is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) TestString(data string) bool {
	return f.Test([]byte(data))
}

// TestLocations returns true if all locations are set in the BloomFilter, false
// otherwise.
func (f *BloomFilter) TestLocations(locs []uint64) bool {
	for i := 0; i < len(locs); i++ {
		if !f.b.Test(uint(locs[i] % uint64(f.m))) {
			return false
		}
	}
	return true
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// TestAndAddString is the equivalent to calling Test(string) then Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestAndAddString(data string) bool {
	return f.TestAndAdd([]byte(data))
}

// ClearAll clears all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// EstimateFalsePositiveRate returns, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// while storing n entries; runs 100,000 tests. This is an empirical
// test using integers as keys. As a side-effect, it clears the BloomFilter.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}

// Equal tests for the equality of two Bloom filters
func (f *BloomFilter) Equal(g *BloomFilter) bool {
	return f.m == g.m && f.k == g.k && f.b.Equal(g.b)
}

// Locations returns a list of hash locations representing a data item.
func Locations(data []byte, k uint) []uint64 {
	locs := make([]uint64, k)

	// calculate locations
	h := baseHashes(data)
	for i := uint(0); i < k; i++ {
		locs[i] = location(h, i)
	}

	return locs
}
<|endoftext|>


<commit_before>package main

/*
This interface reflects the kinds of validations provided by the validator
package (gopkg.in/go-playground/validator.v9).

This interface is kept distinct from the rest for the sake of code
organization and documentation.

Some of the material in the comments is from the validator package (above)
and other material is from the JSON schema validation RFC
(http://json-schema.org/latest/json-schema-validation.html). The material is
annotated to avoid confusion, and some modification have been made for the
same reason.

For compatibility with the Go JSON spec library (github.com/go-openapi/spec),
the return types they are what they are.
*/
type Validator interface {

	/*
		Validator Package Documentation:

			This validates that the value is not the data types default zero
			value. For numbers ensures value is not zero. For strings ensures
			value is not "". For slices, maps, pointers, interfaces, channels
			and functions ensures the value is not nil.

		JSON Schema Validation RFC:

			6.17. required
				The value of this keyword MUST be an array. Elements of this array, if any, MUST be strings, and MUST be unique.
				An object instance is valid against this keyword if every item in the array is the name of a property in the instance.
				Omitting this keyword has the same behavior as an empty array.

		Returns -1 when no length is enforced.
	*/
	IsRequired() bool

	/*
		Validator Package Documentation:

			For numbers, length will ensure that the value is equal to the
			parameter given. For strings, it checks that the string length is
			exactly that number of characters. For slices, arrays, and maps,
			validates the number of items.

		There doesn't appear to be a JSON schema equivalent to Length() for
		string types. Thus, using this with strings should result in the
		equivalent of both a Min() and Max() returning the same value.
	*/
	Length() float64

	/*
		Validator Package Documentation:

			For numbers, max will ensure that the value is less than or equal to
			the parameter given. For strings, it checks that the string length
			is at most that number of characters. For slices, arrays, and maps,
			validates the number of items.

		JSON Schema Validation RFC:

			6.2. maximum
				The value of "maximum" MUST be a number, representing an inclusive upper limit for a numeric instance.
				If the instance is a number, then this keyword validates only if the instance is less than or exactly equal to "maximum".
			6.6. maxLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is less than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
			6.11. maxItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "maxItems" if its size is less than, or equal to, the value of this keyword.

		Returns -1 when this validation is not enforced.
	*/
	Max() float64

	/*
		Validator Package Documentation:

			For numbers, min will ensure that the value is greater or equal to
			the parameter given. For strings, it checks that the string length
			is at least that number of characters. For slices, arrays, and maps,
			validates the number of items.

		JSON Schema Validation RFC:

			6.4. minimum
				The value of "minimum" MUST be a number, representing an inclusive upper limit for a numeric instance.
				If the instance is a number, then this keyword validates only if the instance is greater than or exactly equal to "minimum".
			6.7. minLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is greater than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
				Omitting this keyword has the same behavior as a value of 0.
			6.12. minItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "minItems" if its size is greater than, or equal to, the value of this keyword.
				Omitting this keyword has the same behavior as a value of 0.

		Returns -1 when this validation is not enforced.
	*/
	Min() float64

	/*
		Validator Package Documentation:

			For strings & numbers, eq will ensure that the value is equal to the
			parameter given. For slices, arrays, and maps, validates the number
			of items.

		The closest JSON schema equivalent to Equals() for string types is the
		regex pattern. Consequently, be careful what you equals for. For numeric
		values, using this should have the equivalent result of setting Min()
		and Max() with the same value.

		The returned boolean indicates if the validation is present from the
		source. A string is returned containing the raw value from the source.
		No guarantees are made about the parsability of the returned value in
		the case of numeric and plural types.
	*/
	Equals() (string, bool)

	/*
		Validator Package Documentation:

			For strings & numbers, this will ensure that the value is not equal
			to the parameter given. For slices, arrays, and maps, validates the
			number of items.

		There doesn't appear to be a JSON schema equivalent to NotEqual(). I
		don't know of a way to emulate the behavior without a complex boolean
		expression. Therefore, this method will not be implemented.
	*/
	// NotEqual() string

	/*
		Validator Package Documentation:

			For numbers, this will ensure that the value is greater than the
			parameter given. For strings, it checks that the string length is
			greater than that number of characters. For slices, arrays and maps
			it validates the number of items.

		JSON Schema Validation RFC:

			6.5. exclusiveMinimum
				The value of "exclusiveMinimum" MUST be number, representing an exclusive upper limit for a numeric instance.
				If the instance is a number, then the instance is valid only if it has a value strictly greater than (not equal to) "exclusiveMinimum".
			6.7. minLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is greater than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
				Omitting this keyword has the same behavior as a value of 0.
			6.12. minItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "minItems" if its size is greater than, or equal to, the value of this keyword.
				Omitting this keyword has the same behavior as a value of 0.

		Returns -1 when this validation is not enforced.
	*/
	GreaterThan() float64

	/*
		Validator Package Documentation:

			For numbers, this will ensure that the value is less than the
			parameter given. For strings, it checks that the string length is
			less than that number of characters. For slices, arrays, and maps it
			validates the number of items.

		JSON Schema Validation RFC:

			6.3. exclusiveMaximum
				The value of "exclusiveMaximum" MUST be number, representing an exclusive upper limit for a numeric instance.
				If the instance is a number, then the instance is valid only if it has a value strictly less than (not equal to) "exclusiveMaximum".
			6.6. maxLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is less than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
			6.11. maxItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "maxItems" if its size is less than, or equal to, the value of this keyword.

		Returns -1 when this validation is not enforced.
	*/
	LessThan() float64

	// The following are redundant, and their equivalent expressions in the
	// Validator package will be interpreted and returned with the Min() and
	// Max() accessors.
	//
	// GreaterThanOrEqual() string // Same as 'min' above. Kept both to make terminology with 'len' easier.
	// LessThanOrEqual() string // Same as 'max' above. Kept both to make terminology with 'len' easier.
}
<commit_msg>notes for more validations.<commit_after>package main

/*
This interface reflects the kinds of validations provided by the validator
package (gopkg.in/go-playground/validator.v9).

This interface is kept distinct from the rest for the sake of code
organization and documentation.

Some of the material in the comments is from the validator package (above)
and other material is from the JSON schema validation RFC
(http://json-schema.org/latest/json-schema-validation.html). The material is
annotated to avoid confusion, and some modification have been made for the
same reason.

For compatibility with the Go JSON spec library (github.com/go-openapi/spec),
the return types they are what they are.

The site,
https://spacetelescope.github.io/understanding-json-schema/reference/string.html,
may have clues to creating more interesting validations (email, urls, etc.).
*/
type Validator interface {

	/*
		Validator Package Documentation:

			This validates that the value is not the data types default zero
			value. For numbers ensures value is not zero. For strings ensures
			value is not "". For slices, maps, pointers, interfaces, channels
			and functions ensures the value is not nil.

		JSON Schema Validation RFC:

			6.17. required
				The value of this keyword MUST be an array. Elements of this array, if any, MUST be strings, and MUST be unique.
				An object instance is valid against this keyword if every item in the array is the name of a property in the instance.
				Omitting this keyword has the same behavior as an empty array.

		Returns -1 when no length is enforced.
	*/
	IsRequired() bool

	/*
		Validator Package Documentation:

			For numbers, length will ensure that the value is equal to the
			parameter given. For strings, it checks that the string length is
			exactly that number of characters. For slices, arrays, and maps,
			validates the number of items.

		There doesn't appear to be a JSON schema equivalent to Length() for
		string types. Thus, using this with strings should result in the
		equivalent of both a Min() and Max() returning the same value.
	*/
	Length() float64

	/*
		Validator Package Documentation:

			For numbers, max will ensure that the value is less than or equal to
			the parameter given. For strings, it checks that the string length
			is at most that number of characters. For slices, arrays, and maps,
			validates the number of items.

		JSON Schema Validation RFC:

			6.2. maximum
				The value of "maximum" MUST be a number, representing an inclusive upper limit for a numeric instance.
				If the instance is a number, then this keyword validates only if the instance is less than or exactly equal to "maximum".
			6.6. maxLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is less than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
			6.11. maxItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "maxItems" if its size is less than, or equal to, the value of this keyword.

		Returns -1 when this validation is not enforced.
	*/
	Max() float64

	/*
		Validator Package Documentation:

			For numbers, min will ensure that the value is greater or equal to
			the parameter given. For strings, it checks that the string length
			is at least that number of characters. For slices, arrays, and maps,
			validates the number of items.

		JSON Schema Validation RFC:

			6.4. minimum
				The value of "minimum" MUST be a number, representing an inclusive upper limit for a numeric instance.
				If the instance is a number, then this keyword validates only if the instance is greater than or exactly equal to "minimum".
			6.7. minLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is greater than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
				Omitting this keyword has the same behavior as a value of 0.
			6.12. minItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "minItems" if its size is greater than, or equal to, the value of this keyword.
				Omitting this keyword has the same behavior as a value of 0.

		Returns -1 when this validation is not enforced.
	*/
	Min() float64

	/*
		Validator Package Documentation:

			For strings & numbers, eq will ensure that the value is equal to the
			parameter given. For slices, arrays, and maps, validates the number
			of items.

		The closest JSON schema equivalent to Equals() for string types is the
		regex pattern. Consequently, be careful what you equals for. For numeric
		values, using this should have the equivalent result of setting Min()
		and Max() with the same value.

		The returned boolean indicates if the validation is present from the
		source. A string is returned containing the raw value from the source.
		No guarantees are made about the parsability of the returned value in
		the case of numeric and plural types.
	*/
	Equals() (string, bool)

	/*
		Validator Package Documentation:

			For strings & numbers, this will ensure that the value is not equal
			to the parameter given. For slices, arrays, and maps, validates the
			number of items.

		There doesn't appear to be a JSON schema equivalent to NotEqual(). I
		don't know of a way to emulate the behavior without a complex boolean
		expression. Therefore, this method will not be implemented.
	*/
	// NotEqual() string

	/*
		Validator Package Documentation:

			For numbers, this will ensure that the value is greater than the
			parameter given. For strings, it checks that the string length is
			greater than that number of characters. For slices, arrays and maps
			it validates the number of items.

		JSON Schema Validation RFC:

			6.5. exclusiveMinimum
				The value of "exclusiveMinimum" MUST be number, representing an exclusive upper limit for a numeric instance.
				If the instance is a number, then the instance is valid only if it has a value strictly greater than (not equal to) "exclusiveMinimum".
			6.7. minLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is greater than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
				Omitting this keyword has the same behavior as a value of 0.
			6.12. minItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "minItems" if its size is greater than, or equal to, the value of this keyword.
				Omitting this keyword has the same behavior as a value of 0.

		Returns -1 when this validation is not enforced.
	*/
	GreaterThan() float64

	/*
		Validator Package Documentation:

			For numbers, this will ensure that the value is less than the
			parameter given. For strings, it checks that the string length is
			less than that number of characters. For slices, arrays, and maps it
			validates the number of items.

		JSON Schema Validation RFC:

			6.3. exclusiveMaximum
				The value of "exclusiveMaximum" MUST be number, representing an exclusive upper limit for a numeric instance.
				If the instance is a number, then the instance is valid only if it has a value strictly less than (not equal to) "exclusiveMaximum".
			6.6. maxLength
				The value of this keyword MUST be a non-negative integer.
				A string instance is valid against this keyword if its length is less than, or equal to, the value of this keyword.
				The length of a string instance is defined as the number of its characters as defined by RFC 7159 [RFC7159].
			6.11. maxItems
				The value of this keyword MUST be a non-negative integer.
				An array instance is valid against "maxItems" if its size is less than, or equal to, the value of this keyword.

		Returns -1 when this validation is not enforced.
	*/
	LessThan() float64

	// The following are redundant, and their equivalent expressions in the
	// Validator package will be interpreted and returned with the Min() and
	// Max() accessors.
	//
	// GreaterThanOrEqual() string // Same as 'min' above. Kept both to make terminology with 'len' easier.
	// LessThanOrEqual() string // Same as 'max' above. Kept both to make terminology with 'len' easier.
}
<|endoftext|>


<commit_before>package accounts

import (
	"net/url"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/globals"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
)

// Account holds configuration information for an account
type Account struct {
	DocID       string                 `json:"_id,omitempty"`
	DocRev      string                 `json:"_rev,omitempty"`
	Name        string                 `json:"name"`
	AccountType string                 `json:"account_type"`
	Basic       *BasicInfo             `json:"auth,omitempty"`
	Oauth       *OauthInfo             `json:"oauth,omitempty"`
	Extras      map[string]interface{} `json:"oauth_callback_results,omitempty"`
}

// OauthInfo holds configuration information for an oauth account
type OauthInfo struct {
	AccessToken  string      `json:"access_token,omitempty"`
	TokenType    string      `json:"token_type,omitempty"`
	ExpiresAt    time.Time   `json:"expires_at,omitempty"`
	RefreshToken string      `json:"refresh_token,omitempty"`
	ClientID     string      `json:"client_id,omitempty"`
	ClientSecret string      `json:"client_secret,omitempty"`
	Query        *url.Values `json:"query,omitempty"`
}

// BasicInfo holds configuration information for an user/pass account
type BasicInfo struct {
	Login    string `json:"login,omitempty"`
	Password string `json:"password,omitempty"`
}

// ID is used to implement the couchdb.Doc interface
func (ac *Account) ID() string { return ac.DocID }

// Rev is used to implement the couchdb.Doc interface
func (ac *Account) Rev() string { return ac.DocRev }

// SetID is used to implement the couchdb.Doc interface
func (ac *Account) SetID(id string) { ac.DocID = id }

// SetRev is used to implement the couchdb.Doc interface
func (ac *Account) SetRev(rev string) { ac.DocRev = rev }

// DocType implements couchdb.Doc
func (ac *Account) DocType() string { return consts.Accounts }

// Clone implements couchdb.Doc
func (ac *Account) Clone() couchdb.Doc {
	cloned := *ac
	if ac.Oauth != nil {
		tmp := *ac.Oauth
		cloned.Oauth = &tmp
	}
	if ac.Basic != nil {
		tmp := *ac.Basic
		cloned.Basic = &tmp
	}
	cloned.Extras = make(map[string]interface{})
	for k, v := range ac.Extras {
		cloned.Extras[k] = v
	}
	return &cloned
}

// Valid implements permissions.Validable
func (ac *Account) Valid(field, expected string) bool {
	return field == "account_type" && expected == ac.AccountType
}

func init() {
	couchdb.AddHook(consts.Accounts, couchdb.EventDelete,
		func(domain string, doc couchdb.Doc, old couchdb.Doc) error {
			sched := globals.GetScheduler()
			broker := globals.GetBroker()

			trigs, err := sched.GetAll(domain)
			if err != nil {
				logger.WithDomain(domain).Error(
					"Failed to fetch triggers after account deletion: ", err)
				return err
			}
			for _, t := range trigs {
				toDelete := false

				if t.Infos().WorkerType != "konnector" {
					continue
				}

				var msg struct {
					Account string `json:"account"`
				}
				err := t.Infos().Message.Unmarshal(&msg)
				if err != nil {
					toDelete = true
				} else if msg.Account == doc.ID() {
					toDelete = true
				}
				if toDelete {
					if err := sched.Delete(domain, t.ID()); err != nil {
						logger.WithDomain(domain).Errorln("failed to delete orphan trigger", err)
					}
				}
			}

			// When an account is deleted, we need to push a new job in order to
			// delete possible data associated with this account. This is done via
			// this hook.
			//
			// This may require additionnal specifications to allow konnectors to
			// define more explicitly when and how they want to be called in order to
			// cleanup or update their associated content. For now we make this
			// process really specific to the deletion of an account, which is our
			// only detailed usecase.
			if old != nil {
				log := logger.WithDomain(domain).
					WithField("account_id", old.ID()).
					WithField("account_rev", old.Rev())

				var konnector string
				switch v := old.(type) {
				case *Account:
					konnector = v.AccountType
				case *couchdb.JSONDoc:
					konnector, _ = v.M["account_type"].(string)
				case couchdb.JSONDoc:
					konnector, _ = v.M["account_type"].(string)
				}
				if konnector == "" {
					log.Info("No associated konnector for account: " +
						"cannot create on_delete job")
					return nil
				}

				logger.
					WithField("konnector", konnector).
					Info("Pusing job for konnector on_delete")

				msg, err := jobs.NewMessage(struct {
					Account        string `json:"account"`
					AccountRev     string `json:"account_rev"`
					Konnector      string `json:"konnector"`
					AccountDeleted bool   `json:"account_deleted"`
				}{
					Account:        old.ID(),
					AccountRev:     old.Rev(),
					Konnector:      konnector,
					AccountDeleted: true,
				})
				if err != nil {
					return err
				}
				if _, err = broker.PushJob(&jobs.JobRequest{
					Domain:     domain,
					WorkerType: "konnector",
					Message:    msg,
				}); err != nil {
					return err
				}
			}

			return nil
		})
}
<commit_msg>Fix typo<commit_after>package accounts

import (
	"net/url"
	"time"

	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/globals"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
)

// Account holds configuration information for an account
type Account struct {
	DocID       string                 `json:"_id,omitempty"`
	DocRev      string                 `json:"_rev,omitempty"`
	Name        string                 `json:"name"`
	AccountType string                 `json:"account_type"`
	Basic       *BasicInfo             `json:"auth,omitempty"`
	Oauth       *OauthInfo             `json:"oauth,omitempty"`
	Extras      map[string]interface{} `json:"oauth_callback_results,omitempty"`
}

// OauthInfo holds configuration information for an oauth account
type OauthInfo struct {
	AccessToken  string      `json:"access_token,omitempty"`
	TokenType    string      `json:"token_type,omitempty"`
	ExpiresAt    time.Time   `json:"expires_at,omitempty"`
	RefreshToken string      `json:"refresh_token,omitempty"`
	ClientID     string      `json:"client_id,omitempty"`
	ClientSecret string      `json:"client_secret,omitempty"`
	Query        *url.Values `json:"query,omitempty"`
}

// BasicInfo holds configuration information for an user/pass account
type BasicInfo struct {
	Login    string `json:"login,omitempty"`
	Password string `json:"password,omitempty"`
}

// ID is used to implement the couchdb.Doc interface
func (ac *Account) ID() string { return ac.DocID }

// Rev is used to implement the couchdb.Doc interface
func (ac *Account) Rev() string { return ac.DocRev }

// SetID is used to implement the couchdb.Doc interface
func (ac *Account) SetID(id string) { ac.DocID = id }

// SetRev is used to implement the couchdb.Doc interface
func (ac *Account) SetRev(rev string) { ac.DocRev = rev }

// DocType implements couchdb.Doc
func (ac *Account) DocType() string { return consts.Accounts }

// Clone implements couchdb.Doc
func (ac *Account) Clone() couchdb.Doc {
	cloned := *ac
	if ac.Oauth != nil {
		tmp := *ac.Oauth
		cloned.Oauth = &tmp
	}
	if ac.Basic != nil {
		tmp := *ac.Basic
		cloned.Basic = &tmp
	}
	cloned.Extras = make(map[string]interface{})
	for k, v := range ac.Extras {
		cloned.Extras[k] = v
	}
	return &cloned
}

// Valid implements permissions.Validable
func (ac *Account) Valid(field, expected string) bool {
	return field == "account_type" && expected == ac.AccountType
}

func init() {
	couchdb.AddHook(consts.Accounts, couchdb.EventDelete,
		func(domain string, doc couchdb.Doc, old couchdb.Doc) error {
			sched := globals.GetScheduler()
			broker := globals.GetBroker()

			trigs, err := sched.GetAll(domain)
			if err != nil {
				logger.WithDomain(domain).Error(
					"Failed to fetch triggers after account deletion: ", err)
				return err
			}
			for _, t := range trigs {
				toDelete := false

				if t.Infos().WorkerType != "konnector" {
					continue
				}

				var msg struct {
					Account string `json:"account"`
				}
				err := t.Infos().Message.Unmarshal(&msg)
				if err != nil {
					toDelete = true
				} else if msg.Account == doc.ID() {
					toDelete = true
				}
				if toDelete {
					if err := sched.Delete(domain, t.ID()); err != nil {
						logger.WithDomain(domain).Errorln("failed to delete orphan trigger", err)
					}
				}
			}

			// When an account is deleted, we need to push a new job in order to
			// delete possible data associated with this account. This is done via
			// this hook.
			//
			// This may require additionnal specifications to allow konnectors to
			// define more explicitly when and how they want to be called in order to
			// cleanup or update their associated content. For now we make this
			// process really specific to the deletion of an account, which is our
			// only detailed usecase.
			if old != nil {
				log := logger.WithDomain(domain).
					WithField("account_id", old.ID()).
					WithField("account_rev", old.Rev())

				var konnector string
				switch v := old.(type) {
				case *Account:
					konnector = v.AccountType
				case *couchdb.JSONDoc:
					konnector, _ = v.M["account_type"].(string)
				case couchdb.JSONDoc:
					konnector, _ = v.M["account_type"].(string)
				}
				if konnector == "" {
					log.Info("No associated konnector for account: " +
						"cannot create on_delete job")
					return nil
				}

				logger.
					WithField("konnector", konnector).
					Info("Pushing job for konnector on_delete")

				msg, err := jobs.NewMessage(struct {
					Account        string `json:"account"`
					AccountRev     string `json:"account_rev"`
					Konnector      string `json:"konnector"`
					AccountDeleted bool   `json:"account_deleted"`
				}{
					Account:        old.ID(),
					AccountRev:     old.Rev(),
					Konnector:      konnector,
					AccountDeleted: true,
				})
				if err != nil {
					return err
				}
				if _, err = broker.PushJob(&jobs.JobRequest{
					Domain:     domain,
					WorkerType: "konnector",
					Message:    msg,
				}); err != nil {
					return err
				}
			}

			return nil
		})
}
<|endoftext|>


<commit_before>package cephfs

import (
	"encoding/base64"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/ceph/ceph-csi/pkg/util"
	"github.com/pkg/errors"
	"k8s.io/klog"
)

type volumeMountCacheEntry struct {
	DriverVersion string `json:"driverVersion"`

	VolumeID    string            `json:"volumeID"`
	Secrets     map[string]string `json:"secrets"`
	StagingPath string            `json:"stagingPath"`
	TargetPaths map[string]bool   `json:"targetPaths"`
	CreateTime  time.Time         `json:"createTime"`
}

type volumeMountCacheMap struct {
	volumes        map[string]volumeMountCacheEntry
	nodeCacheStore util.NodeCache
	metadataStore  util.CachePersister
}

var (
	volumeMountCachePrefix = "cephfs-mount-cache-"
	volumeMountCache       volumeMountCacheMap
	volumeMountCacheMtx    sync.Mutex
)

func initVolumeMountCache(driverName string, mountCacheDir string, cachePersister util.CachePersister) {
	volumeMountCache.volumes = make(map[string]volumeMountCacheEntry)

	volumeMountCache.metadataStore = cachePersister
	volumeMountCache.nodeCacheStore.BasePath = mountCacheDir
	volumeMountCache.nodeCacheStore.CacheDir = driverName
	klog.Infof("mount-cache: name: %s, version: %s, mountCacheDir: %s", driverName, version, mountCacheDir)
}

func remountCachedVolumes() error {
	if err := os.MkdirAll(volumeMountCache.nodeCacheStore.BasePath, 0755); err != nil {
		klog.Errorf("mount-cache: failed to create %s: %v", volumeMountCache.nodeCacheStore.BasePath, err)
		return err
	}
	var remountFailCount, remountSuccCount int64
	me := &volumeMountCacheEntry{}
	ce := &controllerCacheEntry{}
	err := volumeMountCache.nodeCacheStore.ForAll(volumeMountCachePrefix, me, func(identifier string) error {
		volID := me.VolumeID
		if err := volumeMountCache.metadataStore.Get(volID, ce); err != nil {
			if err, ok := err.(*util.CacheEntryNotFound); ok {
				klog.Infof("mount-cache: metadata not found, assuming the volume %s to be already deleted (%v)", volID, err)
				if err := volumeMountCache.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID)); err == nil {
					klog.Infof("mount-cache: metadata not found, delete volume cache entry for volume %s", volID)
				}
			}
		} else {
			if err := mountOneCacheEntry(ce, me); err == nil {
				remountSuccCount++
				volumeMountCache.volumes[me.VolumeID] = *me
				klog.Infof("mount-cache: remount volume %s success", volID)
			} else {
				remountFailCount++
				klog.Errorf("mount-cache: remount volume cache %s fail", volID)
			}
		}
		return nil
	})
	if err != nil {
		klog.Infof("mount-cache: metastore list cache fail %v", err)
		return err
	}
	if remountFailCount > 0 {
		klog.Infof("mount-cache: success remount %d volumes, fail remount %d volumes", remountSuccCount, remountFailCount)
	} else {
		klog.Infof("mount-cache: volume cache num %d, all succ remount", remountSuccCount)
	}
	return nil
}

func mountOneCacheEntry(ce *controllerCacheEntry, me *volumeMountCacheEntry) error {
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	var (
		err error
		cr  *credentials
	)
	volID := ce.VolumeID
	volOptions := ce.VolOptions

	if volOptions.ProvisionVolume {
		volOptions.RootPath = getVolumeRootPathCeph(volID)
		cr, err = getAdminCredentials(decodeCredentials(me.Secrets))
		if err != nil {
			return err
		}
		var entity *cephEntity
		entity, err = getCephUser(&volOptions, cr, volID)
		if err != nil {
			return err
		}
		cr = entity.toCredentials()
	} else {
		cr, err = getUserCredentials(decodeCredentials(me.Secrets))
		if err != nil {
			return err
		}
	}

	err = cleanupMountPoint(me.StagingPath)
	if err != nil {
		klog.Infof("mount-cache: failed to cleanup volume mount point %s, remove it: %s %v", volID, me.StagingPath, err)
		return err
	}

	isMnt, err := isMountPoint(me.StagingPath)
	if err != nil {
		isMnt = false
		klog.Infof("mount-cache: failed to check volume mounted %s: %s %v", volID, me.StagingPath, err)
	}

	if !isMnt {
		m, err := newMounter(&volOptions)
		if err != nil {
			klog.Errorf("mount-cache: failed to create mounter for volume %s: %v", volID, err)
			return err
		}
		if err := m.mount(me.StagingPath, cr, &volOptions); err != nil {
			klog.Errorf("mount-cache: failed to mount volume %s: %v", volID, err)
			return err
		}
	}
	for targetPath, readOnly := range me.TargetPaths {
		if err := cleanupMountPoint(targetPath); err == nil {
			if err := bindMount(me.StagingPath, targetPath, readOnly); err != nil {
				klog.Errorf("mount-cache: failed to bind-mount volume %s: %s %s %v %v",
					volID, me.StagingPath, targetPath, readOnly, err)
			} else {
				klog.Infof("mount-cache: successfully bind-mount volume %s: %s %s %v",
					volID, me.StagingPath, targetPath, readOnly)
			}
		}
	}
	return nil
}

func cleanupMountPoint(mountPoint string) error {
	if _, err := os.Stat(mountPoint); err != nil {
		if isCorruptedMnt(err) {
			klog.Infof("mount-cache: corrupted mount point %s, need unmount", mountPoint)
			err := execCommandErr("umount", mountPoint)
			if err != nil {
				klog.Infof("mount-cache: unmount %s fail %v", mountPoint, err)
				//ignore error return err
			}
		}
	}
	if _, err := os.Stat(mountPoint); err != nil {
		klog.Errorf("mount-cache: mount point %s stat fail %v", mountPoint, err)
		return err
	}
	return nil
}

func isCorruptedMnt(err error) bool {
	var underlyingError error
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		underlyingError = pe.Err
	case *os.LinkError:
		underlyingError = pe.Err
	case *os.SyscallError:
		underlyingError = pe.Err
	default:
		return false
	}

	CorruptedErrors := []error{
		syscall.ENOTCONN, syscall.ESTALE, syscall.EIO, syscall.EACCES}

	for _, v := range CorruptedErrors {
		if underlyingError == v {
			return true
		}
	}
	return false
}

func genVolumeMountCacheFileName(volID string) string {
	cachePath := volumeMountCachePrefix + volID
	return cachePath
}
func (mc *volumeMountCacheMap) isEnable() bool {
	//if mount cache dir unset, disable state
	return mc.nodeCacheStore.BasePath != ""
}

func (mc *volumeMountCacheMap) nodeStageVolume(volID string, stagingTargetPath string, secrets map[string]string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	lastTargetPaths := make(map[string]bool)
	me, ok := volumeMountCache.volumes[volID]
	if ok {
		if me.StagingPath == stagingTargetPath {
			klog.Warningf("mount-cache: node unexpected restage volume for volume %s", volID)
			return nil
		}
		lastTargetPaths = me.TargetPaths
		klog.Warningf("mount-cache: node stage volume ignore last cache entry for volume %s", volID)
	}

	me = volumeMountCacheEntry{DriverVersion: version}

	me.VolumeID = volID
	me.Secrets = encodeCredentials(secrets)
	me.StagingPath = stagingTargetPath
	me.TargetPaths = lastTargetPaths

	me.CreateTime = time.Now()
	volumeMountCache.volumes[volID] = me
	return mc.nodeCacheStore.Create(genVolumeMountCacheFileName(volID), me)
}

func (mc *volumeMountCacheMap) nodeUnStageVolume(volID string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()
	delete(volumeMountCache.volumes, volID)
	return mc.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID))
}

func (mc *volumeMountCacheMap) nodePublishVolume(volID string, targetPath string, readOnly bool) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	_, ok := volumeMountCache.volumes[volID]
	if !ok {
		return errors.New("mount-cache: node publish volume failed to find cache entry for volume")
	}
	volumeMountCache.volumes[volID].TargetPaths[targetPath] = readOnly
	return mc.updateNodeCache(volID)
}

func (mc *volumeMountCacheMap) nodeUnPublishVolume(volID string, targetPath string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	_, ok := volumeMountCache.volumes[volID]
	if !ok {
		return errors.New("mount-cache: node unpublish volume failed to find cache entry for volume")
	}
	delete(volumeMountCache.volumes[volID].TargetPaths, targetPath)
	return mc.updateNodeCache(volID)
}

func (mc *volumeMountCacheMap) updateNodeCache(volID string) error {
	me := volumeMountCache.volumes[volID]
	if err := volumeMountCache.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID)); err == nil {
		klog.Infof("mount-cache: metadata notfound, delete mount cache failed for volume %s", volID)
	}
	return mc.nodeCacheStore.Create(genVolumeMountCacheFileName(volID), me)
}

func encodeCredentials(input map[string]string) (output map[string]string) {
	output = make(map[string]string)
	for key, value := range input {
		nKey := base64.StdEncoding.EncodeToString([]byte(key))
		nValue := base64.StdEncoding.EncodeToString([]byte(value))
		output[nKey] = nValue
	}
	return output
}

func decodeCredentials(input map[string]string) (output map[string]string) {
	output = make(map[string]string)
	for key, value := range input {
		nKey, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			klog.Errorf("mount-cache: decode secret fail")
			continue
		}
		nValue, err := base64.StdEncoding.DecodeString(value)
		if err != nil {
			klog.Errorf("mount-cache: decode secret fail")
			continue
		}
		output[string(nKey)] = string(nValue)
	}
	return output
}
<commit_msg>code style<commit_after>package cephfs

import (
	"encoding/base64"
	"os"
	"sync"
	"syscall"
	"time"

	"github.com/ceph/ceph-csi/pkg/util"
	"github.com/pkg/errors"
	"k8s.io/klog"
)

type volumeMountCacheEntry struct {
	DriverVersion string `json:"driverVersion"`

	VolumeID    string            `json:"volumeID"`
	Secrets     map[string]string `json:"secrets"`
	StagingPath string            `json:"stagingPath"`
	TargetPaths map[string]bool   `json:"targetPaths"`
	CreateTime  time.Time         `json:"createTime"`
}

type volumeMountCacheMap struct {
	volumes        map[string]volumeMountCacheEntry
	nodeCacheStore util.NodeCache
	metadataStore  util.CachePersister
}

var (
	volumeMountCachePrefix = "cephfs-mount-cache-"
	volumeMountCache       volumeMountCacheMap
	volumeMountCacheMtx    sync.Mutex
)

func initVolumeMountCache(driverName string, mountCacheDir string, cachePersister util.CachePersister) {
	volumeMountCache.volumes = make(map[string]volumeMountCacheEntry)

	volumeMountCache.metadataStore = cachePersister
	volumeMountCache.nodeCacheStore.BasePath = mountCacheDir
	volumeMountCache.nodeCacheStore.CacheDir = driverName
	klog.Infof("mount-cache: name: %s, version: %s, mountCacheDir: %s", driverName, version, mountCacheDir)
}

func remountCachedVolumes() error {
	if err := os.MkdirAll(volumeMountCache.nodeCacheStore.BasePath, 0755); err != nil {
		klog.Errorf("mount-cache: failed to create %s: %v", volumeMountCache.nodeCacheStore.BasePath, err)
		return err
	}
	var remountFailCount, remountSuccCount int64
	me := &volumeMountCacheEntry{}
	ce := &controllerCacheEntry{}
	err := volumeMountCache.nodeCacheStore.ForAll(volumeMountCachePrefix, me, func(identifier string) error {
		volID := me.VolumeID
		if err := volumeMountCache.metadataStore.Get(volID, ce); err != nil {
			if err, ok := err.(*util.CacheEntryNotFound); ok {
				klog.Infof("mount-cache: metadata not found, assuming the volume %s to be already deleted (%v)", volID, err)
				if err := volumeMountCache.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID)); err == nil {
					klog.Infof("mount-cache: metadata not found, delete volume cache entry for volume %s", volID)
				}
			}
		} else {
			if err := mountOneCacheEntry(ce, me); err == nil {
				remountSuccCount++
				volumeMountCache.volumes[me.VolumeID] = *me
				klog.Infof("mount-cache: successfully remounted volume %s", volID)
			} else {
				remountFailCount++
				klog.Errorf("mount-cache: failed to remount volume %s", volID)
			}
		}
		return nil
	})
	if err != nil {
		klog.Infof("mount-cache: metastore list cache fail %v", err)
		return err
	}
	if remountFailCount > 0 {
		klog.Infof("mount-cache: successfully remounted %d volumes, failed to remount %d volumes", remountSuccCount, remountFailCount)
	} else {
		klog.Infof("mount-cache: successfully remounted %d volumes", remountSuccCount)
	}
	return nil
}

func mountOneCacheEntry(ce *controllerCacheEntry, me *volumeMountCacheEntry) error {
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	var (
		err error
		cr  *credentials
	)
	volID := ce.VolumeID
	volOptions := ce.VolOptions

	if volOptions.ProvisionVolume {
		volOptions.RootPath = getVolumeRootPathCeph(volID)
		cr, err = getAdminCredentials(decodeCredentials(me.Secrets))
		if err != nil {
			return err
		}
		var entity *cephEntity
		entity, err = getCephUser(&volOptions, cr, volID)
		if err != nil {
			return err
		}
		cr = entity.toCredentials()
	} else {
		cr, err = getUserCredentials(decodeCredentials(me.Secrets))
		if err != nil {
			return err
		}
	}

	err = cleanupMountPoint(me.StagingPath)
	if err != nil {
		klog.Infof("mount-cache: failed to cleanup volume mount point %s, remove it: %s %v", volID, me.StagingPath, err)
		return err
	}

	isMnt, err := isMountPoint(me.StagingPath)
	if err != nil {
		isMnt = false
		klog.Infof("mount-cache: failed to check volume mounted %s: %s %v", volID, me.StagingPath, err)
	}

	if !isMnt {
		m, err := newMounter(&volOptions)
		if err != nil {
			klog.Errorf("mount-cache: failed to create mounter for volume %s: %v", volID, err)
			return err
		}
		if err := m.mount(me.StagingPath, cr, &volOptions); err != nil {
			klog.Errorf("mount-cache: failed to mount volume %s: %v", volID, err)
			return err
		}
	}
	for targetPath, readOnly := range me.TargetPaths {
		if err := cleanupMountPoint(targetPath); err == nil {
			if err := bindMount(me.StagingPath, targetPath, readOnly); err != nil {
				klog.Errorf("mount-cache: failed to bind-mount volume %s: %s %s %v %v",
					volID, me.StagingPath, targetPath, readOnly, err)
			} else {
				klog.Infof("mount-cache: successfully bind-mount volume %s: %s %s %v",
					volID, me.StagingPath, targetPath, readOnly)
			}
		}
	}
	return nil
}

func cleanupMountPoint(mountPoint string) error {
	if _, err := os.Stat(mountPoint); err != nil {
		if isCorruptedMnt(err) {
			klog.Infof("mount-cache: corrupted mount point %s, need unmount", mountPoint)
			err := execCommandErr("umount", mountPoint)
			if err != nil {
				klog.Infof("mount-cache: failed to umount %s %v", mountPoint, err)
				//ignore error return err
			}
		}
	}
	if _, err := os.Stat(mountPoint); err != nil {
		klog.Errorf("mount-cache: failed to stat mount point %s %v", mountPoint, err)
		return err
	}
	return nil
}

func isCorruptedMnt(err error) bool {
	var underlyingError error
	switch pe := err.(type) {
	case nil:
		return false
	case *os.PathError:
		underlyingError = pe.Err
	case *os.LinkError:
		underlyingError = pe.Err
	case *os.SyscallError:
		underlyingError = pe.Err
	default:
		return false
	}

	CorruptedErrors := []error{
		syscall.ENOTCONN, syscall.ESTALE, syscall.EIO, syscall.EACCES}

	for _, v := range CorruptedErrors {
		if underlyingError == v {
			return true
		}
	}
	return false
}

func genVolumeMountCacheFileName(volID string) string {
	cachePath := volumeMountCachePrefix + volID
	return cachePath
}
func (mc *volumeMountCacheMap) isEnable() bool {
	//if mount cache dir unset, disable state
	return mc.nodeCacheStore.BasePath != ""
}

func (mc *volumeMountCacheMap) nodeStageVolume(volID string, stagingTargetPath string, secrets map[string]string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	lastTargetPaths := make(map[string]bool)
	me, ok := volumeMountCache.volumes[volID]
	if ok {
		if me.StagingPath == stagingTargetPath {
			klog.Warningf("mount-cache: node unexpected restage volume for volume %s", volID)
			return nil
		}
		lastTargetPaths = me.TargetPaths
		klog.Warningf("mount-cache: node stage volume ignore last cache entry for volume %s", volID)
	}

	me = volumeMountCacheEntry{DriverVersion: version}

	me.VolumeID = volID
	me.Secrets = encodeCredentials(secrets)
	me.StagingPath = stagingTargetPath
	me.TargetPaths = lastTargetPaths

	me.CreateTime = time.Now()
	volumeMountCache.volumes[volID] = me
	return mc.nodeCacheStore.Create(genVolumeMountCacheFileName(volID), me)
}

func (mc *volumeMountCacheMap) nodeUnStageVolume(volID string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()
	delete(volumeMountCache.volumes, volID)
	return mc.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID))
}

func (mc *volumeMountCacheMap) nodePublishVolume(volID string, targetPath string, readOnly bool) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	_, ok := volumeMountCache.volumes[volID]
	if !ok {
		return errors.New("mount-cache: node publish volume failed to find cache entry for volume")
	}
	volumeMountCache.volumes[volID].TargetPaths[targetPath] = readOnly
	return mc.updateNodeCache(volID)
}

func (mc *volumeMountCacheMap) nodeUnPublishVolume(volID string, targetPath string) error {
	if !mc.isEnable() {
		return nil
	}
	volumeMountCacheMtx.Lock()
	defer volumeMountCacheMtx.Unlock()

	_, ok := volumeMountCache.volumes[volID]
	if !ok {
		return errors.New("mount-cache: node unpublish volume failed to find cache entry for volume")
	}
	delete(volumeMountCache.volumes[volID].TargetPaths, targetPath)
	return mc.updateNodeCache(volID)
}

func (mc *volumeMountCacheMap) updateNodeCache(volID string) error {
	me := volumeMountCache.volumes[volID]
	if err := volumeMountCache.nodeCacheStore.Delete(genVolumeMountCacheFileName(volID)); err == nil {
		klog.Infof("mount-cache: metadata notfound, delete mount cache failed for volume %s", volID)
	}
	return mc.nodeCacheStore.Create(genVolumeMountCacheFileName(volID), me)
}

func encodeCredentials(input map[string]string) (output map[string]string) {
	output = make(map[string]string)
	for key, value := range input {
		nKey := base64.StdEncoding.EncodeToString([]byte(key))
		nValue := base64.StdEncoding.EncodeToString([]byte(value))
		output[nKey] = nValue
	}
	return output
}

func decodeCredentials(input map[string]string) (output map[string]string) {
	output = make(map[string]string)
	for key, value := range input {
		nKey, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			klog.Errorf("mount-cache: decode secret fail")
			continue
		}
		nValue, err := base64.StdEncoding.DecodeString(value)
		if err != nil {
			klog.Errorf("mount-cache: decode secret fail")
			continue
		}
		output[string(nKey)] = string(nValue)
	}
	return output
}
<|endoftext|>


<commit_before>// Copyright 2020 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// It may or may not work on other OSes.
// If you test on another OS and it works, enable it.
//go:build linux
// +build linux

package cover

import (
	"bytes"
	"encoding/csv"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/google/syzkaller/pkg/cover/backend"
	"github.com/google/syzkaller/pkg/mgrconfig"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/symbolizer"
	_ "github.com/google/syzkaller/sys"
	"github.com/google/syzkaller/sys/targets"
)

type Test struct {
	Name     string
	CFlags   []string
	LDFlags  []string
	Progs    []Prog
	AddCover bool
	Result   string
	Supports func(target *targets.Target) bool
}

func TestReportGenerator(t *testing.T) {
	tests := []Test{
		{
			Name:     "no-coverage",
			CFlags:   []string{"-g"},
			AddCover: true,
			Result:   `.* doesn't contain coverage callbacks \(set CONFIG_KCOV=y on linux\)`,
		},
		{
			Name:     "no-debug-info",
			CFlags:   []string{"-fsanitize-coverage=trace-pc"},
			AddCover: true,
			Result:   `failed to parse DWARF.*\(set CONFIG_DEBUG_INFO=y on linux\)`,
		},
		{
			Name:   "no-pcs",
			CFlags: []string{"-fsanitize-coverage=trace-pc", "-g"},
			Result: `no coverage collected so far`,
		},
		{
			Name:   "bad-pcs",
			CFlags: []string{"-fsanitize-coverage=trace-pc", "-g"},
			Progs:  []Prog{{Data: "data", PCs: []uint64{0x1, 0x2}}},
			Result: `coverage doesn't match any coverage callbacks`,
		},
		{
			Name:     "good",
			AddCover: true,
			CFlags:   []string{"-fsanitize-coverage=trace-pc", "-g"},
		},
		{
			Name:     "good-pie",
			AddCover: true,
			CFlags:   []string{"-fsanitize-coverage=trace-pc", "-g", "-fpie"},
			LDFlags:  []string{"-pie", "-Wl,--section-start=.text=0x33300000"},
			Supports: func(target *targets.Target) bool {
				return target.OS == targets.Fuchsia ||
					// Fails with "relocation truncated to fit: R_AARCH64_CALL26 against symbol `memcpy'".
					target.OS == targets.Linux && target.Arch != targets.ARM64
			},
		},
		{
			Name:     "good-pie-relocs",
			AddCover: true,
			// This produces a binary that resembles CONFIG_RANDOMIZE_BASE=y.
			// Symbols and .text section has addresses around 0x33300000,
			// but debug info has all PC ranges around 0 address.
			CFlags:  []string{"-fsanitize-coverage=trace-pc", "-g", "-fpie"},
			LDFlags: []string{"-pie", "-Wl,--section-start=.text=0x33300000,--emit-relocs"},
			Supports: func(target *targets.Target) bool {
				return target.OS == targets.Fuchsia ||
					target.OS == targets.Linux && target.Arch != targets.ARM64 &&
						target.Arch != targets.ARM && target.Arch != targets.I386
			},
		},
	}
	t.Parallel()
	for os, arches := range targets.List {
		if os == targets.TestOS {
			continue
		}
		for _, target := range arches {
			target := targets.Get(target.OS, target.Arch)
			if target.BuildOS != runtime.GOOS {
				continue
			}
			t.Run(target.OS+"-"+target.Arch, func(t *testing.T) {
				t.Parallel()
				if target.BrokenCompiler != "" {
					t.Skip("skipping the test due to broken cross-compiler:\n" + target.BrokenCompiler)
				}
				for _, test := range tests {
					test := test
					t.Run(test.Name, func(t *testing.T) {
						if test.Supports != nil && !test.Supports(target) {
							t.Skip("unsupported target")
						}
						t.Parallel()
						testReportGenerator(t, target, test)
					})
				}
			})
		}
	}
}

func testReportGenerator(t *testing.T, target *targets.Target, test Test) {
	rep, csv, err := generateReport(t, target, test)
	if err != nil {
		if test.Result == "" {
			t.Fatalf("expected no error, but got:\n%v", err)
		}
		if !regexp.MustCompile(test.Result).MatchString(err.Error()) {
			t.Fatalf("expected error %q, but got:\n%v", test.Result, err)
		}
		return
	}
	if test.Result != "" {
		t.Fatalf("got no error, but expected %q", test.Result)
	}
	checkCSVReport(t, csv)
	_ = rep
}

const kcovCode = `
#ifdef ASLR_BASE
#define _GNU_SOURCE
#endif

#include <stdio.h>

#ifdef ASLR_BASE
#include <dlfcn.h>
#include <link.h>
#include <stddef.h>

void* aslr_base() {
       struct link_map* map = NULL;
       void* handle = dlopen(NULL, RTLD_LAZY | RTLD_NOLOAD);
       if (handle != NULL) {
              dlinfo(handle, RTLD_DI_LINKMAP, &map);
              dlclose(handle);
       }
       return map ? map->l_addr : NULL;
}
#else
void* aslr_base() { return NULL; }
#endif

void __sanitizer_cov_trace_pc() { printf("%llu", (long long)(__builtin_return_address(0) - aslr_base())); }
`

func buildTestBinary(t *testing.T, target *targets.Target, test Test, dir string) string {
	kcovSrc := filepath.Join(dir, "kcov.c")
	kcovObj := filepath.Join(dir, "kcov.o")
	if err := osutil.WriteFile(kcovSrc, []byte(kcovCode)); err != nil {
		t.Fatal(err)
	}

	aslrDefine := "-DNO_ASLR_BASE"
	if target.OS == targets.Linux || target.OS == targets.OpenBSD ||
		target.OS == targets.FreeBSD || target.OS == targets.NetBSD {
		aslrDefine = "-DASLR_BASE"
	}
	aslrExtraLibs := []string{}
	if target.OS == targets.Linux {
		aslrExtraLibs = []string{"-ldl"}
	}

	kcovFlags := append([]string{"-c", "-fpie", "-w", "-x", "c", "-o", kcovObj, kcovSrc, aslrDefine}, target.CFlags...)
	src := filepath.Join(dir, "main.c")
	obj := filepath.Join(dir, "main.o")
	bin := filepath.Join(dir, target.KernelObject)
	if err := osutil.WriteFile(src, []byte(`int main() {}`)); err != nil {
		t.Fatal(err)
	}
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, kcovFlags...); err != nil {
		t.Fatal(err)
	}

	// We used to compile and link with a single compiler invocation,
	// but clang has a bug that it tries to link in ubsan runtime when
	// -fsanitize-coverage=trace-pc is provided during linking and
	// ubsan runtime is missing for arm/arm64/riscv arches in the llvm packages.
	// So we first compile with -fsanitize-coverage and then link w/o it.
	cflags := append(append([]string{"-w", "-c", "-o", obj, src}, target.CFlags...), test.CFlags...)
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, cflags...); err != nil {
		errText := err.Error()
		errText = strings.ReplaceAll(errText, "‘", "'")
		errText = strings.ReplaceAll(errText, "’", "'")
		if strings.Contains(errText, "error: unrecognized command line option '-fsanitize-coverage=trace-pc'") &&
			(os.Getenv("SYZ_BIG_ENV") == "" || target.OS == targets.Akaros) {
			t.Skip("skipping test, -fsanitize-coverage=trace-pc is not supported")
		}
		t.Fatal(err)
	}

	ldflags := append(append(append([]string{"-o", bin, obj, kcovObj}, aslrExtraLibs...),
		target.CFlags...), test.LDFlags...)
	staticIdx, pieIdx := -1, -1
	for i, arg := range ldflags {
		switch arg {
		case "-static":
			staticIdx = i
		case "-pie":
			pieIdx = i
		}
	}
	if target.OS == targets.Fuchsia && pieIdx != -1 {
		// Fuchsia toolchain fails when given -pie:
		// clang-12: error: argument unused during compilation: '-pie'
		ldflags[pieIdx] = ldflags[len(ldflags)-1]
		ldflags = ldflags[:len(ldflags)-1]
	} else if pieIdx != -1 && staticIdx != -1 {
		// -static and -pie are incompatible during linking.
		ldflags[staticIdx] = ldflags[len(ldflags)-1]
		ldflags = ldflags[:len(ldflags)-1]
	}
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, ldflags...); err != nil {
		// Arm linker in the big-env image has a bug when linking a clang-produced files.
		if regexp.MustCompile(`arm-linux-gnueabi.* assertion fail`).MatchString(err.Error()) {
			t.Skipf("skipping test, broken arm linker (%v)", err)
		}
		t.Fatal(err)
	}
	return bin
}

func generateReport(t *testing.T, target *targets.Target, test Test) ([]byte, []byte, error) {
	dir := t.TempDir()
	bin := buildTestBinary(t, target, test, dir)
	subsystem := []mgrconfig.Subsystem{
		{
			Name: "sound",
			Paths: []string{
				"sound",
				"techpack/audio",
			},
		},
	}

	rg, err := MakeReportGenerator(target, "", dir, dir, dir, subsystem, nil, nil, false)
	if err != nil {
		return nil, nil, err
	}
	if test.AddCover {
		var pcs []uint64
		// Sanitizers crash when installing signal handlers with static libc.
		const sanitizerOptions = "handle_segv=0:handle_sigbus=0:handle_sigfpe=0"
		cmd := osutil.Command(bin)
		cmd.Env = append([]string{
			"UBSAN_OPTIONS=" + sanitizerOptions,
			"ASAN_OPTIONS=" + sanitizerOptions,
		}, os.Environ()...)
		if output, err := osutil.Run(time.Minute, cmd); err == nil {
			pc, err := strconv.ParseUint(string(output), 10, 64)
			if err != nil {
				t.Fatal(err)
			}
			pcs = append(pcs, backend.PreviousInstructionPC(target, pc))
			t.Logf("using exact coverage PC 0x%x", pcs[0])
		} else if target.OS == runtime.GOOS && (target.Arch == runtime.GOARCH || target.VMArch == runtime.GOARCH) {
			t.Fatal(err)
		} else {
			symb := symbolizer.NewSymbolizer(target)
			text, err := symb.ReadTextSymbols(bin)
			if err != nil {
				t.Fatal(err)
			}
			if nmain := len(text["main"]); nmain != 1 {
				t.Fatalf("got %v main symbols", nmain)
			}
			main := text["main"][0]
			for off := 0; off < main.Size; off++ {
				pcs = append(pcs, main.Addr+uint64(off))
			}
			t.Logf("using inexact coverage range 0x%x-0x%x", main.Addr, main.Addr+uint64(main.Size))
		}
		test.Progs = append(test.Progs, Prog{Data: "main", PCs: pcs})
	}
	html := new(bytes.Buffer)
	if err := rg.DoHTML(html, test.Progs, nil); err != nil {
		return nil, nil, err
	}
	htmlTable := new(bytes.Buffer)
	if err := rg.DoHTMLTable(htmlTable, test.Progs, nil); err != nil {
		return nil, nil, err
	}
	_ = htmlTable
	csv := new(bytes.Buffer)
	if err := rg.DoCSV(csv, test.Progs, nil); err != nil {
		return nil, nil, err
	}
	csvFiles := new(bytes.Buffer)
	if err := rg.DoCSVFiles(csvFiles, test.Progs, nil); err != nil {
		return nil, nil, err
	}
	_ = csvFiles

	return html.Bytes(), csv.Bytes(), nil
}

func checkCSVReport(t *testing.T, CSVReport []byte) {
	csvReader := csv.NewReader(bytes.NewBuffer(CSVReport))
	lines, err := csvReader.ReadAll()
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(lines[0], csvHeader) {
		t.Fatalf("heading line in CSV doesn't match %v", lines[0])
	}

	foundMain := false
	for _, line := range lines {
		if line[2] == "main" {
			foundMain = true
			if line[3] != "1" && line[4] != "1" {
				t.Fatalf("function coverage percentage doesn't match %v vs. %v", line[3], "100")
			}
		}
	}
	if !foundMain {
		t.Fatalf("no main in the CSV report")
	}
}
<commit_msg>pkg/cover: mitigate datarace in generateReport<commit_after>// Copyright 2020 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// It may or may not work on other OSes.
// If you test on another OS and it works, enable it.
//go:build linux
// +build linux

package cover

import (
	"bytes"
	"encoding/csv"
	"os"
	"path/filepath"
	"reflect"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/google/syzkaller/pkg/cover/backend"
	"github.com/google/syzkaller/pkg/mgrconfig"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/pkg/symbolizer"
	_ "github.com/google/syzkaller/sys"
	"github.com/google/syzkaller/sys/targets"
)

type Test struct {
	Name     string
	CFlags   []string
	LDFlags  []string
	Progs    []Prog
	AddCover bool
	Result   string
	Supports func(target *targets.Target) bool
}

func TestReportGenerator(t *testing.T) {
	tests := []Test{
		{
			Name:     "no-coverage",
			CFlags:   []string{"-g"},
			AddCover: true,
			Result:   `.* doesn't contain coverage callbacks \(set CONFIG_KCOV=y on linux\)`,
		},
		{
			Name:     "no-debug-info",
			CFlags:   []string{"-fsanitize-coverage=trace-pc"},
			AddCover: true,
			Result:   `failed to parse DWARF.*\(set CONFIG_DEBUG_INFO=y on linux\)`,
		},
		{
			Name:   "no-pcs",
			CFlags: []string{"-fsanitize-coverage=trace-pc", "-g"},
			Result: `no coverage collected so far`,
		},
		{
			Name:   "bad-pcs",
			CFlags: []string{"-fsanitize-coverage=trace-pc", "-g"},
			Progs:  []Prog{{Data: "data", PCs: []uint64{0x1, 0x2}}},
			Result: `coverage doesn't match any coverage callbacks`,
		},
		{
			Name:     "good",
			AddCover: true,
			CFlags:   []string{"-fsanitize-coverage=trace-pc", "-g"},
		},
		{
			Name:     "good-pie",
			AddCover: true,
			CFlags:   []string{"-fsanitize-coverage=trace-pc", "-g", "-fpie"},
			LDFlags:  []string{"-pie", "-Wl,--section-start=.text=0x33300000"},
			Supports: func(target *targets.Target) bool {
				return target.OS == targets.Fuchsia ||
					// Fails with "relocation truncated to fit: R_AARCH64_CALL26 against symbol `memcpy'".
					target.OS == targets.Linux && target.Arch != targets.ARM64
			},
		},
		{
			Name:     "good-pie-relocs",
			AddCover: true,
			// This produces a binary that resembles CONFIG_RANDOMIZE_BASE=y.
			// Symbols and .text section has addresses around 0x33300000,
			// but debug info has all PC ranges around 0 address.
			CFlags:  []string{"-fsanitize-coverage=trace-pc", "-g", "-fpie"},
			LDFlags: []string{"-pie", "-Wl,--section-start=.text=0x33300000,--emit-relocs"},
			Supports: func(target *targets.Target) bool {
				return target.OS == targets.Fuchsia ||
					target.OS == targets.Linux && target.Arch != targets.ARM64 &&
						target.Arch != targets.ARM && target.Arch != targets.I386
			},
		},
	}
	t.Parallel()
	for os, arches := range targets.List {
		if os == targets.TestOS {
			continue
		}
		for _, target := range arches {
			target := targets.Get(target.OS, target.Arch)
			if target.BuildOS != runtime.GOOS {
				continue
			}
			t.Run(target.OS+"-"+target.Arch, func(t *testing.T) {
				t.Parallel()
				if target.BrokenCompiler != "" {
					t.Skip("skipping the test due to broken cross-compiler:\n" + target.BrokenCompiler)
				}
				for _, test := range tests {
					test := test
					t.Run(test.Name, func(t *testing.T) {
						if test.Supports != nil && !test.Supports(target) {
							t.Skip("unsupported target")
						}
						t.Parallel()
						testReportGenerator(t, target, test)
					})
				}
			})
		}
	}
}

func testReportGenerator(t *testing.T, target *targets.Target, test Test) {
	rep, csv, err := generateReport(t, target, test)
	if err != nil {
		if test.Result == "" {
			t.Fatalf("expected no error, but got:\n%v", err)
		}
		if !regexp.MustCompile(test.Result).MatchString(err.Error()) {
			t.Fatalf("expected error %q, but got:\n%v", test.Result, err)
		}
		return
	}
	if test.Result != "" {
		t.Fatalf("got no error, but expected %q", test.Result)
	}
	checkCSVReport(t, csv)
	_ = rep
}

const kcovCode = `
#ifdef ASLR_BASE
#define _GNU_SOURCE
#endif

#include <stdio.h>

#ifdef ASLR_BASE
#include <dlfcn.h>
#include <link.h>
#include <stddef.h>

void* aslr_base() {
       struct link_map* map = NULL;
       void* handle = dlopen(NULL, RTLD_LAZY | RTLD_NOLOAD);
       if (handle != NULL) {
              dlinfo(handle, RTLD_DI_LINKMAP, &map);
              dlclose(handle);
       }
       return map ? map->l_addr : NULL;
}
#else
void* aslr_base() { return NULL; }
#endif

void __sanitizer_cov_trace_pc() { printf("%llu", (long long)(__builtin_return_address(0) - aslr_base())); }
`

func buildTestBinary(t *testing.T, target *targets.Target, test Test, dir string) string {
	kcovSrc := filepath.Join(dir, "kcov.c")
	kcovObj := filepath.Join(dir, "kcov.o")
	if err := osutil.WriteFile(kcovSrc, []byte(kcovCode)); err != nil {
		t.Fatal(err)
	}

	aslrDefine := "-DNO_ASLR_BASE"
	if target.OS == targets.Linux || target.OS == targets.OpenBSD ||
		target.OS == targets.FreeBSD || target.OS == targets.NetBSD {
		aslrDefine = "-DASLR_BASE"
	}
	aslrExtraLibs := []string{}
	if target.OS == targets.Linux {
		aslrExtraLibs = []string{"-ldl"}
	}

	kcovFlags := append([]string{"-c", "-fpie", "-w", "-x", "c", "-o", kcovObj, kcovSrc, aslrDefine}, target.CFlags...)
	src := filepath.Join(dir, "main.c")
	obj := filepath.Join(dir, "main.o")
	bin := filepath.Join(dir, target.KernelObject)
	if err := osutil.WriteFile(src, []byte(`int main() {}`)); err != nil {
		t.Fatal(err)
	}
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, kcovFlags...); err != nil {
		t.Fatal(err)
	}

	// We used to compile and link with a single compiler invocation,
	// but clang has a bug that it tries to link in ubsan runtime when
	// -fsanitize-coverage=trace-pc is provided during linking and
	// ubsan runtime is missing for arm/arm64/riscv arches in the llvm packages.
	// So we first compile with -fsanitize-coverage and then link w/o it.
	cflags := append(append([]string{"-w", "-c", "-o", obj, src}, target.CFlags...), test.CFlags...)
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, cflags...); err != nil {
		errText := err.Error()
		errText = strings.ReplaceAll(errText, "‘", "'")
		errText = strings.ReplaceAll(errText, "’", "'")
		if strings.Contains(errText, "error: unrecognized command line option '-fsanitize-coverage=trace-pc'") &&
			(os.Getenv("SYZ_BIG_ENV") == "" || target.OS == targets.Akaros) {
			t.Skip("skipping test, -fsanitize-coverage=trace-pc is not supported")
		}
		t.Fatal(err)
	}

	ldflags := append(append(append([]string{"-o", bin, obj, kcovObj}, aslrExtraLibs...),
		target.CFlags...), test.LDFlags...)
	staticIdx, pieIdx := -1, -1
	for i, arg := range ldflags {
		switch arg {
		case "-static":
			staticIdx = i
		case "-pie":
			pieIdx = i
		}
	}
	if target.OS == targets.Fuchsia && pieIdx != -1 {
		// Fuchsia toolchain fails when given -pie:
		// clang-12: error: argument unused during compilation: '-pie'
		ldflags[pieIdx] = ldflags[len(ldflags)-1]
		ldflags = ldflags[:len(ldflags)-1]
	} else if pieIdx != -1 && staticIdx != -1 {
		// -static and -pie are incompatible during linking.
		ldflags[staticIdx] = ldflags[len(ldflags)-1]
		ldflags = ldflags[:len(ldflags)-1]
	}
	if _, err := osutil.RunCmd(time.Hour, "", target.CCompiler, ldflags...); err != nil {
		// Arm linker in the big-env image has a bug when linking a clang-produced files.
		if regexp.MustCompile(`arm-linux-gnueabi.* assertion fail`).MatchString(err.Error()) {
			t.Skipf("skipping test, broken arm linker (%v)", err)
		}
		t.Fatal(err)
	}
	return bin
}

func generateReport(t *testing.T, target *targets.Target, test Test) ([]byte, []byte, error) {
	dir := t.TempDir()
	bin := buildTestBinary(t, target, test, dir)
	subsystem := []mgrconfig.Subsystem{
		{
			Name: "sound",
			Paths: []string{
				"sound",
				"techpack/audio",
			},
		},
	}

	// Deep copy, as we are going to modify progs. Our test generate multiple reports from the same
	// test object in parallel. Without copying we have a datarace here.
	progs := []Prog{}
	for _, p := range test.Progs {
		progs = append(progs, Prog{Sig: p.Sig, Data: p.Data, PCs: append([]uint64{}, p.PCs...)})
	}

	rg, err := MakeReportGenerator(target, "", dir, dir, dir, subsystem, nil, nil, false)
	if err != nil {
		return nil, nil, err
	}
	if test.AddCover {
		var pcs []uint64
		// Sanitizers crash when installing signal handlers with static libc.
		const sanitizerOptions = "handle_segv=0:handle_sigbus=0:handle_sigfpe=0"
		cmd := osutil.Command(bin)
		cmd.Env = append([]string{
			"UBSAN_OPTIONS=" + sanitizerOptions,
			"ASAN_OPTIONS=" + sanitizerOptions,
		}, os.Environ()...)
		if output, err := osutil.Run(time.Minute, cmd); err == nil {
			pc, err := strconv.ParseUint(string(output), 10, 64)
			if err != nil {
				t.Fatal(err)
			}
			pcs = append(pcs, backend.PreviousInstructionPC(target, pc))
			t.Logf("using exact coverage PC 0x%x", pcs[0])
		} else if target.OS == runtime.GOOS && (target.Arch == runtime.GOARCH || target.VMArch == runtime.GOARCH) {
			t.Fatal(err)
		} else {
			symb := symbolizer.NewSymbolizer(target)
			text, err := symb.ReadTextSymbols(bin)
			if err != nil {
				t.Fatal(err)
			}
			if nmain := len(text["main"]); nmain != 1 {
				t.Fatalf("got %v main symbols", nmain)
			}
			main := text["main"][0]
			for off := 0; off < main.Size; off++ {
				pcs = append(pcs, main.Addr+uint64(off))
			}
			t.Logf("using inexact coverage range 0x%x-0x%x", main.Addr, main.Addr+uint64(main.Size))
		}
		progs = append(progs, Prog{Data: "main", PCs: pcs})
	}
	html := new(bytes.Buffer)
	if err := rg.DoHTML(html, progs, nil); err != nil {
		return nil, nil, err
	}
	htmlTable := new(bytes.Buffer)
	if err := rg.DoHTMLTable(htmlTable, progs, nil); err != nil {
		return nil, nil, err
	}
	_ = htmlTable
	csv := new(bytes.Buffer)
	if err := rg.DoCSV(csv, progs, nil); err != nil {
		return nil, nil, err
	}
	csvFiles := new(bytes.Buffer)
	if err := rg.DoCSVFiles(csvFiles, progs, nil); err != nil {
		return nil, nil, err
	}
	_ = csvFiles

	return html.Bytes(), csv.Bytes(), nil
}

func checkCSVReport(t *testing.T, CSVReport []byte) {
	csvReader := csv.NewReader(bytes.NewBuffer(CSVReport))
	lines, err := csvReader.ReadAll()
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(lines[0], csvHeader) {
		t.Fatalf("heading line in CSV doesn't match %v", lines[0])
	}

	foundMain := false
	for _, line := range lines {
		if line[2] == "main" {
			foundMain = true
			if line[3] != "1" && line[4] != "1" {
				t.Fatalf("function coverage percentage doesn't match %v vs. %v", line[3], "100")
			}
		}
	}
	if !foundMain {
		t.Fatalf("no main in the CSV report")
	}
}
<|endoftext|>


<commit_before>package gcc

import (
	"sync"
	"time"

	"github.com/pion/interceptor"
	"github.com/pion/interceptor/internal/cc"
	"github.com/pion/rtcp"
	"github.com/pion/rtp"
)

const (
	transportCCURI = "http://www.ietf.org/id/draft-holmer-rmcat-transport-wide-cc-extensions-01"
	latestBitrate  = 10_000
	minBitrate     = 5_000
	maxBitrate     = 50_000_000
)

// Pacer is the interface implemented by packet pacers
type Pacer interface {
	interceptor.RTPWriter
	AddStream(ssrc uint32, writer interceptor.RTPWriter)
	SetTargetBitrate(int)
	Close() error
}

// Stats contains internal statistics of the bandwidth estimator
type Stats struct {
	LossStats
	DelayStats
}

// SendSideBWE implements a combination of loss and delay based GCC
type SendSideBWE struct {
	pacer           Pacer
	lossController  *lossBasedBandwidthEstimator
	delayController *delayController
	feedbackAdapter *cc.FeedbackAdapter

	onTargetBitrateChange func(bitrate int)

	lock          sync.Mutex
	latestStats   Stats
	latestBitrate int
	minBitrate    int
	maxBitrate    int

	close chan struct{}
}

// Option configures a bandwidth estimator
type Option func(*SendSideBWE) error

// SendSideBWEInitialBitrate sets the initial bitrate of new GCC interceptors
func SendSideBWEInitialBitrate(rate int) Option {
	return func(e *SendSideBWE) error {
		e.latestBitrate = rate
		return nil
	}
}

// SendSideBWEPacer sets the pacing algorithm to use.
func SendSideBWEPacer(p Pacer) Option {
	return func(e *SendSideBWE) error {
		e.pacer = p
		return nil
	}
}

// NewSendSideBWE creates a new sender side bandwidth estimator
func NewSendSideBWE(opts ...Option) (*SendSideBWE, error) {
	e := &SendSideBWE{
		pacer:                 nil,
		lossController:        nil,
		delayController:       nil,
		feedbackAdapter:       cc.NewFeedbackAdapter(),
		onTargetBitrateChange: nil,
		lock:                  sync.Mutex{},
		latestStats:           Stats{},
		latestBitrate:         latestBitrate,
		minBitrate:            minBitrate,
		maxBitrate:            maxBitrate,
		close:                 make(chan struct{}),
	}
	for _, opt := range opts {
		if err := opt(e); err != nil {
			return nil, err
		}
	}
	if e.pacer == nil {
		e.pacer = NewLeakyBucketPacer(e.latestBitrate)
	}
	e.lossController = newLossBasedBWE(e.latestBitrate)
	e.delayController = newDelayController(delayControllerConfig{
		nowFn:          time.Now,
		initialBitrate: e.latestBitrate,
		minBitrate:     e.minBitrate,
		maxBitrate:     e.maxBitrate,
	})

	e.delayController.onUpdate(e.onDelayUpdate)

	return e, nil
}

// AddStream adds a new stream to the bandwidth estimator
func (e *SendSideBWE) AddStream(info *interceptor.StreamInfo, writer interceptor.RTPWriter) interceptor.RTPWriter {
	var hdrExtID uint8
	for _, e := range info.RTPHeaderExtensions {
		if e.URI == transportCCURI {
			hdrExtID = uint8(e.ID)
			break
		}
	}

	e.pacer.AddStream(info.SSRC, interceptor.RTPWriterFunc(func(header *rtp.Header, payload []byte, attributes interceptor.Attributes) (int, error) {
		if hdrExtID != 0 {
			if attributes == nil {
				attributes = make(interceptor.Attributes)
			}
			attributes.Set(cc.TwccExtensionAttributesKey, hdrExtID)
		}
		if err := e.feedbackAdapter.OnSent(time.Now(), header, len(payload), attributes); err != nil {
			return 0, err
		}
		return writer.Write(header, payload, attributes)
	}))
	return e.pacer
}

// WriteRTCP adds some RTCP feedback to the bandwidth estimator
func (e *SendSideBWE) WriteRTCP(pkts []rtcp.Packet, attributes interceptor.Attributes) error {
	for _, pkt := range pkts {
		if fb, ok := pkt.(*rtcp.TransportLayerCC); ok {
			acks, err := e.feedbackAdapter.OnTransportCCFeedback(time.Now(), fb)
			if err != nil {
				return err
			}
			e.lossController.updateLossEstimate(acks)
			e.delayController.updateDelayEstimate(acks)
		}
	}
	return nil
}

// GetTargetBitrate returns the current target bitrate
func (e *SendSideBWE) GetTargetBitrate() int {
	e.lock.Lock()
	defer e.lock.Unlock()

	return e.latestBitrate
}

// GetStats returns some internal statistics of the bandwidth estimator
func (e *SendSideBWE) GetStats() map[string]interface{} {
	e.lock.Lock()
	defer e.lock.Unlock()

	return map[string]interface{}{
		"lossTargetBitrate":  e.latestStats.LossStats.TargetBitrate,
		"averageLoss":        e.latestStats.AverageLoss,
		"delayTargetBitrate": e.latestStats.DelayStats.TargetBitrate,
		"delayMeasurement":   float64(e.latestStats.Measurement.Microseconds()) / 1000.0,
		"delayEstimate":      float64(e.latestStats.Estimate.Microseconds()) / 1000.0,
		"delayThreshold":     float64(e.latestStats.Threshold.Microseconds()) / 1000.0,
		"rtt":                float64(e.latestStats.RTT.Microseconds()) / 1000.0,
		"usage":              e.latestStats.Usage.String(),
		"state":              e.latestStats.State.String(),
	}
}

// OnTargetBitrateChange sets the callback that is called when the target
// bitrate changes
func (e *SendSideBWE) OnTargetBitrateChange(f func(bitrate int)) {
	e.onTargetBitrateChange = f
}

// Close stops and closes the bandwidth estimator
func (e *SendSideBWE) Close() error {
	if err := e.delayController.Close(); err != nil {
		return err
	}
	close(e.close)
	return e.pacer.Close()
}

func (e *SendSideBWE) onDelayUpdate(delayStats DelayStats) {
	e.lock.Lock()
	defer e.lock.Unlock()

	lossStats := e.lossController.getEstimate(delayStats.TargetBitrate)
	bitrateChanged := false
	bitrate := minInt(delayStats.TargetBitrate, lossStats.TargetBitrate)
	if bitrate != e.latestBitrate {
		bitrateChanged = true
		e.latestBitrate = bitrate
		e.pacer.SetTargetBitrate(e.latestBitrate)
	}

	if bitrateChanged && e.onTargetBitrateChange != nil {
		go e.onTargetBitrateChange(bitrate)
	}

	e.latestStats = Stats{
		LossStats:  lossStats,
		DelayStats: delayStats,
	}
}
<commit_msg>Add units to GetTargetBitrate<commit_after>package gcc

import (
	"sync"
	"time"

	"github.com/pion/interceptor"
	"github.com/pion/interceptor/internal/cc"
	"github.com/pion/rtcp"
	"github.com/pion/rtp"
)

const (
	transportCCURI = "http://www.ietf.org/id/draft-holmer-rmcat-transport-wide-cc-extensions-01"
	latestBitrate  = 10_000
	minBitrate     = 5_000
	maxBitrate     = 50_000_000
)

// Pacer is the interface implemented by packet pacers
type Pacer interface {
	interceptor.RTPWriter
	AddStream(ssrc uint32, writer interceptor.RTPWriter)
	SetTargetBitrate(int)
	Close() error
}

// Stats contains internal statistics of the bandwidth estimator
type Stats struct {
	LossStats
	DelayStats
}

// SendSideBWE implements a combination of loss and delay based GCC
type SendSideBWE struct {
	pacer           Pacer
	lossController  *lossBasedBandwidthEstimator
	delayController *delayController
	feedbackAdapter *cc.FeedbackAdapter

	onTargetBitrateChange func(bitrate int)

	lock          sync.Mutex
	latestStats   Stats
	latestBitrate int
	minBitrate    int
	maxBitrate    int

	close chan struct{}
}

// Option configures a bandwidth estimator
type Option func(*SendSideBWE) error

// SendSideBWEInitialBitrate sets the initial bitrate of new GCC interceptors
func SendSideBWEInitialBitrate(rate int) Option {
	return func(e *SendSideBWE) error {
		e.latestBitrate = rate
		return nil
	}
}

// SendSideBWEPacer sets the pacing algorithm to use.
func SendSideBWEPacer(p Pacer) Option {
	return func(e *SendSideBWE) error {
		e.pacer = p
		return nil
	}
}

// NewSendSideBWE creates a new sender side bandwidth estimator
func NewSendSideBWE(opts ...Option) (*SendSideBWE, error) {
	e := &SendSideBWE{
		pacer:                 nil,
		lossController:        nil,
		delayController:       nil,
		feedbackAdapter:       cc.NewFeedbackAdapter(),
		onTargetBitrateChange: nil,
		lock:                  sync.Mutex{},
		latestStats:           Stats{},
		latestBitrate:         latestBitrate,
		minBitrate:            minBitrate,
		maxBitrate:            maxBitrate,
		close:                 make(chan struct{}),
	}
	for _, opt := range opts {
		if err := opt(e); err != nil {
			return nil, err
		}
	}
	if e.pacer == nil {
		e.pacer = NewLeakyBucketPacer(e.latestBitrate)
	}
	e.lossController = newLossBasedBWE(e.latestBitrate)
	e.delayController = newDelayController(delayControllerConfig{
		nowFn:          time.Now,
		initialBitrate: e.latestBitrate,
		minBitrate:     e.minBitrate,
		maxBitrate:     e.maxBitrate,
	})

	e.delayController.onUpdate(e.onDelayUpdate)

	return e, nil
}

// AddStream adds a new stream to the bandwidth estimator
func (e *SendSideBWE) AddStream(info *interceptor.StreamInfo, writer interceptor.RTPWriter) interceptor.RTPWriter {
	var hdrExtID uint8
	for _, e := range info.RTPHeaderExtensions {
		if e.URI == transportCCURI {
			hdrExtID = uint8(e.ID)
			break
		}
	}

	e.pacer.AddStream(info.SSRC, interceptor.RTPWriterFunc(func(header *rtp.Header, payload []byte, attributes interceptor.Attributes) (int, error) {
		if hdrExtID != 0 {
			if attributes == nil {
				attributes = make(interceptor.Attributes)
			}
			attributes.Set(cc.TwccExtensionAttributesKey, hdrExtID)
		}
		if err := e.feedbackAdapter.OnSent(time.Now(), header, len(payload), attributes); err != nil {
			return 0, err
		}
		return writer.Write(header, payload, attributes)
	}))
	return e.pacer
}

// WriteRTCP adds some RTCP feedback to the bandwidth estimator
func (e *SendSideBWE) WriteRTCP(pkts []rtcp.Packet, attributes interceptor.Attributes) error {
	for _, pkt := range pkts {
		if fb, ok := pkt.(*rtcp.TransportLayerCC); ok {
			acks, err := e.feedbackAdapter.OnTransportCCFeedback(time.Now(), fb)
			if err != nil {
				return err
			}
			e.lossController.updateLossEstimate(acks)
			e.delayController.updateDelayEstimate(acks)
		}
	}
	return nil
}

// GetTargetBitrate returns the current target bitrate in bits per second
func (e *SendSideBWE) GetTargetBitrate() int {
	e.lock.Lock()
	defer e.lock.Unlock()

	return e.latestBitrate
}

// GetStats returns some internal statistics of the bandwidth estimator
func (e *SendSideBWE) GetStats() map[string]interface{} {
	e.lock.Lock()
	defer e.lock.Unlock()

	return map[string]interface{}{
		"lossTargetBitrate":  e.latestStats.LossStats.TargetBitrate,
		"averageLoss":        e.latestStats.AverageLoss,
		"delayTargetBitrate": e.latestStats.DelayStats.TargetBitrate,
		"delayMeasurement":   float64(e.latestStats.Measurement.Microseconds()) / 1000.0,
		"delayEstimate":      float64(e.latestStats.Estimate.Microseconds()) / 1000.0,
		"delayThreshold":     float64(e.latestStats.Threshold.Microseconds()) / 1000.0,
		"rtt":                float64(e.latestStats.RTT.Microseconds()) / 1000.0,
		"usage":              e.latestStats.Usage.String(),
		"state":              e.latestStats.State.String(),
	}
}

// OnTargetBitrateChange sets the callback that is called when the target
// bitrate in bits per second changes
func (e *SendSideBWE) OnTargetBitrateChange(f func(bitrate int)) {
	e.onTargetBitrateChange = f
}

// Close stops and closes the bandwidth estimator
func (e *SendSideBWE) Close() error {
	if err := e.delayController.Close(); err != nil {
		return err
	}
	close(e.close)
	return e.pacer.Close()
}

func (e *SendSideBWE) onDelayUpdate(delayStats DelayStats) {
	e.lock.Lock()
	defer e.lock.Unlock()

	lossStats := e.lossController.getEstimate(delayStats.TargetBitrate)
	bitrateChanged := false
	bitrate := minInt(delayStats.TargetBitrate, lossStats.TargetBitrate)
	if bitrate != e.latestBitrate {
		bitrateChanged = true
		e.latestBitrate = bitrate
		e.pacer.SetTargetBitrate(e.latestBitrate)
	}

	if bitrateChanged && e.onTargetBitrateChange != nil {
		go e.onTargetBitrateChange(bitrate)
	}

	e.latestStats = Stats{
		LossStats:  lossStats,
		DelayStats: delayStats,
	}
}
<|endoftext|>