<commit_before>package format

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/colorstring"
)

// Plan is a representation of a plan optimized for display to
// an end-user, as opposed to terraform.Plan which is for internal use.
//
// DisplayPlan excludes implementation details that may otherwise appear
// in the main plan, such as destroy actions on data sources (which are
// there only to clean up the state).
type Plan struct {
	Resources []*InstanceDiff
}

// InstanceDiff is a representation of an instance diff optimized
// for display, in conjunction with DisplayPlan.
type InstanceDiff struct {
	Addr   *terraform.ResourceAddress
	Action terraform.DiffChangeType

	// Attributes describes changes to the attributes of the instance.
	//
	// For destroy diffs this is always nil.
	Attributes []*AttributeDiff

	Tainted bool
	Deposed bool
}

// AttributeDiff is a representation of an attribute diff optimized
// for display, in conjunction with DisplayInstanceDiff.
type AttributeDiff struct {
	// Path is a dot-delimited traversal through possibly many levels of list and map structure,
	// intended for display purposes only.
	Path string

	Action terraform.DiffChangeType

	OldValue string
	NewValue string

	NewComputed bool
	Sensitive   bool
	ForcesNew   bool
}

// PlanStats gives summary counts for a Plan.
type PlanStats struct {
	ToAdd, ToChange, ToDestroy int
}

// NewPlan produces a display-oriented Plan from a terraform.Plan.
func NewPlan(plan *terraform.Plan) *Plan {
	ret := &Plan{}
	if plan == nil || plan.Diff == nil || plan.Diff.Empty() {
		// Nothing to do!
		return ret
	}

	for _, m := range plan.Diff.Modules {
		var modulePath []string
		if !m.IsRoot() {
			// trim off the leading "root" path segment, since it's implied
			// when we use a path in a resource address.
			modulePath = m.Path[1:]
		}

		for k, r := range m.Resources {
			if r.Empty() {
				continue
			}

			addr, err := terraform.ParseResourceAddressForInstanceDiff(modulePath, k)
			if err != nil {
				// should never happen; indicates invalid diff
				panic("invalid resource address in diff")
			}

			dataSource := addr.Mode == config.DataResourceMode

			// We create "destroy" actions for data resources so we can clean
			// up their entries in state, but this is an implementation detail
			// that users shouldn't see.
			if dataSource && r.ChangeType() == terraform.DiffDestroy {
				continue
			}

			did := &InstanceDiff{
				Addr:    addr,
				Action:  r.ChangeType(),
				Tainted: r.DestroyTainted,
				Deposed: r.DestroyDeposed,
			}

			if dataSource && did.Action == terraform.DiffCreate {
				// Use "refresh" as the action for display, since core
				// currently uses Create for this.
				did.Action = terraform.DiffRefresh
			}

			ret.Resources = append(ret.Resources, did)

			if did.Action == terraform.DiffDestroy {
				// Don't show any outputs for destroy actions
				continue
			}

			for k, a := range r.Attributes {
				var action terraform.DiffChangeType
				switch {
				case a.NewRemoved:
					action = terraform.DiffDestroy
				case did.Action == terraform.DiffCreate:
					action = terraform.DiffCreate
				default:
					action = terraform.DiffUpdate
				}

				did.Attributes = append(did.Attributes, &AttributeDiff{
					Path:   k,
					Action: action,

					OldValue: a.Old,
					NewValue: a.New,

					Sensitive:   a.Sensitive,
					ForcesNew:   a.RequiresNew,
					NewComputed: a.NewComputed,
				})
			}

			// Sort the attributes by their paths for display
			sort.Slice(did.Attributes, func(i, j int) bool {
				iPath := did.Attributes[i].Path
				jPath := did.Attributes[j].Path

				// as a special case, "id" is always first
				switch {
				case iPath != jPath && (iPath == "id" || jPath == "id"):
					return iPath == "id"
				default:
					return iPath < jPath
				}
			})

		}
	}

	// Sort the instance diffs by their addresses for display.
	sort.Slice(ret.Resources, func(i, j int) bool {
		iAddr := ret.Resources[i].Addr
		jAddr := ret.Resources[j].Addr
		return iAddr.Less(jAddr)
	})

	return ret
}

// Format produces and returns a text representation of the receiving plan
// intended for display in a terminal.
//
// If color is not nil, it is used to colorize the output.
func (p *Plan) Format(color *colorstring.Colorize) string {
	if p.Empty() {
		return "This plan does nothing."
	}

	if color == nil {
		color = &colorstring.Colorize{
			Colors: colorstring.DefaultColors,
			Reset:  false,
		}
	}

	// Find the longest path length of all the paths that are changing,
	// so we can align them all.
	keyLen := 0
	for _, r := range p.Resources {
		for _, attr := range r.Attributes {
			key := attr.Path

			if len(key) > keyLen {
				keyLen = len(key)
			}
		}
	}

	buf := new(bytes.Buffer)
	for _, r := range p.Resources {
		formatPlanInstanceDiff(buf, r, keyLen, color)
	}

	return strings.TrimSpace(buf.String())
}

// Stats returns statistics about the plan
func (p *Plan) Stats() PlanStats {
	var ret PlanStats
	for _, r := range p.Resources {
		switch r.Action {
		case terraform.DiffCreate:
			ret.ToAdd++
		case terraform.DiffUpdate:
			ret.ToChange++
		case terraform.DiffDestroyCreate:
			ret.ToAdd++
			ret.ToDestroy++
		case terraform.DiffDestroy:
			ret.ToDestroy++
		}
	}
	return ret
}

// Empty returns true if there is at least one resource diff in the receiving plan.
func (p *Plan) Empty() bool {
	return len(p.Resources) == 0
}

// formatPlanInstanceDiff writes the text representation of the given instance diff
// to the given buffer, using the given colorizer.
func formatPlanInstanceDiff(buf *bytes.Buffer, r *InstanceDiff, keyLen int, colorizer *colorstring.Colorize) {
	addrStr := r.Addr.String()

	// Determine the color for the text (green for adding, yellow
	// for change, red for delete), and symbol, and output the
	// resource header.
	color := "yellow"
	symbol := "  ~"
	oldValues := true
	switch r.Action {
	case terraform.DiffDestroyCreate:
		color = "yellow"
		symbol = "[red]-[reset]/[green]+[reset][yellow]"
	case terraform.DiffCreate:
		color = "green"
		symbol = "  +"
		oldValues = false
	case terraform.DiffDestroy:
		color = "red"
		symbol = "  -"
	case terraform.DiffRefresh:
		symbol = " <="
		color = "cyan"
		oldValues = false
	}

	var extraStr string
	if r.Tainted {
		extraStr = extraStr + colorizer.Color(" (tainted)")
	}
	if r.Deposed {
		extraStr = extraStr + colorizer.Color(" (deposed)")
	}
	if r.Action == terraform.DiffDestroyCreate {
		extraStr = extraStr + colorizer.Color(" [red][bold](new resource required)")
	}

	buf.WriteString(
		colorizer.Color(fmt.Sprintf(
			"[%s]%s %s%s\n",
			color, symbol, addrStr, extraStr,
		)),
	)

	for _, attr := range r.Attributes {

		v := attr.NewValue
		var dispV string
		switch {
		case v == "" && attr.NewComputed:
			dispV = "<computed>"
		case attr.Sensitive:
			dispV = "<sensitive>"
		default:
			dispV = fmt.Sprintf("%q", v)
		}

		updateMsg := ""
		switch {
		case attr.ForcesNew && r.Action == terraform.DiffDestroy:
			updateMsg = colorizer.Color(" [red](forces new resource)")
		case attr.Sensitive && oldValues:
			updateMsg = colorizer.Color(" [yellow](attribute changed)")
		}

		if oldValues {
			u := attr.OldValue
			var dispU string
			switch {
			case attr.Sensitive:
				dispU = "<sensitive>"
			default:
				dispU = fmt.Sprintf("%q", u)
			}
			buf.WriteString(fmt.Sprintf(
				"      %s:%s %s => %s%s\n",
				attr.Path,
				strings.Repeat(" ", keyLen-len(attr.Path)),
				dispU, dispV,
				updateMsg,
			))
		} else {
			buf.WriteString(fmt.Sprintf(
				"      %s:%s %s%s\n",
				attr.Path,
				strings.Repeat(" ", keyLen-len(attr.Path)),
				dispV,
				updateMsg,
			))
		}
	}

	// Write the reset color so we don't bleed color into later text
	buf.WriteString(colorizer.Color("[reset]\n"))
}
<commit_msg>command: various adjustments to the diff presentation<commit_after>package format

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/colorstring"
)

// Plan is a representation of a plan optimized for display to
// an end-user, as opposed to terraform.Plan which is for internal use.
//
// DisplayPlan excludes implementation details that may otherwise appear
// in the main plan, such as destroy actions on data sources (which are
// there only to clean up the state).
type Plan struct {
	Resources []*InstanceDiff
}

// InstanceDiff is a representation of an instance diff optimized
// for display, in conjunction with DisplayPlan.
type InstanceDiff struct {
	Addr   *terraform.ResourceAddress
	Action terraform.DiffChangeType

	// Attributes describes changes to the attributes of the instance.
	//
	// For destroy diffs this is always nil.
	Attributes []*AttributeDiff

	Tainted bool
	Deposed bool
}

// AttributeDiff is a representation of an attribute diff optimized
// for display, in conjunction with DisplayInstanceDiff.
type AttributeDiff struct {
	// Path is a dot-delimited traversal through possibly many levels of list and map structure,
	// intended for display purposes only.
	Path string

	Action terraform.DiffChangeType

	OldValue string
	NewValue string

	NewComputed bool
	Sensitive   bool
	ForcesNew   bool
}

// PlanStats gives summary counts for a Plan.
type PlanStats struct {
	ToAdd, ToChange, ToDestroy int
}

// NewPlan produces a display-oriented Plan from a terraform.Plan.
func NewPlan(plan *terraform.Plan) *Plan {
	ret := &Plan{}
	if plan == nil || plan.Diff == nil || plan.Diff.Empty() {
		// Nothing to do!
		return ret
	}

	for _, m := range plan.Diff.Modules {
		var modulePath []string
		if !m.IsRoot() {
			// trim off the leading "root" path segment, since it's implied
			// when we use a path in a resource address.
			modulePath = m.Path[1:]
		}

		for k, r := range m.Resources {
			if r.Empty() {
				continue
			}

			addr, err := terraform.ParseResourceAddressForInstanceDiff(modulePath, k)
			if err != nil {
				// should never happen; indicates invalid diff
				panic("invalid resource address in diff")
			}

			dataSource := addr.Mode == config.DataResourceMode

			// We create "destroy" actions for data resources so we can clean
			// up their entries in state, but this is an implementation detail
			// that users shouldn't see.
			if dataSource && r.ChangeType() == terraform.DiffDestroy {
				continue
			}

			did := &InstanceDiff{
				Addr:    addr,
				Action:  r.ChangeType(),
				Tainted: r.DestroyTainted,
				Deposed: r.DestroyDeposed,
			}

			if dataSource && did.Action == terraform.DiffCreate {
				// Use "refresh" as the action for display, since core
				// currently uses Create for this.
				did.Action = terraform.DiffRefresh
			}

			ret.Resources = append(ret.Resources, did)

			if did.Action == terraform.DiffDestroy {
				// Don't show any outputs for destroy actions
				continue
			}

			for k, a := range r.Attributes {
				var action terraform.DiffChangeType
				switch {
				case a.NewRemoved:
					action = terraform.DiffDestroy
				case did.Action == terraform.DiffCreate:
					action = terraform.DiffCreate
				default:
					action = terraform.DiffUpdate
				}

				did.Attributes = append(did.Attributes, &AttributeDiff{
					Path:   k,
					Action: action,

					OldValue: a.Old,
					NewValue: a.New,

					Sensitive:   a.Sensitive,
					ForcesNew:   a.RequiresNew,
					NewComputed: a.NewComputed,
				})
			}

			// Sort the attributes by their paths for display
			sort.Slice(did.Attributes, func(i, j int) bool {
				iPath := did.Attributes[i].Path
				jPath := did.Attributes[j].Path

				// as a special case, "id" is always first
				switch {
				case iPath != jPath && (iPath == "id" || jPath == "id"):
					return iPath == "id"
				default:
					return iPath < jPath
				}
			})

		}
	}

	// Sort the instance diffs by their addresses for display.
	sort.Slice(ret.Resources, func(i, j int) bool {
		iAddr := ret.Resources[i].Addr
		jAddr := ret.Resources[j].Addr
		return iAddr.Less(jAddr)
	})

	return ret
}

// Format produces and returns a text representation of the receiving plan
// intended for display in a terminal.
//
// If color is not nil, it is used to colorize the output.
func (p *Plan) Format(color *colorstring.Colorize) string {
	if p.Empty() {
		return "This plan does nothing."
	}

	if color == nil {
		color = &colorstring.Colorize{
			Colors: colorstring.DefaultColors,
			Reset:  false,
		}
	}

	// Find the longest path length of all the paths that are changing,
	// so we can align them all.
	keyLen := 0
	for _, r := range p.Resources {
		for _, attr := range r.Attributes {
			key := attr.Path

			if len(key) > keyLen {
				keyLen = len(key)
			}
		}
	}

	buf := new(bytes.Buffer)
	for _, r := range p.Resources {
		formatPlanInstanceDiff(buf, r, keyLen, color)
	}

	return strings.TrimSpace(buf.String())
}

// Stats returns statistics about the plan
func (p *Plan) Stats() PlanStats {
	var ret PlanStats
	for _, r := range p.Resources {
		switch r.Action {
		case terraform.DiffCreate:
			ret.ToAdd++
		case terraform.DiffUpdate:
			ret.ToChange++
		case terraform.DiffDestroyCreate:
			ret.ToAdd++
			ret.ToDestroy++
		case terraform.DiffDestroy:
			ret.ToDestroy++
		}
	}
	return ret
}

// ActionCounts returns the number of diffs for each action type
func (p *Plan) ActionCounts() map[terraform.DiffChangeType]int {
	ret := map[terraform.DiffChangeType]int{}
	for _, r := range p.Resources {
		ret[r.Action]++
	}
	return ret
}

// Empty returns true if there is at least one resource diff in the receiving plan.
func (p *Plan) Empty() bool {
	return len(p.Resources) == 0
}

// DiffActionSymbol returns a string that, once passed through a
// colorstring.Colorize, will produce a result that can be written
// to a terminal to produce a symbol made of three printable
// characters, possibly interspersed with VT100 color codes.
func DiffActionSymbol(action terraform.DiffChangeType) string {
	switch action {
	case terraform.DiffDestroyCreate:
		return "[red]-[reset]/[green]+[reset]"
	case terraform.DiffCreate:
		return "  [green]+[reset]"
	case terraform.DiffDestroy:
		return "  [red]-[reset]"
	case terraform.DiffRefresh:
		return " [cyan]<=[reset]"
	default:
		return "  [yellow]~[reset]"
	}
}

// formatPlanInstanceDiff writes the text representation of the given instance diff
// to the given buffer, using the given colorizer.
func formatPlanInstanceDiff(buf *bytes.Buffer, r *InstanceDiff, keyLen int, colorizer *colorstring.Colorize) {
	addrStr := r.Addr.String()

	// Determine the color for the text (green for adding, yellow
	// for change, red for delete), and symbol, and output the
	// resource header.
	color := "yellow"
	symbol := DiffActionSymbol(r.Action)
	oldValues := true
	switch r.Action {
	case terraform.DiffDestroyCreate:
		color = "yellow"
	case terraform.DiffCreate:
		color = "green"
		oldValues = false
	case terraform.DiffDestroy:
		color = "red"
	case terraform.DiffRefresh:
		color = "cyan"
		oldValues = false
	}

	var extraStr string
	if r.Tainted {
		extraStr = extraStr + " (tainted)"
	}
	if r.Deposed {
		extraStr = extraStr + " (deposed)"
	}
	if r.Action == terraform.DiffDestroyCreate {
		extraStr = extraStr + colorizer.Color(" [red][bold](new resource required)")
	}

	buf.WriteString(
		colorizer.Color(fmt.Sprintf(
			"[%s]%s [%s]%s%s\n",
			color, symbol, color, addrStr, extraStr,
		)),
	)

	for _, attr := range r.Attributes {

		v := attr.NewValue
		var dispV string
		switch {
		case v == "" && attr.NewComputed:
			dispV = "<computed>"
		case attr.Sensitive:
			dispV = "<sensitive>"
		default:
			dispV = fmt.Sprintf("%q", v)
		}

		updateMsg := ""
		switch {
		case attr.ForcesNew && r.Action == terraform.DiffDestroy:
			updateMsg = colorizer.Color(" [red](forces new resource)")
		case attr.Sensitive && oldValues:
			updateMsg = colorizer.Color(" [yellow](attribute changed)")
		}

		if oldValues {
			u := attr.OldValue
			var dispU string
			switch {
			case attr.Sensitive:
				dispU = "<sensitive>"
			default:
				dispU = fmt.Sprintf("%q", u)
			}
			buf.WriteString(fmt.Sprintf(
				"      %s:%s %s => %s%s\n",
				attr.Path,
				strings.Repeat(" ", keyLen-len(attr.Path)),
				dispU, dispV,
				updateMsg,
			))
		} else {
			buf.WriteString(fmt.Sprintf(
				"      %s:%s %s%s\n",
				attr.Path,
				strings.Repeat(" ", keyLen-len(attr.Path)),
				dispV,
				updateMsg,
			))
		}
	}

	// Write the reset color so we don't bleed color into later text
	buf.WriteString(colorizer.Color("[reset]\n"))
}
<|endoftext|>


<commit_before>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// controlFlow tells the main loop what it should do next.
type controlFlow int

// controlFlow values.
const (
	controlStop  controlFlow = 0
	controlBegin controlFlow = 1
	controlNext  controlFlow = 2
)

// dataType specifies which type of data we are going to save for each node.
type dataType = interface{}

// tree represents a radix tree.
type tree struct {
	prefix    string
	data      dataType
	children  [256]*tree
	parameter *tree
	wildcard  *tree
	kind      byte
}

// add adds a new element to the tree.
func (node *tree) add(path string, data dataType) {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			// This only occurs when the same parameter based route is added twice.
			// node: /post/:id|
			// path: /post/:id|
			if i == len(path) {
				node.data = data
				return
			}

			// When we hit a separator, we'll search for a fitting child.
			if path[i] == separator {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

		default:
			if i == len(path) {
				// The path already exists.
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					node.data = data
					return
				}

				// The path ended but the node prefix is longer.
				// node: /blog|feed
				// path: /blog|
				node.split(i-offset, "", data)
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				node.split(i-offset, path[i:], data)
				return
			}
		}

	next:
		i++
	}
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *tree) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := &tree{
		prefix:    node.prefix[index:],
		data:      node.data,
		children:  node.children,
		parameter: node.parameter,
		wildcard:  node.wildcard,
		kind:      node.kind,
	}

	/// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.children[splitNode.prefix[0]] = splitNode
		return
	}

	node.children[splitNode.prefix[0]] = splitNode

	// Create new nodes with the remaining path
	node.append(path, data)
}

// reset resets the existing node data.
func (node *tree) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.children = [256]*tree{}
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *tree) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.children[separator] != nil {
		return
	}

	node.children[separator] = &tree{
		prefix: "/",
		data:   data,
	}
}

// append appends the given path to the tree.
func (node *tree) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			child := &tree{
				prefix: path,
				data:   data,
			}

			node.children[path[0]] = child
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &tree{
				prefix: path[1:paramEnd],
				kind:   parameter,
			}

			node.parameter = child
			child.addTrailingSlash(data)
			node = child
			path = path[paramEnd:]
			continue
		}

		// We know there's a parameter, but not directly at the start.
		// Add a normal node with the path before the parameter start.
		child := &tree{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.children[path[0]] = child
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *tree) end(path string, data dataType, i int, offset int) (*tree, int, controlFlow) {
	child := node.children[path[i]]

	if child != nil {
		node = child
		offset = i
		return node, offset, controlNext
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.prefix = path
		node.data = data
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// find returns the data for the given path, if available.
func (node *tree) find(path string) dataType {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			if i == len(path) {
				// fmt.Printf("PARAMETER %s IS %s\n", node.prefix, path[offset:i])
				return node.data
			}

			if path[i] == separator {
				// fmt.Printf("PARAMETER %s IS %s\n", node.prefix, path[offset:i])
				node = node.children[separator]
				offset = i
				goto next
			}

		default:
			// We reached the end.
			if i == len(path) {
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					return node.data
				}

				// node: /blog|feed
				// path: /blog|
				return nil
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				child := node.children[path[i]]

				if child != nil {
					node = child
					offset = i
					goto next
				}

				// node: /|:id
				// path: /|blog
				if node.parameter != nil {
					node = node.parameter
					offset = i
					goto begin
				}

				return nil
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				return nil
			}
		}

	next:
		i++
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *tree) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *tree) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<commit_msg>Cleaner split implementation<commit_after>package aero

import (
	"fmt"
	"io"
	"strings"

	"github.com/akyoto/color"
)

// node types
const (
	separator = '/'
	parameter = ':'
	wildcard  = '*'
)

// controlFlow tells the main loop what it should do next.
type controlFlow int

// controlFlow values.
const (
	controlStop  controlFlow = 0
	controlBegin controlFlow = 1
	controlNext  controlFlow = 2
)

// dataType specifies which type of data we are going to save for each node.
type dataType = interface{}

// tree represents a radix tree.
type tree struct {
	prefix    string
	data      dataType
	children  [256]*tree
	parameter *tree
	wildcard  *tree
	kind      byte
}

// add adds a new element to the tree.
func (node *tree) add(path string, data dataType) {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			// This only occurs when the same parameter based route is added twice.
			// node: /post/:id|
			// path: /post/:id|
			if i == len(path) {
				node.data = data
				return
			}

			// When we hit a separator, we'll search for a fitting child.
			if path[i] == separator {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

		default:
			if i == len(path) {
				// The path already exists.
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					node.data = data
					return
				}

				// The path ended but the node prefix is longer.
				// node: /blog|feed
				// path: /blog|
				node.split(i-offset, "", data)
				return
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				var control controlFlow
				node, offset, control = node.end(path, data, i, offset)

				switch control {
				case controlStop:
					return
				case controlBegin:
					goto begin
				case controlNext:
					goto next
				}
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				node.split(i-offset, path[i:], data)
				return
			}
		}

	next:
		i++
	}
}

// split splits the node at the given index and inserts
// a new child node with the given path and data.
// If path is empty, it will not create another child node
// and instead assign the data directly to the node.
func (node *tree) split(index int, path string, data dataType) {
	// Create split node with the remaining string
	splitNode := node.clone(node.prefix[index:])

	/// The existing data must be removed
	node.reset(node.prefix[:index])

	// If the path is empty, it means we don't create a 2nd child node.
	// Just assign the data for the existing node and store a single child node.
	if path == "" {
		node.data = data
		node.children[splitNode.prefix[0]] = splitNode
		return
	}

	node.children[splitNode.prefix[0]] = splitNode

	// Create new nodes with the remaining path
	node.append(path, data)
}

// clone clones the node with a new prefix.
func (node *tree) clone(prefix string) *tree {
	return &tree{
		prefix:    prefix,
		data:      node.data,
		children:  node.children,
		parameter: node.parameter,
		wildcard:  node.wildcard,
		kind:      node.kind,
	}
}

// reset resets the existing node data.
func (node *tree) reset(prefix string) {
	node.prefix = prefix
	node.data = nil
	node.parameter = nil
	node.wildcard = nil
	node.kind = 0
	node.children = [256]*tree{}
}

// addTrailingSlash adds a trailing slash with the same data.
func (node *tree) addTrailingSlash(data dataType) {
	if strings.HasSuffix(node.prefix, "/") || node.children[separator] != nil {
		return
	}

	node.children[separator] = &tree{
		prefix: "/",
		data:   data,
	}
}

// append appends the given path to the tree.
func (node *tree) append(path string, data dataType) {
	// At this point, all we know is that somewhere
	// in the remaining string we have parameters.
	// node: /user|
	// path: /user|/:userid
	for {
		if path == "" {
			node.data = data
			return
		}

		paramStart := strings.IndexByte(path, parameter)

		if paramStart == -1 {
			paramStart = strings.IndexByte(path, wildcard)
		}

		// If it's a static route we are adding,
		// just add the remainder as a normal node.
		if paramStart == -1 {
			child := &tree{
				prefix: path,
				data:   data,
			}

			node.children[path[0]] = child
			child.addTrailingSlash(data)
			return
		}

		// If we're directly in front of a parameter,
		// add a parameter node.
		if paramStart == 0 {
			paramEnd := strings.IndexByte(path, separator)

			if paramEnd == -1 {
				paramEnd = len(path)
			}

			child := &tree{
				prefix: path[1:paramEnd],
				kind:   parameter,
			}

			node.parameter = child
			child.addTrailingSlash(data)
			node = child
			path = path[paramEnd:]
			continue
		}

		// We know there's a parameter, but not directly at the start.
		// Add a normal node with the path before the parameter start.
		child := &tree{
			prefix: path[:paramStart],
		}

		// Allow trailing slashes to return
		// the same content as their parent node.
		if child.prefix == "/" {
			child.data = node.data
		}

		node.children[path[0]] = child
		node = child
		path = path[paramStart:]
	}
}

// end is called when the node was fully parsed
// and needs to decide the next control flow.
func (node *tree) end(path string, data dataType, i int, offset int) (*tree, int, controlFlow) {
	child := node.children[path[i]]

	if child != nil {
		node = child
		offset = i
		return node, offset, controlNext
	}

	// No fitting children found, does this node even contain a prefix yet?
	// If no prefix is set, this is the starting node.
	if node.prefix == "" {
		node.prefix = path
		node.data = data
		return node, offset, controlStop
	}

	// node: /user/|:id
	// path: /user/|:id/profile
	if node.parameter != nil {
		node = node.parameter
		offset = i
		return node, offset, controlBegin
	}

	node.append(path[i:], data)
	return node, offset, controlStop
}

// find returns the data for the given path, if available.
func (node *tree) find(path string) dataType {
	// Search tree for equal parts until we can no longer proceed
	i := 0
	offset := 0

	for {
	begin:
		switch node.kind {
		case parameter:
			if i == len(path) {
				// fmt.Printf("PARAMETER %s IS %s\n", node.prefix, path[offset:i])
				return node.data
			}

			if path[i] == separator {
				// fmt.Printf("PARAMETER %s IS %s\n", node.prefix, path[offset:i])
				node = node.children[separator]
				offset = i
				goto next
			}

		default:
			// We reached the end.
			if i == len(path) {
				// node: /blog|
				// path: /blog|
				if i-offset == len(node.prefix) {
					return node.data
				}

				// node: /blog|feed
				// path: /blog|
				return nil
			}

			// The node we just checked is entirely included in our path.
			// node: /|
			// path: /|blog
			if i-offset == len(node.prefix) {
				child := node.children[path[i]]

				if child != nil {
					node = child
					offset = i
					goto next
				}

				// node: /|:id
				// path: /|blog
				if node.parameter != nil {
					node = node.parameter
					offset = i
					goto begin
				}

				return nil
			}

			// We got a conflict.
			// node: /b|ag
			// path: /b|riefcase
			if path[i] != node.prefix[i-offset] {
				return nil
			}
		}

	next:
		i++
	}
}

// PrettyPrint prints a human-readable form of the tree to the given writer.
func (node *tree) PrettyPrint(writer io.Writer) {
	node.prettyPrint(writer, -1)
}

// prettyPrint
func (node *tree) prettyPrint(writer io.Writer, level int) {
	prefix := ""

	if level >= 0 {
		prefix = strings.Repeat("  ", level) + "|_ "
	}

	colorFunc := color.CyanString

	switch node.kind {
	case ':':
		prefix += ":"
		colorFunc = color.YellowString
	case '*':
		prefix += "*"
		colorFunc = color.GreenString
	}

	fmt.Fprintf(writer, "%s%s [%t]\n", prefix, colorFunc(node.prefix), node.data != nil)

	for _, child := range node.children {
		if child == nil {
			continue
		}

		child.prettyPrint(writer, level+1)
	}

	if node.parameter != nil {
		node.parameter.prettyPrint(writer, level+1)
	}

	if node.wildcard != nil {
		node.wildcard.prettyPrint(writer, level+1)
	}
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/base64"
	"net/http"
	"net/url"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/mssola/user_agent"
	"github.com/usefathom/fathom/pkg/aggregator"
	"github.com/usefathom/fathom/pkg/datastore"
	"github.com/usefathom/fathom/pkg/models"
)

type Collector struct {
	Store     datastore.Datastore
	Pageviews chan *models.Pageview

	// buffer vars
	updates []*models.Pageview
	inserts []*models.Pageview
	sizeu   int
	sizei   int
}

func NewCollector(store datastore.Datastore) *Collector {
	bufferCap := 100                         // persist every 100 pageviews, see https://github.com/usefathom/fathom/issues/132
	bufferTimeout := 1000 * time.Millisecond // or every 1000 ms, whichever comes first

	c := &Collector{
		Store:     store,
		Pageviews: make(chan *models.Pageview),
		updates:   make([]*models.Pageview, bufferCap),
		inserts:   make([]*models.Pageview, bufferCap),
		sizeu:     0,
		sizei:     0,
	}
	go c.aggregate()
	go c.worker(bufferCap, bufferTimeout)
	return c
}

func (c *Collector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !shouldCollect(r) {
		return
	}

	q := r.URL.Query()
	now := time.Now()

	pageview := &models.Pageview{
		ID:             q.Get("id"),
		SiteTrackingID: q.Get("sid"),
		Hostname:       parseHostname(q.Get("h")),
		Pathname:       parsePathname(q.Get("p")),
		IsNewVisitor:   q.Get("nv") == "1",
		IsNewSession:   q.Get("ns") == "1",
		IsUnique:       q.Get("u") == "1",
		Referrer:       q.Get("r"),
		IsFinished:     false,
		IsBounce:       true,
		Duration:       0,
		Timestamp:      now,
	}

	// push pageview onto channel to be inserted (in batch) later
	c.Pageviews <- pageview

	// indicate that we're not tracking user data, see https://github.com/usefathom/fathom/issues/65
	w.Header().Set("Tk", "N")

	// headers to prevent caching
	w.Header().Set("Content-Type", "image/gif")
	w.Header().Set("Expires", "Mon, 01 Jan 1990 00:00:00 GMT")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Pragma", "no-cache")

	// response, 1x1 px transparent GIF
	w.WriteHeader(http.StatusOK)
	b, _ := base64.StdEncoding.DecodeString("R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7")
	w.Write(b)

	// find previous pageview by same visitor
	previousPageviewID := q.Get("pid")
	if !pageview.IsNewSession && previousPageviewID != "" {
		previousPageview, err := c.Store.GetPageview(previousPageviewID)
		if err != nil && err != datastore.ErrNoResults {
			log.Errorf("error getting previous pageview: %s", err)
			return
		}

		// if we have a recent pageview that is less than 30 minutes old
		if previousPageview != nil && previousPageview.Timestamp.After(now.Add(-30*time.Minute)) {
			previousPageview.Duration = (now.Unix() - previousPageview.Timestamp.Unix())
			previousPageview.IsBounce = false
			previousPageview.IsFinished = true

			// push onto channel to be updated (in batch) later
			c.Pageviews <- previousPageview
		}
	}
}

func (c *Collector) aggregate() {
	agg := aggregator.New(c.Store)
	timeout := 1 * time.Minute

	agg.Run()

	for {
		select {
		case <-time.After(timeout):
			agg.Run()
		}
	}
}

func (c *Collector) worker(cap int, timeout time.Duration) {
	var size int

	for {
		select {
		// persist pageviews in buffer when buffer at capacity
		case p := <-c.Pageviews:
			size = c.buffer(p)
			if size >= cap {
				c.persist()
			}

		// or after timeout passed
		case <-time.After(timeout):
			c.persist()
		}
	}
}

func (c *Collector) buffer(p *models.Pageview) int {
	if !p.IsFinished {
		c.inserts[c.sizei] = p
		c.sizei++
	} else {
		c.updates[c.sizeu] = p
		c.sizeu++
	}

	return (c.sizeu + c.sizei)
}

func (c *Collector) persist() {
	if (c.sizeu + c.sizei) == 0 {
		return
	}

	log.Debugf("persisting %d pageviews (%d inserts, %d updates)", (c.sizeu + c.sizei), c.sizei, c.sizeu)

	if err := c.Store.InsertPageviews(c.inserts[0:c.sizei]); err != nil {
		log.Errorf("error inserting pageviews: %s", err)
	}

	if err := c.Store.UpdatePageviews(c.updates[0:c.sizeu]); err != nil {
		log.Errorf("error updating pageviews: %s", err)
	}

	// reset buffer
	c.sizei = 0
	c.sizeu = 0
}

func shouldCollect(r *http.Request) bool {
	// abort if DNT header is set to "1" (these should have been filtered client-side already)
	if r.Header.Get("DNT") == "1" {
		return false
	}

	// don't track prerendered pages, see https://github.com/usefathom/fathom/issues/13
	if r.Header.Get("X-Moz") == "prefetch" || r.Header.Get("X-Purpose") == "preview" {
		return false
	}

	// abort if this is a bot.
	ua := user_agent.New(r.UserAgent())
	if ua.Bot() {
		return false
	}

	return true
}

func parsePathname(p string) string {
	return "/" + strings.TrimLeft(p, "/")
}

func parseHostname(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	return u.Scheme + "://" + u.Host
}
<commit_msg>discard requests to /collect with missing requried query vars<commit_after>package api

import (
	"encoding/base64"
	"net/http"
	"net/url"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/mssola/user_agent"
	"github.com/usefathom/fathom/pkg/aggregator"
	"github.com/usefathom/fathom/pkg/datastore"
	"github.com/usefathom/fathom/pkg/models"
)

type Collector struct {
	Store     datastore.Datastore
	Pageviews chan *models.Pageview

	// buffer vars
	updates []*models.Pageview
	inserts []*models.Pageview
	sizeu   int
	sizei   int
}

func NewCollector(store datastore.Datastore) *Collector {
	bufferCap := 100                         // persist every 100 pageviews, see https://github.com/usefathom/fathom/issues/132
	bufferTimeout := 1000 * time.Millisecond // or every 1000 ms, whichever comes first

	c := &Collector{
		Store:     store,
		Pageviews: make(chan *models.Pageview),
		updates:   make([]*models.Pageview, bufferCap),
		inserts:   make([]*models.Pageview, bufferCap),
		sizeu:     0,
		sizei:     0,
	}
	go c.aggregate()
	go c.worker(bufferCap, bufferTimeout)
	return c
}

func (c *Collector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !shouldCollect(r) {
		return
	}

	q := r.URL.Query()
	now := time.Now()

	pageview := &models.Pageview{
		ID:             q.Get("id"),
		SiteTrackingID: q.Get("sid"),
		Hostname:       parseHostname(q.Get("h")),
		Pathname:       parsePathname(q.Get("p")),
		IsNewVisitor:   q.Get("nv") == "1",
		IsNewSession:   q.Get("ns") == "1",
		IsUnique:       q.Get("u") == "1",
		Referrer:       q.Get("r"),
		IsFinished:     false,
		IsBounce:       true,
		Duration:       0,
		Timestamp:      now,
	}

	// push pageview onto channel to be inserted (in batch) later
	c.Pageviews <- pageview

	// indicate that we're not tracking user data, see https://github.com/usefathom/fathom/issues/65
	w.Header().Set("Tk", "N")

	// headers to prevent caching
	w.Header().Set("Content-Type", "image/gif")
	w.Header().Set("Expires", "Mon, 01 Jan 1990 00:00:00 GMT")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Pragma", "no-cache")

	// response, 1x1 px transparent GIF
	w.WriteHeader(http.StatusOK)
	b, _ := base64.StdEncoding.DecodeString("R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7")
	w.Write(b)

	// find previous pageview by same visitor
	previousPageviewID := q.Get("pid")
	if !pageview.IsNewSession && previousPageviewID != "" {
		previousPageview, err := c.Store.GetPageview(previousPageviewID)
		if err != nil && err != datastore.ErrNoResults {
			log.Errorf("error getting previous pageview: %s", err)
			return
		}

		// if we have a recent pageview that is less than 30 minutes old
		if previousPageview != nil && previousPageview.Timestamp.After(now.Add(-30*time.Minute)) {
			previousPageview.Duration = (now.Unix() - previousPageview.Timestamp.Unix())
			previousPageview.IsBounce = false
			previousPageview.IsFinished = true

			// push onto channel to be updated (in batch) later
			c.Pageviews <- previousPageview
		}
	}
}

func (c *Collector) aggregate() {
	agg := aggregator.New(c.Store)
	timeout := 1 * time.Minute

	agg.Run()

	for {
		select {
		case <-time.After(timeout):
			agg.Run()
		}
	}
}

func (c *Collector) worker(cap int, timeout time.Duration) {
	var size int

	for {
		select {
		// persist pageviews in buffer when buffer at capacity
		case p := <-c.Pageviews:
			size = c.buffer(p)
			if size >= cap {
				c.persist()
			}

		// or after timeout passed
		case <-time.After(timeout):
			c.persist()
		}
	}
}

func (c *Collector) buffer(p *models.Pageview) int {
	if !p.IsFinished {
		c.inserts[c.sizei] = p
		c.sizei++
	} else {
		c.updates[c.sizeu] = p
		c.sizeu++
	}

	return (c.sizeu + c.sizei)
}

func (c *Collector) persist() {
	if (c.sizeu + c.sizei) == 0 {
		return
	}

	log.Debugf("persisting %d pageviews (%d inserts, %d updates)", (c.sizeu + c.sizei), c.sizei, c.sizeu)

	if err := c.Store.InsertPageviews(c.inserts[0:c.sizei]); err != nil {
		log.Errorf("error inserting pageviews: %s", err)
	}

	if err := c.Store.UpdatePageviews(c.updates[0:c.sizeu]); err != nil {
		log.Errorf("error updating pageviews: %s", err)
	}

	// reset buffer
	c.sizei = 0
	c.sizeu = 0
}

func shouldCollect(r *http.Request) bool {
	// abort if DNT header is set to "1" (these should have been filtered client-side already)
	if r.Header.Get("DNT") == "1" {
		return false
	}

	// don't track prerendered pages, see https://github.com/usefathom/fathom/issues/13
	if r.Header.Get("X-Moz") == "prefetch" || r.Header.Get("X-Purpose") == "preview" {
		return false
	}

	// abort if this is a bot.
	ua := user_agent.New(r.UserAgent())
	if ua.Bot() {
		return false
	}

	// discard if required query vars are missing
	requiredQueryVars := []string{"id", "h", "p"}
	q := r.URL.Query()
	for _, k := range requiredQueryVars {
		if q.Get(k) == "" {
			return false
		}
	}

	return true
}

func parsePathname(p string) string {
	return "/" + strings.TrimLeft(p, "/")
}

func parseHostname(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	return u.Scheme + "://" + u.Host
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 cache

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"time"

	"github.com/GoogleContainerTools/kaniko/pkg/config"
	"github.com/GoogleContainerTools/kaniko/pkg/creds"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// LayerCache is the layer cache
type LayerCache interface {
	RetrieveLayer(string) (v1.Image, error)
}

// RegistryCache is the registry cache
type RegistryCache struct {
	Opts *config.KanikoOptions
}

// RetrieveLayer retrieves a layer from the cache given the cache key ck.
func (rc *RegistryCache) RetrieveLayer(ck string) (v1.Image, error) {
	cache, err := Destination(rc.Opts, ck)
	if err != nil {
		return nil, errors.Wrap(err, "getting cache destination")
	}
	logrus.Infof("Checking for cached layer %s...", cache)

	cacheRef, err := name.NewTag(cache, name.WeakValidation)
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("getting reference for %s", cache))
	}

	registryName := cacheRef.Repository.Registry.Name()
	if rc.Opts.InsecureRegistries.Contains(registryName) {
		newReg, err := name.NewInsecureRegistry(registryName, name.WeakValidation)
		if err != nil {
			return nil, err
		}
		cacheRef.Repository.Registry = newReg
	}

	tr := http.DefaultTransport.(*http.Transport)
	if rc.Opts.SkipTLSVerifyRegistries.Contains(registryName) {
		tr.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	img, err := remote.Image(cacheRef, remote.WithTransport(tr), remote.WithAuthFromKeychain(creds.GetKeychain()))
	if err != nil {
		return nil, err
	}

	cf, err := img.ConfigFile()
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("retrieving config file for %s", cache))
	}

	expiry := cf.Created.Add(rc.Opts.CacheTTL)
	// Layer is stale, rebuild it.
	if expiry.Before(time.Now()) {
		logrus.Infof("Cache entry expired: %s", cache)
		return nil, errors.New(fmt.Sprintf("Cache entry expired: %s", cache))
	}

	// Force the manifest to be populated
	if _, err := img.RawManifest(); err != nil {
		return nil, err
	}
	return img, nil
}

// Destination returns the repo where the layer should be stored
// If no cache is specified, one is inferred from the destination provided
func Destination(opts *config.KanikoOptions, cacheKey string) (string, error) {
	cache := opts.CacheRepo
	if cache == "" {
		destination := opts.Destinations[0]
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return "", errors.Wrap(err, "getting tag for destination")
		}
		return fmt.Sprintf("%s/cache:%s", destRef.Context(), cacheKey), nil
	}
	return fmt.Sprintf("%s:%s", cache, cacheKey), nil
}

// LocalSource retieves a source image from a local cache given cacheKey
func LocalSource(opts *config.KanikoOptions, cacheKey string) (v1.Image, error) {
	cache := opts.CacheDir
	if cache == "" {
		return nil, nil
	}

	path := path.Join(cache, cacheKey)

	fi, err := os.Stat(path)
	if err != nil {
		return nil, errors.Wrap(err, "geting file info")
	}

	// A stale cache is a bad cache
	expiry := fi.ModTime().Add(opts.CacheTTL)
	if expiry.Before(time.Now()) {
		logrus.Debugf("Cached image is too old: %v", fi.ModTime())
		return nil, nil
	}

	logrus.Infof("Found %s in local cache", cacheKey)
	return cachedImageFromPath(path)
}

// cachedImage represents a v1.Tarball that is cached locally in a CAS.
// Computing the digest for a v1.Tarball is very expensive. If the tarball
// is named with the digest we can store this and return it directly rather
// than recompute it.
type cachedImage struct {
	digest string
	v1.Image
}

func (c *cachedImage) Digest() (v1.Hash, error) {
	return v1.NewHash(c.digest)
}

func cachedImageFromPath(p string) (v1.Image, error) {
	imgTar, err := tarball.ImageFromPath(p, nil)
	if err != nil {
		return nil, errors.Wrap(err, "getting image from path")
	}

	return &cachedImage{
		digest: filepath.Base(p),
		Image:  imgTar,
	}, nil
}
<commit_msg>Fix typo<commit_after>/*
Copyright 2018 Google LLC

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 cache

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"time"

	"github.com/GoogleContainerTools/kaniko/pkg/config"
	"github.com/GoogleContainerTools/kaniko/pkg/creds"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/tarball"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// LayerCache is the layer cache
type LayerCache interface {
	RetrieveLayer(string) (v1.Image, error)
}

// RegistryCache is the registry cache
type RegistryCache struct {
	Opts *config.KanikoOptions
}

// RetrieveLayer retrieves a layer from the cache given the cache key ck.
func (rc *RegistryCache) RetrieveLayer(ck string) (v1.Image, error) {
	cache, err := Destination(rc.Opts, ck)
	if err != nil {
		return nil, errors.Wrap(err, "getting cache destination")
	}
	logrus.Infof("Checking for cached layer %s...", cache)

	cacheRef, err := name.NewTag(cache, name.WeakValidation)
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("getting reference for %s", cache))
	}

	registryName := cacheRef.Repository.Registry.Name()
	if rc.Opts.InsecureRegistries.Contains(registryName) {
		newReg, err := name.NewInsecureRegistry(registryName, name.WeakValidation)
		if err != nil {
			return nil, err
		}
		cacheRef.Repository.Registry = newReg
	}

	tr := http.DefaultTransport.(*http.Transport)
	if rc.Opts.SkipTLSVerifyRegistries.Contains(registryName) {
		tr.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	img, err := remote.Image(cacheRef, remote.WithTransport(tr), remote.WithAuthFromKeychain(creds.GetKeychain()))
	if err != nil {
		return nil, err
	}

	cf, err := img.ConfigFile()
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("retrieving config file for %s", cache))
	}

	expiry := cf.Created.Add(rc.Opts.CacheTTL)
	// Layer is stale, rebuild it.
	if expiry.Before(time.Now()) {
		logrus.Infof("Cache entry expired: %s", cache)
		return nil, errors.New(fmt.Sprintf("Cache entry expired: %s", cache))
	}

	// Force the manifest to be populated
	if _, err := img.RawManifest(); err != nil {
		return nil, err
	}
	return img, nil
}

// Destination returns the repo where the layer should be stored
// If no cache is specified, one is inferred from the destination provided
func Destination(opts *config.KanikoOptions, cacheKey string) (string, error) {
	cache := opts.CacheRepo
	if cache == "" {
		destination := opts.Destinations[0]
		destRef, err := name.NewTag(destination, name.WeakValidation)
		if err != nil {
			return "", errors.Wrap(err, "getting tag for destination")
		}
		return fmt.Sprintf("%s/cache:%s", destRef.Context(), cacheKey), nil
	}
	return fmt.Sprintf("%s:%s", cache, cacheKey), nil
}

// LocalSource retieves a source image from a local cache given cacheKey
func LocalSource(opts *config.KanikoOptions, cacheKey string) (v1.Image, error) {
	cache := opts.CacheDir
	if cache == "" {
		return nil, nil
	}

	path := path.Join(cache, cacheKey)

	fi, err := os.Stat(path)
	if err != nil {
		return nil, errors.Wrap(err, "getting file info")
	}

	// A stale cache is a bad cache
	expiry := fi.ModTime().Add(opts.CacheTTL)
	if expiry.Before(time.Now()) {
		logrus.Debugf("Cached image is too old: %v", fi.ModTime())
		return nil, nil
	}

	logrus.Infof("Found %s in local cache", cacheKey)
	return cachedImageFromPath(path)
}

// cachedImage represents a v1.Tarball that is cached locally in a CAS.
// Computing the digest for a v1.Tarball is very expensive. If the tarball
// is named with the digest we can store this and return it directly rather
// than recompute it.
type cachedImage struct {
	digest string
	v1.Image
}

func (c *cachedImage) Digest() (v1.Hash, error) {
	return v1.NewHash(c.digest)
}

func cachedImageFromPath(p string) (v1.Image, error) {
	imgTar, err := tarball.ImageFromPath(p, nil)
	if err != nil {
		return nil, errors.Wrap(err, "getting image from path")
	}

	return &cachedImage{
		digest: filepath.Base(p),
		Image:  imgTar,
	}, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services []string
	files    []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, filepath.Base(f))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<commit_msg>Dump external-dns and dns-controller pod logs<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package dump

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

// logDumper gets all the nodes from a kubernetes cluster and dumps a well-known set of logs
type logDumper struct {
	sshClientFactory sshClientFactory

	artifactsDir string

	services     []string
	files        []string
	podSelectors []string
}

// NewLogDumper is the constructor for a logDumper
func NewLogDumper(sshConfig *ssh.ClientConfig, artifactsDir string) *logDumper {
	sshClientFactory := &sshClientFactoryImplementation{
		sshConfig: sshConfig,
	}

	d := &logDumper{
		sshClientFactory: sshClientFactory,
		artifactsDir:     artifactsDir,
	}

	d.services = []string{
		"node-problem-detector",
		"kubelet",
		"containerd",
		"docker",
		"kops-configuration",
		"protokube",
	}
	d.files = []string{
		"kube-apiserver",
		"kube-scheduler",
		"rescheduler",
		"kube-controller-manager",
		"etcd",
		"etcd-events",
		"etcd-cilium",
		"glbc",
		"cluster-autoscaler",
		"kube-addon-manager",
		"fluentd",
		"kube-proxy",
		"node-problem-detector",
		"cloud-init-output",
		"startupscript",
		"kern",
		"docker",
	}
	d.podSelectors = []string{
		"k8s-app=external-dns",
		"k8s-app=dns-controller",
	}

	return d
}

// DumpAllNodes connects to every node from kubectl get nodes and dumps the logs.
// additionalIPs holds IP addresses of instances found by the deployment tool;
// if the IPs are not found from kubectl get nodes, then these will be dumped also.
// This allows for dumping log on nodes even if they don't register as a kubernetes
// node, or if a node fails to register, or if the whole cluster fails to start.
func (d *logDumper) DumpAllNodes(ctx context.Context, nodes corev1.NodeList, additionalIPs []string) error {
	var dumped []*corev1.Node

	for i := range nodes.Items {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		node := &nodes.Items[i]

		ip := ""
		for _, address := range node.Status.Addresses {
			if address.Type == "ExternalIP" {
				ip = address.Address
				break
			}
		}

		err := d.dumpNode(ctx, node.Name, ip)
		if err != nil {
			log.Printf("could not dump node %s (%s): %v", node.Name, ip, err)
		} else {
			dumped = append(dumped, node)
		}
	}

	notDumped := findInstancesNotDumped(additionalIPs, dumped)
	for _, ip := range notDumped {
		if ctx.Err() != nil {
			log.Printf("stopping dumping nodes: %v", ctx.Err())
			return ctx.Err()
		}

		log.Printf("dumping node not registered in kubernetes: %s", ip)
		err := d.dumpNode(ctx, ip, ip)
		if err != nil {
			log.Printf("error dumping node %s: %v", ip, err)
		}
	}

	return nil
}

// findInstancesNotDumped returns ips from the slice that do not appear as any address of the nodes
func findInstancesNotDumped(ips []string, dumped []*corev1.Node) []string {
	var notDumped []string
	dumpedAddresses := make(map[string]bool)
	for _, node := range dumped {
		for _, address := range node.Status.Addresses {
			dumpedAddresses[address.Address] = true
		}
	}

	for _, ip := range ips {
		if !dumpedAddresses[ip] {
			notDumped = append(notDumped, ip)
		}
	}
	return notDumped
}

// DumpNode connects to a node and dumps the logs.
func (d *logDumper) dumpNode(ctx context.Context, name string, ip string) error {
	if ip == "" {
		return fmt.Errorf("could not find address for %v, ", name)
	}

	log.Printf("Dumping node %s", name)

	n, err := d.connectToNode(ctx, name, ip)
	if err != nil {
		return fmt.Errorf("could not connect: %v", err)
	}

	// As long as we connect to the node we will not return an error;
	// a failure to collect a log (or even any logs at all) is not
	// considered an error in dumping the node.
	// TODO(justinsb): clean up / rationalize
	errors := n.dump(ctx)
	for _, e := range errors {
		log.Printf("error dumping node %s: %v", name, e)
	}

	if err := n.Close(); err != nil {
		log.Printf("error closing connection: %v", err)
	}

	return nil
}

// sshClient is an interface abstracting *ssh.Client, which allows us to test it
type sshClient interface {
	io.Closer

	// ExecPiped runs the command, piping stdout & stderr
	ExecPiped(ctx context.Context, command string, stdout io.Writer, stderr io.Writer) error
}

// sshClientFactory is an interface abstracting to a node over SSH
type sshClientFactory interface {
	Dial(ctx context.Context, host string) (sshClient, error)
}

// logDumperNode holds state for a particular node we are dumping
type logDumperNode struct {
	client sshClient
	dumper *logDumper

	dir string
}

// connectToNode makes an SSH connection to the node and returns a logDumperNode
func (d *logDumper) connectToNode(ctx context.Context, nodeName string, host string) (*logDumperNode, error) {
	client, err := d.sshClientFactory.Dial(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("unable to SSH to %q: %v", host, err)
	}
	return &logDumperNode{
		client: client,
		dir:    filepath.Join(d.artifactsDir, nodeName),
		dumper: d,
	}, nil
}

// logDumperNode cleans up any state in the logDumperNode
func (n *logDumperNode) Close() error {
	return n.client.Close()
}

// dump captures the well-known set of logs
func (n *logDumperNode) dump(ctx context.Context) []error {
	if ctx.Err() != nil {
		return []error{ctx.Err()}
	}

	var errors []error

	// Capture kernel log
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise -k", filepath.Join(n.dir, "kern.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture full journal - needed so we can see e.g. disk mounts
	// This does duplicate the other files, but ensures we have all output
	if err := n.shellToFile(ctx, "sudo journalctl --output=short-precise", filepath.Join(n.dir, "journal.log")); err != nil {
		errors = append(errors, err)
	}

	// Capture logs from any systemd services in our list that are registered
	services, err := n.listSystemdUnits(ctx)
	if err != nil {
		errors = append(errors, fmt.Errorf("error listing systemd services: %v", err))
	}
	for _, s := range n.dumper.services {
		name := s + ".service"
		for _, service := range services {
			if service == name {
				if err := n.shellToFile(ctx, "sudo journalctl --output=cat -u "+name, filepath.Join(n.dir, s+".log")); err != nil {
					errors = append(errors, err)
				}
			}
		}
	}

	// Capture any file logs where the files exist
	fileList, err := n.findFiles(ctx, "/var/log")
	if err != nil {
		errors = append(errors, fmt.Errorf("error reading /var/log: %v", err))
	}
	for _, name := range n.dumper.files {
		prefix := "/var/log/" + name + ".log"
		for _, f := range fileList {
			if !strings.HasPrefix(f, prefix) {
				continue
			}
			if err := n.shellToFile(ctx, "sudo cat '"+strings.ReplaceAll(f, "'", "'\\''")+"'", filepath.Join(n.dir, filepath.Base(f))); err != nil {
				errors = append(errors, err)
			}
		}
	}

	for _, selector := range n.dumper.podSelectors {
		kv := strings.Split(selector, "=")
		app := kv[len(kv)-1]
		if err := n.shellToFile(ctx, "if command -v kubectl &> /dev/null; then kubectl logs -n kube-system --all-containers -l \""+selector+"\"; fi", filepath.Join(n.dir, app)); err != nil {
			errors = append(errors, err)
		}
	}

	return errors
}

// findFiles lists files under the specified directory (recursively)
func (n *logDumperNode) findFiles(ctx context.Context, dir string) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo find "+dir+" -print0", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing %q: %v", dir, err)
	}

	paths := []string{}
	for _, b := range bytes.Split(stdout.Bytes(), []byte{0}) {
		if len(b) == 0 {
			// Likely the last value
			continue
		}
		paths = append(paths, string(b))
	}
	return paths, nil
}

// listSystemdUnits returns the list of systemd units on the node
func (n *logDumperNode) listSystemdUnits(ctx context.Context) ([]string, error) {
	var stdout bytes.Buffer
	var stderr bytes.Buffer

	err := n.client.ExecPiped(ctx, "sudo systemctl list-units -t service --no-pager --no-legend --all", &stdout, &stderr)
	if err != nil {
		return nil, fmt.Errorf("error listing systemd units: %v", err)
	}

	var services []string
	for _, line := range strings.Split(stdout.String(), "\n") {
		tokens := strings.Fields(line)
		if len(tokens) == 0 || tokens[0] == "" {
			continue
		}
		services = append(services, tokens[0])
	}
	return services, nil
}

// shellToFile executes a command and copies the output to a file
func (n *logDumperNode) shellToFile(ctx context.Context, command string, destPath string) error {
	if err := os.MkdirAll(filepath.Dir(destPath), 0755); err != nil {
		log.Printf("unable to mkdir on %q: %v", filepath.Dir(destPath), err)
	}

	f, err := os.Create(destPath)
	if err != nil {
		return fmt.Errorf("error creating file %q: %v", destPath, err)
	}
	defer f.Close()

	if err := n.client.ExecPiped(ctx, command, f, f); err != nil {
		return fmt.Errorf("error executing command %q: %v", command, err)
	}

	return nil
}

// sshClientImplementation is the default implementation of sshClient, binding to a *ssh.Client
type sshClientImplementation struct {
	client *ssh.Client
}

var _ sshClient = &sshClientImplementation{}

// ExecPiped implements sshClientImplementation::ExecPiped
func (s *sshClientImplementation) ExecPiped(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error {
	if ctx.Err() != nil {
		return ctx.Err()
	}

	finished := make(chan error)
	go func() {
		session, err := s.client.NewSession()
		if err != nil {
			finished <- fmt.Errorf("error creating ssh session: %v", err)
			return
		}
		defer session.Close()

		klog.V(2).Infof("running SSH command: %v", cmd)

		session.Stdout = stdout
		session.Stderr = stderr

		finished <- session.Run(cmd)
	}()

	select {
	case <-ctx.Done():
		log.Print("closing SSH tcp connection due to context completion")

		// terminate the TCP connection to force a disconnect - we assume everyone is using the same context.
		// We could make this better by sending a signal on the session, waiting and then closing the session,
		// and only if we still haven't succeeded then closing the TCP connection.  This is sufficient for our
		// current usage though - and hopefully that logic will be implemented in the SSH package itself.
		s.Close()

		<-finished // Wait for cancellation
		return ctx.Err()

	case err := <-finished:
		return err
	}
}

// Close implements sshClientImplementation::Close
func (s *sshClientImplementation) Close() error {
	return s.client.Close()
}

// sshClientFactoryImplementation is the default implementation of sshClientFactory
type sshClientFactoryImplementation struct {
	sshConfig *ssh.ClientConfig
}

var _ sshClientFactory = &sshClientFactoryImplementation{}

// Dial implements sshClientFactory::Dial
func (f *sshClientFactoryImplementation) Dial(ctx context.Context, host string) (sshClient, error) {
	addr := host + ":22"
	d := net.Dialer{}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}

	// We have a TCP connection; we will force-close it to support context cancellation

	var client *ssh.Client
	finished := make(chan error)
	go func() {
		c, chans, reqs, err := ssh.NewClientConn(conn, addr, f.sshConfig)
		if err == nil {
			client = ssh.NewClient(c, chans, reqs)
		}
		finished <- err
	}()

	select {
	case <-ctx.Done():
		log.Print("cancelling SSH tcp connection due to context completion")
		conn.Close() // Close the TCP connection to force cancellation
		<-finished   // Wait for cancellation
		return nil, ctx.Err()
	case err := <-finished:
		if err != nil {
			return nil, err
		}
		return &sshClientImplementation{
			client: client,
		}, nil
	}
}
<|endoftext|>


<commit_before>// Package tros provides methods to sort slices of structs using reflection
package tros

import (
	"fmt"
	"reflect"
	"sort"
)

// SortInterface returns a sort.Interface suitable for sorting, reversing,
// searching, etc., a slice of structs by with the given field name fn
//
// i must be a slice of homogeneous structs. That is, all structs must have an
// exported field with the name fn, and those fields must all be of the same
// type.
//
// Operations using this sort.Interface are significantly (~2.5x) slower than the
// standard Go sorting idiom, because it relies heavily on reflection.
func SortInterface(i interface{}, fn string) (sort.Interface, error) {
	sval := reflect.ValueOf(i)
	if sval.Kind() != reflect.Slice {
		return nil, fmt.Errorf("non-slice interface, got %q", sval.Kind())
	}
	l := sval.Len()
	if l == 0 {
		return nil, fmt.Errorf("slice is empty")
	}

	if fn == "" {
		return nil, fmt.Errorf("must specify field name")
	}
	reverse := false
	if fn[0] == '-' {
		reverse = true
		fn = fn[1:]
	}

	vals := make([]reflect.Value, l)
	var fs []reflect.Value
	var ls []Lesser

	f0 := sval.Index(0).FieldByName(fn)
	k := f0.Kind()
	if k == reflect.Invalid {
		return nil, fmt.Errorf("no field with name %q", fn)
	}
	if !f0.CanSet() {
		return nil, fmt.Errorf("field %q is not exported", fn)
	}
	if k > reflect.Float64 && k != reflect.String && k != reflect.Struct {
		return nil, fmt.Errorf("unsupported kind %q", k)
	}
	for i := 0; i < l; i++ {
		v := sval.Index(i)
		f := v.FieldByName(fn)
		if f.Kind() != k {
			return nil, fmt.Errorf("unmatched field kinds, %q vs %q", f.Kind(), k)
		}
		if less, ok := f.Interface().(Lesser); ok {
			if ls == nil {
				ls = make([]Lesser, l)
			}
			ls[i] = less
		} else if f.Kind() == reflect.Struct {
			return nil, fmt.Errorf("struct field %q does not implement Lesser", fn)
		} else {
			if fs == nil {
				fs = make([]reflect.Value, l)
			}
			fs[i] = f
		}
		vals[i] = v
	}
	tmp := reflect.New(vals[0].Type()).Elem()
	return sort.Interface(&sortable{vals, k, tmp, fs, ls, reverse}), nil
}

// Sort sorts a slice of structs based on the values of the structs' fields with
// the given field name fn
//
// i must be a slice of homogeneous structs. That is, all structs must have an
// exported field with the name fn, and those fields must all be of the same
// type.
//
// Sort is significantly (~10x) slower than the standard Go sort.Sort function,
// because it relies heavily on reflection.
func Sort(i interface{}, fn string) error {
	s, err := SortInterface(i, fn)
	if err != nil {
		return err
	}
	sort.Sort(s)
	return nil
}

type sortable struct {
	vals []reflect.Value
	k    reflect.Kind
	tmp  reflect.Value   // reused for swapping
	fs   []reflect.Value // used for comparing structs by non-Lesser fields
	ls   []Lesser        // used for comparing structs by Lesser fields
	rev  bool            // if true, reverse sort order
}

func (s *sortable) Len() int { return len(s.vals) }

func (s *sortable) Swap(i, j int) {
	a, b := s.vals[i], s.vals[j]
	s.tmp.Set(a)
	a.Set(b)
	b.Set(s.tmp)

	if s.ls != nil {
		s.ls[i], s.ls[j] = s.ls[j], s.ls[i]
	}
}

func (s *sortable) Less(i, j int) bool {
	if s.ls != nil {
		return s.ls[i].Less(s.ls[j])
	}

	af, bf := s.fs[i], s.fs[j]
	var r bool
	switch s.k {
	case reflect.Bool:
		r = !af.Bool() && bf.Bool()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		r = af.Int() < bf.Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		r = af.Uint() < bf.Uint()
	case reflect.Float32, reflect.Float64:
		r = af.Float() < bf.Float()
	case reflect.String:
		r = af.String() < bf.String()
	default:
		panic("unreachable: invalid Kind") // Check in Sort should prevent this
	}
	if s.rev {
		r = !r
	}
	return r
}

// Fields implementing this interface may be used to sort structs using the
// field's implementation of Less.
type Lesser interface {
	Less(other Lesser) bool
}
<commit_msg>pick nit about slice initialization<commit_after>// Package tros provides methods to sort slices of structs using reflection
package tros

import (
	"fmt"
	"reflect"
	"sort"
)

// SortInterface returns a sort.Interface suitable for sorting, reversing,
// searching, etc., a slice of structs by with the given field name fn
//
// i must be a slice of homogeneous structs. That is, all structs must have an
// exported field with the name fn, and those fields must all be of the same
// type.
//
// Operations using this sort.Interface are significantly (~2.5x) slower than the
// standard Go sorting idiom, because it relies heavily on reflection.
func SortInterface(i interface{}, fn string) (sort.Interface, error) {
	sval := reflect.ValueOf(i)
	if sval.Kind() != reflect.Slice {
		return nil, fmt.Errorf("non-slice interface, got %q", sval.Kind())
	}
	l := sval.Len()
	if l == 0 {
		return nil, fmt.Errorf("slice is empty")
	}

	if fn == "" {
		return nil, fmt.Errorf("must specify field name")
	}
	reverse := false
	if fn[0] == '-' {
		reverse = true
		fn = fn[1:]
	}

	vals := make([]reflect.Value, l)
	var fs []reflect.Value
	var ls []Lesser

	f0 := sval.Index(0).FieldByName(fn)
	k := f0.Kind()
	if k == reflect.Invalid {
		return nil, fmt.Errorf("no field with name %q", fn)
	}
	if !f0.CanSet() {
		return nil, fmt.Errorf("field %q is not exported", fn)
	}
	if k > reflect.Float64 && k != reflect.String && k != reflect.Struct {
		return nil, fmt.Errorf("unsupported kind %q", k)
	}
	for i := 0; i < l; i++ {
		v := sval.Index(i)
		f := v.FieldByName(fn)
		if f.Kind() != k {
			return nil, fmt.Errorf("unmatched field kinds, %q vs %q", f.Kind(), k)
		}
		if less, ok := f.Interface().(Lesser); ok {
			ls = append(ls, less)
		} else if f.Kind() == reflect.Struct {
			return nil, fmt.Errorf("struct field %q does not implement Lesser", fn)
		} else {
			fs = append(fs, f)
		}
		vals[i] = v
	}
	tmp := reflect.New(vals[0].Type()).Elem()
	return sort.Interface(&sortable{vals, k, tmp, fs, ls, reverse}), nil
}

// Sort sorts a slice of structs based on the values of the structs' fields with
// the given field name fn
//
// i must be a slice of homogeneous structs. That is, all structs must have an
// exported field with the name fn, and those fields must all be of the same
// type.
//
// Sort is significantly (~10x) slower than the standard Go sort.Sort function,
// because it relies heavily on reflection.
func Sort(i interface{}, fn string) error {
	s, err := SortInterface(i, fn)
	if err != nil {
		return err
	}
	sort.Sort(s)
	return nil
}

type sortable struct {
	vals []reflect.Value
	k    reflect.Kind
	tmp  reflect.Value   // reused for swapping
	fs   []reflect.Value // used for comparing structs by non-Lesser fields
	ls   []Lesser        // used for comparing structs by Lesser fields
	rev  bool            // if true, reverse sort order
}

func (s *sortable) Len() int { return len(s.vals) }

func (s *sortable) Swap(i, j int) {
	a, b := s.vals[i], s.vals[j]
	s.tmp.Set(a)
	a.Set(b)
	b.Set(s.tmp)

	if s.ls != nil {
		s.ls[i], s.ls[j] = s.ls[j], s.ls[i]
	}
}

func (s *sortable) Less(i, j int) bool {
	if s.ls != nil {
		return s.ls[i].Less(s.ls[j])
	}

	af, bf := s.fs[i], s.fs[j]
	var r bool
	switch s.k {
	case reflect.Bool:
		r = !af.Bool() && bf.Bool()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		r = af.Int() < bf.Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		r = af.Uint() < bf.Uint()
	case reflect.Float32, reflect.Float64:
		r = af.Float() < bf.Float()
	case reflect.String:
		r = af.String() < bf.String()
	default:
		panic("unreachable: invalid Kind") // Check in Sort should prevent this
	}
	if s.rev {
		r = !r
	}
	return r
}

// Fields implementing this interface may be used to sort structs using the
// field's implementation of Less.
type Lesser interface {
	Less(other Lesser) bool
}
<|endoftext|>


<commit_before>package store

import (
	"bufio"
	"fmt"
	"io"
	"strings"

	"github.com/oklog/ulid"
)

func mergeRecords(w io.Writer, readers ...io.Reader) (low, high ulid.ULID, n int64, err error) {
	// Optimization and safety.
	if len(readers) == 0 {
		return low, high, 0, nil
	}

	// Initialize our state.
	var (
		first   = true
		scanner = make([]*bufio.Scanner, len(readers))
		ok      = make([]bool, len(readers))
		record  = make([]string, len(readers))
		id      = make([]string, len(readers))
	)

	// Advance moves the next record from a reader into our state.
	advance := func(i int) error {
		if ok[i] = scanner[i].Scan(); ok[i] {
			record[i] = scanner[i].Text()
			id[i] = strings.Fields(record[i])[0] // TODO(pb): profiling reveals this to be $expensive$
		} else if err := scanner[i].Err(); err != nil && err != io.EOF {
			return err
		}
		return nil
	}

	// Initialize all of the scanners and their first record.
	for i := 0; i < len(readers); i++ {
		scanner[i] = bufio.NewScanner(readers[i])
		if err := advance(i); err != nil {
			return low, high, n, err
		}
	}

	// Loop until all readers are drained.
	var chosen int
	for {
		// Choose the smallest ULID.
		chosen = -1
		for i := 0; i < len(readers); i++ {
			if !ok[i] {
				continue
			}
			if chosen < 0 || id[i] < id[chosen] {
				chosen = i
			}
		}
		if chosen < 0 {
			break // drained all the scanners
		}

		// Record the first (lowest) and last (highest) ULIDs.
		id := ulid.MustParse(id[chosen])
		if first {
			low, first = id, false
		} else if !first && id == high {
			// Duplicate detected!
			if err := advance(chosen); err != nil {
				return low, high, n, err
			}
			continue
		}
		high = id

		// Write the record.
		n0, err := fmt.Fprintf(w, "%s\n", record[chosen])
		if err != nil {
			return low, high, n, err
		}
		n += int64(n0)

		// Advance the chosen scanner by 1 record.
		if err := advance(chosen); err != nil {
			return low, high, n, err
		}
	}

	return low, high, n, nil
}

// mergeRecordsToLog is a specialization of mergeRecords.
// It enforces segmentTargetSize by creating WriteSegments as necessary.
// It has best-effort semantics, e.g. it won't split large records.
func mergeRecordsToLog(dst Log, segmentTargetSize int64, readers ...io.Reader) (n int64, err error) {
	// Optimization and safety.
	if len(readers) == 0 {
		return 0, nil
	}

	// Per-segment state.
	writeSegment, err := dst.Create()
	if err != nil {
		return n, err
	}
	defer func() {
		// Don't leak active segments.
		if writeSegment != nil {
			if err := writeSegment.Delete(); err != nil {
				panic(err)
			}
		}
	}()
	var (
		nSegment int64
		first    = true
		low      ulid.ULID
		high     ulid.ULID
	)

	// Global state.
	var (
		scanner = make([]*bufio.Scanner, len(readers))
		ok      = make([]bool, len(readers))
		record  = make([]string, len(readers))
		id      = make([]string, len(readers))
	)

	// Advance moves the next record from a reader into our state.
	advance := func(i int) error {
		if ok[i] = scanner[i].Scan(); ok[i] {
			record[i] = scanner[i].Text()
			id[i] = strings.Fields(record[i])[0]
		} else if err := scanner[i].Err(); err != nil && err != io.EOF {
			return err
		}
		return nil
	}

	// Initialize all of the scanners and their first record.
	for i := 0; i < len(readers); i++ {
		scanner[i] = bufio.NewScanner(readers[i])
		if err := advance(i); err != nil {
			return n, err
		}
	}

	// Loop until all readers are drained.
	var chosen int
	for {
		// Choose the smallest ULID.
		chosen = -1
		for i := 0; i < len(readers); i++ {
			if !ok[i] {
				continue
			}
			if chosen < 0 || id[i] < id[chosen] {
				chosen = i
			}
		}
		if chosen < 0 {
			break // drained all the scanners
		}

		// Record the first (lowest) and last (highest) ULIDs.
		id := ulid.MustParse(id[chosen])
		if first {
			low, first = id, false
		} else if !first && id == high {
			// Duplicate detected!
			if err := advance(chosen); err != nil {
				return n, err
			}
			continue
		}
		high = id

		// Write the record.
		n0, err := fmt.Fprintf(writeSegment, "%s\n", record[chosen])
		if err != nil {
			return n, err
		}
		n += int64(n0)
		nSegment += int64(n0)

		if nSegment >= segmentTargetSize {
			// We've met our target size. Close the segment.
			if err := writeSegment.Close(low, high); err != nil {
				return n, err
			}

			// Create a new segment, and reset our per-segment state.
			writeSegment, err = dst.Create()
			if err != nil {
				return n, err
			}
			nSegment = 0
			first = true
		}

		// Advance the chosen scanner by 1 record.
		if err := advance(chosen); err != nil {
			return n, err
		}
	}

	if nSegment > 0 {
		if err = writeSegment.Close(low, high); err != nil {
			return n, err
		}
	} else {
		err = writeSegment.Delete()
	}
	writeSegment = nil
	return n, err
}
<commit_msg>Mark for improvements pt2<commit_after>package store

import (
	"bufio"
	"fmt"
	"io"
	"strings"

	"github.com/oklog/ulid"
)

func mergeRecords(w io.Writer, readers ...io.Reader) (low, high ulid.ULID, n int64, err error) {
	// Optimization and safety.
	if len(readers) == 0 {
		return low, high, 0, nil
	}

	// Initialize our state.
	var (
		first   = true
		scanner = make([]*bufio.Scanner, len(readers))
		ok      = make([]bool, len(readers))
		record  = make([]string, len(readers))
		id      = make([]string, len(readers))
	)

	// Advance moves the next record from a reader into our state.
	advance := func(i int) error {
		if ok[i] = scanner[i].Scan(); ok[i] {
			record[i] = scanner[i].Text()
			id[i] = strings.Fields(record[i])[0] // TODO(pb): profiling reveals this to be $expensive$
		} else if err := scanner[i].Err(); err != nil && err != io.EOF {
			return err
		}
		return nil
	}

	// Initialize all of the scanners and their first record.
	for i := 0; i < len(readers); i++ {
		scanner[i] = bufio.NewScanner(readers[i])
		if err := advance(i); err != nil {
			return low, high, n, err
		}
	}

	// Loop until all readers are drained.
	var chosen int
	for {
		// Choose the smallest ULID.
		chosen = -1
		for i := 0; i < len(readers); i++ {
			if !ok[i] {
				continue
			}
			if chosen < 0 || id[i] < id[chosen] {
				chosen = i
			}
		}
		if chosen < 0 {
			break // drained all the scanners
		}

		// Record the first (lowest) and last (highest) ULIDs.
		id := ulid.MustParse(id[chosen])
		if first {
			low, first = id, false
		} else if !first && id == high {
			// Duplicate detected!
			if err := advance(chosen); err != nil {
				return low, high, n, err
			}
			continue
		}
		high = id

		// Write the record.
		n0, err := fmt.Fprintf(w, "%s\n", record[chosen])
		if err != nil {
			return low, high, n, err
		}
		n += int64(n0)

		// Advance the chosen scanner by 1 record.
		if err := advance(chosen); err != nil {
			return low, high, n, err
		}
	}

	return low, high, n, nil
}

// mergeRecordsToLog is a specialization of mergeRecords.
// It enforces segmentTargetSize by creating WriteSegments as necessary.
// It has best-effort semantics, e.g. it won't split large records.
func mergeRecordsToLog(dst Log, segmentTargetSize int64, readers ...io.Reader) (n int64, err error) {
	// Optimization and safety.
	if len(readers) == 0 {
		return 0, nil
	}

	// Per-segment state.
	writeSegment, err := dst.Create()
	if err != nil {
		return n, err
	}
	defer func() {
		// Don't leak active segments.
		if writeSegment != nil {
			if err := writeSegment.Delete(); err != nil {
				panic(err)
			}
		}
	}()
	var (
		nSegment int64
		first    = true
		low      ulid.ULID
		high     ulid.ULID
	)

	// Global state.
	var (
		scanner = make([]*bufio.Scanner, len(readers))
		ok      = make([]bool, len(readers))
		record  = make([]string, len(readers))
		id      = make([]string, len(readers))
	)

	// Advance moves the next record from a reader into our state.
	advance := func(i int) error {
		if ok[i] = scanner[i].Scan(); ok[i] {
			record[i] = scanner[i].Text()
			id[i] = strings.Fields(record[i])[0] // TODO(pb): profiling reveals this to be $expensive$
		} else if err := scanner[i].Err(); err != nil && err != io.EOF {
			return err
		}
		return nil
	}

	// Initialize all of the scanners and their first record.
	for i := 0; i < len(readers); i++ {
		scanner[i] = bufio.NewScanner(readers[i])
		if err := advance(i); err != nil {
			return n, err
		}
	}

	// Loop until all readers are drained.
	var chosen int
	for {
		// Choose the smallest ULID.
		chosen = -1
		for i := 0; i < len(readers); i++ {
			if !ok[i] {
				continue
			}
			if chosen < 0 || id[i] < id[chosen] {
				chosen = i
			}
		}
		if chosen < 0 {
			break // drained all the scanners
		}

		// Record the first (lowest) and last (highest) ULIDs.
		id := ulid.MustParse(id[chosen])
		if first {
			low, first = id, false
		} else if !first && id == high {
			// Duplicate detected!
			if err := advance(chosen); err != nil {
				return n, err
			}
			continue
		}
		high = id

		// Write the record.
		n0, err := fmt.Fprintf(writeSegment, "%s\n", record[chosen])
		if err != nil {
			return n, err
		}
		n += int64(n0)
		nSegment += int64(n0)

		if nSegment >= segmentTargetSize {
			// We've met our target size. Close the segment.
			if err := writeSegment.Close(low, high); err != nil {
				return n, err
			}

			// Create a new segment, and reset our per-segment state.
			writeSegment, err = dst.Create()
			if err != nil {
				return n, err
			}
			nSegment = 0
			first = true
		}

		// Advance the chosen scanner by 1 record.
		if err := advance(chosen); err != nil {
			return n, err
		}
	}

	if nSegment > 0 {
		if err = writeSegment.Close(low, high); err != nil {
			return n, err
		}
	} else {
		err = writeSegment.Delete()
	}
	writeSegment = nil
	return n, err
}
<|endoftext|>


<commit_before>package utils

import (
	log "github.com/Sirupsen/logrus"
	"github.com/docker/libcompose/docker"
	"github.com/docker/libcompose/project"
	"io/ioutil"
	"os"
	"runtime/debug"
)

func ParseDockerCompose(filePath string) (composeObject *project.Project, err error) {
	content, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	sanitizedCompose := []byte(os.ExpandEnv(string(content)))
	tmpfile, err := ioutil.TempFile("/tmp", "sanitized-")
	if err != nil {
		return nil, err
	}
	defer os.Remove(tmpfile.Name())

	if _, err := tmpfile.Write(sanitizedCompose); err != nil {
		log.Fatal(err)
	}
	if err := tmpfile.Close(); err != nil {
		log.Fatal(err)
	}

	return libcomposeParser(tmpfile.Name())
}

func libcomposeParser(filePath string) (composeObject *project.Project, err error) {
	defer func() {
		if r := recover(); r != nil {
			log.Infof("Recovered in pkg/utils/utils.go libcomposeParser(filePath string) %#v\n", r)
			debug.PrintStack()
			err = r.(error)
		}
	}()
	context := &docker.Context{}
	if filePath == "" {
		filePath = "docker-compose.yml"
	}
	context.ComposeFiles = []string{filePath}
	composeObject = project.NewProject(&context.Context, nil, nil)
	err = composeObject.Parse()
	if err != nil {
		log.Fatalf("Failed to load compose file", err)
		return nil, err
	}
	log.Infof("Post-parsing: %#v\n", composeObject)
	return composeObject, nil
}
<commit_msg>Use libcompose lookup<commit_after>package utils

import (
	log "github.com/Sirupsen/logrus"
	"github.com/docker/libcompose/config"
	"github.com/docker/libcompose/docker"
	"github.com/docker/libcompose/lookup"
	"github.com/docker/libcompose/project"
	"path/filepath"
	// "io/ioutil"
	"os"
	"runtime/debug"
)

func ParseDockerCompose(filePath string) (composeObject *project.Project, err error) {
	defer func() {
		if r := recover(); r != nil {
			log.Infof("Recovered in pkg/utils/utils.go libcomposeParser(filePath string) %#v\n", r)
			debug.PrintStack()
			err = r.(error)
		}
	}()
	context := &docker.Context{}
	if filePath == "" {
		filePath = "docker-compose.yml"
	}
	context.ComposeFiles = []string{filePath}
	if context.Context.EnvironmentLookup == nil {
		cwd, err := os.Getwd()
		if err != nil {
			return nil, err
		}
		context.EnvironmentLookup = &lookup.ComposableEnvLookup{
			Lookups: []config.EnvironmentLookup{
				&lookup.EnvfileLookup{
					Path: filepath.Join(cwd, ".env"),
				},
				&lookup.OsEnvLookup{},
			},
		}
	}
	composeObject = project.NewProject(&context.Context, nil, nil)
	err = composeObject.Parse()
	if err != nil {
		log.Fatalf("Failed to load compose file", err)
		return nil, err
	}
	log.Infof("Post-parsing: %#v\n", composeObject)
	return composeObject, nil
}
<|endoftext|>