<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"text/template"

	"github.com/gorilla/mux"
	. "github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/ipam"
	"github.com/weaveworks/weave/mesh"
	"github.com/weaveworks/weave/nameserver"
	"github.com/weaveworks/weave/net/address"
	weave "github.com/weaveworks/weave/router"
)

var rootTemplate = template.New("root").Funcs(map[string]interface{}{
	"countDNSEntries": func(entries []nameserver.EntryStatus) int {
		count := 0
		for _, entry := range entries {
			if entry.Tombstone == 0 {
				count++
			}
		}
		return count
	},
	"upstreamServers": func(servers []string) string {
		if len(servers) == 0 {
			return "none"
		}
		return strings.Join(servers, ", ")
	},
	"printIPAMRanges": func(router weave.NetworkRouterStatus, status ipam.Status) string {
		var buffer bytes.Buffer

		type stats struct {
			ips       uint32
			nickname  string
			reachable bool
		}

		peerStats := make(map[string]*stats)

		for _, entry := range status.Entries {
			s, found := peerStats[entry.Peer]
			if !found {
				s = &stats{nickname: entry.Nickname, reachable: entry.IsKnownPeer}
				peerStats[entry.Peer] = s
			}
			s.ips += entry.Size
		}

		printOwned := func(name string, nickName string, reachable bool, ips uint32) {
			reachableStr := ""
			if !reachable {
				reachableStr = "- unreachable!"
			}
			percentageRanges := float32(ips) * 100.0 / float32(status.RangeNumIPs)

			displayName := name + "(" + nickName + ")"
			fmt.Fprintf(&buffer, "%-37v %8d IPs (%04.1f%% of total) %s\n",
				displayName, ips, percentageRanges, reachableStr)
		}

		// print the local info first
		if ourStats := peerStats[router.Name]; ourStats != nil {
			printOwned(router.Name, ourStats.nickname, true, ourStats.ips)
		}

		// and then the rest
		for peer, stats := range peerStats {
			if peer != router.Name {
				printOwned(peer, stats.nickname, stats.reachable, stats.ips)
			}
		}

		return buffer.String()
	},
	"allIPAMOwnersUnreachable": func(status ipam.Status) bool {
		for _, entry := range status.Entries {
			if entry.Size > 0 && entry.IsKnownPeer {
				return false
			}
		}
		return true
	},
	"printConnectionCounts": func(conns []mesh.LocalConnectionStatus) string {
		counts := make(map[string]int)
		for _, conn := range conns {
			counts[conn.State]++
		}
		return printCounts(counts, []string{"established", "pending", "retrying", "failed", "connecting"})
	},
	"printPeerConnectionCounts": func(peers []mesh.PeerStatus) string {
		counts := make(map[string]int)
		for _, peer := range peers {
			for _, conn := range peer.Connections {
				if conn.Established {
					counts["established"]++
				} else {
					counts["pending"]++
				}
			}
		}
		return printCounts(counts, []string{"established", "pending"})
	},
	"printState": func(enabled bool) string {
		if enabled {
			return "enabled"
		}
		return "disabled"
	},
	"trimSuffix": strings.TrimSuffix,
})

// Print counts in a specified order
func printCounts(counts map[string]int, keys []string) string {
	var stringCounts []string
	for _, key := range keys {
		if count, ok := counts[key]; ok {
			stringCounts = append(stringCounts, fmt.Sprintf("%d %s", count, key))
		}
	}
	return strings.Join(stringCounts, ", ")
}

// Strip escaped newlines from template
func escape(template string) string {
	return strings.Replace(template, "\\\n", "", -1)
}

// Define a named template panicking on error
func defTemplate(name string, text string) *template.Template {
	return template.Must(rootTemplate.New(name).Parse(escape(text)))
}

var statusTemplate = defTemplate("status", `\
       Version: {{.Version}}

       Service: router
      Protocol: {{.Router.Protocol}} \
{{if eq .Router.ProtocolMinVersion .Router.ProtocolMaxVersion}}\
{{.Router.ProtocolMaxVersion}}\
{{else}}\
{{.Router.ProtocolMinVersion}}..{{.Router.ProtocolMaxVersion}}\
{{end}}
          Name: {{.Router.Name}}({{.Router.NickName}})
    Encryption: {{printState .Router.Encryption}}
 PeerDiscovery: {{printState .Router.PeerDiscovery}}
       Targets: {{len .Router.Targets}}
   Connections: {{len .Router.Connections}}{{with printConnectionCounts .Router.Connections}} ({{.}}){{end}}
         Peers: {{len .Router.Peers}}{{with printPeerConnectionCounts .Router.Peers}} (with {{.}} connections){{end}}
{{if .IPAM}}\

       Service: ipam
{{if .IPAM.Entries}}\
{{if allIPAMOwnersUnreachable .IPAM}}\
        Status: all IP ranges owned by unreachable peers - use 'rmpeer' if they are dead
{{else if len .IPAM.PendingAllocates}}\
        Status: waiting for IP range grant from peers
{{else}}\
        Status: ready
{{end}}\
{{else if .IPAM.Paxos}}\
        Status: awaiting consensus (quorum: {{.IPAM.Paxos.Quorum}}, known: {{.IPAM.Paxos.KnownNodes}})
{{else}}\
        Status: idle
{{end}}\
         Range: {{.IPAM.Range}}
 DefaultSubnet: {{.IPAM.DefaultSubnet}}
{{end}}\
{{if .DNS}}\

       Service: dns
        Domain: {{.DNS.Domain}}
      Upstream: {{upstreamServers .DNS.Upstream}}
           TTL: {{.DNS.TTL}}
       Entries: {{countDNSEntries .DNS.Entries}}
{{end}}\
`)

var targetsTemplate = defTemplate("targetsTemplate", `\
{{range .Router.Targets}}{{.}}
{{end}}\
`)

var connectionsTemplate = defTemplate("connectionsTemplate", `\
{{range .Router.Connections}}\
{{if .Outbound}}->{{else}}<-{{end}} {{printf "%-21v" .Address}} {{printf "%-11v" .State}} {{.Info}}
{{end}}\
`)

var peersTemplate = defTemplate("peers", `\
{{range .Router.Peers}}\
{{.Name}}({{.NickName}})
{{range .Connections}}\
   {{if .Outbound}}->{{else}}<-{{end}} {{printf "%-21v" .Address}} \
{{$nameNickName := printf "%v(%v)" .Name .NickName}}{{printf "%-37v" $nameNickName}} \
{{if .Established}}established{{else}}pending{{end}}
{{end}}\
{{end}}\
`)

var dnsEntriesTemplate = defTemplate("dnsEntries", `\
{{$domain := printf ".%v" .DNS.Domain}}\
{{range .DNS.Entries}}\
{{if eq .Tombstone 0}}\
{{$hostname := trimSuffix .Hostname $domain}}\
{{printf "%-12v" $hostname}} {{printf "%-15v" .Address}} {{printf "%12.12v" .ContainerID}} {{.Origin}}
{{end}}\
{{end}}\
`)

var ipamTemplate = defTemplate("ipamTemplate", `{{printIPAMRanges .Router .IPAM}}`)

type WeaveStatus struct {
	Version string
	Router  *weave.NetworkRouterStatus `json:"Router,omitempty"`
	IPAM    *ipam.Status               `json:"IPAM,omitempty"`
	DNS     *nameserver.Status         `json:"DNS,omitempty"`
}

func HandleHTTP(muxRouter *mux.Router, version string, router *weave.NetworkRouter, allocator *ipam.Allocator, defaultSubnet address.CIDR, ns *nameserver.Nameserver, dnsserver *nameserver.DNSServer) {
	status := func() WeaveStatus {
		return WeaveStatus{
			version,
			weave.NewNetworkRouterStatus(router),
			ipam.NewStatus(allocator, defaultSubnet),
			nameserver.NewStatus(ns, dnsserver)}
	}
	muxRouter.Methods("GET").Path("/report").Headers("Accept", "application/json").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			json, err := json.MarshalIndent(status(), "", "    ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				Log.Error("Error during report marshalling: ", err)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			w.Write(json)
		})

	muxRouter.Methods("GET").Path("/report").Queries("format", "{format}").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			formatTemplate, err := template.New("format").Parse(mux.Vars(r)["format"])
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			if err := formatTemplate.Execute(w, status()); err != nil {
				http.Error(w, "error during template execution", http.StatusInternalServerError)
				Log.Error(err)
			}
		})

	defHandler := func(path string, template *template.Template) {
		muxRouter.Methods("GET").Path(path).HandlerFunc(
			func(w http.ResponseWriter, r *http.Request) {
				if err := template.Execute(w, status()); err != nil {
					http.Error(w, "error during template execution", http.StatusInternalServerError)
					Log.Error(err)
				}
			})
	}

	defHandler("/status", statusTemplate)
	defHandler("/status/targets", targetsTemplate)
	defHandler("/status/connections", connectionsTemplate)
	defHandler("/status/peers", peersTemplate)
	defHandler("/status/dns", dnsEntriesTemplate)
	defHandler("/status/ipam", ipamTemplate)

}
<commit_msg>Add a 'json' function for 'weave report -f'<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"text/template"

	"github.com/gorilla/mux"
	. "github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/ipam"
	"github.com/weaveworks/weave/mesh"
	"github.com/weaveworks/weave/nameserver"
	"github.com/weaveworks/weave/net/address"
	weave "github.com/weaveworks/weave/router"
)

var rootTemplate = template.New("root").Funcs(map[string]interface{}{
	"countDNSEntries": func(entries []nameserver.EntryStatus) int {
		count := 0
		for _, entry := range entries {
			if entry.Tombstone == 0 {
				count++
			}
		}
		return count
	},
	"upstreamServers": func(servers []string) string {
		if len(servers) == 0 {
			return "none"
		}
		return strings.Join(servers, ", ")
	},
	"printIPAMRanges": func(router weave.NetworkRouterStatus, status ipam.Status) string {
		var buffer bytes.Buffer

		type stats struct {
			ips       uint32
			nickname  string
			reachable bool
		}

		peerStats := make(map[string]*stats)

		for _, entry := range status.Entries {
			s, found := peerStats[entry.Peer]
			if !found {
				s = &stats{nickname: entry.Nickname, reachable: entry.IsKnownPeer}
				peerStats[entry.Peer] = s
			}
			s.ips += entry.Size
		}

		printOwned := func(name string, nickName string, reachable bool, ips uint32) {
			reachableStr := ""
			if !reachable {
				reachableStr = "- unreachable!"
			}
			percentageRanges := float32(ips) * 100.0 / float32(status.RangeNumIPs)

			displayName := name + "(" + nickName + ")"
			fmt.Fprintf(&buffer, "%-37v %8d IPs (%04.1f%% of total) %s\n",
				displayName, ips, percentageRanges, reachableStr)
		}

		// print the local info first
		if ourStats := peerStats[router.Name]; ourStats != nil {
			printOwned(router.Name, ourStats.nickname, true, ourStats.ips)
		}

		// and then the rest
		for peer, stats := range peerStats {
			if peer != router.Name {
				printOwned(peer, stats.nickname, stats.reachable, stats.ips)
			}
		}

		return buffer.String()
	},
	"allIPAMOwnersUnreachable": func(status ipam.Status) bool {
		for _, entry := range status.Entries {
			if entry.Size > 0 && entry.IsKnownPeer {
				return false
			}
		}
		return true
	},
	"printConnectionCounts": func(conns []mesh.LocalConnectionStatus) string {
		counts := make(map[string]int)
		for _, conn := range conns {
			counts[conn.State]++
		}
		return printCounts(counts, []string{"established", "pending", "retrying", "failed", "connecting"})
	},
	"printPeerConnectionCounts": func(peers []mesh.PeerStatus) string {
		counts := make(map[string]int)
		for _, peer := range peers {
			for _, conn := range peer.Connections {
				if conn.Established {
					counts["established"]++
				} else {
					counts["pending"]++
				}
			}
		}
		return printCounts(counts, []string{"established", "pending"})
	},
	"printState": func(enabled bool) string {
		if enabled {
			return "enabled"
		}
		return "disabled"
	},
	"trimSuffix": strings.TrimSuffix,
})

// Print counts in a specified order
func printCounts(counts map[string]int, keys []string) string {
	var stringCounts []string
	for _, key := range keys {
		if count, ok := counts[key]; ok {
			stringCounts = append(stringCounts, fmt.Sprintf("%d %s", count, key))
		}
	}
	return strings.Join(stringCounts, ", ")
}

// Strip escaped newlines from template
func escape(template string) string {
	return strings.Replace(template, "\\\n", "", -1)
}

// Define a named template panicking on error
func defTemplate(name string, text string) *template.Template {
	return template.Must(rootTemplate.New(name).Parse(escape(text)))
}

var statusTemplate = defTemplate("status", `\
       Version: {{.Version}}

       Service: router
      Protocol: {{.Router.Protocol}} \
{{if eq .Router.ProtocolMinVersion .Router.ProtocolMaxVersion}}\
{{.Router.ProtocolMaxVersion}}\
{{else}}\
{{.Router.ProtocolMinVersion}}..{{.Router.ProtocolMaxVersion}}\
{{end}}
          Name: {{.Router.Name}}({{.Router.NickName}})
    Encryption: {{printState .Router.Encryption}}
 PeerDiscovery: {{printState .Router.PeerDiscovery}}
       Targets: {{len .Router.Targets}}
   Connections: {{len .Router.Connections}}{{with printConnectionCounts .Router.Connections}} ({{.}}){{end}}
         Peers: {{len .Router.Peers}}{{with printPeerConnectionCounts .Router.Peers}} (with {{.}} connections){{end}}
{{if .IPAM}}\

       Service: ipam
{{if .IPAM.Entries}}\
{{if allIPAMOwnersUnreachable .IPAM}}\
        Status: all IP ranges owned by unreachable peers - use 'rmpeer' if they are dead
{{else if len .IPAM.PendingAllocates}}\
        Status: waiting for IP range grant from peers
{{else}}\
        Status: ready
{{end}}\
{{else if .IPAM.Paxos}}\
        Status: awaiting consensus (quorum: {{.IPAM.Paxos.Quorum}}, known: {{.IPAM.Paxos.KnownNodes}})
{{else}}\
        Status: idle
{{end}}\
         Range: {{.IPAM.Range}}
 DefaultSubnet: {{.IPAM.DefaultSubnet}}
{{end}}\
{{if .DNS}}\

       Service: dns
        Domain: {{.DNS.Domain}}
      Upstream: {{upstreamServers .DNS.Upstream}}
           TTL: {{.DNS.TTL}}
       Entries: {{countDNSEntries .DNS.Entries}}
{{end}}\
`)

var targetsTemplate = defTemplate("targetsTemplate", `\
{{range .Router.Targets}}{{.}}
{{end}}\
`)

var connectionsTemplate = defTemplate("connectionsTemplate", `\
{{range .Router.Connections}}\
{{if .Outbound}}->{{else}}<-{{end}} {{printf "%-21v" .Address}} {{printf "%-11v" .State}} {{.Info}}
{{end}}\
`)

var peersTemplate = defTemplate("peers", `\
{{range .Router.Peers}}\
{{.Name}}({{.NickName}})
{{range .Connections}}\
   {{if .Outbound}}->{{else}}<-{{end}} {{printf "%-21v" .Address}} \
{{$nameNickName := printf "%v(%v)" .Name .NickName}}{{printf "%-37v" $nameNickName}} \
{{if .Established}}established{{else}}pending{{end}}
{{end}}\
{{end}}\
`)

var dnsEntriesTemplate = defTemplate("dnsEntries", `\
{{$domain := printf ".%v" .DNS.Domain}}\
{{range .DNS.Entries}}\
{{if eq .Tombstone 0}}\
{{$hostname := trimSuffix .Hostname $domain}}\
{{printf "%-12v" $hostname}} {{printf "%-15v" .Address}} {{printf "%12.12v" .ContainerID}} {{.Origin}}
{{end}}\
{{end}}\
`)

var ipamTemplate = defTemplate("ipamTemplate", `{{printIPAMRanges .Router .IPAM}}`)

type WeaveStatus struct {
	Version string
	Router  *weave.NetworkRouterStatus `json:"Router,omitempty"`
	IPAM    *ipam.Status               `json:"IPAM,omitempty"`
	DNS     *nameserver.Status         `json:"DNS,omitempty"`
}

func HandleHTTP(muxRouter *mux.Router, version string, router *weave.NetworkRouter, allocator *ipam.Allocator, defaultSubnet address.CIDR, ns *nameserver.Nameserver, dnsserver *nameserver.DNSServer) {
	status := func() WeaveStatus {
		return WeaveStatus{
			version,
			weave.NewNetworkRouterStatus(router),
			ipam.NewStatus(allocator, defaultSubnet),
			nameserver.NewStatus(ns, dnsserver)}
	}
	muxRouter.Methods("GET").Path("/report").Headers("Accept", "application/json").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			json, err := json.MarshalIndent(status(), "", "    ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				Log.Error("Error during report marshalling: ", err)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			w.Write(json)
		})

	muxRouter.Methods("GET").Path("/report").Queries("format", "{format}").HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			funcs := template.FuncMap{
				"json": func(v interface{}) string {
					a, _ := json.Marshal(v)
					return string(a)
				},
			}
			formatTemplate, err := template.New("format").Funcs(funcs).Parse(mux.Vars(r)["format"])
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			if err := formatTemplate.Execute(w, status()); err != nil {
				http.Error(w, "error during template execution", http.StatusInternalServerError)
				Log.Error(err)
			}
		})

	defHandler := func(path string, template *template.Template) {
		muxRouter.Methods("GET").Path(path).HandlerFunc(
			func(w http.ResponseWriter, r *http.Request) {
				if err := template.Execute(w, status()); err != nil {
					http.Error(w, "error during template execution", http.StatusInternalServerError)
					Log.Error(err)
				}
			})
	}

	defHandler("/status", statusTemplate)
	defHandler("/status/targets", targetsTemplate)
	defHandler("/status/connections", connectionsTemplate)
	defHandler("/status/peers", peersTemplate)
	defHandler("/status/dns", dnsEntriesTemplate)
	defHandler("/status/ipam", ipamTemplate)

}
<|endoftext|>


<commit_before>// Copyright 2015 Matthew Holt and The Caddy 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 headers

import (
	"net/http"
	"strings"

	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)

func init() {
	httpcaddyfile.RegisterHandlerDirective("header", parseCaddyfile)
	httpcaddyfile.RegisterHandlerDirective("request_header", parseReqHdrCaddyfile)
}

// parseCaddyfile sets up the handler for response headers from
// Caddyfile tokens. Syntax:
//
//     header [<matcher>] [[+|-]<field> [<value|regexp>] [<replacement>]] {
//         [+]<field> [<value|regexp> [<replacement>]]
//         -<field>
//         [defer]
//     }
//
// Either a block can be opened or a single header field can be configured
// in the first line, but not both in the same directive. Header operations
// are deferred to write-time if any headers are being deleted or if the
// 'defer' subdirective is used.
func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	hdr := new(Handler)

	makeResponseOps := func() {
		if hdr.Response == nil {
			hdr.Response = &RespHeaderOps{
				HeaderOps: new(HeaderOps),
			}
		}
	}

	for h.Next() {
		// first see if headers are in the initial line
		var hasArgs bool
		if h.NextArg() {
			hasArgs = true
			field := h.Val()
			var value, replacement string
			if h.NextArg() {
				value = h.Val()
			}
			if h.NextArg() {
				replacement = h.Val()
			}
			makeResponseOps()
			CaddyfileHeaderOp(hdr.Response.HeaderOps, field, value, replacement)
			if len(hdr.Response.HeaderOps.Delete) > 0 {
				hdr.Response.Deferred = true
			}
		}

		// if not, they should be in a block
		for h.NextBlock(0) {
			field := h.Val()
			if field == "defer" {
				hdr.Response.Deferred = true
				continue
			}
			if hasArgs {
				return nil, h.Err("cannot specify headers in both arguments and block")
			}
			var value, replacement string
			if h.NextArg() {
				value = h.Val()
			}
			if h.NextArg() {
				replacement = h.Val()
			}
			makeResponseOps()
			CaddyfileHeaderOp(hdr.Response.HeaderOps, field, value, replacement)
			if len(hdr.Response.HeaderOps.Delete) > 0 {
				hdr.Response.Deferred = true
			}
		}
	}

	return hdr, nil
}

// parseReqHdrCaddyfile sets up the handler for request headers
// from Caddyfile tokens. Syntax:
//
//     request_header [<matcher>] [[+|-]<field> [<value|regexp>] [<replacement>]]
//
func parseReqHdrCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	hdr := new(Handler)
	for h.Next() {
		if !h.NextArg() {
			return nil, h.ArgErr()
		}
		field := h.Val()
		var value, replacement string
		if h.NextArg() {
			value = h.Val()
		}
		if h.NextArg() {
			replacement = h.Val()
			if h.NextArg() {
				return nil, h.ArgErr()
			}
		}

		if hdr.Request == nil {
			hdr.Request = new(HeaderOps)
		}
		CaddyfileHeaderOp(hdr.Request, field, value, replacement)

		if h.NextArg() {
			return nil, h.ArgErr()
		}
	}
	return hdr, nil
}

// CaddyfileHeaderOp applies a new header operation according to
// field, value, and replacement. The field can be prefixed with
// "+" or "-" to specify adding or removing; otherwise, the value
// will be set (overriding any previous value). If replacement is
// non-empty, value will be treated as a regular expression which
// will be used to search and then replacement will be used to
// complete the substring replacement; in that case, any + or -
// prefix to field will be ignored.
func CaddyfileHeaderOp(ops *HeaderOps, field, value, replacement string) {
	if strings.HasPrefix(field, "+") {
		if ops.Add == nil {
			ops.Add = make(http.Header)
		}
		ops.Add.Set(field[1:], value)
	} else if strings.HasPrefix(field, "-") {
		ops.Delete = append(ops.Delete, field[1:])
	} else {
		if replacement == "" {
			if ops.Set == nil {
				ops.Set = make(http.Header)
			}
			ops.Set.Set(field, value)
		} else {
			if ops.Replace == nil {
				ops.Replace = make(map[string][]Replacement)
			}
			field = strings.TrimLeft(field, "+-")
			ops.Replace[field] = append(
				ops.Replace[field],
				Replacement{
					SearchRegexp: value,
					Replace:      replacement,
				},
			)
		}
	}
}
<commit_msg>headers: Trim any trailing colon from field names as a courtesy<commit_after>// Copyright 2015 Matthew Holt and The Caddy 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 headers

import (
	"net/http"
	"strings"

	"github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
	"github.com/caddyserver/caddy/v2/modules/caddyhttp"
)

func init() {
	httpcaddyfile.RegisterHandlerDirective("header", parseCaddyfile)
	httpcaddyfile.RegisterHandlerDirective("request_header", parseReqHdrCaddyfile)
}

// parseCaddyfile sets up the handler for response headers from
// Caddyfile tokens. Syntax:
//
//     header [<matcher>] [[+|-]<field> [<value|regexp>] [<replacement>]] {
//         [+]<field> [<value|regexp> [<replacement>]]
//         -<field>
//         [defer]
//     }
//
// Either a block can be opened or a single header field can be configured
// in the first line, but not both in the same directive. Header operations
// are deferred to write-time if any headers are being deleted or if the
// 'defer' subdirective is used.
func parseCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	hdr := new(Handler)

	makeResponseOps := func() {
		if hdr.Response == nil {
			hdr.Response = &RespHeaderOps{
				HeaderOps: new(HeaderOps),
			}
		}
	}

	for h.Next() {
		// first see if headers are in the initial line
		var hasArgs bool
		if h.NextArg() {
			hasArgs = true
			field := h.Val()
			var value, replacement string
			if h.NextArg() {
				value = h.Val()
			}
			if h.NextArg() {
				replacement = h.Val()
			}
			makeResponseOps()
			CaddyfileHeaderOp(hdr.Response.HeaderOps, field, value, replacement)
			if len(hdr.Response.HeaderOps.Delete) > 0 {
				hdr.Response.Deferred = true
			}
		}

		// if not, they should be in a block
		for h.NextBlock(0) {
			field := h.Val()
			if field == "defer" {
				hdr.Response.Deferred = true
				continue
			}
			if hasArgs {
				return nil, h.Err("cannot specify headers in both arguments and block")
			}
			var value, replacement string
			if h.NextArg() {
				value = h.Val()
			}
			if h.NextArg() {
				replacement = h.Val()
			}
			makeResponseOps()
			CaddyfileHeaderOp(hdr.Response.HeaderOps, field, value, replacement)
			if len(hdr.Response.HeaderOps.Delete) > 0 {
				hdr.Response.Deferred = true
			}
		}
	}

	return hdr, nil
}

// parseReqHdrCaddyfile sets up the handler for request headers
// from Caddyfile tokens. Syntax:
//
//     request_header [<matcher>] [[+|-]<field> [<value|regexp>] [<replacement>]]
//
func parseReqHdrCaddyfile(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
	hdr := new(Handler)
	for h.Next() {
		if !h.NextArg() {
			return nil, h.ArgErr()
		}
		field := h.Val()

		// sometimes it is habitual for users to suffix a field name with a colon,
		// as if they were writing a curl command or something; see
		// https://caddy.community/t/v2-reverse-proxy-please-add-cors-example-to-the-docs/7349
		field = strings.TrimSuffix(field, ":")

		var value, replacement string
		if h.NextArg() {
			value = h.Val()
		}
		if h.NextArg() {
			replacement = h.Val()
			if h.NextArg() {
				return nil, h.ArgErr()
			}
		}

		if hdr.Request == nil {
			hdr.Request = new(HeaderOps)
		}
		CaddyfileHeaderOp(hdr.Request, field, value, replacement)

		if h.NextArg() {
			return nil, h.ArgErr()
		}
	}
	return hdr, nil
}

// CaddyfileHeaderOp applies a new header operation according to
// field, value, and replacement. The field can be prefixed with
// "+" or "-" to specify adding or removing; otherwise, the value
// will be set (overriding any previous value). If replacement is
// non-empty, value will be treated as a regular expression which
// will be used to search and then replacement will be used to
// complete the substring replacement; in that case, any + or -
// prefix to field will be ignored.
func CaddyfileHeaderOp(ops *HeaderOps, field, value, replacement string) {
	if strings.HasPrefix(field, "+") {
		if ops.Add == nil {
			ops.Add = make(http.Header)
		}
		ops.Add.Set(field[1:], value)
	} else if strings.HasPrefix(field, "-") {
		ops.Delete = append(ops.Delete, field[1:])
	} else {
		if replacement == "" {
			if ops.Set == nil {
				ops.Set = make(http.Header)
			}
			ops.Set.Set(field, value)
		} else {
			if ops.Replace == nil {
				ops.Replace = make(map[string][]Replacement)
			}
			field = strings.TrimLeft(field, "+-")
			ops.Replace[field] = append(
				ops.Replace[field],
				Replacement{
					SearchRegexp: value,
					Replace:      replacement,
				},
			)
		}
	}
}
<|endoftext|>


<commit_before>package parse

import (
	"fmt"
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/elves/elvish/errutil"
)

type fs map[string]interface{}
type ast struct {
	name   string
	fields fs
}

func a(c ...interface{}) ast {
	return ast{"Chunk/Pipeline/Form", fs{"Head": "a", "Args": c}}
}

var goodCases = []struct {
	src string
	ast ast
}{
	// Chunk
	{"a;b|c\n;d", ast{"Chunk", fs{"Pipelines": []string{"a", "b|c", "d"}}}},
	// Empty chunk
	{"", ast{"Chunk", nil}},
	// Lots of unnecessary whitespaces
	{"  ;\n\n  ls \t ;\n", ast{"Chunk", fs{"Pipelines": []string{"ls \t "}}}},

	// Form
	{"ls x y", ast{"Chunk/Pipeline/Form", fs{
		"Head": "ls",
		"Args": []string{"x", "y"}}}},
	// Redirections
	{"a>b", ast{"Chunk/Pipeline/Form", fs{
		"Head": "a",
		"Redirs": []ast{
			ast{"Redir", fs{"Mode": Write, "Source": "b"}}},
	}}},
	// More redirections
	{"a>>b 2>b 3>&- 4>&1", ast{"Chunk/Pipeline/Form", fs{
		"Head": "a",
		"Redirs": []ast{
			ast{"Redir", fs{"Mode": Append, "Source": "b"}},
			ast{"Redir", fs{"Mode": Write, "Source": "b"}},
			ast{"Redir", fs{"Mode": Write, "SourceIsFd": true, "Source": "-"}},
			ast{"Redir", fs{"Mode": Write, "SourceIsFd": true, "Source": "1"}},
		},
	}}},
	// Exitus redirection
	{"a ?>$e", ast{"Chunk/Pipeline/Form", fs{
		"Head":        "a",
		"ExitusRedir": ast{"ExitusRedir", fs{"Dest": "$e"}},
	}}},
	// Assignments:
	{"k=v k[a][b]=v a", ast{"Chunk/Pipeline/Form", fs{
		"Assignments": []string{"k=v", "k[a][b]=v"},
		"Head":        "a"}}},

	// Compound
	{`a b"foo"$c'xyz'`, a(ast{"Compound", fs{
		"Indexings": []string{"b", `"foo"`, "$c", "'xyz'"}}})},

	// Indexing
	{"a $b[c][d][e]", a(ast{"Compound/Indexing", fs{
		"Head": "$b", "Indicies": []string{"c", "d", "e"},
	}})},

	// Primary
	//
	// Single quote
	{"a 'b'", a(ast{"Compound/Indexing/Primary", fs{
		"text": "'b'", "Type": SingleQuoted,
	}})},
	// Double quote
	{`a "b"`, a(ast{"Compound/Indexing/Primary", fs{
		"text": `"b"`, "Type": DoubleQuoted,
	}})},
	// List
	{"a [] [ ] [1] [ 2] [3 ] [ 4 5 6 7 ]", a(
		ast{"Compound/Indexing/Primary", fs{"Type": List}},
		ast{"Compound/Indexing/Primary", fs{"Type": List}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"1"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"2"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"3"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{
				"Compounds": []string{"4", "5", "6", "7"}}}}},
	)},
	// Map
	{"a [&k v] [ &k v] [&k v ] [ &k v ] [&a b &c d &e f]", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": Map,
			"MapPairs": []ast{
				ast{"MapPair", fs{"Key": "a", "Value": "b"}},
				ast{"MapPair", fs{"Key": "c", "Value": "d"}},
				ast{"MapPair", fs{"Key": "e", "Value": "f"}},
			}}},
	)},
	// Empty map
	{"a [&] [ &] [& ] [ & ]", a(
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[&]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[ &]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[& ]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[ & ]"}},
	)},
	// Lambda
	{"a []{} [ ]{ } []{ echo 233 } [ $x $y ]{puts $x $y} { put $1}", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type": Lambda,
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": Lambda,
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"Chunk": " echo 233 ",
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"List":  "$x $y ",
			"Chunk": "puts $x $y",
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"Chunk": " put $1",
		}},
	)},
	// Output capture
	{"a () (b;c)", a(
		ast{"Compound/Indexing/Primary", fs{"Type": OutputCapture}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": OutputCapture, "Chunk": "b;c",
		}})},
	// Output capture with backquotes
	{"a `` `b;c` `e>f`", a("``", "`b;c`", "`e>f`")},
	// Backquotes may be nested with unclosed parens and braces
	{"a `a (b `c`)` `d [`e`]`", a("`a (b `c`)`", "`d [`e`]`")},
	// Exitus capture
	{"a ?() ?(b;c)", a(
		ast{"Compound/Indexing/Primary", fs{"Type": ErrorCapture}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": ErrorCapture, "Chunk": "b;c",
		}})},
	// Braced
	{"a {a,c-f}", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type":    Braced,
			"Braced":  []string{"a", "c", "f"},
			"IsRange": []bool{false, true}}})},
	// Tilde
	{"a ~xiaq/go", a(
		ast{"Compound", fs{
			"Indexings": []ast{
				{"Indexing/Primary", fs{"Type": Tilde}},
				{"Indexing/Primary", fs{"Type": Bareword, "Value": "xiaq/go"}},
			},
		}},
	)},
}

func TestParse(t *testing.T) {
	for _, tc := range goodCases {
		bn, err := Parse(tc.src)
		if err != nil {
			t.Errorf("Parse(%q) returns error: %v", tc.src, err)
		}
		err = checkParseTree(bn)
		if err != nil {
			t.Errorf("Parse(%q) returns bad parse tree: %v", tc.src, err)
			fmt.Fprintf(os.Stderr, "Parse tree of %q:\n", tc.src)
			PprintParseTree(bn, os.Stderr)
		}
		err = checkAST(bn, tc.ast)
		if err != nil {
			t.Errorf("Parse(%q) returns bad AST: %v", tc.src, err)
			fmt.Fprintf(os.Stderr, "AST of %q:\n", tc.src)
			PprintAST(bn, os.Stderr)
		}
	}
}

// checkParseTree checks whether the parse tree part of a Node is well-formed.
func checkParseTree(n Node) error {
	children := n.Children()
	if len(children) == 0 {
		return nil
	}

	// Parent pointers of all children should point to me.
	for i, ch := range children {
		if ch.Parent() != n {
			return fmt.Errorf("parent of child %d (%s) is wrong: %s", i, summary(ch), summary(n))
		}
	}

	// The Begin of the first child should be equal to mine.
	if children[0].Begin() != n.Begin() {
		return fmt.Errorf("gap between node and first child: %s", summary(n))
	}
	// The End of the last child should be equal to mine.
	nch := len(children)
	if children[nch-1].End() != n.End() {
		return fmt.Errorf("gap between node and last child: %s", summary(n))
	}
	// Consecutive children have consecutive position ranges.
	for i := 0; i < nch-1; i++ {
		if children[i].End() != children[i+1].Begin() {
			return fmt.Errorf("gap beteen child %d and %d of: %s", i, i+1, summary(n))
		}
	}

	// Check children recursively.
	for _, ch := range n.Children() {
		err := checkParseTree(ch)
		if err != nil {
			return err
		}
	}
	return nil
}

// checkAST checks whether the AST part of a Node matches a specification.
func checkAST(n Node, want ast) error {
	// TODO: Check fields present in struct but not in ast
	wantnames := strings.Split(want.name, "/")
	// Check coalesced levels
	for i, wantname := range wantnames {
		name := reflect.TypeOf(n).Elem().Name()
		if wantname != name {
			return fmt.Errorf("want %s, got %s (%s)", wantname, name, summary(n))
		}
		if i == len(wantnames)-1 {
			break
		}
		fields := n.Children()
		if len(fields) != 1 {
			return fmt.Errorf("want exactly 1 child, got %d (%s)", len(fields), summary(n))
		}
		n = fields[0]
	}

	if want.fields == nil && len(n.Children()) != 0 {
		return fmt.Errorf("want leaf, got inner node (%s)", summary(n))
	}
	nv := reflect.ValueOf(n).Elem()

	for fieldname, wantfield := range want.fields {
		if fieldname == "text" {
			if n.SourceText() != wantfield.(string) {
				return fmt.Errorf("want %q, got %q (%s)", wantfield, n.SourceText())
			}
		} else {
			fv := nv.FieldByName(fieldname)
			err := checkAny(fv.Interface(), wantfield, summary(n))
			if err != nil {
				return err
			}
		}
	}
	return nil
}

var nodeType = reflect.TypeOf((*Node)(nil)).Elem()

func checkAny(got interface{}, want interface{}, ctx string) error {
	if got, ok := got.(Node); ok {
		// A Node
		return checkNode(got.(Node), want)
	}
	tgot := reflect.TypeOf(got)
	if tgot.Kind() == reflect.Slice && tgot.Elem().Implements(nodeType) {
		// A slice of Nodes
		vgot := reflect.ValueOf(got)
		vwant := reflect.ValueOf(want)
		if vgot.Len() != vwant.Len() {
			return fmt.Errorf("want %d, got %d (%s)", vwant.Len(), vgot.Len(), ctx)
		}
		for i := 0; i < vgot.Len(); i++ {
			err := checkNode(vgot.Index(i).Interface().(Node),
				vwant.Index(i).Interface())
			if err != nil {
				return err
			}
		}
		return nil
	}

	if !reflect.DeepEqual(want, got) {
		return fmt.Errorf("want %v, got %v (%s)", want, got, ctx)
	}
	return nil
}

func checkNode(got Node, want interface{}) error {
	switch want := want.(type) {
	case string:
		text := got.SourceText()
		if want != text {
			return fmt.Errorf("want %q, got %q (%s)", want, text, summary(got))
		}
		return nil
	case ast:
		return checkAST(got, want)
	default:
		panic(fmt.Sprintf("bad want type %T (%s)", want, summary(got)))
	}
}

var badCases = []struct {
	src string
	pos int // expected Begin position of first error
}{
	// Unopened parens.
	{")", 0}, {"]", 0}, {"}", 0},
	// Unclosed parens.
	{"a (", 3}, {"a [", 3}, {"a {", 3},
	// Ampersand
	{"a &", 3}, {"a [&", 4},
}

func TestParseError(t *testing.T) {
	for _, tc := range badCases {
		_, err := Parse(tc.src)
		if err == nil {
			t.Errorf("Parse(%q) returns no error", tc.src)
			continue
		}
		posErr0 := err.(*errutil.Errors).Errors[0].(*errutil.PosError)
		if posErr0.Begin != tc.pos {
			t.Errorf("Parse(%q) first error begins at %d, want %d. Errors are:%s\n", tc.src, posErr0.Begin, tc.pos, err)
		}
	}
}
<commit_msg>Add a test case against parse.Parse.<commit_after>package parse

import (
	"fmt"
	"os"
	"reflect"
	"strings"
	"testing"

	"github.com/elves/elvish/errutil"
)

type fs map[string]interface{}
type ast struct {
	name   string
	fields fs
}

func a(c ...interface{}) ast {
	return ast{"Chunk/Pipeline/Form", fs{"Head": "a", "Args": c}}
}

var goodCases = []struct {
	src string
	ast ast
}{
	// Chunk
	{"a;b|c\n;d", ast{"Chunk", fs{"Pipelines": []string{"a", "b|c", "d"}}}},
	// Empty chunk
	{"", ast{"Chunk", nil}},
	// Lots of unnecessary whitespaces
	{"  ;\n\n  ls \t ;\n", ast{"Chunk", fs{"Pipelines": []string{"ls \t "}}}},

	// Form
	{"ls x y", ast{"Chunk/Pipeline/Form", fs{
		"Head": "ls",
		"Args": []string{"x", "y"}}}},
	// Redirections
	{"a>b", ast{"Chunk/Pipeline/Form", fs{
		"Head": "a",
		"Redirs": []ast{
			ast{"Redir", fs{"Mode": Write, "Source": "b"}}},
	}}},
	// More redirections
	{"a>>b 2>b 3>&- 4>&1", ast{"Chunk/Pipeline/Form", fs{
		"Head": "a",
		"Redirs": []ast{
			ast{"Redir", fs{"Mode": Append, "Source": "b"}},
			ast{"Redir", fs{"Mode": Write, "Source": "b"}},
			ast{"Redir", fs{"Mode": Write, "SourceIsFd": true, "Source": "-"}},
			ast{"Redir", fs{"Mode": Write, "SourceIsFd": true, "Source": "1"}},
		},
	}}},
	// Exitus redirection
	{"a ?>$e", ast{"Chunk/Pipeline/Form", fs{
		"Head":        "a",
		"ExitusRedir": ast{"ExitusRedir", fs{"Dest": "$e"}},
	}}},
	// Assignments:
	{"k=v k[a][b]=v a", ast{"Chunk/Pipeline/Form", fs{
		"Assignments": []string{"k=v", "k[a][b]=v"},
		"Head":        "a"}}},

	// Compound
	{`a b"foo"$c'xyz'`, a(ast{"Compound", fs{
		"Indexings": []string{"b", `"foo"`, "$c", "'xyz'"}}})},

	// Indexing
	{"a $b[c][d][e]", a(ast{"Compound/Indexing", fs{
		"Head": "$b", "Indicies": []string{"c", "d", "e"},
	}})},

	// Primary
	//
	// Single quote
	{"a 'b'", a(ast{"Compound/Indexing/Primary", fs{
		"text": "'b'", "Type": SingleQuoted,
	}})},
	// Double quote
	{`a "b"`, a(ast{"Compound/Indexing/Primary", fs{
		"text": `"b"`, "Type": DoubleQuoted,
	}})},
	// List
	{"a [] [ ] [1] [ 2] [3 ] [ 4 5 6 7 ]", a(
		ast{"Compound/Indexing/Primary", fs{"Type": List}},
		ast{"Compound/Indexing/Primary", fs{"Type": List}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"1"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"2"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{"Compounds": []string{"3"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": List,
			"List": ast{"Array", fs{
				"Compounds": []string{"4", "5", "6", "7"}}}}},
	)},
	// Map
	{"a [&k v] [ &k v] [&k v ] [ &k v ] [&a b &c d &e f]", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":     Map,
			"MapPairs": []ast{ast{"MapPair", fs{"Key": "k", "Value": "v"}}}}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": Map,
			"MapPairs": []ast{
				ast{"MapPair", fs{"Key": "a", "Value": "b"}},
				ast{"MapPair", fs{"Key": "c", "Value": "d"}},
				ast{"MapPair", fs{"Key": "e", "Value": "f"}},
			}}},
	)},
	// Empty map
	{"a [&] [ &] [& ] [ & ]", a(
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[&]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[ &]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[& ]"}},
		ast{"Compound/Indexing/Primary", fs{"Type": Map, "text": "[ & ]"}},
	)},
	// Lambda
	{"a []{} [ ]{ } []{ echo 233 } [ $x $y ]{puts $x $y} { put $1}", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type": Lambda,
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": Lambda,
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"Chunk": " echo 233 ",
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"List":  "$x $y ",
			"Chunk": "puts $x $y",
		}},
		ast{"Compound/Indexing/Primary", fs{
			"Type":  Lambda,
			"Chunk": " put $1",
		}},
	)},
	// Output capture
	{"a () (b;c)", a(
		ast{"Compound/Indexing/Primary", fs{"Type": OutputCapture}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": OutputCapture, "Chunk": "b;c",
		}})},
	// Output capture with backquotes
	{"a `` `b;c` `e>f`", a("``", "`b;c`", "`e>f`")},
	// Backquotes may be nested with unclosed parens and braces
	{"a `a (b `c`)` `d [`e`]`", a("`a (b `c`)`", "`d [`e`]`")},
	// Exitus capture
	{"a ?() ?(b;c)", a(
		ast{"Compound/Indexing/Primary", fs{"Type": ErrorCapture}},
		ast{"Compound/Indexing/Primary", fs{
			"Type": ErrorCapture, "Chunk": "b;c",
		}})},
	// Braced
	{"a {a,c-f}", a(
		ast{"Compound/Indexing/Primary", fs{
			"Type":    Braced,
			"Braced":  []string{"a", "c", "f"},
			"IsRange": []bool{false, true}}})},
	// Tilde
	{"a ~xiaq/go", a(
		ast{"Compound", fs{
			"Indexings": []ast{
				{"Indexing/Primary", fs{"Type": Tilde}},
				{"Indexing/Primary", fs{"Type": Bareword, "Value": "xiaq/go"}},
			},
		}},
	)},
}

func TestParse(t *testing.T) {
	for _, tc := range goodCases {
		bn, err := Parse(tc.src)
		if err != nil {
			t.Errorf("Parse(%q) returns error: %v", tc.src, err)
		}
		err = checkParseTree(bn)
		if err != nil {
			t.Errorf("Parse(%q) returns bad parse tree: %v", tc.src, err)
			fmt.Fprintf(os.Stderr, "Parse tree of %q:\n", tc.src)
			PprintParseTree(bn, os.Stderr)
		}
		err = checkAST(bn, tc.ast)
		if err != nil {
			t.Errorf("Parse(%q) returns bad AST: %v", tc.src, err)
			fmt.Fprintf(os.Stderr, "AST of %q:\n", tc.src)
			PprintAST(bn, os.Stderr)
		}
	}
}

// checkParseTree checks whether the parse tree part of a Node is well-formed.
func checkParseTree(n Node) error {
	children := n.Children()
	if len(children) == 0 {
		return nil
	}

	// Parent pointers of all children should point to me.
	for i, ch := range children {
		if ch.Parent() != n {
			return fmt.Errorf("parent of child %d (%s) is wrong: %s", i, summary(ch), summary(n))
		}
	}

	// The Begin of the first child should be equal to mine.
	if children[0].Begin() != n.Begin() {
		return fmt.Errorf("gap between node and first child: %s", summary(n))
	}
	// The End of the last child should be equal to mine.
	nch := len(children)
	if children[nch-1].End() != n.End() {
		return fmt.Errorf("gap between node and last child: %s", summary(n))
	}
	// Consecutive children have consecutive position ranges.
	for i := 0; i < nch-1; i++ {
		if children[i].End() != children[i+1].Begin() {
			return fmt.Errorf("gap beteen child %d and %d of: %s", i, i+1, summary(n))
		}
	}

	// Check children recursively.
	for _, ch := range n.Children() {
		err := checkParseTree(ch)
		if err != nil {
			return err
		}
	}
	return nil
}

// checkAST checks whether the AST part of a Node matches a specification.
func checkAST(n Node, want ast) error {
	// TODO: Check fields present in struct but not in ast
	wantnames := strings.Split(want.name, "/")
	// Check coalesced levels
	for i, wantname := range wantnames {
		name := reflect.TypeOf(n).Elem().Name()
		if wantname != name {
			return fmt.Errorf("want %s, got %s (%s)", wantname, name, summary(n))
		}
		if i == len(wantnames)-1 {
			break
		}
		fields := n.Children()
		if len(fields) != 1 {
			return fmt.Errorf("want exactly 1 child, got %d (%s)", len(fields), summary(n))
		}
		n = fields[0]
	}

	if want.fields == nil && len(n.Children()) != 0 {
		return fmt.Errorf("want leaf, got inner node (%s)", summary(n))
	}
	nv := reflect.ValueOf(n).Elem()

	for fieldname, wantfield := range want.fields {
		if fieldname == "text" {
			if n.SourceText() != wantfield.(string) {
				return fmt.Errorf("want %q, got %q (%s)", wantfield, n.SourceText())
			}
		} else {
			fv := nv.FieldByName(fieldname)
			err := checkAny(fv.Interface(), wantfield, summary(n))
			if err != nil {
				return err
			}
		}
	}
	return nil
}

var nodeType = reflect.TypeOf((*Node)(nil)).Elem()

func checkAny(got interface{}, want interface{}, ctx string) error {
	if got, ok := got.(Node); ok {
		// A Node
		return checkNode(got.(Node), want)
	}
	tgot := reflect.TypeOf(got)
	if tgot.Kind() == reflect.Slice && tgot.Elem().Implements(nodeType) {
		// A slice of Nodes
		vgot := reflect.ValueOf(got)
		vwant := reflect.ValueOf(want)
		if vgot.Len() != vwant.Len() {
			return fmt.Errorf("want %d, got %d (%s)", vwant.Len(), vgot.Len(), ctx)
		}
		for i := 0; i < vgot.Len(); i++ {
			err := checkNode(vgot.Index(i).Interface().(Node),
				vwant.Index(i).Interface())
			if err != nil {
				return err
			}
		}
		return nil
	}

	if !reflect.DeepEqual(want, got) {
		return fmt.Errorf("want %v, got %v (%s)", want, got, ctx)
	}
	return nil
}

func checkNode(got Node, want interface{}) error {
	switch want := want.(type) {
	case string:
		text := got.SourceText()
		if want != text {
			return fmt.Errorf("want %q, got %q (%s)", want, text, summary(got))
		}
		return nil
	case ast:
		return checkAST(got, want)
	default:
		panic(fmt.Sprintf("bad want type %T (%s)", want, summary(got)))
	}
}

var badCases = []struct {
	src string
	pos int // expected Begin position of first error
}{
	// Empty form.
	{"a|", 2},
	// Unopened parens.
	{")", 0}, {"]", 0}, {"}", 0},
	// Unclosed parens.
	{"a (", 3}, {"a [", 3}, {"a {", 3},
	// Ampersand
	{"a &", 3}, {"a [&", 4},
}

func TestParseError(t *testing.T) {
	for _, tc := range badCases {
		_, err := Parse(tc.src)
		if err == nil {
			t.Errorf("Parse(%q) returns no error", tc.src)
			continue
		}
		posErr0 := err.(*errutil.Errors).Errors[0].(*errutil.PosError)
		if posErr0.Begin != tc.pos {
			t.Errorf("Parse(%q) first error begins at %d, want %d. Errors are:%s\n", tc.src, posErr0.Begin, tc.pos, err)
		}
	}
}
<|endoftext|>


<commit_before>package aggregator

import (
	"github.com/CapillarySoftware/gostat/stat"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"time"
)

var _ = Describe("Aggregator", func() {

	Describe("Aggregate", func() {
		It("should return all 0 values if a nil slice is received", func() {
			a := Aggregate(nil)
			Expect(a).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return all 0 values if an empty slice is received", func() {
			a := Aggregate([]*stat.Stat{})
			Expect(a).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return all the same values if the slice contains just one Stat", func() {
			const value = 123.456
			stats := []*stat.Stat{{"foo", time.Now().UTC(), value}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: value, Min: value, Max: value, Count: 1}))
		})

		It("should return the expected values for a collection of more than one Stat", func() {
			stats := []*stat.Stat{
				{"foo", time.Now().UTC(), 1},
				{"foo", time.Now().UTC(), 2},
				{"foo", time.Now().UTC(), 3},
				{"foo", time.Now().UTC(), 4},
				{"foo", time.Now().UTC(), 5},
				{"foo", time.Now().UTC(), 6}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}))
		})

		It("should ignore stat names", func() {
			stats := []*stat.Stat{
				{"each",   time.Now().UTC(), 6},
				{"stat",   time.Now().UTC(), 5},
				{"name",   time.Now().UTC(), 4},
				{"here",   time.Now().UTC(), 3},
				{"is",     time.Now().UTC(), 2},
				{"unique", time.Now().UTC(), 1}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}))
		})

		It("should correctly handle negative values", func() {
			stats := []*stat.Stat{
				{"foo", time.Now().UTC(), -1},
				{"foo", time.Now().UTC(), -2},
				{"foo", time.Now().UTC(), -3},
				{"foo", time.Now().UTC(), -4},
				{"foo", time.Now().UTC(), -5},
				{"foo", time.Now().UTC(), -6}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: -3.5, Min: -6, Max: -1, Count: 6}))
		})
	})

	Describe("AppendStatsAggregate", func() {
		It("should return all 0 values if both parameters are zeroed-out StatsAggregates", func() {
			appended := AppendStatsAggregate(StatsAggregate{}, StatsAggregate{})
			Expect(appended).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return b if a's count is zero", func() {
			a := StatsAggregate{}
			b := StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}
			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(b))
		})

		It("should return a if b's count is zero", func() {
			a := StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}
			b := StatsAggregate{}
			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(a))
		})

		It("should compute a correct aggregate when both a and b have non-zero counts", func() {
			a := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), 1},
				{"foo", time.Now().UTC(), 2},
				{"foo", time.Now().UTC(), 3},
				{"foo", time.Now().UTC(), 4},
				{"foo", time.Now().UTC(), 5}})
			Expect(a).To(Equal(StatsAggregate{Average: 3, Min: 1, Max: 5, Count: 5}))

			b := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), 5},
				{"foo", time.Now().UTC(), 7}})
			Expect(b).To(Equal(StatsAggregate{Average: 6, Min: 5, Max: 7, Count: 2}))

			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(StatsAggregate{Average: 3.857142857142857, Min: 1, Max: 7, Count: 7}))
		})

		It("should compute a correct aggregate for negative values", func() {
			a := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), -1},
				{"foo", time.Now().UTC(), -2},
				{"foo", time.Now().UTC(), -3},
				{"foo", time.Now().UTC(), -4},
				{"foo", time.Now().UTC(), -5}})
			Expect(a).To(Equal(StatsAggregate{Average: -3, Min: -5, Max: -1, Count: 5}))

			b := Aggregate([]*stat.Stat{{"foo", time.Now().UTC(), -6}})
			Expect(b).To(Equal(StatsAggregate{Average: -6, Min: -6, Max: -6, Count: 1}))

			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(StatsAggregate{Average: -3.5, Min: -6, Max: -1, Count: 6}))
		})
	})
})
<commit_msg>changed aggregator to use a slice of stat pointers instead of a slice of stats<commit_after>package aggregator

import (
	"github.com/CapillarySoftware/gostat/stat"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"time"
)

var _ = Describe("Aggregator", func() {

	Describe("Aggregate", func() {
		It("should return all 0 values if a nil slice is received", func() {
			a := Aggregate(nil)
			Expect(a).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return all 0 values if an empty slice is received", func() {
			a := Aggregate([]*stat.Stat{})
			Expect(a).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return all the same values if the slice contains just one Stat", func() {
			const value = 123.456
			stats := []*stat.Stat{{"foo", time.Now().UTC(), value}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: value, Min: value, Max: value, Count: 1}))
		})

		It("should return the expected values for a collection of more than one Stat", func() {
			stats := []*stat.Stat{
				&{"foo", time.Now().UTC(), 1},
				&{"foo", time.Now().UTC(), 2},
				&{"foo", time.Now().UTC(), 3},
				&{"foo", time.Now().UTC(), 4},
				&{"foo", time.Now().UTC(), 5},
				&{"foo", time.Now().UTC(), 6}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}))
		})

		It("should ignore stat names", func() {
			stats := []*stat.Stat{
				{"each",   time.Now().UTC(), 6},
				{"stat",   time.Now().UTC(), 5},
				{"name",   time.Now().UTC(), 4},
				{"here",   time.Now().UTC(), 3},
				{"is",     time.Now().UTC(), 2},
				{"unique", time.Now().UTC(), 1}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}))
		})

		It("should correctly handle negative values", func() {
			stats := []*stat.Stat{
				{"foo", time.Now().UTC(), -1},
				{"foo", time.Now().UTC(), -2},
				{"foo", time.Now().UTC(), -3},
				{"foo", time.Now().UTC(), -4},
				{"foo", time.Now().UTC(), -5},
				{"foo", time.Now().UTC(), -6}}

			a := Aggregate(stats)
			Expect(a).To(Equal(StatsAggregate{Average: -3.5, Min: -6, Max: -1, Count: 6}))
		})
	})

	Describe("AppendStatsAggregate", func() {
		It("should return all 0 values if both parameters are zeroed-out StatsAggregates", func() {
			appended := AppendStatsAggregate(StatsAggregate{}, StatsAggregate{})
			Expect(appended).To(Equal(StatsAggregate{Average: 0, Min: 0, Max: 0, Count: 0}))
		})

		It("should return b if a's count is zero", func() {
			a := StatsAggregate{}
			b := StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}
			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(b))
		})

		It("should return a if b's count is zero", func() {
			a := StatsAggregate{Average: 3.5, Min: 1, Max: 6, Count: 6}
			b := StatsAggregate{}
			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(a))
		})

		It("should compute a correct aggregate when both a and b have non-zero counts", func() {
			a := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), 1},
				{"foo", time.Now().UTC(), 2},
				{"foo", time.Now().UTC(), 3},
				{"foo", time.Now().UTC(), 4},
				{"foo", time.Now().UTC(), 5}})
			Expect(a).To(Equal(StatsAggregate{Average: 3, Min: 1, Max: 5, Count: 5}))

			b := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), 5},
				{"foo", time.Now().UTC(), 7}})
			Expect(b).To(Equal(StatsAggregate{Average: 6, Min: 5, Max: 7, Count: 2}))

			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(StatsAggregate{Average: 3.857142857142857, Min: 1, Max: 7, Count: 7}))
		})

		It("should compute a correct aggregate for negative values", func() {
			a := Aggregate([]*stat.Stat{
				{"foo", time.Now().UTC(), -1},
				{"foo", time.Now().UTC(), -2},
				{"foo", time.Now().UTC(), -3},
				{"foo", time.Now().UTC(), -4},
				{"foo", time.Now().UTC(), -5}})
			Expect(a).To(Equal(StatsAggregate{Average: -3, Min: -5, Max: -1, Count: 5}))

			b := Aggregate([]*stat.Stat{{"foo", time.Now().UTC(), -6}})
			Expect(b).To(Equal(StatsAggregate{Average: -6, Min: -6, Max: -6, Count: 1}))

			appended := AppendStatsAggregate(a, b)
			Expect(appended).To(Equal(StatsAggregate{Average: -3.5, Min: -6, Max: -1, Count: 6}))
		})
	})
})
<|endoftext|>


<commit_before>package term

import (
	"io"

	"github.com/gdamore/tcell"
	"github.com/mibk/syd/ui"
)

type UI struct {
	screen tcell.Screen
	frame  *Frame

	p0, p1 int // cursor position
	x, y   int // current position

	wasBtnPressed bool
}

func (t *UI) Init() error {
	sc, err := tcell.NewScreen()
	if err != nil {
		return err
	}
	if err := sc.Init(); err != nil {
		return err
	}
	sc.EnableMouse()
	t.screen = sc
	t.frame = new(Frame)
	go t.translateEvents()
	return nil
}

func (t *UI) Close() error {
	t.screen.Fini()
	return nil
}

func (t *UI) Size() (w, h int) { return t.screen.Size() }

func (t *UI) Clear() {
	t.screen.Clear()
	*t.frame = Frame{
		lines:   make([][]rune, 1),
		wantCol: t.frame.wantCol,
	}
	t.x, t.y = 0, 0
	t.checkSelection()
}

func (t *UI) Select(p0, p1 int) { t.p0, t.p1 = p0, p1 }

func (t *UI) WriteRune(r rune) error {
	if r != '\n' {
		t.frame.lines[t.y] = append(t.frame.lines[t.y], r)
	}

	w, h := t.Size()
	if t.x >= w || r == '\n' {
		t.y++
		t.x = 0
		t.frame.lines = append(t.frame.lines, nil)
		if t.y == h {
			return io.EOF
		}
	} else if r == '\t' {
		t.x += tabWidthForCol(t.x)
	} else {
		t.x++
	}
	t.frame.nchars++
	t.checkSelection()
	return nil
}

// checkSelection tries to line0, line1, and wantCol.
func (t *UI) checkSelection() {
	if t.p0 == t.frame.nchars {
		t.frame.line0 = t.y
		if t.frame.wantCol == ui.ColQ0 {
			t.frame.wantCol = t.x
		}
	}
	if t.p1 == t.frame.nchars {
		t.frame.line1 = t.y
		if t.frame.wantCol == ui.ColQ1 {
			t.frame.wantCol = t.x
		}
	}
}

func (t *UI) Flush() {
	st := tcell.StyleDefault
	selText := func(p, x, y int) {
		if p == t.p0 {
			if t.p0 == t.p1 {
				t.screen.ShowCursor(x, y)
			} else {
				st = st.Reverse(true)
			}
		} else if p == t.p1 {
			st = st.Reverse(false)
		}
	}
	t.screen.HideCursor()
	p := 0
	for y, l := range t.frame.lines {
		x := 0
		for _, r := range l {
			selText(p, x, y)
			t.screen.SetContent(x, y, r, nil, st)
			p++
			if r == '\t' {
				x += tabWidthForCol(x)
			} else {
				x++
			}
		}
		selText(p, x, y)
		p++
	}
	t.screen.Show()
}

func (t *UI) Frame() ui.Frame { return t.frame }

type Frame struct {
	lines   [][]rune
	line0   int
	line1   int
	wantCol int
	nchars  int
}

func (f *Frame) Nchars() int                { return f.nchars }
func (f *Frame) SelectionLines() (int, int) { return f.line0, f.line1 }

func (f *Frame) CharsUntilXY(x, y int) int {
	if y >= len(f.lines) {
		return f.nchars
	}
	var p int
	for n, l := range f.lines {
		if n == y {
			return p + charsUntilX(l, x)
		}
		p += len(l) + 1 // + '\n'
	}
	panic("shouldn't happen")
}

func charsUntilX(s []rune, x int) int {
	var w int
	for i, r := range s {
		if r == '\t' {
			w += tabWidthForCol(w)
		} else {
			w += 1
		}
		if w > x {
			return i
		}
	}
	return len(s)
}

const tabStop = 8

func tabWidthForCol(col int) int {
	w := tabStop - col%tabStop
	if w == 0 {
		return tabStop
	}
	return w
}

func (f *Frame) MaxLines() int { panic("not implemented") }
func (f *Frame) Lines() int    { return len(f.lines) }

func (f *Frame) WantCol() int       { return f.wantCol }
func (f *Frame) SetWantCol(col int) { f.wantCol = col }
<commit_msg>Fix highlighting of text<commit_after>package term

import (
	"io"

	"github.com/gdamore/tcell"
	"github.com/mibk/syd/ui"
)

type UI struct {
	screen tcell.Screen
	frame  *Frame

	p0, p1 int // cursor position
	x, y   int // current position

	wasBtnPressed bool
}

func (t *UI) Init() error {
	sc, err := tcell.NewScreen()
	if err != nil {
		return err
	}
	if err := sc.Init(); err != nil {
		return err
	}
	sc.EnableMouse()
	t.screen = sc
	t.frame = new(Frame)
	go t.translateEvents()
	return nil
}

func (t *UI) Close() error {
	t.screen.Fini()
	return nil
}

func (t *UI) Size() (w, h int) { return t.screen.Size() }

func (t *UI) Clear() {
	t.screen.Clear()
	*t.frame = Frame{
		lines:   make([][]rune, 1),
		wantCol: t.frame.wantCol,
	}
	t.x, t.y = 0, 0
	t.checkSelection()
}

func (t *UI) Select(p0, p1 int) { t.p0, t.p1 = p0, p1 }

func (t *UI) WriteRune(r rune) error {
	if r != '\n' {
		t.frame.lines[t.y] = append(t.frame.lines[t.y], r)
	}

	w, h := t.Size()
	if t.x >= w || r == '\n' {
		t.y++
		t.x = 0
		t.frame.lines = append(t.frame.lines, nil)
		if t.y == h {
			return io.EOF
		}
	} else if r == '\t' {
		t.x += tabWidthForCol(t.x)
	} else {
		t.x++
	}
	t.frame.nchars++
	t.checkSelection()
	return nil
}

// checkSelection tries to line0, line1, and wantCol.
func (t *UI) checkSelection() {
	if t.p0 == t.frame.nchars {
		t.frame.line0 = t.y
		if t.frame.wantCol == ui.ColQ0 {
			t.frame.wantCol = t.x
		}
	}
	if t.p1 == t.frame.nchars {
		t.frame.line1 = t.y
		if t.frame.wantCol == ui.ColQ1 {
			t.frame.wantCol = t.x
		}
	}
}

func (t *UI) Flush() {
	width, _ := t.Size()
	st := tcell.StyleDefault
	selText := func(p, x, y int) {
		if p == t.p0 {
			if t.p0 == t.p1 {
				t.screen.ShowCursor(x, y)
			} else {
				st = st.Reverse(true)
			}
		} else if p == t.p1 {
			st = st.Reverse(false)
		}
	}
	t.screen.HideCursor()
	p := 0
	for y, l := range t.frame.lines {
		x := 0
		for _, r := range l {
			selText(p, x, y)
			w := 1
			if r == '\t' {
				r = ' '
				w = tabWidthForCol(x)

			}
			for i := 0; i < w; i++ {
				t.screen.SetContent(x, y, r, nil, st)
				x += 1
			}
			p++
		}
		selText(p, x, y)
		for ; x < width; x++ {
			t.screen.SetContent(x, y, ' ', nil, st)
		}
		p++
	}
	t.screen.Show()
}

func (t *UI) Frame() ui.Frame { return t.frame }

type Frame struct {
	lines   [][]rune
	line0   int
	line1   int
	wantCol int
	nchars  int
}

func (f *Frame) Nchars() int                { return f.nchars }
func (f *Frame) SelectionLines() (int, int) { return f.line0, f.line1 }

func (f *Frame) CharsUntilXY(x, y int) int {
	if y >= len(f.lines) {
		return f.nchars
	}
	var p int
	for n, l := range f.lines {
		if n == y {
			return p + charsUntilX(l, x)
		}
		p += len(l) + 1 // + '\n'
	}
	panic("shouldn't happen")
}

func charsUntilX(s []rune, x int) int {
	var w int
	for i, r := range s {
		if r == '\t' {
			w += tabWidthForCol(w)
		} else {
			w += 1
		}
		if w > x {
			return i
		}
	}
	return len(s)
}

const tabStop = 8

func tabWidthForCol(col int) int {
	w := tabStop - col%tabStop
	if w == 0 {
		return tabStop
	}
	return w
}

func (f *Frame) MaxLines() int { panic("not implemented") }
func (f *Frame) Lines() int    { return len(f.lines) }

func (f *Frame) WantCol() int       { return f.wantCol }
func (f *Frame) SetWantCol(col int) { f.wantCol = col }
<|endoftext|>


<commit_before>package link

import (
	"errors"
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"path"
	"testing"

	qt "github.com/frankban/quicktest"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal/testutils"
	"github.com/cilium/ebpf/internal/unix"
)

var (
	bashEx, _ = OpenExecutable("/bin/bash")
	bashSym   = "main"
)

func TestExecutable(t *testing.T) {
	_, err := OpenExecutable("")
	if err == nil {
		t.Fatal("create executable: expected error on empty path")
	}

	if bashEx.path != "/bin/bash" {
		t.Fatalf("create executable: unexpected path '%s'", bashEx.path)
	}

	_, err = bashEx.offset(bashSym)
	if err != nil {
		t.Fatalf("find offset: %v", err)
	}

	_, err = bashEx.offset("bogus")
	if err == nil {
		t.Fatal("find symbol: expected error")
	}
}

func TestUprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uprobe(bashSym, prog, nil)
	c.Assert(err, qt.IsNil)
	defer up.Close()

	testLink(t, up, testLinkOptions{
		prog: prog,
	})
}

func TestUprobeExtNotFound(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// This symbol will not be present in Executable (elf.SHN_UNDEF).
	_, err = bashEx.Uprobe("open", prog, nil)
	if err == nil {
		t.Fatal("expected error")
	}
}

func TestUprobeExtWithOpts(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// This Uprobe is broken and will not work because the offset is not
	// correct. This is expected since the offset is provided by the user.
	up, err := bashEx.Uprobe("open", prog, &UprobeOptions{Offset: 0x1})
	if err != nil {
		t.Fatal(err)
	}
	defer up.Close()
}

func TestUprobeWithPID(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uprobe(bashSym, prog, &UprobeOptions{PID: os.Getpid()})
	if err != nil {
		t.Fatal(err)
	}
	defer up.Close()
}

func TestUprobeWithNonExistentPID(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// trying to open a perf event on a non-existent PID will return ESRCH.
	_, err = bashEx.Uprobe(bashSym, prog, &UprobeOptions{PID: -2})
	if !errors.Is(err, unix.ESRCH) {
		t.Fatalf("expected ESRCH, got %v", err)
	}
}

func TestUretprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uretprobe(bashSym, prog, nil)
	c.Assert(err, qt.IsNil)
	defer up.Close()

	testLink(t, up, testLinkOptions{
		prog: prog,
	})
}

// Test u(ret)probe creation using perf_uprobe PMU.
func TestUprobeCreatePMU(t *testing.T) {
	// Requires at least 4.17 (e12f03d7031a "perf/core: Implement the 'perf_kprobe' PMU")
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")

	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// uprobe PMU
	pu, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer pu.Close()

	c.Assert(pu.typ, qt.Equals, uprobeEvent)

	// uretprobe PMU
	pr, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, true)
	c.Assert(err, qt.IsNil)
	defer pr.Close()

	c.Assert(pr.typ, qt.Equals, uretprobeEvent)
}

// Test fallback behaviour on kernels without perf_uprobe PMU available.
func TestUprobePMUUnavailable(t *testing.T) {
	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	pk, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, false)
	if err == nil {
		pk.Close()
		t.Skipf("Kernel supports perf_uprobe PMU, not asserting error.")
	}

	// Expect ErrNotSupported.
	c.Assert(errors.Is(err, ErrNotSupported), qt.IsTrue, qt.Commentf("got error: %s", err))
}

// Test tracefs u(ret)probe creation on all kernel versions.
func TestUprobeTraceFS(t *testing.T) {
	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// Sanitize the symbol in order to be used in tracefs API.
	ssym := uprobeSanitizedSymbol(bashSym)

	// Open and close tracefs u(ret)probes, checking all errors.
	up, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	c.Assert(up.Close(), qt.IsNil)
	c.Assert(up.typ, qt.Equals, uprobeEvent)

	up, err = tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, true)
	c.Assert(err, qt.IsNil)
	c.Assert(up.Close(), qt.IsNil)
	c.Assert(up.typ, qt.Equals, uretprobeEvent)

	// Create two identical trace events, ensure their IDs differ.
	u1, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer u1.Close()
	c.Assert(u1.tracefsID, qt.Not(qt.Equals), 0)

	u2, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer u2.Close()
	c.Assert(u2.tracefsID, qt.Not(qt.Equals), 0)

	// Compare the uprobes' tracefs IDs.
	c.Assert(u1.tracefsID, qt.Not(qt.CmpEquals()), u2.tracefsID)
}

// Test u(ret)probe creation writing directly to <tracefs>/uprobe_events.
// Only runs on 5.0 and over. Earlier versions ignored writes of duplicate
// events, while 5.0 started returning -EEXIST when a uprobe event already
// exists.
func TestUprobeCreateTraceFS(t *testing.T) {
	testutils.SkipOnOldKernel(t, "5.0", "<tracefs>/uprobe_events doesn't reject duplicate events")

	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// Sanitize the symbol in order to be used in tracefs API.
	ssym := uprobeSanitizedSymbol(bashSym)

	pg, _ := randomGroup("ebpftest")
	rg, _ := randomGroup("ebpftest")

	// Tee up cleanups in case any of the Asserts abort the function.
	defer func() {
		_ = closeTraceFSProbeEvent(uprobeType, pg, ssym)
		_ = closeTraceFSProbeEvent(uprobeType, rg, ssym)
	}()

	// Create a uprobe.
	err = createTraceFSProbeEvent(uprobeType, pg, ssym, bashEx.path, off, false)
	c.Assert(err, qt.IsNil)

	// Attempt to create an identical uprobe using tracefs,
	// expect it to fail with os.ErrExist.
	err = createTraceFSProbeEvent(uprobeType, pg, ssym, bashEx.path, off, false)
	c.Assert(errors.Is(err, os.ErrExist), qt.IsTrue,
		qt.Commentf("expected consecutive uprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kprobe.
	c.Assert(closeTraceFSProbeEvent(uprobeType, pg, ssym), qt.IsNil)

	// Same test for a kretprobe.
	err = createTraceFSProbeEvent(uprobeType, rg, ssym, bashEx.path, off, true)
	c.Assert(err, qt.IsNil)

	err = createTraceFSProbeEvent(uprobeType, rg, ssym, bashEx.path, off, true)
	c.Assert(os.IsExist(err), qt.IsFalse,
		qt.Commentf("expected consecutive uretprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the uretprobe.
	c.Assert(closeTraceFSProbeEvent(uprobeType, rg, ssym), qt.IsNil)
}

func TestUprobeSanitizedSymbol(t *testing.T) {
	tests := []struct {
		symbol   string
		expected string
	}{
		{"readline", "readline"},
		{"main.Func", "main_Func"},
		{"a.....a", "a_a"},
		{"./;'{}[]a", "_a"},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprint(i), func(t *testing.T) {
			sanitized := uprobeSanitizedSymbol(tt.symbol)
			if tt.expected != sanitized {
				t.Errorf("Expected sanitized symbol to be '%s', got '%s'", tt.expected, sanitized)
			}
		})
	}
}

func TestUprobePathOffset(t *testing.T) {
	tests := []struct {
		path     string
		offset   uint64
		expected string
	}{
		{"/bin/bash", 0, "/bin/bash:0x0"},
		{"/bin/bash", 1, "/bin/bash:0x1"},
		{"/bin/bash", 65535, "/bin/bash:0xffff"},
		{"/bin/bash", 65536, "/bin/bash:0x10000"},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprint(i), func(t *testing.T) {
			po := uprobePathOffset(tt.path, tt.offset)
			if tt.expected != po {
				t.Errorf("Expected path:offset to be '%s', got '%s'", tt.expected, po)
			}
		})
	}
}

func TestUprobeProgramCall(t *testing.T) {
	tests := []struct {
		name string
		elf  string
		args []string
		sym  string
	}{
		{
			"bash",
			"/bin/bash",
			[]string{"--help"},
			"main",
		},
		{
			"go-binary",
			path.Join(build.Default.GOROOT, "bin/go"),
			[]string{"version"},
			"main.main",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			m, p := newUpdaterMapProg(t, ebpf.Kprobe)

			// Load the executable.
			ex, err := OpenExecutable(tt.elf)
			if err != nil {
				t.Fatal(err)
			}

			// Open Uprobe on the executable for the given symbol
			// and attach it to the ebpf program created above.
			u, err := ex.Uprobe(tt.sym, p, nil)
			if err != nil {
				t.Fatal(err)
			}

			// Trigger ebpf program call.
			trigger := func(t *testing.T) {
				if err := exec.Command(tt.elf, tt.args...).Run(); err != nil {
					t.Fatal(err)
				}
			}
			trigger(t)

			// Assert that the value at index 0 has been updated to 1.
			assertMapValue(t, m, 0, 1)

			// Detach the Uprobe.
			if err := u.Close(); err != nil {
				t.Fatal(err)
			}

			// Reset map value to 0 at index 0.
			if err := m.Update(uint32(0), uint32(0), ebpf.UpdateExist); err != nil {
				t.Fatal(err)
			}

			// Retrigger the ebpf program call.
			trigger(t)

			// Assert that this time the value has not been updated.
			assertMapValue(t, m, 0, 0)
		})
	}
}

func TestUprobeProgramWrongPID(t *testing.T) {
	m, p := newUpdaterMapProg(t, ebpf.Kprobe)

	// Load the '/bin/bash' executable.
	ex, err := OpenExecutable("/bin/bash")
	if err != nil {
		t.Fatal(err)
	}

	// Open Uprobe on '/bin/bash' for the symbol 'main'
	// and attach it to the ebpf program created above.
	// Create the perf-event with the current process' PID
	// to make sure the event is not fired when we will try
	// to trigger the program execution via exec.
	u, err := ex.Uprobe("main", p, &UprobeOptions{PID: os.Getpid()})
	if err != nil {
		t.Fatal(err)
	}
	defer u.Close()

	// Trigger ebpf program call.
	if err := exec.Command("/bin/bash", "--help").Run(); err != nil {
		t.Fatal(err)
	}

	// Assert that the value at index 0 is still 0.
	assertMapValue(t, m, 0, 0)
}
<commit_msg>link/uprobe: skip Go binary test on kernels < 4.14<commit_after>package link

import (
	"errors"
	"fmt"
	"go/build"
	"os"
	"os/exec"
	"path"
	"testing"

	qt "github.com/frankban/quicktest"

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/internal/testutils"
	"github.com/cilium/ebpf/internal/unix"
)

var (
	bashEx, _ = OpenExecutable("/bin/bash")
	bashSym   = "main"
)

func TestExecutable(t *testing.T) {
	_, err := OpenExecutable("")
	if err == nil {
		t.Fatal("create executable: expected error on empty path")
	}

	if bashEx.path != "/bin/bash" {
		t.Fatalf("create executable: unexpected path '%s'", bashEx.path)
	}

	_, err = bashEx.offset(bashSym)
	if err != nil {
		t.Fatalf("find offset: %v", err)
	}

	_, err = bashEx.offset("bogus")
	if err == nil {
		t.Fatal("find symbol: expected error")
	}
}

func TestUprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uprobe(bashSym, prog, nil)
	c.Assert(err, qt.IsNil)
	defer up.Close()

	testLink(t, up, testLinkOptions{
		prog: prog,
	})
}

func TestUprobeExtNotFound(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// This symbol will not be present in Executable (elf.SHN_UNDEF).
	_, err = bashEx.Uprobe("open", prog, nil)
	if err == nil {
		t.Fatal("expected error")
	}
}

func TestUprobeExtWithOpts(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// This Uprobe is broken and will not work because the offset is not
	// correct. This is expected since the offset is provided by the user.
	up, err := bashEx.Uprobe("open", prog, &UprobeOptions{Offset: 0x1})
	if err != nil {
		t.Fatal(err)
	}
	defer up.Close()
}

func TestUprobeWithPID(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uprobe(bashSym, prog, &UprobeOptions{PID: os.Getpid()})
	if err != nil {
		t.Fatal(err)
	}
	defer up.Close()
}

func TestUprobeWithNonExistentPID(t *testing.T) {
	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	// trying to open a perf event on a non-existent PID will return ESRCH.
	_, err = bashEx.Uprobe(bashSym, prog, &UprobeOptions{PID: -2})
	if !errors.Is(err, unix.ESRCH) {
		t.Fatalf("expected ESRCH, got %v", err)
	}
}

func TestUretprobe(t *testing.T) {
	c := qt.New(t)

	prog, err := ebpf.NewProgram(&kprobeSpec)
	if err != nil {
		t.Fatal(err)
	}
	defer prog.Close()

	up, err := bashEx.Uretprobe(bashSym, prog, nil)
	c.Assert(err, qt.IsNil)
	defer up.Close()

	testLink(t, up, testLinkOptions{
		prog: prog,
	})
}

// Test u(ret)probe creation using perf_uprobe PMU.
func TestUprobeCreatePMU(t *testing.T) {
	// Requires at least 4.17 (e12f03d7031a "perf/core: Implement the 'perf_kprobe' PMU")
	testutils.SkipOnOldKernel(t, "4.17", "perf_kprobe PMU")

	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// uprobe PMU
	pu, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer pu.Close()

	c.Assert(pu.typ, qt.Equals, uprobeEvent)

	// uretprobe PMU
	pr, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, true)
	c.Assert(err, qt.IsNil)
	defer pr.Close()

	c.Assert(pr.typ, qt.Equals, uretprobeEvent)
}

// Test fallback behaviour on kernels without perf_uprobe PMU available.
func TestUprobePMUUnavailable(t *testing.T) {
	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	pk, err := pmuUprobe(bashSym, bashEx.path, off, perfAllThreads, false)
	if err == nil {
		pk.Close()
		t.Skipf("Kernel supports perf_uprobe PMU, not asserting error.")
	}

	// Expect ErrNotSupported.
	c.Assert(errors.Is(err, ErrNotSupported), qt.IsTrue, qt.Commentf("got error: %s", err))
}

// Test tracefs u(ret)probe creation on all kernel versions.
func TestUprobeTraceFS(t *testing.T) {
	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// Sanitize the symbol in order to be used in tracefs API.
	ssym := uprobeSanitizedSymbol(bashSym)

	// Open and close tracefs u(ret)probes, checking all errors.
	up, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	c.Assert(up.Close(), qt.IsNil)
	c.Assert(up.typ, qt.Equals, uprobeEvent)

	up, err = tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, true)
	c.Assert(err, qt.IsNil)
	c.Assert(up.Close(), qt.IsNil)
	c.Assert(up.typ, qt.Equals, uretprobeEvent)

	// Create two identical trace events, ensure their IDs differ.
	u1, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer u1.Close()
	c.Assert(u1.tracefsID, qt.Not(qt.Equals), 0)

	u2, err := tracefsUprobe(ssym, bashEx.path, off, perfAllThreads, false)
	c.Assert(err, qt.IsNil)
	defer u2.Close()
	c.Assert(u2.tracefsID, qt.Not(qt.Equals), 0)

	// Compare the uprobes' tracefs IDs.
	c.Assert(u1.tracefsID, qt.Not(qt.CmpEquals()), u2.tracefsID)
}

// Test u(ret)probe creation writing directly to <tracefs>/uprobe_events.
// Only runs on 5.0 and over. Earlier versions ignored writes of duplicate
// events, while 5.0 started returning -EEXIST when a uprobe event already
// exists.
func TestUprobeCreateTraceFS(t *testing.T) {
	testutils.SkipOnOldKernel(t, "5.0", "<tracefs>/uprobe_events doesn't reject duplicate events")

	c := qt.New(t)

	// Fetch the offset from the /bin/bash Executable already defined.
	off, err := bashEx.offset(bashSym)
	c.Assert(err, qt.IsNil)

	// Sanitize the symbol in order to be used in tracefs API.
	ssym := uprobeSanitizedSymbol(bashSym)

	pg, _ := randomGroup("ebpftest")
	rg, _ := randomGroup("ebpftest")

	// Tee up cleanups in case any of the Asserts abort the function.
	defer func() {
		_ = closeTraceFSProbeEvent(uprobeType, pg, ssym)
		_ = closeTraceFSProbeEvent(uprobeType, rg, ssym)
	}()

	// Create a uprobe.
	err = createTraceFSProbeEvent(uprobeType, pg, ssym, bashEx.path, off, false)
	c.Assert(err, qt.IsNil)

	// Attempt to create an identical uprobe using tracefs,
	// expect it to fail with os.ErrExist.
	err = createTraceFSProbeEvent(uprobeType, pg, ssym, bashEx.path, off, false)
	c.Assert(errors.Is(err, os.ErrExist), qt.IsTrue,
		qt.Commentf("expected consecutive uprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the kprobe.
	c.Assert(closeTraceFSProbeEvent(uprobeType, pg, ssym), qt.IsNil)

	// Same test for a kretprobe.
	err = createTraceFSProbeEvent(uprobeType, rg, ssym, bashEx.path, off, true)
	c.Assert(err, qt.IsNil)

	err = createTraceFSProbeEvent(uprobeType, rg, ssym, bashEx.path, off, true)
	c.Assert(os.IsExist(err), qt.IsFalse,
		qt.Commentf("expected consecutive uretprobe creation to contain os.ErrExist, got: %v", err))

	// Expect a successful close of the uretprobe.
	c.Assert(closeTraceFSProbeEvent(uprobeType, rg, ssym), qt.IsNil)
}

func TestUprobeSanitizedSymbol(t *testing.T) {
	tests := []struct {
		symbol   string
		expected string
	}{
		{"readline", "readline"},
		{"main.Func", "main_Func"},
		{"a.....a", "a_a"},
		{"./;'{}[]a", "_a"},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprint(i), func(t *testing.T) {
			sanitized := uprobeSanitizedSymbol(tt.symbol)
			if tt.expected != sanitized {
				t.Errorf("Expected sanitized symbol to be '%s', got '%s'", tt.expected, sanitized)
			}
		})
	}
}

func TestUprobePathOffset(t *testing.T) {
	tests := []struct {
		path     string
		offset   uint64
		expected string
	}{
		{"/bin/bash", 0, "/bin/bash:0x0"},
		{"/bin/bash", 1, "/bin/bash:0x1"},
		{"/bin/bash", 65535, "/bin/bash:0xffff"},
		{"/bin/bash", 65536, "/bin/bash:0x10000"},
	}

	for i, tt := range tests {
		t.Run(fmt.Sprint(i), func(t *testing.T) {
			po := uprobePathOffset(tt.path, tt.offset)
			if tt.expected != po {
				t.Errorf("Expected path:offset to be '%s', got '%s'", tt.expected, po)
			}
		})
	}
}

func TestUprobeProgramCall(t *testing.T) {
	tests := []struct {
		name string
		elf  string
		args []string
		sym  string
	}{
		{
			"bash",
			"/bin/bash",
			[]string{"--help"},
			"main",
		},
		{
			"go-binary",
			path.Join(build.Default.GOROOT, "bin/go"),
			[]string{"version"},
			"main.main",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if tt.name == "go-binary" {
				// https://github.com/cilium/ebpf/issues/406
				testutils.SkipOnOldKernel(t, "4.14", "uprobes on Go binaries silently fail on kernel < 4.14")
			}

			m, p := newUpdaterMapProg(t, ebpf.Kprobe)

			// Load the executable.
			ex, err := OpenExecutable(tt.elf)
			if err != nil {
				t.Fatal(err)
			}

			// Open Uprobe on the executable for the given symbol
			// and attach it to the ebpf program created above.
			u, err := ex.Uprobe(tt.sym, p, nil)
			if err != nil {
				t.Fatal(err)
			}

			// Trigger ebpf program call.
			trigger := func(t *testing.T) {
				if err := exec.Command(tt.elf, tt.args...).Run(); err != nil {
					t.Fatal(err)
				}
			}
			trigger(t)

			// Assert that the value at index 0 has been updated to 1.
			assertMapValue(t, m, 0, 1)

			// Detach the Uprobe.
			if err := u.Close(); err != nil {
				t.Fatal(err)
			}

			// Reset map value to 0 at index 0.
			if err := m.Update(uint32(0), uint32(0), ebpf.UpdateExist); err != nil {
				t.Fatal(err)
			}

			// Retrigger the ebpf program call.
			trigger(t)

			// Assert that this time the value has not been updated.
			assertMapValue(t, m, 0, 0)
		})
	}
}

func TestUprobeProgramWrongPID(t *testing.T) {
	m, p := newUpdaterMapProg(t, ebpf.Kprobe)

	// Load the '/bin/bash' executable.
	ex, err := OpenExecutable("/bin/bash")
	if err != nil {
		t.Fatal(err)
	}

	// Open Uprobe on '/bin/bash' for the symbol 'main'
	// and attach it to the ebpf program created above.
	// Create the perf-event with the current process' PID
	// to make sure the event is not fired when we will try
	// to trigger the program execution via exec.
	u, err := ex.Uprobe("main", p, &UprobeOptions{PID: os.Getpid()})
	if err != nil {
		t.Fatal(err)
	}
	defer u.Close()

	// Trigger ebpf program call.
	if err := exec.Command("/bin/bash", "--help").Run(); err != nil {
		t.Fatal(err)
	}

	// Assert that the value at index 0 is still 0.
	assertMapValue(t, m, 0, 0)
}
<|endoftext|>


<commit_before>package sync

import (
	"bytes"
	"encoding/binary"
	"encoding/gob"
	"os/user"
	"path/filepath"
	"time"

	"github.com/boltdb/bolt"
)

// buckets
var (
	downloadItemsBucket   = []byte("download-items")
	watchedTorrentsBucket = []byte("watched-torrents")
	defaultsBucket        = []byte("defaults")
)

type Error string

func (e Error) Error() string { return string(e) }

const (
	ErrStateNotFound   = Error("state not found")
	ErrConfigNotFound  = Error("configuration not found")
	ErrSaveStateFailed = Error("state could not be saved")
)

type Store struct {
	path string
	db   *bolt.DB
}

func NewStore(path string) *Store {
	return &Store{path: path}
}

func (s *Store) Open() error {
	db, err := bolt.Open(s.path, 0666, &bolt.Options{Timeout: 10 * time.Second})
	if err != nil {
		return err
	}
	s.db = db

	err = s.db.Update(func(tx *bolt.Tx) error {
		_, err = tx.CreateBucketIfNotExists(defaultsBucket)
		return err
	})
	if err != nil {
		return s.db.Close()
	}
	return nil
}

func (s *Store) Close() error { return s.db.Close() }

func (s *Store) Path() string { return s.path }

func (s *Store) CreateBuckets(forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt, err := tx.CreateBucketIfNotExists([]byte(forUser))
		if err != nil {
			return err
		}

		buckets := [][]byte{
			downloadItemsBucket,
			watchedTorrentsBucket,
		}

		for _, bucket := range buckets {
			_, err = userBkt.CreateBucketIfNotExists(bucket)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

// SaveState inserts or updates the given state.
func (s *Store) SaveState(state *State, forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)

		key := itob(state.FileID)
		var value bytes.Buffer

		err := gob.NewEncoder(&value).Encode(state)
		if err != nil {
			return err
		}

		return downloadsBkt.Put(key, value.Bytes())
	})
}

// State returns a state by the given file ID.
func (s *Store) State(id int64, forUser string) (*State, error) {
	var state State
	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)
		fileID := itob(id)

		value := downloadsBkt.Get(fileID)
		if value == nil {
			return ErrStateNotFound
		}

		return gob.NewDecoder(bytes.NewReader(value)).Decode(&state)
	})
	return &state, err
}

// States returns all the states in the store.
func (s *Store) States(forUser string) ([]*State, error) {
	states := make([]*State, 0)

	if forUser == "" {
		return states, nil
	}

	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)

		cursor := downloadsBkt.Cursor()
		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
			var state State
			err := gob.NewDecoder(bytes.NewReader(v)).Decode(&state)
			if err != nil {
				return err
			}
			// dont include hidden downloads
			if state.IsHidden {
				continue
			}
			states = append(states, &state)
		}
		return nil
	})

	return states, err
}

func (s *Store) Config(forUser string) (*Config, error) {
	if forUser == "" {
		return s.DefaultConfig()
	}

	var cfg *Config
	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))

		key := []byte("config")
		value := userBkt.Get(key)

		if value == nil {
			return ErrConfigNotFound
		}

		return gob.NewDecoder(bytes.NewReader(value)).Decode(cfg)
	})

	if err == ErrConfigNotFound {
		cfg, err = s.DefaultConfig()
	}

	return cfg, err
}

func (s *Store) SaveConfig(cfg *Config, forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))

		key := []byte("config")
		var value bytes.Buffer

		err := gob.NewEncoder(&value).Encode(cfg)
		if err != nil {
			return err
		}

		return userBkt.Put(key, value.Bytes())
	})
}

func (s *Store) DefaultConfig() (*Config, error) {
	u, err := user.Current()
	if err != nil {
		return nil, err
	}
	return &Config{
		PollInterval:        Duration(defaultPollInterval),
		DownloadTo:          filepath.Join(u.HomeDir, "putio-sync"),
		DownloadFrom:        defaultDownloadFrom,
		SegmentsPerFile:     defaultSegmentsPerFile,
		MaxParallelFiles:    defaultMaxParallelFiles,
		IsPaused:            true,
		WatchTorrentsFolder: false,
		TorrentsFolder:      "",
	}, nil
}

// CurrentUser returns the last login user.
func (s *Store) CurrentUser() (string, error) {
	var username []byte
	err := s.db.View(func(tx *bolt.Tx) error {
		bkt := tx.Bucket(defaultsBucket)
		username = bkt.Get([]byte("current-user"))
		return nil
	})
	if err != nil {
		return "", err
	}

	if username == nil {
		return "", nil
	}
	return string(username), nil
}

func (s *Store) SaveCurrentUser(username string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		bkt := tx.Bucket(defaultsBucket)
		key := []byte("current-user")
		return bkt.Put(key, []byte(username))
	})
}

func itob(v int64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(v))
	return b
}
<commit_msg>sync: more idiomatic go<commit_after>package sync

import (
	"bytes"
	"encoding/binary"
	"encoding/gob"
	"os/user"
	"path/filepath"
	"time"

	"github.com/boltdb/bolt"
)

// buckets
var (
	downloadItemsBucket   = []byte("download-items")
	watchedTorrentsBucket = []byte("watched-torrents")
	defaultsBucket        = []byte("defaults")
)

type Error string

func (e Error) Error() string { return string(e) }

const (
	ErrStateNotFound   = Error("state not found")
	ErrConfigNotFound  = Error("configuration not found")
	ErrSaveStateFailed = Error("state could not be saved")
)

type Store struct {
	path string
	db   *bolt.DB
}

func NewStore(path string) *Store {
	return &Store{path: path}
}

func (s *Store) Open() error {
	db, err := bolt.Open(s.path, 0666, &bolt.Options{Timeout: 10 * time.Second})
	if err != nil {
		return err
	}
	s.db = db

	err = s.db.Update(func(tx *bolt.Tx) error {
		_, err = tx.CreateBucketIfNotExists(defaultsBucket)
		return err
	})
	if err != nil {
		return s.db.Close()
	}
	return nil
}

func (s *Store) Close() error { return s.db.Close() }

func (s *Store) Path() string { return s.path }

func (s *Store) CreateBuckets(forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt, err := tx.CreateBucketIfNotExists([]byte(forUser))
		if err != nil {
			return err
		}

		buckets := [][]byte{
			downloadItemsBucket,
			watchedTorrentsBucket,
		}

		for _, bucket := range buckets {
			_, err = userBkt.CreateBucketIfNotExists(bucket)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

// SaveState inserts or updates the given state.
func (s *Store) SaveState(state *State, forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)

		key := itob(state.FileID)
		var value bytes.Buffer

		err := gob.NewEncoder(&value).Encode(state)
		if err != nil {
			return err
		}

		return downloadsBkt.Put(key, value.Bytes())
	})
}

// State returns a state by the given file ID.
func (s *Store) State(id int64, forUser string) (*State, error) {
	var state State
	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)
		fileID := itob(id)

		value := downloadsBkt.Get(fileID)
		if value == nil {
			return ErrStateNotFound
		}

		return gob.NewDecoder(bytes.NewReader(value)).Decode(&state)
	})
	return &state, err
}

// States returns all the states in the store.
func (s *Store) States(forUser string) ([]*State, error) {
	states := make([]*State, 0)

	if forUser == "" {
		return states, nil
	}

	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))
		downloadsBkt := userBkt.Bucket(downloadItemsBucket)

		cursor := downloadsBkt.Cursor()
		for k, v := cursor.First(); k != nil; k, v = cursor.Next() {
			var state State
			err := gob.NewDecoder(bytes.NewReader(v)).Decode(&state)
			if err != nil {
				return err
			}
			// dont include hidden downloads
			if state.IsHidden {
				continue
			}
			states = append(states, &state)
		}
		return nil
	})

	return states, err
}

func (s *Store) Config(forUser string) (*Config, error) {
	if forUser == "" {
		return s.DefaultConfig()
	}

	var cfg Config
	err := s.db.View(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))

		key := []byte("config")
		value := userBkt.Get(key)

		if value == nil {
			return ErrConfigNotFound
		}

		return gob.NewDecoder(bytes.NewReader(value)).Decode(&cfg)
	})

	if err == ErrConfigNotFound {
		return s.DefaultConfig()
	}

	return &cfg, err
}

func (s *Store) SaveConfig(cfg *Config, forUser string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		userBkt := tx.Bucket([]byte(forUser))

		key := []byte("config")
		var value bytes.Buffer

		err := gob.NewEncoder(&value).Encode(cfg)
		if err != nil {
			return err
		}

		return userBkt.Put(key, value.Bytes())
	})
}

func (s *Store) DefaultConfig() (*Config, error) {
	u, err := user.Current()
	if err != nil {
		return nil, err
	}
	return &Config{
		PollInterval:        Duration(defaultPollInterval),
		DownloadTo:          filepath.Join(u.HomeDir, "putio-sync"),
		DownloadFrom:        defaultDownloadFrom,
		SegmentsPerFile:     defaultSegmentsPerFile,
		MaxParallelFiles:    defaultMaxParallelFiles,
		IsPaused:            true,
		WatchTorrentsFolder: false,
		TorrentsFolder:      "",
	}, nil
}

// CurrentUser returns the last login user.
func (s *Store) CurrentUser() (string, error) {
	var username []byte
	err := s.db.View(func(tx *bolt.Tx) error {
		bkt := tx.Bucket(defaultsBucket)
		username = bkt.Get([]byte("current-user"))
		return nil
	})
	if err != nil {
		return "", err
	}

	if username == nil {
		return "", nil
	}
	return string(username), nil
}

func (s *Store) SaveCurrentUser(username string) error {
	return s.db.Update(func(tx *bolt.Tx) error {
		bkt := tx.Bucket(defaultsBucket)
		key := []byte("current-user")
		return bkt.Put(key, []byte(username))
	})
}

func itob(v int64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, uint64(v))
	return b
}
<|endoftext|>