<commit_before>// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/gologit"
	"github.com/gorilla/mux"
	httpclient "github.com/mreiferson/go-httpclient"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Server name used in Headers and Via checks
	ServerName string
	// Default headers to add to each response
	AddHeaders map[string]string
}

// ProxyMetrics interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client *http.Client
	config *Config
	// compiled allow list regex
	allowList []*regexp.Regexp
	metrics   ProxyMetrics
}

// NotFound handles client requests that did not match a valid
// mux route.
func (p *Proxy) NotFoundHandler() http.Handler {
	handler := func(w http.ResponseWriter, req *http.Request) {
		h := w.Header()
		for k, v := range p.config.AddHeaders {
			h.Set(k, v)
		}
		h.Set("Server", p.config.ServerName)
		h.Set("Date", formattedDate.String())
		http.Error(w, "404 Not Found", 404)
	}
	return http.HandlerFunc(handler)
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	// add default headers
	wh := w.Header()
	for k, v := range p.config.AddHeaders {
		wh.Set(k, v)
	}
	wh.Set("Server", p.config.ServerName)
	wh.Set("Date", formattedDate.String())

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	vars := mux.Vars(req)
	sURL, ok := encoding.DecodeURL(p.config.HMACKey, vars["sigHash"], vars["encodedURL"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", sURL)
	gologit.Debugln("Client request:", req)

	u, err := url.Parse(sURL)
	if err != nil {
		gologit.Debugln("url parse error:", err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("user-agent", p.config.ServerName)
	nreq.Header.Add("via", p.config.ServerName)

	gologit.Debugln("Built outgoing request:", nreq)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()
	gologit.Debugln("Response from upstream:", resp)

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		gologit.Debugln("Content length exceeded", sURL)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		opErr, ok := err.(*net.OpError)
		if ok {
			switch opErr.Err {
			case syscall.EPIPE:
				// broken pipe - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			case syscall.ECONNRESET:
				// connection reset by peer - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			default:
				// log anything else normally
				gologit.Println("OpError writing response:", err)
			}
		} else {
			// unknown error and not an OpError. just debug log for now.
			gologit.Debugln("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln("Response to client:", w)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// SetMetricsCollector sets a proxy metrics (ProxyMetrics interface) for
// the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// New returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &httpclient.Transport{
		MaxIdleConnsPerHost: 8,
		ConnectTimeout:      2 * time.Second,
		RequestTimeout:      pc.RequestTimeout}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	client := &http.Client{Transport: tr}
	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			return errors.New("Too many redirects")
		}
		return nil
	}

	var allow []*regexp.Regexp
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(strings.TrimSpace(v))
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		config:    &pc,
		allowList: allow}, nil
}
<commit_msg>log less common errors at standard level<commit_after>// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/gologit"
	"github.com/gorilla/mux"
	httpclient "github.com/mreiferson/go-httpclient"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Server name used in Headers and Via checks
	ServerName string
	// Default headers to add to each response
	AddHeaders map[string]string
}

// ProxyMetrics interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client *http.Client
	config *Config
	// compiled allow list regex
	allowList []*regexp.Regexp
	metrics   ProxyMetrics
}

// NotFound handles client requests that did not match a valid
// mux route.
func (p *Proxy) NotFoundHandler() http.Handler {
	handler := func(w http.ResponseWriter, req *http.Request) {
		h := w.Header()
		for k, v := range p.config.AddHeaders {
			h.Set(k, v)
		}
		h.Set("Server", p.config.ServerName)
		h.Set("Date", formattedDate.String())
		http.Error(w, "404 Not Found", 404)
	}
	return http.HandlerFunc(handler)
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	// add default headers
	wh := w.Header()
	for k, v := range p.config.AddHeaders {
		wh.Set(k, v)
	}
	wh.Set("Server", p.config.ServerName)
	wh.Set("Date", formattedDate.String())

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	vars := mux.Vars(req)
	sURL, ok := encoding.DecodeURL(p.config.HMACKey, vars["sigHash"], vars["encodedURL"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", sURL)
	gologit.Debugln("Client request:", req)

	u, err := url.Parse(sURL)
	if err != nil {
		gologit.Debugln("url parse error:", err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url", http.StatusNotFound)
		return
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("user-agent", p.config.ServerName)
	nreq.Header.Add("via", p.config.ServerName)

	gologit.Debugln("Built outgoing request:", nreq)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		if strings.Contains(err.Error(), "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()
	gologit.Debugln("Response from upstream:", resp)

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		gologit.Debugln("Content length exceeded", sURL)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		opErr, ok := err.(*net.OpError)
		if ok {
			switch opErr.Err {
			case syscall.EPIPE:
				// broken pipe - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			case syscall.ECONNRESET:
				// connection reset by peer - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			default:
				// log anything else normally
				gologit.Println("OpError writing response:", err)
			}
		} else {
			// unknown error and not an OpError.
			gologit.Println("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln("Response to client:", w)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// SetMetricsCollector sets a proxy metrics (ProxyMetrics interface) for
// the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// New returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &httpclient.Transport{
		MaxIdleConnsPerHost: 8,
		ConnectTimeout:      2 * time.Second,
		RequestTimeout:      pc.RequestTimeout}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	client := &http.Client{Transport: tr}
	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			return errors.New("Too many redirects")
		}
		return nil
	}

	var allow []*regexp.Regexp
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(strings.TrimSpace(v))
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		config:    &pc,
		allowList: allow}, nil
}
<|endoftext|>


<commit_before>package day11

import (
	"fmt"
	"strings"

	"regexp"

	"sort"

	"github.com/gsmcwhirter/advent2016/lib"
)

type Piece struct {
	Element string
	Type    string
}

type Move struct {
	TargetFloor int
	Pieces      []Piece
}

func (m *Move) HasPiece(p Piece) bool {
	for _, piece := range m.Pieces {
		if piece.Element == p.Element && piece.Type == p.Type {
			return true
		}
	}

	return false
}

func (m *Move) GeneratorTypes() []string {
	types := []string{}
	for _, piece := range m.Pieces {
		if piece.Type == "generator" {
			types = append(types, piece.Element)
		}
	}

	return types
}

func (m *Move) ChipTypes() []string {
	types := []string{}
	for _, piece := range m.Pieces {
		if piece.Type == "chip" {
			types = append(types, piece.Element)
		}
	}

	return types
}

type Floor struct {
	Num    int
	Pieces []Piece
}

func (f *Floor) Equal(other Floor) bool {
	if f.Num != other.Num {
		return false
	}

	for _, piece := range append(f.Pieces, other.Pieces...) {
		if !f.HasPiece(piece) || !other.HasPiece(piece) {
			return false
		}
	}

	return true
}

func (f *Floor) HasPiece(p Piece) bool {
	for _, piece := range f.Pieces {
		if piece.Element == p.Element && piece.Type == p.Type {
			return true
		}
	}

	return false
}

func (f *Floor) HasGenerator(elem string) bool {
	for _, piece := range f.Pieces {
		if piece.Element == elem && piece.Type == "generator" {
			return true
		}
	}

	return false
}

func (f *Floor) GeneratorTypes(excepts ...string) []string {
	types := []string{}
	for _, piece := range f.Pieces {
		if piece.Type == "generator" {
			toAdd := true
			for _, except := range excepts {
				if except == piece.Element {
					toAdd = false
				}
			}

			if toAdd {
				types = append(types, piece.Element)
			}
		}
	}

	return types
}

func (f *Floor) HasChip(elem string) bool {
	for _, piece := range f.Pieces {
		if piece.Element == elem && piece.Type == "chip" {
			return true
		}
	}

	return false
}

func (f *Floor) ChipTypes(excepts ...string) []string {
	types := []string{}
	for _, piece := range f.Pieces {
		if piece.Type == "chip" {
			toAdd := true
			for _, except := range excepts {
				if except == piece.Element {
					toAdd = false
				}
			}

			if toAdd {
				types = append(types, piece.Element)
			}
		}
	}

	return types
}

type State struct {
	Elevator  int
	Floors    []Floor
	Elements  []string
	MoveCount int
}

func (s *State) GenerateMoves(history []State) []Move {
	moves := []Move{}

	floor := s.Floors[s.Elevator-1]

	var move Move
	for i, piece := range floor.Pieces {
		move = Move{s.Elevator + 1, []Piece{piece}}
		if s.MoveIsValid(move, history) {
			moves = append(moves, move)
		}

		move = Move{s.Elevator - 1, []Piece{piece}}
		if s.MoveIsValid(move, history) {
			moves = append(moves, move)
		}

		for j, piece2 := range floor.Pieces {
			if i == j {
				continue
			}

			move = Move{s.Elevator + 1, []Piece{piece, piece2}}
			if s.MoveIsValid(move, history) {
				moves = append(moves, move)
			}

			move = Move{s.Elevator - 1, []Piece{piece, piece2}}
			if s.MoveIsValid(move, history) {
				moves = append(moves, move)
			}
		}
	}

	return moves
}

func (s *State) MoveIsValid(move Move, history []State) bool {
	// Are we moving within bounds?
	if move.TargetFloor > 4 || move.TargetFloor < 1 {
		return false
	}

	// Are we moving enough things?
	if len(move.Pieces) < 1 || len(move.Pieces) > 2 {
		return false
	}

	currentFloor := s.Floors[s.Elevator-1]
	targetFloor := s.Floors[move.TargetFloor-1]

	moveGenerators := move.GeneratorTypes()
	moveChips := move.ChipTypes()

	// Can we move these things together?
	if len(moveGenerators) > 0 && len(moveChips) > 0 {
		for _, gen := range moveGenerators {
			for _, chip := range moveChips {
				if gen == chip {
					return true
				}
			}
		}

		return false
	}

	//Will removing these things fry the current floor?
	currentFloorRemainingGenerators := currentFloor.GeneratorTypes(moveGenerators...)
	for _, chip := range currentFloor.ChipTypes(moveChips...) {
		if !currentFloor.HasGenerator(chip) && len(currentFloorRemainingGenerators) > 0 {
			return false
		}
	}

	// Will we fry something on the target floor?
	for _, gen := range moveGenerators {
		if len(targetFloor.ChipTypes(gen)) > 0 {
			return false
		}
	}

	// Will we fry a chip in the elevator on the target floor?
	for _, chip := range moveChips {
		if !targetFloor.HasGenerator(chip) && len(targetFloor.GeneratorTypes(chip)) > 0 {
			return false
		}
	}

	newState := AdvanceState(*s, move)
	if StateInHistory(history, newState) {
		return false
	}

	return true
}

func (s *State) Print() {
	var floor Floor
	for floorNum := len(s.Floors); floorNum > 0; floorNum-- {
		floor = s.Floors[floorNum-1]
		fmt.Printf("F%v", floorNum)
		if s.Elevator == floorNum {
			fmt.Print(" E")
		} else {
			fmt.Print(" .")
		}

		for _, elem := range s.Elements {
			if floor.HasGenerator(elem) {
				fmt.Printf(" %v-G", string([]rune(elem)[0]))
			} else {
				fmt.Print("  . ")
			}

			if floor.HasChip(elem) {
				fmt.Printf(" %v-C", string([]rune(elem)[0]))
			} else {
				fmt.Print("  . ")
			}
		}

		fmt.Println()
	}
	fmt.Printf("\nMoves: %v\n\n", s.MoveCount)
}

func (s *State) GoalReached() bool {
	for _, floor := range s.Floors[:3] {
		if len(floor.Pieces) > 0 {
			return false
		}
	}

	return true
}

func (s *State) Equal(other State) bool {
	if s.Elevator != other.Elevator {
		return false
	}

	for i, sflr := range s.Floors {
		oflr := other.Floors[i]

		if !sflr.Equal(oflr) {
			return false
		}
	}

	return true
}

func StateInHistory(history []State, ns State) bool {
	for _, oldS := range history {
		if oldS.Equal(ns) {
			return true
		}
	}

	return false
}

func AdvanceState(oldState State, move Move) State {
	// oldState.Print()

	newFloors := []Floor{}

	for i, floor := range oldState.Floors {
		// fmt.Print(i, " ")
		// fmt.Println(floor)
		switch i + 1 {
		case oldState.Elevator:
			// fmt.Println("Cleaning out old floor")
			newPieces := []Piece{}
			for _, piece := range floor.Pieces {
				if !move.HasPiece(piece) {
					newPieces = append(newPieces, piece)
				}
			}

			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: newPieces,
			})
		case move.TargetFloor:
			// fmt.Println("Creating new floor")
			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: append(floor.Pieces[:], move.Pieces...),
			})
		default:
			// fmt.Println("Floor unchanged")
			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: floor.Pieces[:],
			})
		}
	}

	newState := State{
		Elevator:  move.TargetFloor,
		Floors:    newFloors,
		Elements:  oldState.Elements,
		MoveCount: oldState.MoveCount + 1,
	}

	// newState.Print()

	return newState
}

type MoveStatePair struct {
	Move  Move
	State State
}

type MoveStateQueue struct {
	Items []*MoveStatePair
}

func (q *MoveStateQueue) push(move Move, state State) {
	msp := &MoveStatePair{move, state}
	q.Items = append(q.Items, msp)
}

func (q *MoveStateQueue) pop() *MoveStatePair {
	if len(q.Items) == 0 {
		return nil
	}

	item := q.Items[0]
	q.Items = q.Items[1:]

	return item
}

func BFS(startState State) {
	state := startState
	queue := MoveStateQueue{[]*MoveStatePair{}}
	history := []State{state}

	newMoves := state.GenerateMoves(history)
	// fmt.Printf("Adding %v moves to the queue (current length %v)\n", len(newMoves), len(newMoves)+len(queue.Items))
	for _, move := range newMoves {
		queue.push(move, state)
	}

	item := queue.pop()
	var newState State
	for item != nil {
		// fmt.Println(item.Move)
		newState = AdvanceState(item.State, item.Move)
		// newState.Print()

		if StateInHistory(history, newState) {
			// fmt.Println("Skipping state")
			item = queue.pop()
			continue
		} else {
			history = append(history, newState)
		}

		if newState.GoalReached() {
			newState.Print()
			break
		}

		newMoves := newState.GenerateMoves(history)
		// fmt.Printf("Adding %v moves to the queue (current length %v)\n", len(newMoves), len(newMoves)+len(queue.Items))
		for _, move := range newMoves {
			queue.push(move, newState)
		}

		item = queue.pop()
	}
}

func ParseDataString(dat string) ([]Floor, []string) {
	floors := []Floor{}
	typesMap := map[string]bool{}

	generatorRegex := regexp.MustCompile("a ([^\\s]+) generator")
	chipRegex := regexp.MustCompile("a ([^\\s]+)-compatible microchip")

	for num, line := range strings.Split(dat, "\n") {
		pieces := []Piece{}
		generatorMatches := generatorRegex.FindAllStringSubmatch(line, -1)
		for _, match := range generatorMatches {
			pieces = append(pieces, Piece{
				Element: match[1],
				Type:    "generator",
			})
			typesMap[match[1]] = true
		}

		chipMatches := chipRegex.FindAllStringSubmatch(line, -1)
		for _, match := range chipMatches {
			pieces = append(pieces, Piece{
				Element: match[1],
				Type:    "chip",
			})
			typesMap[match[1]] = true
		}

		floors = append(floors, Floor{
			Num:    num + 1,
			Pieces: pieces,
		})
	}

	types := make([]string, len(typesMap))
	i := -1

	for key := range typesMap {
		i++
		types[i] = key
	}

	sort.Strings(types)

	return floors, types
}

func LoadData(filename string) ([]Floor, []string) {
	dat := lib.ReadFileData(filename)

	return ParseDataString(strings.Trim(string(dat), "\n"))
}

func RunPartA(filename string) {
	floors, allTypes := LoadData(filename)

	state := State{
		Elevator:  1,
		Floors:    floors,
		Elements:  allTypes,
		MoveCount: 0,
	}

	BFS(state)
}

func RunPartB(filename string) {

}
<commit_msg>Try to identify isomorphic states instead of simply equal states<commit_after>package day11

import (
	"fmt"
	"strings"

	"regexp"

	"sort"

	"github.com/gsmcwhirter/advent2016/lib"
)

type Piece struct {
	Element string
	Type    string
}

type Move struct {
	TargetFloor int
	Pieces      []Piece
}

func (m *Move) HasPiece(p Piece) bool {
	for _, piece := range m.Pieces {
		if piece.Element == p.Element && piece.Type == p.Type {
			return true
		}
	}

	return false
}

func (m *Move) GeneratorTypes() []string {
	types := []string{}
	for _, piece := range m.Pieces {
		if piece.Type == "generator" {
			types = append(types, piece.Element)
		}
	}

	return types
}

func (m *Move) ChipTypes() []string {
	types := []string{}
	for _, piece := range m.Pieces {
		if piece.Type == "chip" {
			types = append(types, piece.Element)
		}
	}

	return types
}

type Floor struct {
	Num    int
	Pieces []Piece
}

func (f *Floor) Equal(other Floor) bool {
	if f.Num != other.Num {
		return false
	}

	for _, piece := range append(f.Pieces, other.Pieces...) {
		if !f.HasPiece(piece) || !other.HasPiece(piece) {
			return false
		}
	}

	return true
}

func (f *Floor) HasPiece(p Piece) bool {
	for _, piece := range f.Pieces {
		if piece.Element == p.Element && piece.Type == p.Type {
			return true
		}
	}

	return false
}

func (f *Floor) HasGenerator(elem string) bool {
	for _, piece := range f.Pieces {
		if piece.Element == elem && piece.Type == "generator" {
			return true
		}
	}

	return false
}

func (f *Floor) GeneratorTypes(excepts ...string) []string {
	types := []string{}
	for _, piece := range f.Pieces {
		if piece.Type == "generator" {
			toAdd := true
			for _, except := range excepts {
				if except == piece.Element {
					toAdd = false
				}
			}

			if toAdd {
				types = append(types, piece.Element)
			}
		}
	}

	return types
}

func (f *Floor) HasChip(elem string) bool {
	for _, piece := range f.Pieces {
		if piece.Element == elem && piece.Type == "chip" {
			return true
		}
	}

	return false
}

func (f *Floor) ChipTypes(excepts ...string) []string {
	types := []string{}
	for _, piece := range f.Pieces {
		if piece.Type == "chip" {
			toAdd := true
			for _, except := range excepts {
				if except == piece.Element {
					toAdd = false
				}
			}

			if toAdd {
				types = append(types, piece.Element)
			}
		}
	}

	return types
}

func (f *Floor) NumPairs() int {
	gens := f.GeneratorTypes("")

	count := 0
	for _, gen := range gens {
		if f.HasChip(gen) {
			count++
		}
	}

	return count
}

func (f *Floor) NumUnpairedChips() int {
	chips := f.ChipTypes("")

	count := 0
	for _, chip := range chips {
		if !f.HasGenerator(chip) {
			count++
		}
	}

	return count
}

func (f *Floor) NumUnpairedGenerators() int {
	gens := f.GeneratorTypes("")

	count := 0
	for _, gen := range gens {
		if !f.HasChip(gen) {
			count++
		}
	}

	return count
}

type State struct {
	Elevator  int
	Floors    []Floor
	Elements  []string
	MoveCount int
}

func (s *State) GenerateMoves(history []State) []Move {
	moves := []Move{}

	floor := s.Floors[s.Elevator-1]

	var move Move
	for i, piece := range floor.Pieces {
		move = Move{s.Elevator + 1, []Piece{piece}}
		if s.MoveIsValid(move, history) {
			moves = append(moves, move)
		}

		move = Move{s.Elevator - 1, []Piece{piece}}
		if s.MoveIsValid(move, history) {
			moves = append(moves, move)
		}

		for j, piece2 := range floor.Pieces {
			if i == j {
				continue
			}

			move = Move{s.Elevator + 1, []Piece{piece, piece2}}
			if s.MoveIsValid(move, history) {
				moves = append(moves, move)
			}

			move = Move{s.Elevator - 1, []Piece{piece, piece2}}
			if s.MoveIsValid(move, history) {
				moves = append(moves, move)
			}
		}
	}

	return moves
}

func (s *State) MoveIsValid(move Move, history []State) bool {
	// Are we moving within bounds?
	if move.TargetFloor > 4 || move.TargetFloor < 1 {
		return false
	}

	// Are we moving enough things?
	if len(move.Pieces) < 1 || len(move.Pieces) > 2 {
		return false
	}

	currentFloor := s.Floors[s.Elevator-1]
	targetFloor := s.Floors[move.TargetFloor-1]

	moveGenerators := move.GeneratorTypes()
	moveChips := move.ChipTypes()

	// Can we move these things together?
	if len(moveGenerators) > 0 && len(moveChips) > 0 {
		for _, gen := range moveGenerators {
			for _, chip := range moveChips {
				if gen == chip {
					return true
				}
			}
		}

		return false
	}

	//Will removing these things fry the current floor?
	currentFloorRemainingGenerators := currentFloor.GeneratorTypes(moveGenerators...)
	for _, chip := range currentFloor.ChipTypes(moveChips...) {
		if !currentFloor.HasGenerator(chip) && len(currentFloorRemainingGenerators) > 0 {
			return false
		}
	}

	// Will we fry something on the target floor?
	for _, gen := range moveGenerators {
		if len(targetFloor.ChipTypes(gen)) > 0 {
			return false
		}
	}

	// Will we fry a chip in the elevator on the target floor?
	for _, chip := range moveChips {
		if !targetFloor.HasGenerator(chip) && len(targetFloor.GeneratorTypes(chip)) > 0 {
			return false
		}
	}

	newState := AdvanceState(*s, move)
	if StateInHistory(history, newState) {
		return false
	}

	return true
}

func (s *State) Print() {
	var floor Floor
	for floorNum := len(s.Floors); floorNum > 0; floorNum-- {
		floor = s.Floors[floorNum-1]
		fmt.Printf("F%v", floorNum)
		if s.Elevator == floorNum {
			fmt.Print(" E")
		} else {
			fmt.Print(" .")
		}

		for _, elem := range s.Elements {
			if floor.HasGenerator(elem) {
				fmt.Printf(" %v-G", string([]rune(elem)[0]))
			} else {
				fmt.Print("  . ")
			}

			if floor.HasChip(elem) {
				fmt.Printf(" %v-C", string([]rune(elem)[0]))
			} else {
				fmt.Print("  . ")
			}
		}

		fmt.Println()
	}
	fmt.Printf("\nMoves: %v\n\n", s.MoveCount)
}

func (s *State) GoalReached() bool {
	for _, floor := range s.Floors[:3] {
		if len(floor.Pieces) > 0 {
			return false
		}
	}

	return true
}

func (s *State) Equal(other State) bool {
	if s.Elevator != other.Elevator {
		return false
	}

	for i, sflr := range s.Floors {
		oflr := other.Floors[i]

		if !sflr.Equal(oflr) {
			return false
		}
	}

	return true
}

func (s *State) Isomorphic(other State) bool {
	if s.Elevator != other.Elevator {
		return false
	}

	for i, sflr := range s.Floors {
		oflr := other.Floors[i]

		if len(sflr.Pieces) != len(oflr.Pieces) {
			return false
		}

		if sflr.NumPairs() != oflr.NumPairs() {
			return false
		}

		if sflr.NumUnpairedChips() != oflr.NumUnpairedChips() {
			return false
		}

		// don't need to check unpaired generators at this point -- they must match
	}

	return true
}

func StateInHistory(history []State, ns State) bool {
	for _, oldS := range history {
		if oldS.Isomorphic(ns) { // was Equal
			return true
		}
	}

	return false
}

func AdvanceState(oldState State, move Move) State {
	// oldState.Print()

	newFloors := []Floor{}

	for i, floor := range oldState.Floors {
		// fmt.Print(i, " ")
		// fmt.Println(floor)
		switch i + 1 {
		case oldState.Elevator:
			// fmt.Println("Cleaning out old floor")
			newPieces := []Piece{}
			for _, piece := range floor.Pieces {
				if !move.HasPiece(piece) {
					newPieces = append(newPieces, piece)
				}
			}

			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: newPieces,
			})
		case move.TargetFloor:
			// fmt.Println("Creating new floor")
			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: append(floor.Pieces[:], move.Pieces...),
			})
		default:
			// fmt.Println("Floor unchanged")
			newFloors = append(newFloors, Floor{
				Num:    floor.Num,
				Pieces: floor.Pieces[:],
			})
		}
	}

	newState := State{
		Elevator:  move.TargetFloor,
		Floors:    newFloors,
		Elements:  oldState.Elements,
		MoveCount: oldState.MoveCount + 1,
	}

	// newState.Print()

	return newState
}

type MoveStatePair struct {
	Move  Move
	State State
}

type MoveStateQueue struct {
	Items []*MoveStatePair
}

func (q *MoveStateQueue) push(move Move, state State) {
	msp := &MoveStatePair{move, state}
	q.Items = append(q.Items, msp)
}

func (q *MoveStateQueue) pop() *MoveStatePair {
	if len(q.Items) == 0 {
		return nil
	}

	item := q.Items[0]
	q.Items = q.Items[1:]

	return item
}

func BFS(startState State) {
	state := startState
	queue := MoveStateQueue{[]*MoveStatePair{}}
	history := []State{state}

	newMoves := state.GenerateMoves(history)
	// fmt.Printf("Adding %v moves to the queue (current length %v)\n", len(newMoves), len(newMoves)+len(queue.Items))
	for _, move := range newMoves {
		queue.push(move, state)
	}

	item := queue.pop()
	var newState State
	for item != nil {
		// fmt.Println(item.Move)
		newState = AdvanceState(item.State, item.Move)
		// newState.Print()

		if StateInHistory(history, newState) {
			// fmt.Println("Skipping state")
			item = queue.pop()
			continue
		} else {
			history = append(history, newState)
		}

		if newState.GoalReached() {
			newState.Print()
			break
		}

		newMoves := newState.GenerateMoves(history)
		// fmt.Printf("Adding %v moves to the queue (current length %v)\n", len(newMoves), len(newMoves)+len(queue.Items))
		for _, move := range newMoves {
			queue.push(move, newState)
		}

		item = queue.pop()
	}
}

func ParseDataString(dat string) ([]Floor, []string) {
	floors := []Floor{}
	typesMap := map[string]bool{}

	generatorRegex := regexp.MustCompile("a ([^\\s]+) generator")
	chipRegex := regexp.MustCompile("a ([^\\s]+)-compatible microchip")

	for num, line := range strings.Split(dat, "\n") {
		pieces := []Piece{}
		generatorMatches := generatorRegex.FindAllStringSubmatch(line, -1)
		for _, match := range generatorMatches {
			pieces = append(pieces, Piece{
				Element: match[1],
				Type:    "generator",
			})
			typesMap[match[1]] = true
		}

		chipMatches := chipRegex.FindAllStringSubmatch(line, -1)
		for _, match := range chipMatches {
			pieces = append(pieces, Piece{
				Element: match[1],
				Type:    "chip",
			})
			typesMap[match[1]] = true
		}

		floors = append(floors, Floor{
			Num:    num + 1,
			Pieces: pieces,
		})
	}

	types := make([]string, len(typesMap))
	i := -1

	for key := range typesMap {
		i++
		types[i] = key
	}

	sort.Strings(types)

	return floors, types
}

func LoadData(filename string) ([]Floor, []string) {
	dat := lib.ReadFileData(filename)

	return ParseDataString(strings.Trim(string(dat), "\n"))
}

func RunPartA(filename string) {
	floors, allTypes := LoadData(filename)

	state := State{
		Elevator:  1,
		Floors:    floors,
		Elements:  allTypes,
		MoveCount: 0,
	}

	BFS(state)
}

func RunPartB(filename string) {

}
<|endoftext|>


<commit_before>package db

import (
	"fmt"
	"strings"
	"time"
)

type Archive struct {
	UUID           string `json:"uuid"            mbus:"uuid"`
	TenantUUID     string `json:"tenant_uuid"     mbus:"tenant_uuid"`
	TargetUUID     string `json:"target_uuid"     mbus:"target_uuid"`
	StoreUUID      string `json:"store_uuid"      mbus:"store_uuid"`
	StoreKey       string `json:"key"             mbus:"key"`
	TakenAt        int64  `json:"taken_at"        mbus:"taken_at"`
	ExpiresAt      int64  `json:"expires_at"      mbus:"expires_at"`
	Notes          string `json:"notes"           mbus:"notes"`
	Status         string `json:"status"          mbus:"status"`
	PurgeReason    string `json:"purge_reason"    mbus:"purge_reason"`
	EncryptionType string `json:"encryption_type" mbus:"encryption_type"`
	Compression    string `json:"compression"     mbus:"compression"`
	Size           int64  `json:"size"            mbus:"size"`

	TargetName     string `json:"target_name"`
	TargetPlugin   string `json:"target_plugin"`
	TargetEndpoint string `json:"target_endpoint"`
	StoreName      string `json:"store_name"`
	StorePlugin    string `json:"store_plugin"`
	StoreEndpoint  string `json:"store_endpoint"`
	StoreAgent     string `json:"store_agent"`
	Job            string `json:"job"`
}

type ArchiveFilter struct {
	UUID          string
	ExactMatch    bool
	ForTarget     string
	ForStore      string
	Before        *time.Time
	After         *time.Time
	ExpiresBefore *time.Time
	ExpiresAfter  *time.Time
	WithStatus    []string
	WithOutStatus []string
	ForTenant     string
	Limit         int
}

func (f *ArchiveFilter) Query() (string, []interface{}) {
	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}
	if f.UUID != "" {
		if f.ExactMatch {
			wheres = append(wheres, "a.uuid = ?")
			args = append(args, f.UUID)
		} else {
			wheres = append(wheres, "a.uuid LIKE ? ESCAPE '/'")
			args = append(args, PatternPrefix(f.UUID))
		}
	}
	if f.ForTarget != "" {
		wheres = append(wheres, "target_uuid = ?")
		args = append(args, f.ForTarget)
	}
	if f.ForStore != "" {
		wheres = append(wheres, "store_uuid = ?")
		args = append(args, f.ForStore)
	}
	if f.Before != nil {
		wheres = append(wheres, "taken_at <= ?")
		args = append(args, f.Before.Unix())
	}
	if f.After != nil {
		wheres = append(wheres, "taken_at >= ?")
		args = append(args, f.After.Unix())
	}
	if len(f.WithStatus) > 0 {
		var params []string
		for _, e := range f.WithStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status IN (%s)", strings.Join(params, ", ")))
	}
	if len(f.WithOutStatus) > 0 {
		var params []string
		for _, e := range f.WithOutStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status NOT IN (%s)", strings.Join(params, ", ")))
	}
	if f.ExpiresBefore != nil {
		wheres = append(wheres, "expires_at < ?")
		args = append(args, f.ExpiresBefore.Unix())
	}

	if f.ForTenant != "" {
		wheres = append(wheres, "a.tenant_uuid = ?")
		args = append(args, f.ForTenant)
	}
	limit := ""
	if f.Limit > 0 {
		limit = " LIMIT ?"
		args = append(args, f.Limit)
	}

	return `
		SELECT a.uuid, a.store_key,
		       a.taken_at, a.expires_at, a.notes,
		       s.uuid, s.name, s.plugin, s.endpoint,
		       a.status, a.purge_reason, a.job, a.encryption_type,
		       a.compression, a.tenant_uuid, a.size

		FROM archives a
		   LEFT JOIN targets t   ON t.uuid = a.target_uuid
		   INNER JOIN stores  s   ON s.uuid = a.store_uuid

		WHERE ` + strings.Join(wheres, " AND ") + `
		ORDER BY a.taken_at DESC, a.uuid ASC
	` + limit, args
}

func (db *DB) CountArchives(filter *ArchiveFilter) (int, error) {
	if filter == nil {
		filter = &ArchiveFilter{}
	}

	query, args := filter.Query()
	r, err := db.query(query, args...)
	if err != nil {
		return 0, err
	}
	defer r.Close()

	i := 0
	for r.Next() {
		i++
	}
	return i, nil
}

func (db *DB) GetAllArchives(filter *ArchiveFilter) ([]*Archive, error) {
	if filter == nil {
		filter = &ArchiveFilter{}
	}

	l := []*Archive{}
	query, args := filter.Query()
	r, err := db.query(query, args...)
	if err != nil {
		return l, err
	}
	defer r.Close()

	for r.Next() {
		a := &Archive{}

		var takenAt, expiresAt, size *int64
		var targetName, storeName *string
		if err = r.Scan(
			&a.UUID, &a.StoreKey, &takenAt, &expiresAt, &a.Notes,
			&a.TargetUUID, &targetName, &a.TargetPlugin, &a.TargetEndpoint,
			&a.StoreUUID, &storeName, &a.StorePlugin, &a.StoreEndpoint,
			&a.Status, &a.PurgeReason, &a.Job, &a.EncryptionType,
			&a.Compression, &a.TenantUUID, &size); err != nil {

			return l, err
		}
		if takenAt != nil {
			a.TakenAt = *takenAt
		}
		if expiresAt != nil {
			a.ExpiresAt = *expiresAt
		}
		if targetName != nil {
			a.TargetName = *targetName
		}
		if storeName != nil {
			a.StoreName = *storeName
		}
		if size != nil {
			a.Size = *size
		}

		l = append(l, a)
	}

	return l, nil
}

func (db *DB) GetArchive(id string) (*Archive, error) {
	r, err := db.query(`
		SELECT a.uuid, a.store_key,
		       a.taken_at, a.expires_at, a.notes,
		       t.uuid, t.name, t.plugin, t.endpoint,
		       s.uuid, s.name, s.plugin, s.endpoint, s.agent,
		       a.status, a.purge_reason, a.job, a.encryption_type,
		       a.compression, a.tenant_uuid, a.size

		FROM archives a
		   INNER JOIN targets t   ON t.uuid = a.target_uuid
		   INNER JOIN stores  s   ON s.uuid = a.store_uuid

		WHERE a.uuid = ?`, id)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	if !r.Next() {
		return nil, nil
	}
	a := &Archive{}

	var takenAt, expiresAt, size *int64
	var targetName, storeName *string
	if err = r.Scan(
		&a.UUID, &a.StoreKey, &takenAt, &expiresAt, &a.Notes,
		&a.TargetUUID, &targetName, &a.TargetPlugin, &a.TargetEndpoint,
		&a.StoreUUID, &storeName, &a.StorePlugin, &a.StoreEndpoint, &a.StoreAgent,
		&a.Status, &a.PurgeReason, &a.Job, &a.EncryptionType,
		&a.Compression, &a.TenantUUID, &size); err != nil {

		return nil, err
	}
	if takenAt != nil {
		a.TakenAt = *takenAt
	}
	if expiresAt != nil {
		a.ExpiresAt = *expiresAt
	}
	if targetName != nil {
		a.TargetName = *targetName
	}
	if storeName != nil {
		a.StoreName = *storeName
	}
	if size != nil {
		a.Size = *size
	}

	return a, nil
}

func (db *DB) UpdateArchive(update *Archive) error {
	return db.exec(
		`UPDATE archives SET notes = ? WHERE uuid = ?`,
		update.Notes, update.UUID,
	)
}

func (db *DB) AnnotateTargetArchive(target, id, notes string) error {
	return db.exec(
		`UPDATE archives SET notes = ? WHERE uuid = ? AND target_uuid = ?`,
		notes, id, target,
	)
}

func (db *DB) GetArchivesNeedingPurge() ([]*Archive, error) {
	filter := &ArchiveFilter{
		WithOutStatus: []string{"purged", "valid"},
	}
	return db.GetAllArchives(filter)
}

func (db *DB) GetExpiredArchives() ([]*Archive, error) {
	now := time.Now()
	filter := &ArchiveFilter{
		ExpiresBefore: &now,
		WithStatus:    []string{"valid"},
	}
	return db.GetAllArchives(filter)
}

func (db *DB) InvalidateArchive(id string) error {
	return db.exec(`UPDATE archives SET status = 'invalid' WHERE uuid = ?`, id)
}

func (db *DB) PurgeArchive(id string) error {
	a, err := db.GetArchive(id)
	if err != nil {
		return fmt.Errorf("Error grabbing archive %s", id)
	}

	if a.Status == "valid" {
		return fmt.Errorf("Invalid attempt to purge a 'valid' archive detected")
	}

	err = db.exec(`UPDATE archives SET purge_reason = status WHERE uuid = ?`, id)
	if err != nil {
		return err
	}

	return db.exec(`UPDATE archives SET status = 'purged' WHERE uuid = ?`, id)
}

func (db *DB) ExpireArchive(id string) error {
	return db.exec(`UPDATE archives SET status = 'expired' WHERE uuid = ?`, id)
}

func (db *DB) DeleteArchive(id string) (bool, error) {
	return true, db.exec(`DELETE FROM archives WHERE uuid = ?`, id)
}

func (db *DB) ArchiveStorageFootprint(filter *ArchiveFilter) (int64, error) {
	var i int64

	if filter == nil {
		filter = &ArchiveFilter{}
	}

	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}
	if filter.ForTarget != "" {
		wheres = append(wheres, "target_uuid = ?")
		args = append(args, filter.ForTarget)
	}
	if filter.ForStore != "" {
		wheres = append(wheres, "store_uuid = ?")
		args = append(args, filter.ForStore)
	}
	if filter.Before != nil {
		wheres = append(wheres, "taken_at <= ?")
		args = append(args, filter.Before.Unix())
	}
	if filter.After != nil {
		wheres = append(wheres, "taken_at >= ?")
		args = append(args, filter.After.Unix())
	}
	if len(filter.WithStatus) > 0 {
		var params []string
		for _, e := range filter.WithStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status IN (%s)", strings.Join(params, ", ")))
	}
	if len(filter.WithOutStatus) > 0 {
		var params []string
		for _, e := range filter.WithOutStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status NOT IN (%s)", strings.Join(params, ", ")))
	}
	if filter.ExpiresBefore != nil {
		wheres = append(wheres, "expires_at <= ?")
		args = append(args, filter.ExpiresBefore.Unix())
	}
	if filter.ExpiresAfter != nil {
		wheres = append(wheres, "expires_at >= ?")
		args = append(args, filter.ExpiresAfter.Unix())
	}
	if filter.ForTenant != "" {
		wheres = append(wheres, "a.tenant_uuid = ?")
		args = append(args, filter.ForTenant)
	}
	limit := ""
	if filter.Limit > 0 {
		limit = " LIMIT ?"
		args = append(args, filter.Limit)
	}

	r, err := db.query(`
		SELECT SUM(a.size)
		FROM archives a
			INNER JOIN targets t   ON t.uuid = a.target_uuid
			INNER JOIN stores  s   ON s.uuid = a.store_uuid
		WHERE `+strings.Join(wheres, " AND ")+limit, args...)
	if err != nil {
		return i, err
	}
	defer r.Close()

	var p *int64
	if !r.Next() {
		return 0, fmt.Errorf("no results from SUM(size) query...")
	}

	if err = r.Scan(&p); err != nil {
		return 0, err
	}
	if p != nil {
		i = *p
	}
	return i, nil
}

func (db *DB) CleanArchives() error {
	return db.exec(`
	   UPDATE archives
	      SET status = "expired"
	    WHERE uuid IN (SELECT a.uuid
	                     FROM archives a
	                LEFT JOIN tenants  t
	                       ON t.uuid = a.tenant_uuid
	                    WHERE t.uuid IS NULL)`)
}
<commit_msg>Fix missing target selects in db/archives.go<commit_after>package db

import (
	"fmt"
	"strings"
	"time"
)

type Archive struct {
	UUID           string `json:"uuid"            mbus:"uuid"`
	TenantUUID     string `json:"tenant_uuid"     mbus:"tenant_uuid"`
	TargetUUID     string `json:"target_uuid"     mbus:"target_uuid"`
	StoreUUID      string `json:"store_uuid"      mbus:"store_uuid"`
	StoreKey       string `json:"key"             mbus:"key"`
	TakenAt        int64  `json:"taken_at"        mbus:"taken_at"`
	ExpiresAt      int64  `json:"expires_at"      mbus:"expires_at"`
	Notes          string `json:"notes"           mbus:"notes"`
	Status         string `json:"status"          mbus:"status"`
	PurgeReason    string `json:"purge_reason"    mbus:"purge_reason"`
	EncryptionType string `json:"encryption_type" mbus:"encryption_type"`
	Compression    string `json:"compression"     mbus:"compression"`
	Size           int64  `json:"size"            mbus:"size"`

	TargetName     string `json:"target_name"`
	TargetPlugin   string `json:"target_plugin"`
	TargetEndpoint string `json:"target_endpoint"`
	StoreName      string `json:"store_name"`
	StorePlugin    string `json:"store_plugin"`
	StoreEndpoint  string `json:"store_endpoint"`
	StoreAgent     string `json:"store_agent"`
	Job            string `json:"job"`
}

type ArchiveFilter struct {
	UUID          string
	ExactMatch    bool
	ForTarget     string
	ForStore      string
	Before        *time.Time
	After         *time.Time
	ExpiresBefore *time.Time
	ExpiresAfter  *time.Time
	WithStatus    []string
	WithOutStatus []string
	ForTenant     string
	Limit         int
}

func (f *ArchiveFilter) Query() (string, []interface{}) {
	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}
	if f.UUID != "" {
		if f.ExactMatch {
			wheres = append(wheres, "a.uuid = ?")
			args = append(args, f.UUID)
		} else {
			wheres = append(wheres, "a.uuid LIKE ? ESCAPE '/'")
			args = append(args, PatternPrefix(f.UUID))
		}
	}
	if f.ForTarget != "" {
		wheres = append(wheres, "target_uuid = ?")
		args = append(args, f.ForTarget)
	}
	if f.ForStore != "" {
		wheres = append(wheres, "store_uuid = ?")
		args = append(args, f.ForStore)
	}
	if f.Before != nil {
		wheres = append(wheres, "taken_at <= ?")
		args = append(args, f.Before.Unix())
	}
	if f.After != nil {
		wheres = append(wheres, "taken_at >= ?")
		args = append(args, f.After.Unix())
	}
	if len(f.WithStatus) > 0 {
		var params []string
		for _, e := range f.WithStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status IN (%s)", strings.Join(params, ", ")))
	}
	if len(f.WithOutStatus) > 0 {
		var params []string
		for _, e := range f.WithOutStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status NOT IN (%s)", strings.Join(params, ", ")))
	}
	if f.ExpiresBefore != nil {
		wheres = append(wheres, "expires_at < ?")
		args = append(args, f.ExpiresBefore.Unix())
	}

	if f.ForTenant != "" {
		wheres = append(wheres, "a.tenant_uuid = ?")
		args = append(args, f.ForTenant)
	}
	limit := ""
	if f.Limit > 0 {
		limit = " LIMIT ?"
		args = append(args, f.Limit)
	}

	return `
		SELECT a.uuid, a.store_key,
		       a.taken_at, a.expires_at, a.notes,
		       t.uuid, t.name, t.plugin, t.endpoint,
		       s.uuid, s.name, s.plugin, s.endpoint,
		       a.status, a.purge_reason, a.job, a.encryption_type,
		       a.compression, a.tenant_uuid, a.size

		FROM archives a
		   LEFT JOIN targets t   ON t.uuid = a.target_uuid
		   INNER JOIN stores  s   ON s.uuid = a.store_uuid

		WHERE ` + strings.Join(wheres, " AND ") + `
		ORDER BY a.taken_at DESC, a.uuid ASC
	` + limit, args
}

func (db *DB) CountArchives(filter *ArchiveFilter) (int, error) {
	if filter == nil {
		filter = &ArchiveFilter{}
	}

	query, args := filter.Query()
	r, err := db.query(query, args...)
	if err != nil {
		return 0, err
	}
	defer r.Close()

	i := 0
	for r.Next() {
		i++
	}
	return i, nil
}

func (db *DB) GetAllArchives(filter *ArchiveFilter) ([]*Archive, error) {
	if filter == nil {
		filter = &ArchiveFilter{}
	}

	l := []*Archive{}
	query, args := filter.Query()
	r, err := db.query(query, args...)
	if err != nil {
		return l, err
	}
	defer r.Close()

	for r.Next() {
		a := &Archive{}

		var takenAt, expiresAt, size *int64
		var targetName, storeName *string
		if err = r.Scan(
			&a.UUID, &a.StoreKey, &takenAt, &expiresAt, &a.Notes,
			&a.TargetUUID, &targetName, &a.TargetPlugin, &a.TargetEndpoint,
			&a.StoreUUID, &storeName, &a.StorePlugin, &a.StoreEndpoint,
			&a.Status, &a.PurgeReason, &a.Job, &a.EncryptionType,
			&a.Compression, &a.TenantUUID, &size); err != nil {

			return l, err
		}
		if takenAt != nil {
			a.TakenAt = *takenAt
		}
		if expiresAt != nil {
			a.ExpiresAt = *expiresAt
		}
		if targetName != nil {
			a.TargetName = *targetName
		}
		if storeName != nil {
			a.StoreName = *storeName
		}
		if size != nil {
			a.Size = *size
		}

		l = append(l, a)
	}

	return l, nil
}

func (db *DB) GetArchive(id string) (*Archive, error) {
	r, err := db.query(`
		SELECT a.uuid, a.store_key,
		       a.taken_at, a.expires_at, a.notes,
		       t.uuid, t.name, t.plugin, t.endpoint,
		       s.uuid, s.name, s.plugin, s.endpoint, s.agent,
		       a.status, a.purge_reason, a.job, a.encryption_type,
		       a.compression, a.tenant_uuid, a.size

		FROM archives a
		   INNER JOIN targets t   ON t.uuid = a.target_uuid
		   INNER JOIN stores  s   ON s.uuid = a.store_uuid

		WHERE a.uuid = ?`, id)
	if err != nil {
		return nil, err
	}
	defer r.Close()

	if !r.Next() {
		return nil, nil
	}
	a := &Archive{}

	var takenAt, expiresAt, size *int64
	var targetName, storeName *string
	if err = r.Scan(
		&a.UUID, &a.StoreKey, &takenAt, &expiresAt, &a.Notes,
		&a.TargetUUID, &targetName, &a.TargetPlugin, &a.TargetEndpoint,
		&a.StoreUUID, &storeName, &a.StorePlugin, &a.StoreEndpoint, &a.StoreAgent,
		&a.Status, &a.PurgeReason, &a.Job, &a.EncryptionType,
		&a.Compression, &a.TenantUUID, &size); err != nil {

		return nil, err
	}
	if takenAt != nil {
		a.TakenAt = *takenAt
	}
	if expiresAt != nil {
		a.ExpiresAt = *expiresAt
	}
	if targetName != nil {
		a.TargetName = *targetName
	}
	if storeName != nil {
		a.StoreName = *storeName
	}
	if size != nil {
		a.Size = *size
	}

	return a, nil
}

func (db *DB) UpdateArchive(update *Archive) error {
	return db.exec(
		`UPDATE archives SET notes = ? WHERE uuid = ?`,
		update.Notes, update.UUID,
	)
}

func (db *DB) AnnotateTargetArchive(target, id, notes string) error {
	return db.exec(
		`UPDATE archives SET notes = ? WHERE uuid = ? AND target_uuid = ?`,
		notes, id, target,
	)
}

func (db *DB) GetArchivesNeedingPurge() ([]*Archive, error) {
	filter := &ArchiveFilter{
		WithOutStatus: []string{"purged", "valid"},
	}
	return db.GetAllArchives(filter)
}

func (db *DB) GetExpiredArchives() ([]*Archive, error) {
	now := time.Now()
	filter := &ArchiveFilter{
		ExpiresBefore: &now,
		WithStatus:    []string{"valid"},
	}
	return db.GetAllArchives(filter)
}

func (db *DB) InvalidateArchive(id string) error {
	return db.exec(`UPDATE archives SET status = 'invalid' WHERE uuid = ?`, id)
}

func (db *DB) PurgeArchive(id string) error {
	a, err := db.GetArchive(id)
	if err != nil {
		return fmt.Errorf("unable to retrieve archive [%s]: %s", id, err)
	}

	if a.Status == "valid" {
		return fmt.Errorf("invalid attempt to purge a 'valid' archive detected")
	}

	err = db.exec(`UPDATE archives SET purge_reason = status WHERE uuid = ?`, id)
	if err != nil {
		return err
	}

	return db.exec(`UPDATE archives SET status = 'purged' WHERE uuid = ?`, id)
}

func (db *DB) ExpireArchive(id string) error {
	return db.exec(`UPDATE archives SET status = 'expired' WHERE uuid = ?`, id)
}

func (db *DB) DeleteArchive(id string) (bool, error) {
	return true, db.exec(`DELETE FROM archives WHERE uuid = ?`, id)
}

func (db *DB) ArchiveStorageFootprint(filter *ArchiveFilter) (int64, error) {
	var i int64

	if filter == nil {
		filter = &ArchiveFilter{}
	}

	wheres := []string{"a.uuid = a.uuid"}
	var args []interface{}
	if filter.ForTarget != "" {
		wheres = append(wheres, "target_uuid = ?")
		args = append(args, filter.ForTarget)
	}
	if filter.ForStore != "" {
		wheres = append(wheres, "store_uuid = ?")
		args = append(args, filter.ForStore)
	}
	if filter.Before != nil {
		wheres = append(wheres, "taken_at <= ?")
		args = append(args, filter.Before.Unix())
	}
	if filter.After != nil {
		wheres = append(wheres, "taken_at >= ?")
		args = append(args, filter.After.Unix())
	}
	if len(filter.WithStatus) > 0 {
		var params []string
		for _, e := range filter.WithStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status IN (%s)", strings.Join(params, ", ")))
	}
	if len(filter.WithOutStatus) > 0 {
		var params []string
		for _, e := range filter.WithOutStatus {
			params = append(params, "?")
			args = append(args, e)
		}
		wheres = append(wheres, fmt.Sprintf("status NOT IN (%s)", strings.Join(params, ", ")))
	}
	if filter.ExpiresBefore != nil {
		wheres = append(wheres, "expires_at <= ?")
		args = append(args, filter.ExpiresBefore.Unix())
	}
	if filter.ExpiresAfter != nil {
		wheres = append(wheres, "expires_at >= ?")
		args = append(args, filter.ExpiresAfter.Unix())
	}
	if filter.ForTenant != "" {
		wheres = append(wheres, "a.tenant_uuid = ?")
		args = append(args, filter.ForTenant)
	}
	limit := ""
	if filter.Limit > 0 {
		limit = " LIMIT ?"
		args = append(args, filter.Limit)
	}

	r, err := db.query(`
		SELECT SUM(a.size)
		FROM archives a
			INNER JOIN targets t   ON t.uuid = a.target_uuid
			INNER JOIN stores  s   ON s.uuid = a.store_uuid
		WHERE `+strings.Join(wheres, " AND ")+limit, args...)
	if err != nil {
		return i, err
	}
	defer r.Close()

	var p *int64
	if !r.Next() {
		return 0, fmt.Errorf("no results from SUM(size) query...")
	}

	if err = r.Scan(&p); err != nil {
		return 0, err
	}
	if p != nil {
		i = *p
	}
	return i, nil
}

func (db *DB) CleanArchives() error {
	return db.exec(`
	   UPDATE archives
	      SET status = "expired"
	    WHERE uuid IN (SELECT a.uuid
	                     FROM archives a
	                LEFT JOIN tenants  t
	                       ON t.uuid = a.tenant_uuid
	                    WHERE t.uuid IS NULL)`)
}
<|endoftext|>


<commit_before>// Copyright 2017 The Prometheus 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 promhttp

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

func makeInstrumentedClient() (*http.Client, *prometheus.Registry) {
	client := http.DefaultClient
	client.Timeout = 1 * time.Second

	reg := prometheus.NewRegistry()

	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "client_in_flight_requests",
		Help: "A gauge of in-flight requests for the wrapped client.",
	})

	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "client_api_requests_total",
			Help: "A counter for requests from the wrapped client.",
		},
		[]string{"code", "method"},
	)

	dnsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "dns_duration_seconds",
			Help:    "Trace dns latency histogram.",
			Buckets: []float64{.005, .01, .025, .05},
		},
		[]string{"event"},
	)

	tlsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "tls_duration_seconds",
			Help:    "Trace tls latency histogram.",
			Buckets: []float64{.05, .1, .25, .5},
		},
		[]string{"event"},
	)

	histVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "request_duration_seconds",
			Help:    "A histogram of request latencies.",
			Buckets: prometheus.DefBuckets,
		},
		[]string{"method"},
	)

	reg.MustRegister(counter, tlsLatencyVec, dnsLatencyVec, histVec, inFlightGauge)

	trace := &InstrumentTrace{
		DNSStart: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_start").Observe(t)
		},
		DNSDone: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_done").Observe(t)
		},
		TLSHandshakeStart: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_start").Observe(t)
		},
		TLSHandshakeDone: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_done").Observe(t)
		},
	}

	client.Transport = InstrumentRoundTripperInFlight(inFlightGauge,
		InstrumentRoundTripperCounter(counter,
			InstrumentRoundTripperTrace(trace,
				InstrumentRoundTripperDuration(histVec, http.DefaultTransport),
			),
		),
	)
	return client, reg
}

func TestClientMiddlewareAPI(t *testing.T) {
	client, reg := makeInstrumentedClient()
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	resp, err := client.Get(backend.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	mfs, err := reg.Gather()
	if err != nil {
		t.Fatal(err)
	}
	if want, got := 3, len(mfs); want != got {
		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
	}
	for _, mf := range mfs {
		if len(mf.Metric) == 0 {
			t.Errorf("metric family %s must not be empty", mf.GetName())
		}
	}
}

func TestClientMiddlewareAPIWithRequestContext(t *testing.T) {
	client, reg := makeInstrumentedClient()
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	req, err := http.NewRequest("GET", backend.URL, nil)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Set a context with a long timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	mfs, err := reg.Gather()
	if err != nil {
		t.Fatal(err)
	}
	if want, got := 3, len(mfs); want != got {
		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
	}
	for _, mf := range mfs {
		if len(mf.Metric) == 0 {
			t.Errorf("metric family %s must not be empty", mf.GetName())
		}
	}
}

func TestClientMiddlewareAPIWithRequestContextTimeout(t *testing.T) {
	client, _ := makeInstrumentedClient()

	// Slow testserver responding in 100ms.
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(100 * time.Millisecond)
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	req, err := http.NewRequest("GET", backend.URL, nil)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Set a context with a short timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	req = req.WithContext(ctx)

	_, err = client.Do(req)
	if err == nil {
		t.Fatal("did not get timeout error")
	}
	if want, got := fmt.Sprintf("Get %s: context deadline exceeded", backend.URL), err.Error(); want != got {
		t.Fatalf("want error %q, got %q", want, got)
	}
}

func ExampleInstrumentRoundTripperDuration() {
	client := http.DefaultClient
	client.Timeout = 1 * time.Second

	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "client_in_flight_requests",
		Help: "A gauge of in-flight requests for the wrapped client.",
	})

	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "client_api_requests_total",
			Help: "A counter for requests from the wrapped client.",
		},
		[]string{"code", "method"},
	)

	// dnsLatencyVec uses custom buckets based on expected dns durations.
	// It has an instance label "event", which is set in the
	// DNSStart and DNSDonehook functions defined in the
	// InstrumentTrace struct below.
	dnsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "dns_duration_seconds",
			Help:    "Trace dns latency histogram.",
			Buckets: []float64{.005, .01, .025, .05},
		},
		[]string{"event"},
	)

	// tlsLatencyVec uses custom buckets based on expected tls durations.
	// It has an instance label "event", which is set in the
	// TLSHandshakeStart and TLSHandshakeDone hook functions defined in the
	// InstrumentTrace struct below.
	tlsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "tls_duration_seconds",
			Help:    "Trace tls latency histogram.",
			Buckets: []float64{.05, .1, .25, .5},
		},
		[]string{"event"},
	)

	// histVec has no labels, making it a zero-dimensional ObserverVec.
	histVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "request_duration_seconds",
			Help:    "A histogram of request latencies.",
			Buckets: prometheus.DefBuckets,
		},
		[]string{},
	)

	// Register all of the metrics in the standard registry.
	prometheus.MustRegister(counter, tlsLatencyVec, dnsLatencyVec, histVec, inFlightGauge)

	// Define functions for the available httptrace.ClientTrace hook
	// functions that we want to instrument.
	trace := &InstrumentTrace{
		DNSStart: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_start").Observe(t)
		},
		DNSDone: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_done").Observe(t)
		},
		TLSHandshakeStart: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_start").Observe(t)
		},
		TLSHandshakeDone: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_done").Observe(t)
		},
	}

	// Wrap the default RoundTripper with middleware.
	roundTripper := InstrumentRoundTripperInFlight(inFlightGauge,
		InstrumentRoundTripperCounter(counter,
			InstrumentRoundTripperTrace(trace,
				InstrumentRoundTripperDuration(histVec, http.DefaultTransport),
			),
		),
	)

	// Set the RoundTripper on our client.
	client.Transport = roundTripper

	resp, err := client.Get("http://google.com")
	if err != nil {
		log.Printf("error: %v", err)
	}
	defer resp.Body.Close()
}
<commit_msg>Fix unit test failing with Go1.14+<commit_after>// Copyright 2017 The Prometheus 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 promhttp

import (
	"context"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

func makeInstrumentedClient() (*http.Client, *prometheus.Registry) {
	client := http.DefaultClient
	client.Timeout = 1 * time.Second

	reg := prometheus.NewRegistry()

	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "client_in_flight_requests",
		Help: "A gauge of in-flight requests for the wrapped client.",
	})

	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "client_api_requests_total",
			Help: "A counter for requests from the wrapped client.",
		},
		[]string{"code", "method"},
	)

	dnsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "dns_duration_seconds",
			Help:    "Trace dns latency histogram.",
			Buckets: []float64{.005, .01, .025, .05},
		},
		[]string{"event"},
	)

	tlsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "tls_duration_seconds",
			Help:    "Trace tls latency histogram.",
			Buckets: []float64{.05, .1, .25, .5},
		},
		[]string{"event"},
	)

	histVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "request_duration_seconds",
			Help:    "A histogram of request latencies.",
			Buckets: prometheus.DefBuckets,
		},
		[]string{"method"},
	)

	reg.MustRegister(counter, tlsLatencyVec, dnsLatencyVec, histVec, inFlightGauge)

	trace := &InstrumentTrace{
		DNSStart: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_start").Observe(t)
		},
		DNSDone: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_done").Observe(t)
		},
		TLSHandshakeStart: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_start").Observe(t)
		},
		TLSHandshakeDone: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_done").Observe(t)
		},
	}

	client.Transport = InstrumentRoundTripperInFlight(inFlightGauge,
		InstrumentRoundTripperCounter(counter,
			InstrumentRoundTripperTrace(trace,
				InstrumentRoundTripperDuration(histVec, http.DefaultTransport),
			),
		),
	)
	return client, reg
}

func TestClientMiddlewareAPI(t *testing.T) {
	client, reg := makeInstrumentedClient()
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	resp, err := client.Get(backend.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	mfs, err := reg.Gather()
	if err != nil {
		t.Fatal(err)
	}
	if want, got := 3, len(mfs); want != got {
		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
	}
	for _, mf := range mfs {
		if len(mf.Metric) == 0 {
			t.Errorf("metric family %s must not be empty", mf.GetName())
		}
	}
}

func TestClientMiddlewareAPIWithRequestContext(t *testing.T) {
	client, reg := makeInstrumentedClient()
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	req, err := http.NewRequest("GET", backend.URL, nil)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Set a context with a long timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	req = req.WithContext(ctx)

	resp, err := client.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()

	mfs, err := reg.Gather()
	if err != nil {
		t.Fatal(err)
	}
	if want, got := 3, len(mfs); want != got {
		t.Fatalf("unexpected number of metric families gathered, want %d, got %d", want, got)
	}
	for _, mf := range mfs {
		if len(mf.Metric) == 0 {
			t.Errorf("metric family %s must not be empty", mf.GetName())
		}
	}
}

func TestClientMiddlewareAPIWithRequestContextTimeout(t *testing.T) {
	client, _ := makeInstrumentedClient()

	// Slow testserver responding in 100ms.
	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(100 * time.Millisecond)
		w.WriteHeader(http.StatusOK)
	}))
	defer backend.Close()

	req, err := http.NewRequest("GET", backend.URL, nil)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Set a context with a short timeout.
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	req = req.WithContext(ctx)

	_, err = client.Do(req)
	if err == nil {
		t.Fatal("did not get timeout error")
	}
	expectedMsg := "context deadline exceeded"
	if !strings.Contains(err.Error(), expectedMsg) {
		t.Fatalf("unexpected error: %q, expect error: %q", err.Error(), expectedMsg)
	}
}

func ExampleInstrumentRoundTripperDuration() {
	client := http.DefaultClient
	client.Timeout = 1 * time.Second

	inFlightGauge := prometheus.NewGauge(prometheus.GaugeOpts{
		Name: "client_in_flight_requests",
		Help: "A gauge of in-flight requests for the wrapped client.",
	})

	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "client_api_requests_total",
			Help: "A counter for requests from the wrapped client.",
		},
		[]string{"code", "method"},
	)

	// dnsLatencyVec uses custom buckets based on expected dns durations.
	// It has an instance label "event", which is set in the
	// DNSStart and DNSDonehook functions defined in the
	// InstrumentTrace struct below.
	dnsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "dns_duration_seconds",
			Help:    "Trace dns latency histogram.",
			Buckets: []float64{.005, .01, .025, .05},
		},
		[]string{"event"},
	)

	// tlsLatencyVec uses custom buckets based on expected tls durations.
	// It has an instance label "event", which is set in the
	// TLSHandshakeStart and TLSHandshakeDone hook functions defined in the
	// InstrumentTrace struct below.
	tlsLatencyVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "tls_duration_seconds",
			Help:    "Trace tls latency histogram.",
			Buckets: []float64{.05, .1, .25, .5},
		},
		[]string{"event"},
	)

	// histVec has no labels, making it a zero-dimensional ObserverVec.
	histVec := prometheus.NewHistogramVec(
		prometheus.HistogramOpts{
			Name:    "request_duration_seconds",
			Help:    "A histogram of request latencies.",
			Buckets: prometheus.DefBuckets,
		},
		[]string{},
	)

	// Register all of the metrics in the standard registry.
	prometheus.MustRegister(counter, tlsLatencyVec, dnsLatencyVec, histVec, inFlightGauge)

	// Define functions for the available httptrace.ClientTrace hook
	// functions that we want to instrument.
	trace := &InstrumentTrace{
		DNSStart: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_start").Observe(t)
		},
		DNSDone: func(t float64) {
			dnsLatencyVec.WithLabelValues("dns_done").Observe(t)
		},
		TLSHandshakeStart: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_start").Observe(t)
		},
		TLSHandshakeDone: func(t float64) {
			tlsLatencyVec.WithLabelValues("tls_handshake_done").Observe(t)
		},
	}

	// Wrap the default RoundTripper with middleware.
	roundTripper := InstrumentRoundTripperInFlight(inFlightGauge,
		InstrumentRoundTripperCounter(counter,
			InstrumentRoundTripperTrace(trace,
				InstrumentRoundTripperDuration(histVec, http.DefaultTransport),
			),
		),
	)

	// Set the RoundTripper on our client.
	client.Transport = roundTripper

	resp, err := client.Get("http://google.com")
	if err != nil {
		log.Printf("error: %v", err)
	}
	defer resp.Body.Close()
}
<|endoftext|>


<commit_before>package builder

// nginx
const (
	NGINX_VERSION             = "1.11.1"
	NGINX_DOWNLOAD_URL_PREFIX = "http://nginx.org/download"
)

// pcre
const (
	PCRE_VERSION             = "8.38"
	PCRE_DOWNLOAD_URL_PREFIX = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OPENSSL_VERSION             = "1.0.2h"
	OPENSSL_DOWNLOAD_URL_PREFIX = "http://www.openssl.org/source"
)

// zlib
const (
	ZLIB_VERSION             = "1.2.8"
	ZLIB_DOWNLOAD_URL_PREFIX = "http://zlib.net"
)

// openResty
const (
	OPENRESTY_VERSION             = "1.9.7.4"
	OPENRESTY_DOWNLOAD_URL_PREFIX = "https://openresty.org/download"
)

// tengine
const (
	TENGINE_VERSION             = "2.1.2"
	TENGINE_DOWNLOAD_URL_PREFIX = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	COMPONENT_NGINX = iota
	COMPONENT_OPENRESTY
	COMPONENT_TENGINE
	COMPONENT_PCRE
	COMPONENT_OPENSSL
	COMPONENT_ZLIB
	COMPONENT_MAX
)
<commit_msg>bumped default openresty version to 1.9.7.5.<commit_after>package builder

// nginx
const (
	NGINX_VERSION             = "1.11.1"
	NGINX_DOWNLOAD_URL_PREFIX = "http://nginx.org/download"
)

// pcre
const (
	PCRE_VERSION             = "8.38"
	PCRE_DOWNLOAD_URL_PREFIX = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OPENSSL_VERSION             = "1.0.2h"
	OPENSSL_DOWNLOAD_URL_PREFIX = "http://www.openssl.org/source"
)

// zlib
const (
	ZLIB_VERSION             = "1.2.8"
	ZLIB_DOWNLOAD_URL_PREFIX = "http://zlib.net"
)

// openResty
const (
	OPENRESTY_VERSION             = "1.9.7.5"
	OPENRESTY_DOWNLOAD_URL_PREFIX = "https://openresty.org/download"
)

// tengine
const (
	TENGINE_VERSION             = "2.1.2"
	TENGINE_DOWNLOAD_URL_PREFIX = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	COMPONENT_NGINX = iota
	COMPONENT_OPENRESTY
	COMPONENT_TENGINE
	COMPONENT_PCRE
	COMPONENT_OPENSSL
	COMPONENT_ZLIB
	COMPONENT_MAX
)
<|endoftext|>


<commit_before>package builder

// nginx
const (
	NginxVersion           = "1.15.6"
	NginxDownloadURLPrefix = "https://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "8.42"
	PcreDownloadURLPrefix = "https://ftp.pcre.org/pub/pcre"
)

// openssl
const (
	OpenSSLVersion           = "1.0.2p"
	OpenSSLDownloadURLPrefix = "https://www.openssl.org/source"
)

// libressl
const (
	LibreSSLVersion           = "2.7.4"
	LibreSSLDownloadURLPrefix = "https://ftp.openbsd.org/pub/OpenBSD/LibreSSL"
)

// zlib
const (
	ZlibVersion           = "1.2.11"
	ZlibDownloadURLPrefix = "https://zlib.net/fossils"
)

// openResty
const (
	OpenRestyVersion           = "1.13.6.2"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.2.2"
	TengineDownloadURLPrefix = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentLibreSSL
	ComponentZlib
	ComponentMax
)
<commit_msg>bumped default openssl version to 1.0.2q.<commit_after>package builder

// nginx
const (
	NginxVersion           = "1.15.6"
	NginxDownloadURLPrefix = "https://nginx.org/download"
)

// pcre
const (
	PcreVersion           = "8.42"
	PcreDownloadURLPrefix = "https://ftp.pcre.org/pub/pcre"
)

// openssl
const (
	OpenSSLVersion           = "1.0.2q"
	OpenSSLDownloadURLPrefix = "https://www.openssl.org/source"
)

// libressl
const (
	LibreSSLVersion           = "2.7.4"
	LibreSSLDownloadURLPrefix = "https://ftp.openbsd.org/pub/OpenBSD/LibreSSL"
)

// zlib
const (
	ZlibVersion           = "1.2.11"
	ZlibDownloadURLPrefix = "https://zlib.net/fossils"
)

// openResty
const (
	OpenRestyVersion           = "1.13.6.2"
	OpenRestyDownloadURLPrefix = "https://openresty.org/download"
)

// tengine
const (
	TengineVersion           = "2.2.2"
	TengineDownloadURLPrefix = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	ComponentNginx = iota
	ComponentOpenResty
	ComponentTengine
	ComponentPcre
	ComponentOpenSSL
	ComponentLibreSSL
	ComponentZlib
	ComponentMax
)
<|endoftext|>


<commit_before>package logmetrics

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"
)

type readStats struct {
	key_pushed  int64
	byte_pushed int64
	last_report time.Time
}

func (f *readStats) inc(data_read int) {
	f.key_pushed++
	f.byte_pushed += int64(data_read)
}

func (f *readStats) getStats() string {
	line_sec := int(f.key_pushed / int64(time.Now().Sub(f.last_report)/time.Second))
	mbyte_sec := float64(f.byte_pushed) / 1024 / 1024 / float64(time.Now().Sub(f.last_report)/time.Second)

	f.key_pushed = 0
	f.byte_pushed = 0
	f.last_report = time.Now()

	return fmt.Sprintf("%d line/s. %.3f Mb/s.",
		line_sec, mbyte_sec)
}

func (f *readStats) isTimeForStats(interval int) bool {
	return (time.Now().Sub(f.last_report) > time.Duration(interval)*time.Second)
}

func parserTest(filename string, logGroup *LogGroup, perfInfo bool) {
	maxMatches := logGroup.expected_matches + 1

	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Unable to tail %s: %s", filename, err)
		return
	}

	scanner := bufio.NewScanner(file)

	log.Printf("Parsing %s", filename)

	read_stats := readStats{last_report: time.Now()}
	for scanner.Scan() {
		line := scanner.Text()

		//Test out all the regexp, pick the first one that matches
		match_one := false
		for _, re := range logGroup.re {
           if matches := re.FindStringSubmatch(line.Text); len(matches) == maxMatches {
				match_one = true
			}
		}

		if match_one {
			read_stats.inc(len(line))
		}

		if read_stats.isTimeForStats(5) {
			log.Print(read_stats.getStats())
		}
	}

	log.Printf("Finished tailling %s.", filename)
}

func startLogGroupParserTest(logGroup *LogGroup, perfInfo bool) {

	newFiles := make(map[string]bool)
	for _, glob := range logGroup.globFiles {
		files, err := filepath.Glob(glob)
		if err != nil {
			log.Fatalf("Unable to find files for log group %s: %s", logGroup.name, err)
		}

		for _, v := range files {
			newFiles[v] = true
		}
	}

	//Start tailing new files!
	for file, _ := range newFiles {
		parserTest(file, logGroup, perfInfo)
	}

}

func StartParserTest(config *Config, selectedLogGroup string, perfInfo bool) {
	for logGroupName, logGroup := range config.logGroups {
		if selectedLogGroup == "" || logGroupName == selectedLogGroup {
			startLogGroupParserTest(logGroup, perfInfo)
		}
	}
}
<commit_msg>Another leftover.<commit_after>package logmetrics

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"
)

type readStats struct {
	key_pushed  int64
	byte_pushed int64
	last_report time.Time
}

func (f *readStats) inc(data_read int) {
	f.key_pushed++
	f.byte_pushed += int64(data_read)
}

func (f *readStats) getStats() string {
	line_sec := int(f.key_pushed / int64(time.Now().Sub(f.last_report)/time.Second))
	mbyte_sec := float64(f.byte_pushed) / 1024 / 1024 / float64(time.Now().Sub(f.last_report)/time.Second)

	f.key_pushed = 0
	f.byte_pushed = 0
	f.last_report = time.Now()

	return fmt.Sprintf("%d line/s. %.3f Mb/s.",
		line_sec, mbyte_sec)
}

func (f *readStats) isTimeForStats(interval int) bool {
	return (time.Now().Sub(f.last_report) > time.Duration(interval)*time.Second)
}

func parserTest(filename string, logGroup *LogGroup, perfInfo bool) {
	maxMatches := logGroup.expected_matches + 1

	file, err := os.Open(filename)
	if err != nil {
		log.Fatalf("Unable to tail %s: %s", filename, err)
		return
	}

	scanner := bufio.NewScanner(file)

	log.Printf("Parsing %s", filename)

	read_stats := readStats{last_report: time.Now()}
	for scanner.Scan() {
		line := scanner.Text()

		//Test out all the regexp, pick the first one that matches
		match_one := false
		for _, re := range logGroup.re {
           if matches := re.FindStringSubmatch(line); len(matches) == maxMatches {
				match_one = true
			}
		}

		if match_one {
			read_stats.inc(len(line))
		}

		if read_stats.isTimeForStats(5) {
			log.Print(read_stats.getStats())
		}
	}

	log.Printf("Finished tailling %s.", filename)
}

func startLogGroupParserTest(logGroup *LogGroup, perfInfo bool) {

	newFiles := make(map[string]bool)
	for _, glob := range logGroup.globFiles {
		files, err := filepath.Glob(glob)
		if err != nil {
			log.Fatalf("Unable to find files for log group %s: %s", logGroup.name, err)
		}

		for _, v := range files {
			newFiles[v] = true
		}
	}

	//Start tailing new files!
	for file, _ := range newFiles {
		parserTest(file, logGroup, perfInfo)
	}

}

func StartParserTest(config *Config, selectedLogGroup string, perfInfo bool) {
	for logGroupName, logGroup := range config.logGroups {
		if selectedLogGroup == "" || logGroupName == selectedLogGroup {
			startLogGroupParserTest(logGroup, perfInfo)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"fmt"
	"flag"
	"bytes"
	"strings"
	"io/ioutil"
	"path/filepath"

	"github.com/cloudfoundry/cli/plugin"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	"github.com/cloudfoundry/cli/cf/configuration/config_helpers"
)

type TargetsPlugin struct {
	configPath string
	targetsPath string
	currentPath string
	suffix string
}

func newTargetsPlugin() *TargetsPlugin {
	targetsPath := filepath.Join(filepath.Dir(config_helpers.DefaultFilePath()), "targets")
	os.Mkdir(targetsPath, 0700)
	return &TargetsPlugin {
		configPath: config_helpers.DefaultFilePath(),
		targetsPath: targetsPath,
		currentPath: filepath.Join(targetsPath, "current"),
		suffix: "." + filepath.Base(config_helpers.DefaultFilePath()),
	}
}

func (c *TargetsPlugin) GetMetadata() plugin.PluginMetadata {
	return plugin.PluginMetadata{
		Name: "cf-targets",
		Version: plugin.VersionType{
			Major: 1,
			Minor: 0,
			Build: 0,
		},
		Commands: []plugin.Command{
			{
				Name:     "targets",
				HelpText: "List available targets",
				UsageDetails: plugin.Usage {
					Usage: "targets\n   cf targets",
				},
			},
			{
				Name:     "set-target",
				HelpText: "Set current target",
				UsageDetails: plugin.Usage {
					Usage: "set-target\n   cf set-target [-f] NAME",
					Options: map[string]string{
						"f": "replace the current target even if it has not been saved",
					},
				},
			},
			{
				Name:     "save-target",
				HelpText: "Save current target",
				UsageDetails: plugin.Usage {
					Usage: "save-target\n   cf save-target [-f] NAME",
					Options: map[string]string{
						"f": "save the target even if the specified name already exists",
					},
				},
			},
			{
				Name:     "delete-target",
				HelpText: "Delete a saved target",
				UsageDetails: plugin.Usage {
					Usage: "delete-target\n   cf delete-target NAME",
				},
			},
		},
	}
}

func main() {
	plugin.Start(newTargetsPlugin())
}

func (c *TargetsPlugin) Run(cliConnection plugin.CliConnection, args []string) {
	if args[0] == "targets" {
		c.TargetsCommand(args)
	} else if args[0] == "set-target" {
		c.SetTargetCommand(args)
	} else if args[0] == "save-target" {
		c.SaveTargetCommand(args)
	} else if args[0] == "delete-target" {
		c.DeleteTargetCommand(args)
	}
}

func (c *TargetsPlugin) TargetsCommand(args []string) {
	if len(args) != 1 {
		c.exitWithUsage("targets")
	}
	changed, needsUpdate := c.compareCurrent()
	current := c.getCurrent()
	targets := c.getTargets()
	if len(targets) < 1 {
		fmt.Println("No targets have been saved yet. To save the current target, use:")
		fmt.Println("   cf save-target NAME")
	} else {
		for _,target := range targets {
			var qualifier string
			if target == current {
				qualifier = "(current"
				if changed {
					qualifier += ", modified"
				} else if needsUpdate {
					qualifier += "*"
				}
				qualifier += ")"
			}
			fmt.Println(target, qualifier)
		}
	}
}

func (c *TargetsPlugin) SetTargetCommand(args []string) {
	flagSet := flag.NewFlagSet("set-target", flag.ContinueOnError)
	force := flagSet.Bool("f", false, "force")
	err := flagSet.Parse(args[1:])
	if err != nil || len(flagSet.Args()) != 1 {
		c.exitWithUsage("set-target")
	}
	targetName := flagSet.Arg(0)
	targetPath := c.targetPath(targetName)
	changed, needsUpdate := c.compareCurrent()
	if *force || !changed {
		if needsUpdate {
			c.copyContents(c.configPath, c.currentPath)
		}
		c.copyContents(targetPath, c.configPath)
		c.linkCurrent(targetPath)
	} else {
		fmt.Println("Your current target has not been saved. Use save-target first, or use -f to discard your changes.")
		os.Exit(1)
	}
	fmt.Println("Set target to", targetName);
}

func (c *TargetsPlugin) SaveTargetCommand(args []string) {
	flagSet := flag.NewFlagSet("save-target", flag.ContinueOnError)
	force := flagSet.Bool("f", false, "force")
	err := flagSet.Parse(args[1:])
	if err != nil || len(flagSet.Args()) != 1 {
		c.exitWithUsage("save-target")
	}
	targetName := flagSet.Arg(0)
	targetPath := c.targetPath(targetName)
	if *force || !c.targetExists(targetPath) {
		c.copyContents(c.configPath, targetPath)
		c.linkCurrent(targetPath)
	} else {
		fmt.Println("Target", targetName, "already exists. Use -f to overwrite it.")
		os.Exit(1)
	}
	fmt.Println("Saved current target as", targetName)
}

func (c *TargetsPlugin) DeleteTargetCommand(args []string) {
	if len(args) != 2 {
		c.exitWithUsage("delete-target")
	}
	targetName := args[1]
	targetPath := c.targetPath(targetName)
	if !c.targetExists(targetPath) {
		fmt.Println("Target", targetName, "does not exist")
		os.Exit(1);
	}
	os.Remove(targetPath)
	fmt.Println("Deleted target", targetName);
}

func (c *TargetsPlugin) getTargets() []string {
	var targets []string
	files,_ := ioutil.ReadDir(c.targetsPath);
	for _,file := range files {
		filename := file.Name()
		if strings.HasSuffix(filename, c.suffix) {
			targets = append(targets, strings.TrimSuffix(filename, c.suffix))
		}
	}
	return targets
}

func (c *TargetsPlugin) targetExists(targetPath string) bool {
	target := configuration.NewDiskPersistor(targetPath)
	return target.Exists()
}

func (c *TargetsPlugin) compareCurrent() (bool, bool) {
	currentConfig := configuration.NewDiskPersistor(c.configPath)
	currentTarget := configuration.NewDiskPersistor(c.currentPath)
	if !currentTarget.Exists() {
		return true, true;
	}

	configData := core_config.NewData()
	targetData := core_config.NewData()

	err := currentConfig.Load(configData)
	c.checkError(err)
	err = currentTarget.Load(targetData)
	c.checkError(err)

	// Ignore the access-token field, as it changes frequently
	needsUpdate := targetData.AccessToken != configData.AccessToken
	targetData.AccessToken = configData.AccessToken

	currentContent, err := configData.JsonMarshalV3()
	c.checkError(err)
	savedContent, err := targetData.JsonMarshalV3()
	c.checkError(err)
	return !bytes.Equal(currentContent, savedContent), needsUpdate
}

func (c *TargetsPlugin) copyContents(sourcePath, targetPath string) {
	content, err := ioutil.ReadFile(sourcePath)
	c.checkError(err)
	err = ioutil.WriteFile(targetPath, content, 0600)
	c.checkError(err)
}

func (c *TargetsPlugin) linkCurrent(targetPath string) {
	_ = os.Remove(c.currentPath)
	err := os.Symlink(targetPath, c.currentPath)
	c.checkError(err)
}

func (c *TargetsPlugin) targetPath(targetName string) string {
	return filepath.Join(c.targetsPath, targetName + c.suffix)
}

func (c *TargetsPlugin) checkError(err error) {
	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}

func (c *TargetsPlugin) exitWithUsage(command string) {
	metadata := c.GetMetadata()
	for _,candidate := range metadata.Commands {
		if (candidate.Name == command) {
			fmt.Println("Usage: " + candidate.UsageDetails.Usage)
			os.Exit(1);
		}
	}
}

func (c *TargetsPlugin) getCurrent() string {
	targetPath, err := filepath.EvalSymlinks(c.currentPath)
	c.checkError(err)
	return strings.TrimSuffix(filepath.Base(targetPath), c.suffix)
}
<commit_msg>Corrected the end case where no targets yet exist<commit_after>package main

import (
	"os"
	"fmt"
	"flag"
	"bytes"
	"strings"
	"io/ioutil"
	"path/filepath"

	"github.com/cloudfoundry/cli/plugin"
	"github.com/cloudfoundry/cli/cf/configuration"
	"github.com/cloudfoundry/cli/cf/configuration/core_config"
	"github.com/cloudfoundry/cli/cf/configuration/config_helpers"
)

type TargetsPlugin struct {
	configPath string
	targetsPath string
	currentPath string
	suffix string
}

func newTargetsPlugin() *TargetsPlugin {
	targetsPath := filepath.Join(filepath.Dir(config_helpers.DefaultFilePath()), "targets")
	os.Mkdir(targetsPath, 0700)
	return &TargetsPlugin {
		configPath: config_helpers.DefaultFilePath(),
		targetsPath: targetsPath,
		currentPath: filepath.Join(targetsPath, "current"),
		suffix: "." + filepath.Base(config_helpers.DefaultFilePath()),
	}
}

func (c *TargetsPlugin) GetMetadata() plugin.PluginMetadata {
	return plugin.PluginMetadata{
		Name: "cf-targets",
		Version: plugin.VersionType{
			Major: 1,
			Minor: 0,
			Build: 0,
		},
		Commands: []plugin.Command{
			{
				Name:     "targets",
				HelpText: "List available targets",
				UsageDetails: plugin.Usage {
					Usage: "targets\n   cf targets",
				},
			},
			{
				Name:     "set-target",
				HelpText: "Set current target",
				UsageDetails: plugin.Usage {
					Usage: "set-target\n   cf set-target [-f] NAME",
					Options: map[string]string{
						"f": "replace the current target even if it has not been saved",
					},
				},
			},
			{
				Name:     "save-target",
				HelpText: "Save current target",
				UsageDetails: plugin.Usage {
					Usage: "save-target\n   cf save-target [-f] NAME",
					Options: map[string]string{
						"f": "save the target even if the specified name already exists",
					},
				},
			},
			{
				Name:     "delete-target",
				HelpText: "Delete a saved target",
				UsageDetails: plugin.Usage {
					Usage: "delete-target\n   cf delete-target NAME",
				},
			},
		},
	}
}

func main() {
	plugin.Start(newTargetsPlugin())
}

func (c *TargetsPlugin) Run(cliConnection plugin.CliConnection, args []string) {
	if args[0] == "targets" {
		c.TargetsCommand(args)
	} else if args[0] == "set-target" {
		c.SetTargetCommand(args)
	} else if args[0] == "save-target" {
		c.SaveTargetCommand(args)
	} else if args[0] == "delete-target" {
		c.DeleteTargetCommand(args)
	}
}

func (c *TargetsPlugin) TargetsCommand(args []string) {
	if len(args) != 1 {
		c.exitWithUsage("targets")
	}
	changed, needsUpdate := c.compareCurrent()
	targets := c.getTargets()
	if len(targets) < 1 {
		fmt.Println("No targets have been saved yet. To save the current target, use:")
		fmt.Println("   cf save-target NAME")
	} else {
		current := c.getCurrent()
		for _,target := range targets {
			var qualifier string
			if target == current {
				qualifier = "(current"
				if changed {
					qualifier += ", modified"
				} else if needsUpdate {
					qualifier += "*"
				}
				qualifier += ")"
			}
			fmt.Println(target, qualifier)
		}
	}
}

func (c *TargetsPlugin) SetTargetCommand(args []string) {
	flagSet := flag.NewFlagSet("set-target", flag.ContinueOnError)
	force := flagSet.Bool("f", false, "force")
	err := flagSet.Parse(args[1:])
	if err != nil || len(flagSet.Args()) != 1 {
		c.exitWithUsage("set-target")
	}
	targetName := flagSet.Arg(0)
	targetPath := c.targetPath(targetName)
	changed, needsUpdate := c.compareCurrent()
	if *force || !changed {
		if needsUpdate {
			c.copyContents(c.configPath, c.currentPath)
		}
		c.copyContents(targetPath, c.configPath)
		c.linkCurrent(targetPath)
	} else {
		fmt.Println("Your current target has not been saved. Use save-target first, or use -f to discard your changes.")
		os.Exit(1)
	}
	fmt.Println("Set target to", targetName);
}

func (c *TargetsPlugin) SaveTargetCommand(args []string) {
	flagSet := flag.NewFlagSet("save-target", flag.ContinueOnError)
	force := flagSet.Bool("f", false, "force")
	err := flagSet.Parse(args[1:])
	if err != nil || len(flagSet.Args()) != 1 {
		c.exitWithUsage("save-target")
	}
	targetName := flagSet.Arg(0)
	targetPath := c.targetPath(targetName)
	if *force || !c.targetExists(targetPath) {
		c.copyContents(c.configPath, targetPath)
		c.linkCurrent(targetPath)
	} else {
		fmt.Println("Target", targetName, "already exists. Use -f to overwrite it.")
		os.Exit(1)
	}
	fmt.Println("Saved current target as", targetName)
}

func (c *TargetsPlugin) DeleteTargetCommand(args []string) {
	if len(args) != 2 {
		c.exitWithUsage("delete-target")
	}
	targetName := args[1]
	targetPath := c.targetPath(targetName)
	if !c.targetExists(targetPath) {
		fmt.Println("Target", targetName, "does not exist")
		os.Exit(1);
	}
	os.Remove(targetPath)
	fmt.Println("Deleted target", targetName);
}

func (c *TargetsPlugin) getTargets() []string {
	var targets []string
	files,_ := ioutil.ReadDir(c.targetsPath);
	for _,file := range files {
		filename := file.Name()
		if strings.HasSuffix(filename, c.suffix) {
			targets = append(targets, strings.TrimSuffix(filename, c.suffix))
		}
	}
	return targets
}

func (c *TargetsPlugin) targetExists(targetPath string) bool {
	target := configuration.NewDiskPersistor(targetPath)
	return target.Exists()
}

func (c *TargetsPlugin) compareCurrent() (bool, bool) {
	currentConfig := configuration.NewDiskPersistor(c.configPath)
	currentTarget := configuration.NewDiskPersistor(c.currentPath)
	if !currentTarget.Exists() {
		return true, true;
	}

	configData := core_config.NewData()
	targetData := core_config.NewData()

	err := currentConfig.Load(configData)
	c.checkError(err)
	err = currentTarget.Load(targetData)
	c.checkError(err)

	// Ignore the access-token field, as it changes frequently
	needsUpdate := targetData.AccessToken != configData.AccessToken
	targetData.AccessToken = configData.AccessToken

	currentContent, err := configData.JsonMarshalV3()
	c.checkError(err)
	savedContent, err := targetData.JsonMarshalV3()
	c.checkError(err)
	return !bytes.Equal(currentContent, savedContent), needsUpdate
}

func (c *TargetsPlugin) copyContents(sourcePath, targetPath string) {
	content, err := ioutil.ReadFile(sourcePath)
	c.checkError(err)
	err = ioutil.WriteFile(targetPath, content, 0600)
	c.checkError(err)
}

func (c *TargetsPlugin) linkCurrent(targetPath string) {
	_ = os.Remove(c.currentPath)
	err := os.Symlink(targetPath, c.currentPath)
	c.checkError(err)
}

func (c *TargetsPlugin) targetPath(targetName string) string {
	return filepath.Join(c.targetsPath, targetName + c.suffix)
}

func (c *TargetsPlugin) checkError(err error) {
	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}

func (c *TargetsPlugin) exitWithUsage(command string) {
	metadata := c.GetMetadata()
	for _,candidate := range metadata.Commands {
		if (candidate.Name == command) {
			fmt.Println("Usage: " + candidate.UsageDetails.Usage)
			os.Exit(1);
		}
	}
}

func (c *TargetsPlugin) getCurrent() string {
	targetPath, err := filepath.EvalSymlinks(c.currentPath)
	c.checkError(err)
	return strings.TrimSuffix(filepath.Base(targetPath), c.suffix)
}
<|endoftext|>