<commit_before>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code,string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code,string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3,string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name,string,omitempty"`   // Name of the country
	Region         string  `json:"region,string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code,integer,omitempty"`    // Metro code within the country
	City           string  `json:"city,string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code,integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code,string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude,float"`                  // Latitude
	Longitude      float64 `json:"longitude,float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	normalized := make([]string, len(ips))
	for i := range ips {
		if ips[i] == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		var err error
		normalized[i], err = web100.NormalizeIPv6(ips[i])
		if err != nil {
			log.Println(err)
		}
		reqData = append(reqData, RequestData{normalized[i], 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for i := range normalized {
		data, ok := annotationData[normalized[i]+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationErrorCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[i] = *data.Geo
	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<commit_msg>PR fix<commit_after>package annotation

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"time"

	"github.com/m-lab/etl/web100"

	"github.com/m-lab/etl/metrics"
	"github.com/prometheus/client_golang/prometheus"
)

var IPAnnotationEnabled = false

func init() {
	getFlagValues()
}

func getFlagValues() {
	// Check for ANNOTATE_IP = 'true'
	flag, ok := os.LookupEnv("ANNOTATE_IP")
	if ok {
		IPAnnotationEnabled, _ = strconv.ParseBool(flag)
		// If parse fails, then ipAnn will be set to false.
	}
}

// EnableAnnotation is used only for testing.  Should be placed in whitebox _test file.
func EnableAnnotation() {
	os.Setenv("ANNOTATE_IP", "True")
	getFlagValues()
}

// The GeolocationIP struct contains all the information needed for the
// geolocation data that will be inserted into big query. The fiels are
// capitalized for exporting, although the originals in the DB schema
// are not.
type GeolocationIP struct {
	Continent_code string  `json:"continent_code,string,omitempty"` // Gives a shorthand for the continent
	Country_code   string  `json:"country_code,string,omitempty"`   // Gives a shorthand for the country
	Country_code3  string  `json:"country_code3,string,omitempty"`  // Gives a shorthand for the country
	Country_name   string  `json:"country_name,string,omitempty"`   // Name of the country
	Region         string  `json:"region,string,omitempty"`         // Region or State within the country
	Metro_code     int64   `json:"metro_code,integer,omitempty"`    // Metro code within the country
	City           string  `json:"city,string,omitempty"`           // City within the region
	Area_code      int64   `json:"area_code,integer,omitempty"`     // Area code, similar to metro code
	Postal_code    string  `json:"postal_code,string,omitempty"`    // Postal code, again similar to metro
	Latitude       float64 `json:"latitude,float"`                  // Latitude
	Longitude      float64 `json:"longitude,float"`                 // Longitude

}

// The struct that will hold the IP/ASN data when it gets added to the
// schema. Currently empty and unused.
type IPASNData struct{}

// The main struct for the geo metadata, which holds pointers to the
// Geolocation data and the IP/ASN data. This is what we parse the JSON
// response from the annotator into.
type GeoData struct {
	Geo *GeolocationIP // Holds the geolocation data
	ASN *IPASNData     // Holds the IP/ASN data
}

// The RequestData schema is the schema for the json that we will send
// down the pipe to the annotation service.
type RequestData struct {
	IP        string    // Holds the IP from an incoming request
	IPFormat  int       // Holds the ip format, 4 or 6
	Timestamp time.Time // Holds the timestamp from an incoming request
}

// AnnotatorURL holds the https address of the annotator.
// TODO(gfr) See if there is a better way of determining
// where to send the request (there almost certainly is)
var AnnotatorURL = "https://annotator-dot-" +
	os.Getenv("GCLOUD_PROJECT") +
	".appspot.com"

// BaseURL provides the base URL for single annotation requests
var BaseURL = AnnotatorURL + "/annotate?"

// BatchURL provides the base URL for batch annotation requests
var BatchURL = AnnotatorURL + "/batch_annotate"

// FetchGeoAnnotations takes a slice of strings
// containing ip addresses, a timestamp, and a slice of pointers to
// the GeolocationIP structs that correspond to the ip addresses. A
// precondition assumed by this function is that both slices are the
// same length. It will then make a call to the batch annotator, using
// the ip addresses and the timestamp. Then, it uses that data to fill
// in the structs pointed to by the slice of GeolocationIP pointers.
func FetchGeoAnnotations(ips []string, timestamp time.Time, geoDest []*GeolocationIP) {
	reqData := make([]RequestData, 0, len(ips))
	normalized := make([]string, len(ips))
	for i := range ips {
		if ips[i] == "" {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationWarningCount.With(prometheus.
				Labels{"source": "Empty IP Address!!!"}).Inc()
			continue
		}
		var err error
		normalized[i], err = web100.NormalizeIPv6(ips[i])
		if err != nil {
			log.Println(err)
			metrics.AnnotationWarningCount.With(prometheus.
				Labels{"source": "NormalizeIPv6 Error"}).Inc()
		}
		reqData = append(reqData, RequestData{normalized[i], 0, timestamp})
	}
	annotationData := GetBatchGeoData(BatchURL, reqData)
	timeString := strconv.FormatInt(timestamp.Unix(), 36)
	for i := range normalized {
		data, ok := annotationData[normalized[i]+timeString]
		if !ok || data.Geo == nil {
			// TODO(gfr) These should be warning, else we have error > request
			metrics.AnnotationWarningCount.With(prometheus.
				Labels{"source": "Missing or empty data for IP Address!!!"}).Inc()
			continue
		}
		*geoDest[i] = *data.Geo
	}
}

// GetAndInsertGeolocationIPStruct takes a NON-NIL pointer to a
// pre-allocated GeolocationIP struct, an IP address, and a
// timestamp. It will connect to the annotation service, get the
// geo data, and insert the geo data into the reigion pointed to by
// the GeolocationIP pointer.
func GetAndInsertGeolocationIPStruct(geo *GeolocationIP, ip string, timestamp time.Time) {
	url := BaseURL + "ip_addr=" + url.QueryEscape(ip) +
		"&since_epoch=" + strconv.FormatInt(timestamp.Unix(), 10)
	annotationData := GetGeoData(url)
	if annotationData != nil && annotationData.Geo != nil {
		*geo = *annotationData.Geo
	}
}

// GetGeoData combines the functionality of QueryAnnotationService and
// ParseJSONGeoDataResponse to query the annotator service and return
// the corresponding GeoData if it can, or a nil pointer if it
// encounters any error and cannot get the data for any reason
func GetGeoData(url string) *GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := QueryAnnotationService(url)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := ParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		return nil
	}
	return geoDataFromResponse
}

// QueryAnnotationService will connect to the annotation service and
// copy the body of a valid response to a byte slice and return it to a
// user, returning an error if any occurs
func QueryAnnotationService(url string) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()
	// Make the actual request
	resp, err := http.Get(url)

	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Request to Annotator failed"}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Webserver gave non-ok response"}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// ParseJSONGeoDataResponse takes a byte slice containing the test of
// the JSON from the annotator service and parses it into a GeoData
// struct, for easy manipulation. It returns a pointer to the struct on
// success and an error if an error occurs.
func ParseJSONGeoDataResponse(jsonBuffer []byte) (*GeoData, error) {
	parsedJSON := &GeoData{}
	err := json.Unmarshal(jsonBuffer, parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}

// GetBatchGeoData combines the functionality of
// BatchQueryAnnotationService and BatchParseJSONGeoDataResponse to
// query the annotator service and return the corresponding map of
// ip-timestamp strings to GeoData structs, or a nil map if it
// encounters any error and cannot get the data for any reason
// TODO - dedup common code in GetGeoData
func GetBatchGeoData(url string, data []RequestData) map[string]GeoData {
	// Query the service and grab the response safely
	annotatorResponse, err := BatchQueryAnnotationService(url, data)
	if err != nil {
		log.Println(err)
		return nil
	}

	// Safely parse the JSON response and pass it back to the caller
	geoDataFromResponse, err := BatchParseJSONGeoDataResponse(annotatorResponse)
	if err != nil {
		metrics.AnnotationErrorCount.With(prometheus.
			Labels{"source": "Failed to parse JSON"}).Inc()
		log.Println(err)
		log.Printf("%+v\n", data)
		return nil
	}
	return geoDataFromResponse
}

// BatchQueryAnnotationService takes a url to POST the request to and
// a slice of RequestDatas to be sent in the body in a JSON
// format. It will copy the response into a []byte and return it to
// the user, returning an error if any occurs
// TODO(gfr) Should pass the annotator's request context through and use it here.
func BatchQueryAnnotationService(url string, data []RequestData) ([]byte, error) {
	metrics.AnnotationRequestCount.Inc()

	encodedData, err := json.Marshal(data)
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": "Couldn't Marshal Data"}).Inc()
		return nil, err
	}

	var netClient = &http.Client{
		// Median response time is < 10 msec, but 99th percentile is 0.6 seconds.
		Timeout: 2 * time.Second,
	}

	// Make the actual request
	resp, err := netClient.Post(url, "raw", bytes.NewReader(encodedData))
	// Catch http errors
	if err != nil {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": err.Error()}).Inc()
		return nil, err
	}

	// Catch errors reported by the service
	if resp.StatusCode != http.StatusOK {
		metrics.AnnotationErrorCount.
			With(prometheus.Labels{"source": http.StatusText(resp.StatusCode)}).Inc()
		return nil, errors.New("URL:" + url + " gave response code " + resp.Status)
	}
	defer resp.Body.Close()

	// Copy response into a byte slice
	return ioutil.ReadAll(resp.Body)
}

// BatchParseJSONGeoDataResponse takes a byte slice containing the
// text of the JSON from the annoator service's batch request endpoint
// and parses it into a map of strings to GeoData structs, for
// easy manipulation. It returns a pointer to the struct on success
// and an error if one occurs.
// TODO - is there duplicate code with ParseJSON... ?
func BatchParseJSONGeoDataResponse(jsonBuffer []byte) (map[string]GeoData, error) {
	parsedJSON := make(map[string]GeoData)
	err := json.Unmarshal(jsonBuffer, &parsedJSON)
	if err != nil {
		return nil, err
	}
	return parsedJSON, nil
}
<|endoftext|>


<commit_before>package api

import (
	"fmt"
	"github.com/nais/naisd/api/app"
	v1 "k8s.io/api/core/v1"
	k8sextensions "k8s.io/api/extensions/v1beta1"
	k8smeta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"
)

const (
	defaultRedisPort          = 6379
	defaultRedisExporterPort  = 9121
	defaultRedisExporterImage = "oliver006/redis_exporter:v1.0.3-alpine"
	defaultRedisImage         = "redis:5-alpine"
)

type Redis struct {
	Enabled  bool
	Image    string
	Limits   ResourceList
	Requests ResourceList
}

func updateDefaultRedisValues(redis Redis) Redis {
	if redis.Image == "" {
		redis.Image = defaultRedisImage
	}
	if len(redis.Limits.Cpu) == 0 {
		redis.Limits.Cpu = "100m"
	}
	if len(redis.Limits.Memory) == 0 {
		redis.Limits.Memory = "128Mi"
	}
	if len(redis.Requests.Cpu) == 0 {
		redis.Requests.Cpu = "100m"
	}
	if len(redis.Requests.Memory) == 0 {
		redis.Requests.Memory = "128Mi"
	}
	return redis
}

func createRedisPodSpec(redis Redis) v1.PodSpec {
	return v1.PodSpec{
		Containers: []v1.Container{
			{
				Name:  "redis",
				Image: redis.Image,
				Resources: createResourceLimits(redis.Requests.Cpu, redis.Requests.Memory,
					redis.Limits.Cpu, redis.Limits.Memory),
				ImagePullPolicy: v1.PullIfNotPresent,
				Ports: []v1.ContainerPort{
					{
						ContainerPort: int32(defaultRedisPort),
						Name:          DefaultPortName,
						Protocol:      v1.ProtocolTCP,
					},
				},
			},
			{
				Name:  "exporter",
				Image: defaultRedisExporterImage,
				Resources: createResourceLimits("100m", "100Mi",
					"100m", "100Mi"),
				ImagePullPolicy: v1.PullIfNotPresent,
				Ports: []v1.ContainerPort{
					{
						ContainerPort: int32(defaultRedisExporterPort),
						Name:          DefaultPortName,
						Protocol:      v1.ProtocolTCP,
					},
				},
			},
		},
	}
}

func createRedisDeploymentSpec(resourceName string, spec app.Spec, redis Redis) k8sextensions.DeploymentSpec {
	objectMeta := generateObjectMeta(spec)
	objectMeta.Name = resourceName
	objectMeta.Annotations = map[string]string{
		"prometheus.io/scrape": "true",
		"prometheus.io/port":   string(defaultRedisExporterPort),
		"prometheus.io/path":   "/metrics",
	}

	return k8sextensions.DeploymentSpec{
		Replicas: int32p(1),
		Selector: &k8smeta.LabelSelector{
			MatchLabels: createPodSelector(spec),
		},
		Strategy: k8sextensions.DeploymentStrategy{
			Type: k8sextensions.RecreateDeploymentStrategyType,
		},
		ProgressDeadlineSeconds: int32p(300),
		RevisionHistoryLimit:    int32p(10),
		Template: v1.PodTemplateSpec{
			ObjectMeta: objectMeta,
			Spec:       createRedisPodSpec(redis),
		},
	}
}

func createRedisDeploymentDef(resourceName string, spec app.Spec, redis Redis, existingDeployment *k8sextensions.Deployment) *k8sextensions.Deployment {
	existingDeployment.Spec = createRedisDeploymentSpec(resourceName, spec, redis)
	return existingDeployment
}

func createOrUpdateRedisInstance(spec app.Spec, redis Redis, k8sClient kubernetes.Interface) (*k8sextensions.Deployment, error) {
	redisName := fmt.Sprintf("%s-redis", spec.ResourceName())
	existingDeployment, err := getExistingDeployment(redisName, spec.Namespace, k8sClient)

	if err != nil {
		return nil, fmt.Errorf("unable to get existing deployment: %s", err)
	}

	deploymentDef := createRedisDeploymentDef(redisName, spec, redis, existingDeployment)

	return createOrUpdateDeploymentResource(deploymentDef, spec.Namespace, k8sClient)
}

func createRedisServiceDef(spec app.Spec) *v1.Service {
	return &v1.Service{
		TypeMeta: k8smeta.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: generateObjectMeta(spec),
		Spec: v1.ServiceSpec{
			Type:     v1.ServiceTypeClusterIP,
			Selector: createPodSelector(spec),
			Ports: []v1.ServicePort{
				{
					Name:     DefaultPortName,
					Protocol: v1.ProtocolTCP,
					Port:     6379,
					TargetPort: intstr.IntOrString{
						Type:   intstr.String,
						StrVal: DefaultPortName,
					},
				},
			},
		},
	}
}

func createOrUpdateRedisService(spec app.Spec, k8sClient kubernetes.Interface) (*v1.Service, error) {
	redisName := fmt.Sprintf("%s-redis", spec.ResourceName())
	service, err := getExistingService(redisName, spec.Namespace, k8sClient)

	if err != nil {
		return nil, fmt.Errorf("unable to get existing service: %s", err)
	} else if service == nil {
		service = createRedisServiceDef(spec)
		service.Name = redisName
	}

	service.ObjectMeta = addLabelsToObjectMeta(service.ObjectMeta, spec)
	return createOrUpdateServiceResource(service, spec.Namespace, k8sClient)
}
<commit_msg>Riktig navn på deployment<commit_after>package api

import (
	"fmt"
	"github.com/nais/naisd/api/app"
	v1 "k8s.io/api/core/v1"
	k8sextensions "k8s.io/api/extensions/v1beta1"
	k8smeta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"
)

const (
	defaultRedisPort          = 6379
	defaultRedisExporterPort  = 9121
	defaultRedisExporterImage = "oliver006/redis_exporter:v1.0.3-alpine"
	defaultRedisImage         = "redis:5-alpine"
)

type Redis struct {
	Enabled  bool
	Image    string
	Limits   ResourceList
	Requests ResourceList
}

func updateDefaultRedisValues(redis Redis) Redis {
	if redis.Image == "" {
		redis.Image = defaultRedisImage
	}
	if len(redis.Limits.Cpu) == 0 {
		redis.Limits.Cpu = "100m"
	}
	if len(redis.Limits.Memory) == 0 {
		redis.Limits.Memory = "128Mi"
	}
	if len(redis.Requests.Cpu) == 0 {
		redis.Requests.Cpu = "100m"
	}
	if len(redis.Requests.Memory) == 0 {
		redis.Requests.Memory = "128Mi"
	}
	return redis
}

func createRedisPodSpec(redis Redis) v1.PodSpec {
	return v1.PodSpec{
		Containers: []v1.Container{
			{
				Name:  "redis",
				Image: redis.Image,
				Resources: createResourceLimits(redis.Requests.Cpu, redis.Requests.Memory,
					redis.Limits.Cpu, redis.Limits.Memory),
				ImagePullPolicy: v1.PullIfNotPresent,
				Ports: []v1.ContainerPort{
					{
						ContainerPort: int32(defaultRedisPort),
						Name:          DefaultPortName,
						Protocol:      v1.ProtocolTCP,
					},
				},
			},
			{
				Name:  "exporter",
				Image: defaultRedisExporterImage,
				Resources: createResourceLimits("100m", "100Mi",
					"100m", "100Mi"),
				ImagePullPolicy: v1.PullIfNotPresent,
				Ports: []v1.ContainerPort{
					{
						ContainerPort: int32(defaultRedisExporterPort),
						Name:          DefaultPortName,
						Protocol:      v1.ProtocolTCP,
					},
				},
			},
		},
	}
}

func createRedisDeploymentSpec(resourceName string, spec app.Spec, redis Redis) k8sextensions.DeploymentSpec {
	objectMeta := generateObjectMeta(spec)
	objectMeta.Name = resourceName
	objectMeta.Annotations = map[string]string{
		"prometheus.io/scrape": "true",
		"prometheus.io/port":   string(defaultRedisExporterPort),
		"prometheus.io/path":   "/metrics",
	}

	return k8sextensions.DeploymentSpec{
		Replicas: int32p(1),
		Selector: &k8smeta.LabelSelector{
			MatchLabels: createPodSelector(spec),
		},
		Strategy: k8sextensions.DeploymentStrategy{
			Type: k8sextensions.RecreateDeploymentStrategyType,
		},
		ProgressDeadlineSeconds: int32p(300),
		RevisionHistoryLimit:    int32p(10),
		Template: v1.PodTemplateSpec{
			ObjectMeta: objectMeta,
			Spec:       createRedisPodSpec(redis),
		},
	}
}

func createRedisDeploymentDef(resourceName string, spec app.Spec, redis Redis, existingDeployment *k8sextensions.Deployment) *k8sextensions.Deployment {
	existingDeployment.Spec = createRedisDeploymentSpec(resourceName, spec, redis)
	return existingDeployment
}

func createOrUpdateRedisInstance(spec app.Spec, redis Redis, k8sClient kubernetes.Interface) (*k8sextensions.Deployment, error) {
	redisName := fmt.Sprintf("%s-redis", spec.ResourceName())
	existingDeployment, err := getExistingDeployment(redisName, spec.Namespace, k8sClient)

	if err != nil {
		return nil, fmt.Errorf("unable to get existing deployment: %s", err)
	}

	deploymentDef := createRedisDeploymentDef(redisName, spec, redis, existingDeployment)
	deploymentDef.Name = fmt.Sprintf("%s-redis", spec.ResourceName())

	return createOrUpdateDeploymentResource(deploymentDef, spec.Namespace, k8sClient)
}

func createRedisServiceDef(spec app.Spec) *v1.Service {
	return &v1.Service{
		TypeMeta: k8smeta.TypeMeta{
			Kind:       "Service",
			APIVersion: "v1",
		},
		ObjectMeta: generateObjectMeta(spec),
		Spec: v1.ServiceSpec{
			Type:     v1.ServiceTypeClusterIP,
			Selector: createPodSelector(spec),
			Ports: []v1.ServicePort{
				{
					Name:     DefaultPortName,
					Protocol: v1.ProtocolTCP,
					Port:     6379,
					TargetPort: intstr.IntOrString{
						Type:   intstr.String,
						StrVal: DefaultPortName,
					},
				},
			},
		},
	}
}

func createOrUpdateRedisService(spec app.Spec, k8sClient kubernetes.Interface) (*v1.Service, error) {
	redisName := fmt.Sprintf("%s-redis", spec.ResourceName())
	service, err := getExistingService(redisName, spec.Namespace, k8sClient)

	if err != nil {
		return nil, fmt.Errorf("unable to get existing service: %s", err)
	} else if service == nil {
		service = createRedisServiceDef(spec)
		service.Name = redisName
	}

	service.ObjectMeta = addLabelsToObjectMeta(service.ObjectMeta, spec)
	return createOrUpdateServiceResource(service, spec.Namespace, k8sClient)
}
<|endoftext|>


<commit_before>//go:generate swagger generate spec
// Package main GalAirCheck.
//
// the purpose of this application is to provide an Galera Check application
// that will show state of Galera Cluster in real Time
//
// Terms Of Service:
//
// there are no TOS at this moment, use at your own risk we take no responsibility
//
//     Schemes:
//     Host:
//     BasePath:
//     Version: 0.0.1
//     License: MIT http://opensource.org/licenses/MIT
//     Contact: Julien SENON <julien.senon@gmail.com>
package main

import (
	"api"
	// "fmt"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"net/http"
	"web"
)

// TO FIX

func main() {
	r := mux.NewRouter()

	// Remove CORS Header check to allow swagger and application on same host and port
	headersOk := handlers.AllowedHeaders([]string{"X-Requested-With", "Content-Type"})
	// To be changed
	originsOk := handlers.AllowedOrigins([]string{"*"})
	methodsOk := handlers.AllowedMethods([]string{"GET", "HEAD", "POST", "PUT", "OPTIONS", "PATCH"})

	// Web Part
	r.HandleFunc("/index", web.Index)
	r.HandleFunc("/toto", web.Toto)

	r.HandleFunc("/login", web.Login)

	// Static dir
	r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("templates/static/"))))

	// Health Check
	r.HandleFunc("/healthy/am-i-up", api.Statusamiup).Methods("GET")
	r.HandleFunc("/healthy/about", api.Statusabout).Methods("GET")

	http.ListenAndServe(":9030", handlers.CORS(originsOk, headersOk, methodsOk)(r))
}
<commit_msg>cleanning<commit_after>//go:generate swagger generate spec
// Package main GalAirCheck.
//
// the purpose of this application is to provide an Galera Check application
// that will show state of Galera Cluster in real Time
//
// Terms Of Service:
//
// there are no TOS at this moment, use at your own risk we take no responsibility
//
//     Schemes:
//     Host:
//     BasePath:
//     Version: 0.0.1
//     License: MIT http://opensource.org/licenses/MIT
//     Contact: Julien SENON <julien.senon@gmail.com>
package main

import (
	"api"
	// "fmt"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"net/http"
	"web"
)

// TO FIX

func main() {
	r := mux.NewRouter()

	// Remove CORS Header check to allow swagger and application on same host and port
	headersOk := handlers.AllowedHeaders([]string{"X-Requested-With", "Content-Type"})
	// To be changed
	originsOk := handlers.AllowedOrigins([]string{"*"})
	methodsOk := handlers.AllowedMethods([]string{"GET", "HEAD", "POST", "PUT", "OPTIONS", "PATCH"})

	// Web Part
	r.HandleFunc("/index", web.Index)

	// Static dir
	r.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("templates/static/"))))

	// Health Check
	r.HandleFunc("/healthy/am-i-up", api.Statusamiup).Methods("GET")
	r.HandleFunc("/healthy/about", api.Statusabout).Methods("GET")

	http.ListenAndServe(":9030", handlers.CORS(originsOk, headersOk, methodsOk)(r))
}
<|endoftext|>


<commit_before>package vault

import (
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
	"time"
)

var (
	TEST_DIR = os.Getenv("GOPATH") + "/src/github.com/ieee0824/thor/test"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randStringRunes(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func TestIsSecret(t *testing.T) {
	plainTextPath := TEST_DIR + "/vaultTestFiles/plain.json"
	chipherTextPath := TEST_DIR + "/vaultTestFiles/chipher.json"

	plain, err := ioutil.ReadFile(plainTextPath)
	if err != nil {
		t.Error(err)
	} else if IsSecret(plain) {
		t.Errorf("The result is illegal. I want %v, but it is actually %v.", false, IsSecret(plain))
	}

	chipher, err := ioutil.ReadFile(chipherTextPath)
	if err != nil {
		t.Error(err)
	} else if !IsSecret(chipher) {
		t.Errorf("The result is illegal. I want %v, but it is actually %v.", true, IsSecret(chipher))
	}
}

func falsification(d []byte) []byte {

	for i, v := range d {
		if i%2 == 0 {
			d[i] <<= uint(rand.Int())
		} else {
			d[i] >>= uint(rand.Int())
		}
	}
	return d
}

func TestEncryption(t *testing.T) {
	var randomStr = randStringRunes(65536)
	var key = "test"
	var invalidKey = "johnDoe"

	encrypter := NewString(randomStr)

	if err := encrypter.Encrypt(key); err != nil {
		t.Error(err)
	} else if string(encrypter.Chipher) == randomStr {
		t.Errorf("cipher text and plain text is match")
	}

	if result, err := encrypter.Decrypt(key); err != nil {
		t.Error(err)
	} else if string(result) != randomStr {
		t.Errorf("cipher text and plain text is not match")
	}

	if _, err := encrypter.Decrypt(invalidKey); err == nil {
		t.Errorf("There is no error with an invalid key.")
	}
	encrypter.Chipher = falsification(encrypter.Chipher)
	if _, err := encrypter.Decrypt(key); err == nil {
		t.Errorf("Tampering is not detected.")
	}

}
<commit_msg>データの破損検出<commit_after>package vault

import (
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
	"time"
)

var (
	TEST_DIR = os.Getenv("GOPATH") + "/src/github.com/ieee0824/thor/test"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

func randStringRunes(n int) string {
	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func TestIsSecret(t *testing.T) {
	plainTextPath := TEST_DIR + "/vaultTestFiles/plain.json"
	chipherTextPath := TEST_DIR + "/vaultTestFiles/chipher.json"

	plain, err := ioutil.ReadFile(plainTextPath)
	if err != nil {
		t.Error(err)
	} else if IsSecret(plain) {
		t.Errorf("The result is illegal. I want %v, but it is actually %v.", false, IsSecret(plain))
	}

	chipher, err := ioutil.ReadFile(chipherTextPath)
	if err != nil {
		t.Error(err)
	} else if !IsSecret(chipher) {
		t.Errorf("The result is illegal. I want %v, but it is actually %v.", true, IsSecret(chipher))
	}
}

func falsification(d []byte) []byte {

	for i, v := range d {
		if i%2 == 0 {
			d[i] = v << uint(rand.Int())
		} else {
			d[i] = v >> uint(rand.Int())
		}
	}
	return d
}

func TestEncryption(t *testing.T) {
	var randomStr = randStringRunes(65536)
	var key = "test"
	var invalidKey = "johnDoe"

	encrypter := NewString(randomStr)

	if err := encrypter.Encrypt(key); err != nil {
		t.Error(err)
	} else if string(encrypter.Chipher) == randomStr {
		t.Errorf("cipher text and plain text is match")
	}

	if result, err := encrypter.Decrypt(key); err != nil {
		t.Error(err)
	} else if string(result) != randomStr {
		t.Errorf("cipher text and plain text is not match")
	}

	if _, err := encrypter.Decrypt(invalidKey); err == nil {
		t.Errorf("There is no error with an invalid key.")
	}
	encrypter.Chipher = falsification(encrypter.Chipher)
	if _, err := encrypter.Decrypt(key); err == nil {
		t.Errorf("Tampering is not detected.")
	}

}
<|endoftext|>


<commit_before>package check

import (
	"fmt"
	"plaid/parser"
)

// Scope tracks the symbol table and other data used during the check
type Scope struct {
	parent        *Scope
	errs          []error
	variables     []string
	values        map[string]Type
	pendingReturn Type
	queue         []struct {
		ret  Type
		expr parser.FunctionExpr
	}
}

func (s *Scope) hasParent() bool {
	return (s.parent != nil)
}

// Errors returns a list of errors detected during the check
func (s *Scope) Errors() []error {
	if s.hasParent() {
		return s.parent.Errors()
	}

	return s.errs
}

func (s *Scope) addError(err error) {
	if s.hasParent() {
		s.parent.addError(err)
	} else {
		s.errs = append(s.errs, err)
	}
}

func (s *Scope) hasVariable(name string) bool {
	_, exists := s.values[name]
	return exists
}

func (s *Scope) registerVariable(name string, typ Type) {
	s.variables = append(s.variables, name)
	s.values[name] = typ
}

func (s *Scope) getVariable(name string) Type {
	return s.values[name]
}

func (s *Scope) hasPendingReturnType() bool {
	return (s.pendingReturn != nil)
}

func (s *Scope) getPendingReturnType() Type {
	return s.pendingReturn
}

func (s *Scope) setPendingReturnType(typ Type) {
	s.pendingReturn = typ
}

func (s *Scope) hasBodyQueue() bool {
	return len(s.queue) > 0
}

func (s *Scope) enqueueBody(ret Type, expr parser.FunctionExpr) {
	body := struct {
		ret  Type
		expr parser.FunctionExpr
	}{ret, expr}
	s.queue = append(s.queue, body)
}

func (s *Scope) dequeueBody() (Type, parser.FunctionExpr) {
	body := s.queue[0]
	s.queue = s.queue[1:]
	return body.ret, body.expr
}

func (s *Scope) String() string {
	var out string
	for _, name := range s.variables {
		out += fmt.Sprintf("%s : %s\n", name, s.values[name])
	}
	return out
}

func makeScope(parent *Scope, ret Type) *Scope {
	return &Scope{
		parent,
		[]error{},
		[]string{},
		make(map[string]Type),
		ret,
		[]struct {
			ret  Type
			expr parser.FunctionExpr
		}{},
	}
}
<commit_msg>if variable is not local, check parent scope<commit_after>package check

import (
	"fmt"
	"plaid/parser"
)

// Scope tracks the symbol table and other data used during the check
type Scope struct {
	parent        *Scope
	errs          []error
	variables     []string
	values        map[string]Type
	pendingReturn Type
	queue         []struct {
		ret  Type
		expr parser.FunctionExpr
	}
}

func (s *Scope) hasParent() bool {
	return (s.parent != nil)
}

// Errors returns a list of errors detected during the check
func (s *Scope) Errors() []error {
	if s.hasParent() {
		return s.parent.Errors()
	}

	return s.errs
}

func (s *Scope) addError(err error) {
	if s.hasParent() {
		s.parent.addError(err)
	} else {
		s.errs = append(s.errs, err)
	}
}

func (s *Scope) hasVariable(name string) bool {
	_, exists := s.values[name]

	if exists == false && s.hasParent() {
		return s.parent.hasVariable(name)
	}

	return exists
}

func (s *Scope) registerVariable(name string, typ Type) {
	s.variables = append(s.variables, name)
	s.values[name] = typ
}

func (s *Scope) getVariable(name string) Type {
	return s.values[name]
}

func (s *Scope) hasPendingReturnType() bool {
	return (s.pendingReturn != nil)
}

func (s *Scope) getPendingReturnType() Type {
	return s.pendingReturn
}

func (s *Scope) setPendingReturnType(typ Type) {
	s.pendingReturn = typ
}

func (s *Scope) hasBodyQueue() bool {
	return len(s.queue) > 0
}

func (s *Scope) enqueueBody(ret Type, expr parser.FunctionExpr) {
	body := struct {
		ret  Type
		expr parser.FunctionExpr
	}{ret, expr}
	s.queue = append(s.queue, body)
}

func (s *Scope) dequeueBody() (Type, parser.FunctionExpr) {
	body := s.queue[0]
	s.queue = s.queue[1:]
	return body.ret, body.expr
}

func (s *Scope) String() string {
	var out string
	for _, name := range s.variables {
		out += fmt.Sprintf("%s : %s\n", name, s.values[name])
	}
	return out
}

func makeScope(parent *Scope, ret Type) *Scope {
	return &Scope{
		parent,
		[]error{},
		[]string{},
		make(map[string]Type),
		ret,
		[]struct {
			ret  Type
			expr parser.FunctionExpr
		}{},
	}
}
<|endoftext|>


<commit_before>// +build gofig

package config

import (
	"os"
	"runtime"

	log "github.com/Sirupsen/logrus"
	gofigCore "github.com/akutz/gofig"
	gofig "github.com/akutz/gofig/types"

	"github.com/codedellemc/libstorage/api/context"
	"github.com/codedellemc/libstorage/api/registry"
	"github.com/codedellemc/libstorage/api/types"
)

const (
	logStdoutDesc = "The file to which to log os.Stdout"
	logStderrDesc = "The file to which to log os.Stderr"
)

func init() {
	gofigCore.LogGetAndSet = false
	gofigCore.LogSecureKey = false
	gofigCore.LogFlattenEnvVars = false

	registry.RegisterConfigReg(
		"libStorage",
		func(ctx types.Context, r gofig.ConfigRegistration) {

			pathConfig := context.MustPathConfig(ctx)

			var lvl log.Level
			if types.Debug {
				lvl = log.DebugLevel
			} else {
				ll, err := log.ParseLevel(os.Getenv("LIBSTORAGE_LOGGING_LEVEL"))
				if err != nil {
					ll = log.WarnLevel
				}
				lvl = ll
			}

			rk := func(
				keyType gofig.ConfigKeyTypes,
				defaultVal interface{},
				description string,
				keyVal types.ConfigKey,
				args ...interface{}) {

				if args == nil {
					args = []interface{}{keyVal}
				} else {
					args = append([]interface{}{keyVal}, args...)
				}
				r.Key(keyType, "", defaultVal, description, args...)
			}

			defaultAEM := types.UnixEndpoint.String()
			defaultStorageDriver := types.LibStorageDriverName
			defaultLogLevel := lvl.String()
			defaultClientType := types.IntegrationClient.String()

			rk(gofig.String, "", "", types.ConfigHost)
			rk(gofig.String, "", "", types.ConfigService)
			rk(gofig.String, defaultAEM, "", types.ConfigServerAutoEndpointMode)
			rk(gofig.String, runtime.GOOS, "", types.ConfigOSDriver)
			rk(gofig.String, defaultStorageDriver, "",
				types.ConfigStorageDriver)
			rk(gofig.String, defaultIntDriver, "",
				types.ConfigIntegrationDriver)
			rk(gofig.String, defaultClientType, "", types.ConfigClientType)
			rk(gofig.String, defaultLogLevel, "", types.ConfigLogLevel)
			rk(gofig.String, "", logStdoutDesc, types.ConfigLogStderr)
			rk(gofig.String, "", logStderrDesc, types.ConfigLogStdout)
			rk(gofig.Bool, types.Debug, "", types.ConfigLogHTTPRequests)
			rk(gofig.Bool, types.Debug, "", types.ConfigLogHTTPResponses)
			rk(gofig.Bool, false, "", types.ConfigHTTPDisableKeepAlive)
			rk(gofig.Int, 300, "", types.ConfigHTTPWriteTimeout)
			rk(gofig.Int, 300, "", types.ConfigHTTPReadTimeout)

			rk(gofig.String, pathConfig.LSX, "", types.ConfigExecutorPath)

			rk(gofig.Bool, false, "", types.ConfigExecutorNoDownload)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsMountPreempt)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsCreateDisable)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsRemoveDisable)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsUnmountIgnoreUsed)
			rk(gofig.Bool, true, "", types.ConfigIgVolOpsPathCacheEnabled)
			rk(gofig.Bool, true, "", types.ConfigIgVolOpsPathCacheAsync)
			rk(gofig.String, "30m", "", types.ConfigClientCacheInstanceID)
			rk(gofig.String, "30s", "", types.ConfigDeviceAttachTimeout)
			rk(gofig.Int, 0, "", types.ConfigDeviceScanType)
			rk(gofig.Bool, false, "", types.ConfigEmbedded)
			rk(gofig.String, "1m", "", types.ConfigServerTasksExeTimeout)
			rk(gofig.String, "0s", "", types.ConfigServerTasksLogTimeout)
			rk(gofig.Bool, false, "", types.ConfigServerParseRequestOpts)

			// tls config
			rk(
				gofig.String,
				pathConfig.DefaultTLSCertFile,
				"",
				types.ConfigTLSCertFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSKeyFile,
				"",
				types.ConfigTLSKeyFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSTrustedRootsFile,
				"",
				types.ConfigTLSTrustedCertsFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSKnownHosts,
				"",
				types.ConfigTLSKnownHosts)
			rk(gofig.String, "", "", types.ConfigTLSServerName)
			rk(gofig.Bool, false, "", types.ConfigTLSDisabled)
			rk(gofig.Bool, false, "", types.ConfigTLSInsecure)
			rk(gofig.Bool, false, "", types.ConfigTLSClientCertRequired)

			// auth config - client
			rk(gofig.String, "", "", types.ConfigClientAuthToken)

			// auth config - server
			rk(gofig.String, "", "", types.ConfigServerAuthKey)
			rk(gofig.String, "HS256", "", types.ConfigServerAuthAlg)
			rk(gofig.String, "", "", types.ConfigServerAuthAllow)
			rk(gofig.String, "", "", types.ConfigServerAuthDeny)
			rk(gofig.Bool, false, "", types.ConfigServerAuthDisabled)
		})
}
<commit_msg>TLS Boolean Property Fix<commit_after>// +build gofig

package config

import (
	"os"
	"runtime"

	log "github.com/Sirupsen/logrus"
	gofigCore "github.com/akutz/gofig"
	gofig "github.com/akutz/gofig/types"

	"github.com/codedellemc/libstorage/api/context"
	"github.com/codedellemc/libstorage/api/registry"
	"github.com/codedellemc/libstorage/api/types"
)

const (
	logStdoutDesc = "The file to which to log os.Stdout"
	logStderrDesc = "The file to which to log os.Stderr"
)

func init() {
	gofigCore.LogGetAndSet = false
	gofigCore.LogSecureKey = false
	gofigCore.LogFlattenEnvVars = false

	registry.RegisterConfigReg(
		"libStorage",
		func(ctx types.Context, r gofig.ConfigRegistration) {

			pathConfig := context.MustPathConfig(ctx)

			var lvl log.Level
			if types.Debug {
				lvl = log.DebugLevel
			} else {
				ll, err := log.ParseLevel(os.Getenv("LIBSTORAGE_LOGGING_LEVEL"))
				if err != nil {
					ll = log.WarnLevel
				}
				lvl = ll
			}

			rk := func(
				keyType gofig.ConfigKeyTypes,
				defaultVal interface{},
				description string,
				keyVal types.ConfigKey,
				args ...interface{}) {

				if args == nil {
					args = []interface{}{keyVal}
				} else {
					args = append([]interface{}{keyVal}, args...)
				}
				r.Key(keyType, "", defaultVal, description, args...)
			}

			defaultAEM := types.UnixEndpoint.String()
			defaultStorageDriver := types.LibStorageDriverName
			defaultLogLevel := lvl.String()
			defaultClientType := types.IntegrationClient.String()

			rk(gofig.String, "", "", types.ConfigHost)
			rk(gofig.String, "", "", types.ConfigService)
			rk(gofig.String, defaultAEM, "", types.ConfigServerAutoEndpointMode)
			rk(gofig.String, runtime.GOOS, "", types.ConfigOSDriver)
			rk(gofig.String, defaultStorageDriver, "",
				types.ConfigStorageDriver)
			rk(gofig.String, defaultIntDriver, "",
				types.ConfigIntegrationDriver)
			rk(gofig.String, defaultClientType, "", types.ConfigClientType)
			rk(gofig.String, defaultLogLevel, "", types.ConfigLogLevel)
			rk(gofig.String, "", logStdoutDesc, types.ConfigLogStderr)
			rk(gofig.String, "", logStderrDesc, types.ConfigLogStdout)
			rk(gofig.Bool, types.Debug, "", types.ConfigLogHTTPRequests)
			rk(gofig.Bool, types.Debug, "", types.ConfigLogHTTPResponses)
			rk(gofig.Bool, false, "", types.ConfigHTTPDisableKeepAlive)
			rk(gofig.Int, 300, "", types.ConfigHTTPWriteTimeout)
			rk(gofig.Int, 300, "", types.ConfigHTTPReadTimeout)

			rk(gofig.String, pathConfig.LSX, "", types.ConfigExecutorPath)

			rk(gofig.Bool, false, "", types.ConfigExecutorNoDownload)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsMountPreempt)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsCreateDisable)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsRemoveDisable)
			rk(gofig.Bool, false, "", types.ConfigIgVolOpsUnmountIgnoreUsed)
			rk(gofig.Bool, true, "", types.ConfigIgVolOpsPathCacheEnabled)
			rk(gofig.Bool, true, "", types.ConfigIgVolOpsPathCacheAsync)
			rk(gofig.String, "30m", "", types.ConfigClientCacheInstanceID)
			rk(gofig.String, "30s", "", types.ConfigDeviceAttachTimeout)
			rk(gofig.Int, 0, "", types.ConfigDeviceScanType)
			rk(gofig.Bool, false, "", types.ConfigEmbedded)
			rk(gofig.String, "1m", "", types.ConfigServerTasksExeTimeout)
			rk(gofig.String, "0s", "", types.ConfigServerTasksLogTimeout)
			rk(gofig.Bool, false, "", types.ConfigServerParseRequestOpts)

			// tls config
			rk(
				gofig.String,
				pathConfig.DefaultTLSCertFile,
				"",
				types.ConfigTLSCertFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSKeyFile,
				"",
				types.ConfigTLSKeyFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSTrustedRootsFile,
				"",
				types.ConfigTLSTrustedCertsFile)
			rk(
				gofig.String,
				pathConfig.DefaultTLSKnownHosts,
				"",
				types.ConfigTLSKnownHosts)
			rk(gofig.String, "", "", types.ConfigTLSServerName)
			rk(gofig.String, "", "", types.ConfigTLSDisabled)
			rk(gofig.String, "", "", types.ConfigTLSInsecure)
			rk(gofig.String, "", "", types.ConfigTLSClientCertRequired)

			// auth config - client
			rk(gofig.String, "", "", types.ConfigClientAuthToken)

			// auth config - server
			rk(gofig.String, "", "", types.ConfigServerAuthKey)
			rk(gofig.String, "HS256", "", types.ConfigServerAuthAlg)
			rk(gofig.String, "", "", types.ConfigServerAuthAllow)
			rk(gofig.String, "", "", types.ConfigServerAuthDeny)
			rk(gofig.Bool, false, "", types.ConfigServerAuthDisabled)
		})
}
<|endoftext|>


<commit_before>////////////////////////////////////////////////////////////////////////////
// Porgram: FfCvt
// Purpose: ffmpeg convert wrapper tool
// Authors: Tong Sun (c) 2015, All rights reserved
////////////////////////////////////////////////////////////////////////////

/*

Transcodes all episodes in the given directory and all of it's subdirectories
using ffmpeg.

Initial version based (but also heavily hacked) on
https://gist.github.com/mmstick/3182c1c8596c1f830c7e
by Michael Murphy (mmstick)

*/

////////////////////////////////////////////////////////////////////////////
// Program start

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

const (
	STATIC_PARAMS = "me=star:subme=7:bframes=16:b-adapt=2:ref=16:rc-lookahead=60:max-merge=5:tu-intra-depth=4:tu-inter-depth=4"
)

// Contains information about each episode
type Episode struct {
	name           string
	directory      string
	originalSize   int64
	transcodedSize int64
	sizeDifference int64
	time           time.Duration
	stat           error
}

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var (
	sprintf = fmt.Sprintf
)

////////////////////////////////////////////////////////////////////////////
// Main

func main() {
	flag.Usage = Usage
	flag.Parse()

	// One mandatory arguments, either -d or -f
	if len(Opts.Directory)+len(Opts.File) < 1 {
		Usage()
	}

	startTime := time.Now()
	if Opts.Directory != "" {
		transcodeEpisodes(scanEpisodes(scanDirectory(Opts.Directory), Opts.Directory))
	} else if Opts.File != "" {
		outputName := getOutputName(Opts.File)
		fmt.Printf("\n== Transcoding: %s\n", Opts.File)
		transcodeFile(Opts.File, outputName)
	}
	fmt.Printf("Transcoding completed in %s\n", time.Since(startTime))
}

////////////////////////////////////////////////////////////////////////////
// Function definitions

//==========================================================================
// Directory & Episodes handling

// Transcodes all episodes in the episode list
func transcodeEpisodes(episodeList *[]Episode) {
	files := len(*episodeList)
	for index, ep := range *episodeList {
		ep.transcodeEpisode(index+1, files)
		ep.status()
	}
}

// Print the status of the transcoded episode
func (episode Episode) status() {
	if episode.stat != nil {
		fmt.Println("Failed to transcode", episode.name)
	} else {
		fmt.Println("Transcoded", episode.name)
		fmt.Println("Original Size:", episode.originalSize, "MB")
		fmt.Println("New Size:", episode.transcodedSize, "MB")
		fmt.Println("Difference:", episode.sizeDifference, "MB")
		fmt.Println("Time:", episode.time)
	}
}

// Recurse through each subdirectory and adds each episode to the episode list
func scanEpisodes(directoryList []os.FileInfo, directory string) *[]Episode {
	list := []Episode{}
	for _, file := range directoryList {
		if file.IsDir() {
			recurseDirectory(&directory, file.Name())
		} else {
			appendEpisode(&list, file, &directory)
		}
	}
	return &list
}

// Returns a list of files in the current directory
func scanDirectory(path string) []os.FileInfo {
	directory, _ := ioutil.ReadDir(path)
	return directory
}

// If the file is a directory, recurse through the directory.
func recurseDirectory(directory *string, filename string) {
	subdirectory := sprintf("%s/%s", *directory, filename)
	scanEpisodes(scanDirectory(subdirectory), subdirectory)
}

// Append the current episode to the episode list, unless it's encoded already
func appendEpisode(list *[]Episode, file os.FileInfo, directory *string) {
	fname := file.Name()
	if fname[len(fname)-5:] == "_.mkv" {
		return
	}

	*list = append(*list, Episode{
		name:         fname,
		directory:    *directory,
		originalSize: file.Size() / 1000000,
	})
}

//==========================================================================
// Transcode handling

// Transcode the current episode
func (ep Episode) transcodeEpisode(index, files int) {
	inputName := sprintf("%s/%s", ep.directory, ep.name)
	outputName := getOutputName(inputName)
	fmt.Printf("\n== Transcoding [%d/%d]: %s\n", index, files, ep.name)
	ep.time, ep.stat = transcodeFile(inputName, outputName)
	ep.transcodedSize = transcodeSize(outputName)
	ep.sizeDifference = ep.originalSize - ep.transcodedSize
}

func transcodeFile(inputName, outputName string) (time.Duration, error) {
	startTime := time.Now()

	args := encodeParametersV(encodeParametersA(
		[]string{"-i", inputName}))
	args = append(args, os.Args...)
	args = append(args, outputName)
	debug(Opts.FFMpeg)
	debug(strings.Join(args, " "))

	cmd := exec.Command(Opts.FFMpeg, args...)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Printf("%s: Exec error - %s", progname, err.Error())
	}
	//fmt.Printf("\n== Out:\n%s\n", out.String())
	time := time.Since(startTime)
	return time, err
}

// Returns the encode parameters for Audio
func encodeParametersA(args []string) []string {
	if Opts.AC {
		args = append(args, "-c:a", "copy")
		return args
	}
	return args
}

// Returns the encode parameters for Video
func encodeParametersV(args []string) []string {
	if Opts.VC {
		args = append(args, "-c:v", "copy")
		return args
	}
	return args
}

// Returns the size of the newly transcoded episode, if it exists.
func transcodeSize(transcodedEpisode string) int64 {
	file, err := os.Open(transcodedEpisode)
	if err == nil {
		stat, _ := file.Stat()
		return stat.Size() / 1000000
	} else {
		log.Printf("%s: Open error - %s", progname, err.Error())
		return 0
	}
}

// Replaces the file extension from the input string with _.mkv
func getOutputName(input string) string {
	for index := len(input) - 1; index >= 0; index-- {
		if input[index] == '.' {
			input = input[:index]
		}
	}
	return input + "_.mkv"
}

func debug(input string) {
	if Opts.Debug == 0 {
		return
	}
	print("] ")
	print(input)
	print("\n")
}
<commit_msg>- [!] use flag.Args() instead to catch not all but the rest of non-flag arguments<commit_after>////////////////////////////////////////////////////////////////////////////
// Porgram: FfCvt
// Purpose: ffmpeg convert wrapper tool
// Authors: Tong Sun (c) 2015, All rights reserved
////////////////////////////////////////////////////////////////////////////

/*

Transcodes all episodes in the given directory and all of it's subdirectories
using ffmpeg.

Initial version based (but also heavily hacked) on
https://gist.github.com/mmstick/3182c1c8596c1f830c7e
by Michael Murphy (mmstick)

*/

////////////////////////////////////////////////////////////////////////////
// Program start

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

////////////////////////////////////////////////////////////////////////////
// Constant and data type/structure definitions

const (
	STATIC_PARAMS = "me=star:subme=7:bframes=16:b-adapt=2:ref=16:rc-lookahead=60:max-merge=5:tu-intra-depth=4:tu-inter-depth=4"
)

// Contains information about each episode
type Episode struct {
	name           string
	directory      string
	originalSize   int64
	transcodedSize int64
	sizeDifference int64
	time           time.Duration
	stat           error
}

////////////////////////////////////////////////////////////////////////////
// Global variables definitions

var (
	sprintf = fmt.Sprintf
)

////////////////////////////////////////////////////////////////////////////
// Main

func main() {
	flag.Usage = Usage
	flag.Parse()

	// One mandatory arguments, either -d or -f
	if len(Opts.Directory)+len(Opts.File) < 1 {
		Usage()
	}

	startTime := time.Now()
	if Opts.Directory != "" {
		transcodeEpisodes(scanEpisodes(scanDirectory(Opts.Directory), Opts.Directory))
	} else if Opts.File != "" {
		outputName := getOutputName(Opts.File)
		fmt.Printf("\n== Transcoding: %s\n", Opts.File)
		transcodeFile(Opts.File, outputName)
	}
	fmt.Printf("Transcoding completed in %s\n", time.Since(startTime))
}

////////////////////////////////////////////////////////////////////////////
// Function definitions

//==========================================================================
// Directory & Episodes handling

// Transcodes all episodes in the episode list
func transcodeEpisodes(episodeList *[]Episode) {
	files := len(*episodeList)
	for index, ep := range *episodeList {
		ep.transcodeEpisode(index+1, files)
		ep.status()
	}
}

// Print the status of the transcoded episode
func (episode Episode) status() {
	if episode.stat != nil {
		fmt.Println("Failed to transcode", episode.name)
	} else {
		fmt.Println("Transcoded", episode.name)
		fmt.Println("Original Size:", episode.originalSize, "MB")
		fmt.Println("New Size:", episode.transcodedSize, "MB")
		fmt.Println("Difference:", episode.sizeDifference, "MB")
		fmt.Println("Time:", episode.time)
	}
}

// Recurse through each subdirectory and adds each episode to the episode list
func scanEpisodes(directoryList []os.FileInfo, directory string) *[]Episode {
	list := []Episode{}
	for _, file := range directoryList {
		if file.IsDir() {
			recurseDirectory(&directory, file.Name())
		} else {
			appendEpisode(&list, file, &directory)
		}
	}
	return &list
}

// Returns a list of files in the current directory
func scanDirectory(path string) []os.FileInfo {
	directory, _ := ioutil.ReadDir(path)
	return directory
}

// If the file is a directory, recurse through the directory.
func recurseDirectory(directory *string, filename string) {
	subdirectory := sprintf("%s/%s", *directory, filename)
	scanEpisodes(scanDirectory(subdirectory), subdirectory)
}

// Append the current episode to the episode list, unless it's encoded already
func appendEpisode(list *[]Episode, file os.FileInfo, directory *string) {
	fname := file.Name()
	if fname[len(fname)-5:] == "_.mkv" {
		return
	}

	*list = append(*list, Episode{
		name:         fname,
		directory:    *directory,
		originalSize: file.Size() / 1000000,
	})
}

//==========================================================================
// Transcode handling

// Transcode the current episode
func (ep Episode) transcodeEpisode(index, files int) {
	inputName := sprintf("%s/%s", ep.directory, ep.name)
	outputName := getOutputName(inputName)
	fmt.Printf("\n== Transcoding [%d/%d]: %s\n", index, files, ep.name)
	ep.time, ep.stat = transcodeFile(inputName, outputName)
	ep.transcodedSize = transcodeSize(outputName)
	ep.sizeDifference = ep.originalSize - ep.transcodedSize
}

func transcodeFile(inputName, outputName string) (time.Duration, error) {
	startTime := time.Now()

	args := encodeParametersV(encodeParametersA(
		[]string{"-i", inputName}))
	args = append(args, flag.Args()...)
	args = append(args, outputName)
	debug(Opts.FFMpeg)
	debug(strings.Join(args, " "))

	cmd := exec.Command(Opts.FFMpeg, args...)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Printf("%s: Exec error - %s", progname, err.Error())
	}
	fmt.Printf("\n%s\n", out.String())
	time := time.Since(startTime)
	return time, err
}

// Returns the encode parameters for Audio
func encodeParametersA(args []string) []string {
	if Opts.AC {
		args = append(args, "-c:a", "copy")
		return args
	}
	return args
}

// Returns the encode parameters for Video
func encodeParametersV(args []string) []string {
	if Opts.VC {
		args = append(args, "-c:v", "copy")
		return args
	}
	return args
}

// Returns the size of the newly transcoded episode, if it exists.
func transcodeSize(transcodedEpisode string) int64 {
	file, err := os.Open(transcodedEpisode)
	if err == nil {
		stat, _ := file.Stat()
		return stat.Size() / 1000000
	} else {
		log.Printf("%s: Open error - %s", progname, err.Error())
		return 0
	}
}

// Replaces the file extension from the input string with _.mkv
func getOutputName(input string) string {
	for index := len(input) - 1; index >= 0; index-- {
		if input[index] == '.' {
			input = input[:index]
		}
	}
	return input + "_.mkv"
}

func debug(input string) {
	if Opts.Debug == 0 {
		return
	}
	print("] ")
	print(input)
	print("\n")
}
<|endoftext|>


<commit_before>package compress

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"runtime"

	"github.com/klauspost/pgzip"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
	"github.com/pierrec/lz4"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// Fields from config file
	OutputPath        string `mapstructure:"output"`
	CompressionLevel  int    `mapstructure:"compression_level"`
	KeepInputArtifact bool   `mapstructure:"keep_input_artifact"`

	// Derived fields
	Archive   string
	Algorithm string

	ctx *interpolate.Context
}

type PostProcessor struct {
	config *Config
}

var (
	// ErrInvalidCompressionLevel is returned when the compression level passed
	// to gzip is not in the expected range. See compress/flate for details.
	ErrInvalidCompressionLevel = fmt.Errorf(
		"Invalid compression level. Expected an integer from -1 to 9.")

	ErrWrongInputCount = fmt.Errorf(
		"Can only have 1 input file when not using tar/zip")

	filenamePattern = regexp.MustCompile(`(?:\.([a-z0-9]+))`)
)

func (p *PostProcessor) Configure(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate: true,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)

	errs := new(packer.MultiError)

	if p.config.OutputPath == "" {
		p.config.OutputPath = "packer_{{.BuildName}}_{{.Provider}}"
	}

	if err = interpolate.Validate(p.config.OutputPath, p.config.ctx); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"output": &p.config.OutputPath,
	}

	if p.config.CompressionLevel > pgzip.BestCompression {
		p.config.CompressionLevel = pgzip.BestCompression
	}
	// Technically 0 means "don't compress" but I don't know how to
	// differentiate between "user entered zero" and "user entered nothing".
	// Also, why bother creating a compressed file with zero compression?
	if p.config.CompressionLevel == -1 || p.config.CompressionLevel == 0 {
		p.config.CompressionLevel = pgzip.DefaultCompression
	}

	for key, ptr := range templates {
		if *ptr == "" {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("%s must be set", key))
		}

		*ptr, err = interpolate.Render(p.config.OutputPath, p.config.ctx)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", key, err))
		}
	}

	p.config.detectFromFilename()

	if len(errs.Errors) > 0 {
		return errs
	}

	return nil

}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {

	target := p.config.OutputPath
	keep := p.config.KeepInputArtifact
	newArtifact := &Artifact{Path: target}

	outputFile, err := os.Create(target)
	if err != nil {
		return nil, false, fmt.Errorf(
			"Unable to create archive %s: %s", target, err)
	}
	defer outputFile.Close()

	// Setup output interface. If we're using compression, output is a
	// compression writer. Otherwise it's just a file.
	var output io.WriteCloser
	switch p.config.Algorithm {
	case "lz4":
		ui.Say(fmt.Sprintf("Preparing lz4 compression for %s", target))
		output, err = makeLZ4Writer(outputFile, p.config.CompressionLevel)
		defer output.Close()
	case "pgzip":
		ui.Say(fmt.Sprintf("Preparing gzip compression for %s", target))
		output, err = makePgzipWriter(outputFile, p.config.CompressionLevel)
		defer output.Close()
	default:
		output = outputFile
	}

	compression := p.config.Algorithm
	if compression == "" {
		compression = "no"
	}

	// Build an archive, if we're supposed to do that.
	switch p.config.Archive {
	case "tar":
		ui.Say(fmt.Sprintf("Tarring %s with %s compression", target, compression))
		err = createTarArchive(artifact.Files(), output)
		if err != nil {
			return nil, keep, fmt.Errorf("Error creating tar: %s", err)
		}
	case "zip":
		ui.Say(fmt.Sprintf("Zipping %s", target))
		err = createZipArchive(artifact.Files(), output)
		if err != nil {
			return nil, keep, fmt.Errorf("Error creating zip: %s", err)
		}
	default:
		ui.Say(fmt.Sprintf("Copying %s with %s compression", target, compression))
		// Filename indicates no tarball (just compress) so we'll do an io.Copy
		// into our compressor.
		if len(artifact.Files()) != 1 {
			return nil, keep, fmt.Errorf(
				"Can only have 1 input file when not using tar/zip. Found %d "+
					"files: %v", len(artifact.Files()), artifact.Files())
		}

		source, err := os.Open(artifact.Files()[0])
		if err != nil {
			return nil, keep, fmt.Errorf(
				"Failed to open source file %s for reading: %s",
				artifact.Files()[0], err)
		}
		defer source.Close()

		if _, err = io.Copy(output, source); err != nil {
			return nil, keep, fmt.Errorf("Failed to compress %s: %s",
				artifact.Files()[0], err)
		}
	}

	ui.Say(fmt.Sprintf("Archive %s completed", target))

	return newArtifact, keep, nil
}

func (config *Config) detectFromFilename() {

	extensions := map[string]string{
		"tar": "tar",
		"zip": "zip",
		"gz":  "pgzip",
		"lz4": "lz4",
	}

	result := filenamePattern.FindAllStringSubmatch(config.OutputPath, -1)

	// No dots. Bail out with defaults.
	if len(result) == 0 {
		config.Algorithm = "pgzip"
		config.Archive = "tar"
		return
	}

	// Parse the last two .groups, if they're there
	lastItem := result[len(result)-1][1]
	var nextToLastItem string
	if len(result) == 1 {
		nextToLastItem = ""
	} else {
		nextToLastItem = result[len(result)-2][1]
	}

	// Should we make an archive? E.g. tar or zip?
	if nextToLastItem == "tar" {
		config.Archive = "tar"
	}
	if lastItem == "zip" || lastItem == "tar" {
		config.Archive = lastItem
		// Tar or zip is our final artifact. Bail out.
		return
	}

	// Should we compress the artifact?
	algorithm, ok := extensions[lastItem]
	if ok {
		config.Algorithm = algorithm
		// We found our compression algorithm. Bail out.
		return
	}

	// We didn't match a known compression format. Default to tar + pgzip
	config.Algorithm = "pgzip"
	config.Archive = "tar"
	return
}

func makeLZ4Writer(output io.WriteCloser, compressionLevel int) (io.WriteCloser, error) {
	lzwriter := lz4.NewWriter(output)
	if compressionLevel > gzip.DefaultCompression {
		lzwriter.Header.HighCompression = true
	}
	return lzwriter, nil
}

func makePgzipWriter(output io.WriteCloser, compressionLevel int) (io.WriteCloser, error) {
	gzipWriter, err := pgzip.NewWriterLevel(output, compressionLevel)
	if err != nil {
		return nil, ErrInvalidCompressionLevel
	}
	gzipWriter.SetConcurrency(500000, runtime.GOMAXPROCS(-1))
	return gzipWriter, nil
}

func createTarArchive(files []string, output io.WriteCloser) error {
	archive := tar.NewWriter(output)
	defer archive.Close()

	for _, path := range files {
		file, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Unable to read file %s: %s", path, err)
		}
		defer file.Close()

		fi, err := file.Stat()
		if err != nil {
			return fmt.Errorf("Unable to get fileinfo for %s: %s", path, err)
		}

		header, err := tar.FileInfoHeader(fi, path)
		if err != nil {
			return fmt.Errorf("Failed to create tar header for %s: %s", path, err)
		}

		if err := archive.WriteHeader(header); err != nil {
			return fmt.Errorf("Failed to write tar header for %s: %s", path, err)
		}

		if _, err := io.Copy(archive, file); err != nil {
			return fmt.Errorf("Failed to copy %s data to archive: %s", path, err)
		}
	}
	return nil
}

func createZipArchive(files []string, output io.WriteCloser) error {
	archive := zip.NewWriter(output)
	defer archive.Close()

	for _, path := range files {
		path = filepath.ToSlash(path)

		source, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Unable to read file %s: %s", path, err)
		}
		defer source.Close()

		target, err := archive.Create(path)
		if err != nil {
			return fmt.Errorf("Failed to add zip header for %s: %s", path, err)
		}

		_, err = io.Copy(target, source)
		if err != nil {
			return fmt.Errorf("Failed to copy %s data to archive: %s", path, err)
		}
	}
	return nil
}
<commit_msg>Tweaked some of the UI/UX around GOMAXPROCS<commit_after>package compress

import (
	"archive/tar"
	"archive/zip"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"runtime"

	"github.com/klauspost/pgzip"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/config"
	"github.com/mitchellh/packer/packer"
	"github.com/mitchellh/packer/template/interpolate"
	"github.com/pierrec/lz4"
)

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// Fields from config file
	OutputPath        string `mapstructure:"output"`
	CompressionLevel  int    `mapstructure:"compression_level"`
	KeepInputArtifact bool   `mapstructure:"keep_input_artifact"`

	// Derived fields
	Archive   string
	Algorithm string

	ctx *interpolate.Context
}

type PostProcessor struct {
	config *Config
}

var (
	// ErrInvalidCompressionLevel is returned when the compression level passed
	// to gzip is not in the expected range. See compress/flate for details.
	ErrInvalidCompressionLevel = fmt.Errorf(
		"Invalid compression level. Expected an integer from -1 to 9.")

	ErrWrongInputCount = fmt.Errorf(
		"Can only have 1 input file when not using tar/zip")

	filenamePattern = regexp.MustCompile(`(?:\.([a-z0-9]+))`)
)

func (p *PostProcessor) Configure(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate: true,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{},
		},
	}, raws...)

	errs := new(packer.MultiError)

	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	if p.config.OutputPath == "" {
		p.config.OutputPath = "packer_{{.BuildName}}_{{.Provider}}"
	}

	if err = interpolate.Validate(p.config.OutputPath, p.config.ctx); err != nil {
		errs = packer.MultiErrorAppend(
			errs, fmt.Errorf("Error parsing target template: %s", err))
	}

	templates := map[string]*string{
		"output": &p.config.OutputPath,
	}

	if p.config.CompressionLevel > pgzip.BestCompression {
		p.config.CompressionLevel = pgzip.BestCompression
	}
	// Technically 0 means "don't compress" but I don't know how to
	// differentiate between "user entered zero" and "user entered nothing".
	// Also, why bother creating a compressed file with zero compression?
	if p.config.CompressionLevel == -1 || p.config.CompressionLevel == 0 {
		p.config.CompressionLevel = pgzip.DefaultCompression
	}

	for key, ptr := range templates {
		if *ptr == "" {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("%s must be set", key))
		}

		*ptr, err = interpolate.Render(p.config.OutputPath, p.config.ctx)
		if err != nil {
			errs = packer.MultiErrorAppend(
				errs, fmt.Errorf("Error processing %s: %s", key, err))
		}
	}

	p.config.detectFromFilename()

	if len(errs.Errors) > 0 {
		return errs
	}

	return nil

}

func (p *PostProcessor) PostProcess(ui packer.Ui, artifact packer.Artifact) (packer.Artifact, bool, error) {

	target := p.config.OutputPath
	keep := p.config.KeepInputArtifact
	newArtifact := &Artifact{Path: target}

	outputFile, err := os.Create(target)
	if err != nil {
		return nil, false, fmt.Errorf(
			"Unable to create archive %s: %s", target, err)
	}
	defer outputFile.Close()

	// Setup output interface. If we're using compression, output is a
	// compression writer. Otherwise it's just a file.
	var output io.WriteCloser
	switch p.config.Algorithm {
	case "lz4":
		ui.Say(fmt.Sprintf("Using lz4 compression with %d cores for %s",
			runtime.GOMAXPROCS(-1), target))
		output, err = makeLZ4Writer(outputFile, p.config.CompressionLevel)
		defer output.Close()
	case "pgzip":
		ui.Say(fmt.Sprintf("Using pgzip compression with %d cores for %s",
			runtime.GOMAXPROCS(-1), target))
		output, err = makePgzipWriter(outputFile, p.config.CompressionLevel)
		defer output.Close()
	default:
		output = outputFile
	}

	compression := p.config.Algorithm
	if compression == "" {
		compression = "no compression"
	}

	// Build an archive, if we're supposed to do that.
	switch p.config.Archive {
	case "tar":
		ui.Say(fmt.Sprintf("Tarring %s with %s", target, compression))
		err = createTarArchive(artifact.Files(), output)
		if err != nil {
			return nil, keep, fmt.Errorf("Error creating tar: %s", err)
		}
	case "zip":
		ui.Say(fmt.Sprintf("Zipping %s", target))
		err = createZipArchive(artifact.Files(), output)
		if err != nil {
			return nil, keep, fmt.Errorf("Error creating zip: %s", err)
		}
	default:
		// Filename indicates no tarball (just compress) so we'll do an io.Copy
		// into our compressor.
		if len(artifact.Files()) != 1 {
			return nil, keep, fmt.Errorf(
				"Can only have 1 input file when not using tar/zip. Found %d "+
					"files: %v", len(artifact.Files()), artifact.Files())
		}
		archiveFile := artifact.Files()[0]
		ui.Say(fmt.Sprintf("Archiving %s with %s", archiveFile, compression))

		source, err := os.Open(archiveFile)
		if err != nil {
			return nil, keep, fmt.Errorf(
				"Failed to open source file %s for reading: %s",
				archiveFile, err)
		}
		defer source.Close()

		if _, err = io.Copy(output, source); err != nil {
			return nil, keep, fmt.Errorf("Failed to compress %s: %s",
				archiveFile, err)
		}
	}

	ui.Say(fmt.Sprintf("Archive %s completed", target))

	return newArtifact, keep, nil
}

func (config *Config) detectFromFilename() {

	extensions := map[string]string{
		"tar": "tar",
		"zip": "zip",
		"gz":  "pgzip",
		"lz4": "lz4",
	}

	result := filenamePattern.FindAllStringSubmatch(config.OutputPath, -1)

	// No dots. Bail out with defaults.
	if len(result) == 0 {
		config.Algorithm = "pgzip"
		config.Archive = "tar"
		return
	}

	// Parse the last two .groups, if they're there
	lastItem := result[len(result)-1][1]
	var nextToLastItem string
	if len(result) == 1 {
		nextToLastItem = ""
	} else {
		nextToLastItem = result[len(result)-2][1]
	}

	// Should we make an archive? E.g. tar or zip?
	if nextToLastItem == "tar" {
		config.Archive = "tar"
	}
	if lastItem == "zip" || lastItem == "tar" {
		config.Archive = lastItem
		// Tar or zip is our final artifact. Bail out.
		return
	}

	// Should we compress the artifact?
	algorithm, ok := extensions[lastItem]
	if ok {
		config.Algorithm = algorithm
		// We found our compression algorithm. Bail out.
		return
	}

	// We didn't match a known compression format. Default to tar + pgzip
	config.Algorithm = "pgzip"
	config.Archive = "tar"
	return
}

func makeLZ4Writer(output io.WriteCloser, compressionLevel int) (io.WriteCloser, error) {
	lzwriter := lz4.NewWriter(output)
	if compressionLevel > gzip.DefaultCompression {
		lzwriter.Header.HighCompression = true
	}
	return lzwriter, nil
}

func makePgzipWriter(output io.WriteCloser, compressionLevel int) (io.WriteCloser, error) {
	gzipWriter, err := pgzip.NewWriterLevel(output, compressionLevel)
	if err != nil {
		return nil, ErrInvalidCompressionLevel
	}
	gzipWriter.SetConcurrency(500000, runtime.GOMAXPROCS(-1))
	return gzipWriter, nil
}

func createTarArchive(files []string, output io.WriteCloser) error {
	archive := tar.NewWriter(output)
	defer archive.Close()

	for _, path := range files {
		file, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Unable to read file %s: %s", path, err)
		}
		defer file.Close()

		fi, err := file.Stat()
		if err != nil {
			return fmt.Errorf("Unable to get fileinfo for %s: %s", path, err)
		}

		header, err := tar.FileInfoHeader(fi, path)
		if err != nil {
			return fmt.Errorf("Failed to create tar header for %s: %s", path, err)
		}

		if err := archive.WriteHeader(header); err != nil {
			return fmt.Errorf("Failed to write tar header for %s: %s", path, err)
		}

		if _, err := io.Copy(archive, file); err != nil {
			return fmt.Errorf("Failed to copy %s data to archive: %s", path, err)
		}
	}
	return nil
}

func createZipArchive(files []string, output io.WriteCloser) error {
	archive := zip.NewWriter(output)
	defer archive.Close()

	for _, path := range files {
		path = filepath.ToSlash(path)

		source, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("Unable to read file %s: %s", path, err)
		}
		defer source.Close()

		target, err := archive.Create(path)
		if err != nil {
			return fmt.Errorf("Failed to add zip header for %s: %s", path, err)
		}

		_, err = io.Copy(target, source)
		if err != nil {
			return fmt.Errorf("Failed to copy %s data to archive: %s", path, err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// +build kvdb_etcd

package etcd

import (
	"testing"

	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/stretchr/testify/require"
)

func TestReadCursorEmptyInterval(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("alma"))
		require.NoError(t, err)
		require.NotNil(t, b)

		return nil
	})
	require.NoError(t, err)

	err = db.View(func(tx walletdb.ReadTx) error {
		b := tx.ReadBucket([]byte("alma"))
		require.NotNil(t, b)

		cursor := b.ReadCursor()
		k, v := cursor.First()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Next()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Last()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Prev()
		require.Nil(t, k)
		require.Nil(t, v)

		return nil
	})
	require.NoError(t, err)
}

func TestReadCursorNonEmptyInterval(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	testKeyValues := []KV{
		{"b", "1"},
		{"c", "2"},
		{"da", "3"},
		{"e", "4"},
	}

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("alma"))
		require.NoError(t, err)
		require.NotNil(t, b)

		for _, kv := range testKeyValues {
			require.NoError(t, b.Put([]byte(kv.key), []byte(kv.val)))
		}
		return nil
	})

	require.NoError(t, err)

	err = db.View(func(tx walletdb.ReadTx) error {
		b := tx.ReadBucket([]byte("alma"))
		require.NotNil(t, b)

		// Iterate from the front.
		var kvs []KV
		cursor := b.ReadCursor()
		k, v := cursor.First()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Next()
		}
		require.Equal(t, testKeyValues, kvs)

		// Iterate from the back.
		kvs = []KV{}
		k, v = cursor.Last()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Prev()
		}
		require.Equal(t, reverseKVs(testKeyValues), kvs)

		// Random access
		perm := []int{3, 0, 2, 1}
		for _, i := range perm {
			k, v := cursor.Seek([]byte(testKeyValues[i].key))
			require.Equal(t, []byte(testKeyValues[i].key), k)
			require.Equal(t, []byte(testKeyValues[i].val), v)
		}

		// Seek to nonexisting key.
		k, v = cursor.Seek(nil)
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Seek([]byte("x"))
		require.Nil(t, k)
		require.Nil(t, v)

		return nil
	})

	require.NoError(t, err)
}

func TestReadWriteCursor(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	testKeyValues := []KV{
		{"b", "1"},
		{"c", "2"},
		{"da", "3"},
		{"e", "4"},
	}

	count := len(testKeyValues)

	// Pre-store the first half of the interval.
	require.NoError(t, db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("apple"))
		require.NoError(t, err)
		require.NotNil(t, b)

		for i := 0; i < count/2; i++ {
			err = b.Put(
				[]byte(testKeyValues[i].key),
				[]byte(testKeyValues[i].val),
			)
			require.NoError(t, err)
		}
		return nil
	}))

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b := tx.ReadWriteBucket([]byte("apple"))
		require.NotNil(t, b)

		// Store the second half of the interval.
		for i := count / 2; i < count; i++ {
			err = b.Put(
				[]byte(testKeyValues[i].key),
				[]byte(testKeyValues[i].val),
			)
			require.NoError(t, err)
		}

		cursor := b.ReadWriteCursor()

		// First on valid interval.
		fk, fv := cursor.First()
		require.Equal(t, []byte("b"), fk)
		require.Equal(t, []byte("1"), fv)

		// Prev(First()) = nil
		k, v := cursor.Prev()
		require.Nil(t, k)
		require.Nil(t, v)

		// Last on valid interval.
		lk, lv := cursor.Last()
		require.Equal(t, []byte("e"), lk)
		require.Equal(t, []byte("4"), lv)

		// Next(Last()) = nil
		k, v = cursor.Next()
		require.Nil(t, k)
		require.Nil(t, v)

		// Delete first item, then add an item before the
		// deleted one. Check that First/Next will "jump"
		// over the deleted item and return the new first.
		_, _ = cursor.First()
		require.NoError(t, cursor.Delete())
		require.NoError(t, b.Put([]byte("a"), []byte("0")))
		fk, fv = cursor.First()

		require.Equal(t, []byte("a"), fk)
		require.Equal(t, []byte("0"), fv)

		k, v = cursor.Next()
		require.Equal(t, []byte("c"), k)
		require.Equal(t, []byte("2"), v)

		// Similarly test that a new end is returned if
		// the old end is deleted first.
		_, _ = cursor.Last()
		require.NoError(t, cursor.Delete())
		require.NoError(t, b.Put([]byte("f"), []byte("5")))

		lk, lv = cursor.Last()
		require.Equal(t, []byte("f"), lk)
		require.Equal(t, []byte("5"), lv)

		k, v = cursor.Prev()
		require.Equal(t, []byte("da"), k)
		require.Equal(t, []byte("3"), v)

		// Overwrite k/v in the middle of the interval.
		require.NoError(t, b.Put([]byte("c"), []byte("3")))
		k, v = cursor.Prev()
		require.Equal(t, []byte("c"), k)
		require.Equal(t, []byte("3"), v)

		// Insert new key/values.
		require.NoError(t, b.Put([]byte("cx"), []byte("x")))
		require.NoError(t, b.Put([]byte("cy"), []byte("y")))

		k, v = cursor.Next()
		require.Equal(t, []byte("cx"), k)
		require.Equal(t, []byte("x"), v)

		k, v = cursor.Next()
		require.Equal(t, []byte("cy"), k)
		require.Equal(t, []byte("y"), v)

		expected := []KV{
			{"a", "0"},
			{"c", "3"},
			{"cx", "x"},
			{"cy", "y"},
			{"da", "3"},
			{"f", "5"},
		}

		// Iterate from the front.
		var kvs []KV
		k, v = cursor.First()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Next()
		}
		require.Equal(t, expected, kvs)

		// Iterate from the back.
		kvs = []KV{}
		k, v = cursor.Last()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Prev()
		}
		require.Equal(t, reverseKVs(expected), kvs)

		return nil
	})

	require.NoError(t, err)

	expected := map[string]string{
		bkey("apple"):       bval("apple"),
		vkey("a", "apple"):  "0",
		vkey("c", "apple"):  "3",
		vkey("cx", "apple"): "x",
		vkey("cy", "apple"): "y",
		vkey("da", "apple"): "3",
		vkey("f", "apple"):  "5",
	}
	require.Equal(t, expected, f.Dump())
}
<commit_msg>kvdb: s/hu/en/g<commit_after>// +build kvdb_etcd

package etcd

import (
	"testing"

	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/stretchr/testify/require"
)

func TestReadCursorEmptyInterval(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("apple"))
		require.NoError(t, err)
		require.NotNil(t, b)

		return nil
	})
	require.NoError(t, err)

	err = db.View(func(tx walletdb.ReadTx) error {
		b := tx.ReadBucket([]byte("apple"))
		require.NotNil(t, b)

		cursor := b.ReadCursor()
		k, v := cursor.First()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Next()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Last()
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Prev()
		require.Nil(t, k)
		require.Nil(t, v)

		return nil
	})
	require.NoError(t, err)
}

func TestReadCursorNonEmptyInterval(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	testKeyValues := []KV{
		{"b", "1"},
		{"c", "2"},
		{"da", "3"},
		{"e", "4"},
	}

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("apple"))
		require.NoError(t, err)
		require.NotNil(t, b)

		for _, kv := range testKeyValues {
			require.NoError(t, b.Put([]byte(kv.key), []byte(kv.val)))
		}
		return nil
	})

	require.NoError(t, err)

	err = db.View(func(tx walletdb.ReadTx) error {
		b := tx.ReadBucket([]byte("apple"))
		require.NotNil(t, b)

		// Iterate from the front.
		var kvs []KV
		cursor := b.ReadCursor()
		k, v := cursor.First()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Next()
		}
		require.Equal(t, testKeyValues, kvs)

		// Iterate from the back.
		kvs = []KV{}
		k, v = cursor.Last()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Prev()
		}
		require.Equal(t, reverseKVs(testKeyValues), kvs)

		// Random access
		perm := []int{3, 0, 2, 1}
		for _, i := range perm {
			k, v := cursor.Seek([]byte(testKeyValues[i].key))
			require.Equal(t, []byte(testKeyValues[i].key), k)
			require.Equal(t, []byte(testKeyValues[i].val), v)
		}

		// Seek to nonexisting key.
		k, v = cursor.Seek(nil)
		require.Nil(t, k)
		require.Nil(t, v)

		k, v = cursor.Seek([]byte("x"))
		require.Nil(t, k)
		require.Nil(t, v)

		return nil
	})

	require.NoError(t, err)
}

func TestReadWriteCursor(t *testing.T) {
	t.Parallel()

	f := NewEtcdTestFixture(t)
	defer f.Cleanup()

	db, err := newEtcdBackend(f.BackendConfig())
	require.NoError(t, err)

	testKeyValues := []KV{
		{"b", "1"},
		{"c", "2"},
		{"da", "3"},
		{"e", "4"},
	}

	count := len(testKeyValues)

	// Pre-store the first half of the interval.
	require.NoError(t, db.Update(func(tx walletdb.ReadWriteTx) error {
		b, err := tx.CreateTopLevelBucket([]byte("apple"))
		require.NoError(t, err)
		require.NotNil(t, b)

		for i := 0; i < count/2; i++ {
			err = b.Put(
				[]byte(testKeyValues[i].key),
				[]byte(testKeyValues[i].val),
			)
			require.NoError(t, err)
		}
		return nil
	}))

	err = db.Update(func(tx walletdb.ReadWriteTx) error {
		b := tx.ReadWriteBucket([]byte("apple"))
		require.NotNil(t, b)

		// Store the second half of the interval.
		for i := count / 2; i < count; i++ {
			err = b.Put(
				[]byte(testKeyValues[i].key),
				[]byte(testKeyValues[i].val),
			)
			require.NoError(t, err)
		}

		cursor := b.ReadWriteCursor()

		// First on valid interval.
		fk, fv := cursor.First()
		require.Equal(t, []byte("b"), fk)
		require.Equal(t, []byte("1"), fv)

		// Prev(First()) = nil
		k, v := cursor.Prev()
		require.Nil(t, k)
		require.Nil(t, v)

		// Last on valid interval.
		lk, lv := cursor.Last()
		require.Equal(t, []byte("e"), lk)
		require.Equal(t, []byte("4"), lv)

		// Next(Last()) = nil
		k, v = cursor.Next()
		require.Nil(t, k)
		require.Nil(t, v)

		// Delete first item, then add an item before the
		// deleted one. Check that First/Next will "jump"
		// over the deleted item and return the new first.
		_, _ = cursor.First()
		require.NoError(t, cursor.Delete())
		require.NoError(t, b.Put([]byte("a"), []byte("0")))
		fk, fv = cursor.First()

		require.Equal(t, []byte("a"), fk)
		require.Equal(t, []byte("0"), fv)

		k, v = cursor.Next()
		require.Equal(t, []byte("c"), k)
		require.Equal(t, []byte("2"), v)

		// Similarly test that a new end is returned if
		// the old end is deleted first.
		_, _ = cursor.Last()
		require.NoError(t, cursor.Delete())
		require.NoError(t, b.Put([]byte("f"), []byte("5")))

		lk, lv = cursor.Last()
		require.Equal(t, []byte("f"), lk)
		require.Equal(t, []byte("5"), lv)

		k, v = cursor.Prev()
		require.Equal(t, []byte("da"), k)
		require.Equal(t, []byte("3"), v)

		// Overwrite k/v in the middle of the interval.
		require.NoError(t, b.Put([]byte("c"), []byte("3")))
		k, v = cursor.Prev()
		require.Equal(t, []byte("c"), k)
		require.Equal(t, []byte("3"), v)

		// Insert new key/values.
		require.NoError(t, b.Put([]byte("cx"), []byte("x")))
		require.NoError(t, b.Put([]byte("cy"), []byte("y")))

		k, v = cursor.Next()
		require.Equal(t, []byte("cx"), k)
		require.Equal(t, []byte("x"), v)

		k, v = cursor.Next()
		require.Equal(t, []byte("cy"), k)
		require.Equal(t, []byte("y"), v)

		expected := []KV{
			{"a", "0"},
			{"c", "3"},
			{"cx", "x"},
			{"cy", "y"},
			{"da", "3"},
			{"f", "5"},
		}

		// Iterate from the front.
		var kvs []KV
		k, v = cursor.First()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Next()
		}
		require.Equal(t, expected, kvs)

		// Iterate from the back.
		kvs = []KV{}
		k, v = cursor.Last()

		for k != nil && v != nil {
			kvs = append(kvs, KV{string(k), string(v)})
			k, v = cursor.Prev()
		}
		require.Equal(t, reverseKVs(expected), kvs)

		return nil
	})

	require.NoError(t, err)

	expected := map[string]string{
		bkey("apple"):       bval("apple"),
		vkey("a", "apple"):  "0",
		vkey("c", "apple"):  "3",
		vkey("cx", "apple"): "x",
		vkey("cy", "apple"): "y",
		vkey("da", "apple"): "3",
		vkey("f", "apple"):  "5",
	}
	require.Equal(t, expected, f.Dump())
}
<|endoftext|>