<commit_before>package app

import (
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/magnusbaeck/logstash-filter-verifier/v2/internal/app/standalone"
	"github.com/magnusbaeck/logstash-filter-verifier/v2/internal/logging"
)

func makeStandaloneCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:           "standalone [<flags>] <testcases> <config>...",
		Short:         "Run logstash-filter-verifier in standalone mode",
		RunE:          runE,
		Args:          validateArgs,
		SilenceErrors: true,
	}

	// TODO: This flag makes sense with daemon mode test command as well.
	cmd.Flags().String("diff-command", "diff -u", "Set the command to run to compare two events. The command will receive the two files to compare as arguments.")
	_ = viper.BindPFlag("diff-command", cmd.Flags().Lookup("diff-command"))

	// TODO: Move default values to some sort of global lookup like defaultKeptEnvVars.
	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("keep-env", []string{"PATH"}, "Add this environment variable to the list of variables that will be preserved from the calling process's environment.")
	_ = viper.BindPFlag("keep-envs", cmd.Flags().Lookup("keep-env"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("logstash-arg", nil, "Command line arguments, which are passed to Logstash. Flag and value have to be provided as a flag each, e.g.: --logstash-arg=-n --logstash-arg=InstanceName")
	_ = viper.BindPFlag("logstash-args", cmd.Flags().Lookup("logstash-arg"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().Bool("logstash-output", false, "Print the debug output of logstash.")
	_ = viper.BindPFlag("logstash-output", cmd.Flags().Lookup("logstash-output"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("logstash-path", nil, "Add a path to the list of Logstash executable paths that will be tried in order (first match is used).")
	_ = viper.BindPFlag("logstash-paths", cmd.Flags().Lookup("logstash-path"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().String("logstash-version", "auto", "The version of Logstash that's being targeted.")
	_ = viper.BindPFlag("logstash-version", cmd.Flags().Lookup("logstash-version"))

	cmd.Flags().Bool("sockets", false, "Use Unix domain sockets for the communication with Logstash.")
	_ = viper.BindPFlag("sockets", cmd.Flags().Lookup("sockets"))

	cmd.Flags().Duration("sockets-timeout", 60*time.Second, "Timeout (duration) for the communication with Logstash via Unix domain sockets. Has no effect unless --sockets is used.")
	_ = viper.BindPFlag("sockets-timeout", cmd.Flags().Lookup("sockets-timeout"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().Bool("quiet", false, "Omit test progress messages and event diffs.")
	_ = viper.BindPFlag("quiet", cmd.Flags().Lookup("quiet"))

	return cmd
}

func runE(_ *cobra.Command, args []string) error {
	s := standalone.New(
		viper.GetBool("quiet"),
		viper.GetString("diff-command"),
		args[0],
		viper.GetStringSlice("keep-envs"),
		viper.GetStringSlice("logstash-paths"),
		viper.GetString("logstash-version"),
		viper.GetStringSlice("logstash-args"),
		viper.GetBool("logstash-output"),
		args[1:],
		viper.GetBool("sockets"),
		viper.GetDuration("sockets-timeout"),
		viper.Get("logger").(logging.Logger),
	)

	return s.Run()
}

func validateArgs(cmd *cobra.Command, args []string) error {
	if len(args) < 1 {
		return errors.New("required argument 'testcases' not provided, try --help")
	}
	if len(args) < 2 {
		return errors.New("required argument 'config' not provided, try --help")
	}
	for _, arg := range args {
		_, err := os.Stat(arg)
		if os.IsNotExist(err) {
			return fmt.Errorf("path '%s' does not exist, try --help", arg)
		}
	}
	return nil
}
<commit_msg>Format quoted argument with %q<commit_after>package app

import (
	"errors"
	"fmt"
	"os"
	"time"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/magnusbaeck/logstash-filter-verifier/v2/internal/app/standalone"
	"github.com/magnusbaeck/logstash-filter-verifier/v2/internal/logging"
)

func makeStandaloneCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:           "standalone [<flags>] <testcases> <config>...",
		Short:         "Run logstash-filter-verifier in standalone mode",
		RunE:          runE,
		Args:          validateArgs,
		SilenceErrors: true,
	}

	// TODO: This flag makes sense with daemon mode test command as well.
	cmd.Flags().String("diff-command", "diff -u", "Set the command to run to compare two events. The command will receive the two files to compare as arguments.")
	_ = viper.BindPFlag("diff-command", cmd.Flags().Lookup("diff-command"))

	// TODO: Move default values to some sort of global lookup like defaultKeptEnvVars.
	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("keep-env", []string{"PATH"}, "Add this environment variable to the list of variables that will be preserved from the calling process's environment.")
	_ = viper.BindPFlag("keep-envs", cmd.Flags().Lookup("keep-env"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("logstash-arg", nil, "Command line arguments, which are passed to Logstash. Flag and value have to be provided as a flag each, e.g.: --logstash-arg=-n --logstash-arg=InstanceName")
	_ = viper.BindPFlag("logstash-args", cmd.Flags().Lookup("logstash-arg"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().Bool("logstash-output", false, "Print the debug output of logstash.")
	_ = viper.BindPFlag("logstash-output", cmd.Flags().Lookup("logstash-output"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().StringSlice("logstash-path", nil, "Add a path to the list of Logstash executable paths that will be tried in order (first match is used).")
	_ = viper.BindPFlag("logstash-paths", cmd.Flags().Lookup("logstash-path"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().String("logstash-version", "auto", "The version of Logstash that's being targeted.")
	_ = viper.BindPFlag("logstash-version", cmd.Flags().Lookup("logstash-version"))

	cmd.Flags().Bool("sockets", false, "Use Unix domain sockets for the communication with Logstash.")
	_ = viper.BindPFlag("sockets", cmd.Flags().Lookup("sockets"))

	cmd.Flags().Duration("sockets-timeout", 60*time.Second, "Timeout (duration) for the communication with Logstash via Unix domain sockets. Has no effect unless --sockets is used.")
	_ = viper.BindPFlag("sockets-timeout", cmd.Flags().Lookup("sockets-timeout"))

	// TODO: Not yet sure, if this should be global or only in standalone.
	cmd.Flags().Bool("quiet", false, "Omit test progress messages and event diffs.")
	_ = viper.BindPFlag("quiet", cmd.Flags().Lookup("quiet"))

	return cmd
}

func runE(_ *cobra.Command, args []string) error {
	s := standalone.New(
		viper.GetBool("quiet"),
		viper.GetString("diff-command"),
		args[0],
		viper.GetStringSlice("keep-envs"),
		viper.GetStringSlice("logstash-paths"),
		viper.GetString("logstash-version"),
		viper.GetStringSlice("logstash-args"),
		viper.GetBool("logstash-output"),
		args[1:],
		viper.GetBool("sockets"),
		viper.GetDuration("sockets-timeout"),
		viper.Get("logger").(logging.Logger),
	)

	return s.Run()
}

func validateArgs(cmd *cobra.Command, args []string) error {
	if len(args) < 1 {
		return errors.New("required argument 'testcases' not provided, try --help")
	}
	if len(args) < 2 {
		return errors.New("required argument 'config' not provided, try --help")
	}
	for _, arg := range args {
		_, err := os.Stat(arg)
		if os.IsNotExist(err) {
			return fmt.Errorf("path %q does not exist, try --help", arg)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/internal/command/arguments"
	"github.com/hashicorp/terraform/internal/command/views"
	"github.com/hashicorp/terraform/internal/states"
	"github.com/hashicorp/terraform/internal/tfdiags"
)

// OutputCommand is a Command implementation that reads an output
// from a Terraform state and prints it.
type OutputCommand struct {
	Meta
}

func (c *OutputCommand) Run(rawArgs []string) int {
	// Parse and apply global view arguments
	common, rawArgs := arguments.ParseView(rawArgs)
	c.View.Configure(common)

	// Parse and validate flags
	args, diags := arguments.ParseOutput(rawArgs)
	if diags.HasErrors() {
		c.View.Diagnostics(diags)
		c.View.HelpPrompt("output")
		return 1
	}

	view := views.NewOutput(args.ViewType, c.View)

	// Fetch data from state
	outputs, diags := c.Outputs(args.StatePath)
	if diags.HasErrors() {
		view.Diagnostics(diags)
		return 1
	}

	// Render the view
	viewDiags := view.Output(args.Name, outputs)
	diags = diags.Append(viewDiags)

	view.Diagnostics(diags)

	if diags.HasErrors() {
		return 1
	}

	return 0
}

func (c *OutputCommand) Outputs(statePath string) (map[string]*states.OutputValue, tfdiags.Diagnostics) {
	var diags tfdiags.Diagnostics

	// Allow state path override
	if statePath != "" {
		c.Meta.statePath = statePath
	}

	// Load the backend
	b, backendDiags := c.Backend(nil)
	diags = diags.Append(backendDiags)
	if diags.HasErrors() {
		return nil, diags
	}

	// This is a read-only command
	c.ignoreRemoteBackendVersionConflict(b)

	env, err := c.Workspace()
	if err != nil {
		diags = diags.Append(fmt.Errorf("Error selecting workspace: %s", err))
		return nil, diags
	}

	// Get the state
	stateStore, err := b.StateMgr(env)
	if err != nil {
		diags = diags.Append(fmt.Errorf("Failed to load state: %s", err))
		return nil, diags
	}

	if err := stateStore.RefreshState(); err != nil {
		diags = diags.Append(fmt.Errorf("Failed to load state: %s", err))
		return nil, diags
	}

	state := stateStore.State()
	if state == nil {
		state = states.NewState()
	}

	return state.RootModule().OutputValues, nil
}

func (c *OutputCommand) Help() string {
	helpText := `
Usage: terraform [global options] output [options] [NAME]

  Reads an output variable from a Terraform state file and prints
  the value. With no additional arguments, output will display all
  the outputs for the root module.  If NAME is not specified, all
  outputs are printed.

Options:

  -state=path      Path to the state file to read. Defaults to
                   "terraform.tfstate".

  -no-color        If specified, output won't contain any color.

  -json            If specified, machine readable output will be
                   printed in JSON format.

  -raw             For value types that can be automatically
                   converted to a string, will print the raw
                   string directly, rather than a human-oriented
                   representation of the value.
`
	return strings.TrimSpace(helpText)
}

func (c *OutputCommand) Synopsis() string {
	return "Show output values from your root module"
}
<commit_msg>Fix CLI help text for output<commit_after>package command

import (
	"fmt"
	"strings"

	"github.com/hashicorp/terraform/internal/command/arguments"
	"github.com/hashicorp/terraform/internal/command/views"
	"github.com/hashicorp/terraform/internal/states"
	"github.com/hashicorp/terraform/internal/tfdiags"
)

// OutputCommand is a Command implementation that reads an output
// from a Terraform state and prints it.
type OutputCommand struct {
	Meta
}

func (c *OutputCommand) Run(rawArgs []string) int {
	// Parse and apply global view arguments
	common, rawArgs := arguments.ParseView(rawArgs)
	c.View.Configure(common)

	// Parse and validate flags
	args, diags := arguments.ParseOutput(rawArgs)
	if diags.HasErrors() {
		c.View.Diagnostics(diags)
		c.View.HelpPrompt("output")
		return 1
	}

	view := views.NewOutput(args.ViewType, c.View)

	// Fetch data from state
	outputs, diags := c.Outputs(args.StatePath)
	if diags.HasErrors() {
		view.Diagnostics(diags)
		return 1
	}

	// Render the view
	viewDiags := view.Output(args.Name, outputs)
	diags = diags.Append(viewDiags)

	view.Diagnostics(diags)

	if diags.HasErrors() {
		return 1
	}

	return 0
}

func (c *OutputCommand) Outputs(statePath string) (map[string]*states.OutputValue, tfdiags.Diagnostics) {
	var diags tfdiags.Diagnostics

	// Allow state path override
	if statePath != "" {
		c.Meta.statePath = statePath
	}

	// Load the backend
	b, backendDiags := c.Backend(nil)
	diags = diags.Append(backendDiags)
	if diags.HasErrors() {
		return nil, diags
	}

	// This is a read-only command
	c.ignoreRemoteBackendVersionConflict(b)

	env, err := c.Workspace()
	if err != nil {
		diags = diags.Append(fmt.Errorf("Error selecting workspace: %s", err))
		return nil, diags
	}

	// Get the state
	stateStore, err := b.StateMgr(env)
	if err != nil {
		diags = diags.Append(fmt.Errorf("Failed to load state: %s", err))
		return nil, diags
	}

	if err := stateStore.RefreshState(); err != nil {
		diags = diags.Append(fmt.Errorf("Failed to load state: %s", err))
		return nil, diags
	}

	state := stateStore.State()
	if state == nil {
		state = states.NewState()
	}

	return state.RootModule().OutputValues, nil
}

func (c *OutputCommand) Help() string {
	helpText := `
Usage: terraform [global options] output [options] [NAME]

  Reads an output variable from a Terraform state file and prints
  the value. With no additional arguments, output will display all
  the outputs for the root module.  If NAME is not specified, all
  outputs are printed.

Options:

  -state=path      Path to the state file to read. Defaults to
                   "terraform.tfstate". Ignored when remote
                   state is used.

  -no-color        If specified, output won't contain any color.

  -json            If specified, machine readable output will be
                   printed in JSON format.

  -raw             For value types that can be automatically
                   converted to a string, will print the raw
                   string directly, rather than a human-oriented
                   representation of the value.
`
	return strings.TrimSpace(helpText)
}

func (c *OutputCommand) Synopsis() string {
	return "Show output values from your root module"
}
<|endoftext|>


<commit_before>package api

import (
	"context"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/datawire/teleproxy/internal/pkg/dns"
	"github.com/datawire/teleproxy/internal/pkg/interceptor"
	"github.com/datawire/teleproxy/internal/pkg/route"
)

type APIServer struct {
	listener net.Listener
	server   http.Server
}

func NewAPIServer(iceptor *interceptor.Interceptor) (*APIServer, error) {
	handler := http.NewServeMux()
	tables := "/api/tables/"
	handler.HandleFunc(tables, func(w http.ResponseWriter, r *http.Request) {
		table := r.URL.Path[len(tables):]

		switch r.Method {
		case http.MethodGet:
			result := iceptor.Render(table)
			if result == "" {
				http.NotFound(w, r)
			} else {
				w.Write(append([]byte(result), '\n'))
			}
		case http.MethodPost:
			d := json.NewDecoder(r.Body)
			var table []route.Table
			err := d.Decode(&table)
			if err != nil {
				http.Error(w, err.Error(), 400)
			} else {
				for _, t := range table {
					iceptor.Update(t)
				}
				dns.Flush()
			}
		case http.MethodDelete:
			iceptor.Delete(table)
		}
	})
	handler.HandleFunc("/api/search", func(w http.ResponseWriter, r *http.Request) {
		var paths []string
		switch r.Method {
		case http.MethodGet:
			paths = iceptor.GetSearchPath()
			result, err := json.Marshal(paths)
			if err != nil {
				panic(err)
			} else {
				w.Write(result)
			}
		case http.MethodPost:
			d := json.NewDecoder(r.Body)
			err := d.Decode(&paths)
			if err != nil {
				http.Error(w, err.Error(), 400)
			} else {
				iceptor.SetSearchPath(paths)
			}
		}
	})
	handler.HandleFunc("/api/shutdown", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Goodbye!\n"))
		p, err := os.FindProcess(os.Getpid())
		if err != nil {
			panic(err)
		}
		p.Signal(os.Interrupt)
	})

	ln, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}

	return &APIServer{
		listener: ln,
		server: http.Server{
			Handler: handler,
		},
	}, nil
}

func (a *APIServer) Port() string {
	_, port, err := net.SplitHostPort(a.listener.Addr().String())
	if err != nil {
		panic(err)
	}
	return port
}

func (a *APIServer) Start() {
	go func() {
		if err := a.server.Serve(a.listener); err != http.ErrServerClosed {
			// Error starting or closing listener:
			log.Printf("API Server: %v", err)
		}
	}()
}

func (a *APIServer) Stop() {
	if err := a.server.Shutdown(context.Background()); err != nil {
		// Error from closing listeners, or context timeout:
		log.Printf("API Server Shutdown: %v", err)
	}
}
<commit_msg>bind api server to 127.0.0.1 explicitly<commit_after>package api

import (
	"context"
	"encoding/json"
	"log"
	"net"
	"net/http"
	"os"

	"github.com/datawire/teleproxy/internal/pkg/dns"
	"github.com/datawire/teleproxy/internal/pkg/interceptor"
	"github.com/datawire/teleproxy/internal/pkg/route"
)

type APIServer struct {
	listener net.Listener
	server   http.Server
}

func NewAPIServer(iceptor *interceptor.Interceptor) (*APIServer, error) {
	handler := http.NewServeMux()
	tables := "/api/tables/"
	handler.HandleFunc(tables, func(w http.ResponseWriter, r *http.Request) {
		table := r.URL.Path[len(tables):]

		switch r.Method {
		case http.MethodGet:
			result := iceptor.Render(table)
			if result == "" {
				http.NotFound(w, r)
			} else {
				w.Write(append([]byte(result), '\n'))
			}
		case http.MethodPost:
			d := json.NewDecoder(r.Body)
			var table []route.Table
			err := d.Decode(&table)
			if err != nil {
				http.Error(w, err.Error(), 400)
			} else {
				for _, t := range table {
					iceptor.Update(t)
				}
				dns.Flush()
			}
		case http.MethodDelete:
			iceptor.Delete(table)
		}
	})
	handler.HandleFunc("/api/search", func(w http.ResponseWriter, r *http.Request) {
		var paths []string
		switch r.Method {
		case http.MethodGet:
			paths = iceptor.GetSearchPath()
			result, err := json.Marshal(paths)
			if err != nil {
				panic(err)
			} else {
				w.Write(result)
			}
		case http.MethodPost:
			d := json.NewDecoder(r.Body)
			err := d.Decode(&paths)
			if err != nil {
				http.Error(w, err.Error(), 400)
			} else {
				iceptor.SetSearchPath(paths)
			}
		}
	})
	handler.HandleFunc("/api/shutdown", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Goodbye!\n"))
		p, err := os.FindProcess(os.Getpid())
		if err != nil {
			panic(err)
		}
		p.Signal(os.Interrupt)
	})

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}

	return &APIServer{
		listener: ln,
		server: http.Server{
			Handler: handler,
		},
	}, nil
}

func (a *APIServer) Port() string {
	_, port, err := net.SplitHostPort(a.listener.Addr().String())
	if err != nil {
		panic(err)
	}
	return port
}

func (a *APIServer) Start() {
	go func() {
		if err := a.server.Serve(a.listener); err != http.ErrServerClosed {
			// Error starting or closing listener:
			log.Printf("API Server: %v", err)
		}
	}()
}

func (a *APIServer) Stop() {
	if err := a.server.Shutdown(context.Background()); err != nil {
		// Error from closing listeners, or context timeout:
		log.Printf("API Server Shutdown: %v", err)
	}
}
<|endoftext|>


<commit_before>package main_test

import (
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"net/http"
	"testing"
)

func TestGetPing_OK(t *testing.T) {

	req, err := http.NewRequest("GET", "http://localhost:7007/mbu/admin/ping", nil)
	res, err := http.DefaultClient.Do(req)

	assert := assert.New(t)
	assert.Nil(err)
	statusCode := res.StatusCode
	assert.Equal(200, statusCode)

	body, err := ioutil.ReadAll(res.Body)
	assert.Nil(err)
	assert.Equal("pong", string(body))
}
<commit_msg>Remove main_test.go file.<commit_after><|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

func TestLbpkrSelfBdist(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "self", "bdist")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running bdist: %v", err)
	}
}

func TestLbpkrSelfBdistRpm(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	if _, err := exec.LookPath("rpmbuild"); err != nil {
		t.Skip("no rpmbuild installed")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "self", "bdist-rpm")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running bdist-rpm: %v", err)
	}
}

func TestLbpkrInstallLbpkr(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "install", "-siteroot="+tmpdir, "lbpkr")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running install: %v", err)
	}
}

func TestRPMSplit(t *testing.T) {
	for _, table := range []struct {
		rpm  string
		want [3]string
	}{
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-10.20.30-1",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "10.20.30", "1"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-1",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", "1"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-71",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", "71"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0-71",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0", "71"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-10.20.30",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "10.20.30", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0-1",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", "1"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-10.20.30-1",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "10.20.30", "1"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0-71",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", "71"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-10.20.30",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "10.20.30", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "", ""},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0-21",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", "21"},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0-1",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", "1"},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", ""},
		},
		{
			rpm:  "BRUNEL_v45r1",
			want: [3]string{"BRUNEL_v45r1", "", ""},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0-21",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", "21"},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0-1",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", "1"},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", ""},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "", ""},
		},
	} {
		rpm := splitRPM(table.rpm)
		if rpm != table.want {
			t.Errorf(
				"%s: error.\nwant=[name=%q version=%q release=%q].\n got=[name=%q version=%q release=%q]\n",
				table.rpm,
				table.want[0], table.want[1], table.want[2],
				rpm[0], rpm[1], rpm[2],
			)
		}
	}
}
<commit_msg>test: mark a few tests as Parallel<commit_after>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

func TestLbpkrSelfBdist(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "self", "bdist")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running bdist: %v", err)
	}
}

func TestLbpkrSelfBdistRpm(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	if _, err := exec.LookPath("rpmbuild"); err != nil {
		t.Skip("no rpmbuild installed")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "self", "bdist-rpm")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running bdist-rpm: %v", err)
	}
}

func TestLbpkrInstallLbpkr(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	tmpdir, err := ioutil.TempDir("", "test-lbpkr-")
	if err != nil {
		t.Fatalf("error creating temporary directory: %v", err)
	}
	defer os.RemoveAll(tmpdir)

	cmd := newCommand("lbpkr", "install", "-siteroot="+tmpdir, "lbpkr")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Dir = tmpdir

	err = cmd.Run()
	if err != nil {
		t.Fatalf("error running install: %v", err)
	}
}

func TestRPMSplit(t *testing.T) {
	t.Parallel()
	for _, table := range []struct {
		rpm  string
		want [3]string
	}{
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-10.20.30-1",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "10.20.30", "1"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-1",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", "1"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0-71",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", "71"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0-71",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0", "71"},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-10.20.30",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "10.20.30", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt-1.0.0",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "1.0.0", ""},
		},
		{
			rpm:  "AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt",
			want: [3]string{"AIDA-3fe9f_3.2.1_x86_64_slc6_gcc49_opt", "", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0-1",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", "1"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-10.20.30-1",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "10.20.30", "1"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0-71",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", "71"},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-10.20.30",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "10.20.30", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt-1.0.0",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "1.0.0", ""},
		},
		{
			rpm:  "LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt",
			want: [3]string{"LCG_67_AIDA_3.2.1_x86_64_slc6_gcc47_opt", "", ""},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0-21",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", "21"},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0-1",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", "1"},
		},
		{
			rpm:  "BRUNEL_v45r1-1.0.0",
			want: [3]string{"BRUNEL_v45r1", "1.0.0", ""},
		},
		{
			rpm:  "BRUNEL_v45r1",
			want: [3]string{"BRUNEL_v45r1", "", ""},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0-21",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", "21"},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0-1",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", "1"},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt-1.0.0",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "1.0.0", ""},
		},
		{
			rpm:  "BRUNEL_v45r1_x86_64_slc6_gcc48_opt",
			want: [3]string{"BRUNEL_v45r1_x86_64_slc6_gcc48_opt", "", ""},
		},
	} {
		rpm := splitRPM(table.rpm)
		if rpm != table.want {
			t.Errorf(
				"%s: error.\nwant=[name=%q version=%q release=%q].\n got=[name=%q version=%q release=%q]\n",
				table.rpm,
				table.want[0], table.want[1], table.want[2],
				rpm[0], rpm[1], rpm[2],
			)
		}
	}
}
<|endoftext|>


<commit_before>package xenserver

import (
	"github.com/mitchellh/multistep"
	"time"
)

type InterruptibleWait struct {
	Predicate         func() (result bool, err error)
	PredicateInterval time.Duration
	Timeout           time.Duration
}

type TimeoutError struct{}

func (err TimeoutError) Error() string {
	return "Timed out"
}

type InterruptedError struct{}

func (err InterruptedError) Error() string {
	return "Interrupted"
}

type PredicateResult struct {
	complete bool
	err      error
}

func (wait InterruptibleWait) Wait(state multistep.StateBag) error {
	predicateResult := make(chan PredicateResult, 1)
	stopWaiting := make(chan struct{})
	defer close(stopWaiting)

	go func() {
		for {
			select {
			case <-time.After(wait.PredicateInterval):
				if complete, err := wait.Predicate(); err != nil || complete {
					predicateResult <- PredicateResult{complete, err}
					return
				}

			case <-stopWaiting:
				return
			}
		}
	}()

	timeout := time.After(wait.Timeout)
	for {
		// wait for either install to complete/error,
		// an interrupt to come through, or a timeout to occur
		select {
		case result := <-predicateResult:
			return result.err

		case <-time.After(1 * time.Second):
			if _, ok := state.GetOk(multistep.StateCancelled); ok {
				return InterruptedError{}
			}

		case <-timeout:
			return TimeoutError{}
		}
	}
}
<commit_msg>interruptible_wait: evaluate Predicate immediately<commit_after>package xenserver

import (
	"github.com/mitchellh/multistep"
	"time"
)

type InterruptibleWait struct {
	Predicate         func() (result bool, err error)
	PredicateInterval time.Duration
	Timeout           time.Duration
}

type TimeoutError struct{}

func (err TimeoutError) Error() string {
	return "Timed out"
}

type InterruptedError struct{}

func (err InterruptedError) Error() string {
	return "Interrupted"
}

type PredicateResult struct {
	complete bool
	err      error
}

func (wait InterruptibleWait) Wait(state multistep.StateBag) error {
	predicateResult := make(chan PredicateResult, 1)
	stopWaiting := make(chan struct{})
	defer close(stopWaiting)

	go func() {
		for {
			if complete, err := wait.Predicate(); err != nil || complete {
				predicateResult <- PredicateResult{complete, err}
				return
			}

			select {
			case <-time.After(wait.PredicateInterval):
				// do nothing; loop again
			case <-stopWaiting:
				return
			}
		}
	}()

	timeout := time.After(wait.Timeout)
	for {
		// wait for either install to complete/error,
		// an interrupt to come through, or a timeout to occur
		select {
		case result := <-predicateResult:
			return result.err

		case <-time.After(1 * time.Second):
			if _, ok := state.GetOk(multistep.StateCancelled); ok {
				return InterruptedError{}
			}

		case <-timeout:
			return TimeoutError{}
		}
	}
}
<|endoftext|>


<commit_before>package main

import "testing"

func TestGetNextRune(t *testing.T) {
	var got, want rune

	got = getNextRune('9')
	want = 'a'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = getNextRune('z')
	want = 'A'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = getNextRune('Z')
	want = '0'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}
}
<commit_msg>add TestAdvanceRunes<commit_after>package main

import (
	"reflect"
	"testing"
)

func TestGetNextRune(t *testing.T) {
	var got, want rune

	got = getNextRune('9')
	want = 'a'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = getNextRune('z')
	want = 'A'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = getNextRune('Z')
	want = '0'
	if got != want {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}
}

func TestAdvanceRunes(t *testing.T) {
	var got, want []rune

	got = advanceRunes([]rune{'Z'})
	want = []rune{'0', '0'}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = advanceRunes([]rune{'0', '0'})
	want = []rune{'0', '1'}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = advanceRunes([]rune{'9', '9'})
	want = []rune{'9', 'a'}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}

	got = advanceRunes([]rune{'Z', 'Z'})
	want = []rune{'0', '0', '0'}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("error: got rune %s, want %s.", string(got), string(want))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/PuerkitoBio/goquery"
	"github.com/avelino/awesome-go/pkg/slug"
)

type Link struct {
	Title       string
	Url         string
	Description string
}

type Object struct {
	Title       string
	Slug        string
	Description string
	Items       []Link
}

func main() {
	err := GenerateHTML()
	if err != nil {
		panic(err)
	}
	input, err := os.ReadFile("./tmpl/index.html")
	if err != nil {
		panic(err)
	}
	buf := bytes.NewBuffer(input)
	query, err := goquery.NewDocumentFromReader(buf)
	if err != nil {
		panic(err)
	}

	objs := make(map[string]*Object)
	query.Find("body #contents").NextFiltered("ul").Find("ul").Each(func(_ int, s *goquery.Selection) {
		s.Find("li a").Each(func(_ int, s *goquery.Selection) {
			selector, exists := s.Attr("href")
			if !exists {
				return
			}
			obj := makeObjById(selector, query.Find("body"))
			if obj == nil {
				return
			}
			objs[selector] = obj
		})
	})

	makeSiteStruct(objs)
	changeLinksInIndex(string(input), query, objs)

	makeSitemap(objs)
}

func makeSiteStruct(objs map[string]*Object) {
	for _, obj := range objs {
		folder := fmt.Sprintf("tmpl/%s", obj.Slug)
		err := os.Mkdir(folder, 0755)
		if err != nil {
			log.Println(err)
		}

		t := template.Must(template.ParseFiles("tmpl/cat-tmpl.html"))
		f, _ := os.Create(fmt.Sprintf("%s/index.html", folder))
		t.Execute(f, obj)
	}
}

func makeSitemap(objs map[string]*Object) {
	t := template.Must(template.ParseFiles("tmpl/sitemap-tmpl.xml"))
	f, _ := os.Create("tmpl/sitemap.xml")
	t.Execute(f, objs)
}

func makeObjById(selector string, s *goquery.Selection) (obj *Object) {
	s.Find(selector).Each(func(_ int, s *goquery.Selection) {
		desc := s.NextFiltered("p")
		ul := s.NextFilteredUntil("ul", "h2")

		links := []Link{}
		ul.Find("li").Each(func(_ int, s *goquery.Selection) {
			url, _ := s.Find("a").Attr("href")
			link := Link{
				Title:       s.Find("a").Text(),
				Description: s.Text(),
				Url:         url,
			}
			links = append(links, link)
		})
		if len(links) == 0 {
			return
		}
		obj = &Object{
			Slug:        slug.Generate(s.Text()),
			Title:       s.Text(),
			Description: desc.Text(),
			Items:       links,
		}
	})
	return
}

func changeLinksInIndex(html string, query *goquery.Document, objs map[string]*Object) {
	query.Find("body #content ul li ul li a").Each(func(_ int, s *goquery.Selection) {
		href, hrefExists := s.Attr("href")
		if !hrefExists {
			return
		}

		// do not replace links if no page has been created for it
		_, objExists := objs[href]
		if !objExists {
			return
		}

		uri := strings.SplitAfter(href, "#")
		if len(uri) >= 2 && uri[1] != "contents" {
			html = strings.ReplaceAll(
				html, fmt.Sprintf(`href="%s"`, href), fmt.Sprintf(`href="%s"`, uri[1]))
		}
	})

	os.WriteFile("./tmpl/index.html", []byte(html), 0644)
}
<commit_msg>fix code style<commit_after>package main

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/PuerkitoBio/goquery"
	"github.com/avelino/awesome-go/pkg/slug"
)

type Link struct {
	Title       string
	Url         string
	Description string
}

type Object struct {
	Title       string
	Slug        string
	Description string
	Items       []Link
}

func main() {
	err := GenerateHTML()
	if err != nil {
		panic(err)
	}
	input, err := os.ReadFile("./tmpl/index.html")
	if err != nil {
		panic(err)
	}
	buf := bytes.NewBuffer(input)
	query, err := goquery.NewDocumentFromReader(buf)
	if err != nil {
		panic(err)
	}

	objs := make(map[string]*Object)
	query.Find("body #contents").NextFiltered("ul").Find("ul").Each(func(_ int, s *goquery.Selection) {
		s.Find("li a").Each(func(_ int, s *goquery.Selection) {
			selector, exists := s.Attr("href")
			if !exists {
				return
			}
			obj := makeObjByID(selector, query.Find("body"))
			if obj == nil {
				return
			}
			objs[selector] = obj
		})
	})

	makeSiteStruct(objs)
	changeLinksInIndex(string(input), query, objs)

	makeSitemap(objs)
}

func makeSiteStruct(objs map[string]*Object) {
	for _, obj := range objs {
		folder := fmt.Sprintf("tmpl/%s", obj.Slug)
		err := os.Mkdir(folder, 0755)
		if err != nil {
			log.Println(err)
		}

		t := template.Must(template.ParseFiles("tmpl/cat-tmpl.html"))
		f, _ := os.Create(fmt.Sprintf("%s/index.html", folder))
		t.Execute(f, obj)
	}
}

func makeSitemap(objs map[string]*Object) {
	t := template.Must(template.ParseFiles("tmpl/sitemap-tmpl.xml"))
	f, _ := os.Create("tmpl/sitemap.xml")
	t.Execute(f, objs)
}

func makeObjByID(selector string, s *goquery.Selection) (obj *Object) {
	s.Find(selector).Each(func(_ int, s *goquery.Selection) {
		desc := s.NextFiltered("p")
		ul := s.NextFilteredUntil("ul", "h2")

		links := []Link{}
		ul.Find("li").Each(func(_ int, s *goquery.Selection) {
			url, _ := s.Find("a").Attr("href")
			link := Link{
				Title:       s.Find("a").Text(),
				Description: s.Text(),
				Url:         url,
			}
			links = append(links, link)
		})
		if len(links) == 0 {
			return
		}
		obj = &Object{
			Slug:        slug.Generate(s.Text()),
			Title:       s.Text(),
			Description: desc.Text(),
			Items:       links,
		}
	})
	return
}

func changeLinksInIndex(html string, query *goquery.Document, objs map[string]*Object) {
	query.Find("body #content ul li ul li a").Each(func(_ int, s *goquery.Selection) {
		href, hrefExists := s.Attr("href")
		if !hrefExists {
			return
		}

		// do not replace links if no page has been created for it
		_, objExists := objs[href]
		if !objExists {
			return
		}

		uri := strings.SplitAfter(href, "#")
		if len(uri) >= 2 && uri[1] != "contents" {
			html = strings.ReplaceAll(
				html, fmt.Sprintf(`href="%s"`, href), fmt.Sprintf(`href="%s"`, uri[1]))
		}
	})

	os.WriteFile("./tmpl/index.html", []byte(html), 0644)
}
<|endoftext|>


<commit_before>package capn_test

import (
	"bytes"
	"fmt"
	"os"
	"testing"

	capn "github.com/glycerine/go-capnproto"
	air "github.com/glycerine/go-capnproto/aircraftlib"
	cv "github.com/smartystreets/goconvey/convey"
)

func ExampleAirplaneWrite() string {

	fname := "out.write_test.airplane.cpz"

	// make a brand new, empty segment (message)
	seg := capn.NewBuffer(nil)

	// If you want runtime-type identification, this is easily obtained. Just
	// wrap everything in a struct that contains a single anoymous union (e.g. struct Z).
	// Then always set a Z as the root object in you message/first segment.
	// The cost of the extra word of storage is usually worth it, as
	// then human readalbe output is easily obtained via a shell command such as
	//
	// $ cat binary.cpz | capnp decode aircraft.capnp Z
	//
	// If you need to conserve space, and known your content in advance, it
	// isn't necessary to use an anonymous union. Just supply the type name
	// in place of 'Z' in the decode command above.

	z := air.NewRootZ(seg) // root should be allocated first.
	// There must be only one root.

	// then non-root objects:
	aircraft := air.NewAircraft(seg)
	b737 := air.NewB737(seg)
	planebase := air.NewPlaneBase(seg)

	// how to create a list. Requires a cast at the moment.
	homes := air.NewAirportList(seg, 2)
	uint16list := capn.UInt16List(homes) // cast to the underlying type
	uint16list.Set(0, air.AIRPORT_JFK)
	uint16list.Set(1, air.AIRPORT_LAX)

	// set the primitive fields
	planebase.SetCanFly(true)
	planebase.SetName("Henrietta")
	planebase.SetRating(100)

	// set a list field
	planebase.SetHomes(homes)

	// wire up the pointers between objects
	b737.SetBase(planebase)
	aircraft.SetB737(b737)
	z.SetAircraft(aircraft)

	// ready to write

	// example of writing to memory
	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	// example of writing to file. Just use WriteTo().
	// We could have used SegToFile(seg, fname) from
	// util_test.go intead, but this makes it clear how easy it is.
	file, err := os.Create(fname)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	seg.WriteTo(file)

	// readback and view that file in human readable format. Defined in util_test.go
	text, err := CapnFileToText(fname, "aircraftlib/aircraft.capnp", "")
	if err != nil {
		panic(err)
	}
	fmt.Printf("here is our aircraft:\n")
	fmt.Printf("%s\n", text)

	return text
}

func TestAircraftWrite(t *testing.T) {

	observedText := ExampleAirplaneWrite()
	expectedText := `(aircraft = (b737 = (base = (name = "Henrietta", homes = [jfk, lax], rating = 100, canFly = true, capacity = 0, maxSpeed = 0))))`

	cv.Convey("When we run the ExampleAirplaneWrite() function in write_test.go", t, func() {
		cv.Convey("Then we should see the human readable B737 example struct we expect", func() {
			cv.So(observedText, cv.ShouldEqual, expectedText)
		})
	})

}
<commit_msg>typo fix<commit_after>package capn_test

import (
	"bytes"
	"fmt"
	"os"
	"testing"

	capn "github.com/glycerine/go-capnproto"
	air "github.com/glycerine/go-capnproto/aircraftlib"
	cv "github.com/smartystreets/goconvey/convey"
)

func ExampleAirplaneWrite() string {

	fname := "out.write_test.airplane.cpz"

	// make a brand new, empty segment (message)
	seg := capn.NewBuffer(nil)

	// If you want runtime-type identification, this is easily obtained. Just
	// wrap everything in a struct that contains a single anoymous union (e.g. struct Z).
	// Then always set a Z as the root object in you message/first segment.
	// The cost of the extra word of storage is usually worth it, as
	// then human readalbe output is easily obtained via a shell command such as
	//
	// $ cat binary.cpz | capnp decode aircraft.capnp Z
	//
	// If you need to conserve space, and know your content in advance, it
	// isn't necessary to use an anonymous union. Just supply the type name
	// in place of 'Z' in the decode command above.

	z := air.NewRootZ(seg) // root should be allocated first.
	// There must be only one root.

	// then non-root objects:
	aircraft := air.NewAircraft(seg)
	b737 := air.NewB737(seg)
	planebase := air.NewPlaneBase(seg)

	// how to create a list. Requires a cast at the moment.
	homes := air.NewAirportList(seg, 2)
	uint16list := capn.UInt16List(homes) // cast to the underlying type
	uint16list.Set(0, air.AIRPORT_JFK)
	uint16list.Set(1, air.AIRPORT_LAX)

	// set the primitive fields
	planebase.SetCanFly(true)
	planebase.SetName("Henrietta")
	planebase.SetRating(100)

	// set a list field
	planebase.SetHomes(homes)

	// wire up the pointers between objects
	b737.SetBase(planebase)
	aircraft.SetB737(b737)
	z.SetAircraft(aircraft)

	// ready to write

	// example of writing to memory
	buf := bytes.Buffer{}
	seg.WriteTo(&buf)

	// example of writing to file. Just use WriteTo().
	// We could have used SegToFile(seg, fname) from
	// util_test.go intead, but this makes it clear how easy it is.
	file, err := os.Create(fname)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	seg.WriteTo(file)

	// readback and view that file in human readable format. Defined in util_test.go
	text, err := CapnFileToText(fname, "aircraftlib/aircraft.capnp", "")
	if err != nil {
		panic(err)
	}
	fmt.Printf("here is our aircraft:\n")
	fmt.Printf("%s\n", text)

	return text
}

func TestAircraftWrite(t *testing.T) {

	observedText := ExampleAirplaneWrite()
	expectedText := `(aircraft = (b737 = (base = (name = "Henrietta", homes = [jfk, lax], rating = 100, canFly = true, capacity = 0, maxSpeed = 0))))`

	cv.Convey("When we run the ExampleAirplaneWrite() function in write_test.go", t, func() {
		cv.Convey("Then we should see the human readable B737 example struct we expect", func() {
			cv.So(observedText, cv.ShouldEqual, expectedText)
		})
	})

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
)

func parseNumber(s string) float64 {
	var value float64
	fmt.Sscanf(s, "%f", &value)
	return value
}

const (
	Lakh  = 100000.0
	Crore = 10000000.0
)

type InrAmount struct {
	Value float64
}

const (
	Million  = 1000000.0
	Billion  = 1000000000.0
	Trillion = 1000000000000.0
)

type UsdAmount struct {
	Value float64
}

type Amount interface {
	Convert(usdToInr float64) Amount
	FormatValue() string
}

func NewInrAmount(s string) *InrAmount {
	units := regexp.MustCompile(`lakh|crore`)
	unit := units.FindString(s)
	number := parseNumber(s)

	switch unit {
	case "lakh":
		return &InrAmount{Value: number * Lakh}
	case "crore":
		return &InrAmount{Value: number * Crore}
	default:
		return &InrAmount{Value: number}
	}
}

func (amount *InrAmount) Convert(usdToInr float64) Amount {
	return &UsdAmount{Value: amount.Value / usdToInr}
}

func (amount *InrAmount) FormatValue() string {
	if v := amount.Value / Crore; v >= 1.0 {
		return fmt.Sprintf("₹ %.1f crore", v)
	} else if v := amount.Value / Lakh; v >= 1.0 {
		return fmt.Sprintf("₹ %.1f lakh", v)
	} else {
		return fmt.Sprintf("₹ %.1f", amount.Value)
	}
}

func NewUsdAmount(s string) *UsdAmount {
	units := regexp.MustCompile("million|billion|trillion")
	unit := units.FindString(s)
	number := parseNumber(s)

	switch unit {
	case "million":
		return &UsdAmount{Value: number * Million}
	case "billion":
		return &UsdAmount{Value: number * Billion}
	case "trillion":
		return &UsdAmount{Value: number * Trillion}
	default:
		return &UsdAmount{Value: number}
	}
}

func (amount *UsdAmount) Convert(usdToInr float64) Amount {
	return &InrAmount{Value: amount.Value * usdToInr}
}

func (amount *UsdAmount) FormatValue() string {
	if v := amount.Value / Trillion; v >= 1.0 {
		return fmt.Sprintf("$ %.1f trillion", v)
	} else if v := amount.Value / Billion; v >= 1.0 {
		return fmt.Sprintf("$ %.1f billion", v)
	} else if v := amount.Value / Million; v >= 1.0 {
		return fmt.Sprintf("$ %.1f million", v)
	} else {
		return fmt.Sprintf("$ %.1f", amount.Value)
	}
}

var inrSignifiers = regexp.MustCompile(`lakh|crore|rs|inr|₹|rupee`)
var usdSignifiers = regexp.MustCompile(`million|billion|trillion|\$|usd|dollar`)

func parseAmount(s string) (Amount, error) {
	if inrSignifiers.MatchString(s) {
		return NewInrAmount(s), nil
	} else if usdSignifiers.MatchString(s) {
		return NewUsdAmount(s), nil
	} else {
		return nil, errors.New("no currency recognized")
	}
}

type FixerResponse struct {
	Rates map[string]float64 `json:"rates"`
}

func getUsdToInr() float64 {
	defaultRate := 62.0

	r, err := http.Get("http://api.fixer.io/latest?base=USD")
	if err != nil {
		return defaultRate
	}
	defer r.Body.Close()

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return defaultRate
	}

	var response FixerResponse
	if err = json.Unmarshal(body, &response); err != nil {
		return defaultRate
	}

	return response.Rates["INR"]
}

func main() {
	input := strings.Join(os.Args[1:], " ")
	amount, err := parseAmount(input)
	if err != nil {
		fmt.Println(err.Error())
	}

	fmt.Println(amount.Convert(getUsdToInr()).FormatValue())
}
<commit_msg>Better error handling<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"
)

func parseNumber(s string) (float64, error) {
	var value float64
	n, err := fmt.Sscanf(s, "%f", &value)

	if n != 1 || err != nil {
		return 0, errors.New("invalid number")
	} else {
		return value, nil
	}
}

const (
	Lakh  = 100000.0
	Crore = 10000000.0
)

type InrAmount struct {
	Value float64
}

const (
	Million  = 1000000.0
	Billion  = 1000000000.0
	Trillion = 1000000000000.0
)

type UsdAmount struct {
	Value float64
}

type Amount interface {
	Convert(usdToInr float64) Amount
	FormatValue() string
}

func NewInrAmount(s string) (*InrAmount, error) {
	units := regexp.MustCompile(`lakh|crore`)
	unit := units.FindString(s)
	number, err := parseNumber(s)

	if err != nil {
		return nil, err
	}

	switch unit {
	case "lakh":
		return &InrAmount{Value: number * Lakh}, nil
	case "crore":
		return &InrAmount{Value: number * Crore}, nil
	default:
		return &InrAmount{Value: number}, nil
	}
}

func (amount *InrAmount) Convert(usdToInr float64) Amount {
	return &UsdAmount{Value: amount.Value / usdToInr}
}

func (amount *InrAmount) FormatValue() string {
	if v := amount.Value / Crore; v >= 1.0 {
		return fmt.Sprintf("₹ %.1f crore", v)
	} else if v := amount.Value / Lakh; v >= 1.0 {
		return fmt.Sprintf("₹ %.1f lakh", v)
	} else {
		return fmt.Sprintf("₹ %.1f", amount.Value)
	}
}

func NewUsdAmount(s string) (*UsdAmount, error) {
	units := regexp.MustCompile("million|billion|trillion")
	unit := units.FindString(s)
	number, err := parseNumber(s)

	if err != nil {
		return nil, err
	}

	switch unit {
	case "million":
		return &UsdAmount{Value: number * Million}, nil
	case "billion":
		return &UsdAmount{Value: number * Billion}, nil
	case "trillion":
		return &UsdAmount{Value: number * Trillion}, nil
	default:
		return &UsdAmount{Value: number}, nil
	}
}

func (amount *UsdAmount) Convert(usdToInr float64) Amount {
	return &InrAmount{Value: amount.Value * usdToInr}
}

func (amount *UsdAmount) FormatValue() string {
	if v := amount.Value / Trillion; v >= 1.0 {
		return fmt.Sprintf("$ %.1f trillion", v)
	} else if v := amount.Value / Billion; v >= 1.0 {
		return fmt.Sprintf("$ %.1f billion", v)
	} else if v := amount.Value / Million; v >= 1.0 {
		return fmt.Sprintf("$ %.1f million", v)
	} else {
		return fmt.Sprintf("$ %.1f", amount.Value)
	}
}

var inrSignifiers = regexp.MustCompile(`lakh|crore|rs|inr|₹|rupee`)
var usdSignifiers = regexp.MustCompile(`million|billion|trillion|\$|usd|dollar`)

func parseAmount(s string) (Amount, error) {
	if inrSignifiers.MatchString(s) {
		return NewInrAmount(s)
	} else if usdSignifiers.MatchString(s) {
		return NewUsdAmount(s)
	} else {
		return nil, errors.New("no currency recognized")
	}
}

type FixerResponse struct {
	Rates map[string]float64 `json:"rates"`
}

func getUsdToInr() float64 {
	defaultRate := 62.0

	r, err := http.Get("http://api.fixer.io/latest?base=USD")
	if err != nil {
		return defaultRate
	}
	defer r.Body.Close()

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return defaultRate
	}

	var response FixerResponse
	if err = json.Unmarshal(body, &response); err != nil {
		return defaultRate
	}

	return response.Rates["INR"]
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("no input given")
		os.Exit(1)
	}

	input := strings.Join(os.Args[1:], " ")
	amount, err := parseAmount(input)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}

	fmt.Println(amount.Convert(getUsdToInr()).FormatValue())
}
<|endoftext|>


<commit_before>// Package mongo is a REST Layer resource storage handler for MongoDB using mgo
package mongo

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/rs/rest-layer/resource"
	"github.com/rs/rest-layer/schema/query"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// mongoItem is a bson representation of a resource.Item.
type mongoItem struct {
	ID      interface{}            `bson:"_id"`
	ETag    string                 `bson:"_etag"`
	Updated time.Time              `bson:"_updated"`
	Payload map[string]interface{} `bson:",inline"`
}

// newMongoItem converts a resource.Item into a mongoItem.
func newMongoItem(i *resource.Item) *mongoItem {
	// Filter out id from the payload so we don't store it twice
	p := map[string]interface{}{}
	for k, v := range i.Payload {
		if k != "id" {
			p[k] = v
		}
	}
	return &mongoItem{
		ID:      i.ID,
		ETag:    i.ETag,
		Updated: i.Updated,
		Payload: p,
	}
}

// newItem converts a back mongoItem into a resource.Item.
func newItem(i *mongoItem) *resource.Item {
	// Add the id back (we use the same map hoping the mongoItem won't be stored back)
	i.Payload["id"] = i.ID
	item := &resource.Item{
		ID:      i.ID,
		ETag:    i.ETag,
		Updated: i.Updated,
		Payload: i.Payload,
	}

	if item.ETag == "" {
		if v, ok := i.ID.(bson.ObjectId); ok {
			item.ETag = "p-" + v.Hex()
		} else {
			item.ETag = "p-" + fmt.Sprint(i.ID)
		}
	}
	return item
}

// Handler handles resource storage in a MongoDB collection.
type Handler func(ctx context.Context) (*mgo.Collection, error)

// NewHandler creates an new mongo handler
func NewHandler(s *mgo.Session, db, collection string) Handler {
	c := func() *mgo.Collection {
		return s.DB(db).C(collection)
	}
	return func(ctx context.Context) (*mgo.Collection, error) {
		return c(), nil
	}
}

// C returns the mongo collection managed by this storage handler
// from a Copy() of the mgo session.
func (m Handler) c(ctx context.Context) (*mgo.Collection, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	c, err := m(ctx)
	if err != nil {
		return nil, err
	}
	// With mgo, session.Copy() pulls a connection from the connection pool
	s := c.Database.Session.Copy()
	// Ensure safe mode is enabled in order to get errors
	s.EnsureSafe(&mgo.Safe{})
	// Set a timeout to match the context deadline if any
	if deadline, ok := ctx.Deadline(); ok {
		timeout := deadline.Sub(time.Now())
		if timeout <= 0 {
			timeout = 0
		}
		s.SetSocketTimeout(timeout)
		s.SetSyncTimeout(timeout)
	}
	c.Database.Session = s
	return c, nil
}

// close returns a mgo.Collection's session to the connection pool.
func (m Handler) close(c *mgo.Collection) {
	c.Database.Session.Close()
}

// Insert inserts new items in the mongo collection.
func (m Handler) Insert(ctx context.Context, items []*resource.Item) error {
	mItems := make([]interface{}, len(items))
	for i, item := range items {
		mItems[i] = newMongoItem(item)
	}
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	err = c.Insert(mItems...)
	if mgo.IsDup(err) {
		// Duplicate ID key
		err = resource.ErrConflict
	}
	if ctx.Err() != nil {
		return ctx.Err()
	}
	return err
}

// Update replace an item by a new one in the mongo collection.
func (m Handler) Update(ctx context.Context, item *resource.Item, original *resource.Item) error {
	mItem := newMongoItem(item)
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	s := bson.M{"_id": original.ID}
	if strings.HasPrefix(original.ETag, "p-") {
		// If the original ETag is in "p-[id]" format,
		// then _etag field must be absent from the resource in DB
		s["_etag"] = bson.M{"$exists": false}
	} else {
		s["_etag"] = original.ETag
	}
	err = c.Update(s, mItem)
	if err == mgo.ErrNotFound {
		// Determine if the item is not found or if the item is found but etag missmatch
		var count int
		count, err = c.FindId(original.ID).Count()
		if err != nil {
			// The find returned an unexpected err, just forward it with no mapping
		} else if count == 0 {
			err = resource.ErrNotFound
		} else if ctx.Err() != nil {
			err = ctx.Err()
		} else {
			// If the item were found, it means that its etag didn't match
			err = resource.ErrConflict
		}
	}
	return err
}

// Delete deletes an item from the mongo collection.
func (m Handler) Delete(ctx context.Context, item *resource.Item) error {
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	s := bson.M{"_id": item.ID}
	if strings.HasPrefix(item.ETag, "p-") {
		// If the item ETag is in "p-[id]" format,
		// then _etag field must be absent from the resource in DB
		s["_etag"] = bson.M{"$exists": false}
	} else {
		s["_etag"] = item.ETag
	}
	err = c.Remove(s)
	if err == mgo.ErrNotFound {
		// Determine if the item is not found or if the item is found but etag missmatch
		var count int
		count, err = c.FindId(item.ID).Count()
		if err != nil {
			// The find returned an unexpected err, just forward it with no mapping
		} else if count == 0 {
			err = resource.ErrNotFound
		} else if ctx.Err() != nil {
			err = ctx.Err()
		} else {
			// If the item were found, it means that its etag didn't match
			err = resource.ErrConflict
		}
	}
	return err
}

// Clear clears all items from the mongo collection matching the query. Note
// that when q.Window != nil, the current implementation may error if the BSON
// encoding of all matching IDs according to the q.Window length gets close to
// the maximum document size in MongDB (usually 16MiB):
// https://docs.mongodb.com/manual/reference/limits/#bson-documents
func (m Handler) Clear(ctx context.Context, q *query.Query) (int, error) {
	// When not applying windowing, qry will be passed directly to RemoveAll.
	qry, err := getQuery(q)
	if err != nil {
		return 0, err
	}

	c, err := m.c(ctx)
	if err != nil {
		return 0, err
	}
	defer m.close(c)

	if q.Window != nil {
		// RemoveAll does not allow skip and limit to be set. To workaround
		// this we do an additional pre-query to retrieve a sorted and sliced
		// list of the IDs for all items to be deleted.
		//
		// This solution does not handle the case where a query containg all
		// IDs is larger than the maximum BSON document size in MongoDB:
		// https://docs.mongodb.com/manual/reference/limits/#bson-documents
		srt := getSort(q)
		mq := applyWindow(c.Find(qry).Sort(srt...), *q.Window)

		if ids, err := selectIDs(c, mq); err == nil {
			qry = bson.M{"_id": bson.M{"$in": ids}}
		} else {
			return 0, err
		}
	}

	// We handle the potential of partial failure by returning both the number
	// of removed items and an error, if both are present.
	info, err := c.RemoveAll(qry)
	if err == nil {
		err = ctx.Err()
	}
	if info == nil {
		return 0, err
	}
	return info.Removed, err
}

// Find items from the mongo collection matching the provided query.
func (m Handler) Find(ctx context.Context, q *query.Query) (*resource.ItemList, error) {
	// MongoDB will return all records on Limit=0. Workaround that behavior.
	// https://docs.mongodb.com/manual/reference/method/cursor.limit/#zero-value
	if q.Window != nil && q.Window.Limit == 0 {
		n, err := m.Count(ctx, q)
		if err != nil {
			return nil, err
		}
		list := &resource.ItemList{
			Total: n,
			Limit: q.Window.Limit,
			Items: []*resource.Item{},
		}
		return list, err
	}

	qry, err := getQuery(q)
	if err != nil {
		return nil, err
	}
	srt := getSort(q)

	c, err := m.c(ctx)
	if err != nil {
		return nil, err
	}
	defer m.close(c)

	mq := c.Find(qry).Sort(srt...)
	limit := -1
	if q.Window != nil {
		mq = applyWindow(mq, *q.Window)
		limit = q.Window.Limit
	}

	// Apply context deadline if any
	if dl, ok := ctx.Deadline(); ok {
		dur := dl.Sub(time.Now())
		if dur < 0 {
			dur = 0
		}
		mq.SetMaxTime(dur)
	}

	// Perform request
	iter := mq.Iter()
	// Total is set to -1 because we have no easy way with MongoDB to to compute
	// this value without performing two requests.
	list := &resource.ItemList{
		Total: -1,
		Limit: limit,
		Items: []*resource.Item{},
	}

	var mItem mongoItem
	for iter.Next(&mItem) {
		// Check if context is still ok before to continue
		if err = ctx.Err(); err != nil {
			// TODO bench this as net/context is using mutex under the hood
			iter.Close()
			return nil, err
		}
		list.Items = append(list.Items, newItem(&mItem))
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}
	// If the number of returned elements is lower than requested limit, or no
	// limit is requested, we can deduce the total number of element for free.
	if limit < 0 || len(list.Items) < limit {
		if q.Window != nil && q.Window.Offset > 0 {
			if len(list.Items) > 0 {
				list.Total = q.Window.Offset + len(list.Items)
			}
			// If there are no items returned when Offset > 0, we may be out-of-bounds,
			// and therefore cannot deduce the total count of items.
		} else {
			list.Total = len(list.Items)
		}
	}
	return list, err
}

// Count counts the number items matching the lookup filter
func (m Handler) Count(ctx context.Context, query *query.Query) (int, error) {
	q, err := getQuery(query)
	if err != nil {
		return -1, err
	}
	c, err := m.c(ctx)
	if err != nil {
		return -1, err
	}
	defer m.close(c)
	mq := c.Find(q)
	// Apply context deadline if any
	if dl, ok := ctx.Deadline(); ok {
		dur := dl.Sub(time.Now())
		if dur < 0 {
			dur = 0
		}
		mq.SetMaxTime(dur)
	}
	return mq.Count()
}
<commit_msg>Bugfix: prevent nilmap assignment<commit_after>// Package mongo is a REST Layer resource storage handler for MongoDB using mgo
package mongo

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/rs/rest-layer/resource"
	"github.com/rs/rest-layer/schema/query"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// mongoItem is a bson representation of a resource.Item.
type mongoItem struct {
	ID      interface{}            `bson:"_id"`
	ETag    string                 `bson:"_etag"`
	Updated time.Time              `bson:"_updated"`
	Payload map[string]interface{} `bson:",inline"`
}

// newMongoItem converts a resource.Item into a mongoItem.
func newMongoItem(i *resource.Item) *mongoItem {
	// Filter out id from the payload so we don't store it twice
	p := map[string]interface{}{}
	for k, v := range i.Payload {
		if k != "id" {
			p[k] = v
		}
	}
	return &mongoItem{
		ID:      i.ID,
		ETag:    i.ETag,
		Updated: i.Updated,
		Payload: p,
	}
}

// newItem converts a back mongoItem into a resource.Item.
func newItem(i *mongoItem) *resource.Item {
	// If there is no field except those defined in mongoItem, Payload could be nil
	// when just fetched from the database.
	if i.Payload == nil {
		i.Payload = make(map[string]interface{})
	}
	// Add the id back (we use the same map hoping the mongoItem won't be stored back)
	i.Payload["id"] = i.ID
	item := &resource.Item{
		ID:      i.ID,
		ETag:    i.ETag,
		Updated: i.Updated,
		Payload: i.Payload,
	}

	if item.ETag == "" {
		if v, ok := i.ID.(bson.ObjectId); ok {
			item.ETag = "p-" + v.Hex()
		} else {
			item.ETag = "p-" + fmt.Sprint(i.ID)
		}
	}
	return item
}

// Handler handles resource storage in a MongoDB collection.
type Handler func(ctx context.Context) (*mgo.Collection, error)

// NewHandler creates an new mongo handler
func NewHandler(s *mgo.Session, db, collection string) Handler {
	c := func() *mgo.Collection {
		return s.DB(db).C(collection)
	}
	return func(ctx context.Context) (*mgo.Collection, error) {
		return c(), nil
	}
}

// C returns the mongo collection managed by this storage handler
// from a Copy() of the mgo session.
func (m Handler) c(ctx context.Context) (*mgo.Collection, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	c, err := m(ctx)
	if err != nil {
		return nil, err
	}
	// With mgo, session.Copy() pulls a connection from the connection pool
	s := c.Database.Session.Copy()
	// Ensure safe mode is enabled in order to get errors
	s.EnsureSafe(&mgo.Safe{})
	// Set a timeout to match the context deadline if any
	if deadline, ok := ctx.Deadline(); ok {
		timeout := deadline.Sub(time.Now())
		if timeout <= 0 {
			timeout = 0
		}
		s.SetSocketTimeout(timeout)
		s.SetSyncTimeout(timeout)
	}
	c.Database.Session = s
	return c, nil
}

// close returns a mgo.Collection's session to the connection pool.
func (m Handler) close(c *mgo.Collection) {
	c.Database.Session.Close()
}

// Insert inserts new items in the mongo collection.
func (m Handler) Insert(ctx context.Context, items []*resource.Item) error {
	mItems := make([]interface{}, len(items))
	for i, item := range items {
		mItems[i] = newMongoItem(item)
	}
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	err = c.Insert(mItems...)
	if mgo.IsDup(err) {
		// Duplicate ID key
		err = resource.ErrConflict
	}
	if ctx.Err() != nil {
		return ctx.Err()
	}
	return err
}

// Update replace an item by a new one in the mongo collection.
func (m Handler) Update(ctx context.Context, item *resource.Item, original *resource.Item) error {
	mItem := newMongoItem(item)
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	s := bson.M{"_id": original.ID}
	if strings.HasPrefix(original.ETag, "p-") {
		// If the original ETag is in "p-[id]" format,
		// then _etag field must be absent from the resource in DB
		s["_etag"] = bson.M{"$exists": false}
	} else {
		s["_etag"] = original.ETag
	}
	err = c.Update(s, mItem)
	if err == mgo.ErrNotFound {
		// Determine if the item is not found or if the item is found but etag missmatch
		var count int
		count, err = c.FindId(original.ID).Count()
		if err != nil {
			// The find returned an unexpected err, just forward it with no mapping
		} else if count == 0 {
			err = resource.ErrNotFound
		} else if ctx.Err() != nil {
			err = ctx.Err()
		} else {
			// If the item were found, it means that its etag didn't match
			err = resource.ErrConflict
		}
	}
	return err
}

// Delete deletes an item from the mongo collection.
func (m Handler) Delete(ctx context.Context, item *resource.Item) error {
	c, err := m.c(ctx)
	if err != nil {
		return err
	}
	defer m.close(c)
	s := bson.M{"_id": item.ID}
	if strings.HasPrefix(item.ETag, "p-") {
		// If the item ETag is in "p-[id]" format,
		// then _etag field must be absent from the resource in DB
		s["_etag"] = bson.M{"$exists": false}
	} else {
		s["_etag"] = item.ETag
	}
	err = c.Remove(s)
	if err == mgo.ErrNotFound {
		// Determine if the item is not found or if the item is found but etag missmatch
		var count int
		count, err = c.FindId(item.ID).Count()
		if err != nil {
			// The find returned an unexpected err, just forward it with no mapping
		} else if count == 0 {
			err = resource.ErrNotFound
		} else if ctx.Err() != nil {
			err = ctx.Err()
		} else {
			// If the item were found, it means that its etag didn't match
			err = resource.ErrConflict
		}
	}
	return err
}

// Clear clears all items from the mongo collection matching the query. Note
// that when q.Window != nil, the current implementation may error if the BSON
// encoding of all matching IDs according to the q.Window length gets close to
// the maximum document size in MongDB (usually 16MiB):
// https://docs.mongodb.com/manual/reference/limits/#bson-documents
func (m Handler) Clear(ctx context.Context, q *query.Query) (int, error) {
	// When not applying windowing, qry will be passed directly to RemoveAll.
	qry, err := getQuery(q)
	if err != nil {
		return 0, err
	}

	c, err := m.c(ctx)
	if err != nil {
		return 0, err
	}
	defer m.close(c)

	if q.Window != nil {
		// RemoveAll does not allow skip and limit to be set. To workaround
		// this we do an additional pre-query to retrieve a sorted and sliced
		// list of the IDs for all items to be deleted.
		//
		// This solution does not handle the case where a query containg all
		// IDs is larger than the maximum BSON document size in MongoDB:
		// https://docs.mongodb.com/manual/reference/limits/#bson-documents
		srt := getSort(q)
		mq := applyWindow(c.Find(qry).Sort(srt...), *q.Window)

		if ids, err := selectIDs(c, mq); err == nil {
			qry = bson.M{"_id": bson.M{"$in": ids}}
		} else {
			return 0, err
		}
	}

	// We handle the potential of partial failure by returning both the number
	// of removed items and an error, if both are present.
	info, err := c.RemoveAll(qry)
	if err == nil {
		err = ctx.Err()
	}
	if info == nil {
		return 0, err
	}
	return info.Removed, err
}

// Find items from the mongo collection matching the provided query.
func (m Handler) Find(ctx context.Context, q *query.Query) (*resource.ItemList, error) {
	// MongoDB will return all records on Limit=0. Workaround that behavior.
	// https://docs.mongodb.com/manual/reference/method/cursor.limit/#zero-value
	if q.Window != nil && q.Window.Limit == 0 {
		n, err := m.Count(ctx, q)
		if err != nil {
			return nil, err
		}
		list := &resource.ItemList{
			Total: n,
			Limit: q.Window.Limit,
			Items: []*resource.Item{},
		}
		return list, err
	}

	qry, err := getQuery(q)
	if err != nil {
		return nil, err
	}
	srt := getSort(q)

	c, err := m.c(ctx)
	if err != nil {
		return nil, err
	}
	defer m.close(c)

	mq := c.Find(qry).Sort(srt...)
	limit := -1
	if q.Window != nil {
		mq = applyWindow(mq, *q.Window)
		limit = q.Window.Limit
	}

	// Apply context deadline if any
	if dl, ok := ctx.Deadline(); ok {
		dur := dl.Sub(time.Now())
		if dur < 0 {
			dur = 0
		}
		mq.SetMaxTime(dur)
	}

	// Perform request
	iter := mq.Iter()
	// Total is set to -1 because we have no easy way with MongoDB to to compute
	// this value without performing two requests.
	list := &resource.ItemList{
		Total: -1,
		Limit: limit,
		Items: []*resource.Item{},
	}

	var mItem mongoItem
	for iter.Next(&mItem) {
		// Check if context is still ok before to continue
		if err = ctx.Err(); err != nil {
			// TODO bench this as net/context is using mutex under the hood
			iter.Close()
			return nil, err
		}
		list.Items = append(list.Items, newItem(&mItem))
	}
	if err := iter.Close(); err != nil {
		return nil, err
	}
	// If the number of returned elements is lower than requested limit, or no
	// limit is requested, we can deduce the total number of element for free.
	if limit < 0 || len(list.Items) < limit {
		if q.Window != nil && q.Window.Offset > 0 {
			if len(list.Items) > 0 {
				list.Total = q.Window.Offset + len(list.Items)
			}
			// If there are no items returned when Offset > 0, we may be out-of-bounds,
			// and therefore cannot deduce the total count of items.
		} else {
			list.Total = len(list.Items)
		}
	}
	return list, err
}

// Count counts the number items matching the lookup filter
func (m Handler) Count(ctx context.Context, query *query.Query) (int, error) {
	q, err := getQuery(query)
	if err != nil {
		return -1, err
	}
	c, err := m.c(ctx)
	if err != nil {
		return -1, err
	}
	defer m.close(c)
	mq := c.Find(q)
	// Apply context deadline if any
	if dl, ok := ctx.Deadline(); ok {
		dur := dl.Sub(time.Now())
		if dur < 0 {
			dur = 0
		}
		mq.SetMaxTime(dur)
	}
	return mq.Count()
}
<|endoftext|>


<commit_before>package middleware

import (
	"crypto/subtle"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/pkg/errors"
	"golang.org/x/crypto/bcrypt"
)

const (
	cost = 12
)

func hasher(s string) []byte {
	hash, err := bcrypt.GenerateFromPassword([]byte(s), cost)
	if err != nil {
		log.Fatal(errors.Wrap(err, "basic auth handler hashing error"))
	}

	return hash
}

func BasicAuthHandler(handler http.Handler) http.Handler {
	userHash := hasher(os.Getenv("LIBERTY_USER"))
	passHash := hasher(os.Getenv("LIBERTY_PASS"))
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		user, pass, ok := r.BasicAuth()
		fmt.Println("BASIC_AUTH", user, pass)
		fmt.Println("BASIC_HASHED", string(userHash), string(passHash))
		fmt.Println("CREDS_HASHED", string(hasher(user)), string(hasher(pass)))
		if !ok || subtle.ConstantTimeCompare(hasher(user),
			userHash) != 1 || subtle.ConstantTimeCompare(hasher(pass), passHash) != 1 {
			w.Header().Set("WWW-Authenticate", `Basic realm=Username and Password`)
			http.Error(w, "Unauthorized.", http.StatusUnauthorized)
			return
		}
		handler.ServeHTTP(w, r)
	})
}
<commit_msg>fix: enable basic auth handler<commit_after>package middleware

import (
	"log"
	"net/http"
	"os"

	"github.com/pkg/errors"
	"golang.org/x/crypto/bcrypt"
)

const (
	cost = 12
)

func hasher(s string) []byte {
	hash, err := bcrypt.GenerateFromPassword([]byte(s), cost)
	if err != nil {
		log.Fatal(errors.Wrap(err, "basic auth handler hashing error"))
	}

	return hash
}

var userHash = hasher(os.Getenv("LIBERTY_USER"))
var passHash = hasher(os.Getenv("LIBERTY_PASS"))

func BasicAuthHandler(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		user, pass, ok := r.BasicAuth()
		if !ok || bcrypt.CompareHashAndPassword(userHash, []byte(user)) != nil || bcrypt.CompareHashAndPassword(passHash, []byte(pass)) != nil {
			w.Header().Set("WWW-Authenticate", `Basic realm=Username and Password`)
			http.Error(w, "Unauthorized.", http.StatusUnauthorized)
			return
		}
		handler.ServeHTTP(w, r)
	})
}
<|endoftext|>


<commit_before>package infrastructure

import (
	"encoding/json"

	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type fileMetadataService struct {
	metaDataFilePath string
	userDataFilePath string
	settingsFilePath string
	fs               boshsys.FileSystem

	logger boshlog.Logger
	logTag string
}

func NewFileMetadataService(
	metaDataFilePath string,
	userDataFilePath string,
	settingsFilePath string,
	fs boshsys.FileSystem,
	logger boshlog.Logger,
) MetadataService {
	return fileMetadataService{
		metaDataFilePath: metaDataFilePath,
		userDataFilePath: userDataFilePath,
		settingsFilePath: settingsFilePath,
		fs:               fs,

		logTag: "fileMetadataService",
		logger: logger,
	}
}

func (ms fileMetadataService) Load() error {
	return nil
}

func (ms fileMetadataService) GetPublicKey() (string, error) {
	return "", nil
}

func (ms fileMetadataService) GetInstanceID() (string, error) {
	var metadata MetadataContentsType

	contents, err := ms.fs.ReadFile(ms.metaDataFilePath)
	if err != nil {
		return "", bosherr.WrapError(err, "Reading metadata file")
	}

	err = json.Unmarshal([]byte(contents), &metadata)
	if err != nil {
		return "", bosherr.WrapError(err, "Unmarshalling metadata")
	}

	ms.logger.Debug(ms.logTag, "Read metadata '%#v'", metadata)

	return metadata.InstanceID, nil
}

func (ms fileMetadataService) GetServerName() (string, error) {
	return "", nil
}

func (ms fileMetadataService) GetRegistryEndpoint() (string, error) {
	var userData UserDataContentsType

	contents, err := ms.fs.ReadFile(ms.userDataFilePath)
	if err != nil {
		// Older versions of bosh-warden-cpi placed
		// full settings file at a specific location.
		return ms.settingsFilePath, nil
	}

	err = json.Unmarshal([]byte(contents), &userData)
	if err != nil {
		return "", bosherr.WrapError(err, "Unmarshalling user data")
	}

	ms.logger.Debug(ms.logTag, "Read user data '%#v'", userData)

	return userData.Registry.Endpoint, nil
}

func (ms fileMetadataService) GetNetworks() (boshsettings.Networks, error) {
	var userData UserDataContentsType

	contents, err := ms.fs.ReadFile(ms.userDataFilePath)
	if err != nil {
		return nil, bosherr.WrapError(err, "Reading user data")
	}

	err = json.Unmarshal([]byte(contents), &userData)
	if err != nil {
		return nil, bosherr.WrapError(err, "Unmarshalling user data")
	}

	ms.logger.Debug(ms.logTag, "Read user data '%#v'", userData)

	return userData.Networks, nil
}

func (ms fileMetadataService) IsAvailable() bool {
	return ms.fs.FileExists(ms.settingsFilePath)
}
<commit_msg>Allow the BOSH agent to look for an SSH public key in the initial agent settings<commit_after>package infrastructure

import (
	"encoding/json"

	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	bosherr "github.com/cloudfoundry/bosh-utils/errors"
	boshlog "github.com/cloudfoundry/bosh-utils/logger"
	boshsys "github.com/cloudfoundry/bosh-utils/system"
)

type PublicKeyContent struct {
	PublicKey string `json:"public_key"`
}

type fileMetadataService struct {
	metaDataFilePath string
	userDataFilePath string
	settingsFilePath string
	fs               boshsys.FileSystem

	logger boshlog.Logger
	logTag string
}

func NewFileMetadataService(
	metaDataFilePath string,
	userDataFilePath string,
	settingsFilePath string,
	fs boshsys.FileSystem,
	logger boshlog.Logger,
) MetadataService {
	return fileMetadataService{
		metaDataFilePath: metaDataFilePath,
		userDataFilePath: userDataFilePath,
		settingsFilePath: settingsFilePath,
		fs:               fs,

		logTag: "fileMetadataService",
		logger: logger,
	}
}

func (ms fileMetadataService) Load() error {
	return nil
}

func (ms fileMetadataService) GetPublicKey() (string, error) {
	var p PublicKeyContent

	contents, err := ms.fs.ReadFile(ms.settingsFilePath)
	if err != nil {
		return "", bosherr.WrapError(err, "Reading metadata file")
	}

	err = json.Unmarshal([]byte(contents), &p)
	if err != nil {
		return "", bosherr.WrapError(err, "Unmarshalling metadata")
	}

	return p.PublicKey, nil
}

func (ms fileMetadataService) GetInstanceID() (string, error) {
	var metadata MetadataContentsType

	contents, err := ms.fs.ReadFile(ms.metaDataFilePath)
	if err != nil {
		return "", bosherr.WrapError(err, "Reading metadata file")
	}

	err = json.Unmarshal([]byte(contents), &metadata)
	if err != nil {
		return "", bosherr.WrapError(err, "Unmarshalling metadata")
	}

	ms.logger.Debug(ms.logTag, "Read metadata '%#v'", metadata)

	return metadata.InstanceID, nil
}

func (ms fileMetadataService) GetServerName() (string, error) {
	return "", nil
}

func (ms fileMetadataService) GetRegistryEndpoint() (string, error) {
	var userData UserDataContentsType

	contents, err := ms.fs.ReadFile(ms.userDataFilePath)
	if err != nil {
		// Older versions of bosh-warden-cpi placed
		// full settings file at a specific location.
		return ms.settingsFilePath, nil
	}

	err = json.Unmarshal([]byte(contents), &userData)
	if err != nil {
		return "", bosherr.WrapError(err, "Unmarshalling user data")
	}

	ms.logger.Debug(ms.logTag, "Read user data '%#v'", userData)

	return userData.Registry.Endpoint, nil
}

func (ms fileMetadataService) GetNetworks() (boshsettings.Networks, error) {
	var userData UserDataContentsType

	contents, err := ms.fs.ReadFile(ms.userDataFilePath)
	if err != nil {
		return nil, bosherr.WrapError(err, "Reading user data")
	}

	err = json.Unmarshal([]byte(contents), &userData)
	if err != nil {
		return nil, bosherr.WrapError(err, "Unmarshalling user data")
	}

	ms.logger.Debug(ms.logTag, "Read user data '%#v'", userData)

	return userData.Networks, nil
}

func (ms fileMetadataService) IsAvailable() bool {
	return ms.fs.FileExists(ms.settingsFilePath)
}
<|endoftext|>


<commit_before>package load

import (
	"bufio"
	"bytes"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"

	"github.com/pkg/errors"
)

var frozenImgDir = "/docker-frozen-images"

// FrozenImagesLinux loads the frozen image set for the integration suite
// If the images are not available locally it will download them
// TODO: This loads whatever is in the frozen image dir, regardless of what
// images were passed in. If the images need to be downloaded, then it will respect
// the passed in images
func FrozenImagesLinux(dockerBinary string, images ...string) error {
	imgNS := os.Getenv("TEST_IMAGE_NAMESPACE")
	var loadImages []struct{ srcName, destName string }
	for _, img := range images {
		if err := exec.Command(dockerBinary, "inspect", "--type=image", img).Run(); err != nil {
			srcName := img
			// hello-world:latest gets re-tagged as hello-world:frozen
			// there are some tests that use hello-world:latest specifically so it pulls
			// the image and hello-world:frozen is used for when we just want a super
			// small image
			if img == "hello-world:frozen" {
				srcName = "hello-world:latest"
			}
			if imgNS != "" {
				srcName = imgNS + "/" + srcName
			}
			loadImages = append(loadImages, struct{ srcName, destName string }{
				srcName:  srcName,
				destName: img,
			})
		}
	}
	if len(loadImages) == 0 {
		// everything is loaded, we're done
		return nil
	}

	fi, err := os.Stat(frozenImgDir)
	if err != nil || !fi.IsDir() {
		srcImages := make([]string, 0, len(loadImages))
		for _, img := range loadImages {
			srcImages = append(srcImages, img.srcName)
		}
		if err := pullImages(dockerBinary, srcImages); err != nil {
			return errors.Wrap(err, "error pulling image list")
		}
	} else {
		if err := loadFrozenImages(dockerBinary); err != nil {
			return err
		}
	}

	for _, img := range loadImages {
		if img.srcName != img.destName {
			if out, err := exec.Command(dockerBinary, "tag", img.srcName, img.destName).CombinedOutput(); err != nil {
				return errors.Errorf("%v: %s", err, string(out))
			}
			if out, err := exec.Command(dockerBinary, "rmi", img.srcName).CombinedOutput(); err != nil {
				return errors.Errorf("%v: %s", err, string(out))
			}
		}
	}
	return nil
}

func loadFrozenImages(dockerBinary string) error {
	tar, err := exec.LookPath("tar")
	if err != nil {
		return errors.Wrap(err, "could not find tar binary")
	}
	tarCmd := exec.Command(tar, "-cC", frozenImgDir, ".")
	out, err := tarCmd.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "error getting stdout pipe for tar command")
	}

	errBuf := bytes.NewBuffer(nil)
	tarCmd.Stderr = errBuf
	tarCmd.Start()
	defer tarCmd.Wait()

	cmd := exec.Command(dockerBinary, "load")
	cmd.Stdin = out
	if out, err := cmd.CombinedOutput(); err != nil {
		return errors.Errorf("%v: %s", err, string(out))
	}
	return nil
}

func pullImages(dockerBinary string, images []string) error {
	cwd, err := os.Getwd()
	if err != nil {
		return errors.Wrap(err, "error getting path to dockerfile")
	}
	dockerfile := os.Getenv("DOCKERFILE")
	if dockerfile == "" {
		dockerfile = "Dockerfile"
	}
	dockerfilePath := filepath.Join(filepath.Dir(filepath.Clean(cwd)), dockerfile)
	pullRefs, err := readFrozenImageList(dockerfilePath, images)
	if err != nil {
		return errors.Wrap(err, "error reading frozen image list")
	}

	var wg sync.WaitGroup
	chErr := make(chan error, len(images))
	for tag, ref := range pullRefs {
		wg.Add(1)
		go func(tag, ref string) {
			defer wg.Done()
			if out, err := exec.Command(dockerBinary, "pull", ref).CombinedOutput(); err != nil {
				chErr <- errors.Errorf("%v: %s", string(out), err)
				return
			}
			if out, err := exec.Command(dockerBinary, "tag", ref, tag).CombinedOutput(); err != nil {
				chErr <- errors.Errorf("%v: %s", string(out), err)
				return
			}
			if out, err := exec.Command(dockerBinary, "rmi", ref).CombinedOutput(); err != nil {
				chErr <- errors.Errorf("%v: %s", string(out), err)
				return
			}
		}(tag, ref)
	}
	wg.Wait()
	close(chErr)
	return <-chErr
}

func readFrozenImageList(dockerfilePath string, images []string) (map[string]string, error) {
	f, err := os.Open(dockerfilePath)
	if err != nil {
		return nil, errors.Wrap(err, "error reading dockerfile")
	}
	defer f.Close()
	ls := make(map[string]string)

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.Fields(scanner.Text())
		if len(line) < 3 {
			continue
		}
		if !(line[0] == "RUN" && line[1] == "./contrib/download-frozen-image-v2.sh") {
			continue
		}

		frozenImgDir = line[2]
		if line[2] == frozenImgDir {
			frozenImgDir = filepath.Join(os.Getenv("DEST"), "frozen-images")
		}

		for scanner.Scan() {
			img := strings.TrimSpace(scanner.Text())
			img = strings.TrimSuffix(img, "\\")
			img = strings.TrimSpace(img)
			split := strings.Split(img, "@")
			if len(split) < 2 {
				break
			}

			for _, i := range images {
				if split[0] == i {
					ls[i] = img
					break
				}
			}
		}
	}
	return ls, nil
}
<commit_msg>Update fixtures/load to use the APIClient<commit_after>package load

import (
	"bufio"
	"bytes"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"

	"context"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/client"
	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/docker/docker/pkg/term"
	"github.com/pkg/errors"
)

const frozenImgDir = "/docker-frozen-images"

// FrozenImagesLinux loads the frozen image set for the integration suite
// If the images are not available locally it will download them
// TODO: This loads whatever is in the frozen image dir, regardless of what
// images were passed in. If the images need to be downloaded, then it will respect
// the passed in images
func FrozenImagesLinux(client client.APIClient, images ...string) error {
	imgNS := os.Getenv("TEST_IMAGE_NAMESPACE")
	var loadImages []struct{ srcName, destName string }
	for _, img := range images {
		if !imageExists(client, img) {
			srcName := img
			// hello-world:latest gets re-tagged as hello-world:frozen
			// there are some tests that use hello-world:latest specifically so it pulls
			// the image and hello-world:frozen is used for when we just want a super
			// small image
			if img == "hello-world:frozen" {
				srcName = "hello-world:latest"
			}
			if imgNS != "" {
				srcName = imgNS + "/" + srcName
			}
			loadImages = append(loadImages, struct{ srcName, destName string }{
				srcName:  srcName,
				destName: img,
			})
		}
	}
	if len(loadImages) == 0 {
		// everything is loaded, we're done
		return nil
	}

	ctx := context.Background()
	fi, err := os.Stat(frozenImgDir)
	if err != nil || !fi.IsDir() {
		srcImages := make([]string, 0, len(loadImages))
		for _, img := range loadImages {
			srcImages = append(srcImages, img.srcName)
		}
		if err := pullImages(ctx, client, srcImages); err != nil {
			return errors.Wrap(err, "error pulling image list")
		}
	} else {
		if err := loadFrozenImages(ctx, client); err != nil {
			return err
		}
	}

	for _, img := range loadImages {
		if img.srcName != img.destName {
			if err := client.ImageTag(ctx, img.srcName, img.destName); err != nil {
				return errors.Wrapf(err, "failed to tag %s as %s", img.srcName, img.destName)
			}
			if _, err := client.ImageRemove(ctx, img.srcName, types.ImageRemoveOptions{}); err != nil {
				return errors.Wrapf(err, "failed to remove %s", img.srcName)
			}
		}
	}
	return nil
}

func imageExists(client client.APIClient, name string) bool {
	_, _, err := client.ImageInspectWithRaw(context.Background(), name)
	return err == nil
}

func loadFrozenImages(ctx context.Context, client client.APIClient) error {
	tar, err := exec.LookPath("tar")
	if err != nil {
		return errors.Wrap(err, "could not find tar binary")
	}
	tarCmd := exec.Command(tar, "-cC", frozenImgDir, ".")
	out, err := tarCmd.StdoutPipe()
	if err != nil {
		return errors.Wrap(err, "error getting stdout pipe for tar command")
	}

	errBuf := bytes.NewBuffer(nil)
	tarCmd.Stderr = errBuf
	tarCmd.Start()
	defer tarCmd.Wait()

	resp, err := client.ImageLoad(ctx, out, true)
	if err != nil {
		return errors.Wrap(err, "failed to load frozen images")
	}
	defer resp.Body.Close()
	fd, isTerminal := term.GetFdInfo(os.Stdout)
	return jsonmessage.DisplayJSONMessagesStream(resp.Body, os.Stdout, fd, isTerminal, nil)
}

func pullImages(ctx context.Context, client client.APIClient, images []string) error {
	cwd, err := os.Getwd()
	if err != nil {
		return errors.Wrap(err, "error getting path to dockerfile")
	}
	dockerfile := os.Getenv("DOCKERFILE")
	if dockerfile == "" {
		dockerfile = "Dockerfile"
	}
	dockerfilePath := filepath.Join(filepath.Dir(filepath.Clean(cwd)), dockerfile)
	pullRefs, err := readFrozenImageList(dockerfilePath, images)
	if err != nil {
		return errors.Wrap(err, "error reading frozen image list")
	}

	var wg sync.WaitGroup
	chErr := make(chan error, len(images))
	for tag, ref := range pullRefs {
		wg.Add(1)
		go func(tag, ref string) {
			defer wg.Done()
			if err := pullTagAndRemove(ctx, client, ref, tag); err != nil {
				chErr <- err
				return
			}
		}(tag, ref)
	}
	wg.Wait()
	close(chErr)
	return <-chErr
}

func pullTagAndRemove(ctx context.Context, client client.APIClient, ref string, tag string) error {
	resp, err := client.ImagePull(ctx, ref, types.ImagePullOptions{})
	if err != nil {
		return errors.Wrapf(err, "failed to pull %s", ref)
	}
	defer resp.Close()
	fd, isTerminal := term.GetFdInfo(os.Stdout)
	if err := jsonmessage.DisplayJSONMessagesStream(resp, os.Stdout, fd, isTerminal, nil); err != nil {
		return err
	}

	if err := client.ImageTag(ctx, ref, tag); err != nil {
		return errors.Wrapf(err, "failed to tag %s as %s", ref, tag)
	}
	_, err = client.ImageRemove(ctx, ref, types.ImageRemoveOptions{})
	return errors.Wrapf(err, "failed to remove %s", ref)

}

func readFrozenImageList(dockerfilePath string, images []string) (map[string]string, error) {
	f, err := os.Open(dockerfilePath)
	if err != nil {
		return nil, errors.Wrap(err, "error reading dockerfile")
	}
	defer f.Close()
	ls := make(map[string]string)

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := strings.Fields(scanner.Text())
		if len(line) < 3 {
			continue
		}
		if !(line[0] == "RUN" && line[1] == "./contrib/download-frozen-image-v2.sh") {
			continue
		}

		for scanner.Scan() {
			img := strings.TrimSpace(scanner.Text())
			img = strings.TrimSuffix(img, "\\")
			img = strings.TrimSpace(img)
			split := strings.Split(img, "@")
			if len(split) < 2 {
				break
			}

			for _, i := range images {
				if split[0] == i {
					ls[i] = img
					break
				}
			}
		}
	}
	return ls, nil
}
<|endoftext|>