<commit_before>// Copyright © 2017 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"github.com/ciao-project/ciao/ciao-controller/types"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List objects",
}

var cnciListCmd = &cobra.Command{
	Use:  "cncis",
	Long: `List CNCIs`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		if !c.IsPrivileged() {
			return errors.New("Listing CNCIs is limited to privileged users")
		}

		cncis, err := c.ListCNCIs()
		if err != nil {
			return errors.Wrap(err, "Error listing CNCIs")
		}

		return render(cmd, cncis.CNCIs)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var eventListCmd = &cobra.Command{
	Use:  "events [TENANT]",
	Long: `List events for the provided tenant. If no tenant is specified and the user is privileged events for all tenants will be returned otherwise returns the current tenants events.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		tenantID := ""
		if len(args) == 1 {
			tenantID = args[0]
		}

		if !c.IsPrivileged() {
			if tenantID == "" {
				tenantID = c.TenantID
			}
		}

		events, err := c.ListEvents(tenantID)
		if err != nil {
			return errors.Wrap(err, "Error listing events")
		}

		return render(cmd, events.Events)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var externalipListCmd = &cobra.Command{
	Use:  "external-ips",
	Long: `List external IP addresses.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		IPs, err := c.ListExternalIPs()
		if err != nil {
			return errors.Wrap(err, "Error listing external IPs")
		}

		return render(cmd, IPs)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ExternalIP" "InternalIP" "InstanceID" "PoolName")}}`},
}

var imageListCmd = &cobra.Command{
	Use:  "images",
	Long: `List images.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		images, err := c.ListImages()
		if err != nil {
			return errors.Wrap(err, "Error getting list of images")
		}

		return render(cmd, images)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var instanceListCmd = &cobra.Command{
	Use:  "instances [WORKLOAD]",
	Long: `List instances. If the optional workload ID is provided then only show instances matching that ID.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		workloadID := ""
		if len(args) == 1 {
			workloadID = args[0]
		}

		servers, err := c.ListInstancesByWorkload(c.TenantID, workloadID)
		if err != nil {
			return errors.Wrap(err, "Error listing instances")
		}

		return render(cmd, servers.Servers)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "Name" "ID" "SSHIP" "SSHPort" "Status") }}`},
}

var nodeListFlags = struct {
	computeNodesOnly bool
	networkNodesOnly bool
}{}

var nodeListCmd = &cobra.Command{
	Use:  "nodes",
	Long: `Lists nodes. Node type can be limited by flags.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		if !c.IsPrivileged() {
			return errors.New("Listing nodes is limited to privileged users")
		}

		var n types.CiaoNodes
		var err error
		if nodeListFlags.computeNodesOnly {
			n, err = c.ListComputeNodes()
		} else if nodeListFlags.networkNodesOnly {
			n, err = c.ListComputeNodes()
		} else {
			n, err = c.ListNodes()
		}

		if err != nil {
			return errors.Wrap(err, "Error getting nodes")
		}

		return render(cmd, n.Nodes)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ID" "Hostname" "Status")}}`},
}

var poolListCmd = &cobra.Command{
	Use:  "pools",
	Long: `List external IP pools.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		p, err := c.ListExternalIPPools()
		if err != nil {
			return errors.Wrap(err, "Error getting external IP pools")
		}

		return render(cmd, p.Pools)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "Name" "Free" "TotalIPs")}}`},
}

var quotasListCmd = &cobra.Command{
	Use:  "quotas [TENANT ID]",
	Long: `List quotas for the current tenant or the supplied tenant if the current user is privileged.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		tenantID := ""
		if len(args) == 1 {
			if !c.IsPrivileged() {
				return errors.New("Listing quotas for other tenants is for privileged users only")
			}
			tenantID = args[0]
		}

		quotas, err := c.ListQuotas(tenantID)
		if err != nil {
			return errors.Wrap(err, "Error getting quotas")
		}

		return render(cmd, quotas)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var tenantListCmd = &cobra.Command{
	Use:  "tenants",
	Long: `List tenants available to the user or if privileged those on the cluster.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		var tenants []types.TenantSummary
		if c.IsPrivileged() {
			t, err := c.ListTenants()
			if err != nil {
				return errors.Wrap(err, "Error listing tenants")
			}
			tenants = t.Tenants
		} else {
			for _, t := range c.Tenants {
				tenants = append(tenants, types.TenantSummary{
					ID: t,
				})
			}
		}

		return render(cmd, tenants)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ID" "Name")}}`},
}

var traceListCmd = &cobra.Command{
	Use:  "traces",
	Long: `List trace labels.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		t, err := c.ListTraceLabels()
		if err != nil {
			return errors.Wrap(err, "Error getting trace labels")
		}

		return render(cmd, t.Summaries)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var volumeListCmd = &cobra.Command{
	Use:  "volumes",
	Long: `List volumes.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		volumes, err := c.ListVolumes()
		if err != nil {
			return errors.Wrap(err, "Error listing volumes")
		}

		return render(cmd, volumes)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

type workload struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	CPUs int    `json:"vcpus"`
	Mem  int    `json:"ram"`
}

var workloadListCmd = &cobra.Command{
	Use:  "workloads",
	Long: `List workloads.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		wls, err := c.ListWorkloads()
		if err != nil {
			return errors.Wrap(err, "Error listing workloads")
		}

		var workloads []workload
		for _, wl := range wls {
			workloads = append(workloads, workload{
				Name: wl.Description,
				ID:   wl.ID,
				Mem:  wl.Requirements.MemMB,
				CPUs: wl.Requirements.VCPUs,
			})
		}

		return render(cmd, workloads)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var listCmds = []*cobra.Command{
	cnciListCmd,
	eventListCmd,
	externalipListCmd,
	imageListCmd,
	instanceListCmd,
	nodeListCmd,
	poolListCmd,
	quotasListCmd,
	tenantListCmd,
	traceListCmd,
	volumeListCmd,
	workloadListCmd,
}

func init() {

	for _, cmd := range listCmds {
		listCmd.AddCommand(cmd)
	}

	nodeListCmd.Flags().BoolVar(&nodeListFlags.computeNodesOnly, "compute-nodes", false, "Only show compute nodes")
	nodeListCmd.Flags().BoolVar(&nodeListFlags.networkNodesOnly, "network-nodes", false, "Only show network nodes")

	rootCmd.AddCommand(listCmd)
}
<commit_msg>ciao: Build "list volumes" on top of "show volume" output<commit_after>// Copyright © 2017 Intel Corporation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"github.com/ciao-project/ciao/ciao-controller/types"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List objects",
}

var cnciListCmd = &cobra.Command{
	Use:  "cncis",
	Long: `List CNCIs`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		if !c.IsPrivileged() {
			return errors.New("Listing CNCIs is limited to privileged users")
		}

		cncis, err := c.ListCNCIs()
		if err != nil {
			return errors.Wrap(err, "Error listing CNCIs")
		}

		return render(cmd, cncis.CNCIs)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var eventListCmd = &cobra.Command{
	Use:  "events [TENANT]",
	Long: `List events for the provided tenant. If no tenant is specified and the user is privileged events for all tenants will be returned otherwise returns the current tenants events.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		tenantID := ""
		if len(args) == 1 {
			tenantID = args[0]
		}

		if !c.IsPrivileged() {
			if tenantID == "" {
				tenantID = c.TenantID
			}
		}

		events, err := c.ListEvents(tenantID)
		if err != nil {
			return errors.Wrap(err, "Error listing events")
		}

		return render(cmd, events.Events)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var externalipListCmd = &cobra.Command{
	Use:  "external-ips",
	Long: `List external IP addresses.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		IPs, err := c.ListExternalIPs()
		if err != nil {
			return errors.Wrap(err, "Error listing external IPs")
		}

		return render(cmd, IPs)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ExternalIP" "InternalIP" "InstanceID" "PoolName")}}`},
}

var imageListCmd = &cobra.Command{
	Use:  "images",
	Long: `List images.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		images, err := c.ListImages()
		if err != nil {
			return errors.Wrap(err, "Error getting list of images")
		}

		return render(cmd, images)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var instanceListCmd = &cobra.Command{
	Use:  "instances [WORKLOAD]",
	Long: `List instances. If the optional workload ID is provided then only show instances matching that ID.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		workloadID := ""
		if len(args) == 1 {
			workloadID = args[0]
		}

		servers, err := c.ListInstancesByWorkload(c.TenantID, workloadID)
		if err != nil {
			return errors.Wrap(err, "Error listing instances")
		}

		return render(cmd, servers.Servers)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "Name" "ID" "SSHIP" "SSHPort" "Status") }}`},
}

var nodeListFlags = struct {
	computeNodesOnly bool
	networkNodesOnly bool
}{}

var nodeListCmd = &cobra.Command{
	Use:  "nodes",
	Long: `Lists nodes. Node type can be limited by flags.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		if !c.IsPrivileged() {
			return errors.New("Listing nodes is limited to privileged users")
		}

		var n types.CiaoNodes
		var err error
		if nodeListFlags.computeNodesOnly {
			n, err = c.ListComputeNodes()
		} else if nodeListFlags.networkNodesOnly {
			n, err = c.ListComputeNodes()
		} else {
			n, err = c.ListNodes()
		}

		if err != nil {
			return errors.Wrap(err, "Error getting nodes")
		}

		return render(cmd, n.Nodes)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ID" "Hostname" "Status")}}`},
}

var poolListCmd = &cobra.Command{
	Use:  "pools",
	Long: `List external IP pools.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		p, err := c.ListExternalIPPools()
		if err != nil {
			return errors.Wrap(err, "Error getting external IP pools")
		}

		return render(cmd, p.Pools)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "Name" "Free" "TotalIPs")}}`},
}

var quotasListCmd = &cobra.Command{
	Use:  "quotas [TENANT ID]",
	Long: `List quotas for the current tenant or the supplied tenant if the current user is privileged.`,
	Args: cobra.MaximumNArgs(1),
	RunE: func(cmd *cobra.Command, args []string) error {
		tenantID := ""
		if len(args) == 1 {
			if !c.IsPrivileged() {
				return errors.New("Listing quotas for other tenants is for privileged users only")
			}
			tenantID = args[0]
		}

		quotas, err := c.ListQuotas(tenantID)
		if err != nil {
			return errors.Wrap(err, "Error getting quotas")
		}

		return render(cmd, quotas)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var tenantListCmd = &cobra.Command{
	Use:  "tenants",
	Long: `List tenants available to the user or if privileged those on the cluster.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		var tenants []types.TenantSummary
		if c.IsPrivileged() {
			t, err := c.ListTenants()
			if err != nil {
				return errors.Wrap(err, "Error listing tenants")
			}
			tenants = t.Tenants
		} else {
			for _, t := range c.Tenants {
				tenants = append(tenants, types.TenantSummary{
					ID: t,
				})
			}
		}

		return render(cmd, tenants)
	},
	Annotations: map[string]string{"default_template": `{{ table (cols . "ID" "Name")}}`},
}

var traceListCmd = &cobra.Command{
	Use:  "traces",
	Long: `List trace labels.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		t, err := c.ListTraceLabels()
		if err != nil {
			return errors.Wrap(err, "Error getting trace labels")
		}

		return render(cmd, t.Summaries)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var volumeListTemplate = `{{ range . }}` + volumeShowTemplate + `
{{ end }}`

var volumeListCmd = &cobra.Command{
	Use:  "volumes",
	Long: `List volumes.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		volumes, err := c.ListVolumes()
		if err != nil {
			return errors.Wrap(err, "Error listing volumes")
		}

		return render(cmd, volumes)
	},
	Annotations: map[string]string{"default_template": volumeListTemplate},
}

type workload struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	CPUs int    `json:"vcpus"`
	Mem  int    `json:"ram"`
}

var workloadListCmd = &cobra.Command{
	Use:  "workloads",
	Long: `List workloads.`,
	Args: cobra.ExactArgs(0),
	RunE: func(cmd *cobra.Command, args []string) error {
		wls, err := c.ListWorkloads()
		if err != nil {
			return errors.Wrap(err, "Error listing workloads")
		}

		var workloads []workload
		for _, wl := range wls {
			workloads = append(workloads, workload{
				Name: wl.Description,
				ID:   wl.ID,
				Mem:  wl.Requirements.MemMB,
				CPUs: wl.Requirements.VCPUs,
			})
		}

		return render(cmd, workloads)
	},
	Annotations: map[string]string{"default_template": "{{ table .}}"},
}

var listCmds = []*cobra.Command{
	cnciListCmd,
	eventListCmd,
	externalipListCmd,
	imageListCmd,
	instanceListCmd,
	nodeListCmd,
	poolListCmd,
	quotasListCmd,
	tenantListCmd,
	traceListCmd,
	volumeListCmd,
	workloadListCmd,
}

func init() {

	for _, cmd := range listCmds {
		listCmd.AddCommand(cmd)
	}

	nodeListCmd.Flags().BoolVar(&nodeListFlags.computeNodesOnly, "compute-nodes", false, "Only show compute nodes")
	nodeListCmd.Flags().BoolVar(&nodeListFlags.networkNodesOnly, "network-nodes", false, "Only show network nodes")

	rootCmd.AddCommand(listCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
	"unicode"
)

type dockerPullStat struct {
	ExitCode int
	Duration time.Duration
}

type dockerRunStat struct {
	ExitCode int
	Duration time.Duration
}

type estafettePipelineStat struct {
	DockerPullStat dockerPullStat
	DockerRunStat  dockerRunStat
}

func (c *estafettePipelineStat) ExitCode() int {
	if c.DockerPullStat.ExitCode > 0 {
		return c.DockerPullStat.ExitCode
	}
	if c.DockerRunStat.ExitCode > 0 {
		return c.DockerRunStat.ExitCode
	}
	return 0
}

func runDockerPull(p estafettePipeline) (stat dockerPullStat, err error) {

	start := time.Now()

	cmd := "docker"

	// add docker command and options
	argsSlice := make([]string, 0)
	argsSlice = append(argsSlice, "pull")
	argsSlice = append(argsSlice, p.ContainerImage)

	fmt.Printf("[estafette] Running command '%v %v'\n", cmd, strings.Join(argsSlice, " "))
	dockerPullCmd := exec.Command(cmd, argsSlice...)

	// run and wait until completion
	if err := dockerPullCmd.Run(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok && status.ExitStatus() > 0 {
				stat.ExitCode = status.ExitStatus()
				return stat, err
			}
		} else {
			return stat, err
		}
	}

	stat.Duration = time.Since(start)

	return
}

func runDockerRun(dir string, envvars map[string]string, p estafettePipeline) (stat dockerRunStat, err error) {

	// run docker with image and commands from yaml
	start := time.Now()

	cmd := "docker"

	// add docker command and options
	argsSlice := make([]string, 0)
	argsSlice = append(argsSlice, "run")
	argsSlice = append(argsSlice, "--privileged=\"true\"")
	argsSlice = append(argsSlice, "--rm=\"true\"")
	argsSlice = append(argsSlice, "--entrypoint=\"\"")
	argsSlice = append(argsSlice, fmt.Sprintf("--volume=\"%v:%v\"", dir, p.WorkingDirectory))
	argsSlice = append(argsSlice, "--volume=\"/var/run/docker.sock:/var/run/docker.sock\"")
	argsSlice = append(argsSlice, fmt.Sprintf("--workdir==\"%v\"", p.WorkingDirectory))
	if envvars != nil && len(envvars) > 0 {
		for k, v := range envvars {
			argsSlice = append(argsSlice, fmt.Sprintf("--env=\"%v=%v\"", k, v))
		}
	}

	// the actual container to run
	argsSlice = append(argsSlice, p.ContainerImage)

	// the commands to execute in the container
	argsSlice = append(argsSlice, p.Shell)
	argsSlice = append(argsSlice, "-c")
	argsSlice = append(argsSlice, strings.Join(p.Commands, ";"))

	fmt.Printf("[estafette] Running command '%v %v'\n", cmd, strings.Join(argsSlice, " "))
	dockerRunCmd := exec.Command(cmd, argsSlice...)

	// pipe logs
	stdout, err := dockerRunCmd.StdoutPipe()
	if err != nil {
		return stat, err
	}
	stderr, err := dockerRunCmd.StderrPipe()
	if err != nil {
		return stat, err
	}

	// start
	if err := dockerRunCmd.Start(); err != nil {
		return stat, err
	}

	// tail logs
	multi := io.MultiReader(stdout, stderr)

	in := bufio.NewScanner(multi)

	for in.Scan() {
		log.Printf("[%v] %s", p.Name, in.Text()) // write each line to your log, or anything you need
	}
	if err := in.Err(); err != nil {
		log.Printf("[%v] Error: %s", p.Name, err)
	}

	// wait for completion
	if err := dockerRunCmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok && status.ExitStatus() > 0 {
				stat.ExitCode = status.ExitStatus()
				return stat, err
			}
		} else {
			return stat, err
		}
	}

	stat.Duration = time.Since(start)

	return
}

// https://gist.github.com/elwinar/14e1e897fdbe4d3432e1
func toUpperSnake(in string) string {
	runes := []rune(in)
	length := len(runes)

	var out []rune
	for i := 0; i < length; i++ {
		if i > 0 && unicode.IsUpper(runes[i]) && ((i+1 < length && unicode.IsLower(runes[i+1])) || unicode.IsLower(runes[i-1])) {
			out = append(out, '_')
		}
		out = append(out, unicode.ToUpper(runes[i]))
	}

	return string(out)
}

func collectEstafetteEnvvars(m estafetteManifest) (envvars map[string]string) {

	envvars = map[string]string{}

	for _, e := range os.Environ() {
		kvPair := strings.Split(e, "=")
		if len(kvPair) == 2 {
			envvarName := kvPair[0]
			envvarValue := kvPair[1]

			if strings.HasPrefix(envvarName, "ESTAFETTE_") {
				envvars[envvarName] = envvarValue
			}
		}
	}

	// add the labels as envvars
	if m.Labels != nil && len(m.Labels) > 0 {
		for key, value := range m.Labels {

			envvarName := "ESTAFETTE_LABEL_" + toUpperSnake(key)
			envvars[envvarName] = value
		}
	}

	return
}

func runPipeline(dir string, envvars map[string]string, p estafettePipeline) (stat estafettePipelineStat, err error) {

	fmt.Printf("[estafette] Starting pipeline '%v'\n", p.Name)

	// pull docker image
	stat.DockerPullStat, err = runDockerPull(p)
	if err != nil {
		return
	}

	stat.DockerRunStat, err = runDockerRun(dir, envvars, p)
	if err != nil {
		return
	}

	fmt.Printf("[estafette] Finished pipeline '%v' successfully\n", p.Name)

	return
}
<commit_msg>removed true from privileged and rm argument<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
	"unicode"
)

type dockerPullStat struct {
	ExitCode int
	Duration time.Duration
}

type dockerRunStat struct {
	ExitCode int
	Duration time.Duration
}

type estafettePipelineStat struct {
	DockerPullStat dockerPullStat
	DockerRunStat  dockerRunStat
}

func (c *estafettePipelineStat) ExitCode() int {
	if c.DockerPullStat.ExitCode > 0 {
		return c.DockerPullStat.ExitCode
	}
	if c.DockerRunStat.ExitCode > 0 {
		return c.DockerRunStat.ExitCode
	}
	return 0
}

func runDockerPull(p estafettePipeline) (stat dockerPullStat, err error) {

	start := time.Now()

	cmd := "docker"

	// add docker command and options
	argsSlice := make([]string, 0)
	argsSlice = append(argsSlice, "pull")
	argsSlice = append(argsSlice, p.ContainerImage)

	fmt.Printf("[estafette] Running command '%v %v'\n", cmd, strings.Join(argsSlice, " "))
	dockerPullCmd := exec.Command(cmd, argsSlice...)

	// run and wait until completion
	if err := dockerPullCmd.Run(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok && status.ExitStatus() > 0 {
				stat.ExitCode = status.ExitStatus()
				return stat, err
			}
		} else {
			return stat, err
		}
	}

	stat.Duration = time.Since(start)

	return
}

func runDockerRun(dir string, envvars map[string]string, p estafettePipeline) (stat dockerRunStat, err error) {

	// run docker with image and commands from yaml
	start := time.Now()

	cmd := "docker"

	// add docker command and options
	argsSlice := make([]string, 0)
	argsSlice = append(argsSlice, "run")
	argsSlice = append(argsSlice, "--privileged")
	argsSlice = append(argsSlice, "--rm")
	argsSlice = append(argsSlice, "--entrypoint=\"\"")
	argsSlice = append(argsSlice, fmt.Sprintf("--volume=\"%v:%v\"", dir, p.WorkingDirectory))
	argsSlice = append(argsSlice, "--volume=\"/var/run/docker.sock:/var/run/docker.sock\"")
	argsSlice = append(argsSlice, fmt.Sprintf("--workdir==\"%v\"", p.WorkingDirectory))
	if envvars != nil && len(envvars) > 0 {
		for k, v := range envvars {
			argsSlice = append(argsSlice, fmt.Sprintf("--env=\"%v=%v\"", k, v))
		}
	}

	// the actual container to run
	argsSlice = append(argsSlice, p.ContainerImage)

	// the commands to execute in the container
	argsSlice = append(argsSlice, p.Shell)
	argsSlice = append(argsSlice, "-c")
	argsSlice = append(argsSlice, strings.Join(p.Commands, ";"))

	fmt.Printf("[estafette] Running command '%v %v'\n", cmd, strings.Join(argsSlice, " "))
	dockerRunCmd := exec.Command(cmd, argsSlice...)

	// pipe logs
	stdout, err := dockerRunCmd.StdoutPipe()
	if err != nil {
		return stat, err
	}
	stderr, err := dockerRunCmd.StderrPipe()
	if err != nil {
		return stat, err
	}

	// start
	if err := dockerRunCmd.Start(); err != nil {
		return stat, err
	}

	// tail logs
	multi := io.MultiReader(stdout, stderr)

	in := bufio.NewScanner(multi)

	for in.Scan() {
		log.Printf("[%v] %s", p.Name, in.Text()) // write each line to your log, or anything you need
	}
	if err := in.Err(); err != nil {
		log.Printf("[%v] Error: %s", p.Name, err)
	}

	// wait for completion
	if err := dockerRunCmd.Wait(); err != nil {
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok && status.ExitStatus() > 0 {
				stat.ExitCode = status.ExitStatus()
				return stat, err
			}
		} else {
			return stat, err
		}
	}

	stat.Duration = time.Since(start)

	return
}

// https://gist.github.com/elwinar/14e1e897fdbe4d3432e1
func toUpperSnake(in string) string {
	runes := []rune(in)
	length := len(runes)

	var out []rune
	for i := 0; i < length; i++ {
		if i > 0 && unicode.IsUpper(runes[i]) && ((i+1 < length && unicode.IsLower(runes[i+1])) || unicode.IsLower(runes[i-1])) {
			out = append(out, '_')
		}
		out = append(out, unicode.ToUpper(runes[i]))
	}

	return string(out)
}

func collectEstafetteEnvvars(m estafetteManifest) (envvars map[string]string) {

	envvars = map[string]string{}

	for _, e := range os.Environ() {
		kvPair := strings.Split(e, "=")
		if len(kvPair) == 2 {
			envvarName := kvPair[0]
			envvarValue := kvPair[1]

			if strings.HasPrefix(envvarName, "ESTAFETTE_") {
				envvars[envvarName] = envvarValue
			}
		}
	}

	// add the labels as envvars
	if m.Labels != nil && len(m.Labels) > 0 {
		for key, value := range m.Labels {

			envvarName := "ESTAFETTE_LABEL_" + toUpperSnake(key)
			envvars[envvarName] = value
		}
	}

	return
}

func runPipeline(dir string, envvars map[string]string, p estafettePipeline) (stat estafettePipelineStat, err error) {

	fmt.Printf("[estafette] Starting pipeline '%v'\n", p.Name)

	// pull docker image
	stat.DockerPullStat, err = runDockerPull(p)
	if err != nil {
		return
	}

	stat.DockerRunStat, err = runDockerRun(dir, envvars, p)
	if err != nil {
		return
	}

	fmt.Printf("[estafette] Finished pipeline '%v' successfully\n", p.Name)

	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc.
//
// 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 pixiecore

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
)

func (s *Server) httpError(w http.ResponseWriter, r *http.Request, status int, format string, args ...interface{}) {
	s.logHTTP(r, format, args...)
	http.Error(w, fmt.Sprintf(format, args...), status)
}

func (s *Server) handleIpxe(w http.ResponseWriter, r *http.Request) {
	args := r.URL.Query()
	mac, err := net.ParseMAC(args.Get("mac"))
	if err != nil {
		s.httpError(w, r, http.StatusBadRequest, "invalid MAC address %q: %s\n", args.Get("mac"), err)
		return
	}

	i, err := strconv.Atoi(args.Get("arch"))
	if err != nil {
		s.httpError(w, r, http.StatusBadRequest, "invalid architecture %q: %s\n", args.Get("arch"), err)
		return

	}
	arch := Architecture(i)
	switch arch {
	case ArchIA32, ArchX64:
	default:
		s.httpError(w, r, http.StatusBadRequest, "Unknown architecture %q\n", arch)
		return
	}

	mach := Machine{
		MAC:  mac,
		Arch: arch,
	}
	spec, err := s.Booter.BootSpec(mach)
	if err != nil {
		// TODO: maybe don't send this error over the network?
		s.logHTTP(r, "error getting bootspec for %#v: %s", mach, err)
		http.Error(w, "couldn't get a bootspec", http.StatusInternalServerError)
		return
	}
	if spec == nil {
		// TODO: make ipxe abort netbooting so it can fall through to
		// other boot options - unsure if that's possible.
		s.httpError(w, r, http.StatusNotFound, "no bootspec found for %q", mach.MAC)
		return
	}
	if spec.Kernel == "" {
		// TODO: maybe don't send this error over the network?
		s.logHTTP(r, "invalid bootspec for %q: missing kernel", mach.MAC)
		http.Error(w, "couldn't get a bootspec", http.StatusInternalServerError)
		return
	}

	// All is well, assemble the iPXE script.

	urlPrefix := fmt.Sprintf("http://%s/_/file?name=", r.Host)

	var b bytes.Buffer
	b.WriteString("#!ipxe\n")
	fmt.Fprintf(&b, "kernel --name kernel %s%s\n", urlPrefix, url.QueryEscape(string(spec.Kernel)))
	for i, initrd := range spec.Initrd {
		fmt.Fprintf(&b, "initrd --name initrd%d %s%s\n", i, urlPrefix, url.QueryEscape(string(initrd)))
	}
	b.WriteString("boot kernel ")
	for i := range spec.Initrd {
		fmt.Fprintf(&b, "initrd=initrd%d ", i)
	}
	for k, v := range spec.Cmdline {
		switch val := v.(type) {
		case string:
			fmt.Fprintf(&b, "%s=%s ", k, val)
		case int, int32, int64, uint32, uint64:
			fmt.Fprintf(&b, "%s=%d ", k, v)
		case ID:
			fmt.Fprintf(&b, "%s=%s%s ", k, urlPrefix, url.QueryEscape(string(val)))
		default:
			s.logHTTP(r, "invalid bootspec for %q: unknown cmdline type for key %q", mach.MAC, k)
			http.Error(w, "couldn't get a bootspec", http.StatusInternalServerError)
			return
		}
	}
	b.WriteByte('\n')
	w.Header().Set("Content-Type", "text/plain")
	b.WriteTo(w)
}

func (s *Server) handleFile(w http.ResponseWriter, r *http.Request) {
	name := r.URL.Query().Get("name")
	f, err := s.Booter.ReadBootFile(ID(name))
	if err != nil {
		s.logHTTP(r, "error getting requested file %q: %s", name, err)
		http.Error(w, "couldn't get file", http.StatusInternalServerError)
		return
	}
	defer f.Close()
	if _, err = io.Copy(w, f); err != nil {
		s.logHTTP(r, "copy of file %q failed: %s", name, err)
	}
}
<commit_msg>pixiecore: factor the ipxe script generation out of the HTTP handler.<commit_after>// Copyright 2016 Google Inc.
//
// 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 pixiecore

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
)

func (s *Server) httpError(w http.ResponseWriter, r *http.Request, status int, format string, args ...interface{}) {
	s.logHTTP(r, format, args...)
	http.Error(w, fmt.Sprintf(format, args...), status)
}

func (s *Server) handleIpxe(w http.ResponseWriter, r *http.Request) {
	args := r.URL.Query()
	mac, err := net.ParseMAC(args.Get("mac"))
	if err != nil {
		s.httpError(w, r, http.StatusBadRequest, "invalid MAC address %q: %s\n", args.Get("mac"), err)
		return
	}

	i, err := strconv.Atoi(args.Get("arch"))
	if err != nil {
		s.httpError(w, r, http.StatusBadRequest, "invalid architecture %q: %s\n", args.Get("arch"), err)
		return

	}
	arch := Architecture(i)
	switch arch {
	case ArchIA32, ArchX64:
	default:
		s.httpError(w, r, http.StatusBadRequest, "Unknown architecture %q\n", arch)
		return
	}

	mach := Machine{
		MAC:  mac,
		Arch: arch,
	}
	spec, err := s.Booter.BootSpec(mach)
	if err != nil {
		// TODO: maybe don't send this error over the network?
		s.logHTTP(r, "error getting bootspec for %#v: %s", mach, err)
		http.Error(w, "couldn't get a bootspec", http.StatusInternalServerError)
		return
	}
	if spec == nil {
		// TODO: make ipxe abort netbooting so it can fall through to
		// other boot options - unsure if that's possible.
		s.httpError(w, r, http.StatusNotFound, "no bootspec found for %q", mach.MAC)
		return
	}
	script, err := ipxeScript(spec, r.Host)
	if err != nil {
		s.logHTTP(r, "failed to assemble ipxe script: %s", err)
		http.Error(w, "couldn't get a bootspec", http.StatusInternalServerError)
	}

	w.Header().Set("Content-Type", "text/plain")
	w.Write(script)
}

func (s *Server) handleFile(w http.ResponseWriter, r *http.Request) {
	name := r.URL.Query().Get("name")
	f, err := s.Booter.ReadBootFile(ID(name))
	if err != nil {
		s.logHTTP(r, "error getting requested file %q: %s", name, err)
		http.Error(w, "couldn't get file", http.StatusInternalServerError)
		return
	}
	defer f.Close()
	if _, err = io.Copy(w, f); err != nil {
		s.logHTTP(r, "copy of file %q failed: %s", name, err)
	}
}

func ipxeScript(spec *Spec, serverHost string) ([]byte, error) {
	if spec.Kernel == "" {
		return nil, errors.New("spec is missing Kernel")
	}

	urlPrefix := fmt.Sprintf("http://%s/_/file?name=", serverHost)
	var b bytes.Buffer
	b.WriteString("#!ipxe\n")
	fmt.Fprintf(&b, "kernel --name kernel %s%s\n", urlPrefix, url.QueryEscape(string(spec.Kernel)))
	for i, initrd := range spec.Initrd {
		fmt.Fprintf(&b, "initrd --name initrd%d %s%s\n", i, urlPrefix, url.QueryEscape(string(initrd)))
	}
	b.WriteString("boot kernel ")
	for i := range spec.Initrd {
		fmt.Fprintf(&b, "initrd=initrd%d ", i)
	}
	for k, v := range spec.Cmdline {
		switch val := v.(type) {
		case string:
			fmt.Fprintf(&b, "%s=%s ", k, val)
		case int, int32, int64, uint32, uint64:
			fmt.Fprintf(&b, "%s=%d ", k, v)
		case ID:
			fmt.Fprintf(&b, "%s=%s%s ", k, urlPrefix, url.QueryEscape(string(val)))
		default:
			return nil, fmt.Errorf("unsupported cmdline type %T for key %q", v, k)
		}
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}
<|endoftext|>


<commit_before>package symlink

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
)

const maxLoopCounter = 100

// FollowSymlink will follow an existing link and scope it to the root
// path provided.
func FollowSymlinkInScope(link, root string) (string, error) {
	root, err := filepath.Abs(root)
	if err != nil {
		return "", err
	}

	link, err = filepath.Abs(link)
	if err != nil {
		return "", err
	}

	if link == root {
		return root, nil
	}

	if !strings.HasPrefix(filepath.Dir(link), root) {
		return "", fmt.Errorf("%s is not within %s", link, root)
	}

	prev := "/"

	for _, p := range strings.Split(link, "/") {
		prev = filepath.Join(prev, p)
		prev = filepath.Clean(prev)

		loopCounter := 0
		for {
			loopCounter++

			if loopCounter >= maxLoopCounter {
				return "", fmt.Errorf("loopCounter reached MAX: %v", loopCounter)
			}

			if !strings.HasPrefix(prev, root) {
				// Don't resolve symlinks outside of root. For example,
				// we don't have to check /home in the below.
				//
				//   /home -> usr/home
				//   FollowSymlinkInScope("/home/bob/foo/bar", "/home/bob/foo")
				break
			}

			stat, err := os.Lstat(prev)
			if err != nil {
				if os.IsNotExist(err) {
					break
				}
				return "", err
			}
			if stat.Mode()&os.ModeSymlink == os.ModeSymlink {
				dest, err := os.Readlink(prev)
				if err != nil {
					return "", err
				}

				if path.IsAbs(dest) {
					prev = filepath.Join(root, dest)
				} else {
					prev, _ = filepath.Abs(prev)

					if prev = filepath.Clean(filepath.Join(filepath.Dir(prev), dest)); len(prev) < len(root) {
						prev = filepath.Join(root, filepath.Base(dest))
					}
				}
			} else {
				break
			}
		}
	}
	return prev, nil
}
<commit_msg>removed redundant Clean<commit_after>package symlink

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
)

const maxLoopCounter = 100

// FollowSymlink will follow an existing link and scope it to the root
// path provided.
func FollowSymlinkInScope(link, root string) (string, error) {
	root, err := filepath.Abs(root)
	if err != nil {
		return "", err
	}

	link, err = filepath.Abs(link)
	if err != nil {
		return "", err
	}

	if link == root {
		return root, nil
	}

	if !strings.HasPrefix(filepath.Dir(link), root) {
		return "", fmt.Errorf("%s is not within %s", link, root)
	}

	prev := "/"

	for _, p := range strings.Split(link, "/") {
		prev = filepath.Join(prev, p)

		loopCounter := 0
		for {
			loopCounter++

			if loopCounter >= maxLoopCounter {
				return "", fmt.Errorf("loopCounter reached MAX: %v", loopCounter)
			}

			if !strings.HasPrefix(prev, root) {
				// Don't resolve symlinks outside of root. For example,
				// we don't have to check /home in the below.
				//
				//   /home -> usr/home
				//   FollowSymlinkInScope("/home/bob/foo/bar", "/home/bob/foo")
				break
			}

			stat, err := os.Lstat(prev)
			if err != nil {
				if os.IsNotExist(err) {
					break
				}
				return "", err
			}
			if stat.Mode()&os.ModeSymlink == os.ModeSymlink {
				dest, err := os.Readlink(prev)
				if err != nil {
					return "", err
				}

				if path.IsAbs(dest) {
					prev = filepath.Join(root, dest)
				} else {
					prev, _ = filepath.Abs(prev)

					if prev = filepath.Join(filepath.Dir(prev), dest); len(prev) < len(root) {
						prev = filepath.Join(root, filepath.Base(dest))
					}
				}
			} else {
				break
			}
		}
	}
	return prev, nil
}
<|endoftext|>


<commit_before>package ui

import (
	"strings"
)

// Styling specifies how to change a Style. It can also be applied to a Segment
// or Text.
type Styling interface{ transform(*Style) }

// StyleText returns a new Text with the given Styling's applied. It does not
// modify the given Text.
func StyleText(t Text, ts ...Styling) Text {
	newt := make(Text, len(t))
	for i, seg := range t {
		newt[i] = StyleSegment(seg, ts...)
	}
	return newt
}

// StyleSegment returns a new Segment with the given Styling's applied. It does
// not modify the given Segment.
func StyleSegment(seg *Segment, ts ...Styling) *Segment {
	return &Segment{Text: seg.Text, Style: ApplyStyling(seg.Style, ts...)}
}

// ApplyStyling returns a new Style with the given Styling's applied.
func ApplyStyling(s Style, ts ...Styling) Style {
	for _, t := range ts {
		if t != nil {
			t.transform(&s)
		}
	}
	return s
}

// Stylings joins several transformers into one.
func Stylings(ts ...Styling) Styling { return jointStyling(ts) }

var (
	FgDefault = setForeground{nil}

	FgBlack   = setForeground{Black}
	FgRed     = setForeground{Red}
	FgGreen   = setForeground{Green}
	FgYellow  = setForeground{Yellow}
	FgBlue    = setForeground{Blue}
	FgMagenta = setForeground{Magenta}
	FgCyan    = setForeground{Cyan}
	FgWhite   = setForeground{White}

	FgBrightBlack   = setForeground{BrightBlack}
	FgBrightRed     = setForeground{BrightRed}
	FgBrightGreen   = setForeground{BrightGreen}
	FgBrightYellow  = setForeground{BrightYellow}
	FgBrightBlue    = setForeground{BrightBlue}
	FgBrightMagenta = setForeground{BrightMagenta}
	FgBrightCyan    = setForeground{BrightCyan}
	FgBrightWhite   = setForeground{BrightWhite}

	BgDefault = setBackground{nil}

	BgBlack   = setBackground{Black}
	BgRed     = setBackground{Red}
	BgGreen   = setBackground{Green}
	BgYellow  = setBackground{Yellow}
	BgBlue    = setBackground{Blue}
	BgMagenta = setBackground{Magenta}
	BgCyan    = setBackground{Cyan}
	BgWhite   = setBackground{White}

	BgBrightBlack   = setBackground{BrightBlack}
	BgBrightRed     = setBackground{BrightRed}
	BgBrightGreen   = setBackground{BrightGreen}
	BgBrightYellow  = setBackground{BrightYellow}
	BgBrightBlue    = setBackground{BrightBlue}
	BgBrightMagenta = setBackground{BrightMagenta}
	BgBrightCyan    = setBackground{BrightCyan}
	BgBrightWhite   = setBackground{BrightWhite}

	Bold       = boolOn(accessBold)
	Dim        = boolOn(accessDim)
	Italic     = boolOn(accessItalic)
	Underlined = boolOn(accessUnderlined)
	Blink      = boolOn(accessBlink)
	Inverse    = boolOn(accessInverse)

	NoBold       = boolOff(accessBold)
	NoDim        = boolOff(accessDim)
	NoItalic     = boolOff(accessItalic)
	NoUnderlined = boolOff(accessUnderlined)
	NoBlink      = boolOff(accessBlink)
	NoInverse    = boolOff(accessInverse)

	ToggleBold       = boolToggle(accessBold)
	ToggleDim        = boolToggle(accessDim)
	ToggleItalic     = boolToggle(accessItalic)
	ToggleUnderlined = boolToggle(accessUnderlined)
	ToggleBlink      = boolToggle(accessBlink)
	ToggleInverse    = boolToggle(accessInverse)
)

type setForeground struct{ c Color }
type setBackground struct{ c Color }
type boolOn func(*Style) *bool
type boolOff func(*Style) *bool
type boolToggle func(*Style) *bool

func (t setForeground) transform(s *Style) { s.Foreground = t.c }
func (t setBackground) transform(s *Style) { s.Background = t.c }
func (t boolOn) transform(s *Style)        { *t(s) = true }
func (t boolOff) transform(s *Style)       { *t(s) = false }
func (t boolToggle) transform(s *Style)    { p := t(s); *p = !*p }

func accessBold(s *Style) *bool       { return &s.Bold }
func accessDim(s *Style) *bool        { return &s.Dim }
func accessItalic(s *Style) *bool     { return &s.Italic }
func accessUnderlined(s *Style) *bool { return &s.Underlined }
func accessBlink(s *Style) *bool      { return &s.Blink }
func accessInverse(s *Style) *bool    { return &s.Inverse }

type jointStyling []Styling

func (t jointStyling) transform(s *Style) {
	for _, t := range t {
		t.transform(s)
	}
}

// ParseStyling parses a text representation of Styling, which are kebab
// case counterparts to the names of the builtin Styling's. For example,
// ToggleInverse is expressed as "toggle-inverse".
//
// Multiple stylings can be joined by spaces, which is equivalent to calling
// Stylings.
//
// If the given string is invalid, ParseStyling returns nil.
func ParseStyling(s string) Styling {
	if !strings.ContainsRune(s, ' ') {
		return parseOneStyling(s)
	}
	var joint jointStyling
	for _, subs := range strings.Split(s, " ") {
		parsed := parseOneStyling(subs)
		if parsed == nil {
			return nil
		}
		joint = append(joint, parseOneStyling(subs))
	}
	return joint
}

var boolFieldAccessor = map[string]func(*Style) *bool{
	"bold":       accessBold,
	"dim":        accessDim,
	"italic":     accessItalic,
	"underlined": accessUnderlined,
	"blink":      accessBlink,
	"inverse":    accessInverse,
}

func parseOneStyling(name string) Styling {
	switch {
	case name == "default" || name == "fg-default":
		return FgDefault
	case strings.HasPrefix(name, "fg-"):
		if color := parseColor(name[len("fg-"):]); color != nil {
			return setForeground{color}
		}
	case name == "bg-default":
		return BgDefault
	case strings.HasPrefix(name, "bg-"):
		if color := parseColor(name[len("bg-"):]); color != nil {
			return setBackground{color}
		}
	case strings.HasPrefix(name, "no-"):
		if f, ok := boolFieldAccessor[name[len("no-"):]]; ok {
			return boolOff(f)
		}
	case strings.HasPrefix(name, "toggle-"):
		if f, ok := boolFieldAccessor[name[len("toggle-"):]]; ok {
			return boolToggle(f)
		}
	default:
		if f, ok := boolFieldAccessor[name]; ok {
			return boolOn(f)
		}
		if color := parseColor(name); color != nil {
			return setForeground{color}
		}
	}
	return nil
}
<commit_msg>pkg/ui: Improve the doc for styling variables.<commit_after>package ui

import (
	"strings"
)

// Styling specifies how to change a Style. It can also be applied to a Segment
// or Text.
type Styling interface{ transform(*Style) }

// StyleText returns a new Text with the given Styling's applied. It does not
// modify the given Text.
func StyleText(t Text, ts ...Styling) Text {
	newt := make(Text, len(t))
	for i, seg := range t {
		newt[i] = StyleSegment(seg, ts...)
	}
	return newt
}

// StyleSegment returns a new Segment with the given Styling's applied. It does
// not modify the given Segment.
func StyleSegment(seg *Segment, ts ...Styling) *Segment {
	return &Segment{Text: seg.Text, Style: ApplyStyling(seg.Style, ts...)}
}

// ApplyStyling returns a new Style with the given Styling's applied.
func ApplyStyling(s Style, ts ...Styling) Style {
	for _, t := range ts {
		if t != nil {
			t.transform(&s)
		}
	}
	return s
}

// Stylings joins several transformers into one.
func Stylings(ts ...Styling) Styling { return jointStyling(ts) }

// Common stylings.
var (
	FgDefault Styling = setForeground{nil}

	FgBlack   = setForeground{Black}
	FgRed     = setForeground{Red}
	FgGreen   = setForeground{Green}
	FgYellow  = setForeground{Yellow}
	FgBlue    = setForeground{Blue}
	FgMagenta = setForeground{Magenta}
	FgCyan    = setForeground{Cyan}
	FgWhite   = setForeground{White}

	FgBrightBlack   = setForeground{BrightBlack}
	FgBrightRed     = setForeground{BrightRed}
	FgBrightGreen   = setForeground{BrightGreen}
	FgBrightYellow  = setForeground{BrightYellow}
	FgBrightBlue    = setForeground{BrightBlue}
	FgBrightMagenta = setForeground{BrightMagenta}
	FgBrightCyan    = setForeground{BrightCyan}
	FgBrightWhite   = setForeground{BrightWhite}

	BgDefault = setBackground{nil}

	BgBlack   = setBackground{Black}
	BgRed     = setBackground{Red}
	BgGreen   = setBackground{Green}
	BgYellow  = setBackground{Yellow}
	BgBlue    = setBackground{Blue}
	BgMagenta = setBackground{Magenta}
	BgCyan    = setBackground{Cyan}
	BgWhite   = setBackground{White}

	BgBrightBlack   = setBackground{BrightBlack}
	BgBrightRed     = setBackground{BrightRed}
	BgBrightGreen   = setBackground{BrightGreen}
	BgBrightYellow  = setBackground{BrightYellow}
	BgBrightBlue    = setBackground{BrightBlue}
	BgBrightMagenta = setBackground{BrightMagenta}
	BgBrightCyan    = setBackground{BrightCyan}
	BgBrightWhite   = setBackground{BrightWhite}

	Bold       = boolOn(accessBold)
	Dim        = boolOn(accessDim)
	Italic     = boolOn(accessItalic)
	Underlined = boolOn(accessUnderlined)
	Blink      = boolOn(accessBlink)
	Inverse    = boolOn(accessInverse)

	NoBold       = boolOff(accessBold)
	NoDim        = boolOff(accessDim)
	NoItalic     = boolOff(accessItalic)
	NoUnderlined = boolOff(accessUnderlined)
	NoBlink      = boolOff(accessBlink)
	NoInverse    = boolOff(accessInverse)

	ToggleBold       = boolToggle(accessBold)
	ToggleDim        = boolToggle(accessDim)
	ToggleItalic     = boolToggle(accessItalic)
	ToggleUnderlined = boolToggle(accessUnderlined)
	ToggleBlink      = boolToggle(accessBlink)
	ToggleInverse    = boolToggle(accessInverse)
)

type setForeground struct{ c Color }
type setBackground struct{ c Color }
type boolOn func(*Style) *bool
type boolOff func(*Style) *bool
type boolToggle func(*Style) *bool

func (t setForeground) transform(s *Style) { s.Foreground = t.c }
func (t setBackground) transform(s *Style) { s.Background = t.c }
func (t boolOn) transform(s *Style)        { *t(s) = true }
func (t boolOff) transform(s *Style)       { *t(s) = false }
func (t boolToggle) transform(s *Style)    { p := t(s); *p = !*p }

func accessBold(s *Style) *bool       { return &s.Bold }
func accessDim(s *Style) *bool        { return &s.Dim }
func accessItalic(s *Style) *bool     { return &s.Italic }
func accessUnderlined(s *Style) *bool { return &s.Underlined }
func accessBlink(s *Style) *bool      { return &s.Blink }
func accessInverse(s *Style) *bool    { return &s.Inverse }

type jointStyling []Styling

func (t jointStyling) transform(s *Style) {
	for _, t := range t {
		t.transform(s)
	}
}

// ParseStyling parses a text representation of Styling, which are kebab
// case counterparts to the names of the builtin Styling's. For example,
// ToggleInverse is expressed as "toggle-inverse".
//
// Multiple stylings can be joined by spaces, which is equivalent to calling
// Stylings.
//
// If the given string is invalid, ParseStyling returns nil.
func ParseStyling(s string) Styling {
	if !strings.ContainsRune(s, ' ') {
		return parseOneStyling(s)
	}
	var joint jointStyling
	for _, subs := range strings.Split(s, " ") {
		parsed := parseOneStyling(subs)
		if parsed == nil {
			return nil
		}
		joint = append(joint, parseOneStyling(subs))
	}
	return joint
}

var boolFieldAccessor = map[string]func(*Style) *bool{
	"bold":       accessBold,
	"dim":        accessDim,
	"italic":     accessItalic,
	"underlined": accessUnderlined,
	"blink":      accessBlink,
	"inverse":    accessInverse,
}

func parseOneStyling(name string) Styling {
	switch {
	case name == "default" || name == "fg-default":
		return FgDefault
	case strings.HasPrefix(name, "fg-"):
		if color := parseColor(name[len("fg-"):]); color != nil {
			return setForeground{color}
		}
	case name == "bg-default":
		return BgDefault
	case strings.HasPrefix(name, "bg-"):
		if color := parseColor(name[len("bg-"):]); color != nil {
			return setBackground{color}
		}
	case strings.HasPrefix(name, "no-"):
		if f, ok := boolFieldAccessor[name[len("no-"):]]; ok {
			return boolOff(f)
		}
	case strings.HasPrefix(name, "toggle-"):
		if f, ok := boolFieldAccessor[name[len("toggle-"):]]; ok {
			return boolToggle(f)
		}
	default:
		if f, ok := boolFieldAccessor[name]; ok {
			return boolOn(f)
		}
		if color := parseColor(name); color != nil {
			return setForeground{color}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package anonymous_messaging

import (
	packet_format "anonymous-messaging/packet_format"
	"fmt"
	"net"
	"os"
	"math/rand"
	"time"
	"anonymous-messaging/publics"
	"anonymous-messaging/pki"
	"github.com/jmoiron/sqlx"
)

const (
	desiredRateParameter = 5
	pathLength = 2
)

type Client struct {
	Id string
	Host string
	Port string
	PubKey int
	PrvKey int
	ActiveMixes []publics.MixPubs

	listener *net.TCPListener
}

type ClientOperations interface {
	EncodeMessage(message string) string
	DecodeMessage(message string) string
}

func (c Client) EncodeMessage(message string, path []publics.MixPubs, delays []float64) packet_format.Packet {
	return packet_format.Encode(message, path, delays)
}

func (c Client) DecodeMessage(packet packet_format.Packet) packet_format.Packet {
	return packet_format.Decode(packet)
}

func (c Client) SendMessage(message string, recipientHost string, recipientPort string) {
	path := c.GetRandomMixSequence(c.ActiveMixes, pathLength)
	path = append(path, publics.MixPubs{Id:c.Id, Host:c.Host, Port:c.Port, PubKey:0})
	delays := c.GenerateDelaySequence(desiredRateParameter, pathLength)
	packet := c.EncodeMessage(message, path, delays)
	c.Send(packet_format.ToString(packet), recipientHost, recipientPort)
}

func (c Client) GenerateDelaySequence(desiredRateParameter float64, length int) []float64{
	rand.Seed(time.Now().UTC().UnixNano())
	var delays []float64
	for i := 0; i < length; i++{
		sample := rand.ExpFloat64() / desiredRateParameter
		delays = append(delays, sample)
	}
	return delays
}

func (c Client) GetRandomMixSequence(data []publics.MixPubs, length int) []publics.MixPubs {
	rand.Seed(time.Now().UTC().UnixNano())
	permutedData := make([]publics.MixPubs, len(data))
	permutation := rand.Perm(len(data))
	for i, v := range permutation {
		permutedData[v] = data[i]
	}
	return permutedData[:length]
}

func (c Client) Send(packet string, host string, port string) {
	conn, err := net.Dial("tcp", host + ":" + port)
	defer conn.Close()

	if err != nil {
		fmt.Print("Error in Client connect", err.Error())
		os.Exit(1)
	}

	conn.Write([]byte(packet))

	buff := make([]byte, 1024)
	n, _ := conn.Read(buff)
	fmt.Println("Received answer: ", string(buff[:n]))
}

func (c Client) listenForConnections() {
	for {
		conn, err := c.listener.Accept()

		if err != nil {
			fmt.Println("Error in connection accepting:", err.Error())
			os.Exit(1)
		}
		fmt.Println(conn)
		//fmt.Println("Received connection from : ", conn.RemoteAddr())
		go c.handleConnection(conn)
	}
}

func (c Client) handleConnection(conn net.Conn) {
	fmt.Println("> Handle Connection")

	buff := make([]byte, 1024)
	reqLen, err := conn.Read(buff)

	if err != nil {
		fmt.Println()
	}

	c.ProcessPacket(packet_format.FromString(string(buff[:reqLen])))
	conn.Close()
}

func (c Client) ProcessPacket(packet packet_format.Packet) string{
	fmt.Println("Processing packet")
	return packet.Message
}

func (c Client) Start() {

	defer c.Run()

	c.ReadInMixnetPKI()
}

func (c Client) Run() {
	fmt.Println("> Client is running")
	defer c.listener.Close()
	finish := make(chan bool)

	go func() {
		fmt.Println("Listening on " + c.Host + ":" + c.Port)
		c.listenForConnections()
	}()

	//c.SendMessage("Hello world", "localhost", "3330")
	<-finish
}

func (c Client) ReadInMixnetPKI() {
	fmt.Println("Reading network")
	db := c.ConnectToPKI()
	records := pki.QueryDatabase(db,"Mixes")

	for records.Next() {
		results := make(map[string]interface{})
		err := records.MapScan(results)
		if err != nil {
			panic(err)
		}
		id := string(results["MixId"].([]byte))
		host := string(results["Host"].([]byte))
		port := string(results["Port"].([]byte))
		pubK := results["PubKey"].(int64)
		pubs := publics.NewMixPubs(id, host, port, pubK)
		c.ActiveMixes = append(c.ActiveMixes, pubs)
	}
	fmt.Println("> The mix network data is uploaded.")
	fmt.Println(c.ActiveMixes)

}

func (c Client) ConnectToPKI() *sqlx.DB{
	db := pki.CreateAndOpenDatabase("./pki/database.db", "./pki/database.db", "sqlite3")
	return db
}

func NewClient(id, host, port string, pubKey, prvKey int) Client{
	c := Client{Id:id, Host:host, Port:port, PubKey:pubKey, PrvKey:prvKey}

	addr, err := net.ResolveTCPAddr("tcp", c.Host + ":" + c.Port)

	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	c.listener, err = net.ListenTCP("tcp", addr)
	return c
}
<commit_msg>Added function for saving clients' pub info in database pki<commit_after>package anonymous_messaging

import (
	packet_format "anonymous-messaging/packet_format"
	"fmt"
	"net"
	"os"
	"math/rand"
	"time"
	"anonymous-messaging/publics"
	"anonymous-messaging/pki"
	"github.com/jmoiron/sqlx"
)

const (
	desiredRateParameter = 5
	pathLength = 2
)

type Client struct {
	Id string
	Host string
	Port string
	PubKey int
	PrvKey int
	ActiveMixes []publics.MixPubs

	listener *net.TCPListener
}

type ClientOperations interface {
	EncodeMessage(message string) string
	DecodeMessage(message string) string
}

func (c Client) EncodeMessage(message string, path []publics.MixPubs, delays []float64) packet_format.Packet {
	return packet_format.Encode(message, path, delays)
}

func (c Client) DecodeMessage(packet packet_format.Packet) packet_format.Packet {
	return packet_format.Decode(packet)
}

func (c Client) SendMessage(message string, recipientHost string, recipientPort string) {
	path := c.GetRandomMixSequence(c.ActiveMixes, pathLength)
	path = append(path, publics.MixPubs{Id:c.Id, Host:c.Host, Port:c.Port, PubKey:0})
	delays := c.GenerateDelaySequence(desiredRateParameter, pathLength)
	packet := c.EncodeMessage(message, path, delays)
	c.Send(packet_format.ToString(packet), recipientHost, recipientPort)
}

func (c Client) GenerateDelaySequence(desiredRateParameter float64, length int) []float64{
	rand.Seed(time.Now().UTC().UnixNano())
	var delays []float64
	for i := 0; i < length; i++{
		sample := rand.ExpFloat64() / desiredRateParameter
		delays = append(delays, sample)
	}
	return delays
}

func (c Client) GetRandomMixSequence(data []publics.MixPubs, length int) []publics.MixPubs {
	rand.Seed(time.Now().UTC().UnixNano())
	permutedData := make([]publics.MixPubs, len(data))
	permutation := rand.Perm(len(data))
	for i, v := range permutation {
		permutedData[v] = data[i]
	}
	return permutedData[:length]
}

func (c Client) Send(packet string, host string, port string) {
	conn, err := net.Dial("tcp", host + ":" + port)
	defer conn.Close()

	if err != nil {
		fmt.Print("Error in Client connect", err.Error())
		os.Exit(1)
	}

	conn.Write([]byte(packet))

	buff := make([]byte, 1024)
	n, _ := conn.Read(buff)
	fmt.Println("Received answer: ", string(buff[:n]))
}

func (c Client) listenForConnections() {
	for {
		conn, err := c.listener.Accept()

		if err != nil {
			fmt.Println("Error in connection accepting:", err.Error())
			os.Exit(1)
		}
		fmt.Println(conn)
		//fmt.Println("Received connection from : ", conn.RemoteAddr())
		go c.handleConnection(conn)
	}
}

func (c Client) handleConnection(conn net.Conn) {
	fmt.Println("> Handle Connection")

	buff := make([]byte, 1024)
	reqLen, err := conn.Read(buff)

	if err != nil {
		fmt.Println()
	}

	c.ProcessPacket(packet_format.FromString(string(buff[:reqLen])))
	conn.Close()
}

func (c Client) ProcessPacket(packet packet_format.Packet) string{
	fmt.Println("Processing packet")
	return packet.Message
}

func (c Client) Start() {

	defer c.Run()

	c.SaveInPKI()
	c.ReadInMixnetPKI()
}

func (c Client) Run() {
	fmt.Println("> Client is running")
	defer c.listener.Close()
	finish := make(chan bool)

	go func() {
		fmt.Println("Listening on " + c.Host + ":" + c.Port)
		c.listenForConnections()
	}()

	//c.SendMessage("Hello world", "localhost", "3330")
	<-finish
}

func (c Client) ReadInMixnetPKI() {
	fmt.Println("Reading network")

	db := c.ConnectToPKI()
	records := pki.QueryDatabase(db,"Mixes")

	for records.Next() {
		results := make(map[string]interface{})
		err := records.MapScan(results)

		if err != nil {
			panic(err)

		}

		pubs := publics.NewMixPubs(string(results["MixId"].([]byte)), string(results["Host"].([]byte)),
									string(results["Port"].([]byte)), results["PubKey"].(int64))

		c.ActiveMixes = append(c.ActiveMixes, pubs)
	}
	fmt.Println("> The mix network data is uploaded.")
}

func (c Client) SaveInPKI() {
	fmt.Println("> Saving Client Public Info into Database")

	db := c.ConnectToPKI()

	pubInfo := make(map[string]interface{})
	pubInfo["MixId"] = c.Id
	pubInfo["Host"] = c.Host
	pubInfo["Port"] = c.Port
	pubInfo["PubKey"] = c.PubKey
	pki.InsertToTable(db, "Clients", pubInfo)

	fmt.Println("> Public info of the mixserver saved in database")
}

func (c Client) ConnectToPKI() *sqlx.DB{
	db := pki.CreateAndOpenDatabase("./pki/database.db", "./pki/database.db", "sqlite3")

	params := make(map[string]string)
	params["ClientId"] = "TEXT"
	params["Host"] = "TEXT"
	params["Port"] = "TEXT"
	params["PubKey"] = "NUM"
	pki.CreateTable(db, "Clients", params)

	return db
}

func NewClient(id, host, port string, pubKey, prvKey int) Client{
	c := Client{Id:id, Host:host, Port:port, PubKey:pubKey, PrvKey:prvKey}

	addr, err := net.ResolveTCPAddr("tcp", c.Host + ":" + c.Port)

	if err != nil {
		fmt.Println("Error listening:", err.Error())
		os.Exit(1)
	}
	c.listener, err = net.ListenTCP("tcp", addr)
	return c
}
<|endoftext|>


<commit_before>// Package client implements a low-level OAuth2 client to perform the various
// request/response flows against a OAuth2 authentication server.
package client

import (
	"net/http"

	"github.com/256dpi/oauth2"
)

// Config is used to configure a client.
type Config struct {
	BaseURI               string
	TokenEndpoint         string
	IntrospectionEndpoint string
	RevocationEndpoint    string
}

// Default will return a default configuration.
func Default(baseURI string) Config {
	return Config{
		BaseURI:               baseURI,
		TokenEndpoint:         "/oauth2/token",
		IntrospectionEndpoint: "/oauth2/introspect",
		RevocationEndpoint:    "/oauth2/revoke",
	}
}

// Client is a low-level OAuth2 client.
type Client struct {
	config Config
	client *http.Client
}

// New will create and return a new client.
func New(config Config) *Client {
	return NewWithClient(config, new(http.Client))
}

// NewWithClient will create and return an new client using the provided client.
func NewWithClient(config Config, client *http.Client) *Client {
	return &Client{
		config: config,
		client: client,
	}
}

// Authenticate will send the provided token request and return the servers
// token response or an error if failed.
func (c *Client) Authenticate(trq oauth2.TokenRequest) (*oauth2.TokenResponse, error) {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.TokenEndpoint

	// build request
	req, err := BuildTokenRequest(endpoint, trq)
	if err != nil {
		return nil, err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	// check status
	if res.StatusCode != http.StatusOK {
		return nil, ParseRequestError(res)
	}

	// parse response
	trs, err := ParseTokenResponse(res)
	if err != nil {
		return nil, err
	}

	return trs, nil
}

// Introspect will send the provided introspection request and return the servers
// response of an error if failed.
func (c *Client) Introspect(irq oauth2.IntrospectionRequest) (*oauth2.IntrospectionResponse, error) {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.IntrospectionEndpoint

	// build request
	req, err := BuildIntrospectionRequest(endpoint, irq)
	if err != nil {
		return nil, err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	// check status
	if res.StatusCode != http.StatusOK {
		return nil, ParseRequestError(res)
	}

	// parse response
	irs, err := ParseIntrospectionResponse(res)
	if err != nil {
		return nil, err
	}

	return irs, nil
}

// Revoke will send the provided revocation request and return and error if it
// failed.
func (c *Client) Revoke(rrq oauth2.RevocationRequest) error {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.RevocationEndpoint

	// build request
	req, err := BuildRevocationRequest(endpoint, rrq)
	if err != nil {
		return err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return err
	}

	// check status
	if res.StatusCode != http.StatusOK {
		return ParseRequestError(res)
	}

	return nil
}
<commit_msg>ensure body is closed<commit_after>// Package client implements a low-level OAuth2 client to perform the various
// request/response flows against a OAuth2 authentication server.
package client

import (
	"net/http"

	"github.com/256dpi/oauth2"
)

// Config is used to configure a client.
type Config struct {
	BaseURI               string
	TokenEndpoint         string
	IntrospectionEndpoint string
	RevocationEndpoint    string
}

// Default will return a default configuration.
func Default(baseURI string) Config {
	return Config{
		BaseURI:               baseURI,
		TokenEndpoint:         "/oauth2/token",
		IntrospectionEndpoint: "/oauth2/introspect",
		RevocationEndpoint:    "/oauth2/revoke",
	}
}

// Client is a low-level OAuth2 client.
type Client struct {
	config Config
	client *http.Client
}

// New will create and return a new client.
func New(config Config) *Client {
	return NewWithClient(config, new(http.Client))
}

// NewWithClient will create and return an new client using the provided client.
func NewWithClient(config Config, client *http.Client) *Client {
	return &Client{
		config: config,
		client: client,
	}
}

// Authenticate will send the provided token request and return the servers
// token response or an error if failed.
func (c *Client) Authenticate(trq oauth2.TokenRequest) (*oauth2.TokenResponse, error) {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.TokenEndpoint

	// build request
	req, err := BuildTokenRequest(endpoint, trq)
	if err != nil {
		return nil, err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	// ensure body is closed
	defer res.Body.Close()

	// check status
	if res.StatusCode != http.StatusOK {
		return nil, ParseRequestError(res)
	}

	// parse response
	trs, err := ParseTokenResponse(res)
	if err != nil {
		return nil, err
	}

	return trs, nil
}

// Introspect will send the provided introspection request and return the servers
// response of an error if failed.
func (c *Client) Introspect(irq oauth2.IntrospectionRequest) (*oauth2.IntrospectionResponse, error) {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.IntrospectionEndpoint

	// build request
	req, err := BuildIntrospectionRequest(endpoint, irq)
	if err != nil {
		return nil, err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}

	// ensure body is closed
	defer res.Body.Close()

	// check status
	if res.StatusCode != http.StatusOK {
		return nil, ParseRequestError(res)
	}

	// parse response
	irs, err := ParseIntrospectionResponse(res)
	if err != nil {
		return nil, err
	}

	return irs, nil
}

// Revoke will send the provided revocation request and return and error if it
// failed.
func (c *Client) Revoke(rrq oauth2.RevocationRequest) error {
	// prepare endpoint
	endpoint := c.config.BaseURI + c.config.RevocationEndpoint

	// build request
	req, err := BuildRevocationRequest(endpoint, rrq)
	if err != nil {
		return err
	}

	// perform request
	res, err := c.client.Do(req)
	if err != nil {
		return err
	}

	// ensure body is closed
	defer res.Body.Close()

	// check status
	if res.StatusCode != http.StatusOK {
		return ParseRequestError(res)
	}

	return nil
}
<|endoftext|>


<commit_before>package backend

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

type S3Backend struct {
	bucket     string
	path       string
	maxRetries int
	svc        s3iface.S3API
}

func NewS3Backend(bucket string, s3path string, maxRetries int, svc s3iface.S3API) *S3Backend {
	return &S3Backend{
		bucket:     bucket,
		path:       strings.TrimPrefix(path.Clean(s3path), "/"),
		maxRetries: maxRetries,
		svc:        svc,
	}
}

func (s *S3Backend) ListDBs() ([]string, error) {
	return s.listDirs(s.path, "")
}

func (s *S3Backend) ListVersions(db, after string, checkForSuccess bool) ([]string, error) {
	versions, err := s.listDirs(path.Join(s.path, db), after)
	if err != nil {
		return nil, err
	}

	if checkForSuccess {
		var filtered []string
		for _, version := range versions {
			successFile := path.Join(s.path, db, version, "_SUCCESS")
			symlinkFile := path.Join(s.path, db, version, "_SYMLINK")

			if s.exists(successFile) || s.exists(symlinkFile) {
				filtered = append(filtered, version)
			}
		}

		versions = filtered
	}

	return versions, nil
}

func (s *S3Backend) listDirs(dir, after string) ([]string, error) {
	// This code assumes you're using S3 like a filesystem, with directories
	// separated by /'s. It also ignores the trailing slash on a prefix (for the
	// purposes of sorting lexicographically), to be consistent with other
	// backends.
	var res []string

	for {
		params := &s3.ListObjectsInput{
			Bucket:    aws.String(s.bucket),
			Delimiter: aws.String("/"),
			Marker:    aws.String(after),
			MaxKeys:   aws.Int64(1000),
			Prefix:    aws.String(dir + "/"),
		}
		resp, err := s.svc.ListObjects(params)

		if err != nil {
			return nil, s.s3error(err)
		} else if resp.CommonPrefixes == nil {
			break
		}

		for _, p := range resp.CommonPrefixes {
			prefix := strings.TrimSuffix(*p.Prefix, "/")

			// List the prefix, to make sure it's a "directory"
			isDir := false
			params := &s3.ListObjectsInput{
				Bucket:    aws.String(s.bucket),
				Delimiter: aws.String(""),
				Marker:    aws.String(after),
				MaxKeys:   aws.Int64(3),
				Prefix:    aws.String(prefix),
			}
			resp, err := s.svc.ListObjects(params)
			if err != nil {
				return nil, err
			}

			for _, key := range resp.Contents {
				if strings.TrimSpace(path.Base(*key.Key)) != "" {
					isDir = true
					break
				}
			}

			if isDir {
				res = append(res, path.Base(prefix))
			}
		}

		if !*resp.IsTruncated || len(resp.CommonPrefixes) == 0 {
			break
		} else {
			after = resp.CommonPrefixes[len(resp.CommonPrefixes)-1].String()
		}
	}

	sort.Strings(res)
	return res, nil
}

func (s *S3Backend) ListFiles(db, version string) ([]string, error) {
	// If this directory is a symlink (denoted with the _SYMLINK file), we
	// should return the contents of the directory it is linking to.
	targetVersion, err := s.getTargetVersion(db, version)
	if err != nil {
		return nil, err
	}
	versionPrefix := path.Join(s.path, db, targetVersion)

	// We use a set here because S3 sometimes returns duplicate keys.
	res := make(map[string]bool)

	params := &s3.ListObjectsInput{
		Bucket:    aws.String(s.bucket),
		Delimiter: aws.String(""),
		MaxKeys:   aws.Int64(1000),
		Prefix:    aws.String(versionPrefix),
	}

	datasetSize := int64(0)
	numFiles := int64(0)

	err = s.svc.ListObjectsPages(params, func(page *s3.ListObjectsOutput, isLastPage bool) bool {
		for _, key := range page.Contents {
			name := path.Base(*key.Key)
			// S3 sometimes has keys that are the same as the "directory"
			if strings.TrimSpace(name) != "" && !strings.HasPrefix(name, "_") && !strings.HasPrefix(name, ".") {
				datasetSize += *key.Size
				numFiles++
				res[name] = true
			}
		}
		return true
	})

	if err != nil {
		return nil, s.s3error(err)
	}

	if version != targetVersion {
		log.Printf("call_site=s3.ListFiles sequins_db=%q sequins_db_version=%q symlinked_version=%q dataset_size=%d file_count=%d", db, version, targetVersion, datasetSize, numFiles)
	} else {
		log.Printf("call_site=s3.ListFiles sequins_db=%q sequins_db_version=%q symlinked_version=\"nil\" dataset_size=%d file_count=%d", db, version, datasetSize, numFiles)
	}

	sorted := make([]string, 0, len(res))
	for name := range res {
		sorted = append(sorted, name)
	}

	sort.Strings(sorted)
	return sorted, nil
}

func (s *S3Backend) open(db, version, file string) (io.ReadCloser, error) {
	src := path.Join(s.path, db, version, file)
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(src),
	}
	resp, err := s.svc.GetObject(params)

	// If the download failed, due to the key not being found, retry
	// maxRetries number of times with an exponential backoff as it may
	// have been due to latency.
	backoff := time.Duration(1)
	for i := 0; i < s.maxRetries && err != nil; i++ {
		aerr, ok := err.(awserr.Error)
		if ok && aerr.Code() == s3.ErrCodeNoSuchKey {
			time.Sleep(backoff * time.Second)
			resp, err = s.svc.GetObject(params)
			backoff *= 2
		} else {
			break
		}
	}

	if err != nil {
		return nil, fmt.Errorf("error opening S3 path %s: %s", s.path, err)
	}

	return resp.Body, nil
}

func (s *S3Backend) Open(db, version, file string) (io.ReadCloser, error) {
	// If this is a symlinked version, we should open the file at the
	// target version directory.
	targetVersion, err := s.getTargetVersion(db, version)
	if err != nil {
		return nil, err
	}

	return s.open(db, targetVersion, file)
}

// getTargetVersion returns the target version of a symlinked version if the
// _SYMLINK file exists or the same version otherwise.
func (s *S3Backend) getTargetVersion(db, version string) (string, error) {
	src := path.Join(s.path, db, version, "_SYMLINK")
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(src),
	}

	resp, err := s.svc.GetObject(params)
	if err != nil {
		return version, nil
	}
	defer resp.Body.Close()

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(b), nil
}

func (s *S3Backend) DisplayPath(parts ...string) string {
	allParts := append([]string{s.path}, parts...)
	return s.displayURL(allParts...)
}

func (s *S3Backend) displayURL(parts ...string) string {
	key := strings.TrimPrefix(path.Join(parts...), "/")
	return fmt.Sprintf("s3://%s/%s", s.bucket, key)
}

func (s *S3Backend) exists(key string) bool {
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	}

	resp, err := s.svc.GetObject(params)
	if err != nil {
		return false
	}

	defer resp.Body.Close()

	return true
}

func (s *S3Backend) s3error(err error) error {
	return fmt.Errorf("unexpected S3 error on bucket %s: %s", s.bucket, err)
}
<commit_msg>strip contents of _SYMLINK file<commit_after>package backend

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

type S3Backend struct {
	bucket     string
	path       string
	maxRetries int
	svc        s3iface.S3API
}

func NewS3Backend(bucket string, s3path string, maxRetries int, svc s3iface.S3API) *S3Backend {
	return &S3Backend{
		bucket:     bucket,
		path:       strings.TrimPrefix(path.Clean(s3path), "/"),
		maxRetries: maxRetries,
		svc:        svc,
	}
}

func (s *S3Backend) ListDBs() ([]string, error) {
	return s.listDirs(s.path, "")
}

func (s *S3Backend) ListVersions(db, after string, checkForSuccess bool) ([]string, error) {
	versions, err := s.listDirs(path.Join(s.path, db), after)
	if err != nil {
		return nil, err
	}

	if checkForSuccess {
		var filtered []string
		for _, version := range versions {
			successFile := path.Join(s.path, db, version, "_SUCCESS")
			symlinkFile := path.Join(s.path, db, version, "_SYMLINK")

			if s.exists(successFile) || s.exists(symlinkFile) {
				filtered = append(filtered, version)
			}
		}

		versions = filtered
	}

	return versions, nil
}

func (s *S3Backend) listDirs(dir, after string) ([]string, error) {
	// This code assumes you're using S3 like a filesystem, with directories
	// separated by /'s. It also ignores the trailing slash on a prefix (for the
	// purposes of sorting lexicographically), to be consistent with other
	// backends.
	var res []string

	for {
		params := &s3.ListObjectsInput{
			Bucket:    aws.String(s.bucket),
			Delimiter: aws.String("/"),
			Marker:    aws.String(after),
			MaxKeys:   aws.Int64(1000),
			Prefix:    aws.String(dir + "/"),
		}
		resp, err := s.svc.ListObjects(params)

		if err != nil {
			return nil, s.s3error(err)
		} else if resp.CommonPrefixes == nil {
			break
		}

		for _, p := range resp.CommonPrefixes {
			prefix := strings.TrimSuffix(*p.Prefix, "/")

			// List the prefix, to make sure it's a "directory"
			isDir := false
			params := &s3.ListObjectsInput{
				Bucket:    aws.String(s.bucket),
				Delimiter: aws.String(""),
				Marker:    aws.String(after),
				MaxKeys:   aws.Int64(3),
				Prefix:    aws.String(prefix),
			}
			resp, err := s.svc.ListObjects(params)
			if err != nil {
				return nil, err
			}

			for _, key := range resp.Contents {
				if strings.TrimSpace(path.Base(*key.Key)) != "" {
					isDir = true
					break
				}
			}

			if isDir {
				res = append(res, path.Base(prefix))
			}
		}

		if !*resp.IsTruncated || len(resp.CommonPrefixes) == 0 {
			break
		} else {
			after = resp.CommonPrefixes[len(resp.CommonPrefixes)-1].String()
		}
	}

	sort.Strings(res)
	return res, nil
}

func (s *S3Backend) ListFiles(db, version string) ([]string, error) {
	// If this directory is a symlink (denoted with the _SYMLINK file), we
	// should return the contents of the directory it is linking to.
	targetVersion, err := s.getTargetVersion(db, version)
	if err != nil {
		return nil, err
	}
	versionPrefix := path.Join(s.path, db, targetVersion)

	// We use a set here because S3 sometimes returns duplicate keys.
	res := make(map[string]bool)

	params := &s3.ListObjectsInput{
		Bucket:    aws.String(s.bucket),
		Delimiter: aws.String(""),
		MaxKeys:   aws.Int64(1000),
		Prefix:    aws.String(versionPrefix),
	}

	datasetSize := int64(0)
	numFiles := int64(0)

	err = s.svc.ListObjectsPages(params, func(page *s3.ListObjectsOutput, isLastPage bool) bool {
		for _, key := range page.Contents {
			name := path.Base(*key.Key)
			// S3 sometimes has keys that are the same as the "directory"
			if strings.TrimSpace(name) != "" && !strings.HasPrefix(name, "_") && !strings.HasPrefix(name, ".") {
				datasetSize += *key.Size
				numFiles++
				res[name] = true
			}
		}
		return true
	})

	if err != nil {
		return nil, s.s3error(err)
	}

	if version != targetVersion {
		log.Printf("call_site=s3.ListFiles sequins_db=%q sequins_db_version=%q symlinked_version=%q dataset_size=%d file_count=%d", db, version, targetVersion, datasetSize, numFiles)
	} else {
		log.Printf("call_site=s3.ListFiles sequins_db=%q sequins_db_version=%q symlinked_version=\"nil\" dataset_size=%d file_count=%d", db, version, datasetSize, numFiles)
	}

	sorted := make([]string, 0, len(res))
	for name := range res {
		sorted = append(sorted, name)
	}

	sort.Strings(sorted)
	return sorted, nil
}

func (s *S3Backend) open(db, version, file string) (io.ReadCloser, error) {
	src := path.Join(s.path, db, version, file)
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(src),
	}
	resp, err := s.svc.GetObject(params)

	// If the download failed, due to the key not being found, retry
	// maxRetries number of times with an exponential backoff as it may
	// have been due to latency.
	backoff := time.Duration(1)
	for i := 0; i < s.maxRetries && err != nil; i++ {
		aerr, ok := err.(awserr.Error)
		if ok && aerr.Code() == s3.ErrCodeNoSuchKey {
			time.Sleep(backoff * time.Second)
			resp, err = s.svc.GetObject(params)
			backoff *= 2
		} else {
			break
		}
	}

	if err != nil {
		return nil, fmt.Errorf("error opening S3 path %s: %s", s.path, err)
	}

	return resp.Body, nil
}

func (s *S3Backend) Open(db, version, file string) (io.ReadCloser, error) {
	// If this is a symlinked version, we should open the file at the
	// target version directory.
	targetVersion, err := s.getTargetVersion(db, version)
	if err != nil {
		return nil, err
	}

	return s.open(db, targetVersion, file)
}

// getTargetVersion returns the target version of a symlinked version if the
// _SYMLINK file exists or the same version otherwise. Whitespace is trimmed
// when reading the file in case it was manually generated.
func (s *S3Backend) getTargetVersion(db, version string) (string, error) {
	src := path.Join(s.path, db, version, "_SYMLINK")
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(src),
	}

	resp, err := s.svc.GetObject(params)
	if err != nil {
		return version, nil
	}
	defer resp.Body.Close()

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(string(b)), nil
}

func (s *S3Backend) DisplayPath(parts ...string) string {
	allParts := append([]string{s.path}, parts...)
	return s.displayURL(allParts...)
}

func (s *S3Backend) displayURL(parts ...string) string {
	key := strings.TrimPrefix(path.Join(parts...), "/")
	return fmt.Sprintf("s3://%s/%s", s.bucket, key)
}

func (s *S3Backend) exists(key string) bool {
	params := &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(key),
	}

	resp, err := s.svc.GetObject(params)
	if err != nil {
		return false
	}

	defer resp.Body.Close()

	return true
}

func (s *S3Backend) s3error(err error) error {
	return fmt.Errorf("unexpected S3 error on bucket %s: %s", s.bucket, err)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 - The TXTDirect 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 txtdirect

import (
	"fmt"
	"net/http/httptest"
	"testing"
)

func Test_zoneFromPath(t *testing.T) {
	tests := []struct {
		url      string
		from     string
		regex    string
		expected string
		err      error
	}{
		{
			"https://example.com/caddy/v1",
			"",
			"",
			"_redirect.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/1/2",
			"",
			"",
			"_redirect.2.1.example.com",
			nil,
		},
		{
			"https://example.com/",
			"",
			"",
			"_redirect.example.com",
			nil,
		},
		{
			"https://example.com/caddy/pkg/v1/download",
			"/$1/$4/$2/$3",
			"",
			"_redirect.pkg.download.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/caddy/pkg/v1",
			"/$1/$3/$2",
			"",
			"_redirect.pkg.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/path-routing",
			"",
			"",
			"_redirect.path-routing.example.com",
			nil,
		},
		{
			"https://example.com/path-routing/test",
			"",
			"",
			"_redirect.test.path-routing.example.com",
			nil,
		},
		{
			"https://example.com/path_routing/test",
			"",
			"",
			"_redirect.test.path_routing.example.com",
			nil,
		},
		{
			"https://example.com/special-chars/#?%!",
			"",
			"",
			"_redirect.special-chars.example.com",
			nil,
		},
		{
			"https://example.com/path2routing/nested/test/",
			"",
			"",
			"_redirect.test.nested.path2routing.example.com",
			nil,
		},
		{
			"https://example.com/123/test",
			"",
			"",
			"_redirect.test.123.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\/([A-Za-z0-9-._~!$'()*+,;=:@]+)",
			"_redirect.12345-some-path.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"(\\d+)",
			"_redirect.12345.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\?query=([^&]*)",
			"_redirect.string.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b>[^&]+)",
			"_redirect.stuff.string.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"",
			"\\?query=(?P<b>[^&]+)\\&more=(?P<a>[^&]+)\\&test=(?P<c>[^&]+)",
			"_redirect.testing.string.stuff.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b2>[^&]+)\\&test=(?P<b1>[^&]+)",
			"_redirect.stuff.testing.string.example.com",
			nil,
		},
		{
			"https://example.com/test",
			"/$2/$1",
			"",
			"",
			fmt.Errorf("length of path doesn't match with length of from= in record"),
		},
	}
	for _, test := range tests {
		rec := record{}
		rec.Re = test.regex
		rec.From = test.from
		req := httptest.NewRequest("GET", test.url, nil)
		zone, _, _, err := zoneFromPath(req, rec)
		if err != nil {
			// Check negative tests
			if test.err != nil {
				if err.Error() == test.err.Error() {
					continue
				}
				t.Errorf("Expected %s, got %s", test.err.Error(), err.Error())
			}
			t.Errorf("Got error: %s", err.Error())
		}
		if zone != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, zone)
		}
	}
}

func Test_zoneFromPathRegexPlaceholders(t *testing.T) {
	tests := []struct {
		url      string
		to       string
		from     string
		regex    string
		expected string
	}{
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"\\/([A-Za-z0-9-._~!$'()*+,;=:@]+)",
			"https://example.com/12345-some-path",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"(\\d+)",
			"https://example.com/12345",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"\\?query=([^&]*)",
			"https://example.com/string",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{a}/{b}",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b>[^&]+)",
			"https://example.com/string/stuff",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"https://example.com/{a}/{b}/{c}",
			"",
			"\\?query=(?P<b>[^&]+)\\&more=(?P<a>[^&]+)\\&test=(?P<c>[^&]+)",
			"https://example.com/stuff/string/testing",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"https://example.com/{a}/{b2}/{b1}",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b2>[^&]+)\\&test=(?P<b1>[^&]+)",
			"https://example.com/string/stuff/testing",
		},
	}
	for _, test := range tests {
		rec := record{
			Re:   test.regex,
			From: test.from,
			To:   test.to,
		}
		req := httptest.NewRequest("GET", test.url, nil)
		_, _, _, err := zoneFromPath(req, rec)
		if err != nil {
			t.Errorf("Unexpected error while parsing path: %s", err.Error())
		}
		to, err := parsePlaceholders(rec.To, req, []string{})
		if err != nil {
			t.Errorf("Unexpected error while parsing placeholders: %s", err.Error())
		}
		if to != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, to)
		}
	}
}
<commit_msg>(test): Use regex placeholder prefix in unit tests<commit_after>/*
Copyright 2019 - The TXTDirect 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 txtdirect

import (
	"fmt"
	"net/http/httptest"
	"testing"
)

func Test_zoneFromPath(t *testing.T) {
	tests := []struct {
		url      string
		from     string
		regex    string
		expected string
		err      error
	}{
		{
			"https://example.com/caddy/v1",
			"",
			"",
			"_redirect.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/1/2",
			"",
			"",
			"_redirect.2.1.example.com",
			nil,
		},
		{
			"https://example.com/",
			"",
			"",
			"_redirect.example.com",
			nil,
		},
		{
			"https://example.com/caddy/pkg/v1/download",
			"/$1/$4/$2/$3",
			"",
			"_redirect.pkg.download.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/caddy/pkg/v1",
			"/$1/$3/$2",
			"",
			"_redirect.pkg.v1.caddy.example.com",
			nil,
		},
		{
			"https://example.com/path-routing",
			"",
			"",
			"_redirect.path-routing.example.com",
			nil,
		},
		{
			"https://example.com/path-routing/test",
			"",
			"",
			"_redirect.test.path-routing.example.com",
			nil,
		},
		{
			"https://example.com/path_routing/test",
			"",
			"",
			"_redirect.test.path_routing.example.com",
			nil,
		},
		{
			"https://example.com/special-chars/#?%!",
			"",
			"",
			"_redirect.special-chars.example.com",
			nil,
		},
		{
			"https://example.com/path2routing/nested/test/",
			"",
			"",
			"_redirect.test.nested.path2routing.example.com",
			nil,
		},
		{
			"https://example.com/123/test",
			"",
			"",
			"_redirect.test.123.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\/([A-Za-z0-9-._~!$'()*+,;=:@]+)",
			"_redirect.12345-some-path.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"(\\d+)",
			"_redirect.12345.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\?query=([^&]*)",
			"_redirect.string.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b>[^&]+)",
			"_redirect.stuff.string.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"",
			"\\?query=(?P<b>[^&]+)\\&more=(?P<a>[^&]+)\\&test=(?P<c>[^&]+)",
			"_redirect.testing.string.stuff.example.com",
			nil,
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b2>[^&]+)\\&test=(?P<b1>[^&]+)",
			"_redirect.stuff.testing.string.example.com",
			nil,
		},
		{
			"https://example.com/test",
			"/$2/$1",
			"",
			"",
			fmt.Errorf("length of path doesn't match with length of from= in record"),
		},
	}
	for _, test := range tests {
		rec := record{}
		rec.Re = test.regex
		rec.From = test.from
		req := httptest.NewRequest("GET", test.url, nil)
		zone, _, _, err := zoneFromPath(req, rec)
		if err != nil {
			// Check negative tests
			if test.err != nil {
				if err.Error() == test.err.Error() {
					continue
				}
				t.Errorf("Expected %s, got %s", test.err.Error(), err.Error())
			}
			t.Errorf("Got error: %s", err.Error())
		}
		if zone != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, zone)
		}
	}
}

func Test_zoneFromPathRegexPlaceholders(t *testing.T) {
	tests := []struct {
		url      string
		to       string
		from     string
		regex    string
		expected string
	}{
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"\\/([A-Za-z0-9-._~!$'()*+,;=:@]+)",
			"https://example.com/12345-some-path",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"(\\d+)",
			"https://example.com/12345",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{1}",
			"",
			"\\?query=([^&]*)",
			"https://example.com/string",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff",
			"https://example.com/{$a}/{$b}",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b>[^&]+)",
			"https://example.com/string/stuff",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"https://example.com/{$a}/{$b}/{$c}",
			"",
			"\\?query=(?P<b>[^&]+)\\&more=(?P<a>[^&]+)\\&test=(?P<c>[^&]+)",
			"https://example.com/stuff/string/testing",
		},
		{
			"https://example.com/12345-some-path?query=string&more=stuff&test=testing",
			"https://example.com/{$a}/{$b2}/{$b1}",
			"",
			"\\?query=(?P<a>[^&]+)\\&more=(?P<b2>[^&]+)\\&test=(?P<b1>[^&]+)",
			"https://example.com/string/stuff/testing",
		},
	}
	for _, test := range tests {
		rec := record{
			Re:   test.regex,
			From: test.from,
			To:   test.to,
		}
		req := httptest.NewRequest("GET", test.url, nil)
		_, _, _, err := zoneFromPath(req, rec)
		if err != nil {
			t.Errorf("Unexpected error while parsing path: %s", err.Error())
		}
		to, err := parsePlaceholders(rec.To, req, []string{})
		if err != nil {
			t.Errorf("Unexpected error while parsing placeholders: %s", err.Error())
		}
		if to != test.expected {
			t.Errorf("Expected %s, got %s", test.expected, to)
		}
	}
}
<|endoftext|>