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

// +build !privileged_tests

package api

import (
	"fmt"
	"testing"

	"github.com/cilium/cilium/pkg/checker"
	"github.com/cilium/cilium/pkg/labels"

	. "gopkg.in/check.v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8sLbls "k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/selection"
)

var _ = Suite(&PolicyAPITestSuite{})

func (s *PolicyAPITestSuite) TestSelectsAllEndpoints(c *C) {

	// Empty endpoint selector slice equates to a wildcard.
	selectorSlice := EndpointSelectorSlice{}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, false)

	selectorSlice = EndpointSelectorSlice{WildcardEndpointSelector}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, true)

	// Slice that contains wildcard and other selectors still selects all endpoints.
	selectorSlice = EndpointSelectorSlice{WildcardEndpointSelector, NewESFromLabels(labels.ParseSelectLabel("bar"))}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, true)

	selectorSlice = EndpointSelectorSlice{NewESFromLabels(labels.ParseSelectLabel("bar")), NewESFromLabels(labels.ParseSelectLabel("foo"))}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, false)
}

func (s *PolicyAPITestSuite) TestLabelSelectorToRequirements(c *C) {
	labelSelector := &metav1.LabelSelector{
		MatchLabels: map[string]string{
			"any.foo": "bar",
			"k8s.baz": "alice",
		},
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "any.foo",
				Operator: "NotIn",
				Values:   []string{"default"},
			},
		},
	}

	expRequirements := k8sLbls.Requirements{}
	req, err := k8sLbls.NewRequirement("any.foo", selection.Equals, []string{"bar"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)
	req, err = k8sLbls.NewRequirement("any.foo", selection.NotIn, []string{"default"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)
	req, err = k8sLbls.NewRequirement("k8s.baz", selection.Equals, []string{"alice"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)

	c.Assert(labelSelectorToRequirements(labelSelector), checker.DeepEquals, &expRequirements)
}

func benchmarkMatchesSetup(match string, count int) (EndpointSelector, labels.LabelArray) {
	stringLabels := []string{}
	for i := 0; i < count; i++ {
		stringLabels = append(stringLabels, fmt.Sprintf("%d", i))
	}
	lbls := labels.NewLabelsFromModel(stringLabels)
	return NewESFromLabels(lbls.ToSlice()...), labels.ParseLabelArray(match)
}

func BenchmarkMatchesValid1000(b *testing.B) {
	es, match := benchmarkMatchesSetup("42", 1000)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		es.Matches(match)
	}
}

func BenchmarkMatchesInvalid1000(b *testing.B) {
	es, match := benchmarkMatchesSetup("foo", 1000)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		es.Matches(match)
	}
}

func BenchmarkMatchesValid1000Parallel(b *testing.B) {
	es, match := benchmarkMatchesSetup("42", 1000)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			es.Matches(match)
		}
	})
}

func BenchmarkMatchesInvalid1000Parallel(b *testing.B) {
	es, match := benchmarkMatchesSetup("foo", 1000)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			es.Matches(match)
		}
	})
}
<commit_msg>policy/api: Add test case for EntityAll<commit_after>// Copyright 2018 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build !privileged_tests

package api

import (
	"fmt"
	"testing"

	"github.com/cilium/cilium/pkg/checker"
	"github.com/cilium/cilium/pkg/labels"

	. "gopkg.in/check.v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8sLbls "k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/selection"
)

var _ = Suite(&PolicyAPITestSuite{})

func (s *PolicyAPITestSuite) TestSelectsAllEndpoints(c *C) {

	// Empty endpoint selector slice does NOT equate to a wildcard.
	selectorSlice := EndpointSelectorSlice{}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, false)

	selectorSlice = EndpointSelectorSlice{WildcardEndpointSelector}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, true)

	// Entity "reserved:all" maps to WildcardEndpointSelector
	selectorSlice = EntitySlice{EntityAll}.GetAsEndpointSelectors()
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, true)

	// Slice that contains wildcard and other selectors still selects all endpoints.
	selectorSlice = EndpointSelectorSlice{WildcardEndpointSelector, NewESFromLabels(labels.ParseSelectLabel("bar"))}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, true)

	selectorSlice = EndpointSelectorSlice{NewESFromLabels(labels.ParseSelectLabel("bar")), NewESFromLabels(labels.ParseSelectLabel("foo"))}
	c.Assert(selectorSlice.SelectsAllEndpoints(), Equals, false)
}

func (s *PolicyAPITestSuite) TestLabelSelectorToRequirements(c *C) {
	labelSelector := &metav1.LabelSelector{
		MatchLabels: map[string]string{
			"any.foo": "bar",
			"k8s.baz": "alice",
		},
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      "any.foo",
				Operator: "NotIn",
				Values:   []string{"default"},
			},
		},
	}

	expRequirements := k8sLbls.Requirements{}
	req, err := k8sLbls.NewRequirement("any.foo", selection.Equals, []string{"bar"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)
	req, err = k8sLbls.NewRequirement("any.foo", selection.NotIn, []string{"default"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)
	req, err = k8sLbls.NewRequirement("k8s.baz", selection.Equals, []string{"alice"})
	c.Assert(err, IsNil)
	expRequirements = append(expRequirements, *req)

	c.Assert(labelSelectorToRequirements(labelSelector), checker.DeepEquals, &expRequirements)
}

func benchmarkMatchesSetup(match string, count int) (EndpointSelector, labels.LabelArray) {
	stringLabels := []string{}
	for i := 0; i < count; i++ {
		stringLabels = append(stringLabels, fmt.Sprintf("%d", i))
	}
	lbls := labels.NewLabelsFromModel(stringLabels)
	return NewESFromLabels(lbls.ToSlice()...), labels.ParseLabelArray(match)
}

func BenchmarkMatchesValid1000(b *testing.B) {
	es, match := benchmarkMatchesSetup("42", 1000)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		es.Matches(match)
	}
}

func BenchmarkMatchesInvalid1000(b *testing.B) {
	es, match := benchmarkMatchesSetup("foo", 1000)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		es.Matches(match)
	}
}

func BenchmarkMatchesValid1000Parallel(b *testing.B) {
	es, match := benchmarkMatchesSetup("42", 1000)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			es.Matches(match)
		}
	})
}

func BenchmarkMatchesInvalid1000Parallel(b *testing.B) {
	es, match := benchmarkMatchesSetup("foo", 1000)
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			es.Matches(match)
		}
	})
}
<|endoftext|>


<commit_before>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.GetBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	if cp == nil {
		return nil, errors.New("no pool")
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	if cp == nil {
		return
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens if the key exists.
// The value will be serialized into a JSON document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing happens if the key exists.
// The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := memcached.UnwrapMemcachedError(mc.Add(vb, k, 0, exp, v)); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

// A callback function to update a document
type UpdateFunc func(current []byte) (updated []byte, err error)

// Return this as the error from an UpdateFunc to cancel the Update operation.
const UpdateCancel = memcached.CASQuit

// Safe update of a document, avoiding conflicts by using CAS.
// The callback function will be invoked with the current raw document contents (or nil if the
// document doesn't exist); it should return the updated raw contents (or nil to delete.)
// If it decides not to change anything it can return UpdateCancel as the error.
//
// If another writer modifies the document between the get and the set, the callback will be
// invoked again with the newer value.
func (b *Bucket) Update(k string, exp int, callback UpdateFunc) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		var callbackErr error
		casFunc := func(current []byte) ([]byte, memcached.CasOp) {
			var updated []byte
			updated, callbackErr = callback(current)
			if callbackErr != nil {
				return nil, memcached.CASQuit
			} else if updated == nil {
				return updated, memcached.CASDelete
			} else {
				return updated, memcached.CASStore
			}
			panic("unreachable")
		}
		_, err := mc.CAS(vb, k, casFunc, exp)
		if err == memcached.CASQuit {
			err = callbackErr
		}
		return err
	})
}

// Install a design document.
func (b *Bucket) PutDDoc(docname string, value interface{}) error {
	u, err := b.randomBaseURL()
	if err != nil {
		return err
	}

	j, err := json.Marshal(value)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("/%s/_design/%s", b.Name, docname)
	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(j))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 201 {
		body, _ := ioutil.ReadAll(res.Body)
		return fmt.Errorf("Error installing view: %v / %s",
			res.Status, body)
	}

	return nil
}
<commit_msg>Build fix.  :(<commit_after>/*
A smart client for go.

Usage:

 client, err := couchbase.Connect("http://myserver:8091/")
 handleError(err)
 pool, err := client.GetPool("default")
 handleError(err)
 bucket, err := pool.GetBucket("MyAwesomeBucket")
 handleError(err)
 ...

or a shortcut for the bucket directly

 bucket, err := couchbase.GetBucket("http://myserver:8091/", "default", "default")
*/
package couchbase

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"sync/atomic"

	"github.com/dustin/gomemcached"
	"github.com/dustin/gomemcached/client"
)

type connectionPool struct {
	host        string
	name        string
	connections []*memcached.Client
	mutex       sync.Mutex
}

func (cp *connectionPool) Close() error {
	cp.mutex.Lock()
	defer cp.mutex.Unlock()
	for _, c := range cp.connections {
		c.Close()
	}
	cp.connections = []*memcached.Client{}
	return nil
}

func (cp *connectionPool) Get() (*memcached.Client, error) {
	if cp == nil {
		return nil, errors.New("no pool")
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if len(cp.connections) == 0 {
		conn, err := memcached.Connect("tcp", cp.host)
		if err != nil {
			return nil, err
		}
		if cp.name != "default" {
			conn.Auth(cp.name, "")
		}

		cp.connections = append(cp.connections, conn)
	}

	rv := cp.connections[0]
	cp.connections = cp.connections[1:]

	return rv, nil
}

func (cp *connectionPool) Return(c *memcached.Client) {
	if cp == nil {
		return
	}
	cp.mutex.Lock()
	defer cp.mutex.Unlock()

	if c != nil {
		if c.IsHealthy() {
			cp.connections = append(cp.connections, c)
		} else {
			c.Close()
		}
	}
}

// Execute a function on a memcached connection to the node owning key "k"
//
// Note that this automatically handles transient errors by replaying
// your function on a "not-my-vbucket" error, so don't assume
// your command will only be executed only once.
func (b *Bucket) Do(k string, f func(mc *memcached.Client, vb uint16) error) error {
	vb := b.VBHash(k)
	for {
		masterId := b.VBucketServerMap.VBucketMap[vb][0]
		conn, err := b.connections[masterId].Get()
		defer b.connections[masterId].Return(conn)
		if err != nil {
			return err
		}

		err = f(conn, uint16(vb))
		switch err.(type) {
		default:
			return err
		case gomemcached.MCResponse:
			st := err.(gomemcached.MCResponse).Status
			atomic.AddUint64(&b.pool.client.Statuses[st], 1)
			if st == gomemcached.NOT_MY_VBUCKET {
				b.refresh()
			} else {
				return err
			}
		}
	}
	panic("Unreachable.")
}

type gathered_stats struct {
	sn   string
	vals map[string]string
}

func getStatsParallel(b *Bucket, offset int, which string, ch chan<- gathered_stats) {
	sn := b.VBucketServerMap.ServerList[offset]

	results := map[string]string{}
	conn, err := b.connections[offset].Get()
	defer b.connections[offset].Return(conn)
	if err != nil {
		ch <- gathered_stats{sn, results}
	} else {
		st, err := conn.StatsMap(which)
		if err == nil {
			ch <- gathered_stats{sn, st}
		} else {
			ch <- gathered_stats{sn, results}
		}
	}
}

// Get a set of stats from all servers.
//
// Returns a map of server ID -> map of stat key to map value.
func (b *Bucket) GetStats(which string) map[string]map[string]string {
	rv := map[string]map[string]string{}

	if b.VBucketServerMap.ServerList == nil {
		return rv
	}
	// Go grab all the things at once.
	todo := len(b.VBucketServerMap.ServerList)
	ch := make(chan gathered_stats, todo)

	for offset, _ := range b.VBucketServerMap.ServerList {
		go getStatsParallel(b, offset, which, ch)
	}

	// Gather the results
	for i := 0; i < len(b.VBucketServerMap.ServerList); i++ {
		g := <-ch
		if len(g.vals) > 0 {
			rv[g.sn] = g.vals
		}
	}

	return rv
}

func (b *Bucket) doBulkGet(vb uint16, keys []string,
	ch chan<- map[string]*gomemcached.MCResponse) {

	masterId := b.VBucketServerMap.VBucketMap[vb][0]
	conn, err := b.connections[masterId].Get()
	if err != nil {
		ch <- map[string]*gomemcached.MCResponse{}
	}
	defer b.connections[masterId].Return(conn)

	m, err := conn.GetBulk(vb, keys)
	switch err.(type) {
	default:
		ch <- m
	case *gomemcached.MCResponse:
		fmt.Printf("Got a memcached error")
		st := err.(gomemcached.MCResponse).Status
		atomic.AddUint64(&b.pool.client.Statuses[st], 1)
		if st == gomemcached.NOT_MY_VBUCKET {
			b.refresh()
		}
		ch <- map[string]*gomemcached.MCResponse{}
	}
}

func (b *Bucket) processBulkGet(kdm map[uint16][]string,
	ch chan map[string]*gomemcached.MCResponse) {

	wch := make(chan uint16)

	worker := func() {
		for k := range wch {
			b.doBulkGet(k, kdm[k], ch)
		}
	}

	for i := 0; i < 4; i++ {
		go worker()
	}

	for k := range kdm {
		wch <- k
	}
	close(wch)

}
func (b *Bucket) GetBulk(keys []string) map[string]*gomemcached.MCResponse {
	// Organize by vbucket
	kdm := map[uint16][]string{}
	for _, k := range keys {
		vb := uint16(b.VBHash(k))
		a, ok := kdm[vb]
		if !ok {
			a = []string{}
		}
		kdm[vb] = append(a, k)
	}

	ch := make(chan map[string]*gomemcached.MCResponse)
	defer close(ch)

	go b.processBulkGet(kdm, ch)

	rv := map[string]*gomemcached.MCResponse{}
	for _ = range kdm {
		m := <-ch
		for k, v := range m {
			rv[k] = v
		}
	}

	return rv
}

// Set a value in this bucket.
// The value will be serialized into a JSON document.
func (b *Bucket) Set(k string, exp int, v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return err
	}
	return b.SetRaw(k, exp, data)
}

// Set a value in this bucket.
// The value will be stored as raw bytes.
func (b *Bucket) SetRaw(k string, exp int, v []byte) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Set(vb, k, 0, exp, v)
		if err != nil {
			return err
		}
		return nil
	})
}

// Adds a value to this bucket; like Set except that nothing happens if the key exists.
// The value will be serialized into a JSON document.
func (b *Bucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	data, err := json.Marshal(v)
	if err != nil {
		return false, err
	}
	return b.AddRaw(k, exp, data)
}

// Adds a value to this bucket; like SetRaw except that nothing happens if the key exists.
// The value will be stored as raw bytes.
func (b *Bucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	err = b.Do(k, func(mc *memcached.Client, vb uint16) error {
		switch res, err := memcached.UnwrapMemcachedError(mc.Add(vb, k, 0, exp, v)); {
		case err != nil:
			return err
		case res.Status == gomemcached.SUCCESS:
			added = true
		case res.Status != gomemcached.KEY_EEXISTS:
			return res
		}
		return nil
	})
	return
}

// Get a raw value from this bucket, including its CAS counter.
func (b *Bucket) GetsRaw(k string, cas *uint64) ([]byte, error) {
	var data []byte
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Get(vb, k)
		if err != nil {
			return err
		}
		if cas != nil {
			*cas = res.Cas
		}
		data = res.Body
		return nil
	})
	return data, err
}

// Get a value from this bucket, including its CAS counter.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Gets(k string, rv interface{}, cas *uint64) error {
	data, err := b.GetsRaw(k, cas)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, rv)
}

// Get a value from this bucket.
// The value is expected to be a JSON stream and will be deserialized
// into rv.
func (b *Bucket) Get(k string, rv interface{}) error {
	return b.Gets(k, rv, nil)
}

// Get a raw value from this bucket.
func (b *Bucket) GetRaw(k string) ([]byte, error) {
	return b.GetsRaw(k, nil)
}

// Delete a key from this bucket.
func (b *Bucket) Delete(k string) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		_, err := mc.Del(vb, k)
		if err != nil {
			return err
		}
		return nil
	})
}

// Increment a key
func (b *Bucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	var rv uint64
	err := b.Do(k, func(mc *memcached.Client, vb uint16) error {
		res, err := mc.Incr(vb, k, amt, def, exp)
		if err != nil {
			return err
		}
		rv = res
		return nil
	})
	return rv, err
}

// A callback function to update a document
type UpdateFunc func(current []byte) (updated []byte, err error)

// Return this as the error from an UpdateFunc to cancel the Update operation.
const UpdateCancel = memcached.CASQuit

// Safe update of a document, avoiding conflicts by using CAS.
// The callback function will be invoked with the current raw document contents (or nil if the
// document doesn't exist); it should return the updated raw contents (or nil to delete.)
// If it decides not to change anything it can return UpdateCancel as the error.
//
// If another writer modifies the document between the get and the set, the callback will be
// invoked again with the newer value.
func (b *Bucket) Update(k string, exp int, callback UpdateFunc) error {
	return b.Do(k, func(mc *memcached.Client, vb uint16) error {
		var callbackErr error
		casFunc := func(current []byte) ([]byte, memcached.CasOp) {
			var updated []byte
			updated, callbackErr = callback(current)
			if callbackErr != nil {
				return nil, memcached.CASQuit
			} else if updated == nil {
				return updated, memcached.CASDelete
			} else {
				return updated, memcached.CASStore
			}
			panic("unreachable")
		}
		_, err := mc.CAS(vb, k, casFunc, exp)
		if err == memcached.CASQuit {
			err = callbackErr
		}
		return err
	})
}

// Install a design document.
func (b *Bucket) PutDDoc(docname string, value interface{}) error {
	u, err := b.randomBaseURL()
	if err != nil {
		return err
	}

	j, err := json.Marshal(value)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("/%s/_design/%s", b.Name, docname)
	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(j))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")

	res, err := HttpClient.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode != 201 {
		body, _ := ioutil.ReadAll(res.Body)
		return fmt.Errorf("Error installing view: %v / %s",
			res.Status, body)
	}

	return nil
}
<|endoftext|>


<commit_before>package coap

import (
	"net"
	"time"
)

const RESPONSE_TIMEOUT = time.Second * 2

const RESPONSE_RANDOM_FACTOR = 1.5

const MAX_RETRANSMIT = 4

// A CoAP client connection.
type Conn struct {
	conn *net.UDPConn
	buf  []byte
}

// Get a CoAP client.
func Dial(n, addr string) (*Conn, error) {
	uaddr, err := net.ResolveUDPAddr(n, addr)
	if err != nil {
		return nil, err
	}

	s, err := net.DialUDP("udp", nil, uaddr)
	if err != nil {
		return nil, err
	}

	return &Conn{s, make([]byte, maxPktLen)}, nil
}

// Duration a message.  Get a response if there is one.
func (c *Conn) Send(req Message) (*Message, error) {
	err := Transmit(c.conn, nil, req)
	if err == nil {
		return nil, err
	}

	rv, err := Receive(c.conn, c.buf)

	return &rv, nil
}

// Receive a message.
func (c *Conn) Receive() (*Message, error) {
	rv, err := Receive(c.conn, c.buf)
	return &rv, err
}
<commit_msg>Don't wait for responses when there aren't any<commit_after>package coap

import (
	"net"
	"time"
)

const RESPONSE_TIMEOUT = time.Second * 2

const RESPONSE_RANDOM_FACTOR = 1.5

const MAX_RETRANSMIT = 4

// A CoAP client connection.
type Conn struct {
	conn *net.UDPConn
	buf  []byte
}

// Get a CoAP client.
func Dial(n, addr string) (*Conn, error) {
	uaddr, err := net.ResolveUDPAddr(n, addr)
	if err != nil {
		return nil, err
	}

	s, err := net.DialUDP("udp", nil, uaddr)
	if err != nil {
		return nil, err
	}

	return &Conn{s, make([]byte, maxPktLen)}, nil
}

// Duration a message.  Get a response if there is one.
func (c *Conn) Send(req Message) (*Message, error) {
	err := Transmit(c.conn, nil, req)
	if err == nil {
		return nil, err
	}

	if !req.IsConfirmable() {
		return nil, nil
	}

	rv, err := Receive(c.conn, c.buf)

	return &rv, nil
}

// Receive a message.
func (c *Conn) Receive() (*Message, error) {
	rv, err := Receive(c.conn, c.buf)
	return &rv, err
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package tiller

import (
	"fmt"
	"strings"

	ctx "golang.org/x/net/context"
	"k8s.io/helm/pkg/hooks"
	"k8s.io/helm/pkg/proto/hapi/release"
	"k8s.io/helm/pkg/proto/hapi/services"
	relutil "k8s.io/helm/pkg/releaseutil"
	"k8s.io/helm/pkg/timeconv"
)

// UninstallRelease deletes all of the resources associated with this release, and marks the release DELETED.
func (s *ReleaseServer) UninstallRelease(c ctx.Context, req *services.UninstallReleaseRequest) (*services.UninstallReleaseResponse, error) {
	err := s.env.Releases.LockRelease(req.Name)
	if err != nil {
		return nil, err
	}
	defer s.env.Releases.UnlockRelease(req.Name)

	if !ValidName.MatchString(req.Name) {
		s.Log("uninstall: Release not found: %s", req.Name)
		return nil, errMissingRelease
	}

	if len(req.Name) > releaseNameMaxLen {
		return nil, fmt.Errorf("release name %q exceeds max length of %d", req.Name, releaseNameMaxLen)
	}

	rels, err := s.env.Releases.History(req.Name)
	if err != nil {
		s.Log("uninstall: Release not loaded: %s", req.Name)
		return nil, err
	}
	if len(rels) < 1 {
		return nil, errMissingRelease
	}

	relutil.SortByRevision(rels)
	rel := rels[len(rels)-1]

	// TODO: Are there any cases where we want to force a delete even if it's
	// already marked deleted?
	if rel.Info.Status.Code == release.Status_DELETED {
		if req.Purge {
			if err := s.purgeReleases(rels...); err != nil {
				s.Log("uninstall: Failed to purge the release: %s", err)
				return nil, err
			}
			return &services.UninstallReleaseResponse{Release: rel}, nil
		}
		return nil, fmt.Errorf("the release named %q is already deleted", req.Name)
	}

	s.Log("uninstall: Deleting %s", req.Name)
	rel.Info.Status.Code = release.Status_DELETING
	rel.Info.Deleted = timeconv.Now()
	rel.Info.Description = "Deletion in progress (or silently failed)"
	res := &services.UninstallReleaseResponse{Release: rel}

	if !req.DisableHooks {
		if err := s.execHook(rel.Hooks, rel.Name, rel.Namespace, hooks.PreDelete, req.Timeout); err != nil {
			return res, err
		}
	} else {
		s.Log("delete hooks disabled for %s", req.Name)
	}

	// From here on out, the release is currently considered to be in Status_DELETING
	// state.
	if err := s.env.Releases.Update(rel); err != nil {
		s.Log("uninstall: Failed to store updated release: %s", err)
	}

	kept, errs := s.ReleaseModule.Delete(rel, req, s.env)
	res.Info = kept

	es := make([]string, 0, len(errs))
	for _, e := range errs {
		s.Log("error: %v", e)
		es = append(es, e.Error())
	}

	if !req.DisableHooks {
		if err := s.execHook(rel.Hooks, rel.Name, rel.Namespace, hooks.PostDelete, req.Timeout); err != nil {
			es = append(es, err.Error())
		}
	}

	rel.Info.Status.Code = release.Status_DELETED
	rel.Info.Description = "Deletion complete"

	if req.Purge {
		s.Log("purge requested for %s", req.Name)
		err := s.purgeReleases(rels...)
		if err != nil {
			s.Log("uninstall: Failed to purge the release: %s", err)
		}
		return res, err
	}

	if err := s.env.Releases.Update(rel); err != nil {
		s.Log("uninstall: Failed to store updated release: %s", err)
	}

	if len(es) > 0 {
		return res, fmt.Errorf("deletion completed with %d error(s): %s", len(es), strings.Join(es, "; "))
	}
	return res, nil
}

func (s *ReleaseServer) purgeReleases(rels ...*release.Release) error {
	for _, rel := range rels {
		if _, err := s.env.Releases.Delete(rel.Name, rel.Version); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>adjust the check parameter order for uninstall<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package tiller

import (
	"fmt"
	"strings"

	ctx "golang.org/x/net/context"
	"k8s.io/helm/pkg/hooks"
	"k8s.io/helm/pkg/proto/hapi/release"
	"k8s.io/helm/pkg/proto/hapi/services"
	relutil "k8s.io/helm/pkg/releaseutil"
	"k8s.io/helm/pkg/timeconv"
)

// UninstallRelease deletes all of the resources associated with this release, and marks the release DELETED.
func (s *ReleaseServer) UninstallRelease(c ctx.Context, req *services.UninstallReleaseRequest) (*services.UninstallReleaseResponse, error) {
	if !ValidName.MatchString(req.Name) {
		s.Log("uninstall: Release not found: %s", req.Name)
		return nil, errMissingRelease
	}

	if len(req.Name) > releaseNameMaxLen {
		return nil, fmt.Errorf("release name %q exceeds max length of %d", req.Name, releaseNameMaxLen)
	}

	err := s.env.Releases.LockRelease(req.Name)
	if err != nil {
		return nil, err
	}
	defer s.env.Releases.UnlockRelease(req.Name)

	rels, err := s.env.Releases.History(req.Name)
	if err != nil {
		s.Log("uninstall: Release not loaded: %s", req.Name)
		return nil, err
	}
	if len(rels) < 1 {
		return nil, errMissingRelease
	}

	relutil.SortByRevision(rels)
	rel := rels[len(rels)-1]

	// TODO: Are there any cases where we want to force a delete even if it's
	// already marked deleted?
	if rel.Info.Status.Code == release.Status_DELETED {
		if req.Purge {
			if err := s.purgeReleases(rels...); err != nil {
				s.Log("uninstall: Failed to purge the release: %s", err)
				return nil, err
			}
			return &services.UninstallReleaseResponse{Release: rel}, nil
		}
		return nil, fmt.Errorf("the release named %q is already deleted", req.Name)
	}

	s.Log("uninstall: Deleting %s", req.Name)
	rel.Info.Status.Code = release.Status_DELETING
	rel.Info.Deleted = timeconv.Now()
	rel.Info.Description = "Deletion in progress (or silently failed)"
	res := &services.UninstallReleaseResponse{Release: rel}

	if !req.DisableHooks {
		if err := s.execHook(rel.Hooks, rel.Name, rel.Namespace, hooks.PreDelete, req.Timeout); err != nil {
			return res, err
		}
	} else {
		s.Log("delete hooks disabled for %s", req.Name)
	}

	// From here on out, the release is currently considered to be in Status_DELETING
	// state.
	if err := s.env.Releases.Update(rel); err != nil {
		s.Log("uninstall: Failed to store updated release: %s", err)
	}

	kept, errs := s.ReleaseModule.Delete(rel, req, s.env)
	res.Info = kept

	es := make([]string, 0, len(errs))
	for _, e := range errs {
		s.Log("error: %v", e)
		es = append(es, e.Error())
	}

	if !req.DisableHooks {
		if err := s.execHook(rel.Hooks, rel.Name, rel.Namespace, hooks.PostDelete, req.Timeout); err != nil {
			es = append(es, err.Error())
		}
	}

	rel.Info.Status.Code = release.Status_DELETED
	rel.Info.Description = "Deletion complete"

	if req.Purge {
		s.Log("purge requested for %s", req.Name)
		err := s.purgeReleases(rels...)
		if err != nil {
			s.Log("uninstall: Failed to purge the release: %s", err)
		}
		return res, err
	}

	if err := s.env.Releases.Update(rel); err != nil {
		s.Log("uninstall: Failed to store updated release: %s", err)
	}

	if len(es) > 0 {
		return res, fmt.Errorf("deletion completed with %d error(s): %s", len(es), strings.Join(es, "; "))
	}
	return res, nil
}

func (s *ReleaseServer) purgeReleases(rels ...*release.Release) error {
	for _, rel := range rels {
		if _, err := s.env.Releases.Delete(rel.Name, rel.Version); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"os"
	"path/filepath"

	"github.com/gsamokovarov/jump/cli"
	"github.com/gsamokovarov/jump/config"
	"github.com/gsamokovarov/jump/fuzzy"
	"github.com/gsamokovarov/jump/scoring"
)

const proximity = 5
const osSeparator = string(os.PathSeparator)

func cdCmd(args cli.Args, conf *config.Config) {
	term := args.CommandName()
	entries, err := conf.ReadEntries()

	if err != nil {
		cli.Exitf(1, "%s\n", err)
	}

	// If an auto-completion triggered a full path, just go there.
	if filepath.IsAbs(term) {
		cli.Outf("%s\n", term)
		return
	}

	index := 0

	// If we happen to match the last term, e.g. j is called with no
	// arguments then jump to the previous search.
	if term == "" {
		search := conf.ReadSearch()

		term, index = search.Term, search.Index+1
	}

	fuzzyEntries := scoring.NewFuzzyEntries(entries, term)
	for {
		// Prefer an exact match if it's in a reasonable proximity of the best
		// match. Useful for jumping to 2 to 4 letter directories, which you
		// may just type anyway.
		index = exactMatchInProximity(fuzzyEntries, term, index)

		if entry, empty := fuzzyEntries.Select(index); !empty {
			// Remove the entries that no longer exists.
			if _, err := os.Stat(entry.Path); os.IsNotExist(err) {
				entries.Remove(entry.Path)
				conf.WriteEntries(entries)

				index++
				continue
			}

			// Jump to the next entry, if the jump is going to land on the
			// current directory.
			if cwd, err := os.Getwd(); err == nil && entry.Path == cwd {
				index++
				continue
			}

			cli.Outf("%s\n", entry.Path)
			conf.WriteSearch(term, index)
		}

		break
	}
}

func exactMatchInProximity(entries *scoring.FuzzyEntries, term string, offset int) int {
	norm := fuzzy.NewNormalizer(filepath.Base(term))
	normalizedTerm := norm.NormalizeTerm()

	for index := offset; index <= offset+proximity; index++ {
		if entry, empty := entries.Select(index); !empty {

			// Take only the base part, if you wanna do a deep search
			// like Dev/nes.
			basePath := filepath.Base(norm.NormalizePath(entry.Path))
			if basePath == normalizedTerm {
				return index
			}
		}
	}

	return offset
}

func init() {
	cli.RegisterCommand("cd", "Fuzzy match a directory to jump to.", cdCmd)
}
<commit_msg>Revert "Don't read the search history unless needed"<commit_after>package cmd

import (
	"os"
	"path/filepath"

	"github.com/gsamokovarov/jump/cli"
	"github.com/gsamokovarov/jump/config"
	"github.com/gsamokovarov/jump/fuzzy"
	"github.com/gsamokovarov/jump/scoring"
)

const proximity = 5
const osSeparator = string(os.PathSeparator)

func cdCmd(args cli.Args, conf *config.Config) {
	term := args.CommandName()
	entries, err := conf.ReadEntries()

	if err != nil {
		cli.Exitf(1, "%s\n", err)
	}

	// If an auto-completion triggered a full path, just go there.
	if filepath.IsAbs(term) {
		cli.Outf("%s\n", term)
		return
	}

	index, search := 0, conf.ReadSearch()

	// If we happen to match the last term, e.g. j is called with no
	// arguments then jump to the previous search.
	if len(term) == 0 {
		term, index = search.Term, search.Index+1
	}

	fuzzyEntries := scoring.NewFuzzyEntries(entries, term)
	for {
		// Prefer an exact match if it's in a reasonable proximity of the best
		// match. Useful for jumping to 2 to 4 letter directories, which you
		// may just type anyway.
		index = exactMatchInProximity(fuzzyEntries, term, index)

		if entry, empty := fuzzyEntries.Select(index); !empty {
			// Remove the entries that no longer exists.
			if _, err := os.Stat(entry.Path); os.IsNotExist(err) {
				entries.Remove(entry.Path)
				conf.WriteEntries(entries)

				index++
				continue
			}

			// Jump to the next entry, if the jump is going to land on the
			// current directory.
			if cwd, err := os.Getwd(); err == nil && entry.Path == cwd {
				index++
				continue
			}

			cli.Outf("%s\n", entry.Path)
			conf.WriteSearch(term, index)
		}

		break
	}
}

func exactMatchInProximity(entries *scoring.FuzzyEntries, term string, offset int) int {
	norm := fuzzy.NewNormalizer(filepath.Base(term))
	normalizedTerm := norm.NormalizeTerm()

	for index := offset; index <= offset+proximity; index++ {
		if entry, empty := entries.Select(index); !empty {

			// Take only the base part, if you wanna do a deep search
			// like Dev/nes.
			basePath := filepath.Base(norm.NormalizePath(entry.Path))
			if basePath == normalizedTerm {
				return index
			}
		}
	}

	return offset
}

func init() {
	cli.RegisterCommand("cd", "Fuzzy match a directory to jump to.", cdCmd)
}
<|endoftext|>


<commit_before>package adt

import (
	"bytes"
	"encoding/binary"
	"errors"
	"io"
	"strings"
)

var ErrReadingColumnDescriptor = errors.New("adt: error reading column descriptor")

type Column struct {
	Name          string
	Type          ColumnType
	Length        uint16
	Offset        uint16
	DecimalDigits uint16
}

func ColumnFromReader(r io.Reader) (*Column, error) {
	header := make([]byte, ColumnDescriptorLength)
	if _, err := io.ReadAtLeast(r, header, ColumnDescriptorLength); err != nil {
		if err == io.EOF {
			return nil, ErrReadingColumnDescriptor
		}
		return nil, err
	}
	column := &Column{
		Name: strings.Trim(string(header[:128]), "\x00"),
	}
	var columnType uint8
	if err := binary.Read(bytes.NewBuffer(header[129:130]), binary.LittleEndian, &columnType); err != nil {
		return nil, err
	}
	column.Type = ColumnType(columnType)
	if err := binary.Read(bytes.NewBuffer(header[131:133]), binary.LittleEndian, &column.Offset); err != nil {
		return nil, err
	}
	if err := binary.Read(bytes.NewBuffer(header[134:136]), binary.BigEndian, &column.Length); err != nil {
		return nil, err
	}
	if err := binary.Read(bytes.NewBuffer(header[138:140]), binary.BigEndian, &column.DecimalDigits); err != nil {
		return nil, err
	}
	return column, nil
}

func (c Column) String() string {
	return c.Name
}
<commit_msg>include type in string repr<commit_after>package adt

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"strings"
)

var ErrReadingColumnDescriptor = errors.New("adt: error reading column descriptor")

type Column struct {
	Name          string
	Type          ColumnType
	Length        uint16
	Offset        uint16
	DecimalDigits uint16
}

func ColumnFromReader(r io.Reader) (*Column, error) {
	header := make([]byte, ColumnDescriptorLength)
	if _, err := io.ReadAtLeast(r, header, ColumnDescriptorLength); err != nil {
		if err == io.EOF {
			return nil, ErrReadingColumnDescriptor
		}
		return nil, err
	}
	column := &Column{
		Name: strings.Trim(string(header[:128]), "\x00"),
	}
	var columnType uint8
	if err := binary.Read(bytes.NewBuffer(header[129:130]), binary.LittleEndian, &columnType); err != nil {
		return nil, err
	}
	column.Type = ColumnType(columnType)
	if err := binary.Read(bytes.NewBuffer(header[131:133]), binary.LittleEndian, &column.Offset); err != nil {
		return nil, err
	}
	if err := binary.Read(bytes.NewBuffer(header[134:136]), binary.BigEndian, &column.Length); err != nil {
		return nil, err
	}
	if err := binary.Read(bytes.NewBuffer(header[138:140]), binary.BigEndian, &column.DecimalDigits); err != nil {
		return nil, err
	}
	return column, nil
}

func (c Column) String() string {
	return fmt.Sprintf("%s (%s)", c.Name, c.Type)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package record

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"testing"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

func init() {
	// Don't bother sleeping between retries.
	sleepDuration = 0
}

type testEventRecorder struct {
	OnEvent func(e *api.Event) (*api.Event, error)
}

// CreateEvent records the event for testing.
func (t *testEventRecorder) Create(e *api.Event) (*api.Event, error) {
	if t.OnEvent != nil {
		return t.OnEvent(e)
	}
	return e, nil
}

func (t *testEventRecorder) clearOnEvent() {
	t.OnEvent = nil
}

func TestEventf(t *testing.T) {
	testPod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			SelfLink:  "/api/v1beta1/pods/foo",
			Name:      "foo",
			Namespace: "baz",
			UID:       "bar",
		},
	}
	testRef, err := api.GetPartialReference(testPod, "desiredState.manifest.containers[2]")
	if err != nil {
		t.Fatal(err)
	}
	table := []struct {
		obj        runtime.Object
		reason     string
		messageFmt string
		elements   []interface{}
		expect     *api.Event
		expectLog  string
	}{
		{
			obj:        testRef,
			reason:     "Started",
			messageFmt: "some verbose message: %v",
			elements:   []interface{}{1},
			expect: &api.Event{
				ObjectMeta: api.ObjectMeta{
					Name:      "foo",
					Namespace: "baz",
				},
				InvolvedObject: api.ObjectReference{
					Kind:       "Pod",
					Name:       "foo",
					Namespace:  "baz",
					UID:        "bar",
					APIVersion: "v1beta1",
					FieldPath:  "desiredState.manifest.containers[2]",
				},
				Reason:  "Started",
				Message: "some verbose message: 1",
				Source:  api.EventSource{Component: "eventTest"},
			},
			expectLog: `Event(api.ObjectReference{Kind:"Pod", Namespace:"baz", Name:"foo", UID:"bar", APIVersion:"v1beta1", ResourceVersion:"", FieldPath:"desiredState.manifest.containers[2]"}): reason: 'Started' some verbose message: 1`,
		},
		{
			obj:        testPod,
			reason:     "Started",
			messageFmt: "some verbose message: %v",
			elements:   []interface{}{1},
			expect: &api.Event{
				ObjectMeta: api.ObjectMeta{
					Name:      "foo",
					Namespace: "baz",
				},
				InvolvedObject: api.ObjectReference{
					Kind:       "Pod",
					Name:       "foo",
					Namespace:  "baz",
					UID:        "bar",
					APIVersion: "v1beta1",
				},
				Reason:  "Started",
				Message: "some verbose message: 1",
				Source:  api.EventSource{Component: "eventTest"},
			},
			expectLog: `Event(api.ObjectReference{Kind:"Pod", Namespace:"baz", Name:"foo", UID:"bar", APIVersion:"v1beta1", ResourceVersion:"", FieldPath:""}): reason: 'Started' some verbose message: 1`,
		},
	}

	for _, item := range table {
		called := make(chan struct{})
		testEvents := testEventRecorder{
			OnEvent: func(event *api.Event) (*api.Event, error) {
				a := *event
				// Just check that the timestamp was set.
				if a.Timestamp.IsZero() {
					t.Errorf("timestamp wasn't set")
				}
				a.Timestamp = item.expect.Timestamp
				// Check that name has the right prefix.
				if n, en := a.Name, item.expect.Name; !strings.HasPrefix(n, en) {
					t.Errorf("Name '%v' does not contain prefix '%v'", n, en)
				}
				a.Name = item.expect.Name
				if e, a := item.expect, &a; !reflect.DeepEqual(e, a) {
					t.Errorf("diff: %s", util.ObjectDiff(e, a))
				}
				called <- struct{}{}
				return event, nil
			},
		}
		recorder := StartRecording(&testEvents, api.EventSource{Component: "eventTest"})
		logger := StartLogging(t.Logf) // Prove that it is useful
		logger2 := StartLogging(func(formatter string, args ...interface{}) {
			if e, a := item.expectLog, fmt.Sprintf(formatter, args...); e != a {
				t.Errorf("Expected '%v', got '%v'", e, a)
			}
			called <- struct{}{}
		})

		Eventf(item.obj, item.reason, item.messageFmt, item.elements...)

		<-called
		<-called
		recorder.Stop()
		logger.Stop()
		logger2.Stop()
	}
}

func TestWriteEventError(t *testing.T) {
	ref := &api.ObjectReference{
		Kind:       "Pod",
		Name:       "foo",
		Namespace:  "baz",
		UID:        "bar",
		APIVersion: "v1beta1",
	}
	type entry struct {
		timesToSendError int
		attemptsMade     int
		attemptsWanted   int
		err              error
	}
	table := map[string]*entry{
		"giveUp1": {
			timesToSendError: 1000,
			attemptsWanted:   1,
			err:              &client.RequestConstructionError{},
		},
		"giveUp2": {
			timesToSendError: 1000,
			attemptsWanted:   1,
			err:              &errors.StatusError{},
		},
		"retry1": {
			timesToSendError: 1000,
			attemptsWanted:   12,
			err:              &errors.UnexpectedObjectError{},
		},
		"retry2": {
			timesToSendError: 1000,
			attemptsWanted:   12,
			err:              fmt.Errorf("A weird error"),
		},
		"succeedEventually": {
			timesToSendError: 2,
			attemptsWanted:   2,
			err:              fmt.Errorf("A weird error"),
		},
	}
	done := make(chan struct{})

	defer StartRecording(
		&testEventRecorder{
			OnEvent: func(event *api.Event) (*api.Event, error) {
				if event.Message == "finished" {
					close(done)
					return event, nil
				}
				item, ok := table[event.Message]
				if !ok {
					t.Errorf("Unexpected event: %#v", event)
					return event, nil
				}
				item.attemptsMade++
				if item.attemptsMade < item.timesToSendError {
					return nil, item.err
				}
				return event, nil
			},
		},
		api.EventSource{Component: "eventTest"},
	).Stop()

	for caseName := range table {
		Event(ref, "Reason", caseName)
	}
	Event(ref, "Reason", "finished")
	<-done

	for caseName, item := range table {
		if e, a := item.attemptsWanted, item.attemptsMade; e != a {
			t.Errorf("case %v: wanted %v, got %v attempts", caseName, e, a)
		}
	}
}

func TestLotsOfEvents(t *testing.T) {
	recorderCalled := make(chan struct{})
	loggerCalled := make(chan struct{})

	// Fail each event a few times to ensure there's some load on the tested code.
	var counts [1000]int
	testEvents := testEventRecorder{
		OnEvent: func(event *api.Event) (*api.Event, error) {
			num, err := strconv.Atoi(event.Message)
			if err != nil {
				t.Error(err)
				return event, nil
			}
			counts[num]++
			if counts[num] < 5 {
				return nil, fmt.Errorf("fake error")
			}
			recorderCalled <- struct{}{}
			return event, nil
		},
	}
	recorder := StartRecording(&testEvents, api.EventSource{Component: "eventTest"})
	logger := StartLogging(func(formatter string, args ...interface{}) {
		loggerCalled <- struct{}{}
	})

	ref := &api.ObjectReference{
		Kind:       "Pod",
		Name:       "foo",
		Namespace:  "baz",
		UID:        "bar",
		APIVersion: "v1beta1",
	}
	for i := 0; i < maxQueuedEvents; i++ {
		go Event(ref, "Status", "Reason", strconv.Itoa(i))
	}
	// Make sure no events were dropped by either of the listeners.
	for i := 0; i < maxQueuedEvents; i++ {
		<-recorderCalled
		<-loggerCalled
	}
	// Make sure that every event was attempted 5 times
	for i := 0; i < maxQueuedEvents; i++ {
		if counts[i] < 5 {
			t.Errorf("Only attempted to record event '%d' %d times.", i, counts[i])
		}
	}
	recorder.Stop()
	logger.Stop()
}
<commit_msg>fix build<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package record

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"testing"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/runtime"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

func init() {
	// Don't bother sleeping between retries.
	sleepDuration = 0
}

type testEventRecorder struct {
	OnEvent func(e *api.Event) (*api.Event, error)
}

// CreateEvent records the event for testing.
func (t *testEventRecorder) Create(e *api.Event) (*api.Event, error) {
	if t.OnEvent != nil {
		return t.OnEvent(e)
	}
	return e, nil
}

func (t *testEventRecorder) clearOnEvent() {
	t.OnEvent = nil
}

func TestEventf(t *testing.T) {
	testPod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			SelfLink:  "/api/v1beta1/pods/foo",
			Name:      "foo",
			Namespace: "baz",
			UID:       "bar",
		},
	}
	testRef, err := api.GetPartialReference(testPod, "desiredState.manifest.containers[2]")
	if err != nil {
		t.Fatal(err)
	}
	table := []struct {
		obj        runtime.Object
		reason     string
		messageFmt string
		elements   []interface{}
		expect     *api.Event
		expectLog  string
	}{
		{
			obj:        testRef,
			reason:     "Started",
			messageFmt: "some verbose message: %v",
			elements:   []interface{}{1},
			expect: &api.Event{
				ObjectMeta: api.ObjectMeta{
					Name:      "foo",
					Namespace: "baz",
				},
				InvolvedObject: api.ObjectReference{
					Kind:       "Pod",
					Name:       "foo",
					Namespace:  "baz",
					UID:        "bar",
					APIVersion: "v1beta1",
					FieldPath:  "desiredState.manifest.containers[2]",
				},
				Reason:  "Started",
				Message: "some verbose message: 1",
				Source:  api.EventSource{Component: "eventTest"},
			},
			expectLog: `Event(api.ObjectReference{Kind:"Pod", Namespace:"baz", Name:"foo", UID:"bar", APIVersion:"v1beta1", ResourceVersion:"", FieldPath:"desiredState.manifest.containers[2]"}): reason: 'Started' some verbose message: 1`,
		},
		{
			obj:        testPod,
			reason:     "Started",
			messageFmt: "some verbose message: %v",
			elements:   []interface{}{1},
			expect: &api.Event{
				ObjectMeta: api.ObjectMeta{
					Name:      "foo",
					Namespace: "baz",
				},
				InvolvedObject: api.ObjectReference{
					Kind:       "Pod",
					Name:       "foo",
					Namespace:  "baz",
					UID:        "bar",
					APIVersion: "v1beta1",
				},
				Reason:  "Started",
				Message: "some verbose message: 1",
				Source:  api.EventSource{Component: "eventTest"},
			},
			expectLog: `Event(api.ObjectReference{Kind:"Pod", Namespace:"baz", Name:"foo", UID:"bar", APIVersion:"v1beta1", ResourceVersion:"", FieldPath:""}): reason: 'Started' some verbose message: 1`,
		},
	}

	for _, item := range table {
		called := make(chan struct{})
		testEvents := testEventRecorder{
			OnEvent: func(event *api.Event) (*api.Event, error) {
				a := *event
				// Just check that the timestamp was set.
				if a.Timestamp.IsZero() {
					t.Errorf("timestamp wasn't set")
				}
				a.Timestamp = item.expect.Timestamp
				// Check that name has the right prefix.
				if n, en := a.Name, item.expect.Name; !strings.HasPrefix(n, en) {
					t.Errorf("Name '%v' does not contain prefix '%v'", n, en)
				}
				a.Name = item.expect.Name
				if e, a := item.expect, &a; !reflect.DeepEqual(e, a) {
					t.Errorf("diff: %s", util.ObjectDiff(e, a))
				}
				called <- struct{}{}
				return event, nil
			},
		}
		recorder := StartRecording(&testEvents, api.EventSource{Component: "eventTest"})
		logger := StartLogging(t.Logf) // Prove that it is useful
		logger2 := StartLogging(func(formatter string, args ...interface{}) {
			if e, a := item.expectLog, fmt.Sprintf(formatter, args...); e != a {
				t.Errorf("Expected '%v', got '%v'", e, a)
			}
			called <- struct{}{}
		})

		Eventf(item.obj, item.reason, item.messageFmt, item.elements...)

		<-called
		<-called
		recorder.Stop()
		logger.Stop()
		logger2.Stop()
	}
}

func TestWriteEventError(t *testing.T) {
	ref := &api.ObjectReference{
		Kind:       "Pod",
		Name:       "foo",
		Namespace:  "baz",
		UID:        "bar",
		APIVersion: "v1beta1",
	}
	type entry struct {
		timesToSendError int
		attemptsMade     int
		attemptsWanted   int
		err              error
	}
	table := map[string]*entry{
		"giveUp1": {
			timesToSendError: 1000,
			attemptsWanted:   1,
			err:              &client.RequestConstructionError{},
		},
		"giveUp2": {
			timesToSendError: 1000,
			attemptsWanted:   1,
			err:              &errors.StatusError{},
		},
		"retry1": {
			timesToSendError: 1000,
			attemptsWanted:   12,
			err:              &errors.UnexpectedObjectError{},
		},
		"retry2": {
			timesToSendError: 1000,
			attemptsWanted:   12,
			err:              fmt.Errorf("A weird error"),
		},
		"succeedEventually": {
			timesToSendError: 2,
			attemptsWanted:   2,
			err:              fmt.Errorf("A weird error"),
		},
	}
	done := make(chan struct{})

	defer StartRecording(
		&testEventRecorder{
			OnEvent: func(event *api.Event) (*api.Event, error) {
				if event.Message == "finished" {
					close(done)
					return event, nil
				}
				item, ok := table[event.Message]
				if !ok {
					t.Errorf("Unexpected event: %#v", event)
					return event, nil
				}
				item.attemptsMade++
				if item.attemptsMade < item.timesToSendError {
					return nil, item.err
				}
				return event, nil
			},
		},
		api.EventSource{Component: "eventTest"},
	).Stop()

	for caseName := range table {
		Event(ref, "Reason", caseName)
	}
	Event(ref, "Reason", "finished")
	<-done

	for caseName, item := range table {
		if e, a := item.attemptsWanted, item.attemptsMade; e != a {
			t.Errorf("case %v: wanted %v, got %v attempts", caseName, e, a)
		}
	}
}

func TestLotsOfEvents(t *testing.T) {
	recorderCalled := make(chan struct{})
	loggerCalled := make(chan struct{})

	// Fail each event a few times to ensure there's some load on the tested code.
	var counts [1000]int
	testEvents := testEventRecorder{
		OnEvent: func(event *api.Event) (*api.Event, error) {
			num, err := strconv.Atoi(event.Message)
			if err != nil {
				t.Error(err)
				return event, nil
			}
			counts[num]++
			if counts[num] < 5 {
				return nil, fmt.Errorf("fake error")
			}
			recorderCalled <- struct{}{}
			return event, nil
		},
	}
	recorder := StartRecording(&testEvents, api.EventSource{Component: "eventTest"})
	logger := StartLogging(func(formatter string, args ...interface{}) {
		loggerCalled <- struct{}{}
	})

	ref := &api.ObjectReference{
		Kind:       "Pod",
		Name:       "foo",
		Namespace:  "baz",
		UID:        "bar",
		APIVersion: "v1beta1",
	}
	for i := 0; i < maxQueuedEvents; i++ {
		go Event(ref, "Reason", strconv.Itoa(i))
	}
	// Make sure no events were dropped by either of the listeners.
	for i := 0; i < maxQueuedEvents; i++ {
		<-recorderCalled
		<-loggerCalled
	}
	// Make sure that every event was attempted 5 times
	for i := 0; i < maxQueuedEvents; i++ {
		if counts[i] < 5 {
			t.Errorf("Only attempted to record event '%d' %d times.", i, counts[i])
		}
	}
	recorder.Stop()
	logger.Stop()
}
<|endoftext|>


<commit_before>package provisioningv2

import (
	"embed"

	v1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	rkev1 "github.com/rancher/rancher/pkg/apis/rke.cattle.io/v1"
	"github.com/rancher/rancher/pkg/features"
	"github.com/rancher/wrangler/pkg/crd"
	"github.com/rancher/wrangler/pkg/data"
	"github.com/rancher/wrangler/pkg/yaml"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

var (
	capiCRDs = map[string]bool{
		"Machine":            true,
		"MachineHealthCheck": true,
		"MachineDeployment":  true,
		"MachineSet":         true,
		"Cluster":            true,
	}

	//go:embed capi-crds.yaml capi-webhooks.yaml
	capiData embed.FS
)

func List() (result []crd.CRD) {
	result = append(result, provisioning()...)
	if features.RKE2.Enabled() {
		result = append(result, rke2()...)
	}
	if features.EmbeddedClusterAPI.Enabled() {
		result = append(result, capi()...)
	}
	return
}

func provisioning() []crd.CRD {
	return []crd.CRD{
		newRancherCRD(&v1.Cluster{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"auth.cattle.io/cluster-indexed": "true",
			}
			return c.
				WithColumn("Ready", ".status.ready").
				WithColumn("Kubeconfig", ".status.clientSecretName")
		}),
	}
}

func clusterIndexed(c crd.CRD) crd.CRD {
	newLabels := map[string]string{}
	for k, v := range c.Labels {
		newLabels[k] = v
	}
	newLabels["auth.cattle.io/cluster-indexed"] = "true"
	c.Labels = newLabels
	return c
}

func rke2() []crd.CRD {
	return []crd.CRD{
		newRancherCRD(&v1.Cluster{}, func(c crd.CRD) crd.CRD {
			return clusterIndexed(c).
				WithColumn("Ready", ".status.ready").
				WithColumn("Kubeconfig", ".status.clientSecretName")
		}),
		newRKECRD(&rkev1.RKECluster{}, func(c crd.CRD) crd.CRD {
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEControlPlane{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEBootstrap{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEBootstrapTemplate{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEControlPlane{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.CustomMachine{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
	}
}

func Webhooks() []runtime.Object {
	if features.EmbeddedClusterAPI.Enabled() {
		return capiWebhooks()
	}
	return nil
}

func capiWebhooks() []runtime.Object {
	f, err := capiData.Open("capi-webhooks.yaml")
	if err != nil {
		panic(err)
	}
	defer f.Close()

	objs, err := yaml.ToObjects(f)
	if err != nil {
		panic(err)
	}

	return objs
}

func capi() []crd.CRD {
	f, err := capiData.Open("capi-crds.yaml")
	if err != nil {
		panic(err)
	}
	defer f.Close()

	objs, err := yaml.ToObjects(f)
	if err != nil {
		panic(err)
	}

	var result []crd.CRD
	for _, obj := range objs {
		if obj.GetObjectKind().GroupVersionKind().Kind != "CustomResourceDefinition" {
			continue
		}
		if unstr, ok := obj.(*unstructured.Unstructured); ok &&
			capiCRDs[data.Object(unstr.Object).String("spec", "names", "kind")] {
			labels := unstr.GetLabels()
			if labels == nil {
				labels = map[string]string{}
			}
			labels["auth.cattle.io/cluster-indexed"] = "true"
			unstr.SetLabels(labels)
			result = append(result, crd.CRD{
				Override: obj,
			})
		}
	}

	return result
}

func newRKECRD(obj interface{}, customize func(crd.CRD) crd.CRD) crd.CRD {
	crd := crd.CRD{
		GVK: schema.GroupVersionKind{
			Group:   "rke.cattle.io",
			Version: "v1",
		},
		Status:       true,
		SchemaObject: obj,
	}
	if customize != nil {
		crd = customize(crd)
	}
	return crd
}

func newRancherCRD(obj interface{}, customize func(crd.CRD) crd.CRD) crd.CRD {
	crd := crd.CRD{
		GVK: schema.GroupVersionKind{
			Group:   "provisioning.cattle.io",
			Version: "v1",
		},
		Status:       true,
		SchemaObject: obj,
	}
	if customize != nil {
		crd = customize(crd)
	}
	return crd
}
<commit_msg>add contract labels for rkeclusters.rke.cattle.io<commit_after>package provisioningv2

import (
	"embed"

	v1 "github.com/rancher/rancher/pkg/apis/provisioning.cattle.io/v1"
	rkev1 "github.com/rancher/rancher/pkg/apis/rke.cattle.io/v1"
	"github.com/rancher/rancher/pkg/features"
	"github.com/rancher/wrangler/pkg/crd"
	"github.com/rancher/wrangler/pkg/data"
	"github.com/rancher/wrangler/pkg/yaml"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

var (
	capiCRDs = map[string]bool{
		"Machine":            true,
		"MachineHealthCheck": true,
		"MachineDeployment":  true,
		"MachineSet":         true,
		"Cluster":            true,
	}

	//go:embed capi-crds.yaml capi-webhooks.yaml
	capiData embed.FS
)

func List() (result []crd.CRD) {
	result = append(result, provisioning()...)
	if features.RKE2.Enabled() {
		result = append(result, rke2()...)
	}
	if features.EmbeddedClusterAPI.Enabled() {
		result = append(result, capi()...)
	}
	return
}

func provisioning() []crd.CRD {
	return []crd.CRD{
		newRancherCRD(&v1.Cluster{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"auth.cattle.io/cluster-indexed": "true",
			}
			return c.
				WithColumn("Ready", ".status.ready").
				WithColumn("Kubeconfig", ".status.clientSecretName")
		}),
	}
}

func clusterIndexed(c crd.CRD) crd.CRD {
	newLabels := map[string]string{}
	for k, v := range c.Labels {
		newLabels[k] = v
	}
	newLabels["auth.cattle.io/cluster-indexed"] = "true"
	c.Labels = newLabels
	return c
}

func rke2() []crd.CRD {
	return []crd.CRD{
		newRancherCRD(&v1.Cluster{}, func(c crd.CRD) crd.CRD {
			return clusterIndexed(c).
				WithColumn("Ready", ".status.ready").
				WithColumn("Kubeconfig", ".status.clientSecretName")
		}),
		newRKECRD(&rkev1.RKECluster{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEControlPlane{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEBootstrap{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEBootstrapTemplate{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.RKEControlPlane{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
		newRKECRD(&rkev1.CustomMachine{}, func(c crd.CRD) crd.CRD {
			c.Labels = map[string]string{
				"cluster.x-k8s.io/v1alpha4": "v1",
			}
			return clusterIndexed(c)
		}),
	}
}

func Webhooks() []runtime.Object {
	if features.EmbeddedClusterAPI.Enabled() {
		return capiWebhooks()
	}
	return nil
}

func capiWebhooks() []runtime.Object {
	f, err := capiData.Open("capi-webhooks.yaml")
	if err != nil {
		panic(err)
	}
	defer f.Close()

	objs, err := yaml.ToObjects(f)
	if err != nil {
		panic(err)
	}

	return objs
}

func capi() []crd.CRD {
	f, err := capiData.Open("capi-crds.yaml")
	if err != nil {
		panic(err)
	}
	defer f.Close()

	objs, err := yaml.ToObjects(f)
	if err != nil {
		panic(err)
	}

	var result []crd.CRD
	for _, obj := range objs {
		if obj.GetObjectKind().GroupVersionKind().Kind != "CustomResourceDefinition" {
			continue
		}
		if unstr, ok := obj.(*unstructured.Unstructured); ok &&
			capiCRDs[data.Object(unstr.Object).String("spec", "names", "kind")] {
			labels := unstr.GetLabels()
			if labels == nil {
				labels = map[string]string{}
			}
			labels["auth.cattle.io/cluster-indexed"] = "true"
			unstr.SetLabels(labels)
			result = append(result, crd.CRD{
				Override: obj,
			})
		}
	}

	return result
}

func newRKECRD(obj interface{}, customize func(crd.CRD) crd.CRD) crd.CRD {
	crd := crd.CRD{
		GVK: schema.GroupVersionKind{
			Group:   "rke.cattle.io",
			Version: "v1",
		},
		Status:       true,
		SchemaObject: obj,
	}
	if customize != nil {
		crd = customize(crd)
	}
	return crd
}

func newRancherCRD(obj interface{}, customize func(crd.CRD) crd.CRD) crd.CRD {
	crd := crd.CRD{
		GVK: schema.GroupVersionKind{
			Group:   "provisioning.cattle.io",
			Version: "v1",
		},
		Status:       true,
		SchemaObject: obj,
	}
	if customize != nil {
		crd = customize(crd)
	}
	return crd
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package components

import (
	"fmt"
	"strings"

	"k8s.io/klog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// KubeletOptionsBuilder adds options for kubelets
type KubeletOptionsBuilder struct {
	*OptionsContext
}

var _ loader.OptionsBuilder = &KubeletOptionsBuilder{}

// BuildOptions is responsible for filling the defaults for the kubelet
func (b *KubeletOptionsBuilder) BuildOptions(o interface{}) error {
	clusterSpec := o.(*kops.ClusterSpec)

	if clusterSpec.Kubelet == nil {
		clusterSpec.Kubelet = &kops.KubeletConfigSpec{}
	}
	if clusterSpec.MasterKubelet == nil {
		clusterSpec.MasterKubelet = &kops.KubeletConfigSpec{}
	}

	if clusterSpec.KubeAPIServer != nil && clusterSpec.KubeAPIServer.EnableBootstrapAuthToken != nil {
		if *clusterSpec.KubeAPIServer.EnableBootstrapAuthToken {
			if clusterSpec.Kubelet.BootstrapKubeconfig == "" {
				clusterSpec.Kubelet.BootstrapKubeconfig = "/var/lib/kubelet/bootstrap-kubeconfig"
			}
		}
	}

	// Standard options
	clusterSpec.Kubelet.EnableDebuggingHandlers = fi.Bool(true)
	clusterSpec.Kubelet.PodManifestPath = "/etc/kubernetes/manifests"
	clusterSpec.Kubelet.LogLevel = fi.Int32(2)
	clusterSpec.Kubelet.ClusterDomain = clusterSpec.ClusterDNSDomain
	clusterSpec.Kubelet.NonMasqueradeCIDR = clusterSpec.NonMasqueradeCIDR

	// AllowPrivileged is deprecated and removed in v1.14.
	// See https://github.com/kubernetes/kubernetes/pull/71835
	if b.IsKubernetesGTE("1.14") {
		if clusterSpec.Kubelet.AllowPrivileged != nil {
			// If it is explicitly set to false, return an error, because this
			// behavior is no longer supported in v1.14 (the default was true, prior).
			if !*clusterSpec.Kubelet.AllowPrivileged {
				klog.Warningf("Kubelet's --allow-privileged flag is no longer supported in v1.14.")
			}
			// Explicitly set it to nil, so it won't be passed on the command line.
			clusterSpec.Kubelet.AllowPrivileged = nil
		}
	} else {
		clusterSpec.Kubelet.AllowPrivileged = fi.Bool(true)
	}

	if clusterSpec.Kubelet.ClusterDNS == "" {
		if clusterSpec.KubeDNS != nil && clusterSpec.KubeDNS.NodeLocalDNS != nil && fi.BoolValue(clusterSpec.KubeDNS.NodeLocalDNS.Enabled) &&
			((clusterSpec.KubeProxy != nil && clusterSpec.KubeProxy.ProxyMode == "ipvs") || (clusterSpec.Networking != nil && clusterSpec.Networking.Cilium != nil)) {
			clusterSpec.Kubelet.ClusterDNS = clusterSpec.KubeDNS.NodeLocalDNS.LocalIP
		} else {
			ip, err := WellKnownServiceIP(clusterSpec, 10)
			if err != nil {
				return err
			}
			clusterSpec.Kubelet.ClusterDNS = ip.String()
		}
	}

	clusterSpec.MasterKubelet.RegisterSchedulable = fi.Bool(false)
	// Replace the CIDR with a CIDR allocated by KCM (the default, but included for clarity)
	// We _do_ allow debugging handlers, so we can do logs
	// This does allow more access than we would like though
	clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(true)

	{
		// For pod eviction in low memory or empty disk situations
		if clusterSpec.Kubelet.EvictionHard == nil {
			evictionHard := []string{
				// TODO: Some people recommend 250Mi, but this would hurt small machines
				"memory.available<100Mi",
			}
			// Disk based evictions are not detecting the correct disk capacity in Kubernetes 1.19 and are
			// blocking scheduling by tainting worker nodes with "node.kubernetes.io/disk-pressure:NoSchedule"
			// TODO: Re-enable once the Kubelet issue is fixed
			if b.IsKubernetesLT("1.19") {
				// Disk based eviction (evict old images)
				// We don't need to specify both, but it seems harmless / safer
				evictionHard = append(evictionHard, "nodefs.available<10%")
				evictionHard = append(evictionHard, "nodefs.inodesFree<5%")
				evictionHard = append(evictionHard, "imagefs.available<10%")
				evictionHard = append(evictionHard, "imagefs.inodesFree<5%")
			}
			clusterSpec.Kubelet.EvictionHard = fi.String(strings.Join(evictionHard, ","))
		}
	}

	// use kubeconfig instead of api-servers
	const kubeconfigPath = "/var/lib/kubelet/kubeconfig"
	clusterSpec.Kubelet.KubeconfigPath = kubeconfigPath
	clusterSpec.MasterKubelet.KubeconfigPath = kubeconfigPath

	// IsolateMasters enables the legacy behaviour, where master pods on a separate network
	// In newer versions of kubernetes, most of that functionality has been removed though
	if fi.BoolValue(clusterSpec.IsolateMasters) {
		clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(false)
		clusterSpec.MasterKubelet.HairpinMode = "none"
	}

	cloudProvider := kops.CloudProviderID(clusterSpec.CloudProvider)

	clusterSpec.Kubelet.CgroupRoot = "/"

	klog.V(1).Infof("Cloud Provider: %s", cloudProvider)
	if cloudProvider == kops.CloudProviderAWS {
		clusterSpec.Kubelet.CloudProvider = "aws"

		// Use the hostname from the AWS metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@aws"
		}
	}

	if cloudProvider == kops.CloudProviderDO {
		clusterSpec.Kubelet.CloudProvider = "external"
		clusterSpec.Kubelet.HostnameOverride = "@digitalocean"
	}

	if cloudProvider == kops.CloudProviderGCE {
		clusterSpec.Kubelet.CloudProvider = "gce"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"

		if clusterSpec.CloudConfig == nil {
			clusterSpec.CloudConfig = &kops.CloudConfiguration{}
		}
		clusterSpec.CloudConfig.Multizone = fi.Bool(true)
		clusterSpec.CloudConfig.NodeTags = fi.String(GCETagForRole(b.ClusterName, kops.InstanceGroupRoleNode))

		// Use the hostname from the GCE metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@gce"
		}
	}

	if cloudProvider == kops.CloudProviderOpenstack {
		clusterSpec.Kubelet.CloudProvider = "openstack"
	}

	if cloudProvider == kops.CloudProviderALI {
		clusterSpec.Kubelet.CloudProvider = "alicloud"
		clusterSpec.Kubelet.HostnameOverride = "@alicloud"
	}

	if featureflag.EnableExternalCloudController.Enabled() && clusterSpec.ExternalCloudControllerManager != nil {
		clusterSpec.Kubelet.CloudProvider = "external"
	}

	if clusterSpec.ContainerRuntime == "docker" || clusterSpec.ContainerRuntime == "" {
		networking := clusterSpec.Networking
		if networking == nil {
			return fmt.Errorf("no networking mode set")

		}
		if UsesKubenet(networking) {
			clusterSpec.Kubelet.NetworkPluginName = "kubenet"

			// AWS MTU is 9001
			clusterSpec.Kubelet.NetworkPluginMTU = fi.Int32(9001)
		}

		// Specify our pause image
		image := "k8s.gcr.io/pause:3.2"
		var err error
		if image, err = b.AssetBuilder.RemapImage(image); err != nil {
			return err
		}
		clusterSpec.Kubelet.PodInfraContainerImage = image
	}

	if clusterSpec.Kubelet.FeatureGates == nil {
		clusterSpec.Kubelet.FeatureGates = make(map[string]string)
	}
	if _, found := clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"]; !found {
		if b.IsKubernetesLT("1.16") {
			clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"] = "true"
		}
	}

	return nil
}
<commit_msg>Re-enable disk based evictions for Kubernetes 1.19<commit_after>/*
Copyright 2019 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package components

import (
	"fmt"
	"strings"

	"k8s.io/klog"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/loader"
)

// KubeletOptionsBuilder adds options for kubelets
type KubeletOptionsBuilder struct {
	*OptionsContext
}

var _ loader.OptionsBuilder = &KubeletOptionsBuilder{}

// BuildOptions is responsible for filling the defaults for the kubelet
func (b *KubeletOptionsBuilder) BuildOptions(o interface{}) error {
	clusterSpec := o.(*kops.ClusterSpec)

	if clusterSpec.Kubelet == nil {
		clusterSpec.Kubelet = &kops.KubeletConfigSpec{}
	}
	if clusterSpec.MasterKubelet == nil {
		clusterSpec.MasterKubelet = &kops.KubeletConfigSpec{}
	}

	if clusterSpec.KubeAPIServer != nil && clusterSpec.KubeAPIServer.EnableBootstrapAuthToken != nil {
		if *clusterSpec.KubeAPIServer.EnableBootstrapAuthToken {
			if clusterSpec.Kubelet.BootstrapKubeconfig == "" {
				clusterSpec.Kubelet.BootstrapKubeconfig = "/var/lib/kubelet/bootstrap-kubeconfig"
			}
		}
	}

	// Standard options
	clusterSpec.Kubelet.EnableDebuggingHandlers = fi.Bool(true)
	clusterSpec.Kubelet.PodManifestPath = "/etc/kubernetes/manifests"
	clusterSpec.Kubelet.LogLevel = fi.Int32(2)
	clusterSpec.Kubelet.ClusterDomain = clusterSpec.ClusterDNSDomain
	clusterSpec.Kubelet.NonMasqueradeCIDR = clusterSpec.NonMasqueradeCIDR

	// AllowPrivileged is deprecated and removed in v1.14.
	// See https://github.com/kubernetes/kubernetes/pull/71835
	if b.IsKubernetesGTE("1.14") {
		if clusterSpec.Kubelet.AllowPrivileged != nil {
			// If it is explicitly set to false, return an error, because this
			// behavior is no longer supported in v1.14 (the default was true, prior).
			if !*clusterSpec.Kubelet.AllowPrivileged {
				klog.Warningf("Kubelet's --allow-privileged flag is no longer supported in v1.14.")
			}
			// Explicitly set it to nil, so it won't be passed on the command line.
			clusterSpec.Kubelet.AllowPrivileged = nil
		}
	} else {
		clusterSpec.Kubelet.AllowPrivileged = fi.Bool(true)
	}

	if clusterSpec.Kubelet.ClusterDNS == "" {
		if clusterSpec.KubeDNS != nil && clusterSpec.KubeDNS.NodeLocalDNS != nil && fi.BoolValue(clusterSpec.KubeDNS.NodeLocalDNS.Enabled) &&
			((clusterSpec.KubeProxy != nil && clusterSpec.KubeProxy.ProxyMode == "ipvs") || (clusterSpec.Networking != nil && clusterSpec.Networking.Cilium != nil)) {
			clusterSpec.Kubelet.ClusterDNS = clusterSpec.KubeDNS.NodeLocalDNS.LocalIP
		} else {
			ip, err := WellKnownServiceIP(clusterSpec, 10)
			if err != nil {
				return err
			}
			clusterSpec.Kubelet.ClusterDNS = ip.String()
		}
	}

	clusterSpec.MasterKubelet.RegisterSchedulable = fi.Bool(false)
	// Replace the CIDR with a CIDR allocated by KCM (the default, but included for clarity)
	// We _do_ allow debugging handlers, so we can do logs
	// This does allow more access than we would like though
	clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(true)

	{
		// For pod eviction in low memory or empty disk situations
		if clusterSpec.Kubelet.EvictionHard == nil {
			evictionHard := []string{
				// TODO: Some people recommend 250Mi, but this would hurt small machines
				"memory.available<100Mi",

				// Disk based eviction (evict old images)
				// We don't need to specify both, but it seems harmless / safer
				"nodefs.available<10%",
				"nodefs.inodesFree<5%",
				"imagefs.available<10%",
				"imagefs.inodesFree<5%",
			}
			clusterSpec.Kubelet.EvictionHard = fi.String(strings.Join(evictionHard, ","))
		}
	}

	// use kubeconfig instead of api-servers
	const kubeconfigPath = "/var/lib/kubelet/kubeconfig"
	clusterSpec.Kubelet.KubeconfigPath = kubeconfigPath
	clusterSpec.MasterKubelet.KubeconfigPath = kubeconfigPath

	// IsolateMasters enables the legacy behaviour, where master pods on a separate network
	// In newer versions of kubernetes, most of that functionality has been removed though
	if fi.BoolValue(clusterSpec.IsolateMasters) {
		clusterSpec.MasterKubelet.EnableDebuggingHandlers = fi.Bool(false)
		clusterSpec.MasterKubelet.HairpinMode = "none"
	}

	cloudProvider := kops.CloudProviderID(clusterSpec.CloudProvider)

	clusterSpec.Kubelet.CgroupRoot = "/"

	klog.V(1).Infof("Cloud Provider: %s", cloudProvider)
	if cloudProvider == kops.CloudProviderAWS {
		clusterSpec.Kubelet.CloudProvider = "aws"

		// Use the hostname from the AWS metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@aws"
		}
	}

	if cloudProvider == kops.CloudProviderDO {
		clusterSpec.Kubelet.CloudProvider = "external"
		clusterSpec.Kubelet.HostnameOverride = "@digitalocean"
	}

	if cloudProvider == kops.CloudProviderGCE {
		clusterSpec.Kubelet.CloudProvider = "gce"
		clusterSpec.Kubelet.HairpinMode = "promiscuous-bridge"

		if clusterSpec.CloudConfig == nil {
			clusterSpec.CloudConfig = &kops.CloudConfiguration{}
		}
		clusterSpec.CloudConfig.Multizone = fi.Bool(true)
		clusterSpec.CloudConfig.NodeTags = fi.String(GCETagForRole(b.ClusterName, kops.InstanceGroupRoleNode))

		// Use the hostname from the GCE metadata service
		// if hostnameOverride is not set.
		if clusterSpec.Kubelet.HostnameOverride == "" {
			clusterSpec.Kubelet.HostnameOverride = "@gce"
		}
	}

	if cloudProvider == kops.CloudProviderOpenstack {
		clusterSpec.Kubelet.CloudProvider = "openstack"
	}

	if cloudProvider == kops.CloudProviderALI {
		clusterSpec.Kubelet.CloudProvider = "alicloud"
		clusterSpec.Kubelet.HostnameOverride = "@alicloud"
	}

	if featureflag.EnableExternalCloudController.Enabled() && clusterSpec.ExternalCloudControllerManager != nil {
		clusterSpec.Kubelet.CloudProvider = "external"
	}

	if clusterSpec.ContainerRuntime == "docker" || clusterSpec.ContainerRuntime == "" {
		networking := clusterSpec.Networking
		if networking == nil {
			return fmt.Errorf("no networking mode set")

		}
		if UsesKubenet(networking) {
			clusterSpec.Kubelet.NetworkPluginName = "kubenet"

			// AWS MTU is 9001
			clusterSpec.Kubelet.NetworkPluginMTU = fi.Int32(9001)
		}

		// Specify our pause image
		image := "k8s.gcr.io/pause:3.2"
		var err error
		if image, err = b.AssetBuilder.RemapImage(image); err != nil {
			return err
		}
		clusterSpec.Kubelet.PodInfraContainerImage = image
	}

	if clusterSpec.Kubelet.FeatureGates == nil {
		clusterSpec.Kubelet.FeatureGates = make(map[string]string)
	}
	if _, found := clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"]; !found {
		if b.IsKubernetesLT("1.16") {
			clusterSpec.Kubelet.FeatureGates["ExperimentalCriticalPodAnnotation"] = "true"
		}
	}

	return nil
}
<|endoftext|>


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

package transport

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/name"
)

type challenge string

const (
	anonymous challenge = "anonymous"
	basic     challenge = "basic"
	bearer    challenge = "bearer"
)

type pingResp struct {
	challenge challenge

	// Following the challenge there are often key/value pairs
	// e.g. Bearer service="gcr.io",realm="https://auth.gcr.io/v36/tokenz"
	parameters map[string]string

	// The registry's scheme to use. Communicates whether we fell back to http.
	scheme string
}

func (c challenge) Canonical() challenge {
	return challenge(strings.ToLower(string(c)))
}

func parseChallenge(suffix string) map[string]string {
	kv := make(map[string]string)
	for _, token := range strings.Split(suffix, ",") {
		// Trim any whitespace around each token.
		token = strings.Trim(token, " ")

		// Break the token into a key/value pair
		if parts := strings.SplitN(token, "=", 2); len(parts) == 2 {
			// Unquote the value, if it is quoted.
			kv[parts[0]] = strings.Trim(parts[1], `"`)
		} else {
			// If there was only one part, treat is as a key with an empty value
			kv[token] = ""
		}
	}
	return kv
}

func ping(reg name.Registry, t http.RoundTripper) (*pingResp, error) {
	client := http.Client{Transport: t, Timeout: 2 * time.Second}

	// This first attempts to use "https" for every request, falling back to http
	// if the registry matches our localhost heuristic or if it is intentionally
	// set to insecure via name.NewInsecureRegistry.
	schemes := []string{"https"}
	if reg.Scheme() == "http" {
		schemes = append(schemes, "http")
	}

	var connErr error
	for _, scheme := range schemes {
		url := fmt.Sprintf("%s://%s/v2/", scheme, reg.Name())
		resp, err := client.Get(url)
		if err != nil {
			connErr = err
			// Potentially retry with http.
			continue
		}
		defer func() {
			// By draining the body, make sure to reuse the connection made by
			// the ping for the following access to the registry
			io.Copy(ioutil.Discard, resp.Body)
			resp.Body.Close()
		}()

		switch resp.StatusCode {
		case http.StatusOK:
			// If we get a 200, then no authentication is needed.
			return &pingResp{
				challenge: anonymous,
				scheme:    scheme,
			}, nil
		case http.StatusUnauthorized:
			wac := resp.Header.Get("WWW-Authenticate")
			if parts := strings.SplitN(wac, " ", 2); len(parts) == 2 {
				// If there are two parts, then parse the challenge parameters.
				return &pingResp{
					challenge:  challenge(parts[0]).Canonical(),
					parameters: parseChallenge(parts[1]),
					scheme:     scheme,
				}, nil
			}
			// Otherwise, just return the challenge without parameters.
			return &pingResp{
				challenge: challenge(wac).Canonical(),
				scheme:    scheme,
			}, nil
		default:
			return nil, CheckError(resp, http.StatusOK, http.StatusUnauthorized)
		}
	}
	return nil, connErr
}
<commit_msg>Set longer timeout for ping (#833)<commit_after>// Copyright 2018 Google LLC All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package transport

import (
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"github.com/google/go-containerregistry/pkg/name"
)

type challenge string

const (
	anonymous challenge = "anonymous"
	basic     challenge = "basic"
	bearer    challenge = "bearer"
)

type pingResp struct {
	challenge challenge

	// Following the challenge there are often key/value pairs
	// e.g. Bearer service="gcr.io",realm="https://auth.gcr.io/v36/tokenz"
	parameters map[string]string

	// The registry's scheme to use. Communicates whether we fell back to http.
	scheme string
}

func (c challenge) Canonical() challenge {
	return challenge(strings.ToLower(string(c)))
}

func parseChallenge(suffix string) map[string]string {
	kv := make(map[string]string)
	for _, token := range strings.Split(suffix, ",") {
		// Trim any whitespace around each token.
		token = strings.Trim(token, " ")

		// Break the token into a key/value pair
		if parts := strings.SplitN(token, "=", 2); len(parts) == 2 {
			// Unquote the value, if it is quoted.
			kv[parts[0]] = strings.Trim(parts[1], `"`)
		} else {
			// If there was only one part, treat is as a key with an empty value
			kv[token] = ""
		}
	}
	return kv
}

func ping(reg name.Registry, t http.RoundTripper) (*pingResp, error) {
	client := http.Client{Transport: t, Timeout: 120 * time.Second}

	// This first attempts to use "https" for every request, falling back to http
	// if the registry matches our localhost heuristic or if it is intentionally
	// set to insecure via name.NewInsecureRegistry.
	schemes := []string{"https"}
	if reg.Scheme() == "http" {
		schemes = append(schemes, "http")
	}

	var connErr error
	for _, scheme := range schemes {
		url := fmt.Sprintf("%s://%s/v2/", scheme, reg.Name())
		resp, err := client.Get(url)
		if err != nil {
			connErr = err
			// Potentially retry with http.
			continue
		}
		defer func() {
			// By draining the body, make sure to reuse the connection made by
			// the ping for the following access to the registry
			io.Copy(ioutil.Discard, resp.Body)
			resp.Body.Close()
		}()

		switch resp.StatusCode {
		case http.StatusOK:
			// If we get a 200, then no authentication is needed.
			return &pingResp{
				challenge: anonymous,
				scheme:    scheme,
			}, nil
		case http.StatusUnauthorized:
			wac := resp.Header.Get("WWW-Authenticate")
			if parts := strings.SplitN(wac, " ", 2); len(parts) == 2 {
				// If there are two parts, then parse the challenge parameters.
				return &pingResp{
					challenge:  challenge(parts[0]).Canonical(),
					parameters: parseChallenge(parts[1]),
					scheme:     scheme,
				}, nil
			}
			// Otherwise, just return the challenge without parameters.
			return &pingResp{
				challenge: challenge(wac).Canonical(),
				scheme:    scheme,
			}, nil
		default:
			return nil, CheckError(resp, http.StatusOK, http.StatusUnauthorized)
		}
	}
	return nil, connErr
}
<|endoftext|>