<commit_before>package rbd

// #cgo LDFLAGS: -lrbd
// #include <stdlib.h>
// #include <rbd/librbd.h>
import "C"

import (
	"fmt"
	"unsafe"
)

const (
	// RBD image options.

	// ImageOptionFormat is the representation of RBD_IMAGE_OPTION_FORMAT from
	// librbd
	ImageOptionFormat = C.RBD_IMAGE_OPTION_FORMAT
	// ImageOptionFeatures is the representation of RBD_IMAGE_OPTION_FEATURES
	// from librbd
	ImageOptionFeatures = C.RBD_IMAGE_OPTION_FEATURES
	// ImageOptionOrder is the representation of RBD_IMAGE_OPTION_ORDER from
	// librbd
	ImageOptionOrder = C.RBD_IMAGE_OPTION_ORDER
	// ImageOptionStripeUnit is the representation of
	// RBD_IMAGE_OPTION_STRIPE_UNIT from librbd
	ImageOptionStripeUnit = C.RBD_IMAGE_OPTION_STRIPE_UNIT
	// ImageOptionStripeCount is the representation of
	// RBD_IMAGE_OPTION_STRIPE_COUNT from librbd
	ImageOptionStripeCount = C.RBD_IMAGE_OPTION_STRIPE_COUNT
	// ImageOptionJournalOrder is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_ORDER from librbd
	ImageOptionJournalOrder = C.RBD_IMAGE_OPTION_JOURNAL_ORDER
	// ImageOptionJournalSplayWidth is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH from librbd
	ImageOptionJournalSplayWidth = C.RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH
	// ImageOptionJournalPool is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_POOL from librbd
	ImageOptionJournalPool = C.RBD_IMAGE_OPTION_JOURNAL_POOL
	// ImageOptionFeaturesSet is the representation of
	// RBD_IMAGE_OPTION_FEATURES_SET from librbd
	ImageOptionFeaturesSet = C.RBD_IMAGE_OPTION_FEATURES_SET
	// ImageOptionFeaturesClear is the representation of
	// RBD_IMAGE_OPTION_FEATURES_CLEAR from librbd
	ImageOptionFeaturesClear = C.RBD_IMAGE_OPTION_FEATURES_CLEAR
	// ImageOptionDataPool is the representation of RBD_IMAGE_OPTION_DATA_POOL
	// from librbd
	ImageOptionDataPool = C.RBD_IMAGE_OPTION_DATA_POOL
	// ImageOptionFlatten is the representation of RBD_IMAGE_OPTION_FLATTEN
	// from librbd
	ImageOptionFlatten = C.RBD_IMAGE_OPTION_FLATTEN
	// ImageOptionCloneFormat is the representation of
	// RBD_IMAGE_OPTION_CLONE_FORMAT from librbd
	ImageOptionCloneFormat = C.RBD_IMAGE_OPTION_CLONE_FORMAT

	// RbdImageOptionFormat deprecated alias for ImageOptionFormat
	RbdImageOptionFormat = ImageOptionFormat
	// RbdImageOptionFeatures deprecated alias for ImageOptionFeatures
	RbdImageOptionFeatures = ImageOptionFeatures
	// RbdImageOptionOrder deprecated alias for ImageOptionOrder
	RbdImageOptionOrder = ImageOptionOrder
	// RbdImageOptionStripeUnit deprecated alias for ImageOptionStripeUnit
	RbdImageOptionStripeUnit = ImageOptionStripeUnit
	// RbdImageOptionStripeCount deprecated alias for ImageOptionStripeCount
	RbdImageOptionStripeCount = ImageOptionStripeCount
	// RbdImageOptionJournalOrder deprecated alias for ImageOptionJournalOrder
	RbdImageOptionJournalOrder = ImageOptionJournalOrder
	// RbdImageOptionJournalSplayWidth deprecated alias for
	RbdImageOptionJournalSplayWidth = ImageOptionJournalSplayWidth
	// RbdImageOptionJournalPool deprecated alias for ImageOptionJournalPool
	RbdImageOptionJournalPool = ImageOptionJournalPool
	// RbdImageOptionFeaturesSet deprecated alias for ImageOptionFeaturesSet
	RbdImageOptionFeaturesSet = ImageOptionFeaturesSet
	// RbdImageOptionFeaturesClear deprecated alias for ImageOptionFeaturesClear
	RbdImageOptionFeaturesClear = ImageOptionFeaturesClear
	// RbdImageOptionDataPool deprecated alias for ImageOptionDataPool
	RbdImageOptionDataPool = ImageOptionDataPool
)

// ImageOptions represents a group of configurable image options.
type ImageOptions struct {
	options C.rbd_image_options_t
}

// ImageOption values are unique keys for configurable options.
type ImageOption C.int

// revive:disable:exported Deprecated aliases

// RbdImageOptions deprecated alias for ImageOptions
type RbdImageOptions = ImageOptions

// RbdImageOption is a deprecated alias for ImageOption
type RbdImageOption = ImageOption

//revive:enable:exported

// NewRbdImageOptions creates a new RbdImageOptions struct. Call
// RbdImageOptions.Destroy() to free the resources.
//
// Implements:
//  void rbd_image_options_create(rbd_image_options_t* opts)
func NewRbdImageOptions() *ImageOptions {
	rio := &ImageOptions{}
	C.rbd_image_options_create(&rio.options)
	return rio
}

// Destroy a RbdImageOptions struct and free the associated resources.
//
// Implements:
//  void rbd_image_options_destroy(rbd_image_options_t opts);
func (rio *ImageOptions) Destroy() {
	C.rbd_image_options_destroy(rio.options)
}

// SetString sets the value of the RbdImageOption to the given string.
//
// Implements:
//  int rbd_image_options_set_string(rbd_image_options_t opts, int optname,
//          const char* optval);
func (rio *ImageOptions) SetString(option ImageOption, value string) error {
	c_value := C.CString(value)
	defer C.free(unsafe.Pointer(c_value))

	ret := C.rbd_image_options_set_string(rio.options, C.int(option), c_value)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetString returns the string value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_string(rbd_image_options_t opts, int optname,
//          char* optval, size_t maxlen);
func (rio *ImageOptions) GetString(option ImageOption) (string, error) {
	value := make([]byte, 4096)

	ret := C.rbd_image_options_get_string(rio.options, C.int(option),
		(*C.char)(unsafe.Pointer(&value[0])),
		C.size_t(len(value)))
	if ret != 0 {
		return "", fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return C.GoString((*C.char)(unsafe.Pointer(&value[0]))), nil
}

// SetUint64 sets the value of the RbdImageOption to the given uint64.
//
// Implements:
//  int rbd_image_options_set_uint64(rbd_image_options_t opts, int optname,
//          const uint64_t optval);
func (rio *ImageOptions) SetUint64(option ImageOption, value uint64) error {
	c_value := C.uint64_t(value)

	ret := C.rbd_image_options_set_uint64(rio.options, C.int(option), c_value)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetUint64 returns the uint64 value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_uint64(rbd_image_options_t opts, int optname,
//          uint64_t* optval);
func (rio *ImageOptions) GetUint64(option ImageOption) (uint64, error) {
	var c_value C.uint64_t

	ret := C.rbd_image_options_get_uint64(rio.options, C.int(option), &c_value)
	if ret != 0 {
		return 0, fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return uint64(c_value), nil
}

// IsSet returns a true if the RbdImageOption is set, false otherwise.
//
// Implements:
//  int rbd_image_options_is_set(rbd_image_options_t opts, int optname,
//          bool* is_set);
func (rio *ImageOptions) IsSet(option ImageOption) (bool, error) {
	var c_set C.bool

	ret := C.rbd_image_options_is_set(rio.options, C.int(option), &c_set)
	if ret != 0 {
		return false, fmt.Errorf("%v, could not check option %v", getError(ret), option)
	}

	return bool(c_set), nil
}

// Unset a given RbdImageOption.
//
// Implements:
//  int rbd_image_options_unset(rbd_image_options_t opts, int optname)
func (rio *ImageOptions) Unset(option ImageOption) error {
	ret := C.rbd_image_options_unset(rio.options, C.int(option))
	if ret != 0 {
		return fmt.Errorf("%v, could not unset option %v", getError(ret), option)
	}

	return nil
}

// Clear all options in the RbdImageOptions.
//
// Implements:
//  void rbd_image_options_clear(rbd_image_options_t opts)
func (rio *ImageOptions) Clear() {
	C.rbd_image_options_clear(rio.options)
}

// IsEmpty returns true if there are no options set in the RbdImageOptions,
// false otherwise.
//
// Implements:
//  int rbd_image_options_is_empty(rbd_image_options_t opts)
func (rio *ImageOptions) IsEmpty() bool {
	ret := C.rbd_image_options_is_empty(rio.options)
	return ret != 0
}
<commit_msg>rbd: fix naming convention violations in options.go<commit_after>package rbd

// #cgo LDFLAGS: -lrbd
// #include <stdlib.h>
// #include <rbd/librbd.h>
import "C"

import (
	"fmt"
	"unsafe"
)

const (
	// RBD image options.

	// ImageOptionFormat is the representation of RBD_IMAGE_OPTION_FORMAT from
	// librbd
	ImageOptionFormat = C.RBD_IMAGE_OPTION_FORMAT
	// ImageOptionFeatures is the representation of RBD_IMAGE_OPTION_FEATURES
	// from librbd
	ImageOptionFeatures = C.RBD_IMAGE_OPTION_FEATURES
	// ImageOptionOrder is the representation of RBD_IMAGE_OPTION_ORDER from
	// librbd
	ImageOptionOrder = C.RBD_IMAGE_OPTION_ORDER
	// ImageOptionStripeUnit is the representation of
	// RBD_IMAGE_OPTION_STRIPE_UNIT from librbd
	ImageOptionStripeUnit = C.RBD_IMAGE_OPTION_STRIPE_UNIT
	// ImageOptionStripeCount is the representation of
	// RBD_IMAGE_OPTION_STRIPE_COUNT from librbd
	ImageOptionStripeCount = C.RBD_IMAGE_OPTION_STRIPE_COUNT
	// ImageOptionJournalOrder is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_ORDER from librbd
	ImageOptionJournalOrder = C.RBD_IMAGE_OPTION_JOURNAL_ORDER
	// ImageOptionJournalSplayWidth is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH from librbd
	ImageOptionJournalSplayWidth = C.RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH
	// ImageOptionJournalPool is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_POOL from librbd
	ImageOptionJournalPool = C.RBD_IMAGE_OPTION_JOURNAL_POOL
	// ImageOptionFeaturesSet is the representation of
	// RBD_IMAGE_OPTION_FEATURES_SET from librbd
	ImageOptionFeaturesSet = C.RBD_IMAGE_OPTION_FEATURES_SET
	// ImageOptionFeaturesClear is the representation of
	// RBD_IMAGE_OPTION_FEATURES_CLEAR from librbd
	ImageOptionFeaturesClear = C.RBD_IMAGE_OPTION_FEATURES_CLEAR
	// ImageOptionDataPool is the representation of RBD_IMAGE_OPTION_DATA_POOL
	// from librbd
	ImageOptionDataPool = C.RBD_IMAGE_OPTION_DATA_POOL
	// ImageOptionFlatten is the representation of RBD_IMAGE_OPTION_FLATTEN
	// from librbd
	ImageOptionFlatten = C.RBD_IMAGE_OPTION_FLATTEN
	// ImageOptionCloneFormat is the representation of
	// RBD_IMAGE_OPTION_CLONE_FORMAT from librbd
	ImageOptionCloneFormat = C.RBD_IMAGE_OPTION_CLONE_FORMAT

	// RbdImageOptionFormat deprecated alias for ImageOptionFormat
	RbdImageOptionFormat = ImageOptionFormat
	// RbdImageOptionFeatures deprecated alias for ImageOptionFeatures
	RbdImageOptionFeatures = ImageOptionFeatures
	// RbdImageOptionOrder deprecated alias for ImageOptionOrder
	RbdImageOptionOrder = ImageOptionOrder
	// RbdImageOptionStripeUnit deprecated alias for ImageOptionStripeUnit
	RbdImageOptionStripeUnit = ImageOptionStripeUnit
	// RbdImageOptionStripeCount deprecated alias for ImageOptionStripeCount
	RbdImageOptionStripeCount = ImageOptionStripeCount
	// RbdImageOptionJournalOrder deprecated alias for ImageOptionJournalOrder
	RbdImageOptionJournalOrder = ImageOptionJournalOrder
	// RbdImageOptionJournalSplayWidth deprecated alias for
	RbdImageOptionJournalSplayWidth = ImageOptionJournalSplayWidth
	// RbdImageOptionJournalPool deprecated alias for ImageOptionJournalPool
	RbdImageOptionJournalPool = ImageOptionJournalPool
	// RbdImageOptionFeaturesSet deprecated alias for ImageOptionFeaturesSet
	RbdImageOptionFeaturesSet = ImageOptionFeaturesSet
	// RbdImageOptionFeaturesClear deprecated alias for ImageOptionFeaturesClear
	RbdImageOptionFeaturesClear = ImageOptionFeaturesClear
	// RbdImageOptionDataPool deprecated alias for ImageOptionDataPool
	RbdImageOptionDataPool = ImageOptionDataPool
)

// ImageOptions represents a group of configurable image options.
type ImageOptions struct {
	options C.rbd_image_options_t
}

// ImageOption values are unique keys for configurable options.
type ImageOption C.int

// revive:disable:exported Deprecated aliases

// RbdImageOptions deprecated alias for ImageOptions
type RbdImageOptions = ImageOptions

// RbdImageOption is a deprecated alias for ImageOption
type RbdImageOption = ImageOption

//revive:enable:exported

// NewRbdImageOptions creates a new RbdImageOptions struct. Call
// RbdImageOptions.Destroy() to free the resources.
//
// Implements:
//  void rbd_image_options_create(rbd_image_options_t* opts)
func NewRbdImageOptions() *ImageOptions {
	rio := &ImageOptions{}
	C.rbd_image_options_create(&rio.options)
	return rio
}

// Destroy a RbdImageOptions struct and free the associated resources.
//
// Implements:
//  void rbd_image_options_destroy(rbd_image_options_t opts);
func (rio *ImageOptions) Destroy() {
	C.rbd_image_options_destroy(rio.options)
}

// SetString sets the value of the RbdImageOption to the given string.
//
// Implements:
//  int rbd_image_options_set_string(rbd_image_options_t opts, int optname,
//          const char* optval);
func (rio *ImageOptions) SetString(option ImageOption, value string) error {
	cValue := C.CString(value)
	defer C.free(unsafe.Pointer(cValue))

	ret := C.rbd_image_options_set_string(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetString returns the string value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_string(rbd_image_options_t opts, int optname,
//          char* optval, size_t maxlen);
func (rio *ImageOptions) GetString(option ImageOption) (string, error) {
	value := make([]byte, 4096)

	ret := C.rbd_image_options_get_string(rio.options, C.int(option),
		(*C.char)(unsafe.Pointer(&value[0])),
		C.size_t(len(value)))
	if ret != 0 {
		return "", fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return C.GoString((*C.char)(unsafe.Pointer(&value[0]))), nil
}

// SetUint64 sets the value of the RbdImageOption to the given uint64.
//
// Implements:
//  int rbd_image_options_set_uint64(rbd_image_options_t opts, int optname,
//          const uint64_t optval);
func (rio *ImageOptions) SetUint64(option ImageOption, value uint64) error {
	cValue := C.uint64_t(value)

	ret := C.rbd_image_options_set_uint64(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetUint64 returns the uint64 value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_uint64(rbd_image_options_t opts, int optname,
//          uint64_t* optval);
func (rio *ImageOptions) GetUint64(option ImageOption) (uint64, error) {
	var cValue C.uint64_t

	ret := C.rbd_image_options_get_uint64(rio.options, C.int(option), &cValue)
	if ret != 0 {
		return 0, fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return uint64(cValue), nil
}

// IsSet returns a true if the RbdImageOption is set, false otherwise.
//
// Implements:
//  int rbd_image_options_is_set(rbd_image_options_t opts, int optname,
//          bool* is_set);
func (rio *ImageOptions) IsSet(option ImageOption) (bool, error) {
	var c_set C.bool

	ret := C.rbd_image_options_is_set(rio.options, C.int(option), &c_set)
	if ret != 0 {
		return false, fmt.Errorf("%v, could not check option %v", getError(ret), option)
	}

	return bool(c_set), nil
}

// Unset a given RbdImageOption.
//
// Implements:
//  int rbd_image_options_unset(rbd_image_options_t opts, int optname)
func (rio *ImageOptions) Unset(option ImageOption) error {
	ret := C.rbd_image_options_unset(rio.options, C.int(option))
	if ret != 0 {
		return fmt.Errorf("%v, could not unset option %v", getError(ret), option)
	}

	return nil
}

// Clear all options in the RbdImageOptions.
//
// Implements:
//  void rbd_image_options_clear(rbd_image_options_t opts)
func (rio *ImageOptions) Clear() {
	C.rbd_image_options_clear(rio.options)
}

// IsEmpty returns true if there are no options set in the RbdImageOptions,
// false otherwise.
//
// Implements:
//  int rbd_image_options_is_empty(rbd_image_options_t opts)
func (rio *ImageOptions) IsEmpty() bool {
	ret := C.rbd_image_options_is_empty(rio.options)
	return ret != 0
}
<|endoftext|>


<commit_before>package utils

import (
	"bytes"
	"github.com/cloudflare/goflow/decoders/netflowlegacy"
	flowmessage "github.com/cloudflare/goflow/pb"
	"github.com/cloudflare/goflow/producer"
	"github.com/prometheus/client_golang/prometheus"
	"time"
)

type StateNFLegacy struct {
	Transport Transport
	Logger    Logger
}

func (s *StateNFLegacy) DecodeFlow(msg interface{}) error {
	pkt := msg.(BaseMessage)
	buf := bytes.NewBuffer(pkt.Payload)
	key := pkt.Src.String()
	samplerAddress := pkt.Src
	if samplerAddress.To4() != nil {
		samplerAddress = samplerAddress.To4()
	}

	timeTrackStart := time.Now()
	msgDec, err := netflowlegacy.DecodeMessage(buf)

	if err != nil {
		switch err.(type) {
		case *netflowlegacy.ErrorVersion:
			NetFlowErrors.With(
				prometheus.Labels{
					"router": key,
					"error":  "error_version",
				}).
				Inc()
		}
		return err
	}

	switch msgDec.(type) {
	case netflowlegacy.PacketNetFlowV5:
		NetFlowStats.With(
			prometheus.Labels{
				"router":  key,
				"version": "5",
			}).
			Inc()
	}

	var flowMessageSet []*flowmessage.FlowMessage
	flowMessageSet, err = producer.ProcessMessageNetFlowLegacy(msgDec)

	timeTrackStop := time.Now()
	DecoderTime.With(
		prometheus.Labels{
			"name": "NetFlowV5",
		}).
		Observe(float64((timeTrackStop.Sub(timeTrackStart)).Nanoseconds()) / 1000)

	for _, fmsg := range flowMessageSet {
		fmsg.TimeReceived = uint64(time.Now().UTC().Unix())
		fmsg.SamplerAddress = samplerAddress
	}

	if s.Transport != nil {
		s.Transport.Publish(flowMessageSet)
	}

	return nil
}

func (s *StateNFLegacy) FlowRoutine(workers int, addr string, port int) error {
	return UDPRoutine("NetFlowV5", s.DecodeFlow, workers, addr, port, false, s.Logger)
}
<commit_msg>add netflow v5 flows counter for prometheus stats<commit_after>package utils

import (
	"bytes"
	"github.com/cloudflare/goflow/decoders/netflowlegacy"
	flowmessage "github.com/cloudflare/goflow/pb"
	"github.com/cloudflare/goflow/producer"
	"github.com/prometheus/client_golang/prometheus"
	"time"
)

type StateNFLegacy struct {
	Transport Transport
	Logger    Logger
}

func (s *StateNFLegacy) DecodeFlow(msg interface{}) error {
	pkt := msg.(BaseMessage)
	buf := bytes.NewBuffer(pkt.Payload)
	key := pkt.Src.String()
	samplerAddress := pkt.Src
	if samplerAddress.To4() != nil {
		samplerAddress = samplerAddress.To4()
	}

	timeTrackStart := time.Now()
	msgDec, err := netflowlegacy.DecodeMessage(buf)

	if err != nil {
		switch err.(type) {
		case *netflowlegacy.ErrorVersion:
			NetFlowErrors.With(
				prometheus.Labels{
					"router": key,
					"error":  "error_version",
				}).
				Inc()
		}
		return err
	}

	switch msgDecConv := msgDec.(type) {
	case netflowlegacy.PacketNetFlowV5:
		NetFlowStats.With(
			prometheus.Labels{
				"router":  key,
				"version": "5",
			}).
			Inc()
		NetFlowSetStatsSum.With(
			prometheus.Labels{
				"router":  key,
				"version": "5",
				"type":    "DataFlowSet",
			}).
			Add(float64(msgDecConv.Count))
	}

	var flowMessageSet []*flowmessage.FlowMessage
	flowMessageSet, err = producer.ProcessMessageNetFlowLegacy(msgDec)

	timeTrackStop := time.Now()
	DecoderTime.With(
		prometheus.Labels{
			"name": "NetFlowV5",
		}).
		Observe(float64((timeTrackStop.Sub(timeTrackStart)).Nanoseconds()) / 1000)

	for _, fmsg := range flowMessageSet {
		fmsg.TimeReceived = uint64(time.Now().UTC().Unix())
		fmsg.SamplerAddress = samplerAddress
	}

	if s.Transport != nil {
		s.Transport.Publish(flowMessageSet)
	}

	return nil
}

func (s *StateNFLegacy) FlowRoutine(workers int, addr string, port int) error {
	return UDPRoutine("NetFlowV5", s.DecodeFlow, workers, addr, port, false, s.Logger)
}
<|endoftext|>


<commit_before>package geoip2

import (
	"math/rand"
	"net"
	"testing"

	. "gopkg.in/check.v1"
)

func TestGeoIP2(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.City(net.ParseIP("81.2.69.160"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	m := reader.Metadata()
	c.Assert(m.BinaryFormatMajorVersion, Equals, uint(2))
	c.Assert(m.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(m.BuildEpoch, Equals, uint(1461706823))
	c.Assert(m.DatabaseType, Equals, "GeoIP2-City")
	c.Assert(m.Description, DeepEquals, map[string]string{
		"en": "GeoIP2 City Test Database (fake GeoIP2 data, for example purposes only)",
		"zh": "小型数据库",
	})
	c.Assert(m.IPVersion, Equals, uint(6))
	c.Assert(m.Languages, DeepEquals, []string{"en", "zh"})
	c.Assert(m.NodeCount, Equals, uint(1431))
	c.Assert(m.RecordSize, Equals, uint(28))

	c.Assert(record.City.GeoNameID, Equals, uint(2643743))
	c.Assert(record.City.Names, DeepEquals, map[string]string{
		"de":    "London",
		"en":    "London",
		"es":    "Londres",
		"fr":    "Londres",
		"ja":    "ロンドン",
		"pt-BR": "Londres",
		"ru":    "Лондон",
	})
	c.Assert(record.Continent.GeoNameID, Equals, uint(6255148))
	c.Assert(record.Continent.Code, Equals, "EU")
	c.Assert(record.Continent.Names, DeepEquals, map[string]string{
		"de":    "Europa",
		"en":    "Europe",
		"es":    "Europa",
		"fr":    "Europe",
		"ja":    "ヨーロッパ",
		"pt-BR": "Europa",
		"ru":    "Европа",
		"zh-CN": "欧洲",
	})

	c.Assert(record.Country.GeoNameID, Equals, uint(2635167))
	c.Assert(record.Country.IsoCode, Equals, "GB")
	c.Assert(record.Country.Names, DeepEquals, map[string]string{
		"de":    "Vereinigtes Königreich",
		"en":    "United Kingdom",
		"es":    "Reino Unido",
		"fr":    "Royaume-Uni",
		"ja":    "イギリス",
		"pt-BR": "Reino Unido",
		"ru":    "Великобритания",
		"zh-CN": "英国",
	})

	c.Assert(record.Location.AccuracyRadius, Equals, uint16(100))
	c.Assert(record.Location.Latitude, Equals, 51.5142)
	c.Assert(record.Location.Longitude, Equals, -0.0931)
	c.Assert(record.Location.TimeZone, Equals, "Europe/London")

	c.Assert(record.Subdivisions[0].GeoNameID, Equals, uint(6269131))
	c.Assert(record.Subdivisions[0].IsoCode, Equals, "ENG")
	c.Assert(record.Subdivisions[0].Names, DeepEquals, map[string]string{
		"en":    "England",
		"pt-BR": "Inglaterra",
		"fr":    "Angleterre",
		"es":    "Inglaterra",
	})

	c.Assert(record.RegisteredCountry.GeoNameID, Equals, uint(6252001))
	c.Assert(record.RegisteredCountry.IsoCode, Equals, "US")
	c.Assert(record.RegisteredCountry.Names, DeepEquals, map[string]string{
		"de":    "USA",
		"en":    "United States",
		"es":    "Estados Unidos",
		"fr":    "États-Unis",
		"ja":    "アメリカ合衆国",
		"pt-BR": "Estados Unidos",
		"ru":    "США",
		"zh-CN": "美国",
	})
}

func (s *MySuite) TestMetroCode(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.City(net.ParseIP("216.160.83.56"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(record.Location.MetroCode, Equals, uint(819))
}

func (s *MySuite) TestConnectionType(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Connection-Type-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.ConnectionType(net.ParseIP("1.0.1.0"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	c.Assert(record.ConnectionType, Equals, "Cable/DSL")

}

func (s *MySuite) TestDomain(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Domain-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.Domain(net.ParseIP("1.2.0.0"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	c.Assert(record.Domain, Equals, "maxmind.com")

}

func (s *MySuite) TestISP(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-ISP-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.ISP(net.ParseIP("1.128.0.0"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(record.AutonomousSystemNumber, Equals, uint(1221))

	c.Assert(record.AutonomousSystemOrganization, Equals, "Telstra Pty Ltd")
	c.Assert(record.ISP, Equals, "Telstra Internet")
	c.Assert(record.Organization, Equals, "Telstra Internet")

}

func (s *MySuite) TestAnonymousIP(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Anonymous-IP-Test.mmdb")
	if err != nil {
		c.Log(err)
		c.Fail()
	}
	defer reader.Close()

	record, err := reader.AnonymousIP(net.ParseIP("1.2.0.0"))
	if err != nil {
		c.Log(err)
		c.Fail()
	}

	c.Assert(record.IsAnonymous, Equals, true)

	c.Assert(record.IsAnonymousVPN, Equals, true)
	c.Assert(record.IsHostingProvider, Equals, false)
	c.Assert(record.IsPublicProxy, Equals, false)
	c.Assert(record.IsTorExitNode, Equals, false)

}

// This ensures the compiler does not optimize away the function call
var cityResult *City

func BenchmarkMaxMindDB(b *testing.B) {
	db, err := Open("GeoLite2-City.mmdb")
	if err != nil {
		b.Fatal(err)
	}
	defer db.Close()

	r := rand.New(rand.NewSource(0))

	var city *City

	for i := 0; i < b.N; i++ {
		ip := randomIPv4Address(b, r)
		city, err = db.City(ip)
		if err != nil {
			b.Fatal(err)
		}
	}
	cityResult = city
}

func randomIPv4Address(b *testing.B, r *rand.Rand) net.IP {
	num := r.Uint32()
	return []byte{byte(num >> 24), byte(num >> 16), byte(num >> 8),
		byte(num)}
}
<commit_msg>Clean up test error handling<commit_after>package geoip2

import (
	"math/rand"
	"net"
	"testing"

	. "gopkg.in/check.v1"
)

func TestGeoIP2(t *testing.T) { TestingT(t) }

type MySuite struct{}

var _ = Suite(&MySuite{})

func (s *MySuite) TestReader(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	c.Assert(err, IsNil)

	defer reader.Close()

	record, err := reader.City(net.ParseIP("81.2.69.160"))
	c.Assert(err, IsNil)

	m := reader.Metadata()
	c.Assert(m.BinaryFormatMajorVersion, Equals, uint(2))
	c.Assert(m.BinaryFormatMinorVersion, Equals, uint(0))
	c.Assert(m.BuildEpoch, Equals, uint(1461706823))
	c.Assert(m.DatabaseType, Equals, "GeoIP2-City")
	c.Assert(m.Description, DeepEquals, map[string]string{
		"en": "GeoIP2 City Test Database (fake GeoIP2 data, for example purposes only)",
		"zh": "小型数据库",
	})
	c.Assert(m.IPVersion, Equals, uint(6))
	c.Assert(m.Languages, DeepEquals, []string{"en", "zh"})
	c.Assert(m.NodeCount, Equals, uint(1431))
	c.Assert(m.RecordSize, Equals, uint(28))

	c.Assert(record.City.GeoNameID, Equals, uint(2643743))
	c.Assert(record.City.Names, DeepEquals, map[string]string{
		"de":    "London",
		"en":    "London",
		"es":    "Londres",
		"fr":    "Londres",
		"ja":    "ロンドン",
		"pt-BR": "Londres",
		"ru":    "Лондон",
	})
	c.Assert(record.Continent.GeoNameID, Equals, uint(6255148))
	c.Assert(record.Continent.Code, Equals, "EU")
	c.Assert(record.Continent.Names, DeepEquals, map[string]string{
		"de":    "Europa",
		"en":    "Europe",
		"es":    "Europa",
		"fr":    "Europe",
		"ja":    "ヨーロッパ",
		"pt-BR": "Europa",
		"ru":    "Европа",
		"zh-CN": "欧洲",
	})

	c.Assert(record.Country.GeoNameID, Equals, uint(2635167))
	c.Assert(record.Country.IsoCode, Equals, "GB")
	c.Assert(record.Country.Names, DeepEquals, map[string]string{
		"de":    "Vereinigtes Königreich",
		"en":    "United Kingdom",
		"es":    "Reino Unido",
		"fr":    "Royaume-Uni",
		"ja":    "イギリス",
		"pt-BR": "Reino Unido",
		"ru":    "Великобритания",
		"zh-CN": "英国",
	})

	c.Assert(record.Location.AccuracyRadius, Equals, uint16(100))
	c.Assert(record.Location.Latitude, Equals, 51.5142)
	c.Assert(record.Location.Longitude, Equals, -0.0931)
	c.Assert(record.Location.TimeZone, Equals, "Europe/London")

	c.Assert(record.Subdivisions[0].GeoNameID, Equals, uint(6269131))
	c.Assert(record.Subdivisions[0].IsoCode, Equals, "ENG")
	c.Assert(record.Subdivisions[0].Names, DeepEquals, map[string]string{
		"en":    "England",
		"pt-BR": "Inglaterra",
		"fr":    "Angleterre",
		"es":    "Inglaterra",
	})

	c.Assert(record.RegisteredCountry.GeoNameID, Equals, uint(6252001))
	c.Assert(record.RegisteredCountry.IsoCode, Equals, "US")
	c.Assert(record.RegisteredCountry.Names, DeepEquals, map[string]string{
		"de":    "USA",
		"en":    "United States",
		"es":    "Estados Unidos",
		"fr":    "États-Unis",
		"ja":    "アメリカ合衆国",
		"pt-BR": "Estados Unidos",
		"ru":    "США",
		"zh-CN": "美国",
	})
}

func (s *MySuite) TestMetroCode(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-City-Test.mmdb")
	c.Assert(err, IsNil)
	defer reader.Close()

	record, err := reader.City(net.ParseIP("216.160.83.56"))
	c.Assert(err, IsNil)

	c.Assert(record.Location.MetroCode, Equals, uint(819))
}

func (s *MySuite) TestConnectionType(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Connection-Type-Test.mmdb")
	c.Assert(err, IsNil)

	defer reader.Close()

	record, err := reader.ConnectionType(net.ParseIP("1.0.1.0"))
	c.Assert(err, IsNil)

	c.Assert(record.ConnectionType, Equals, "Cable/DSL")

}

func (s *MySuite) TestDomain(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Domain-Test.mmdb")
	c.Assert(err, IsNil)
	defer reader.Close()

	record, err := reader.Domain(net.ParseIP("1.2.0.0"))
	c.Assert(err, IsNil)
	c.Assert(record.Domain, Equals, "maxmind.com")

}

func (s *MySuite) TestISP(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-ISP-Test.mmdb")
	c.Assert(err, IsNil)
	defer reader.Close()

	record, err := reader.ISP(net.ParseIP("1.128.0.0"))
	c.Assert(err, IsNil)

	c.Assert(record.AutonomousSystemNumber, Equals, uint(1221))

	c.Assert(record.AutonomousSystemOrganization, Equals, "Telstra Pty Ltd")
	c.Assert(record.ISP, Equals, "Telstra Internet")
	c.Assert(record.Organization, Equals, "Telstra Internet")

}

func (s *MySuite) TestAnonymousIP(c *C) {
	reader, err := Open("test-data/test-data/GeoIP2-Anonymous-IP-Test.mmdb")
	c.Assert(err, IsNil)
	defer reader.Close()

	record, err := reader.AnonymousIP(net.ParseIP("1.2.0.0"))
	c.Assert(err, IsNil)

	c.Assert(record.IsAnonymous, Equals, true)

	c.Assert(record.IsAnonymousVPN, Equals, true)
	c.Assert(record.IsHostingProvider, Equals, false)
	c.Assert(record.IsPublicProxy, Equals, false)
	c.Assert(record.IsTorExitNode, Equals, false)

}

// This ensures the compiler does not optimize away the function call
var cityResult *City

func BenchmarkMaxMindDB(b *testing.B) {
	db, err := Open("GeoLite2-City.mmdb")
	if err != nil {
		b.Fatal(err)
	}
	defer db.Close()

	r := rand.New(rand.NewSource(0))

	var city *City

	for i := 0; i < b.N; i++ {
		ip := randomIPv4Address(b, r)
		city, err = db.City(ip)
		if err != nil {
			b.Fatal(err)
		}
	}
	cityResult = city
}

func randomIPv4Address(b *testing.B, r *rand.Rand) net.IP {
	num := r.Uint32()
	return []byte{byte(num >> 24), byte(num >> 16), byte(num >> 8),
		byte(num)}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 versioning

import (
	"io"
	"reflect"

	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

// NewDefaultingCodecForScheme is a convenience method for callers that are using a scheme.
func NewDefaultingCodecForScheme(
	// TODO: I should be a scheme interface?
	scheme *runtime.Scheme,
	encoder runtime.Encoder,
	decoder runtime.Decoder,
	encodeVersion runtime.GroupVersioner,
	decodeVersion runtime.GroupVersioner,
) runtime.Codec {
	return NewCodec(encoder, decoder, runtime.UnsafeObjectConvertor(scheme), scheme, scheme, scheme, encodeVersion, decodeVersion, scheme.Name())
}

// NewCodec takes objects in their internal versions and converts them to external versions before
// serializing them. It assumes the serializer provided to it only deals with external versions.
// This class is also a serializer, but is generally used with a specific version.
func NewCodec(
	encoder runtime.Encoder,
	decoder runtime.Decoder,
	convertor runtime.ObjectConvertor,
	creater runtime.ObjectCreater,
	typer runtime.ObjectTyper,
	defaulter runtime.ObjectDefaulter,
	encodeVersion runtime.GroupVersioner,
	decodeVersion runtime.GroupVersioner,
	originalSchemeName string,
) runtime.Codec {
	internal := &codec{
		encoder:   encoder,
		decoder:   decoder,
		convertor: convertor,
		creater:   creater,
		typer:     typer,
		defaulter: defaulter,

		encodeVersion: encodeVersion,
		decodeVersion: decodeVersion,

		originalSchemeName: originalSchemeName,
	}
	return internal
}

type codec struct {
	encoder   runtime.Encoder
	decoder   runtime.Decoder
	convertor runtime.ObjectConvertor
	creater   runtime.ObjectCreater
	typer     runtime.ObjectTyper
	defaulter runtime.ObjectDefaulter

	encodeVersion runtime.GroupVersioner
	decodeVersion runtime.GroupVersioner

	// originalSchemeName is optional, but when filled in it holds the name of the scheme from which this codec originates
	originalSchemeName string
}

// Decode attempts a decode of the object, then tries to convert it to the internal version. If into is provided and the decoding is
// successful, the returned runtime.Object will be the value passed as into. Note that this may bypass conversion if you pass an
// into that matches the serialized version.
func (c *codec) Decode(data []byte, defaultGVK *schema.GroupVersionKind, into runtime.Object) (runtime.Object, *schema.GroupVersionKind, error) {
	versioned, isVersioned := into.(*runtime.VersionedObjects)
	if isVersioned {
		into = versioned.Last()
	}

	// If the into object is unstructured and expresses an opinion about its group/version,
	// create a new instance of the type so we always exercise the conversion path (skips short-circuiting on `into == obj`)
	decodeInto := into
	if into != nil {
		if _, ok := into.(runtime.Unstructured); ok && !into.GetObjectKind().GroupVersionKind().GroupVersion().Empty() {
			decodeInto = reflect.New(reflect.TypeOf(into).Elem()).Interface().(runtime.Object)
		}
	}

	obj, gvk, err := c.decoder.Decode(data, defaultGVK, decodeInto)
	if err != nil {
		return nil, gvk, err
	}

	if d, ok := obj.(runtime.NestedObjectDecoder); ok {
		if err := d.DecodeNestedObjects(runtime.WithoutVersionDecoder{c.decoder}); err != nil {
			return nil, gvk, err
		}
	}

	// if we specify a target, use generic conversion.
	if into != nil {
		if into == obj {
			if isVersioned {
				return versioned, gvk, nil
			}
			return into, gvk, nil
		}

		// perform defaulting if requested
		if c.defaulter != nil {
			// create a copy to ensure defaulting is not applied to the original versioned objects
			if isVersioned {
				versioned.Objects = []runtime.Object{obj.DeepCopyObject()}
			}
			c.defaulter.Default(obj)
		} else {
			if isVersioned {
				versioned.Objects = []runtime.Object{obj}
			}
		}

		if err := c.convertor.Convert(obj, into, c.decodeVersion); err != nil {
			return nil, gvk, err
		}

		if isVersioned {
			versioned.Objects = append(versioned.Objects, into)
			return versioned, gvk, nil
		}
		return into, gvk, nil
	}

	// Convert if needed.
	if isVersioned {
		// create a copy, because ConvertToVersion does not guarantee non-mutation of objects
		versioned.Objects = []runtime.Object{obj.DeepCopyObject()}
	}

	// perform defaulting if requested
	if c.defaulter != nil {
		c.defaulter.Default(obj)
	}

	out, err := c.convertor.ConvertToVersion(obj, c.decodeVersion)
	if err != nil {
		return nil, gvk, err
	}
	if isVersioned {
		if versioned.Last() != out {
			versioned.Objects = append(versioned.Objects, out)
		}
		return versioned, gvk, nil
	}
	return out, gvk, nil
}

// Encode ensures the provided object is output in the appropriate group and version, invoking
// conversion if necessary. Unversioned objects (according to the ObjectTyper) are output as is.
func (c *codec) Encode(obj runtime.Object, w io.Writer) error {
	switch obj := obj.(type) {
	case *runtime.Unknown:
		return c.encoder.Encode(obj, w)
	case runtime.Unstructured:
		// An unstructured list can contain objects of multiple group version kinds. don't short-circuit just
		// because the top-level type matches our desired destination type. actually send the object to the converter
		// to give it a chance to convert the list items if needed.
		if _, ok := obj.(*unstructured.UnstructuredList); !ok {
			// avoid conversion roundtrip if GVK is the right one already or is empty (yes, this is a hack, but the old behaviour we rely on in kubectl)
			objGVK := obj.GetObjectKind().GroupVersionKind()
			if len(objGVK.Version) == 0 {
				return c.encoder.Encode(obj, w)
			}
			targetGVK, ok := c.encodeVersion.KindForGroupVersionKinds([]schema.GroupVersionKind{objGVK})
			if !ok {
				return runtime.NewNotRegisteredGVKErrForTarget(c.originalSchemeName, objGVK, c.encodeVersion)
			}
			if targetGVK == objGVK {
				return c.encoder.Encode(obj, w)
			}
		}
	}

	gvks, isUnversioned, err := c.typer.ObjectKinds(obj)
	if err != nil {
		return err
	}

	objectKind := obj.GetObjectKind()
	old := objectKind.GroupVersionKind()
	// restore the old GVK after encoding
	defer objectKind.SetGroupVersionKind(old)

	if c.encodeVersion == nil || isUnversioned {
		if e, ok := obj.(runtime.NestedObjectEncoder); ok {
			if err := e.EncodeNestedObjects(runtime.WithVersionEncoder{Encoder: c.encoder, ObjectTyper: c.typer}); err != nil {
				return err
			}
		}
		objectKind.SetGroupVersionKind(gvks[0])
		return c.encoder.Encode(obj, w)
	}

	// Perform a conversion if necessary
	out, err := c.convertor.ConvertToVersion(obj, c.encodeVersion)
	if err != nil {
		return err
	}

	if e, ok := out.(runtime.NestedObjectEncoder); ok {
		if err := e.EncodeNestedObjects(runtime.WithVersionEncoder{Version: c.encodeVersion, Encoder: c.encoder, ObjectTyper: c.typer}); err != nil {
			return err
		}
	}

	// Conversion is responsible for setting the proper group, version, and kind onto the outgoing object
	return c.encoder.Encode(out, w)
}

// DirectEncoder was moved and renamed to runtime.WithVersionEncoder in 1.15.
// TODO: remove in 1.16.
type DirectEncoder = runtime.WithVersionEncoder

// DirectDecoder was moved and renamed to runtime.WithoutVersionDecoder in 1.15.
// TODO: remove in 1.16.
type DirectDecoder = runtime.WithoutVersionDecoder
<commit_msg>Ensure defaulting applies to custom resource items in list response<commit_after>/*
Copyright 2014 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 versioning

import (
	"io"
	"reflect"

	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
)

// NewDefaultingCodecForScheme is a convenience method for callers that are using a scheme.
func NewDefaultingCodecForScheme(
	// TODO: I should be a scheme interface?
	scheme *runtime.Scheme,
	encoder runtime.Encoder,
	decoder runtime.Decoder,
	encodeVersion runtime.GroupVersioner,
	decodeVersion runtime.GroupVersioner,
) runtime.Codec {
	return NewCodec(encoder, decoder, runtime.UnsafeObjectConvertor(scheme), scheme, scheme, scheme, encodeVersion, decodeVersion, scheme.Name())
}

// NewCodec takes objects in their internal versions and converts them to external versions before
// serializing them. It assumes the serializer provided to it only deals with external versions.
// This class is also a serializer, but is generally used with a specific version.
func NewCodec(
	encoder runtime.Encoder,
	decoder runtime.Decoder,
	convertor runtime.ObjectConvertor,
	creater runtime.ObjectCreater,
	typer runtime.ObjectTyper,
	defaulter runtime.ObjectDefaulter,
	encodeVersion runtime.GroupVersioner,
	decodeVersion runtime.GroupVersioner,
	originalSchemeName string,
) runtime.Codec {
	internal := &codec{
		encoder:   encoder,
		decoder:   decoder,
		convertor: convertor,
		creater:   creater,
		typer:     typer,
		defaulter: defaulter,

		encodeVersion: encodeVersion,
		decodeVersion: decodeVersion,

		originalSchemeName: originalSchemeName,
	}
	return internal
}

type codec struct {
	encoder   runtime.Encoder
	decoder   runtime.Decoder
	convertor runtime.ObjectConvertor
	creater   runtime.ObjectCreater
	typer     runtime.ObjectTyper
	defaulter runtime.ObjectDefaulter

	encodeVersion runtime.GroupVersioner
	decodeVersion runtime.GroupVersioner

	// originalSchemeName is optional, but when filled in it holds the name of the scheme from which this codec originates
	originalSchemeName string
}

// Decode attempts a decode of the object, then tries to convert it to the internal version. If into is provided and the decoding is
// successful, the returned runtime.Object will be the value passed as into. Note that this may bypass conversion if you pass an
// into that matches the serialized version.
func (c *codec) Decode(data []byte, defaultGVK *schema.GroupVersionKind, into runtime.Object) (runtime.Object, *schema.GroupVersionKind, error) {
	versioned, isVersioned := into.(*runtime.VersionedObjects)
	if isVersioned {
		into = versioned.Last()
	}

	// If the into object is unstructured and expresses an opinion about its group/version,
	// create a new instance of the type so we always exercise the conversion path (skips short-circuiting on `into == obj`)
	decodeInto := into
	if into != nil {
		if _, ok := into.(runtime.Unstructured); ok && !into.GetObjectKind().GroupVersionKind().GroupVersion().Empty() {
			decodeInto = reflect.New(reflect.TypeOf(into).Elem()).Interface().(runtime.Object)
		}
	}

	obj, gvk, err := c.decoder.Decode(data, defaultGVK, decodeInto)
	if err != nil {
		return nil, gvk, err
	}

	if d, ok := obj.(runtime.NestedObjectDecoder); ok {
		if err := d.DecodeNestedObjects(runtime.WithoutVersionDecoder{c.decoder}); err != nil {
			return nil, gvk, err
		}
	}

	// if we specify a target, use generic conversion.
	if into != nil {
		// perform defaulting if requested
		if c.defaulter != nil {
			// create a copy to ensure defaulting is not applied to the original versioned objects
			if isVersioned {
				versioned.Objects = []runtime.Object{obj.DeepCopyObject()}
			}
			c.defaulter.Default(obj)
		} else {
			if isVersioned {
				versioned.Objects = []runtime.Object{obj}
			}
		}

		// Short-circuit conversion if the into object is same object
		if into == obj {
			if isVersioned {
				return versioned, gvk, nil
			}
			return into, gvk, nil
		}

		if err := c.convertor.Convert(obj, into, c.decodeVersion); err != nil {
			return nil, gvk, err
		}

		if isVersioned {
			versioned.Objects = append(versioned.Objects, into)
			return versioned, gvk, nil
		}
		return into, gvk, nil
	}

	// Convert if needed.
	if isVersioned {
		// create a copy, because ConvertToVersion does not guarantee non-mutation of objects
		versioned.Objects = []runtime.Object{obj.DeepCopyObject()}
	}

	// perform defaulting if requested
	if c.defaulter != nil {
		c.defaulter.Default(obj)
	}

	out, err := c.convertor.ConvertToVersion(obj, c.decodeVersion)
	if err != nil {
		return nil, gvk, err
	}
	if isVersioned {
		if versioned.Last() != out {
			versioned.Objects = append(versioned.Objects, out)
		}
		return versioned, gvk, nil
	}
	return out, gvk, nil
}

// Encode ensures the provided object is output in the appropriate group and version, invoking
// conversion if necessary. Unversioned objects (according to the ObjectTyper) are output as is.
func (c *codec) Encode(obj runtime.Object, w io.Writer) error {
	switch obj := obj.(type) {
	case *runtime.Unknown:
		return c.encoder.Encode(obj, w)
	case runtime.Unstructured:
		// An unstructured list can contain objects of multiple group version kinds. don't short-circuit just
		// because the top-level type matches our desired destination type. actually send the object to the converter
		// to give it a chance to convert the list items if needed.
		if _, ok := obj.(*unstructured.UnstructuredList); !ok {
			// avoid conversion roundtrip if GVK is the right one already or is empty (yes, this is a hack, but the old behaviour we rely on in kubectl)
			objGVK := obj.GetObjectKind().GroupVersionKind()
			if len(objGVK.Version) == 0 {
				return c.encoder.Encode(obj, w)
			}
			targetGVK, ok := c.encodeVersion.KindForGroupVersionKinds([]schema.GroupVersionKind{objGVK})
			if !ok {
				return runtime.NewNotRegisteredGVKErrForTarget(c.originalSchemeName, objGVK, c.encodeVersion)
			}
			if targetGVK == objGVK {
				return c.encoder.Encode(obj, w)
			}
		}
	}

	gvks, isUnversioned, err := c.typer.ObjectKinds(obj)
	if err != nil {
		return err
	}

	objectKind := obj.GetObjectKind()
	old := objectKind.GroupVersionKind()
	// restore the old GVK after encoding
	defer objectKind.SetGroupVersionKind(old)

	if c.encodeVersion == nil || isUnversioned {
		if e, ok := obj.(runtime.NestedObjectEncoder); ok {
			if err := e.EncodeNestedObjects(runtime.WithVersionEncoder{Encoder: c.encoder, ObjectTyper: c.typer}); err != nil {
				return err
			}
		}
		objectKind.SetGroupVersionKind(gvks[0])
		return c.encoder.Encode(obj, w)
	}

	// Perform a conversion if necessary
	out, err := c.convertor.ConvertToVersion(obj, c.encodeVersion)
	if err != nil {
		return err
	}

	if e, ok := out.(runtime.NestedObjectEncoder); ok {
		if err := e.EncodeNestedObjects(runtime.WithVersionEncoder{Version: c.encodeVersion, Encoder: c.encoder, ObjectTyper: c.typer}); err != nil {
			return err
		}
	}

	// Conversion is responsible for setting the proper group, version, and kind onto the outgoing object
	return c.encoder.Encode(out, w)
}

// DirectEncoder was moved and renamed to runtime.WithVersionEncoder in 1.15.
// TODO: remove in 1.16.
type DirectEncoder = runtime.WithVersionEncoder

// DirectDecoder was moved and renamed to runtime.WithoutVersionDecoder in 1.15.
// TODO: remove in 1.16.
type DirectDecoder = runtime.WithoutVersionDecoder
<|endoftext|>


<commit_before>// Package buildtsi reads an in-memory index and exports it as a TSI index.
package buildtsi

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/influxdata/influxdb/logger"
	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb"
	"github.com/influxdata/influxdb/tsdb/engine/tsm1"
	"github.com/influxdata/influxdb/tsdb/index/tsi1"
	"go.uber.org/zap"
)

const defaultBatchSize = 10000

// Command represents the program execution for "influx_inspect buildtsi".
type Command struct {
	Stderr  io.Writer
	Stdout  io.Writer
	Verbose bool
	Logger  *zap.Logger

	databaseFilter  string
	retentionFilter string
	shardFilter     string
	maxLogFileSize  int64
	batchSize       int
}

// NewCommand returns a new instance of Command.
func NewCommand() *Command {
	return &Command{
		Stderr:    os.Stderr,
		Stdout:    os.Stdout,
		Logger:    zap.NewNop(),
		batchSize: defaultBatchSize,
	}
}

// Run executes the command.
func (cmd *Command) Run(args ...string) error {
	fs := flag.NewFlagSet("buildtsi", flag.ExitOnError)
	dataDir := fs.String("datadir", "", "data directory")
	walDir := fs.String("waldir", "", "WAL directory")
	fs.StringVar(&cmd.databaseFilter, "database", "", "optional: database name")
	fs.StringVar(&cmd.retentionFilter, "retention", "", "optional: retention policy")
	fs.StringVar(&cmd.shardFilter, "shard", "", "optional: shard id")
	fs.Int64Var(&cmd.maxLogFileSize, "max-log-file-size", tsdb.DefaultMaxIndexLogFileSize, "optional: maximum log file size")
	fs.IntVar(&cmd.batchSize, "batch-size", defaultBatchSize, "optional: set the size of the batches we write to the index. Setting this can have adverse affects on performance and heap requirements")
	fs.BoolVar(&cmd.Verbose, "v", false, "verbose")
	fs.SetOutput(cmd.Stdout)
	if err := fs.Parse(args); err != nil {
		return err
	} else if fs.NArg() > 0 || *dataDir == "" || *walDir == "" {
		fs.Usage()
		return nil
	}
	cmd.Logger = logger.New(cmd.Stderr)

	finish := startProfiles()
	defer finish()
	return cmd.run(*dataDir, *walDir)
}

func startProfiles() func() {
	runtime.MemProfileRate = 100 // Sample 1% of allocations.

	paths := []string{"/tmp/buildtsi.mem.pprof", "/tmp/buildtsi.cpu.pprof"}
	var files []*os.File
	for _, pth := range paths {
		f, err := os.Create(pth)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing profile to: %s\n", pth)
		files = append(files, f)

	}

	closeFn := func() {
		// Write the memory profile
		if err := pprof.Lookup("heap").WriteTo(files[0], 0); err != nil {
			panic(err)
		}

		// Stop the CPU profile.
		pprof.StopCPUProfile()

		for _, fd := range files {
			if err := fd.Close(); err != nil {
				panic(err)
			}
		}
	}

	if err := pprof.StartCPUProfile(files[1]); err != nil {
		panic(err)
	}
	return closeFn
}

func (cmd *Command) run(dataDir, walDir string) error {
	// Verify the user actually wants to run as root.
	if isRoot() {
		fmt.Println("You are currently running as root. This will build your")
		fmt.Println("index files with root ownership and will be inaccessible")
		fmt.Println("if you run influxd as a non-root user. You should run")
		fmt.Println("buildtsi as the same user you are running influxd.")
		fmt.Print("Are you sure you want to continue? (y/N): ")
		var answer string
		if fmt.Scanln(&answer); !strings.HasPrefix(strings.TrimSpace(strings.ToLower(answer)), "y") {
			return fmt.Errorf("operation aborted")
		}
	}

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		name := fi.Name()
		if !fi.IsDir() {
			continue
		} else if cmd.databaseFilter != "" && name != cmd.databaseFilter {
			continue
		}

		if err := cmd.processDatabase(name, filepath.Join(dataDir, name), filepath.Join(walDir, name)); err != nil {
			return err
		}
	}

	return nil
}

func (cmd *Command) processDatabase(dbName, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding database", zap.String("name", dbName))

	sfile := tsdb.NewSeriesFile(filepath.Join(dataDir, tsdb.SeriesFileDirectory))
	sfile.Logger = cmd.Logger
	if err := sfile.Open(); err != nil {
		return err
	}
	defer sfile.Close()

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		rpName := fi.Name()
		if !fi.IsDir() {
			continue
		} else if rpName == tsdb.SeriesFileDirectory {
			continue
		} else if cmd.retentionFilter != "" && rpName != cmd.retentionFilter {
			continue
		}

		if err := cmd.processRetentionPolicy(sfile, dbName, rpName, filepath.Join(dataDir, rpName), filepath.Join(walDir, rpName)); err != nil {
			return err
		}
	}

	return nil
}

func (cmd *Command) processRetentionPolicy(sfile *tsdb.SeriesFile, dbName, rpName, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding retention policy", logger.Database(dbName), logger.RetentionPolicy(rpName))

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		if !fi.IsDir() {
			continue
		} else if cmd.shardFilter != "" && fi.Name() != cmd.shardFilter {
			continue
		}

		shardID, err := strconv.ParseUint(fi.Name(), 10, 64)
		if err != nil {
			continue
		}

		if err := cmd.processShard(sfile, dbName, rpName, shardID, filepath.Join(dataDir, fi.Name()), filepath.Join(walDir, fi.Name())); err != nil {
			return err
		}
	}
	return nil
}

func (cmd *Command) processShard(sfile *tsdb.SeriesFile, dbName, rpName string, shardID uint64, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding shard", logger.Database(dbName), logger.RetentionPolicy(rpName), logger.Shard(shardID))

	// Check if shard already has a TSI index.
	indexPath := filepath.Join(dataDir, "index")
	cmd.Logger.Info("Checking index path", zap.String("path", indexPath))
	if _, err := os.Stat(indexPath); !os.IsNotExist(err) {
		cmd.Logger.Info("tsi1 index already exists, skipping", zap.String("path", indexPath))
		return nil
	}

	cmd.Logger.Info("Opening shard")

	// Find shard files.
	tsmPaths, err := cmd.collectTSMFiles(dataDir)
	if err != nil {
		return err
	}
	walPaths, err := cmd.collectWALFiles(walDir)
	if err != nil {
		return err
	}

	// Remove temporary index files if this is being re-run.
	tmpPath := filepath.Join(dataDir, ".index")
	cmd.Logger.Info("Cleaning up partial index from previous run, if any")
	if err := os.RemoveAll(tmpPath); err != nil {
		return err
	}

	// Open TSI index in temporary path.
	tsiIndex := tsi1.NewIndex(sfile, dbName,
		tsi1.WithPath(tmpPath),
		tsi1.WithMaximumLogFileSize(cmd.maxLogFileSize),
		tsi1.DisableFsync(),
		// Each new series entry in a log file is ~12 bytes so this should
		// roughly equate to one flush to the file for every batch.
		tsi1.WithLogFileBufferSize(12*cmd.batchSize),
	)

	tsiIndex.WithLogger(cmd.Logger)

	cmd.Logger.Info("Opening tsi index in temporary location", zap.String("path", tmpPath))
	if err := tsiIndex.Open(); err != nil {
		return err
	}
	defer tsiIndex.Close()

	// Write out tsm1 files.
	cmd.Logger.Info("Iterating over tsm files")
	for _, path := range tsmPaths {
		cmd.Logger.Info("Processing tsm file", zap.String("path", path))
		if err := cmd.processTSMFile(tsiIndex, path); err != nil {
			return err
		}
	}

	// Write out wal files.
	cmd.Logger.Info("Building cache from wal files")
	cache := tsm1.NewCache(tsdb.DefaultCacheMaxMemorySize)
	loader := tsm1.NewCacheLoader(walPaths)
	loader.WithLogger(cmd.Logger)
	if err := loader.Load(cache); err != nil {
		return err
	}

	cmd.Logger.Info("Iterating over cache")
	keysBatch := make([][]byte, 0, cmd.batchSize)
	namesBatch := make([][]byte, 0, cmd.batchSize)
	tagsBatch := make([]models.Tags, 0, cmd.batchSize)

	for _, key := range cache.Keys() {
		seriesKey, _ := tsm1.SeriesAndFieldFromCompositeKey(key)
		name, tags := models.ParseKey(seriesKey)

		if cmd.Verbose {
			cmd.Logger.Info("Series", zap.String("name", string(name)), zap.String("tags", tags.String()))
		}

		keysBatch = append(keysBatch, seriesKey)
		namesBatch = append(namesBatch, []byte(name))
		tagsBatch = append(tagsBatch, tags)

		// Flush batch?
		if len(keysBatch) == cmd.batchSize {
			if err := tsiIndex.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
				return fmt.Errorf("problem creating series: (%s)", err)
			}
			keysBatch = keysBatch[:0]
			namesBatch = namesBatch[:0]
			tagsBatch = tagsBatch[:0]
		}
	}

	// Flush any remaining series in the batches
	if len(keysBatch) > 0 {
		if err := tsiIndex.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
			return fmt.Errorf("problem creating series: (%s)", err)
		}
		keysBatch = nil
		namesBatch = nil
		tagsBatch = nil
	}

	// Attempt to compact the index & wait for all compactions to complete.
	cmd.Logger.Info("compacting index")
	tsiIndex.Compact()
	tsiIndex.Wait()

	// Close TSI index.
	cmd.Logger.Info("Closing tsi index")
	if err := tsiIndex.Close(); err != nil {
		return err
	}

	// Rename TSI to standard path.
	cmd.Logger.Info("Moving tsi to permanent location")
	return os.Rename(tmpPath, indexPath)
}

func (cmd *Command) processTSMFile(index *tsi1.Index, path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	r, err := tsm1.NewTSMReader(f)
	if err != nil {
		cmd.Logger.Warn("Unable to read, skipping", zap.String("path", path), zap.Error(err))
		return nil
	}
	defer r.Close()

	keysBatch := make([][]byte, 0, cmd.batchSize)
	namesBatch := make([][]byte, 0, cmd.batchSize)
	tagsBatch := make([]models.Tags, 0, cmd.batchSize)

	for i := 0; i < r.KeyCount(); i++ {
		key, _ := r.KeyAt(i)
		seriesKey, _ := tsm1.SeriesAndFieldFromCompositeKey(key)
		name, tags := models.ParseKey(seriesKey)

		if cmd.Verbose {
			cmd.Logger.Info("Series", zap.String("name", string(name)), zap.String("tags", tags.String()))
		}

		keysBatch = append(keysBatch, seriesKey)
		namesBatch = append(namesBatch, []byte(name))
		tagsBatch = append(tagsBatch, tags)

		// Flush batch?
		if len(keysBatch) == cmd.batchSize {
			if err := index.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
				return fmt.Errorf("problem creating series: (%s)", err)
			}
			keysBatch = keysBatch[:0]
			namesBatch = namesBatch[:0]
			tagsBatch = tagsBatch[:0]
		}
	}

	// Flush any remaining series in the batches
	if len(keysBatch) > 0 {
		if err := index.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
			return fmt.Errorf("problem creating series: (%s)", err)
		}
	}
	return nil
}

func (cmd *Command) collectTSMFiles(path string) ([]string, error) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}

	var paths []string
	for _, fi := range fis {
		if filepath.Ext(fi.Name()) != "."+tsm1.TSMFileExtension {
			continue
		}
		paths = append(paths, filepath.Join(path, fi.Name()))
	}
	return paths, nil
}

func (cmd *Command) collectWALFiles(path string) ([]string, error) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}

	var paths []string
	for _, fi := range fis {
		if filepath.Ext(fi.Name()) != "."+tsm1.WALFileExtension {
			continue
		}
		paths = append(paths, filepath.Join(path, fi.Name()))
	}
	return paths, nil
}

func isRoot() bool {
	user, _ := user.Current()
	return user != nil && user.Username == "root"
}
<commit_msg>Avoid allocating a string for measurment names<commit_after>// Package buildtsi reads an in-memory index and exports it as a TSI index.
package buildtsi

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"

	"github.com/influxdata/influxdb/logger"
	"github.com/influxdata/influxdb/models"
	"github.com/influxdata/influxdb/tsdb"
	"github.com/influxdata/influxdb/tsdb/engine/tsm1"
	"github.com/influxdata/influxdb/tsdb/index/tsi1"
	"go.uber.org/zap"
)

const defaultBatchSize = 10000

// Command represents the program execution for "influx_inspect buildtsi".
type Command struct {
	Stderr  io.Writer
	Stdout  io.Writer
	Verbose bool
	Logger  *zap.Logger

	databaseFilter  string
	retentionFilter string
	shardFilter     string
	maxLogFileSize  int64
	batchSize       int
}

// NewCommand returns a new instance of Command.
func NewCommand() *Command {
	return &Command{
		Stderr:    os.Stderr,
		Stdout:    os.Stdout,
		Logger:    zap.NewNop(),
		batchSize: defaultBatchSize,
	}
}

// Run executes the command.
func (cmd *Command) Run(args ...string) error {
	fs := flag.NewFlagSet("buildtsi", flag.ExitOnError)
	dataDir := fs.String("datadir", "", "data directory")
	walDir := fs.String("waldir", "", "WAL directory")
	fs.StringVar(&cmd.databaseFilter, "database", "", "optional: database name")
	fs.StringVar(&cmd.retentionFilter, "retention", "", "optional: retention policy")
	fs.StringVar(&cmd.shardFilter, "shard", "", "optional: shard id")
	fs.Int64Var(&cmd.maxLogFileSize, "max-log-file-size", tsdb.DefaultMaxIndexLogFileSize, "optional: maximum log file size")
	fs.IntVar(&cmd.batchSize, "batch-size", defaultBatchSize, "optional: set the size of the batches we write to the index. Setting this can have adverse affects on performance and heap requirements")
	fs.BoolVar(&cmd.Verbose, "v", false, "verbose")
	fs.SetOutput(cmd.Stdout)
	if err := fs.Parse(args); err != nil {
		return err
	} else if fs.NArg() > 0 || *dataDir == "" || *walDir == "" {
		fs.Usage()
		return nil
	}
	cmd.Logger = logger.New(cmd.Stderr)

	finish := startProfiles()
	defer finish()
	return cmd.run(*dataDir, *walDir)
}

func startProfiles() func() {
	runtime.MemProfileRate = 100 // Sample 1% of allocations.

	paths := []string{"/tmp/buildtsi.mem.pprof", "/tmp/buildtsi.cpu.pprof"}
	var files []*os.File
	for _, pth := range paths {
		f, err := os.Create(pth)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing profile to: %s\n", pth)
		files = append(files, f)

	}

	closeFn := func() {
		// Write the memory profile
		if err := pprof.Lookup("heap").WriteTo(files[0], 0); err != nil {
			panic(err)
		}

		// Stop the CPU profile.
		pprof.StopCPUProfile()

		for _, fd := range files {
			if err := fd.Close(); err != nil {
				panic(err)
			}
		}
	}

	if err := pprof.StartCPUProfile(files[1]); err != nil {
		panic(err)
	}
	return closeFn
}

func (cmd *Command) run(dataDir, walDir string) error {
	// Verify the user actually wants to run as root.
	if isRoot() {
		fmt.Println("You are currently running as root. This will build your")
		fmt.Println("index files with root ownership and will be inaccessible")
		fmt.Println("if you run influxd as a non-root user. You should run")
		fmt.Println("buildtsi as the same user you are running influxd.")
		fmt.Print("Are you sure you want to continue? (y/N): ")
		var answer string
		if fmt.Scanln(&answer); !strings.HasPrefix(strings.TrimSpace(strings.ToLower(answer)), "y") {
			return fmt.Errorf("operation aborted")
		}
	}

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		name := fi.Name()
		if !fi.IsDir() {
			continue
		} else if cmd.databaseFilter != "" && name != cmd.databaseFilter {
			continue
		}

		if err := cmd.processDatabase(name, filepath.Join(dataDir, name), filepath.Join(walDir, name)); err != nil {
			return err
		}
	}

	return nil
}

func (cmd *Command) processDatabase(dbName, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding database", zap.String("name", dbName))

	sfile := tsdb.NewSeriesFile(filepath.Join(dataDir, tsdb.SeriesFileDirectory))
	sfile.Logger = cmd.Logger
	if err := sfile.Open(); err != nil {
		return err
	}
	defer sfile.Close()

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		rpName := fi.Name()
		if !fi.IsDir() {
			continue
		} else if rpName == tsdb.SeriesFileDirectory {
			continue
		} else if cmd.retentionFilter != "" && rpName != cmd.retentionFilter {
			continue
		}

		if err := cmd.processRetentionPolicy(sfile, dbName, rpName, filepath.Join(dataDir, rpName), filepath.Join(walDir, rpName)); err != nil {
			return err
		}
	}

	return nil
}

func (cmd *Command) processRetentionPolicy(sfile *tsdb.SeriesFile, dbName, rpName, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding retention policy", logger.Database(dbName), logger.RetentionPolicy(rpName))

	fis, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return err
	}

	for _, fi := range fis {
		if !fi.IsDir() {
			continue
		} else if cmd.shardFilter != "" && fi.Name() != cmd.shardFilter {
			continue
		}

		shardID, err := strconv.ParseUint(fi.Name(), 10, 64)
		if err != nil {
			continue
		}

		if err := cmd.processShard(sfile, dbName, rpName, shardID, filepath.Join(dataDir, fi.Name()), filepath.Join(walDir, fi.Name())); err != nil {
			return err
		}
	}
	return nil
}

func (cmd *Command) processShard(sfile *tsdb.SeriesFile, dbName, rpName string, shardID uint64, dataDir, walDir string) error {
	cmd.Logger.Info("Rebuilding shard", logger.Database(dbName), logger.RetentionPolicy(rpName), logger.Shard(shardID))

	// Check if shard already has a TSI index.
	indexPath := filepath.Join(dataDir, "index")
	cmd.Logger.Info("Checking index path", zap.String("path", indexPath))
	if _, err := os.Stat(indexPath); !os.IsNotExist(err) {
		cmd.Logger.Info("tsi1 index already exists, skipping", zap.String("path", indexPath))
		return nil
	}

	cmd.Logger.Info("Opening shard")

	// Find shard files.
	tsmPaths, err := cmd.collectTSMFiles(dataDir)
	if err != nil {
		return err
	}
	walPaths, err := cmd.collectWALFiles(walDir)
	if err != nil {
		return err
	}

	// Remove temporary index files if this is being re-run.
	tmpPath := filepath.Join(dataDir, ".index")
	cmd.Logger.Info("Cleaning up partial index from previous run, if any")
	if err := os.RemoveAll(tmpPath); err != nil {
		return err
	}

	// Open TSI index in temporary path.
	tsiIndex := tsi1.NewIndex(sfile, dbName,
		tsi1.WithPath(tmpPath),
		tsi1.WithMaximumLogFileSize(cmd.maxLogFileSize),
		tsi1.DisableFsync(),
		// Each new series entry in a log file is ~12 bytes so this should
		// roughly equate to one flush to the file for every batch.
		tsi1.WithLogFileBufferSize(12*cmd.batchSize),
	)

	tsiIndex.WithLogger(cmd.Logger)

	cmd.Logger.Info("Opening tsi index in temporary location", zap.String("path", tmpPath))
	if err := tsiIndex.Open(); err != nil {
		return err
	}
	defer tsiIndex.Close()

	// Write out tsm1 files.
	cmd.Logger.Info("Iterating over tsm files")
	for _, path := range tsmPaths {
		cmd.Logger.Info("Processing tsm file", zap.String("path", path))
		if err := cmd.processTSMFile(tsiIndex, path); err != nil {
			return err
		}
	}

	// Write out wal files.
	cmd.Logger.Info("Building cache from wal files")
	cache := tsm1.NewCache(tsdb.DefaultCacheMaxMemorySize)
	loader := tsm1.NewCacheLoader(walPaths)
	loader.WithLogger(cmd.Logger)
	if err := loader.Load(cache); err != nil {
		return err
	}

	cmd.Logger.Info("Iterating over cache")
	keysBatch := make([][]byte, 0, cmd.batchSize)
	namesBatch := make([][]byte, 0, cmd.batchSize)
	tagsBatch := make([]models.Tags, 0, cmd.batchSize)

	for _, key := range cache.Keys() {
		seriesKey, _ := tsm1.SeriesAndFieldFromCompositeKey(key)
		name, tags := models.ParseKeyBytes(seriesKey)

		if cmd.Verbose {
			cmd.Logger.Info("Series", zap.String("name", string(name)), zap.String("tags", tags.String()))
		}

		keysBatch = append(keysBatch, seriesKey)
		namesBatch = append(namesBatch, name)
		tagsBatch = append(tagsBatch, tags)

		// Flush batch?
		if len(keysBatch) == cmd.batchSize {
			if err := tsiIndex.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
				return fmt.Errorf("problem creating series: (%s)", err)
			}
			keysBatch = keysBatch[:0]
			namesBatch = namesBatch[:0]
			tagsBatch = tagsBatch[:0]
		}
	}

	// Flush any remaining series in the batches
	if len(keysBatch) > 0 {
		if err := tsiIndex.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
			return fmt.Errorf("problem creating series: (%s)", err)
		}
		keysBatch = nil
		namesBatch = nil
		tagsBatch = nil
	}

	// Attempt to compact the index & wait for all compactions to complete.
	cmd.Logger.Info("compacting index")
	tsiIndex.Compact()
	tsiIndex.Wait()

	// Close TSI index.
	cmd.Logger.Info("Closing tsi index")
	if err := tsiIndex.Close(); err != nil {
		return err
	}

	// Rename TSI to standard path.
	cmd.Logger.Info("Moving tsi to permanent location")
	return os.Rename(tmpPath, indexPath)
}

func (cmd *Command) processTSMFile(index *tsi1.Index, path string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	r, err := tsm1.NewTSMReader(f)
	if err != nil {
		cmd.Logger.Warn("Unable to read, skipping", zap.String("path", path), zap.Error(err))
		return nil
	}
	defer r.Close()

	keysBatch := make([][]byte, 0, cmd.batchSize)
	namesBatch := make([][]byte, 0, cmd.batchSize)
	tagsBatch := make([]models.Tags, 0, cmd.batchSize)

	for i := 0; i < r.KeyCount(); i++ {
		key, _ := r.KeyAt(i)
		seriesKey, _ := tsm1.SeriesAndFieldFromCompositeKey(key)
		name, tags := models.ParseKeyBytes(seriesKey)

		if cmd.Verbose {
			cmd.Logger.Info("Series", zap.String("name", string(name)), zap.String("tags", tags.String()))
		}

		keysBatch = append(keysBatch, seriesKey)
		namesBatch = append(namesBatch, name)
		tagsBatch = append(tagsBatch, tags)

		// Flush batch?
		if len(keysBatch) == cmd.batchSize {
			if err := index.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
				return fmt.Errorf("problem creating series: (%s)", err)
			}
			keysBatch = keysBatch[:0]
			namesBatch = namesBatch[:0]
			tagsBatch = tagsBatch[:0]
		}
	}

	// Flush any remaining series in the batches
	if len(keysBatch) > 0 {
		if err := index.CreateSeriesListIfNotExists(keysBatch, namesBatch, tagsBatch); err != nil {
			return fmt.Errorf("problem creating series: (%s)", err)
		}
	}
	return nil
}

func (cmd *Command) collectTSMFiles(path string) ([]string, error) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}

	var paths []string
	for _, fi := range fis {
		if filepath.Ext(fi.Name()) != "."+tsm1.TSMFileExtension {
			continue
		}
		paths = append(paths, filepath.Join(path, fi.Name()))
	}
	return paths, nil
}

func (cmd *Command) collectWALFiles(path string) ([]string, error) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}

	var paths []string
	for _, fi := range fis {
		if filepath.Ext(fi.Name()) != "."+tsm1.WALFileExtension {
			continue
		}
		paths = append(paths, filepath.Join(path, fi.Name()))
	}
	return paths, nil
}

func isRoot() bool {
	user, _ := user.Current()
	return user != nil && user.Username == "root"
}
<|endoftext|>


<commit_before>// Copyright 2016 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 cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"text/tabwriter"
	"time"

	"github.com/google/key-transparency/core/keystore"

	"github.com/golang/protobuf/ptypes"
	"github.com/spf13/cobra"

	kmpb "github.com/google/key-transparency/core/proto/keymaster"
)

const (
	keyStoreFile      = ".keystore"
	keyIDTruncatedLen = 8
)

var (
	pubKey      string
	privKey     string
	description string
	activate    bool
)

var store *keystore.KeyStore

// keysCmd represents the authorized-keys command.
var keysCmd = &cobra.Command{
	Use:   "authorized-keys",
	Short: "Manage authorized keys",
	Long: `Manage keys authorized to sign updates, and select the active signing key.

Verifying always happens using the keys listed in the previous epoch.`,
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		if err := readKeyStoreFile(); err != nil {
			log.Fatal(err)
		}
	},
	PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
		buf, err := store.Marshal()
		if err != nil {
			return err
		}
		if err := ioutil.WriteFile(keyStoreFile, buf, 0600); err != nil {
			return err
		}
		return nil
	},
}

// addCmd represents the authorized-keys add command.
var addCmd = &cobra.Command{
	Use:   "add [ --privkey=[path] --active | --pubkey=[path] ] --description=[comment]",
	Short: "Add a key to the list of authorized keys",
	Long: `Provide a pair of public and private keys, already existing on disk, to be added to the list of authorized keys. e.g.:

./key-transparency-client authorized-keys add --pubkey=/path/to/PEM/pubkey --description=[comment]
./key-transparency-client authorized-keys add --privkey=/path/to/PEM/privkey --activate --description=[comment]
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if pubKey == "" && privKey == "" {
			return fmt.Errorf("should provide private or public key")
		}
		if pubKey != "" && privKey != "" {
			return fmt.Errorf("cannot provide public and private key at the same time")
		}

		// Add either a private key, or a public key.
		switch {
		case pubKey != "":
			if err := addPublicKey(pubKey, description); err == nil {
				return err
			}
		case privKey != "":
			if err := addPrivateKey(privKey, description, activate); err == nil {
				return err
			}
		}
		return nil
	},
}

// removeCmd represents the authorized-keys remove command.
var removeCmd = &cobra.Command{
	Use:   "remove [keyid]",
	Short: "Remove a key from the list of authorized keys",
	Long: `Remove a key based on its key ID from the list of authorized keys. e.g.:

./key-transparency-client authorized-keys remove [keyid]

If the list contains a single key, it cannot be removed.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if len(args) < 1 {
			return fmt.Errorf("key ID needs to be provided")
		}

		// First remove verifying key with matching ID then attempt to
		// remove signing key if exist.
		keyID := keyID(args[0])
		if err := store.RemoveVerifyingKey(keyID); err != nil {
			return err
		}
		if err := store.RemoveSigningKey(keyID); err != nil && err != keystore.ErrKeyNotExist {
			return err
		}
		return nil
	},
}

// activateCmd represents the authorized-keys activate command.
var activateCmd = &cobra.Command{
	Use:   "activate [keyid]",
	Short: "Activate a key in the list of authorized keys",
	Long: `Activate a key based on its key ID in the list of authorized keys. e.g.:

./key-transparency-client authorized-keys activate [keyid]
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if len(args) < 1 {
			return fmt.Errorf("key ID needs to be provided")
		}
		keyID := keyID(args[0])
		if err := store.Activate(keyID); err != nil {
			return err
		}
		return nil
	},
}

// listCmd represents the authorized-keys list command.
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List all authorized keys",
	Long: `List metadata about all authorized keys. e.g.:

./key-transparency-client authorized-keys list

The actual keys are not listed, only their corresponding metadata.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		signingInfo, verifyingInfo, err := store.Info()
		if err != nil {
			return err
		}

		w := tabwriter.NewWriter(os.Stdout, 0, 0, 1, ' ', tabwriter.Debug)

		// List signing keys.
		fmt.Fprintln(w, "Signing Keys:")
		fmt.Fprintln(w, "  ID\tAdded At\tStatus\tDescription")
		for _, info := range signingInfo {
			timestamp, err := ptypes.Timestamp(info.Metadata.AddedAt)
			if err != nil {
				return err
			}
			fmt.Fprintf(w, "  %v\t%v\t%v\t%v\n", info.Metadata.KeyId[:keyIDTruncatedLen], timestamp.Format(time.ANSIC), info.Status, info.Metadata.Description)
		}

		// List verifying keys.
		fmt.Fprintln(w, "\nVerifying Keys:")
		fmt.Fprintln(w, "  ID\tAdded At\tStatus\tDescription")
		for _, info := range verifyingInfo {
			timestamp, err := ptypes.Timestamp(info.Metadata.AddedAt)
			if err != nil {
				return err
			}
			fmt.Fprintf(w, "  %v\t%v\t%v\t%v\n", info.Metadata.KeyId[:keyIDTruncatedLen], timestamp.Format(time.ANSIC), info.Status, info.Metadata.Description)
		}

		if err := w.Flush(); err != nil {
			return nil
		}
		return nil
	},
}

func readKeyStoreFile() error {
	store = keystore.New()
	// Authorized keys file might not exist.
	if _, err := os.Stat(keyStoreFile); err == nil {
		data, err := ioutil.ReadFile(keyStoreFile)
		if err != nil {
			return fmt.Errorf("reading keystore file failed: %v", err)
		}
		if err = keystore.Unmarshal(data, store); err != nil {
			return fmt.Errorf("keystore.Unmarshak() failed: %v", err)
		}
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("error checking if keystore file exists: %v", err)
	}
	return nil
}

func addPublicKey(pkFile, description string) error {
	if activate {
		return errors.New("--activate requires a private key")
	}

	pkPEM, err := ioutil.ReadFile(pkFile)
	if err != nil {
		return err
	}
	if _, err := store.AddVerifyingKey(description, pkPEM); err != nil {
		return err
	}
	return nil
}

func addPrivateKey(skFile, description string, activate bool) error {
	skPEM, err := ioutil.ReadFile(skFile)
	if err != nil {
		return err
	}
	status := kmpb.SigningKey_INACTIVE
	if activate {
		status = kmpb.SigningKey_ACTIVE
	}
	keyID, err := store.AddSigningKey(status, description, skPEM)
	if err != nil {
		return err
	}

	// Add the corresponding verifying key.
	signer, err := store.Signer(keyID)
	if err != nil {
		return err
	}
	pkPEM, err := signer.PublicKeyPEM()
	if err != nil {
		return err
	}
	if _, err := store.AddVerifyingKey(description, pkPEM); err != nil {
		return err
	}
	return nil
}

func keyID(hint string) string {
	ids := store.KeyIDs()
	for _, id := range ids {
		if id[:len(hint)] == hint {
			return id
		}
	}
	return ""
}

func init() {
	RootCmd.AddCommand(keysCmd)
	keysCmd.AddCommand(addCmd)
	keysCmd.AddCommand(removeCmd)
	keysCmd.AddCommand(activateCmd)
	keysCmd.AddCommand(listCmd)

	addCmd.PersistentFlags().StringVar(&pubKey, "pubkey", "", "Path to a public key file")
	addCmd.PersistentFlags().StringVar(&privKey, "privkey", "", "Path to a private key file")
	addCmd.PersistentFlags().StringVar(&description, "description", "", "Description of the added authorized key")
	addCmd.PersistentFlags().BoolVar(&activate, "activate", false, "(Optional) Activate the added signing key")
}
<commit_msg>Add authorized-keys generate command in client binary<commit_after>// Copyright 2016 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 cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"text/tabwriter"
	"time"

	"github.com/google/key-transparency/core/keystore"
	"github.com/google/key-transparency/core/signatures/p256"

	"github.com/golang/protobuf/ptypes"
	"github.com/spf13/cobra"

	kmpb "github.com/google/key-transparency/core/proto/keymaster"
)

const (
	keyStoreFile      = ".keystore"
	keyIDTruncatedLen = 8
)

var (
	pubKey      string
	privKey     string
	description string
	activate    bool
	keyType     string
	generate    bool
)

var store *keystore.KeyStore

// keysCmd represents the authorized-keys command.
var keysCmd = &cobra.Command{
	Use:   "authorized-keys",
	Short: "Manage authorized keys",
	Long: `Manage keys authorized to sign updates, and select the active signing key.

Verifying always happens using the keys listed in the previous epoch.`,
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		if err := readKeyStoreFile(); err != nil {
			log.Fatal(err)
		}
	},
	PersistentPostRunE: func(cmd *cobra.Command, args []string) error {
		buf, err := store.Marshal()
		if err != nil {
			return err
		}
		if err := ioutil.WriteFile(keyStoreFile, buf, 0600); err != nil {
			return err
		}
		return nil
	},
}

// addCmd represents the authorized-keys add command.
var addCmd = &cobra.Command{
	Use:   "add [ --privkey=[path] --activate | --pubkey=[path] | --generate --type=[key_type] --activate ] --description=[comment]",
	Short: "Add a key to the list of authorized keys",
	Long: `Provide a pair of public and private keys, already existing on disk, to be added to the list of authorized keys. The --generate flag can be used to generate a random key pair.

./key-transparency-client authorized-keys add --pubkey=/path/to/PEM/pubkey --description=[comment]
./key-transparency-client authorized-keys add --privkey=/path/to/PEM/privkey --activate --description=[comment]
./key-transparency-client authorized-keys add --generate --type=[key_type] --activate --description=[comment]
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if generate {
			if pubKey != "" || privKey != "" {
				return fmt.Errorf("cannot provide public or private key with generate")
			}
			if keyType == "" {
				return fmt.Errorf("must provide key type")
			}
		} else {
			if pubKey == "" && privKey == "" {
				return fmt.Errorf("should provide private or public key")
			}
			if pubKey != "" && privKey != "" {
				return fmt.Errorf("cannot provide public and private key at the same time")
			}
		}

		// Add either a private key, or a public key.
		switch {
		case generate:
			switch keyType {
			case "ecdsa":
				skPEM, _, err := p256.GeneratePEMs()
				if err != nil {
					return err
				}
				if err := addPrivateKey(skPEM, description, activate); err == nil {
					return err
				}
			default:
				return fmt.Errorf("unrecognized key type %v", keyType)
			}
		case privKey != "":
			skPEM, err := ioutil.ReadFile(privKey)
			if err != nil {
				return err
			}
			if err := addPrivateKey(skPEM, description, activate); err == nil {
				return err
			}
		case pubKey != "":
			pkPEM, err := ioutil.ReadFile(pubKey)
			if err != nil {
				return err
			}
			if err := addPublicKey(pkPEM, description); err == nil {
				return err
			}
		}
		return nil
	},
}

// removeCmd represents the authorized-keys remove command.
var removeCmd = &cobra.Command{
	Use:   "remove [keyid]",
	Short: "Remove a key from the list of authorized keys",
	Long: `Remove a key based on its key ID from the list of authorized keys. e.g.:

./key-transparency-client authorized-keys remove [keyid]

If the list contains a single key, it cannot be removed.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if len(args) < 1 {
			return fmt.Errorf("key ID needs to be provided")
		}

		// First remove verifying key with matching ID then attempt to
		// remove signing key if exist.
		keyID := keyID(args[0])
		if err := store.RemoveVerifyingKey(keyID); err != nil {
			return err
		}
		if err := store.RemoveSigningKey(keyID); err != nil && err != keystore.ErrKeyNotExist {
			return err
		}
		return nil
	},
}

// activateCmd represents the authorized-keys activate command.
var activateCmd = &cobra.Command{
	Use:   "activate [keyid]",
	Short: "Activate a key in the list of authorized keys",
	Long: `Activate a key based on its key ID in the list of authorized keys. e.g.:

./key-transparency-client authorized-keys activate [keyid]
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		// Validate input.
		if len(args) < 1 {
			return fmt.Errorf("key ID needs to be provided")
		}
		keyID := keyID(args[0])
		if err := store.Activate(keyID); err != nil {
			return err
		}
		return nil
	},
}

// listCmd represents the authorized-keys list command.
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List all authorized keys",
	Long: `List metadata about all authorized keys. e.g.:

./key-transparency-client authorized-keys list

The actual keys are not listed, only their corresponding metadata.
`,
	RunE: func(cmd *cobra.Command, args []string) error {
		signingInfo, verifyingInfo, err := store.Info()
		if err != nil {
			return err
		}

		w := tabwriter.NewWriter(os.Stdout, 0, 0, 1, ' ', tabwriter.Debug)

		// List signing keys.
		fmt.Fprintln(w, "Signing Keys:")
		fmt.Fprintln(w, "  ID\tAdded At\tStatus\tDescription")
		for _, info := range signingInfo {
			timestamp, err := ptypes.Timestamp(info.Metadata.AddedAt)
			if err != nil {
				return err
			}
			fmt.Fprintf(w, "  %v\t%v\t%v\t%v\n", info.Metadata.KeyId[:keyIDTruncatedLen], timestamp.Format(time.ANSIC), info.Status, info.Metadata.Description)
		}

		// List verifying keys.
		fmt.Fprintln(w, "\nVerifying Keys:")
		fmt.Fprintln(w, "  ID\tAdded At\tStatus\tDescription")
		for _, info := range verifyingInfo {
			timestamp, err := ptypes.Timestamp(info.Metadata.AddedAt)
			if err != nil {
				return err
			}
			fmt.Fprintf(w, "  %v\t%v\t%v\t%v\n", info.Metadata.KeyId[:keyIDTruncatedLen], timestamp.Format(time.ANSIC), info.Status, info.Metadata.Description)
		}

		if err := w.Flush(); err != nil {
			return nil
		}
		return nil
	},
}

func readKeyStoreFile() error {
	store = keystore.New()
	// Authorized keys file might not exist.
	if _, err := os.Stat(keyStoreFile); err == nil {
		data, err := ioutil.ReadFile(keyStoreFile)
		if err != nil {
			return fmt.Errorf("reading keystore file failed: %v", err)
		}
		if err = keystore.Unmarshal(data, store); err != nil {
			return fmt.Errorf("keystore.Unmarshak() failed: %v", err)
		}
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("error checking if keystore file exists: %v", err)
	}
	return nil
}

func addPublicKey(pkPEM []byte, description string) error {
	if activate {
		return errors.New("--activate requires a private key")
	}
	if _, err := store.AddVerifyingKey(description, pkPEM); err != nil {
		return err
	}
	return nil
}

func addPrivateKey(skPEM []byte, description string, activate bool) error {
	status := kmpb.SigningKey_INACTIVE
	if activate {
		status = kmpb.SigningKey_ACTIVE
	}
	keyID, err := store.AddSigningKey(status, description, skPEM)
	if err != nil {
		return err
	}

	// Add the corresponding verifying key.
	signer, err := store.Signer(keyID)
	if err != nil {
		return err
	}
	pkPEM, err := signer.PublicKeyPEM()
	if err != nil {
		return err
	}
	if _, err := store.AddVerifyingKey(description, pkPEM); err != nil {
		return err
	}
	return nil
}

func keyID(hint string) string {
	ids := store.KeyIDs()
	for _, id := range ids {
		if id[:len(hint)] == hint {
			return id
		}
	}
	return ""
}

func init() {
	RootCmd.AddCommand(keysCmd)
	keysCmd.AddCommand(addCmd)
	keysCmd.AddCommand(removeCmd)
	keysCmd.AddCommand(activateCmd)
	keysCmd.AddCommand(listCmd)

	addCmd.PersistentFlags().StringVar(&pubKey, "pubkey", "", "Path to a public key file")
	addCmd.PersistentFlags().StringVar(&privKey, "privkey", "", "Path to a private key file")
	addCmd.PersistentFlags().StringVar(&description, "description", "", "(Optional) Description of the added authorized key")
	addCmd.PersistentFlags().BoolVar(&activate, "activate", false, "(Optional) Activate the added signing key")
	addCmd.PersistentFlags().BoolVar(&generate, "generate", false, "Generate a random public and private key pair")
	addCmd.PersistentFlags().StringVar(&keyType, "type", "", "The key type to be generated, e.g., ecdsa")
}
<|endoftext|>


<commit_before>package gc

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/metric"
	"github.com/concourse/concourse/atc/worker"
	multierror "github.com/hashicorp/go-multierror"
)

const HijackedContainerTimeout = 5 * time.Minute

type containerCollector struct {
	containerRepository         db.ContainerRepository
	jobRunner                   WorkerJobRunner
	missingContainerGracePeriod time.Duration
}

func NewContainerCollector(
	containerRepository db.ContainerRepository,
	jobRunner WorkerJobRunner,
	missingContainerGracePeriod time.Duration,
) Collector {
	return &containerCollector{
		containerRepository:         containerRepository,
		jobRunner:                   jobRunner,
		missingContainerGracePeriod: missingContainerGracePeriod,
	}
}

type job struct {
	JobName string
	RunFunc func(worker.Worker)
}

func (j *job) Name() string {
	return j.JobName
}

func (j *job) Run(w worker.Worker) {
	j.RunFunc(w)
}

func (c *containerCollector) Run(ctx context.Context) error {
	logger := lagerctx.FromContext(ctx).Session("container-collector")

	logger.Debug("start")
	defer logger.Debug("done")

	var errs error

	err := c.cleanupOrphanedContainers(logger.Session("orphaned-containers"))
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-orphaned-containers", err)
	}

	err = c.cleanupFailedContainers(logger.Session("failed-containers"))
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-failed-containers", err)
	}

	_, err = c.containerRepository.RemoveMissingContainers(c.missingContainerGracePeriod)
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-missing-containers", err)
	}

	return errs
}

func (c *containerCollector) cleanupFailedContainers(logger lager.Logger) error {
	failedContainersLen, err := c.containerRepository.DestroyFailedContainers()
	if err != nil {
		logger.Error("failed-to-find-failed-containers-for-deletion", err)
		return err
	}

	if failedContainersLen > 0 {
		logger.Debug("found-failed-containers-for-deletion", lager.Data{
			"number": failedContainersLen,
		})
	}

	metric.FailedContainersToBeGarbageCollected{
		Containers: failedContainersLen,
	}.Emit(logger)

	return nil
}

func (c *containerCollector) cleanupOrphanedContainers(logger lager.Logger) error {
	creatingContainers, createdContainers, destroyingContainers, err := c.containerRepository.FindOrphanedContainers()

	if err != nil {
		logger.Error("failed-to-get-orphaned-containers-for-deletion", err)
		return err
	}

	if len(creatingContainers) > 0 || len(createdContainers) > 0 || len(destroyingContainers) > 0 {
		logger.Debug("found-orphaned-containers-for-deletion", lager.Data{
			"creating-containers-num":   len(creatingContainers),
			"created-containers-num":    len(createdContainers),
			"destroying-containers-num": len(destroyingContainers),
		})
	}

	metric.CreatingContainersToBeGarbageCollected{
		Containers: len(creatingContainers),
	}.Emit(logger)

	metric.CreatedContainersToBeGarbageCollected{
		Containers: len(createdContainers),
	}.Emit(logger)

	metric.DestroyingContainersToBeGarbageCollected{
		Containers: len(destroyingContainers),
	}.Emit(logger)

	var workerCreatedContainers = make(map[string][]db.CreatedContainer)

	for _, createdContainer := range createdContainers {
		containers, ok := workerCreatedContainers[createdContainer.WorkerName()]
		if ok {
			// update existing array
			containers = append(containers, createdContainer)
			workerCreatedContainers[createdContainer.WorkerName()] = containers
		} else {
			// create new array
			workerCreatedContainers[createdContainer.WorkerName()] = []db.CreatedContainer{createdContainer}
		}
		logger.Debug("found-created-containers-for-deletion", lager.Data{
			"handle": createdContainer.Handle(),
			"worker": createdContainer.WorkerName(),
		})
	}

	for worker, createdContainers := range workerCreatedContainers {
		go destroyNonHijackedCreatedContainers(logger, createdContainers)

		// prevent closure from capturing last value of loop
		c.jobRunner.Try(logger,
			worker,
			&job{
				JobName: fmt.Sprintf("destroy-hijacked-containers"),
				RunFunc: destroyHijackedCreatedContainers(logger, createdContainers),
			},
		)
	}

	return nil
}

func destroyNonHijackedCreatedContainers(logger lager.Logger, containers []db.CreatedContainer) {
	cLog := logger.Session("mark-created-as-destroying")

	for _, container := range containers {
		if container.IsHijacked() {
			continue
		}

		_, err := container.Destroying()
		if err != nil {
			cLog.Error("failed-to-transition", err, lager.Data{
				"container": container.Handle(),
			})
			return
		}
	}
}

func destroyHijackedCreatedContainers(logger lager.Logger, containers []db.CreatedContainer) func(worker.Worker) {
	return func(workerClient worker.Worker) {
		cLog := logger.Session("mark-hijacked-container", lager.Data{
			"worker": workerClient.Name(),
		})

		for _, container := range containers {
			if !container.IsHijacked() {
				continue
			}

			_, err := markHijackedContainerAsDestroying(cLog, container, workerClient.GardenClient())
			if err != nil {
				cLog.Error("failed-to-transition", err, lager.Data{
					"container": container.Handle(),
				})
				return
			}
		}
	}
}

func markHijackedContainerAsDestroying(
	logger lager.Logger,
	hijackedContainer db.CreatedContainer,
	gardenClient garden.Client,
) (db.DestroyingContainer, error) {

	gardenContainer, found, err := findContainer(gardenClient, hijackedContainer.Handle())
	if err != nil {
		logger.Error("failed-to-lookup-container-in-garden", err)
		return nil, err
	}

	if !found {
		var destroyingContainer db.DestroyingContainer
		logger.Debug("hijacked-container-not-found-in-garden")

		destroyingContainer, err = hijackedContainer.Destroying()
		if err != nil {
			logger.Error("failed-to-mark-container-as-destroying", err)
			return nil, err
		}
		return destroyingContainer, nil
	}

	err = gardenContainer.SetGraceTime(HijackedContainerTimeout)
	if err != nil {
		logger.Error("failed-to-set-grace-time-on-hijacked-container", err)
		return nil, err
	}

	_, err = hijackedContainer.Discontinue()
	if err != nil {
		logger.Error("failed-to-mark-container-as-destroying", err)
		return nil, err
	}

	return nil, nil
}

func findContainer(gardenClient garden.Client, handle string) (garden.Container, bool, error) {
	gardenContainer, err := gardenClient.Lookup(handle)
	if err != nil {
		if _, ok := err.(garden.ContainerNotFoundError); ok {
			return nil, false, nil
		}
		return nil, false, err
	}
	return gardenContainer, true, nil
}
<commit_msg>atc: reduce logging for deleting containers<commit_after>package gc

import (
	"context"
	"fmt"
	"time"

	"code.cloudfoundry.org/garden"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagerctx"
	"github.com/concourse/concourse/atc/db"
	"github.com/concourse/concourse/atc/metric"
	"github.com/concourse/concourse/atc/worker"
	multierror "github.com/hashicorp/go-multierror"
)

const HijackedContainerTimeout = 5 * time.Minute

type containerCollector struct {
	containerRepository         db.ContainerRepository
	jobRunner                   WorkerJobRunner
	missingContainerGracePeriod time.Duration
}

func NewContainerCollector(
	containerRepository db.ContainerRepository,
	jobRunner WorkerJobRunner,
	missingContainerGracePeriod time.Duration,
) Collector {
	return &containerCollector{
		containerRepository:         containerRepository,
		jobRunner:                   jobRunner,
		missingContainerGracePeriod: missingContainerGracePeriod,
	}
}

type job struct {
	JobName string
	RunFunc func(worker.Worker)
}

func (j *job) Name() string {
	return j.JobName
}

func (j *job) Run(w worker.Worker) {
	j.RunFunc(w)
}

func (c *containerCollector) Run(ctx context.Context) error {
	logger := lagerctx.FromContext(ctx).Session("container-collector")

	logger.Debug("start")
	defer logger.Debug("done")

	var errs error

	err := c.cleanupOrphanedContainers(logger.Session("orphaned-containers"))
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-orphaned-containers", err)
	}

	err = c.cleanupFailedContainers(logger.Session("failed-containers"))
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-failed-containers", err)
	}

	_, err = c.containerRepository.RemoveMissingContainers(c.missingContainerGracePeriod)
	if err != nil {
		errs = multierror.Append(errs, err)
		logger.Error("failed-to-clean-up-missing-containers", err)
	}

	return errs
}

func (c *containerCollector) cleanupFailedContainers(logger lager.Logger) error {
	failedContainersLen, err := c.containerRepository.DestroyFailedContainers()
	if err != nil {
		logger.Error("failed-to-find-failed-containers-for-deletion", err)
		return err
	}

	if failedContainersLen > 0 {
		logger.Debug("found-failed-containers-for-deletion", lager.Data{
			"number": failedContainersLen,
		})
	}

	metric.FailedContainersToBeGarbageCollected{
		Containers: failedContainersLen,
	}.Emit(logger)

	return nil
}

func (c *containerCollector) cleanupOrphanedContainers(logger lager.Logger) error {
	creatingContainers, createdContainers, destroyingContainers, err := c.containerRepository.FindOrphanedContainers()

	if err != nil {
		logger.Error("failed-to-get-orphaned-containers-for-deletion", err)
		return err
	}

	if len(creatingContainers) > 0 || len(createdContainers) > 0 || len(destroyingContainers) > 0 {
		logger.Debug("found-orphaned-containers-for-deletion", lager.Data{
			"creating-containers-num":   len(creatingContainers),
			"created-containers-num":    len(createdContainers),
			"destroying-containers-num": len(destroyingContainers),
		})
	}

	metric.CreatingContainersToBeGarbageCollected{
		Containers: len(creatingContainers),
	}.Emit(logger)

	metric.CreatedContainersToBeGarbageCollected{
		Containers: len(createdContainers),
	}.Emit(logger)

	metric.DestroyingContainersToBeGarbageCollected{
		Containers: len(destroyingContainers),
	}.Emit(logger)

	var workerCreatedContainers = make(map[string][]db.CreatedContainer)

	for _, createdContainer := range createdContainers {
		containers, ok := workerCreatedContainers[createdContainer.WorkerName()]
		if ok {
			// update existing array
			containers = append(containers, createdContainer)
			workerCreatedContainers[createdContainer.WorkerName()] = containers
		} else {
			// create new array
			workerCreatedContainers[createdContainer.WorkerName()] = []db.CreatedContainer{createdContainer}
		}
	}

	logger.Debug("found-created-containers-for-deletion", lager.Data{
		"num-containers": len(createdContainers),
	})

	for worker, createdContainers := range workerCreatedContainers {
		go destroyNonHijackedCreatedContainers(logger, createdContainers)

		// prevent closure from capturing last value of loop
		c.jobRunner.Try(logger,
			worker,
			&job{
				JobName: fmt.Sprintf("destroy-hijacked-containers"),
				RunFunc: destroyHijackedCreatedContainers(logger, createdContainers),
			},
		)
	}

	return nil
}

func destroyNonHijackedCreatedContainers(logger lager.Logger, containers []db.CreatedContainer) {
	cLog := logger.Session("mark-created-as-destroying")

	for _, container := range containers {
		if container.IsHijacked() {
			continue
		}

		_, err := container.Destroying()
		if err != nil {
			cLog.Error("failed-to-transition", err, lager.Data{
				"container": container.Handle(),
			})
			return
		}
	}
}

func destroyHijackedCreatedContainers(logger lager.Logger, containers []db.CreatedContainer) func(worker.Worker) {
	return func(workerClient worker.Worker) {
		cLog := logger.Session("mark-hijacked-container", lager.Data{
			"worker": workerClient.Name(),
		})

		for _, container := range containers {
			if !container.IsHijacked() {
				continue
			}

			_, err := markHijackedContainerAsDestroying(cLog, container, workerClient.GardenClient())
			if err != nil {
				cLog.Error("failed-to-transition", err, lager.Data{
					"container": container.Handle(),
				})
				return
			}
		}
	}
}

func markHijackedContainerAsDestroying(
	logger lager.Logger,
	hijackedContainer db.CreatedContainer,
	gardenClient garden.Client,
) (db.DestroyingContainer, error) {

	gardenContainer, found, err := findContainer(gardenClient, hijackedContainer.Handle())
	if err != nil {
		logger.Error("failed-to-lookup-container-in-garden", err)
		return nil, err
	}

	if !found {
		var destroyingContainer db.DestroyingContainer
		logger.Debug("hijacked-container-not-found-in-garden")

		destroyingContainer, err = hijackedContainer.Destroying()
		if err != nil {
			logger.Error("failed-to-mark-container-as-destroying", err)
			return nil, err
		}
		return destroyingContainer, nil
	}

	err = gardenContainer.SetGraceTime(HijackedContainerTimeout)
	if err != nil {
		logger.Error("failed-to-set-grace-time-on-hijacked-container", err)
		return nil, err
	}

	_, err = hijackedContainer.Discontinue()
	if err != nil {
		logger.Error("failed-to-mark-container-as-destroying", err)
		return nil, err
	}

	return nil, nil
}

func findContainer(gardenClient garden.Client, handle string) (garden.Container, bool, error) {
	gardenContainer, err := gardenClient.Lookup(handle)
	if err != nil {
		if _, ok := err.(garden.ContainerNotFoundError); ok {
			return nil, false, nil
		}
		return nil, false, err
	}
	return gardenContainer, true, nil
}
<|endoftext|>