<commit_before>// Copyright 2018 The gVisor 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 testutil

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/kr/pty"
)

var runtimeType = flag.String("runtime-type", "", "specify which runtime to use: kvm, hostnet, overlay")

func getRuntime() string {
	r, ok := os.LookupEnv("RUNSC_RUNTIME")
	if !ok {
		r = "runsc-test"
	}
	if *runtimeType != "" {
		r += "-" + *runtimeType
	}
	return r
}

// IsPauseResumeSupported returns true if Pause/Resume is supported by runtime.
func IsPauseResumeSupported() bool {
	// Native host network stack can't be saved.
	return !strings.Contains(getRuntime(), "hostnet")
}

// EnsureSupportedDockerVersion checks if correct docker is installed.
func EnsureSupportedDockerVersion() {
	cmd := exec.Command("docker", "version")
	out, err := cmd.CombinedOutput()
	if err != nil {
		log.Fatalf("Error running %q: %v", "docker version", err)
	}
	re := regexp.MustCompile(`Version:\s+(\d+)\.(\d+)\.\d.*`)
	matches := re.FindStringSubmatch(string(out))
	if len(matches) != 3 {
		log.Fatalf("Invalid docker output: %s", out)
	}
	major, _ := strconv.Atoi(matches[1])
	minor, _ := strconv.Atoi(matches[2])
	if major < 17 || (major == 17 && minor < 9) {
		log.Fatalf("Docker version 17.09.0 or greater is required, found: %02d.%02d", major, minor)
	}
}

// MountMode describes if the mount should be ro or rw.
type MountMode int

const (
	// ReadOnly is what the name says.
	ReadOnly MountMode = iota
	// ReadWrite is what the name says.
	ReadWrite
)

// String returns the mount mode argument for this MountMode.
func (m MountMode) String() string {
	switch m {
	case ReadOnly:
		return "ro"
	case ReadWrite:
		return "rw"
	}
	panic(fmt.Sprintf("invalid mode: %d", m))
}

// MountArg formats the volume argument to mount in the container.
func MountArg(source, target string, mode MountMode) string {
	return fmt.Sprintf("-v=%s:%s:%v", source, target, mode)
}

// LinkArg formats the link argument.
func LinkArg(source *Docker, target string) string {
	return fmt.Sprintf("--link=%s:%s", source.Name, target)
}

// PrepareFiles creates temp directory to copy files there. The sandbox doesn't
// have access to files in the test dir.
func PrepareFiles(names ...string) (string, error) {
	dir, err := ioutil.TempDir("", "image-test")
	if err != nil {
		return "", fmt.Errorf("ioutil.TempDir failed: %v", err)
	}
	if err := os.Chmod(dir, 0777); err != nil {
		return "", fmt.Errorf("os.Chmod(%q, 0777) failed: %v", dir, err)
	}
	for _, name := range names {
		src := getLocalPath(name)
		dst := path.Join(dir, name)
		if err := Copy(src, dst); err != nil {
			return "", fmt.Errorf("testutil.Copy(%q, %q) failed: %v", src, dst, err)
		}
	}
	return dir, nil
}

func getLocalPath(file string) string {
	return path.Join(".", file)
}

// do executes docker command.
func do(args ...string) (string, error) {
	log.Printf("Running: docker %s\n", args)
	cmd := exec.Command("docker", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("error executing docker %s: %v\nout: %s", args, err, out)
	}
	return string(out), nil
}

// doWithPty executes docker command with stdio attached to a pty.
func doWithPty(args ...string) (*exec.Cmd, *os.File, error) {
	log.Printf("Running with pty: docker %s\n", args)
	cmd := exec.Command("docker", args...)
	ptmx, err := pty.Start(cmd)
	if err != nil {
		return nil, nil, fmt.Errorf("error executing docker %s with a pty: %v", args, err)
	}
	return cmd, ptmx, nil
}

// Pull pulls a docker image. This is used in tests to isolate the
// time to pull the image off the network from the time to actually
// start the container, to avoid timeouts over slow networks.
func Pull(image string) error {
	_, err := do("pull", image)
	return err
}

// Docker contains the name and the runtime of a docker container.
type Docker struct {
	Runtime string
	Name    string
}

// MakeDocker sets up the struct for a Docker container.
// Names of containers will be unique.
func MakeDocker(namePrefix string) Docker {
	return Docker{Name: RandomName(namePrefix), Runtime: getRuntime()}
}

// logDockerID logs a container id, which is needed to find container runsc logs.
func (d *Docker) logDockerID() {
	id, err := d.ID()
	if err != nil {
		log.Printf("%v\n", err)
	}
	log.Printf("Name: %s ID: %v\n", d.Name, id)
}

// Create calls 'docker create' with the arguments provided.
func (d *Docker) Create(args ...string) error {
	a := []string{"create", "--runtime", d.Runtime, "--name", d.Name}
	a = append(a, args...)
	_, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return err
}

// Start calls 'docker start'.
func (d *Docker) Start() error {
	if _, err := do("start", d.Name); err != nil {
		return fmt.Errorf("error starting container %q: %v", d.Name, err)
	}
	return nil
}

// Stop calls 'docker stop'.
func (d *Docker) Stop() error {
	if _, err := do("stop", d.Name); err != nil {
		return fmt.Errorf("error stopping container %q: %v", d.Name, err)
	}
	return nil
}

// Run calls 'docker run' with the arguments provided. The container starts
// running in the background and the call returns immediately.
func (d *Docker) Run(args ...string) error {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name, "-d"}
	a = append(a, args...)
	_, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return err
}

// RunWithPty is like Run but with an attached pty.
func (d *Docker) RunWithPty(args ...string) (*exec.Cmd, *os.File, error) {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name, "-it"}
	a = append(a, args...)
	return doWithPty(a...)
}

// RunFg calls 'docker run' with the arguments provided in the foreground. It
// blocks until the container exits and returns the output.
func (d *Docker) RunFg(args ...string) (string, error) {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name}
	a = append(a, args...)
	out, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return string(out), err
}

// Logs calls 'docker logs'.
func (d *Docker) Logs() (string, error) {
	return do("logs", d.Name)
}

// Exec calls 'docker exec' with the arguments provided.
func (d *Docker) Exec(args ...string) (string, error) {
	a := []string{"exec", d.Name}
	a = append(a, args...)
	return do(a...)
}

// ExecWithTerminal calls 'docker exec -it' with the arguments provided and
// attaches a pty to stdio.
func (d *Docker) ExecWithTerminal(args ...string) (*exec.Cmd, *os.File, error) {
	a := []string{"exec", "-it", d.Name}
	a = append(a, args...)
	return doWithPty(a...)
}

// Pause calls 'docker pause'.
func (d *Docker) Pause() error {
	if _, err := do("pause", d.Name); err != nil {
		return fmt.Errorf("error pausing container %q: %v", d.Name, err)
	}
	return nil
}

// Unpause calls 'docker pause'.
func (d *Docker) Unpause() error {
	if _, err := do("unpause", d.Name); err != nil {
		return fmt.Errorf("error unpausing container %q: %v", d.Name, err)
	}
	return nil
}

// Checkpoint calls 'docker checkpoint'.
func (d *Docker) Checkpoint(name string) error {
	if _, err := do("checkpoint", "create", d.Name, name); err != nil {
		return fmt.Errorf("error pausing container %q: %v", d.Name, err)
	}
	return nil
}

// Restore calls 'docker start --checkname [name]'.
func (d *Docker) Restore(name string) error {
	if _, err := do("start", "--checkpoint", name, d.Name); err != nil {
		return fmt.Errorf("error starting container %q: %v", d.Name, err)
	}
	return nil
}

// Remove calls 'docker rm'.
func (d *Docker) Remove() error {
	if _, err := do("rm", d.Name); err != nil {
		return fmt.Errorf("error deleting container %q: %v", d.Name, err)
	}
	return nil
}

// CleanUp kills and deletes the container (best effort).
func (d *Docker) CleanUp() {
	d.logDockerID()
	if _, err := do("kill", d.Name); err != nil {
		log.Printf("error killing container %q: %v", d.Name, err)
	}
	if err := d.Remove(); err != nil {
		log.Print(err)
	}
}

// FindPort returns the host port that is mapped to 'sandboxPort'. This calls
// docker to allocate a free port in the host and prevent conflicts.
func (d *Docker) FindPort(sandboxPort int) (int, error) {
	format := fmt.Sprintf(`{{ (index (index .NetworkSettings.Ports "%d/tcp") 0).HostPort }}`, sandboxPort)
	out, err := do("inspect", "-f", format, d.Name)
	if err != nil {
		return -1, fmt.Errorf("error retrieving port: %v", err)
	}
	port, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
	if err != nil {
		return -1, fmt.Errorf("error parsing port %q: %v", out, err)
	}
	return port, nil
}

// SandboxPid returns the PID to the sandbox process.
func (d *Docker) SandboxPid() (int, error) {
	out, err := do("inspect", "-f={{.State.Pid}}", d.Name)
	if err != nil {
		return -1, fmt.Errorf("error retrieving pid: %v", err)
	}
	pid, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
	if err != nil {
		return -1, fmt.Errorf("error parsing pid %q: %v", out, err)
	}
	return pid, nil
}

// ID returns the container ID.
func (d *Docker) ID() (string, error) {
	out, err := do("inspect", "-f={{.Id}}", d.Name)
	if err != nil {
		return "", fmt.Errorf("error retrieving ID: %v", err)
	}
	return strings.TrimSpace(string(out)), nil
}

// Wait waits for container to exit, up to the given timeout. Returns error if
// wait fails or timeout is hit. Returns the application return code otherwise.
// Note that the application may have failed even if err == nil, always check
// the exit code.
func (d *Docker) Wait(timeout time.Duration) (syscall.WaitStatus, error) {
	timeoutChan := time.After(timeout)
	waitChan := make(chan (syscall.WaitStatus))
	errChan := make(chan (error))

	go func() {
		out, err := do("wait", d.Name)
		if err != nil {
			errChan <- fmt.Errorf("error waiting for container %q: %v", d.Name, err)
		}
		exit, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
		if err != nil {
			errChan <- fmt.Errorf("error parsing exit code %q: %v", out, err)
		}
		waitChan <- syscall.WaitStatus(uint32(exit))
	}()

	select {
	case ws := <-waitChan:
		return ws, nil
	case err := <-errChan:
		return syscall.WaitStatus(1), err
	case <-timeoutChan:
		return syscall.WaitStatus(1), fmt.Errorf("timeout waiting for container %q", d.Name)
	}
}

// WaitForOutput calls 'docker logs' to retrieve containers output and searches
// for the given pattern.
func (d *Docker) WaitForOutput(pattern string, timeout time.Duration) (string, error) {
	matches, err := d.WaitForOutputSubmatch(pattern, timeout)
	if err != nil {
		return "", err
	}
	if len(matches) == 0 {
		return "", nil
	}
	return matches[0], nil
}

// WaitForOutputSubmatch calls 'docker logs' to retrieve containers output and
// searches for the given pattern. It returns any regexp submatches as well.
func (d *Docker) WaitForOutputSubmatch(pattern string, timeout time.Duration) ([]string, error) {
	re := regexp.MustCompile(pattern)
	var out string
	for exp := time.Now().Add(timeout); time.Now().Before(exp); {
		var err error
		out, err = d.Logs()
		if err != nil {
			return nil, err
		}
		if matches := re.FindStringSubmatch(out); matches != nil {
			// Success!
			return matches, nil
		}
		time.Sleep(100 * time.Millisecond)
	}
	return nil, fmt.Errorf("timeout waiting for output %q: %s", re.String(), out)
}
<commit_msg>Don't log an error when stopping the container if it is not running.<commit_after>// Copyright 2018 The gVisor 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 testutil

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/kr/pty"
)

var runtimeType = flag.String("runtime-type", "", "specify which runtime to use: kvm, hostnet, overlay")

func getRuntime() string {
	r, ok := os.LookupEnv("RUNSC_RUNTIME")
	if !ok {
		r = "runsc-test"
	}
	if *runtimeType != "" {
		r += "-" + *runtimeType
	}
	return r
}

// IsPauseResumeSupported returns true if Pause/Resume is supported by runtime.
func IsPauseResumeSupported() bool {
	// Native host network stack can't be saved.
	return !strings.Contains(getRuntime(), "hostnet")
}

// EnsureSupportedDockerVersion checks if correct docker is installed.
func EnsureSupportedDockerVersion() {
	cmd := exec.Command("docker", "version")
	out, err := cmd.CombinedOutput()
	if err != nil {
		log.Fatalf("Error running %q: %v", "docker version", err)
	}
	re := regexp.MustCompile(`Version:\s+(\d+)\.(\d+)\.\d.*`)
	matches := re.FindStringSubmatch(string(out))
	if len(matches) != 3 {
		log.Fatalf("Invalid docker output: %s", out)
	}
	major, _ := strconv.Atoi(matches[1])
	minor, _ := strconv.Atoi(matches[2])
	if major < 17 || (major == 17 && minor < 9) {
		log.Fatalf("Docker version 17.09.0 or greater is required, found: %02d.%02d", major, minor)
	}
}

// MountMode describes if the mount should be ro or rw.
type MountMode int

const (
	// ReadOnly is what the name says.
	ReadOnly MountMode = iota
	// ReadWrite is what the name says.
	ReadWrite
)

// String returns the mount mode argument for this MountMode.
func (m MountMode) String() string {
	switch m {
	case ReadOnly:
		return "ro"
	case ReadWrite:
		return "rw"
	}
	panic(fmt.Sprintf("invalid mode: %d", m))
}

// MountArg formats the volume argument to mount in the container.
func MountArg(source, target string, mode MountMode) string {
	return fmt.Sprintf("-v=%s:%s:%v", source, target, mode)
}

// LinkArg formats the link argument.
func LinkArg(source *Docker, target string) string {
	return fmt.Sprintf("--link=%s:%s", source.Name, target)
}

// PrepareFiles creates temp directory to copy files there. The sandbox doesn't
// have access to files in the test dir.
func PrepareFiles(names ...string) (string, error) {
	dir, err := ioutil.TempDir("", "image-test")
	if err != nil {
		return "", fmt.Errorf("ioutil.TempDir failed: %v", err)
	}
	if err := os.Chmod(dir, 0777); err != nil {
		return "", fmt.Errorf("os.Chmod(%q, 0777) failed: %v", dir, err)
	}
	for _, name := range names {
		src := getLocalPath(name)
		dst := path.Join(dir, name)
		if err := Copy(src, dst); err != nil {
			return "", fmt.Errorf("testutil.Copy(%q, %q) failed: %v", src, dst, err)
		}
	}
	return dir, nil
}

func getLocalPath(file string) string {
	return path.Join(".", file)
}

// do executes docker command.
func do(args ...string) (string, error) {
	log.Printf("Running: docker %s\n", args)
	cmd := exec.Command("docker", args...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("error executing docker %s: %v\nout: %s", args, err, out)
	}
	return string(out), nil
}

// doWithPty executes docker command with stdio attached to a pty.
func doWithPty(args ...string) (*exec.Cmd, *os.File, error) {
	log.Printf("Running with pty: docker %s\n", args)
	cmd := exec.Command("docker", args...)
	ptmx, err := pty.Start(cmd)
	if err != nil {
		return nil, nil, fmt.Errorf("error executing docker %s with a pty: %v", args, err)
	}
	return cmd, ptmx, nil
}

// Pull pulls a docker image. This is used in tests to isolate the
// time to pull the image off the network from the time to actually
// start the container, to avoid timeouts over slow networks.
func Pull(image string) error {
	_, err := do("pull", image)
	return err
}

// Docker contains the name and the runtime of a docker container.
type Docker struct {
	Runtime string
	Name    string
}

// MakeDocker sets up the struct for a Docker container.
// Names of containers will be unique.
func MakeDocker(namePrefix string) Docker {
	return Docker{Name: RandomName(namePrefix), Runtime: getRuntime()}
}

// logDockerID logs a container id, which is needed to find container runsc logs.
func (d *Docker) logDockerID() {
	id, err := d.ID()
	if err != nil {
		log.Printf("%v\n", err)
	}
	log.Printf("Name: %s ID: %v\n", d.Name, id)
}

// Create calls 'docker create' with the arguments provided.
func (d *Docker) Create(args ...string) error {
	a := []string{"create", "--runtime", d.Runtime, "--name", d.Name}
	a = append(a, args...)
	_, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return err
}

// Start calls 'docker start'.
func (d *Docker) Start() error {
	if _, err := do("start", d.Name); err != nil {
		return fmt.Errorf("error starting container %q: %v", d.Name, err)
	}
	return nil
}

// Stop calls 'docker stop'.
func (d *Docker) Stop() error {
	if _, err := do("stop", d.Name); err != nil {
		return fmt.Errorf("error stopping container %q: %v", d.Name, err)
	}
	return nil
}

// Run calls 'docker run' with the arguments provided. The container starts
// running in the background and the call returns immediately.
func (d *Docker) Run(args ...string) error {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name, "-d"}
	a = append(a, args...)
	_, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return err
}

// RunWithPty is like Run but with an attached pty.
func (d *Docker) RunWithPty(args ...string) (*exec.Cmd, *os.File, error) {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name, "-it"}
	a = append(a, args...)
	return doWithPty(a...)
}

// RunFg calls 'docker run' with the arguments provided in the foreground. It
// blocks until the container exits and returns the output.
func (d *Docker) RunFg(args ...string) (string, error) {
	a := []string{"run", "--runtime", d.Runtime, "--name", d.Name}
	a = append(a, args...)
	out, err := do(a...)
	if err == nil {
		d.logDockerID()
	}
	return string(out), err
}

// Logs calls 'docker logs'.
func (d *Docker) Logs() (string, error) {
	return do("logs", d.Name)
}

// Exec calls 'docker exec' with the arguments provided.
func (d *Docker) Exec(args ...string) (string, error) {
	a := []string{"exec", d.Name}
	a = append(a, args...)
	return do(a...)
}

// ExecWithTerminal calls 'docker exec -it' with the arguments provided and
// attaches a pty to stdio.
func (d *Docker) ExecWithTerminal(args ...string) (*exec.Cmd, *os.File, error) {
	a := []string{"exec", "-it", d.Name}
	a = append(a, args...)
	return doWithPty(a...)
}

// Pause calls 'docker pause'.
func (d *Docker) Pause() error {
	if _, err := do("pause", d.Name); err != nil {
		return fmt.Errorf("error pausing container %q: %v", d.Name, err)
	}
	return nil
}

// Unpause calls 'docker pause'.
func (d *Docker) Unpause() error {
	if _, err := do("unpause", d.Name); err != nil {
		return fmt.Errorf("error unpausing container %q: %v", d.Name, err)
	}
	return nil
}

// Checkpoint calls 'docker checkpoint'.
func (d *Docker) Checkpoint(name string) error {
	if _, err := do("checkpoint", "create", d.Name, name); err != nil {
		return fmt.Errorf("error pausing container %q: %v", d.Name, err)
	}
	return nil
}

// Restore calls 'docker start --checkname [name]'.
func (d *Docker) Restore(name string) error {
	if _, err := do("start", "--checkpoint", name, d.Name); err != nil {
		return fmt.Errorf("error starting container %q: %v", d.Name, err)
	}
	return nil
}

// Remove calls 'docker rm'.
func (d *Docker) Remove() error {
	if _, err := do("rm", d.Name); err != nil {
		return fmt.Errorf("error deleting container %q: %v", d.Name, err)
	}
	return nil
}

// CleanUp kills and deletes the container (best effort).
func (d *Docker) CleanUp() {
	d.logDockerID()
	if _, err := do("kill", d.Name); err != nil {
		if strings.Contains(err.Error(), "is not running") {
			// Nothing to kill. Don't log the error in this case.
		} else {
			log.Printf("error killing container %q: %v", d.Name, err)
		}
	}
	if err := d.Remove(); err != nil {
		log.Print(err)
	}
}

// FindPort returns the host port that is mapped to 'sandboxPort'. This calls
// docker to allocate a free port in the host and prevent conflicts.
func (d *Docker) FindPort(sandboxPort int) (int, error) {
	format := fmt.Sprintf(`{{ (index (index .NetworkSettings.Ports "%d/tcp") 0).HostPort }}`, sandboxPort)
	out, err := do("inspect", "-f", format, d.Name)
	if err != nil {
		return -1, fmt.Errorf("error retrieving port: %v", err)
	}
	port, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
	if err != nil {
		return -1, fmt.Errorf("error parsing port %q: %v", out, err)
	}
	return port, nil
}

// SandboxPid returns the PID to the sandbox process.
func (d *Docker) SandboxPid() (int, error) {
	out, err := do("inspect", "-f={{.State.Pid}}", d.Name)
	if err != nil {
		return -1, fmt.Errorf("error retrieving pid: %v", err)
	}
	pid, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
	if err != nil {
		return -1, fmt.Errorf("error parsing pid %q: %v", out, err)
	}
	return pid, nil
}

// ID returns the container ID.
func (d *Docker) ID() (string, error) {
	out, err := do("inspect", "-f={{.Id}}", d.Name)
	if err != nil {
		return "", fmt.Errorf("error retrieving ID: %v", err)
	}
	return strings.TrimSpace(string(out)), nil
}

// Wait waits for container to exit, up to the given timeout. Returns error if
// wait fails or timeout is hit. Returns the application return code otherwise.
// Note that the application may have failed even if err == nil, always check
// the exit code.
func (d *Docker) Wait(timeout time.Duration) (syscall.WaitStatus, error) {
	timeoutChan := time.After(timeout)
	waitChan := make(chan (syscall.WaitStatus))
	errChan := make(chan (error))

	go func() {
		out, err := do("wait", d.Name)
		if err != nil {
			errChan <- fmt.Errorf("error waiting for container %q: %v", d.Name, err)
		}
		exit, err := strconv.Atoi(strings.TrimSuffix(string(out), "\n"))
		if err != nil {
			errChan <- fmt.Errorf("error parsing exit code %q: %v", out, err)
		}
		waitChan <- syscall.WaitStatus(uint32(exit))
	}()

	select {
	case ws := <-waitChan:
		return ws, nil
	case err := <-errChan:
		return syscall.WaitStatus(1), err
	case <-timeoutChan:
		return syscall.WaitStatus(1), fmt.Errorf("timeout waiting for container %q", d.Name)
	}
}

// WaitForOutput calls 'docker logs' to retrieve containers output and searches
// for the given pattern.
func (d *Docker) WaitForOutput(pattern string, timeout time.Duration) (string, error) {
	matches, err := d.WaitForOutputSubmatch(pattern, timeout)
	if err != nil {
		return "", err
	}
	if len(matches) == 0 {
		return "", nil
	}
	return matches[0], nil
}

// WaitForOutputSubmatch calls 'docker logs' to retrieve containers output and
// searches for the given pattern. It returns any regexp submatches as well.
func (d *Docker) WaitForOutputSubmatch(pattern string, timeout time.Duration) ([]string, error) {
	re := regexp.MustCompile(pattern)
	var out string
	for exp := time.Now().Add(timeout); time.Now().Before(exp); {
		var err error
		out, err = d.Logs()
		if err != nil {
			return nil, err
		}
		if matches := re.FindStringSubmatch(out); matches != nil {
			// Success!
			return matches, nil
		}
		time.Sleep(100 * time.Millisecond)
	}
	return nil, fmt.Errorf("timeout waiting for output %q: %s", re.String(), out)
}
<|endoftext|>


<commit_before>package again

import (
	"fmt"
	"io"
	"reflect"
)

func jsonToJson(r io.Reader, w io.Writer) {
	turnBothCranks(
		NewJsonDecoder(r),
		NewJsonEncoder(w),
	)
}

func turnBothCranks(tokenSrc TokenSrc, tokenSink TokenSink) error {
	var tok Token
	var srcDone, sinkDone bool
	var err error
	for {
		srcDone, err = tokenSrc.Step(&tok)
		if err != nil {
			return err
		}
		sinkDone, err = tokenSink.Step(&tok)
		if err != nil {
			return err
		}
		if srcDone {
			if sinkDone {
				return nil
			}
			return fmt.Errorf("src at end of item but sink expects more")
		}
	}
}

/*
	Fill with address of primitive (or []byte), or the magic const tokens
	for beginning and ending of maps and arrays.

	Decoder implementations are encouraged to use `util.DecodeBag` to contain
	primitives during decode, then return the address of the relevant
	primitive field from the `DecodeBag` as a `Token`.  This avoids repeated
	pointer allocations.
*/
type Token interface{}

const (
	Token_MapOpen  = '{'
	Token_MapClose = '}'
	Token_ArrOpen  = '['
	Token_ArrClose = ']'
)

type TokenSrc interface {
	Step(fillme *Token) (done bool, err error)
	//Reset()
}

type TokenSink interface {
	Step(consume *Token) (done bool, err error)
	//Reset()
}

//
// Constructors
//

func NewJsonDecoder(r io.Reader /* optional *JsonSchemaNotes */) TokenSrc  { return nil }
func NewJsonEncoder(w io.Writer /* optional *JsonSchemaNotes */) TokenSink { return nil }

func NewVarTokenizer(v interface{} /* TODO visitmagicks */) TokenSrc { return nil }
func NewVarReceiver(v interface{} /* TODO visitmagicks */) TokenSink {
	vr := &varReceiver{}
	vr.step = vr.stepFor(v)
	return vr
}

type varReceiver struct {
	step func(*Token) (done bool, err error)
	done bool
	err  error
}

func (vr *varReceiver) Step(tok *Token) (done bool, err error) {
	return vr.step(tok)
}

// used at initialization to figure out the first step given the type of var
//
func (vr *varReceiver) stepFor(v interface{}) func(*Token) (done bool, err error) {
	switch v.(type) {
	// For total wildcards:
	//  Return a machine that will pick between a literal or `map[string]interface{}`
	//  or `[]interface{}` based on the next token.
	case *interface{}:
		return wildcardStep(v)
	// For single literals:
	//  we have a single machine that handles all these.
	case *string, *[]byte,
		*int, *int8, *int16, *int32, *int64,
		*uint, *uint8, *uint16, *uint32, *uint64:
		dec := &literalDecoderMachine{}
		dec.Reset(v)
		return dec.Step
	// Anything that has real type info:
	//  ... Plaaaay ball!
	default:
		// TODO mustAddressable check goes here.
		if reflect.TypeOf(v).Kind() == reflect.Interface {
			// special path because we can recycle the decoder machines, if they implement resettable.
		}
		// any other concrete type or particular interface:
		//  must have its own visit func defined.
		//  we don't know if it expects to be a map, lit, arr, etc until it takes over.
		//  (the rest of our functions here are the exception: they're half inlined here -- TODO maybe don't be like that; this lookup only makes sense for top level wtf-is-this'es)
		panic("TODO mappersuite lookup")
	}
}

func wildcardStep(target interface{}) func(*Token) (bool, error) {
	return func(tok *Token) (done bool, err error) {
		// If it's a special state, start an object.
		//  (Or, blow up if its a special state that's silly).
		switch *tok {
		case Token_MapOpen:
			// Fill in our wildcard ref with a blank map,
			//  and make a new machine for it; hand off everything.
			target = make(map[string]interface{})
			dec := &wildcardMapDecoderMachine{}
			dec.Reset(target)
			return dec.Step(tok)
		case Token_ArrOpen:
			// TODO same as maps, but with a machine for arrays
			panic("NYI")
		case Token_MapClose:
			return true, fmt.Errorf("unexpected mapClose; expected start of value")
		case Token_ArrClose:
			return true, fmt.Errorf("unexpected arrClose; expected start of value")
		default:
			// If it wasn't the start of composite, shell out to the machine for literals.
			dec := &literalDecoderMachine{}
			dec.Reset(target)
			return dec.Step(tok)
		}
	}
}

type wildcardMapDecoderMachine struct {
	target map[string]interface{}
	step   func(*Token) (done bool, err error)
	key    string // The key consumed by the prev `step_AcceptKey`.
}

func (dm *wildcardMapDecoderMachine) Reset(target interface{}) {
	dm.target = target.(map[string]interface{})
	dm.step = dm.step_Initial
	dm.key = ""
}

func (dm *wildcardMapDecoderMachine) Step(tok *Token) (done bool, err error) {
	return dm.step(tok)
}

func (dm *wildcardMapDecoderMachine) step_Initial(tok *Token) (done bool, err error) {
	// If it's a special state, start an object.
	//  (Or, blow up if its a special state that's silly).
	switch *tok {
	case Token_MapOpen:
		// Great.  Consumed.
		dm.step = dm.step_AcceptKey
		return false, nil
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected start of map")
	case Token_MapClose:
		return true, fmt.Errorf("unexpected mapClose; expected start of map")
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected start of map")
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of map", *tok))
	}
}
func (dm *wildcardMapDecoderMachine) step_AcceptKey(tok *Token) (done bool, err error) {
	switch *tok {
	case Token_MapOpen:
		return true, fmt.Errorf("unexpected mapOpen; expected map key")
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected map key")
	case Token_MapClose:
		// no special checks for ends of wildcard map; no such thing as incomplete.
		return true, nil
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected map key")
	}
	switch k := (*tok).(type) {
	case *string:
		if err = dm.mustAcceptKey(*k); err != nil {
			return true, err
		}
		dm.key = *k
		dm.step = dm.step_AcceptValue
		return false, nil
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok))
	}
}
func (dm *wildcardMapDecoderMachine) mustAcceptKey(k string) error {
	if _, exists := dm.target[k]; exists {
		return fmt.Errorf("repeated key %q", k)
	}
	return nil
}

func (dm *wildcardMapDecoderMachine) step_AcceptValue(tok *Token) (done bool, err error) {
	/*
		driver.Fill(
			tok, // still meant for next person and the real step is to come; we just had to figure out types, here.
			dm.Addr(dm.key),
			dm.step_postValue(), // driver returns to us after the value is done by calling this.
			    // may actually be that we stash that stepfunc, and give driver more general self pointer and Resume func in interface.
		)
	*/
	return false, nil // TODO
}

type literalDecoderMachine struct {
	target interface{}
}

func (dm *literalDecoderMachine) Reset(target interface{}) {
	dm.target = target
}

func (dm *literalDecoderMachine) Step(tok *Token) (done bool, err error) {
	var ok bool
	switch v2 := dm.target.(type) {
	case *string:
		*v2, ok = (*tok).(string)
	case *[]byte:
		panic("TODO")
	case *int:
		*v2, ok = (*tok).(int)
	case *int8, *int16, *int32, *int64:
		panic("TODO")
	case *uint, *uint8, *uint16, *uint32, *uint64:
		panic("TODO")
	default:
		panic(fmt.Errorf("cannot unmarshall into unhandled type %T", dm.target))
	}
	if ok {
		return true, nil
	}
	return true, fmt.Errorf("unexpected token of type %T, expected literal of type %T", *tok, dm.target)
}

/*
	Suppose we have the following var to unmarshal into:

		var thingy SomeType

	Where SomeType is defined as:

		type SomeType struct {
			AnInt int
			Something interface{}
		}

	The flow of a VarReciever working on this will be something like the following:

		- Begin handling a var of type `SomeType`.
		- Look up the hander for that type info.
		- The handler is accepts the val ref, and returns a step function.
		- The step function is called with the token.
		- [Much work ensues.]
		- If the step function returns done, we return entirely;
		  otherwise we hang onto the next stepFunc, and return.

	The flow of the specific handler for SomeType will look like this:

		- Expect a MapOpen token.
		- Expect a MapKey token.  Return a step func expecting that matching value.
		  - When called with the next token, this step func grabs the ref
		    of the struct field matching the name we were primed with...
		  - And calls dispatch on the whole thing.
		  - (Generally this func looks like it needs {fillingName string, rest},
		    so it can tell what value grab the ref to fill, and decide whether
			to return "expect all done" step.)
		- At any point, it may receive MapClose, which will jump to a check
		  that all fields are either noted as filled (requires sidebar) or
		  are tagged as omitEmpty.
*/

// Returns an atlas so we can use this to build the contin-passing machine without bothering you.
func HandleMe(vreal interface{}) (
	vmediate interface{},
	atl *Atlas,
	after func(), /* closure, already has vreal and vmediate refs */
) {
	return nil, nil, nil
}

type Atlas struct{}

type atlasDecoderMachine struct {
	val      reflect.Value // We're filling this.
	atl      *Atlas        // Our directions.
	step     func(*Token)  // The next step.
	key      string        // The key consumed by the prev `step_AcceptKey`.
	keysDone []string      // List of keys we've completed already (repeats aren't wanted).
}

func NewAtlasDecoderMachine(into reflect.Value, atl *Atlas) *atlasDecoderMachine {
	// TODO this return type should prob have some interface that covers it sufficiently.
	dm := &atlasDecoderMachine{
		atl: atl,
	}
	dm.Reset(into)
	return dm
}

func (dm *atlasDecoderMachine) Reset(into reflect.Value) {
	dm.val = into
	dm.step = dm.step_Initial
	dm.key = ""
	dm.keysDone = dm.keysDone[0:0]
}

func (dm *atlasDecoderMachine) step_Initial(tok *Token) {
	switch *tok {
	case Token_MapOpen:
		// Great.  Consumed.
		dm.step = dm.step_AcceptKey
	case Token_ArrOpen:
		panic("unexpected arrOpen; expected start of struct")
	case Token_MapClose:
		panic("unexpected mapClose; expected start of struct")
	case Token_ArrClose:
		panic("unexpected arrClose; expected start of struct")
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok))
	}
}

func (dm *atlasDecoderMachine) step_AcceptKey(tok *Token) {
	switch *tok {
	case Token_MapOpen:
		panic("unexpected mapOpen; expected map key")
	case Token_ArrOpen:
		panic("unexpected arrOpen; expected map key")
	case Token_MapClose:
		dm.handleEnd()
	case Token_ArrClose:
		panic("unexpected arrClose; expected map key")
	}
	switch k := (*tok).(type) {
	case *string:
		dm.key = *k
		dm.mustAcceptKey(*k)
		//dm.step = dm.step_AcceptValue
		// actually we might wanna just push up our plea now --
		//  this saves us from having to see and forward the token at all,
		//  and makes the pattern of fab-var-filler, ret step func(token) consistent.
		//  if you *really* wanted to implement a breakout for known prims, you could still do that branch here.
		//  HANG ON, nope: keep it in the value step and keep the tok passdown.
		//   do it for parity with arrays, which must have that step
		//   and accept that token during it so they can check for end there.
		/*
			driver.Fill(
				tok, // still meant for next person and the real step is to come; we just had to figure out types, here.
				dm.Addr(dm.key),
				dm.step_postValue(), // driver returns to us after the value is done by calling this.
				    // may actually be that we stash that stepfunc, and give driver more general self pointer and Resume func in interface.
			)
		*/
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok))
	}
}
func (dm *atlasDecoderMachine) mustAcceptKey(k string) {
	for _, x := range dm.keysDone {
		if x == k {
			panic(fmt.Errorf("repeated key %q", k))
		}
	}
	dm.keysDone = append(dm.keysDone, k)
}
func (dm *atlasDecoderMachine) addr(k string) interface{} {
	_ = dm.atl
	return nil // TODO
	// n.b. this is one of the spots where i can't decide if &thing or reflect.Value is better
	// but either way we may want to define a `Slot` type alias to make it readable
}

func (dm *atlasDecoderMachine) step_AcceptValue(tok *Token) {
}

func (dm *atlasDecoderMachine) handleEnd() {
	// TODO check for all filled, etc.  then set terminal states.
}
<commit_msg>Changes all steps to new VarUnmarshalStep type, which consistently takes the top level driver as a param so we can support recursing.<commit_after>package again

import (
	"fmt"
	"io"
	"reflect"
)

func jsonToJson(r io.Reader, w io.Writer) {
	turnBothCranks(
		NewJsonDecoder(r),
		NewJsonEncoder(w),
	)
}

func turnBothCranks(tokenSrc TokenSrc, tokenSink TokenSink) error {
	var tok Token
	var srcDone, sinkDone bool
	var err error
	for {
		srcDone, err = tokenSrc.Step(&tok)
		if err != nil {
			return err
		}
		sinkDone, err = tokenSink.Step(&tok)
		if err != nil {
			return err
		}
		if srcDone {
			if sinkDone {
				return nil
			}
			return fmt.Errorf("src at end of item but sink expects more")
		}
	}
}

/*
	Fill with address of primitive (or []byte), or the magic const tokens
	for beginning and ending of maps and arrays.

	Decoder implementations are encouraged to use `util.DecodeBag` to contain
	primitives during decode, then return the address of the relevant
	primitive field from the `DecodeBag` as a `Token`.  This avoids repeated
	pointer allocations.
*/
type Token interface{}

const (
	Token_MapOpen  = '{'
	Token_MapClose = '}'
	Token_ArrOpen  = '['
	Token_ArrClose = ']'
)

type TokenSrc interface {
	Step(fillme *Token) (done bool, err error)
	//Reset()
}

type TokenSink interface {
	Step(consume *Token) (done bool, err error)
	//Reset()
}

//
// Constructors
//

func NewJsonDecoder(r io.Reader /* optional *JsonSchemaNotes */) TokenSrc  { return nil }
func NewJsonEncoder(w io.Writer /* optional *JsonSchemaNotes */) TokenSink { return nil }

func NewVarTokenizer(v interface{} /* TODO visitmagicks */) TokenSrc { return nil }
func NewVarReceiver(v interface{} /* TODO visitmagicks */) TokenSink {
	vr := &VarUnmarshalDriver{}
	vr.stepStack = []VarUnmarshalStep{
		stepFor(v),
	}
	return vr
}

//
// VarUnmarshal
//

type VarUnmarshalDriver struct {
	stepStack []VarUnmarshalStep
}

func (vr *VarUnmarshalDriver) Step(tok *Token) (done bool, err error) {
	nSteps := len(vr.stepStack) - 1
	step := vr.stepStack[nSteps]
	done, err = step(vr, tok)
	if nSteps == 1 {
		return // that's all folks
	}
	if err != nil {
		return true, err
	}
	vr.stepStack = vr.stepStack[0:nSteps]
	return false, nil
}

type VarUnmarshalStep func(*VarUnmarshalDriver, *Token) (done bool, err error)

/*
	Fills `v`,
	first looking up the machine for that type just like it's a new top-level object,
	then pushing the first step with `tok` (the upstream tends to have peeked at it
	in order to decide what to do, but if recursing, it belongs to the next obj),
	then continuing to use this new machine until it returns a done status,
	then finally setting the overall next step to `continueWith`.

	In other words, your decoder machine calls this when it wants to deal with
	an object, and by the time we call back to your `continueWith` state,
	that object will be filled and the stream ready for you to continue.

	The unmarshal driver keeps a stack of `continueWith` step funcs
	to support "recursion".  Your calls will never actually see increases in
	goroutine stack depth.
*/
func (vr *VarUnmarshalDriver) Recurse(tok *Token, v interface{}, continueWith VarUnmarshalStep) error {
	vr.stepStack = append(vr.stepStack, continueWith) // FIXME replace something
	// TODO call something
	return nil
}

// used at initialization to figure out the first step given the type of var
func stepFor(v interface{}) VarUnmarshalStep {
	switch v.(type) {
	// For total wildcards:
	//  Return a machine that will pick between a literal or `map[string]interface{}`
	//  or `[]interface{}` based on the next token.
	case *interface{}:
		return wildcardStep(v)
	// For single literals:
	//  we have a single machine that handles all these.
	case *string, *[]byte,
		*int, *int8, *int16, *int32, *int64,
		*uint, *uint8, *uint16, *uint32, *uint64:
		dec := &literalDecoderMachine{}
		dec.Reset(v)
		return dec.Step
	// Anything that has real type info:
	//  ... Plaaaay ball!
	default:
		// TODO mustAddressable check goes here.
		if reflect.TypeOf(v).Kind() == reflect.Interface {
			// special path because we can recycle the decoder machines, if they implement resettable.
		}
		// any other concrete type or particular interface:
		//  must have its own visit func defined.
		//  we don't know if it expects to be a map, lit, arr, etc until it takes over.
		//  (the rest of our functions here are the exception: they're half inlined here -- TODO maybe don't be like that; this lookup only makes sense for top level wtf-is-this'es)
		panic("TODO mappersuite lookup")
	}
}

func wildcardStep(target interface{}) VarUnmarshalStep {
	return func(vr *VarUnmarshalDriver, tok *Token) (done bool, err error) {
		// If it's a special state, start an object.
		//  (Or, blow up if its a special state that's silly).
		switch *tok {
		case Token_MapOpen:
			// Fill in our wildcard ref with a blank map,
			//  and make a new machine for it; hand off everything.
			target = make(map[string]interface{})
			dec := &wildcardMapDecoderMachine{}
			dec.Reset(target)
			return dec.Step(vr, tok)
		case Token_ArrOpen:
			// TODO same as maps, but with a machine for arrays
			panic("NYI")
		case Token_MapClose:
			return true, fmt.Errorf("unexpected mapClose; expected start of value")
		case Token_ArrClose:
			return true, fmt.Errorf("unexpected arrClose; expected start of value")
		default:
			// If it wasn't the start of composite, shell out to the machine for literals.
			dec := &literalDecoderMachine{}
			dec.Reset(target)
			return dec.Step(vr, tok)
		}
	}
}

type wildcardMapDecoderMachine struct {
	target map[string]interface{}
	step   VarUnmarshalStep
	key    string // The key consumed by the prev `step_AcceptKey`.
}

func (dm *wildcardMapDecoderMachine) Reset(target interface{}) {
	dm.target = target.(map[string]interface{})
	dm.step = dm.step_Initial
	dm.key = ""
}

func (dm *wildcardMapDecoderMachine) Step(vr *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	return dm.step(vr, tok)
}

func (dm *wildcardMapDecoderMachine) step_Initial(_ *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	// If it's a special state, start an object.
	//  (Or, blow up if its a special state that's silly).
	switch *tok {
	case Token_MapOpen:
		// Great.  Consumed.
		dm.step = dm.step_AcceptKey
		return false, nil
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected start of map")
	case Token_MapClose:
		return true, fmt.Errorf("unexpected mapClose; expected start of map")
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected start of map")
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of map", *tok))
	}
}
func (dm *wildcardMapDecoderMachine) step_AcceptKey(_ *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	switch *tok {
	case Token_MapOpen:
		return true, fmt.Errorf("unexpected mapOpen; expected map key")
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected map key")
	case Token_MapClose:
		// no special checks for ends of wildcard map; no such thing as incomplete.
		return true, nil
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected map key")
	}
	switch k := (*tok).(type) {
	case *string:
		if err = dm.mustAcceptKey(*k); err != nil {
			return true, err
		}
		dm.key = *k
		dm.step = dm.step_AcceptValue
		return false, nil
	default:
		panic(fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok))
	}
}
func (dm *wildcardMapDecoderMachine) mustAcceptKey(k string) error {
	if _, exists := dm.target[k]; exists {
		return fmt.Errorf("repeated key %q", k)
	}
	return nil
}

func (dm *wildcardMapDecoderMachine) step_AcceptValue(driver *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	var v interface{}
	dm.target[dm.key] = v
	return false, driver.Recurse(
		tok,
		&v,
		dm.step_AcceptKey,
	)
}

type literalDecoderMachine struct {
	target interface{}
}

func (dm *literalDecoderMachine) Reset(target interface{}) {
	dm.target = target
}

func (dm *literalDecoderMachine) Step(_ *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	var ok bool
	switch v2 := dm.target.(type) {
	case *string:
		*v2, ok = (*tok).(string)
	case *[]byte:
		panic("TODO")
	case *int:
		*v2, ok = (*tok).(int)
	case *int8, *int16, *int32, *int64:
		panic("TODO")
	case *uint, *uint8, *uint16, *uint32, *uint64:
		panic("TODO")
	default:
		panic(fmt.Errorf("cannot unmarshall into unhandled type %T", dm.target))
	}
	if ok {
		return true, nil
	}
	return true, fmt.Errorf("unexpected token of type %T, expected literal of type %T", *tok, dm.target)
}

/*
	Suppose we have the following var to unmarshal into:

		var thingy SomeType

	Where SomeType is defined as:

		type SomeType struct {
			AnInt int
			Something interface{}
		}

	The flow of a VarReciever working on this will be something like the following:

		- Begin handling a var of type `SomeType`.
		- Look up the hander for that type info.
		- The handler is accepts the val ref, and returns a step function.
		- The step function is called with the token.
		- [Much work ensues.]
		- If the step function returns done, we return entirely;
		  otherwise we hang onto the next stepFunc, and return.

	The flow of the specific handler for SomeType will look like this:

		- Expect a MapOpen token.
		- Expect a MapKey token.  Return a step func expecting that matching value.
		  - When called with the next token, this step func grabs the ref
		    of the struct field matching the name we were primed with...
		  - And calls dispatch on the whole thing.
		  - (Generally this func looks like it needs {fillingName string, rest},
		    so it can tell what value grab the ref to fill, and decide whether
			to return "expect all done" step.)
		- At any point, it may receive MapClose, which will jump to a check
		  that all fields are either noted as filled (requires sidebar) or
		  are tagged as omitEmpty.
*/

// Returns an atlas so we can use this to build the contin-passing machine without bothering you.
func HandleMe(vreal interface{}) (
	vmediate interface{},
	atl *Atlas,
	after func(), /* closure, already has vreal and vmediate refs */
) {
	return nil, nil, nil
}

type Atlas struct{}

type atlasDecoderMachine struct {
	val      reflect.Value    // We're filling this.
	atl      *Atlas           // Our directions.
	step     VarUnmarshalStep // The next step.
	key      string           // The key consumed by the prev `step_AcceptKey`.
	keysDone []string         // List of keys we've completed already (repeats aren't wanted).
}

func NewAtlasDecoderMachine(into reflect.Value, atl *Atlas) *atlasDecoderMachine {
	// TODO this return type should prob have some interface that covers it sufficiently.
	dm := &atlasDecoderMachine{
		atl: atl,
	}
	dm.Reset(into)
	return dm
}

func (dm *atlasDecoderMachine) Reset(into reflect.Value) {
	dm.val = into
	dm.step = dm.step_Initial
	dm.key = ""
	dm.keysDone = dm.keysDone[0:0]
}

func (dm *atlasDecoderMachine) step_Initial(_ *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	switch *tok {
	case Token_MapOpen:
		// Great.  Consumed.
		dm.step = dm.step_AcceptKey
		return false, nil
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected start of struct")
	case Token_MapClose:
		return true, fmt.Errorf("unexpected mapClose; expected start of struct")
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected start of struct")
	default:
		return true, fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok)
	}
}

func (dm *atlasDecoderMachine) step_AcceptKey(_ *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	switch *tok {
	case Token_MapOpen:
		return true, fmt.Errorf("unexpected mapOpen; expected map key")
	case Token_ArrOpen:
		return true, fmt.Errorf("unexpected arrOpen; expected map key")
	case Token_MapClose:
		return true, dm.handleEnd()
	case Token_ArrClose:
		return true, fmt.Errorf("unexpected arrClose; expected map key")
	}
	switch k := (*tok).(type) {
	case *string:
		dm.key = *k
		if err = dm.mustAcceptKey(*k); err != nil {
			return true, err
		}
		dm.step = dm.step_AcceptValue
		return false, nil
	default:
		return true, fmt.Errorf("unexpected literal of type %T; expected start of struct", *tok)
	}
}
func (dm *atlasDecoderMachine) mustAcceptKey(k string) error {
	for _, x := range dm.keysDone {
		if x == k {
			return fmt.Errorf("repeated key %q", k)
		}
	}
	dm.keysDone = append(dm.keysDone, k)
	return nil
}
func (dm *atlasDecoderMachine) addr(k string) interface{} {
	_ = dm.atl
	return nil // TODO
}

func (dm *atlasDecoderMachine) step_AcceptValue(driver *VarUnmarshalDriver, tok *Token) (done bool, err error) {
	return false, driver.Recurse(
		tok,
		dm.addr(dm.key),
		dm.step_AcceptKey,
	)
}

func (dm *atlasDecoderMachine) handleEnd() error {
	// TODO check for all filled, etc.
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Google Inc. All rights reserved.

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.
*/

// Command agent forwards requests from an inverting proxy to a backend server.
//
// To build, run:
//
//    $ make
//
// And to use, run:
//
//    $ $(GOPATH)/bin/proxy-forwarding-agent -proxy <proxy-url> -backend <backend-ID>

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/golang/groupcache/lru"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"

	"github.com/google/inverting-proxy/agent/utils"
	"github.com/google/inverting-proxy/agent/websockets"
)

const (
	requestCacheLimit = 1000
	emailScope        = "email"

	maxBackoffDuration = 100 * time.Millisecond
)

var (
	proxy                = flag.String("proxy", "", "URL (including scheme) of the inverting proxy")
	proxyTimeout         = flag.Duration("proxy-timeout", 60*time.Second, "Client timeout when sending requests to the inverting proxy")
	host                 = flag.String("host", "localhost:8080", "Hostname (including port) of the backend server")
	backendID            = flag.String("backend", "", "Unique ID for this backend.")
	debug                = flag.Bool("debug", false, "Whether or not to print debug log messages")
	forwardUserID        = flag.Bool("forward-user-id", false, "Whether or not to include the ID (email address) of the end user in requests to the backend")
	shimWebsockets       = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
	shimPath             = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
	healthCheckPath      = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
	healthCheckFreq      = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
	healthCheckUnhealthy = flag.Int("health-check-unhealthy-threshold", 2, "A so-far healthy backend will be marked unhealthy after this many consecutive failures. The minimum value is 1.")
)

func hostProxy(ctx context.Context, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	hostProxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   host,
	})
	hostProxy.FlushInterval = 100 * time.Millisecond
	if shimPath == "" {
		return hostProxy, nil
	}
	return websockets.Proxy(ctx, hostProxy, host, shimPath, injectShimCode)
}

// forwardRequest forwards the given request from the proxy to
// the backend server and reports the response back to the proxy.
func forwardRequest(client *http.Client, hostProxy http.Handler, request *utils.ForwardedRequest) error {
	httpRequest := request.Contents
	if *forwardUserID {
		httpRequest.Header.Add(utils.HeaderUserID, request.User)
	}
	responseForwarder, err := utils.NewResponseForwarder(client, *proxy, request.BackendID, request.RequestID)
	if err != nil {
		return fmt.Errorf("failed to create the response forwarder: %v", err)
	}
	hostProxy.ServeHTTP(responseForwarder, httpRequest)
	if *debug {
		log.Printf("Backend latency for request %s: %s\n", request.RequestID, time.Since(request.StartTime).String())
	}
	if err := responseForwarder.Close(); err != nil {
		return fmt.Errorf("failed to close the response forwarder: %v", err)
	}
	return nil
}

// healthCheck issues a health check against the backend server
// and returns the result.
func healthCheck() error {
	resp, err := http.Get("http://" + *host + *healthCheckPath)
	if err != nil {
		log.Printf("Health Check request failed: %s", err.Error())
		return err
	}
	if resp.StatusCode != 200 {
		log.Printf("Health Check request had non-200 status code: %d", resp.StatusCode)
		return fmt.Errorf("Bad Health Check Response Code: %s", resp.Status)
	}
	return nil
}

// processOneRequest reads a single request from the proxy and forwards it to the backend server.
func processOneRequest(client *http.Client, hostProxy http.Handler, backendID string, requestID string) {
	requestForwarder := func(client *http.Client, request *utils.ForwardedRequest) error {
		if err := forwardRequest(client, hostProxy, request); err != nil {
			log.Printf("Failure forwarding a request: [%s] %q\n", requestID, err.Error())
			return fmt.Errorf("failed to forward the request %q: %v", requestID, err)
		}
		return nil
	}
	if err := utils.ReadRequest(client, *proxy, backendID, requestID, requestForwarder); err != nil {
		log.Printf("Failed to forward a request: [%s] %q\n", requestID, err.Error())
	}
}

func exponentialBackoffDuration(retryCount uint) time.Duration {
	targetDuration := (1 << retryCount) * time.Millisecond
	if targetDuration > maxBackoffDuration {
		return maxBackoffDuration
	}
	return targetDuration
}

// pollForNewRequests repeatedly reaches out to the proxy server to ask if any pending are available, and then
// processes any newly-seen ones.
func pollForNewRequests(client *http.Client, hostProxy http.Handler, backendID string) {
	previouslySeenRequests := lru.New(requestCacheLimit)
	var retryCount uint
	for {
		if requests, err := utils.ListPendingRequests(client, *proxy, backendID); err != nil {
			log.Printf("Failed to read pending requests: %q\n", err.Error())
			time.Sleep(exponentialBackoffDuration(retryCount))
			retryCount++
		} else {
			retryCount = 0
			for _, requestID := range requests {
				if _, ok := previouslySeenRequests.Get(requestID); !ok {
					previouslySeenRequests.Add(requestID, requestID)
					go processOneRequest(client, hostProxy, backendID, requestID)
				}
			}
		}
	}
}

func getGoogleClient(ctx context.Context) (*http.Client, error) {
	sdkConfig, err := google.NewSDKConfig("")
	if err == nil {
		return sdkConfig.Client(ctx), nil
	}

	client, err := google.DefaultClient(ctx, compute.CloudPlatformScope, emailScope)
	if err != nil {
		return nil, err
	}
	client.Transport = utils.RoundTripperWithVMIdentity(ctx, client.Transport, *proxy)
	return client, nil
}

// waitForHealthy runs health checks against the backend and returns
// the first time it sees a healthy check.
func waitForHealthy() {
	if *healthCheckFreq <= 0 {
		return
	}
	if healthCheck() == nil {
		return
	}
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	for _ = range ticker.C {
		if healthCheck() == nil {
			ticker.Stop()
			return
		}
	}
}

// runHealthChecks runs health checks against the backend and shuts down
// the proxy if the backend is unhealthy.
func runHealthChecks() {
	if *healthCheckFreq <= 0 {
		return
	}
	if *healthCheckUnhealthy < 1 {
		*healthCheckUnhealthy = 1
	}
	// Always start in the unhealthy state, but only require a single positive
	// health check to become healthy for the first time, and do the first check
	// immediately.
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	badHealthChecks := 0
	for _ = range ticker.C {
		if healthCheck() != nil {
			badHealthChecks++
		} else {
			badHealthChecks = 0
		}
		if badHealthChecks >= *healthCheckUnhealthy {
			ticker.Stop()
			log.Fatal("Too many unhealthy checks")
		}
	}
}

// runAdapter sets up the HTTP client for the agent to use (including OAuth credentials),
// and then does the actual work of forwarding requests and responses.
func runAdapter() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client, err := getGoogleClient(ctx)
	if err != nil {
		return err
	}
	client.Timeout = *proxyTimeout

	hostProxy, err := hostProxy(ctx, *host, *shimPath, *shimWebsockets)
	if err != nil {
		return err
	}
	pollForNewRequests(client, hostProxy, *backendID)
	return nil
}

func main() {
	flag.Parse()

	if *proxy == "" {
		log.Fatal("You must specify the address of the proxy")
	}
	if *backendID == "" {
		log.Fatal("You must specify a backend ID")
	}
	if !strings.HasPrefix(*healthCheckPath, "/") {
		*healthCheckPath = "/" + *healthCheckPath
	}

	waitForHealthy()
	go runHealthChecks()

	if err := runAdapter(); err != nil {
		log.Fatal(err.Error())
	}
}
<commit_msg>Fix exponential backoff logic<commit_after>/*
Copyright 2016 Google Inc. All rights reserved.

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.
*/

// Command agent forwards requests from an inverting proxy to a backend server.
//
// To build, run:
//
//    $ make
//
// And to use, run:
//
//    $ $(GOPATH)/bin/proxy-forwarding-agent -proxy <proxy-url> -backend <backend-ID>

package main

import (
	"flag"
	"fmt"
	"log"
	"math"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
	"time"

	"github.com/golang/groupcache/lru"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"
	compute "google.golang.org/api/compute/v1"

	"github.com/google/inverting-proxy/agent/utils"
	"github.com/google/inverting-proxy/agent/websockets"
)

const (
	requestCacheLimit      = 1000
	emailScope             = "email"
	maxBackoffDuration     = 3 * time.Second
	firstRetryWaitDuration = time.Millisecond
)

var (
	proxy                = flag.String("proxy", "", "URL (including scheme) of the inverting proxy")
	proxyTimeout         = flag.Duration("proxy-timeout", 60*time.Second, "Client timeout when sending requests to the inverting proxy")
	host                 = flag.String("host", "localhost:8080", "Hostname (including port) of the backend server")
	backendID            = flag.String("backend", "", "Unique ID for this backend.")
	debug                = flag.Bool("debug", false, "Whether or not to print debug log messages")
	forwardUserID        = flag.Bool("forward-user-id", false, "Whether or not to include the ID (email address) of the end user in requests to the backend")
	shimWebsockets       = flag.Bool("shim-websockets", false, "Whether or not to replace websockets with a shim")
	shimPath             = flag.String("shim-path", "", "Path under which to handle websocket shim requests")
	healthCheckPath      = flag.String("health-check-path", "/", "Path on backend host to issue health checks against.  Defaults to the root.")
	healthCheckFreq      = flag.Int("health-check-interval-seconds", 0, "Wait time in seconds between health checks.  Set to zero to disable health checks.  Checks disabled by default.")
	healthCheckUnhealthy = flag.Int("health-check-unhealthy-threshold", 2, "A so-far healthy backend will be marked unhealthy after this many consecutive failures. The minimum value is 1.")
)

var (
	// compute the max retry count
	maxRetryCount = math.Log2(float64(maxBackoffDuration / firstRetryWaitDuration))
)

func hostProxy(ctx context.Context, host, shimPath string, injectShimCode bool) (http.Handler, error) {
	hostProxy := httputil.NewSingleHostReverseProxy(&url.URL{
		Scheme: "http",
		Host:   host,
	})
	hostProxy.FlushInterval = 100 * time.Millisecond
	if shimPath == "" {
		return hostProxy, nil
	}
	return websockets.Proxy(ctx, hostProxy, host, shimPath, injectShimCode)
}

// forwardRequest forwards the given request from the proxy to
// the backend server and reports the response back to the proxy.
func forwardRequest(client *http.Client, hostProxy http.Handler, request *utils.ForwardedRequest) error {
	httpRequest := request.Contents
	if *forwardUserID {
		httpRequest.Header.Add(utils.HeaderUserID, request.User)
	}
	responseForwarder, err := utils.NewResponseForwarder(client, *proxy, request.BackendID, request.RequestID)
	if err != nil {
		return fmt.Errorf("failed to create the response forwarder: %v", err)
	}
	hostProxy.ServeHTTP(responseForwarder, httpRequest)
	if *debug {
		log.Printf("Backend latency for request %s: %s\n", request.RequestID, time.Since(request.StartTime).String())
	}
	if err := responseForwarder.Close(); err != nil {
		return fmt.Errorf("failed to close the response forwarder: %v", err)
	}
	return nil
}

// healthCheck issues a health check against the backend server
// and returns the result.
func healthCheck() error {
	resp, err := http.Get("http://" + *host + *healthCheckPath)
	if err != nil {
		log.Printf("Health Check request failed: %s", err.Error())
		return err
	}
	if resp.StatusCode != 200 {
		log.Printf("Health Check request had non-200 status code: %d", resp.StatusCode)
		return fmt.Errorf("Bad Health Check Response Code: %s", resp.Status)
	}
	return nil
}

// processOneRequest reads a single request from the proxy and forwards it to the backend server.
func processOneRequest(client *http.Client, hostProxy http.Handler, backendID string, requestID string) {
	requestForwarder := func(client *http.Client, request *utils.ForwardedRequest) error {
		if err := forwardRequest(client, hostProxy, request); err != nil {
			log.Printf("Failure forwarding a request: [%s] %q\n", requestID, err.Error())
			return fmt.Errorf("failed to forward the request %q: %v", requestID, err)
		}
		return nil
	}
	if err := utils.ReadRequest(client, *proxy, backendID, requestID, requestForwarder); err != nil {
		log.Printf("Failed to forward a request: [%s] %q\n", requestID, err.Error())
	}
}

func exponentialBackoffDuration(retryCount uint) time.Duration {
	if retryCount > uint(maxRetryCount) {
		return maxBackoffDuration
	}

	targetDuration := (1 << retryCount) * firstRetryWaitDuration
	return targetDuration
}

// pollForNewRequests repeatedly reaches out to the proxy server to ask if any pending are available, and then
// processes any newly-seen ones.
func pollForNewRequests(client *http.Client, hostProxy http.Handler, backendID string) {
	previouslySeenRequests := lru.New(requestCacheLimit)
	var retryCount uint
	for {
		if requests, err := utils.ListPendingRequests(client, *proxy, backendID); err != nil {
			log.Printf("Failed to read pending requests: %q\n", err.Error())
			time.Sleep(exponentialBackoffDuration(retryCount))
			retryCount++
		} else {
			retryCount = 0
			for _, requestID := range requests {
				if _, ok := previouslySeenRequests.Get(requestID); !ok {
					previouslySeenRequests.Add(requestID, requestID)
					go processOneRequest(client, hostProxy, backendID, requestID)
				}
			}
		}
	}
}

func getGoogleClient(ctx context.Context) (*http.Client, error) {
	sdkConfig, err := google.NewSDKConfig("")
	if err == nil {
		return sdkConfig.Client(ctx), nil
	}

	client, err := google.DefaultClient(ctx, compute.CloudPlatformScope, emailScope)
	if err != nil {
		return nil, err
	}
	client.Transport = utils.RoundTripperWithVMIdentity(ctx, client.Transport, *proxy)
	return client, nil
}

// waitForHealthy runs health checks against the backend and returns
// the first time it sees a healthy check.
func waitForHealthy() {
	if *healthCheckFreq <= 0 {
		return
	}
	if healthCheck() == nil {
		return
	}
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	for _ = range ticker.C {
		if healthCheck() == nil {
			ticker.Stop()
			return
		}
	}
}

// runHealthChecks runs health checks against the backend and shuts down
// the proxy if the backend is unhealthy.
func runHealthChecks() {
	if *healthCheckFreq <= 0 {
		return
	}
	if *healthCheckUnhealthy < 1 {
		*healthCheckUnhealthy = 1
	}
	// Always start in the unhealthy state, but only require a single positive
	// health check to become healthy for the first time, and do the first check
	// immediately.
	ticker := time.NewTicker(time.Duration(*healthCheckFreq) * time.Second)
	badHealthChecks := 0
	for _ = range ticker.C {
		if healthCheck() != nil {
			badHealthChecks++
		} else {
			badHealthChecks = 0
		}
		if badHealthChecks >= *healthCheckUnhealthy {
			ticker.Stop()
			log.Fatal("Too many unhealthy checks")
		}
	}
}

// runAdapter sets up the HTTP client for the agent to use (including OAuth credentials),
// and then does the actual work of forwarding requests and responses.
func runAdapter() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client, err := getGoogleClient(ctx)
	if err != nil {
		return err
	}
	client.Timeout = *proxyTimeout

	hostProxy, err := hostProxy(ctx, *host, *shimPath, *shimWebsockets)
	if err != nil {
		return err
	}
	pollForNewRequests(client, hostProxy, *backendID)
	return nil
}

func main() {
	flag.Parse()

	if *proxy == "" {
		log.Fatal("You must specify the address of the proxy")
	}
	if *backendID == "" {
		log.Fatal("You must specify a backend ID")
	}
	if !strings.HasPrefix(*healthCheckPath, "/") {
		*healthCheckPath = "/" + *healthCheckPath
	}

	waitForHealthy()
	go runHealthChecks()

	if err := runAdapter(); err != nil {
		log.Fatal(err.Error())
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"encoding/json"
	"path"
	"sync"

	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/job"
)

type AgentState struct {
	// used to lock the datastructure for multi-goroutine safety
	mutex sync.Mutex

	// unresolved job offers
	offers map[string]job.JobOffer

	// job names for which a bid has been submitted
	bids []string

	// reverse index of peers that would cause a reassesment of a JobOffer this
	// Agent could not have bid on previously
	// i.e. {"hello.service": ["howareyou.service", "goodbye.service"]}
	peers map[string][]string

	// index of local payload conflicts to the job they belong to
	conflicts map[string][]string
}

func NewState() *AgentState {
	return &AgentState{
		offers:    make(map[string]job.JobOffer),
		bids:      make([]string, 0),
		peers:     make(map[string][]string),
		conflicts: make(map[string][]string, 0),
	}
}

func (self *AgentState) Lock() {
	log.V(2).Infof("Attempting to lock AgentState")
	self.mutex.Lock()
	log.V(2).Infof("AgentState locked")
}

func (self *AgentState) Unlock() {
	log.V(2).Infof("Attempting to unlock AgentState")
	self.mutex.Unlock()
	log.V(2).Infof("AgentState unlocked")
}

func (self *AgentState) MarshalJSON() ([]byte, error) {
	type ds struct {
		Offers    map[string]job.JobOffer
		Conflicts map[string][]string
		Bids      []string
		Peers     map[string][]string
	}
	data := ds{
		Offers:    self.offers,
		Conflicts: self.conflicts,
		Bids:      self.bids,
		Peers:     self.peers,
	}
	return json.Marshal(data)
}

// Store a list of conflicts on behalf of a given Job
func (self *AgentState) TrackJobConflicts(jobName string, conflicts []string) {
	self.conflicts[jobName] = conflicts
}

// Determine whether there are any known conflicts with the given argument
func (self *AgentState) HasConflict(potentialJobName string, potentialConflicts []string) (bool, string) {
	// Iterate through each existing Job, asserting two things:
	for existingJobName, existingConflicts := range self.conflicts {

		// 1. Each tracked Job does not conflict with the potential conflicts
		for _, pc := range potentialConflicts {
			if globMatches(pc, existingJobName) {
				return true, existingJobName
			}
		}

		// 2. The new Job does not conflict with any of the tracked confclits
		for _, ec := range existingConflicts {
			if globMatches(ec, potentialJobName) {
				return true, existingJobName
			}
		}
	}

	return false, ""
}

// Purge all tracked conflicts for a given Job
func (self *AgentState) DropJobConflicts(jobName string) {
	delete(self.conflicts, jobName)
}

// Store a relation of 1 Job -> N Peers
func (self *AgentState) TrackJobPeers(jobName string, peers []string) {
	for _, peer := range peers {
		_, ok := self.peers[peer]
		if !ok {
			self.peers[peer] = make([]string, 0)
		}
		self.peers[peer] = append(self.peers[peer], jobName)
	}
}

// Retrieve all Jobs that share a given Peer
func (self *AgentState) GetJobsByPeer(peerName string) []string {
	peers, ok := self.peers[peerName]
	if ok {
		return peers
	} else {
		return make([]string, 0)
	}
}

// Remove all references to a given Job from all Peer indexes
func (self *AgentState) DropPeersJob(jobName string) {
	for peer, peerIndex := range self.peers {
		var idxs []int

		// Determine which item indexes must be removed from the Peer index
		for idx, record := range peerIndex {
			if jobName == record {
				idxs = append(idxs, idx)
			}
		}

		// Iterate through the item indexes, removing the corresponding Peers
		for i, idx := range idxs {
			self.peers[peer] = append(self.peers[peer][0:idx-i], self.peers[peer][idx-i+1:]...)
		}
	}
}

func (self *AgentState) TrackOffer(offer job.JobOffer) {
	self.offers[offer.Job.Name] = offer
}

// GetOffersWithoutBids returns all tracked JobOffers that have
// no corresponding JobBid tracked in the same AgentState object.
func (self *AgentState) GetOffersWithoutBids() []job.JobOffer {
	offers := make([]job.JobOffer, 0)
	for _, offer := range self.offers {
		if !self.HasBid(offer.Job.Name) {
			offers = append(offers, offer)
		}
	}
	return offers
}

func (self *AgentState) GetOffer(name string) (job.JobOffer, bool) {
	offer, ok := self.offers[name]
	return offer, ok
}

func (self *AgentState) DropOffer(name string) {
	if _, ok := self.offers[name]; !ok {
		log.V(2).Infof("AgentState knows nothing of JobOffer(%s)", name)
		return
	}

	delete(self.offers, name)
}

func (self *AgentState) TrackBid(name string) {
	self.bids = append(self.bids, name)
}

func (self *AgentState) HasBid(name string) bool {
	for _, val := range self.bids {
		if val == name {
			return true
		}
	}
	return false
}

func (self *AgentState) DropBid(name string) {
	for idx, val := range self.bids {
		if val == name {
			self.bids = append(self.bids[0:idx], self.bids[idx+1:]...)
			return
		}
	}
}

func globMatches(pattern, target string) bool {
	matched, err := path.Match(pattern, target)
	if err != nil {
		log.V(2).Infof("Received error while matching pattern '%s': %v", pattern, err)
	}
	return matched
}
<commit_msg>fix(agent): clear empty entries when dropping peers<commit_after>package agent

import (
	"encoding/json"
	"path"
	"sync"

	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/job"
)

type AgentState struct {
	// used to lock the datastructure for multi-goroutine safety
	mutex sync.Mutex

	// unresolved job offers
	offers map[string]job.JobOffer

	// job names for which a bid has been submitted
	bids []string

	// reverse index of peers that would cause a reassesment of a JobOffer this
	// Agent could not have bid on previously
	// i.e. {"hello.service": ["howareyou.service", "goodbye.service"]}
	peers map[string][]string

	// index of local payload conflicts to the job they belong to
	conflicts map[string][]string
}

func NewState() *AgentState {
	return &AgentState{
		offers:    make(map[string]job.JobOffer),
		bids:      make([]string, 0),
		peers:     make(map[string][]string),
		conflicts: make(map[string][]string, 0),
	}
}

func (self *AgentState) Lock() {
	log.V(2).Infof("Attempting to lock AgentState")
	self.mutex.Lock()
	log.V(2).Infof("AgentState locked")
}

func (self *AgentState) Unlock() {
	log.V(2).Infof("Attempting to unlock AgentState")
	self.mutex.Unlock()
	log.V(2).Infof("AgentState unlocked")
}

func (self *AgentState) MarshalJSON() ([]byte, error) {
	type ds struct {
		Offers    map[string]job.JobOffer
		Conflicts map[string][]string
		Bids      []string
		Peers     map[string][]string
	}
	data := ds{
		Offers:    self.offers,
		Conflicts: self.conflicts,
		Bids:      self.bids,
		Peers:     self.peers,
	}
	return json.Marshal(data)
}

// Store a list of conflicts on behalf of a given Job
func (self *AgentState) TrackJobConflicts(jobName string, conflicts []string) {
	self.conflicts[jobName] = conflicts
}

// Determine whether there are any known conflicts with the given argument
func (self *AgentState) HasConflict(potentialJobName string, potentialConflicts []string) (bool, string) {
	// Iterate through each existing Job, asserting two things:
	for existingJobName, existingConflicts := range self.conflicts {

		// 1. Each tracked Job does not conflict with the potential conflicts
		for _, pc := range potentialConflicts {
			if globMatches(pc, existingJobName) {
				return true, existingJobName
			}
		}

		// 2. The new Job does not conflict with any of the tracked confclits
		for _, ec := range existingConflicts {
			if globMatches(ec, potentialJobName) {
				return true, existingJobName
			}
		}
	}

	return false, ""
}

// Purge all tracked conflicts for a given Job
func (self *AgentState) DropJobConflicts(jobName string) {
	delete(self.conflicts, jobName)
}

// Store a relation of 1 Job -> N Peers
func (self *AgentState) TrackJobPeers(jobName string, peers []string) {
	for _, peer := range peers {
		_, ok := self.peers[peer]
		if !ok {
			self.peers[peer] = make([]string, 0)
		}
		self.peers[peer] = append(self.peers[peer], jobName)
	}
}

// Retrieve all Jobs that share a given Peer
func (self *AgentState) GetJobsByPeer(peerName string) []string {
	peers, ok := self.peers[peerName]
	if ok {
		return peers
	} else {
		return make([]string, 0)
	}
}

// Remove all references to a given Job from all Peer indexes
func (self *AgentState) DropPeersJob(jobName string) {
	for peer, peerIndex := range self.peers {
		var idxs []int

		// Determine which item indexes must be removed from the Peer index
		for idx, record := range peerIndex {
			if jobName == record {
				idxs = append(idxs, idx)
			}
		}

		// Iterate through the item indexes, removing the corresponding Peers
		for i, idx := range idxs {
			self.peers[peer] = append(self.peers[peer][0:idx-i], self.peers[peer][idx-i+1:]...)
		}

		// Clean up empty peer relations when possible
		if len(self.peers[peer]) == 0 {
			delete(self.peers, peer)
		}
	}
}

func (self *AgentState) TrackOffer(offer job.JobOffer) {
	self.offers[offer.Job.Name] = offer
}

// GetOffersWithoutBids returns all tracked JobOffers that have
// no corresponding JobBid tracked in the same AgentState object.
func (self *AgentState) GetOffersWithoutBids() []job.JobOffer {
	offers := make([]job.JobOffer, 0)
	for _, offer := range self.offers {
		if !self.HasBid(offer.Job.Name) {
			offers = append(offers, offer)
		}
	}
	return offers
}

func (self *AgentState) GetOffer(name string) (job.JobOffer, bool) {
	offer, ok := self.offers[name]
	return offer, ok
}

func (self *AgentState) DropOffer(name string) {
	if _, ok := self.offers[name]; !ok {
		log.V(2).Infof("AgentState knows nothing of JobOffer(%s)", name)
		return
	}

	delete(self.offers, name)
}

func (self *AgentState) TrackBid(name string) {
	self.bids = append(self.bids, name)
}

func (self *AgentState) HasBid(name string) bool {
	for _, val := range self.bids {
		if val == name {
			return true
		}
	}
	return false
}

func (self *AgentState) DropBid(name string) {
	for idx, val := range self.bids {
		if val == name {
			self.bids = append(self.bids[0:idx], self.bids[idx+1:]...)
			return
		}
	}
}

func globMatches(pattern, target string) bool {
	matched, err := path.Match(pattern, target)
	if err != nil {
		log.V(2).Infof("Received error while matching pattern '%s': %v", pattern, err)
	}
	return matched
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/peterh/liner"
	"github.com/pkg/errors"

	"github.com/johnny-morrice/godless/api"
	"github.com/johnny-morrice/godless/crdt"
	"github.com/johnny-morrice/godless/crypto"
	"github.com/johnny-morrice/godless/log"
	"github.com/johnny-morrice/godless/query"
)

type TerminalOptions struct {
	Client api.Client
}

func RunTerminalConsole(options TerminalOptions) error {
	line := liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)
	for {
		queryText, err := line.Prompt(">")

		if err != nil {
			return err
		}

		query, err := query.CompileQuery(queryText)

		if err != nil {
			fmt.Printf("Compiliation error: %v", err.Error())
			continue
		}

		resp, err := options.Client.SendQuery(query)

		if err != nil {
			fmt.Printf("Error: %v", err.Error())

			if resp.Msg != "" {
				fmt.Printf("Response message: %v", resp.Msg)
			}

			if resp.Err != nil && resp.Err != err {
				fmt.Printf("Response error: %v", resp.Err.Error())
			}
		}

		printResponseTables(resp)
	}

	return nil
}

func printResponseTables(resp api.APIResponse) {
	printNamespaceTables(resp.Namespace)
	printIndexTables(resp.Index)
	printPath(resp.Path)
}

func printIndexTables(index crdt.Index) {
	if index.IsEmpty() {
		return
	}

	table := makeIndexTable(index)
	table.fprint(os.Stdout)
}

func printNamespaceTables(namespace crdt.Namespace) {
	if namespace.IsEmpty() {
		return
	}

	table := makeNamespaceTable(namespace)
	table.fprint(os.Stdout)
}

func printPath(path crdt.IPFSPath) {
	if !crdt.IsNilPath(path) {
		fmt.Println(path)
	}
}

type monospaceTable struct {
	columns      []string
	columnWidths []int
	rows         [][]string
	frozen       bool
	totalWidth   int
}

func (table *monospaceTable) fprint(w io.Writer) {
	table.printline(w)
	table.fprintrow(w, table.columns)
	table.printline(w)

	for _, r := range table.rows {
		table.fprintrow(w, r)
	}
	table.printline(w)
}

func (table *monospaceTable) gettotalwidth() int {
	if table.totalWidth == 0 {
		for _, w := range table.columnWidths {
			table.totalWidth += w
			table.totalWidth += __VALUE_PADDING_COUNT
		}

		table.totalWidth += __ROW_PADDING_COUNT
	}

	return table.totalWidth
}

func (table *monospaceTable) printline(w io.Writer) {
	total := table.gettotalwidth()
	table.repeat(w, "-", total)

	table.newline(w)
}

func (table *monospaceTable) repeat(w io.Writer, text string, count int) {
	for i := 0; i < count; i++ {
		table.write(w, text)
	}
}

func (table *monospaceTable) fprintrow(w io.Writer, row []string) {
	for i, value := range row {
		table.write(w, "| ")
		table.fprintvalue(w, i, value)
		table.write(w, " ")
	}

	table.write(w, " |")
	table.newline(w)
}

func (table *monospaceTable) newline(w io.Writer) {
	table.write(w, "\n")
}

func (table *monospaceTable) write(w io.Writer, text string) {
	w.Write([]byte(text))
}

func (table *monospaceTable) fprintvalue(w io.Writer, columnIndex int, value string) {
	table.write(w, value)
	padding := table.padsize(columnIndex, value)
	table.repeat(w, " ", padding)
}

func (table *monospaceTable) padsize(columnIndex int, text string) int {
	columnWidth := table.columnWidths[columnIndex]
	return columnWidth - len(text)
}

func (table *monospaceTable) addColumn(column string) error {
	if table.frozen {
		return errors.New("don't addColumn after addRow")
	}

	table.columns = append(table.columns, column)
	table.columnWidths = append(table.columnWidths, 0)
	index := len(table.columns) - 1
	table.recalcWidth(index, column)
	return nil
}

func (table *monospaceTable) recalcWidth(columnIndex int, addition string) {
	current := table.columnWidths[columnIndex]
	newWidth := len(addition)
	if current < newWidth {
		table.columnWidths[columnIndex] = newWidth
	}
}

func (table *monospaceTable) addRow(values []string) error {
	if len(values) != len(table.columns) {
		return errors.New("row length does not match column length")
	}

	table.frozen = true

	table.rows = append(table.rows, values)

	for i, value := range values {
		table.recalcWidth(i, value)
	}

	return nil
}

func makeIndexTable(index crdt.Index) *monospaceTable {
	panic("not implemented")
}

func makeNamespaceTable(namespace crdt.Namespace) *monospaceTable {
	table := &monospaceTable{}
	columns := []string{
		"Table",
		"Row",
		"Entry",
		"Point",
		"Signatures",
	}

	for _, c := range columns {
		table.addColumn(c)
	}

	namespace.ForeachEntry(func(t crdt.TableName, r crdt.RowName, e crdt.EntryName, entry crdt.Entry) {
		for _, point := range entry.GetValues() {
			sigText := makeSigText(point.Signatures())
			row := []string{
				string(t),
				string(r),
				string(e),
				string(point.Text()),
				sigText,
			}

			table.addRow(row)
		}
	})

	return table
}

func makeSigText(signatures []crypto.Signature) string {
	sigText := make([]string, 0, len(signatures))

	for _, sig := range signatures {
		text, err := crypto.PrintSignature(sig)

		if err != nil {
			log.Error("Error serializing crypto.Signature: %v", err)
			continue
		}

		sigText = append(sigText, string(text))
	}

	return strings.Join(sigText, " ")
}

const __VALUE_PADDING_COUNT = 3
const __ROW_PADDING_COUNT = 2
<commit_msg>Console enhancements<commit_after>package cli

import (
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/peterh/liner"
	"github.com/pkg/errors"

	"github.com/johnny-morrice/godless/api"
	"github.com/johnny-morrice/godless/crdt"
	"github.com/johnny-morrice/godless/crypto"
	"github.com/johnny-morrice/godless/log"
	"github.com/johnny-morrice/godless/query"
)

type TerminalOptions struct {
	Client api.Client
}

func RunTerminalConsole(options TerminalOptions) error {
	line := liner.NewLiner()
	defer line.Close()

	line.SetCtrlCAborts(true)
	for {
		queryText, err := line.Prompt("> ")

		if err != nil {
			return err
		}

		line.AppendHistory(queryText)

		query, err := query.CompileQuery(queryText)

		if err != nil {
			fmt.Printf("Compiliation error: %v", err.Error())
			continue
		}

		resp, err := options.Client.SendQuery(query)

		if err != nil {
			fmt.Printf("Error: %v", err.Error())

			if resp.Msg != "" {
				fmt.Printf("Response message: %v", resp.Msg)
			}

			if resp.Err != nil && resp.Err != err {
				fmt.Printf("Response error: %v", resp.Err.Error())
			}
		}

		printResponseTables(resp)
	}

	return nil
}

func printResponseTables(resp api.APIResponse) {
	printNamespaceTables(resp.Namespace)
	printIndexTables(resp.Index)
	printPath(resp.Path)
}

func printIndexTables(index crdt.Index) {
	if index.IsEmpty() {
		return
	}

	table := makeIndexTable(index)
	table.fprint(os.Stdout)
}

func printNamespaceTables(namespace crdt.Namespace) {
	if namespace.IsEmpty() {
		fmt.Println("No results returned")
		return
	}

	table := makeNamespaceTable(namespace)
	table.fprint(os.Stdout)
}

func printPath(path crdt.IPFSPath) {
	if !crdt.IsNilPath(path) {
		fmt.Println(path)
	}
}

type monospaceTable struct {
	columns      []string
	columnWidths []int
	rows         [][]string
	frozen       bool
	totalWidth   int
}

func (table *monospaceTable) fprint(w io.Writer) {
	table.printline(w)
	table.fprintrow(w, table.columns)
	table.printline(w)

	for _, r := range table.rows {
		table.fprintrow(w, r)
	}
	table.printline(w)
}

func (table *monospaceTable) gettotalwidth() int {
	if table.totalWidth == 0 {
		for _, w := range table.columnWidths {
			table.totalWidth += w
			table.totalWidth += __VALUE_PADDING_COUNT
		}

		table.totalWidth += __ROW_PADDING_COUNT
	}

	return table.totalWidth
}

func (table *monospaceTable) printline(w io.Writer) {
	total := table.gettotalwidth()
	table.repeat(w, "-", total)

	table.newline(w)
}

func (table *monospaceTable) repeat(w io.Writer, text string, count int) {
	for i := 0; i < count; i++ {
		table.write(w, text)
	}
}

func (table *monospaceTable) fprintrow(w io.Writer, row []string) {
	for i, value := range row {
		table.write(w, "| ")
		table.fprintvalue(w, i, value)
		table.write(w, " ")
	}

	table.write(w, " |")
	table.newline(w)
}

func (table *monospaceTable) newline(w io.Writer) {
	table.write(w, "\n")
}

func (table *monospaceTable) write(w io.Writer, text string) {
	w.Write([]byte(text))
}

func (table *monospaceTable) fprintvalue(w io.Writer, columnIndex int, value string) {
	table.write(w, value)
	padding := table.padsize(columnIndex, value)
	table.repeat(w, " ", padding)
}

func (table *monospaceTable) padsize(columnIndex int, text string) int {
	columnWidth := table.columnWidths[columnIndex]
	return columnWidth - len(text)
}

func (table *monospaceTable) addColumn(column string) error {
	if table.frozen {
		return errors.New("don't addColumn after addRow")
	}

	table.columns = append(table.columns, column)
	table.columnWidths = append(table.columnWidths, 0)
	index := len(table.columns) - 1
	table.recalcWidth(index, column)
	return nil
}

func (table *monospaceTable) recalcWidth(columnIndex int, addition string) {
	current := table.columnWidths[columnIndex]
	newWidth := len(addition)
	if current < newWidth {
		table.columnWidths[columnIndex] = newWidth
	}
}

func (table *monospaceTable) addRow(values []string) error {
	if len(values) != len(table.columns) {
		return errors.New("row length does not match column length")
	}

	table.frozen = true

	table.rows = append(table.rows, values)

	for i, value := range values {
		table.recalcWidth(i, value)
	}

	return nil
}

func makeIndexTable(index crdt.Index) *monospaceTable {
	panic("not implemented")
}

// TODO figure out how to make signatures look nice
func makeNamespaceTable(namespace crdt.Namespace) *monospaceTable {
	table := &monospaceTable{}
	columns := []string{
		"Table",
		"Row",
		"Entry",
		"Point",
		// "Signatures",
	}

	for _, c := range columns {
		table.addColumn(c)
	}

	namespace.ForeachEntry(func(t crdt.TableName, r crdt.RowName, e crdt.EntryName, entry crdt.Entry) {
		for _, point := range entry.GetValues() {
			// sigText := makeSigText(point.Signatures())
			row := []string{
				string(t),
				string(r),
				string(e),
				string(point.Text()),
				// sigText,
			}

			table.addRow(row)
		}
	})

	return table
}

func makeSigText(signatures []crypto.Signature) string {
	sigText := make([]string, 0, len(signatures))

	for _, sig := range signatures {
		text, err := crypto.PrintSignature(sig)

		if err != nil {
			log.Error("Error serializing crypto.Signature: %v", err)
			continue
		}

		sigText = append(sigText, string(text))
	}

	return strings.Join(sigText, " ")
}

const __VALUE_PADDING_COUNT = 3
const __ROW_PADDING_COUNT = 2
<|endoftext|>


<commit_before>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package client

import (
	"html/template"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	"github.com/caixw/gitype/data"
	"github.com/caixw/gitype/helper"
	"github.com/caixw/gitype/vars"
	"github.com/issue9/logs"
	"github.com/issue9/utils"
)

const (
	contentTypeKey = "Content-Type"
	cookieKey      = "Set-Cookie"
)

// 生成一个带编码的 content-type 报头内容
func buildContentTypeContent(mime string) string {
	return mime + ";charset=utf-8"
}

// 设置页面的编码，若已经存在，则不会受影响。
// 要强制指定，请直接使用 w.Header().Set()
func setContentType(w http.ResponseWriter, mime string) {
	h := w.Header()
	if len(h.Get(contentTypeKey)) == 0 {
		h.Set(contentTypeKey, buildContentTypeContent(mime))
	}
}

// 用于描述一个页面的所有无素
type page struct {
	client   *Client
	Info     *info
	template *template.Template // 用于当前页面渲染的模板
	response http.ResponseWriter
	request  *http.Request

	Title       string       // 文章标题，可以为空
	Subtitle    string       // 副标题
	Canonical   string       // 当前页的唯一链接
	Keywords    string       // meta.keywords 的值
	Description string       // meta.description 的值
	PrevPage    *data.Link   // 前一页
	NextPage    *data.Link   // 下一页
	Type        string       // 当前页面类型
	Author      *data.Author // 作者
	License     *data.Link   // 当前页的版本信息，可以为空
	Theme       *data.Theme

	// 以下内容，仅在对应的页面才会有内容
	Q        string          // 搜索关键字
	Tag      *data.Tag       // 标签详细页面，非标签详细页，则为空
	Posts    []*data.Post    // 文章列表，仅标签详情页和搜索页用到。
	Post     *data.Post      // 文章详细内容，仅文章页面用到。
	Archives []*data.Archive // 归档
}

// 页面的附加信息，除非重新加载数据，否则内容不会变。
type info struct {
	AppName    string // 程序名称
	AppURL     string // 程序官网
	AppVersion string // 当前程序的版本号
	GoVersion  string // 编译的 Go 版本号

	SiteName    string     // 网站名称
	URL         string     // 网站地址，若是一个子目录，则需要包含该子目录
	Icon        *data.Icon // 网站图标
	Language    string     // 页面语言
	PostSize    int        // 总文章数量
	Beian       string     // 备案号
	Uptime      time.Time  // 上线时间
	LastUpdated time.Time  // 最后更新时间
	RSS         *data.Link // RSS，NOTICE:指针方便模板判断其值是否为空
	Atom        *data.Link
	Opensearch  *data.Link
	Tags        []*data.Tag  // 标签列表
	Series      []*data.Tag  // 专题列表
	Links       []*data.Link // 友情链接
	Menus       []*data.Link // 导航菜单
}

func (client *Client) newInfo() *info {
	d := client.data

	info := &info{
		AppName:    vars.Name,
		AppURL:     vars.URL,
		AppVersion: vars.Version(),
		GoVersion:  runtime.Version(),

		SiteName:    d.Title,
		URL:         d.URL,
		Icon:        d.Icon,
		Language:    d.Language,
		PostSize:    len(d.Posts),
		Beian:       d.Beian,
		Uptime:      d.Uptime,
		LastUpdated: d.Created,
		Tags:        d.Tags,
		Series:      d.Series,
		Links:       d.Links,
		Menus:       d.Menus,
	}

	if d.RSS != nil {
		info.RSS = &data.Link{
			Title: d.RSS.Title,
			URL:   d.RSS.URL,
			Type:  d.RSS.Type,
		}
	}

	if d.Atom != nil {
		info.Atom = &data.Link{
			Title: d.Atom.Title,
			URL:   d.Atom.URL,
			Type:  d.Atom.Type,
		}
	}

	if d.Opensearch != nil {
		info.Opensearch = &data.Link{
			Title: d.Opensearch.Title,
			URL:   d.Opensearch.URL,
			Type:  d.Opensearch.Type,
		}
	}

	return info
}

func (client *Client) page(typ string, w http.ResponseWriter, r *http.Request) *page {
	theme := client.getRequestTheme(r)
	d := client.data

	return &page{
		client:   client,
		Info:     client.info,
		template: theme.Template,
		response: w,
		request:  r,

		Subtitle:    d.Subtitle,
		Keywords:    d.Keywords,
		Description: d.Description,
		Type:        typ,
		Author:      d.Author,
		License:     d.License,
		Theme:       theme,
	}
}

func (p *page) nextPage(url, text string) {
	if len(text) == 0 {
		text = vars.NextPageText
	}

	p.NextPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "next",
	}
}

func (p *page) prevPage(url, text string) {
	if len(text) == 0 {
		text = vars.PrevPageText
	}

	p.PrevPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "prev",
	}
}

// 输出当前内容到指定模板
func (p *page) render(name string) {
	setContentType(p.response, p.client.data.Type)

	cookie := &http.Cookie{
		Name:     vars.CookieKeyTheme,
		Value:    p.Theme.ID,
		HttpOnly: vars.CookieHTTPOnly,
	}
	if p.Theme.ID != p.client.data.Themes[0].ID {
		cookie.MaxAge = vars.CookieMaxAge
	} else {
		cookie.MaxAge = -1
	}
	cookie.Expires = time.Now().Add(time.Second * time.Duration(vars.CookieMaxAge))
	p.response.Header().Add(cookieKey, cookie.String())

	err := p.template.ExecuteTemplate(p.response, name, p)
	if err != nil {
		logs.Error(err)
		p.client.renderError(p.response, p.request, http.StatusInternalServerError)
		return
	}
}

// 从客户端获取主题内容
func (client *Client) getRequestTheme(r *http.Request) *data.Theme {
	// 获取主题名称
	name := r.FormValue(vars.CookieKeyTheme)
	if len(name) == 0 {
		cookie, err := r.Cookie(vars.CookieKeyTheme)
		if err != nil { // 有记录错误，但不退出
			logs.Error(err)
		}

		if cookie != nil {
			name = cookie.Value
		}
	}

	// 查询对应名称的主题
	for _, t := range client.data.Themes {
		if name == t.ID {
			return t
		}
	}

	return client.data.Themes[0] // 不存在的情况下，返回默认主题
}

// 输出一个特定状态码下的错误页面。
// 若该页面模板不存在，则输出状态码对应的文本内容。
// 只查找当前主题目录下的相关文件。
// 只对状态码大于等于 400 的起作用。
func (client *Client) renderError(w http.ResponseWriter, r *http.Request, code int) {
	if code < 400 {
		return
	}
	logs.Debug("输出非正常状态码：", code)

	// 根据情况输出内容，若不存在模板，则直接输出最简单的状态码对应的文本。
	theme := client.getRequestTheme(r)
	filename := strconv.Itoa(code) + vars.TemplateExtension
	path := filepath.Join(client.path.ThemesDir, theme.ID, filename)
	if !utils.FileExists(path) {
		logs.Debugf("模板文件 %s 不存在\n", path)
		helper.StatusError(w, code)
		return
	}

	data, err := ioutil.ReadFile(path)
	if err != nil {
		logs.Errorf("读取模板文件 %s 时出现以下错误: %v\n", path, err)
		helper.StatusError(w, code)
		return
	}

	setContentType(w, client.data.Type)
	w.WriteHeader(code)
	w.Write(data)
}
<commit_msg>使用 SetCookie 代替 Header().Set() 函数<commit_after>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package client

import (
	"html/template"
	"io/ioutil"
	"net/http"
	"path/filepath"
	"runtime"
	"strconv"
	"time"

	"github.com/caixw/gitype/data"
	"github.com/caixw/gitype/helper"
	"github.com/caixw/gitype/vars"
	"github.com/issue9/logs"
	"github.com/issue9/utils"
)

const contentTypeKey = "Content-Type"

// 生成一个带编码的 content-type 报头内容
func buildContentTypeContent(mime string) string {
	return mime + ";charset=utf-8"
}

// 设置页面的编码，若已经存在，则不会受影响。
// 要强制指定，请直接使用 w.Header().Set()
func setContentType(w http.ResponseWriter, mime string) {
	h := w.Header()
	if len(h.Get(contentTypeKey)) == 0 {
		h.Set(contentTypeKey, buildContentTypeContent(mime))
	}
}

// 用于描述一个页面的所有无素
type page struct {
	client   *Client
	Info     *info
	template *template.Template // 用于当前页面渲染的模板
	response http.ResponseWriter
	request  *http.Request

	Title       string       // 文章标题，可以为空
	Subtitle    string       // 副标题
	Canonical   string       // 当前页的唯一链接
	Keywords    string       // meta.keywords 的值
	Description string       // meta.description 的值
	PrevPage    *data.Link   // 前一页
	NextPage    *data.Link   // 下一页
	Type        string       // 当前页面类型
	Author      *data.Author // 作者
	License     *data.Link   // 当前页的版本信息，可以为空
	Theme       *data.Theme

	// 以下内容，仅在对应的页面才会有内容
	Q        string          // 搜索关键字
	Tag      *data.Tag       // 标签详细页面，非标签详细页，则为空
	Posts    []*data.Post    // 文章列表，仅标签详情页和搜索页用到。
	Post     *data.Post      // 文章详细内容，仅文章页面用到。
	Archives []*data.Archive // 归档
}

// 页面的附加信息，除非重新加载数据，否则内容不会变。
type info struct {
	AppName    string // 程序名称
	AppURL     string // 程序官网
	AppVersion string // 当前程序的版本号
	GoVersion  string // 编译的 Go 版本号

	SiteName    string     // 网站名称
	URL         string     // 网站地址，若是一个子目录，则需要包含该子目录
	Icon        *data.Icon // 网站图标
	Language    string     // 页面语言
	PostSize    int        // 总文章数量
	Beian       string     // 备案号
	Uptime      time.Time  // 上线时间
	LastUpdated time.Time  // 最后更新时间
	RSS         *data.Link // RSS，NOTICE:指针方便模板判断其值是否为空
	Atom        *data.Link
	Opensearch  *data.Link
	Tags        []*data.Tag  // 标签列表
	Series      []*data.Tag  // 专题列表
	Links       []*data.Link // 友情链接
	Menus       []*data.Link // 导航菜单
}

func (client *Client) newInfo() *info {
	d := client.data

	info := &info{
		AppName:    vars.Name,
		AppURL:     vars.URL,
		AppVersion: vars.Version(),
		GoVersion:  runtime.Version(),

		SiteName:    d.Title,
		URL:         d.URL,
		Icon:        d.Icon,
		Language:    d.Language,
		PostSize:    len(d.Posts),
		Beian:       d.Beian,
		Uptime:      d.Uptime,
		LastUpdated: d.Created,
		Tags:        d.Tags,
		Series:      d.Series,
		Links:       d.Links,
		Menus:       d.Menus,
	}

	if d.RSS != nil {
		info.RSS = &data.Link{
			Title: d.RSS.Title,
			URL:   d.RSS.URL,
			Type:  d.RSS.Type,
		}
	}

	if d.Atom != nil {
		info.Atom = &data.Link{
			Title: d.Atom.Title,
			URL:   d.Atom.URL,
			Type:  d.Atom.Type,
		}
	}

	if d.Opensearch != nil {
		info.Opensearch = &data.Link{
			Title: d.Opensearch.Title,
			URL:   d.Opensearch.URL,
			Type:  d.Opensearch.Type,
		}
	}

	return info
}

func (client *Client) page(typ string, w http.ResponseWriter, r *http.Request) *page {
	theme := client.getRequestTheme(r)
	d := client.data

	return &page{
		client:   client,
		Info:     client.info,
		template: theme.Template,
		response: w,
		request:  r,

		Subtitle:    d.Subtitle,
		Keywords:    d.Keywords,
		Description: d.Description,
		Type:        typ,
		Author:      d.Author,
		License:     d.License,
		Theme:       theme,
	}
}

func (p *page) nextPage(url, text string) {
	if len(text) == 0 {
		text = vars.NextPageText
	}

	p.NextPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "next",
	}
}

func (p *page) prevPage(url, text string) {
	if len(text) == 0 {
		text = vars.PrevPageText
	}

	p.PrevPage = &data.Link{
		Text: text,
		URL:  url,
		Rel:  "prev",
	}
}

// 输出当前内容到指定模板
func (p *page) render(name string) {
	setContentType(p.response, p.client.data.Type)

	cookie := &http.Cookie{
		Name:     vars.CookieKeyTheme,
		Value:    p.Theme.ID,
		HttpOnly: vars.CookieHTTPOnly,
	}
	if p.Theme.ID != p.client.data.Themes[0].ID {
		cookie.MaxAge = vars.CookieMaxAge
	} else {
		cookie.MaxAge = -1
	}
	cookie.Expires = time.Now().Add(time.Second * time.Duration(vars.CookieMaxAge))
	http.SetCookie(p.response, cookie)

	err := p.template.ExecuteTemplate(p.response, name, p)
	if err != nil {
		logs.Error(err)
		p.client.renderError(p.response, p.request, http.StatusInternalServerError)
		return
	}
}

// 从客户端获取主题内容
func (client *Client) getRequestTheme(r *http.Request) *data.Theme {
	// 获取主题名称
	name := r.FormValue(vars.CookieKeyTheme)
	if len(name) == 0 {
		cookie, err := r.Cookie(vars.CookieKeyTheme)
		if err != nil { // 有记录错误，但不退出
			logs.Error(err)
		}

		if cookie != nil {
			name = cookie.Value
		}
	}

	// 查询对应名称的主题
	for _, t := range client.data.Themes {
		if name == t.ID {
			return t
		}
	}

	return client.data.Themes[0] // 不存在的情况下，返回默认主题
}

// 输出一个特定状态码下的错误页面。
// 若该页面模板不存在，则输出状态码对应的文本内容。
// 只查找当前主题目录下的相关文件。
// 只对状态码大于等于 400 的起作用。
func (client *Client) renderError(w http.ResponseWriter, r *http.Request, code int) {
	if code < 400 {
		return
	}
	logs.Debug("输出非正常状态码：", code)

	// 根据情况输出内容，若不存在模板，则直接输出最简单的状态码对应的文本。
	theme := client.getRequestTheme(r)
	filename := strconv.Itoa(code) + vars.TemplateExtension
	path := filepath.Join(client.path.ThemesDir, theme.ID, filename)
	if !utils.FileExists(path) {
		logs.Debugf("模板文件 %s 不存在\n", path)
		helper.StatusError(w, code)
		return
	}

	data, err := ioutil.ReadFile(path)
	if err != nil {
		logs.Errorf("读取模板文件 %s 时出现以下错误: %v\n", path, err)
		helper.StatusError(w, code)
		return
	}

	setContentType(w, client.data.Type)
	w.WriteHeader(code)
	w.Write(data)
}
<|endoftext|>