<commit_before>package ansible

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

const (
	// RawFormat is the raw Ansible output formatting
	RawFormat = OutputFormat("raw")
	// JSONLinesFormat is a JSON Lines representation of Ansible events
	JSONLinesFormat = OutputFormat("json_lines")
)

// OutputFormat is used for controlling the STDOUT format of the Ansible runner
type OutputFormat string

// Runner for running Ansible playbooks
type Runner interface {
	// StartPlaybook runs the playbook asynchronously with the given inventory and extra vars.
	// It returns a read-only channel that must be consumed for the playbook execution to proceed.
	StartPlaybook(playbookFile string, inventory Inventory, vars ExtraVars) (<-chan Event, error)
	// WaitPlaybook blocks until the execution of the playbook is complete. If an error occurred,
	// it is returned. Otherwise, returns nil to signal the completion of the playbook.
	WaitPlaybook() error
	// StartPlaybookOnNode runs the playbook asynchronously with the given inventory and extra vars
	// against the specific node.
	// It returns a read-only channel that must be consumed for the playbook execution to proceed.
	StartPlaybookOnNode(playbookFile string, inventory Inventory, vars ExtraVars, node string) (<-chan Event, error)
}

type runner struct {
	// Out is the stdout writer for the Ansible process
	out io.Writer
	// ErrOut is the stderr writer for the Ansible process
	errOut io.Writer

	pythonPath   string
	ansibleDir   string
	waitPlaybook func() error
	namedPipe    string
}

// ExtraVars is a map of variables that are used when executing a playbook
type ExtraVars map[string]string

func (v ExtraVars) commandLineVars() (string, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return "", fmt.Errorf("error marshaling ansible vars")
	}
	return string(b), nil
}

// NewRunner returns a new runner for running Ansible playbooks.
func NewRunner(out, errOut io.Writer, ansibleDir string) (Runner, error) {
	ppath, err := getPythonPath()
	if err != nil {
		return nil, err
	}

	return &runner{
		out:        out,
		errOut:     errOut,
		pythonPath: ppath,
		ansibleDir: ansibleDir,
	}, nil
}

// WaitPlaybook blocks until the ansible process running the playbook exits.
// If the process exits with a non-zero status, it will return an error.
func (r *runner) WaitPlaybook() error {
	if r.waitPlaybook == nil {
		return fmt.Errorf("wait called, but playbook not started")
	}
	execErr := r.waitPlaybook()
	// Process exited, we can clean up named pipe
	removeErr := os.Remove(r.namedPipe)
	if removeErr != nil && execErr != nil {
		return fmt.Errorf("an error occurred running ansible: %v. Removing named pipe at %q failed: %v", execErr, r.namedPipe, removeErr)
	}
	if removeErr != nil {
		return fmt.Errorf("failed to clean up named pipe at %q: %v", r.namedPipe, removeErr)
	}
	if execErr != nil {
		return fmt.Errorf("error running ansible: %v", execErr)
	}
	return nil
}

// RunPlaybook with the given inventory and extra vars
func (r *runner) StartPlaybook(playbookFile string, inv Inventory, vars ExtraVars) (<-chan Event, error) {
	return r.startPlaybook(playbookFile, inv, vars, "") // Don't set the --limit arg
}

// StartPlaybookOnNode runs the playbook asynchronously with the given inventory and extra vars
// against the specific node.
// It returns a read-only channel that must be consumed for the playbook execution to proceed.
func (r *runner) StartPlaybookOnNode(playbookFile string, inv Inventory, vars ExtraVars, node string) (<-chan Event, error) {
	limitArg := node // set the --limit arg to the node we want to target
	return r.startPlaybook(playbookFile, inv, vars, limitArg)
}

func (r *runner) startPlaybook(playbookFile string, inv Inventory, vars ExtraVars, limitArg string) (<-chan Event, error) {
	extraVars, err := vars.commandLineVars()
	if err != nil {
		return nil, fmt.Errorf("error building extra vars: %v", err)
	}

	inventoryFile := filepath.Join(r.ansibleDir, "inventory.ini")
	if err = ioutil.WriteFile(inventoryFile, inv.ToINI(), 0644); err != nil {
		return nil, fmt.Errorf("error writing inventory file to %q: %v", inventoryFile, err)
	}

	playbook := filepath.Join(r.ansibleDir, "playbooks", playbookFile)
	cmd := exec.Command(filepath.Join(r.ansibleDir, "bin", "ansible-playbook"), "-i", inventoryFile, "-s", playbook, "--extra-vars", extraVars)
	cmd.Stdout = r.out
	cmd.Stderr = r.errOut

	if limitArg != "" {
		cmd.Args = append(cmd.Args, "--limit", limitArg)
	}

	os.Setenv("PYTHONPATH", r.pythonPath)
	os.Setenv("ANSIBLE_HOST_KEY_CHECKING", "False")
	os.Setenv("ANSIBLE_CALLBACK_PLUGINS", filepath.Join(r.ansibleDir, "playbooks", "callback"))
	os.Setenv("ANSIBLE_CALLBACK_WHITELIST", "json_lines")
	os.Setenv("TIMEOUT", "30")
	os.Setenv("ANSIBLE_SSH_ARGS", "-C -o ControlMaster=auto -o ControlPersist=1800s -o ControlPath=/tmp/ssh-%r-%h-%p")

	// We always want the most verbose output from Ansible. If it's not going to
	// stdout, it's going to a log file.
	cmd.Args = append(cmd.Args, "-vvvv")

	// Create named pipe for getting JSON lines event stream
	start := time.Now()
	r.namedPipe = filepath.Join(os.TempDir(), fmt.Sprintf("ansible-pipe-%s", start.Format("2006-01-02-15-04-05.99999")))
	if err = syscall.Mkfifo(r.namedPipe, 0644); err != nil {
		return nil, fmt.Errorf("error creating named pipe %q: %v", r.namedPipe, err)
	}
	os.Setenv("ANSIBLE_JSON_LINES_PIPE", r.namedPipe)

	// Print Ansible command
	fmt.Fprint(r.out, strings.Join(cmd.Args, " "))
	// Starts async execution of ansible, which will block until
	// we start reading from the named pipe
	err = cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("error running playbook: %v", err)
	}
	r.waitPlaybook = cmd.Wait

	// Create the event stream out of the named pipe
	eventStreamFile, err := os.OpenFile(r.namedPipe, os.O_RDWR, os.ModeNamedPipe)
	if err != nil {
		return nil, fmt.Errorf("error openning event stream pipe: %v", err)
	}
	eventStream := EventStream(eventStreamFile)
	return eventStream, nil
}

func getPythonPath() (string, error) {
	wd, err := os.Getwd()
	if err != nil {
		return "", fmt.Errorf("error getting working dir: %v", err)
	}
	lib := filepath.Join(wd, "ansible", "lib", "python2.7", "site-packages")
	lib64 := filepath.Join(wd, "ansible", "lib64", "python2.7", "site-packages")
	return fmt.Sprintf("%s:%s", lib, lib64), nil
}
<commit_msg>Remove “-C”<commit_after>package ansible

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

const (
	// RawFormat is the raw Ansible output formatting
	RawFormat = OutputFormat("raw")
	// JSONLinesFormat is a JSON Lines representation of Ansible events
	JSONLinesFormat = OutputFormat("json_lines")
)

// OutputFormat is used for controlling the STDOUT format of the Ansible runner
type OutputFormat string

// Runner for running Ansible playbooks
type Runner interface {
	// StartPlaybook runs the playbook asynchronously with the given inventory and extra vars.
	// It returns a read-only channel that must be consumed for the playbook execution to proceed.
	StartPlaybook(playbookFile string, inventory Inventory, vars ExtraVars) (<-chan Event, error)
	// WaitPlaybook blocks until the execution of the playbook is complete. If an error occurred,
	// it is returned. Otherwise, returns nil to signal the completion of the playbook.
	WaitPlaybook() error
	// StartPlaybookOnNode runs the playbook asynchronously with the given inventory and extra vars
	// against the specific node.
	// It returns a read-only channel that must be consumed for the playbook execution to proceed.
	StartPlaybookOnNode(playbookFile string, inventory Inventory, vars ExtraVars, node string) (<-chan Event, error)
}

type runner struct {
	// Out is the stdout writer for the Ansible process
	out io.Writer
	// ErrOut is the stderr writer for the Ansible process
	errOut io.Writer

	pythonPath   string
	ansibleDir   string
	waitPlaybook func() error
	namedPipe    string
}

// ExtraVars is a map of variables that are used when executing a playbook
type ExtraVars map[string]string

func (v ExtraVars) commandLineVars() (string, error) {
	b, err := json.Marshal(v)
	if err != nil {
		return "", fmt.Errorf("error marshaling ansible vars")
	}
	return string(b), nil
}

// NewRunner returns a new runner for running Ansible playbooks.
func NewRunner(out, errOut io.Writer, ansibleDir string) (Runner, error) {
	ppath, err := getPythonPath()
	if err != nil {
		return nil, err
	}

	return &runner{
		out:        out,
		errOut:     errOut,
		pythonPath: ppath,
		ansibleDir: ansibleDir,
	}, nil
}

// WaitPlaybook blocks until the ansible process running the playbook exits.
// If the process exits with a non-zero status, it will return an error.
func (r *runner) WaitPlaybook() error {
	if r.waitPlaybook == nil {
		return fmt.Errorf("wait called, but playbook not started")
	}
	execErr := r.waitPlaybook()
	// Process exited, we can clean up named pipe
	removeErr := os.Remove(r.namedPipe)
	if removeErr != nil && execErr != nil {
		return fmt.Errorf("an error occurred running ansible: %v. Removing named pipe at %q failed: %v", execErr, r.namedPipe, removeErr)
	}
	if removeErr != nil {
		return fmt.Errorf("failed to clean up named pipe at %q: %v", r.namedPipe, removeErr)
	}
	if execErr != nil {
		return fmt.Errorf("error running ansible: %v", execErr)
	}
	return nil
}

// RunPlaybook with the given inventory and extra vars
func (r *runner) StartPlaybook(playbookFile string, inv Inventory, vars ExtraVars) (<-chan Event, error) {
	return r.startPlaybook(playbookFile, inv, vars, "") // Don't set the --limit arg
}

// StartPlaybookOnNode runs the playbook asynchronously with the given inventory and extra vars
// against the specific node.
// It returns a read-only channel that must be consumed for the playbook execution to proceed.
func (r *runner) StartPlaybookOnNode(playbookFile string, inv Inventory, vars ExtraVars, node string) (<-chan Event, error) {
	limitArg := node // set the --limit arg to the node we want to target
	return r.startPlaybook(playbookFile, inv, vars, limitArg)
}

func (r *runner) startPlaybook(playbookFile string, inv Inventory, vars ExtraVars, limitArg string) (<-chan Event, error) {
	extraVars, err := vars.commandLineVars()
	if err != nil {
		return nil, fmt.Errorf("error building extra vars: %v", err)
	}

	inventoryFile := filepath.Join(r.ansibleDir, "inventory.ini")
	if err = ioutil.WriteFile(inventoryFile, inv.ToINI(), 0644); err != nil {
		return nil, fmt.Errorf("error writing inventory file to %q: %v", inventoryFile, err)
	}

	playbook := filepath.Join(r.ansibleDir, "playbooks", playbookFile)
	cmd := exec.Command(filepath.Join(r.ansibleDir, "bin", "ansible-playbook"), "-i", inventoryFile, "-s", playbook, "--extra-vars", extraVars)
	cmd.Stdout = r.out
	cmd.Stderr = r.errOut

	if limitArg != "" {
		cmd.Args = append(cmd.Args, "--limit", limitArg)
	}

	os.Setenv("PYTHONPATH", r.pythonPath)
	os.Setenv("ANSIBLE_HOST_KEY_CHECKING", "False")
	os.Setenv("ANSIBLE_CALLBACK_PLUGINS", filepath.Join(r.ansibleDir, "playbooks", "callback"))
	os.Setenv("ANSIBLE_CALLBACK_WHITELIST", "json_lines")
	os.Setenv("TIMEOUT", "30")
	os.Setenv("ANSIBLE_SSH_ARGS", "-o ControlMaster=auto -o ControlPersist=1800s -o ControlPath=/tmp/ssh-%r-%h-%p")

	// We always want the most verbose output from Ansible. If it's not going to
	// stdout, it's going to a log file.
	cmd.Args = append(cmd.Args, "-vvvv")

	// Create named pipe for getting JSON lines event stream
	start := time.Now()
	r.namedPipe = filepath.Join(os.TempDir(), fmt.Sprintf("ansible-pipe-%s", start.Format("2006-01-02-15-04-05.99999")))
	if err = syscall.Mkfifo(r.namedPipe, 0644); err != nil {
		return nil, fmt.Errorf("error creating named pipe %q: %v", r.namedPipe, err)
	}
	os.Setenv("ANSIBLE_JSON_LINES_PIPE", r.namedPipe)

	// Print Ansible command
	fmt.Fprint(r.out, strings.Join(cmd.Args, " "))
	// Starts async execution of ansible, which will block until
	// we start reading from the named pipe
	err = cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("error running playbook: %v", err)
	}
	r.waitPlaybook = cmd.Wait

	// Create the event stream out of the named pipe
	eventStreamFile, err := os.OpenFile(r.namedPipe, os.O_RDWR, os.ModeNamedPipe)
	if err != nil {
		return nil, fmt.Errorf("error openning event stream pipe: %v", err)
	}
	eventStream := EventStream(eventStreamFile)
	return eventStream, nil
}

func getPythonPath() (string, error) {
	wd, err := os.Getwd()
	if err != nil {
		return "", fmt.Errorf("error getting working dir: %v", err)
	}
	lib := filepath.Join(wd, "ansible", "lib", "python2.7", "site-packages")
	lib64 := filepath.Join(wd, "ansible", "lib64", "python2.7", "site-packages")
	return fmt.Sprintf("%s:%s", lib, lib64), nil
}
<|endoftext|>


<commit_before>// Copyright 2020 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package menu displays a Terminal UI based text menu to choose boot options
// from.
package menu

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/sh"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/sys/unix"
)

var (
	initialTimeout    = 10 * time.Second
	subsequentTimeout = 60 * time.Second
)

// Entry is a menu entry.
type Entry interface {
	// Label is the string displayed to the user in the menu.
	Label() string

	// Edit the kernel command line if possible. Must be called prior to
	// Load.
	Edit(func(cmdline string) string)

	// Load is called when the entry is chosen, but does not transfer
	// execution to another process or kernel.
	Load() error

	// Exec transfers execution to another process or kernel.
	//
	// Exec either returns an error or does not return at all.
	Exec() error

	// IsDefault indicates that this action should be run by default if the
	// user didn't make an entry choice.
	IsDefault() bool
}

func parseBootNum(choice string, entries []Entry) (int, error) {
	num, err := strconv.Atoi(choice)
	if err != nil {
		return -1, fmt.Errorf("%s is not a valid entry number: %v", choice, err)
	}
	if num < 1 || num > len(entries) {
		return -1, fmt.Errorf("%s is not a valid entry number", choice)
	}
	return num, nil
}

// Choose presents the user a menu on input to choose an entry from and returns that entry.
func Choose(input *os.File, entries ...Entry) Entry {
	fmt.Println("")
	for i, e := range entries {
		fmt.Printf("%02d. %s\n\n", i+1, e.Label())
	}
	fmt.Println("\r")

	oldState, err := terminal.MakeRaw(int(input.Fd()))
	if err != nil {
		log.Printf("BUG: Please report: We cannot actually let you choose from menu (MakeRaw failed): %v", err)
		return nil
	}
	defer terminal.Restore(int(input.Fd()), oldState)

	// TODO(chrisko): reduce this timeout a la GRUB. 3 seconds, and hitting
	// any button resets the timeout. We could save 7 seconds here.
	t := time.NewTimer(initialTimeout)

	boot := make(chan Entry, 1)

	go func() {
		// Read exactly one line.
		term := terminal.NewTerminal(input, "")

		term.AutoCompleteCallback = func(line string, pos int, key rune) (string, int, bool) {
			// We ain't gonna autocomplete, but we'll reset the countdown timer when you press a key.
			t.Reset(subsequentTimeout)
			return "", 0, false
		}

		for {
			term.SetPrompt("Enter an option ('01' is the default, 'e' to edit kernel cmdline):\n > ")
			choice, err := term.ReadLine()
			if err != nil {
				if err != io.EOF {
					fmt.Printf("BUG: Please report: Terminal read error: %v.\n", err)
				}
				boot <- nil
				return
			}

			if choice == "e" {
				// Edit command line.
				term.SetPrompt("Select a boot option to edit:\n > ")
				choice, err := term.ReadLine()
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				num, err := parseBootNum(choice, entries)
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				entries[num-1].Edit(func(cmdline string) string {
					fmt.Fprintf(term, "The current quoted cmdline for option %d is:\n > %q\n", num, cmdline)
					fmt.Fprintln(term, ` * Note the cmdline is c-style quoted. Ex: \n => newline, \\ => \`)
					term.SetPrompt("Enter an option:\n * (a)ppend, (o)verwrite, (r)eturn to main menu\n > ")
					choice, err := term.ReadLine()
					if err != nil {
						fmt.Fprintln(term, err)
						return cmdline
					}
					switch choice {
					case "a":
						term.SetPrompt("Enter unquoted cmdline to append:\n > ")
						appendCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						if appendCmdline != "" {
							cmdline += " " + appendCmdline
						}
					case "o":
						term.SetPrompt("Enter new unquoted cmdline:\n > ")
						newCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						cmdline = newCmdline
					case "r":
					default:
						fmt.Fprintf(term, "Unrecognized choice %q", choice)
					}
					fmt.Fprintf(term, "The new quoted cmdline for option %d is:\n > %q\n", num, cmdline)
					return cmdline
				})
				fmt.Fprintln(term, "Returning to main menu...")
				continue
			}
			if choice == "" {
				// nil will result in the default order.
				boot <- nil
				return
			}
			num, err := parseBootNum(choice, entries)
			if err != nil {
				fmt.Fprintln(term, err)
				continue
			}
			boot <- entries[num-1]
			return
		}
	}()

	select {
	case entry := <-boot:
		if entry != nil {
			fmt.Printf("Chosen option %s.\r\n\r\n", entry.Label())
		}
		return entry

	case <-t.C:
		return nil
	}
}

// ShowMenuAndLoad lets the user choose one of entries and loads it. If no
// entry is chosen by the user, an entry whose IsDefault() is true will be
// returned.
//
// The user is left to call Entry.Exec when this function returns.
func ShowMenuAndLoad(input *os.File, entries ...Entry) Entry {
	// Clear the screen (ANSI terminal escape code for screen clear).
	fmt.Printf("\033[1;1H\033[2J\n\n")
	fmt.Printf("Welcome to NERF's Boot Menu\n\n")
	fmt.Printf("Enter a number to boot a kernel:\n")

	for {
		// Allow the user to choose.
		entry := Choose(input, entries...)
		if entry == nil {
			// This only returns something if the user explicitly
			// entered something.
			//
			// If nothing was entered, fall back to default.
			break
		}
		if err := entry.Load(); err != nil {
			log.Printf("Failed to load %s: %v", entry.Label(), err)
			continue
		}

		// Entry was successfully loaded. Leave it to the caller to
		// exec, so the caller can clean up the OS before rebooting or
		// kexecing (e.g. unmount file systems).
		return entry
	}

	fmt.Println("")

	// We only get one shot at actually booting, so boot the first kernel
	// that can be loaded correctly.
	for _, e := range entries {
		// Only perform actions that are default actions. I.e. don't
		// drop to shell.
		if e.IsDefault() {
			fmt.Printf("Attempting to boot %s.\n\n", e)

			if err := e.Load(); err != nil {
				log.Printf("Failed to load %s: %v", e.Label(), err)
				continue
			}

			// Entry was successfully loaded. Leave it to the
			// caller to exec, so the caller can clean up the OS
			// before rebooting or kexecing (e.g. unmount file
			// systems).
			return e
		}
	}
	return nil
}

// OSImages returns menu entries for the given OSImages.
func OSImages(verbose bool, imgs ...boot.OSImage) []Entry {
	var menu []Entry
	for _, img := range imgs {
		menu = append(menu, &OSImageAction{
			OSImage: img,
			Verbose: verbose,
		})
	}
	return menu
}

// OSImageAction is a menu.Entry that boots an OSImage.
type OSImageAction struct {
	boot.OSImage
	Verbose bool
}

// Load implements Entry.Load by loading the OS image into memory.
func (oia OSImageAction) Load() error {
	if err := oia.OSImage.Load(oia.Verbose); err != nil {
		return fmt.Errorf("could not load image %s: %v", oia.OSImage, err)
	}
	return nil
}

// Exec executes the loaded image.
func (oia OSImageAction) Exec() error {
	return boot.Execute()
}

// IsDefault returns true -- this action should be performed in order by
// default if the user did not choose a boot entry.
func (OSImageAction) IsDefault() bool { return true }

// StartShell is a menu.Entry that starts a LinuxBoot shell.
type StartShell struct{}

// Label is the label to show to the user.
func (StartShell) Label() string {
	return "Enter a LinuxBoot shell"
}

// Edit does nothing.
func (StartShell) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (StartShell) Load() error {
	return nil
}

// Exec implements Entry.Exec by running /bin/defaultsh.
func (StartShell) Exec() error {
	// Reset signal handler for SIGINT to enable user interrupts again
	signal.Reset(syscall.SIGINT)
	return sh.RunWithLogs("/bin/defaultsh")
}

// IsDefault indicates that this should not be run as a default action.
func (StartShell) IsDefault() bool { return false }

// Reboot is a menu.Entry that reboots the machine.
type Reboot struct{}

// Label is the label to show to the user.
func (Reboot) Label() string {
	return "Reboot"
}

// Edit does nothing.
func (Reboot) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (Reboot) Load() error {
	return nil
}

// Exec reboots the machine using sys_reboot.
func (Reboot) Exec() error {
	unix.Sync()
	return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
}

// IsDefault indicates that this should not be run as a default action.
func (Reboot) IsDefault() bool { return false }
<commit_msg>menu: Put \r\n back into menu prints<commit_after>// Copyright 2020 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package menu displays a Terminal UI based text menu to choose boot options
// from.
package menu

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/u-root/u-root/pkg/boot"
	"github.com/u-root/u-root/pkg/sh"
	"golang.org/x/crypto/ssh/terminal"
	"golang.org/x/sys/unix"
)

var (
	initialTimeout    = 10 * time.Second
	subsequentTimeout = 60 * time.Second
)

// Entry is a menu entry.
type Entry interface {
	// Label is the string displayed to the user in the menu.
	Label() string

	// Edit the kernel command line if possible. Must be called prior to
	// Load.
	Edit(func(cmdline string) string)

	// Load is called when the entry is chosen, but does not transfer
	// execution to another process or kernel.
	Load() error

	// Exec transfers execution to another process or kernel.
	//
	// Exec either returns an error or does not return at all.
	Exec() error

	// IsDefault indicates that this action should be run by default if the
	// user didn't make an entry choice.
	IsDefault() bool
}

func parseBootNum(choice string, entries []Entry) (int, error) {
	num, err := strconv.Atoi(choice)
	if err != nil {
		return -1, fmt.Errorf("%s is not a valid entry number: %v", choice, err)
	}
	if num < 1 || num > len(entries) {
		return -1, fmt.Errorf("%s is not a valid entry number", choice)
	}
	return num, nil
}

// Choose presents the user a menu on input to choose an entry from and returns that entry.
func Choose(input *os.File, entries ...Entry) Entry {
	fmt.Println("")
	for i, e := range entries {
		fmt.Printf("%02d. %s\n\n", i+1, e.Label())
	}
	fmt.Println("\r")

	oldState, err := terminal.MakeRaw(int(input.Fd()))
	if err != nil {
		log.Printf("BUG: Please report: We cannot actually let you choose from menu (MakeRaw failed): %v", err)
		return nil
	}
	defer terminal.Restore(int(input.Fd()), oldState)

	// TODO(chrisko): reduce this timeout a la GRUB. 3 seconds, and hitting
	// any button resets the timeout. We could save 7 seconds here.
	t := time.NewTimer(initialTimeout)

	boot := make(chan Entry, 1)

	go func() {
		// Note that term is in raw mode. Write \r\n whenever you would
		// write a \n. When testing in qemu, it might look fine because
		// there might be another tty cooking the newlines. In for
		// example minicom, the behavior is different. And you would
		// see something like:
		//
		//     Select a boot option to edit:
		//                                  >
		//
		// Instead of:
		//
		//     Select a boot option to edit:
		//      >
		term := terminal.NewTerminal(input, "")

		term.AutoCompleteCallback = func(line string, pos int, key rune) (string, int, bool) {
			// We ain't gonna autocomplete, but we'll reset the countdown timer when you press a key.
			t.Reset(subsequentTimeout)
			return "", 0, false
		}

		for {
			term.SetPrompt("Enter an option ('01' is the default, 'e' to edit kernel cmdline):\r\n > ")
			choice, err := term.ReadLine()
			if err != nil {
				if err != io.EOF {
					fmt.Printf("BUG: Please report: Terminal read error: %v.\n", err)
				}
				boot <- nil
				return
			}

			if choice == "e" {
				// Edit command line.
				term.SetPrompt("Select a boot option to edit:\r\n > ")
				choice, err := term.ReadLine()
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				num, err := parseBootNum(choice, entries)
				if err != nil {
					fmt.Fprintln(term, err)
					fmt.Fprintln(term, "Returning to main menu...")
					continue
				}
				entries[num-1].Edit(func(cmdline string) string {
					fmt.Fprintf(term, "The current quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					fmt.Fprintln(term, ` * Note the cmdline is c-style quoted. Ex: \n => newline, \\ => \`)
					term.SetPrompt("Enter an option:\r\n * (a)ppend, (o)verwrite, (r)eturn to main menu\r\n > ")
					choice, err := term.ReadLine()
					if err != nil {
						fmt.Fprintln(term, err)
						return cmdline
					}
					switch choice {
					case "a":
						term.SetPrompt("Enter unquoted cmdline to append:\r\n > ")
						appendCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						if appendCmdline != "" {
							cmdline += " " + appendCmdline
						}
					case "o":
						term.SetPrompt("Enter new unquoted cmdline:\r\n > ")
						newCmdline, err := term.ReadLine()
						if err != nil {
							fmt.Fprintln(term, err)
							return cmdline
						}
						cmdline = newCmdline
					case "r":
					default:
						fmt.Fprintf(term, "Unrecognized choice %q", choice)
					}
					fmt.Fprintf(term, "The new quoted cmdline for option %d is:\r\n > %q\r\n", num, cmdline)
					return cmdline
				})
				fmt.Fprintln(term, "Returning to main menu...")
				continue
			}
			if choice == "" {
				// nil will result in the default order.
				boot <- nil
				return
			}
			num, err := parseBootNum(choice, entries)
			if err != nil {
				fmt.Fprintln(term, err)
				continue
			}
			boot <- entries[num-1]
			return
		}
	}()

	select {
	case entry := <-boot:
		if entry != nil {
			fmt.Printf("Chosen option %s.\r\n\r\n", entry.Label())
		}
		return entry

	case <-t.C:
		return nil
	}
}

// ShowMenuAndLoad lets the user choose one of entries and loads it. If no
// entry is chosen by the user, an entry whose IsDefault() is true will be
// returned.
//
// The user is left to call Entry.Exec when this function returns.
func ShowMenuAndLoad(input *os.File, entries ...Entry) Entry {
	// Clear the screen (ANSI terminal escape code for screen clear).
	fmt.Printf("\033[1;1H\033[2J\n\n")
	fmt.Printf("Welcome to NERF's Boot Menu\n\n")
	fmt.Printf("Enter a number to boot a kernel:\n")

	for {
		// Allow the user to choose.
		entry := Choose(input, entries...)
		if entry == nil {
			// This only returns something if the user explicitly
			// entered something.
			//
			// If nothing was entered, fall back to default.
			break
		}
		if err := entry.Load(); err != nil {
			log.Printf("Failed to load %s: %v", entry.Label(), err)
			continue
		}

		// Entry was successfully loaded. Leave it to the caller to
		// exec, so the caller can clean up the OS before rebooting or
		// kexecing (e.g. unmount file systems).
		return entry
	}

	fmt.Println("")

	// We only get one shot at actually booting, so boot the first kernel
	// that can be loaded correctly.
	for _, e := range entries {
		// Only perform actions that are default actions. I.e. don't
		// drop to shell.
		if e.IsDefault() {
			fmt.Printf("Attempting to boot %s.\n\n", e)

			if err := e.Load(); err != nil {
				log.Printf("Failed to load %s: %v", e.Label(), err)
				continue
			}

			// Entry was successfully loaded. Leave it to the
			// caller to exec, so the caller can clean up the OS
			// before rebooting or kexecing (e.g. unmount file
			// systems).
			return e
		}
	}
	return nil
}

// OSImages returns menu entries for the given OSImages.
func OSImages(verbose bool, imgs ...boot.OSImage) []Entry {
	var menu []Entry
	for _, img := range imgs {
		menu = append(menu, &OSImageAction{
			OSImage: img,
			Verbose: verbose,
		})
	}
	return menu
}

// OSImageAction is a menu.Entry that boots an OSImage.
type OSImageAction struct {
	boot.OSImage
	Verbose bool
}

// Load implements Entry.Load by loading the OS image into memory.
func (oia OSImageAction) Load() error {
	if err := oia.OSImage.Load(oia.Verbose); err != nil {
		return fmt.Errorf("could not load image %s: %v", oia.OSImage, err)
	}
	return nil
}

// Exec executes the loaded image.
func (oia OSImageAction) Exec() error {
	return boot.Execute()
}

// IsDefault returns true -- this action should be performed in order by
// default if the user did not choose a boot entry.
func (OSImageAction) IsDefault() bool { return true }

// StartShell is a menu.Entry that starts a LinuxBoot shell.
type StartShell struct{}

// Label is the label to show to the user.
func (StartShell) Label() string {
	return "Enter a LinuxBoot shell"
}

// Edit does nothing.
func (StartShell) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (StartShell) Load() error {
	return nil
}

// Exec implements Entry.Exec by running /bin/defaultsh.
func (StartShell) Exec() error {
	// Reset signal handler for SIGINT to enable user interrupts again
	signal.Reset(syscall.SIGINT)
	return sh.RunWithLogs("/bin/defaultsh")
}

// IsDefault indicates that this should not be run as a default action.
func (StartShell) IsDefault() bool { return false }

// Reboot is a menu.Entry that reboots the machine.
type Reboot struct{}

// Label is the label to show to the user.
func (Reboot) Label() string {
	return "Reboot"
}

// Edit does nothing.
func (Reboot) Edit(func(cmdline string) string) {
}

// Load does nothing.
func (Reboot) Load() error {
	return nil
}

// Exec reboots the machine using sys_reboot.
func (Reboot) Exec() error {
	unix.Sync()
	return unix.Reboot(unix.LINUX_REBOOT_CMD_RESTART)
}

// IsDefault indicates that this should not be run as a default action.
func (Reboot) IsDefault() bool { return false }
<|endoftext|>


<commit_before>// Copyright (C) 2015 Scaleway. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.md file.

package commands

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Sirupsen/logrus"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/scaleway/scaleway-cli/pkg/api"
	"github.com/scaleway/scaleway-cli/pkg/config"
	"github.com/scaleway/scaleway-cli/pkg/scwversion"
)

// LoginArgs are arguments passed to `RunLogin`
type LoginArgs struct {
	Organization string
	Token        string
	SSHKey       string
	SkipSSHKey   bool
}

// selectKey allows to choice a key in ~/.ssh
func selectKey(args *LoginArgs) error {
	home, err := config.GetHomeDir()
	if err != nil {
		return err
	}
	dir := filepath.Join(home, ".ssh")
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return fmt.Errorf("Unable to open your ~/.ssh: %v", err)
	}
	var pubs []string

	for i := range files {
		if filepath.Ext(files[i].Name()) == ".pub" {
			pubs = append(pubs, files[i].Name())
		}
	}
	if len(pubs) == 0 {
		return nil
	}
	fmt.Println("Do you want to upload an SSH key ?")
	fmt.Println("[0] I don't want to upload a key !")
	for i := range pubs {
		fmt.Printf("[%d] %s\n", i+1, pubs[i])
	}
	for {
		if err := promptUser("Which [id]: ", &args.SSHKey, true); err != nil {
			return err
		}
		id, err := strconv.ParseUint(strings.TrimSpace(args.SSHKey), 10, 32)
		if err != nil {
			fmt.Println(err)
			continue
		}
		if int(id) > len(pubs) {
			fmt.Println("Out of range id must be lower than", len(pubs))
			continue
		}
		args.SSHKey = ""
		if id == 0 {
			break
		}
		buff, err := ioutil.ReadFile(filepath.Join(dir, pubs[id-1]))
		if err != nil {
			return fmt.Errorf("Unable to open your key: %v", err)
		}
		args.SSHKey = string(buff[:len(buff)])
		break
	}
	return nil
}

func getToken(connect api.ScalewayConnect) (string, error) {
	FakeConnection, err := api.NewScalewayAPI("", "", scwversion.UserAgent())
	if err != nil {
		return "", fmt.Errorf("Unable to create a fake ScalewayAPI: %s", err)
	}
	FakeConnection.SetPassword(connect.Password)

	resp, err := FakeConnection.PostResponse(api.AccountAPI, "tokens", connect)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return "", fmt.Errorf("unable to connect %v", err)
	}

	// Succeed POST code
	if resp.StatusCode != 201 {
		return "", fmt.Errorf("[%d] maybe your email or your password is not valid", resp.StatusCode)
	}
	var data api.ScalewayConnectResponse

	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(&data)
	if err != nil {
		return "", err
	}
	return data.Token.ID, nil
}

func getOrganization(token string, email string) (string, error) {
	FakeConnection, err := api.NewScalewayAPI("", token, scwversion.UserAgent())
	if err != nil {
		return "", fmt.Errorf("Unable to create a fake ScalewayAPI: %s", err)
	}
	data, err := FakeConnection.GetOrganization()
	if err != nil {
		return "", err
	}

	orgaID := ""

	for _, orga := range data.Organizations {
		for _, user := range orga.Users {
			if user.Email == email {
				for i := range user.Organizations {
					if user.Organizations[i].Name != "OCS" {
						orgaID = user.Organizations[i].ID
						goto exit
					}
				}
			}
		}
	}
	if orgaID == "" {
		return "", fmt.Errorf("Unable to find your organization")
	}
exit:
	return orgaID, nil
}

func connectAPI() (string, string, error) {
	email := ""
	password := ""
	orga := ""
	token := ""
	hostname, err := os.Hostname()
	if err != nil {
		return "", "", fmt.Errorf("unable to get your Hostname %v", err)
	}
	if err := promptUser("Login (cloud.scaleway.com): ", &email, true); err != nil {
		return "", "", err
	}
	if err := promptUser("Password: ", &password, false); err != nil {
		return "", "", err
	}

	connect := api.ScalewayConnect{
		Email:       strings.Trim(email, "\n"),
		Password:    strings.Trim(password, "\n"),
		Expires:     false,
		Description: strings.Join([]string{"scw", hostname}, "-"),
	}
	token, err = getToken(connect)
	if err != nil {
		return "", "", err
	}
	orga, err = getOrganization(token, connect.Email)
	if err != nil {
		return "", "", err
	}
	return orga, token, nil
}

// uploadSSHKeys uploads an SSH Key
func uploadSSHKeys(apiConnection *api.ScalewayAPI, newKey string) {
	user, err := apiConnection.GetUser()
	if err != nil {
		logrus.Errorf("Unable to contact ScalewayAPI: %s", err)
	} else {
		user.SSHPublicKeys = append(user.SSHPublicKeys, api.ScalewayKeyDefinition{Key: strings.Trim(newKey, "\n")})

		SSHKeys := api.ScalewayUserPatchSSHKeyDefinition{
			SSHPublicKeys: user.SSHPublicKeys,
		}
		for i := range SSHKeys.SSHPublicKeys {
			SSHKeys.SSHPublicKeys[i].Fingerprint = ""
		}

		userID, err := apiConnection.GetUserID()
		if err != nil {
			logrus.Errorf("Unable to get userID: %s", err)
		} else {
			if err = apiConnection.PatchUserSSHKey(userID, SSHKeys); err != nil {
				logrus.Errorf("Unable to patch SSHkey: %v", err)
			}
		}
	}
}

// RunLogin is the handler for 'scw login'
func RunLogin(ctx CommandContext, args LoginArgs) error {
	if args.Organization == "" || args.Token == "" {
		var err error

		args.Organization, args.Token, err = connectAPI()
		if err != nil {
			return err
		}
	}

	cfg := &config.Config{
		Organization: strings.Trim(args.Organization, "\n"),
		Token:        strings.Trim(args.Token, "\n"),
	}

	apiConnection, err := api.NewScalewayAPI(cfg.Organization, cfg.Token, scwversion.UserAgent())
	if err != nil {
		return fmt.Errorf("Unable to create ScalewayAPI: %s", err)
	}
	err = apiConnection.CheckCredentials()
	if err != nil {
		return fmt.Errorf("Unable to contact ScalewayAPI: %s", err)
	}
	if !args.SkipSSHKey {
		if err := selectKey(&args); err != nil {
			logrus.Errorf("Unable to select a key: %v", err)
		} else {
			if args.SSHKey != "" {
				uploadSSHKeys(apiConnection, args.SSHKey)
			}
		}
	}
	return cfg.Save()
}

func promptUser(prompt string, output *string, echo bool) error {
	// FIXME: should use stdin/stdout from command context
	fmt.Fprintf(os.Stdout, prompt)
	os.Stdout.Sync()

	if !echo {
		b, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			return fmt.Errorf("Unable to prompt for password: %s", err)
		}
		*output = string(b)
		fmt.Fprintf(os.Stdout, "\n")
	} else {
		reader := bufio.NewReader(os.Stdin)
		*output, _ = reader.ReadString('\n')
	}
	return nil
}
<commit_msg>windows: fix login<commit_after>// Copyright (C) 2015 Scaleway. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.md file.

package commands

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/Sirupsen/logrus"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/scaleway/scaleway-cli/pkg/api"
	"github.com/scaleway/scaleway-cli/pkg/config"
	"github.com/scaleway/scaleway-cli/pkg/scwversion"
)

// LoginArgs are arguments passed to `RunLogin`
type LoginArgs struct {
	Organization string
	Token        string
	SSHKey       string
	SkipSSHKey   bool
}

// selectKey allows to choice a key in ~/.ssh
func selectKey(args *LoginArgs) error {
	home, err := config.GetHomeDir()
	if err != nil {
		return err
	}
	dir := filepath.Join(home, ".ssh")
	files, err := ioutil.ReadDir(dir)
	if err != nil {
		return fmt.Errorf("Unable to open your ~/.ssh: %v", err)
	}
	var pubs []string

	for i := range files {
		if filepath.Ext(files[i].Name()) == ".pub" {
			pubs = append(pubs, files[i].Name())
		}
	}
	if len(pubs) == 0 {
		return nil
	}
	fmt.Println("Do you want to upload an SSH key ?")
	fmt.Println("[0] I don't want to upload a key !")
	for i := range pubs {
		fmt.Printf("[%d] %s\n", i+1, pubs[i])
	}
	for {
		if err := promptUser("Which [id]: ", &args.SSHKey, true); err != nil {
			return err
		}
		id, err := strconv.ParseUint(strings.TrimSpace(args.SSHKey), 10, 32)
		if err != nil {
			fmt.Println(err)
			continue
		}
		if int(id) > len(pubs) {
			fmt.Println("Out of range id must be lower than", len(pubs))
			continue
		}
		args.SSHKey = ""
		if id == 0 {
			break
		}
		buff, err := ioutil.ReadFile(filepath.Join(dir, pubs[id-1]))
		if err != nil {
			return fmt.Errorf("Unable to open your key: %v", err)
		}
		args.SSHKey = string(buff[:len(buff)])
		break
	}
	return nil
}

func getToken(connect api.ScalewayConnect) (string, error) {
	FakeConnection, err := api.NewScalewayAPI("", "", scwversion.UserAgent())
	if err != nil {
		return "", fmt.Errorf("Unable to create a fake ScalewayAPI: %s", err)
	}
	FakeConnection.SetPassword(connect.Password)

	resp, err := FakeConnection.PostResponse(api.AccountAPI, "tokens", connect)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return "", fmt.Errorf("unable to connect %v", err)
	}

	// Succeed POST code
	if resp.StatusCode != 201 {
		return "", fmt.Errorf("[%d] maybe your email or your password is not valid", resp.StatusCode)
	}
	var data api.ScalewayConnectResponse

	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(&data)
	if err != nil {
		return "", err
	}
	return data.Token.ID, nil
}

func getOrganization(token string, email string) (string, error) {
	FakeConnection, err := api.NewScalewayAPI("", token, scwversion.UserAgent())
	if err != nil {
		return "", fmt.Errorf("Unable to create a fake ScalewayAPI: %s", err)
	}
	data, err := FakeConnection.GetOrganization()
	if err != nil {
		return "", err
	}

	orgaID := ""

	for _, orga := range data.Organizations {
		for _, user := range orga.Users {
			if user.Email == email {
				for i := range user.Organizations {
					if user.Organizations[i].Name != "OCS" {
						orgaID = user.Organizations[i].ID
						goto exit
					}
				}
			}
		}
	}
	if orgaID == "" {
		return "", fmt.Errorf("Unable to find your organization")
	}
exit:
	return orgaID, nil
}

func connectAPI() (string, string, error) {
	email := ""
	password := ""
	orga := ""
	token := ""
	hostname, err := os.Hostname()
	if err != nil {
		return "", "", fmt.Errorf("unable to get your Hostname %v", err)
	}
	if err := promptUser("Login (cloud.scaleway.com): ", &email, true); err != nil {
		return "", "", err
	}
	if err := promptUser("Password: ", &password, false); err != nil {
		return "", "", err
	}

	connect := api.ScalewayConnect{
		Email:       strings.Trim(email, "\r\n"),
		Password:    strings.Trim(password, "\r\n"),
		Expires:     false,
		Description: strings.Join([]string{"scw", hostname}, "-"),
	}
	token, err = getToken(connect)
	if err != nil {
		return "", "", err
	}
	orga, err = getOrganization(token, connect.Email)
	if err != nil {
		return "", "", err
	}
	return orga, token, nil
}

// uploadSSHKeys uploads an SSH Key
func uploadSSHKeys(apiConnection *api.ScalewayAPI, newKey string) {
	user, err := apiConnection.GetUser()
	if err != nil {
		logrus.Errorf("Unable to contact ScalewayAPI: %s", err)
	} else {
		user.SSHPublicKeys = append(user.SSHPublicKeys, api.ScalewayKeyDefinition{Key: strings.Trim(newKey, "\n")})

		SSHKeys := api.ScalewayUserPatchSSHKeyDefinition{
			SSHPublicKeys: user.SSHPublicKeys,
		}
		for i := range SSHKeys.SSHPublicKeys {
			SSHKeys.SSHPublicKeys[i].Fingerprint = ""
		}

		userID, err := apiConnection.GetUserID()
		if err != nil {
			logrus.Errorf("Unable to get userID: %s", err)
		} else {
			if err = apiConnection.PatchUserSSHKey(userID, SSHKeys); err != nil {
				logrus.Errorf("Unable to patch SSHkey: %v", err)
			}
		}
	}
}

// RunLogin is the handler for 'scw login'
func RunLogin(ctx CommandContext, args LoginArgs) error {
	if args.Organization == "" || args.Token == "" {
		var err error

		args.Organization, args.Token, err = connectAPI()
		if err != nil {
			return err
		}
	}

	cfg := &config.Config{
		Organization: strings.Trim(args.Organization, "\n"),
		Token:        strings.Trim(args.Token, "\n"),
	}

	apiConnection, err := api.NewScalewayAPI(cfg.Organization, cfg.Token, scwversion.UserAgent())
	if err != nil {
		return fmt.Errorf("Unable to create ScalewayAPI: %s", err)
	}
	err = apiConnection.CheckCredentials()
	if err != nil {
		return fmt.Errorf("Unable to contact ScalewayAPI: %s", err)
	}
	if !args.SkipSSHKey {
		if err := selectKey(&args); err != nil {
			logrus.Errorf("Unable to select a key: %v", err)
		} else {
			if args.SSHKey != "" {
				uploadSSHKeys(apiConnection, args.SSHKey)
			}
		}
	}
	return cfg.Save()
}

func promptUser(prompt string, output *string, echo bool) error {
	// FIXME: should use stdin/stdout from command context
	fmt.Fprintf(os.Stdout, prompt)
	os.Stdout.Sync()

	if !echo {
		b, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			return fmt.Errorf("Unable to prompt for password: %s", err)
		}
		*output = string(b)
		fmt.Fprintf(os.Stdout, "\n")
	} else {
		reader := bufio.NewReader(os.Stdin)
		*output, _ = reader.ReadString('\n')
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

// Package rand provides utilities related to randomization.
package rand

import (
	"math/rand"
	"sync"
	"time"
)

var rng = struct {
	sync.Mutex
	rand *rand.Rand
}{
	rand: rand.New(rand.NewSource(time.Now().UTC().UnixNano())),
}

// Intn generates an integer in range [0,max).
// By design this should panic if input is invalid, <= 0.
func Intn(max int) int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Intn(max)
}

// IntnRange generates an integer in range [min,max).
// By design this should panic if input is invalid, <= 0.
func IntnRange(min, max int) int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Intn(max-min) + min
}

// IntnRange generates an int64 integer in range [min,max).
// By design this should panic if input is invalid, <= 0.
func Int63nRange(min, max int64) int64 {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Int63n(max-min) + min
}

// Seed seeds the rng with the provided seed.
func Seed(seed int64) {
	rng.Lock()
	defer rng.Unlock()

	rng.rand = rand.New(rand.NewSource(seed))
}

// Perm returns, as a slice of n ints, a pseudo-random permutation of the integers [0,n)
// from the default Source.
func Perm(n int) []int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Perm(n)
}

// We omit vowels from the set of available characters to reduce the chances
// of "bad words" being formed.
var alphanums = []rune("bcdfghjklmnpqrstvwxz2456789")

// String generates a random alphanumeric string, without vowels, which is n
// characters long.  This will panic if n is less than zero.
func String(length int) string {
	b := make([]rune, length)
	for i := range b {
		b[i] = alphanums[Intn(len(alphanums))]
	}
	return string(b)
}

// SafeEncodeString encodes s using the same characters as rand.String. This reduces the chances of bad words and
// ensures that strings generated from hash functions appear consistent throughout the API.
func SafeEncodeString(s string) string {
	r := make([]rune, len(s))
	for i, b := range []rune(s) {
		r[i] = alphanums[(int(b) % len(alphanums))]
	}
	return string(r)
}
<commit_msg>Optimize random string generator to avoid multiple locks & use bitmasking<commit_after>/*
Copyright 2015 The Kubernetes Authors.

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

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

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

// Package rand provides utilities related to randomization.
package rand

import (
	"math/rand"
	"sync"
	"time"
)

var rng = struct {
	sync.Mutex
	rand *rand.Rand
}{
	rand: rand.New(rand.NewSource(time.Now().UTC().UnixNano())),
}

// Intn generates an integer in range [0,max).
// By design this should panic if input is invalid, <= 0.
func Intn(max int) int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Intn(max)
}

// IntnRange generates an integer in range [min,max).
// By design this should panic if input is invalid, <= 0.
func IntnRange(min, max int) int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Intn(max-min) + min
}

// IntnRange generates an int64 integer in range [min,max).
// By design this should panic if input is invalid, <= 0.
func Int63nRange(min, max int64) int64 {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Int63n(max-min) + min
}

// Seed seeds the rng with the provided seed.
func Seed(seed int64) {
	rng.Lock()
	defer rng.Unlock()

	rng.rand = rand.New(rand.NewSource(seed))
}

// Perm returns, as a slice of n ints, a pseudo-random permutation of the integers [0,n)
// from the default Source.
func Perm(n int) []int {
	rng.Lock()
	defer rng.Unlock()
	return rng.rand.Perm(n)
}

const (
	// We omit vowels from the set of available characters to reduce the chances
	// of "bad words" being formed.
	alphanums = "bcdfghjklmnpqrstvwxz2456789"
	// No. of bits required to index into alphanums string.
	alphanumsIdxBits = 5
	// Mask used to extract last alphanumsIdxBits of an int.
	alphanumsIdxMask = 1<<alphanumsIdxBits - 1
	// No. of random letters we can extract from a single int63.
	maxAlphanumsPerInt = 63 / alphanumsIdxBits
)

// String generates a random alphanumeric string, without vowels, which is n
// characters long.  This will panic if n is less than zero.
// How the random string is created:
// - we generate random int63's
// - from each int63, we are extracting multiple random letters by bit-shifting and masking
// - if some index is out of range of alphanums we neglect it (unlikely to happen multiple times in a row)
func String(n int) string {
	b := make([]byte, n)
	rng.Lock()
	defer rng.Unlock()

	randomInt63 := rng.rand.Int63()
	remaining := maxAlphanumsPerInt
	for i := 0; i < n; {
		if remaining == 0 {
			randomInt63, remaining = rng.rand.Int63(), maxAlphanumsPerInt
		}
		if idx := int(randomInt63 & alphanumsIdxMask); idx < len(alphanums) {
			b[i] = alphanums[idx]
			i++
		}
		randomInt63 >>= alphanumsIdxBits
		remaining--
	}
	return string(b)
}

// SafeEncodeString encodes s using the same characters as rand.String. This reduces the chances of bad words and
// ensures that strings generated from hash functions appear consistent throughout the API.
func SafeEncodeString(s string) string {
	r := make([]byte, len(s))
	for i, b := range []rune(s) {
		r[i] = alphanums[(int(b) % len(alphanums))]
	}
	return string(r)
}
<|endoftext|>


<commit_before>package grandcentral

import (
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	datastore "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	dhtpb "github.com/jbenet/go-ipfs/routing/dht/pb"
	proxy "github.com/jbenet/go-ipfs/routing/grandcentral/proxy"
	util "github.com/jbenet/go-ipfs/util"
	errors "github.com/jbenet/go-ipfs/util/debugerror"
)

// Server handles routing queries using a database backend
type Server struct {
	local           peer.ID
	routingBackend  datastore.ThreadSafeDatastore
	peerstore       peer.Peerstore
	*proxy.Loopback // so server can be injected into client
}

// NewServer creates a new GrandCentral routing Server
func NewServer(ds datastore.ThreadSafeDatastore, ps peer.Peerstore, local peer.ID) (*Server, error) {
	s := &Server{local, ds, ps, nil}
	s.Loopback = &proxy.Loopback{
		Handler: s,
		Local:   local,
	}
	return s, nil
}

// HandleLocalRequest implements the proxy.RequestHandler interface. This is
// where requests are received from the outside world.
func (s *Server) HandleRequest(ctx context.Context, p peer.ID, req *dhtpb.Message) *dhtpb.Message {
	_, response := s.handleMessage(ctx, p, req) // ignore response peer. it's local.
	return response
}

// TODO extract backend. backend can be implemented with whatever database we desire
func (s *Server) handleMessage(
	ctx context.Context, p peer.ID, req *dhtpb.Message) (peer.ID, *dhtpb.Message) {

	log.EventBegin(ctx, "routingMessageReceived", req, p, s.local).Done() // TODO may need to differentiate between local and remote

	var response = dhtpb.NewMessage(req.GetType(), req.GetKey(), req.GetClusterLevel())
	switch req.GetType() {

	case dhtpb.Message_GET_VALUE:
		rawRecord, err := getRoutingRecord(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.Record = rawRecord
		// TODO before merging: if we know any providers for the requested value, return those.
		return p, response

	case dhtpb.Message_PUT_VALUE:
		// TODO before merging: verifyRecord(req.GetRecord())
		putRoutingRecord(s.routingBackend, util.Key(req.GetKey()), req.GetRecord())
		return p, req // TODO before merging: verify that we should return record

	case dhtpb.Message_FIND_NODE:
		p := s.peerstore.PeerInfo(peer.ID(req.GetKey()))
		pri := []dhtpb.PeerRoutingInfo{
			dhtpb.PeerRoutingInfo{
				PeerInfo: p,
				// Connectedness: TODO
			},
		}
		response.CloserPeers = dhtpb.PeerRoutingInfosToPBPeers(pri)
		return p.ID, response

	case dhtpb.Message_ADD_PROVIDER:
		// FIXME(btc): do we want to store these locally? I think the
		// storeProvidersToPeerstore behavior is straight from the DHT message
		// handler.
		storeProvidersToPeerstore(s.peerstore, p, req.GetProviderPeers())

		if err := putRoutingProviders(s.routingBackend, util.Key(req.GetKey()), req.GetProviderPeers()); err != nil {
			return "", nil
		}
		return "", nil

	case dhtpb.Message_GET_PROVIDERS:
		providers, err := getRoutingProviders(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.ProviderPeers = providers
		return p, response

	case dhtpb.Message_PING:
		return p, req
	default:
	}
	return "", nil
}

var _ proxy.RequestHandler = &Server{}
var _ proxy.Proxy = &Server{}

func getRoutingRecord(ds datastore.Datastore, k util.Key) (*dhtpb.Record, error) {
	dskey := k.DsKey()
	val, err := ds.Get(dskey)
	if err != nil {
		return nil, errors.Wrap(err)
	}
	recordBytes, ok := val.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}
	var record dhtpb.Record
	if err := proto.Unmarshal(recordBytes, &record); err != nil {
		return nil, errors.New("failed to unmarshal dht record from datastore")
	}
	return &record, nil
}

func putRoutingRecord(ds datastore.Datastore, k util.Key, value *dhtpb.Record) error {
	data, err := proto.Marshal(value)
	if err != nil {
		return err
	}
	dskey := k.DsKey()
	// TODO namespace
	if err := ds.Put(dskey, data); err != nil {
		return err
	}
	return nil
}

func putRoutingProviders(ds datastore.Datastore, k util.Key, newRecords []*dhtpb.Message_Peer) error {
	log.Event(context.Background(), "putRoutingProviders", &k)
	pkey := datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()})
	oldRecords, err := getRoutingProviders(ds, k)
	if err != nil {
		return err
	}
	mergedRecords := make(map[string]*dhtpb.Message_Peer)
	for _, provider := range oldRecords {
		mergedRecords[provider.GetId()] = provider // add original records
	}
	for _, provider := range newRecords {
		mergedRecords[provider.GetId()] = provider // overwrite old record if new exists
	}
	var protomsg dhtpb.Message
	protomsg.ProviderPeers = make([]*dhtpb.Message_Peer, 0, len(mergedRecords))
	for _, provider := range mergedRecords {
		protomsg.ProviderPeers = append(protomsg.ProviderPeers, provider)
	}
	data, err := proto.Marshal(&protomsg)
	if err != nil {
		return err
	}
	return ds.Put(pkey, data)
}

func storeProvidersToPeerstore(ps peer.Peerstore, p peer.ID, providers []*dhtpb.Message_Peer) {
	for _, provider := range providers {
		providerID := peer.ID(provider.GetId())
		if providerID != p {
			log.Errorf("provider message came from third-party %s", p)
			continue
		}
		for _, maddr := range provider.Addresses() {
			// as a router, we want to store addresses for peers who have provided
			ps.AddAddr(p, maddr, peer.AddressTTL)
		}
	}
}

func getRoutingProviders(ds datastore.Datastore, k util.Key) ([]*dhtpb.Message_Peer, error) {
	e := log.EventBegin(context.Background(), "getProviders", &k)
	defer e.Done()
	var providers []*dhtpb.Message_Peer
	pkey := datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()}) // TODO key fmt
	if v, err := ds.Get(pkey); err == nil {
		if data, ok := v.([]byte); ok {
			var msg dhtpb.Message
			if err := proto.Unmarshal(data, &msg); err != nil {
				return nil, err
			}
			providers = append(providers, msg.GetProviderPeers()...)
		}
	}
	return providers, nil
}
<commit_msg>refactor(gcr/s) move declaration<commit_after>package grandcentral

import (
	"fmt"

	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
	proto "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/goprotobuf/proto"
	datastore "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-datastore"
	peer "github.com/jbenet/go-ipfs/p2p/peer"
	dhtpb "github.com/jbenet/go-ipfs/routing/dht/pb"
	proxy "github.com/jbenet/go-ipfs/routing/grandcentral/proxy"
	util "github.com/jbenet/go-ipfs/util"
	errors "github.com/jbenet/go-ipfs/util/debugerror"
)

// Server handles routing queries using a database backend
type Server struct {
	local           peer.ID
	routingBackend  datastore.ThreadSafeDatastore
	peerstore       peer.Peerstore
	*proxy.Loopback // so server can be injected into client
}

// NewServer creates a new GrandCentral routing Server
func NewServer(ds datastore.ThreadSafeDatastore, ps peer.Peerstore, local peer.ID) (*Server, error) {
	s := &Server{local, ds, ps, nil}
	s.Loopback = &proxy.Loopback{
		Handler: s,
		Local:   local,
	}
	return s, nil
}

// HandleLocalRequest implements the proxy.RequestHandler interface. This is
// where requests are received from the outside world.
func (s *Server) HandleRequest(ctx context.Context, p peer.ID, req *dhtpb.Message) *dhtpb.Message {
	_, response := s.handleMessage(ctx, p, req) // ignore response peer. it's local.
	return response
}

// TODO extract backend. backend can be implemented with whatever database we desire
func (s *Server) handleMessage(
	ctx context.Context, p peer.ID, req *dhtpb.Message) (peer.ID, *dhtpb.Message) {

	log.EventBegin(ctx, "routingMessageReceived", req, p, s.local).Done() // TODO may need to differentiate between local and remote

	var response = dhtpb.NewMessage(req.GetType(), req.GetKey(), req.GetClusterLevel())
	switch req.GetType() {

	case dhtpb.Message_GET_VALUE:
		rawRecord, err := getRoutingRecord(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.Record = rawRecord
		// TODO before merging: if we know any providers for the requested value, return those.
		return p, response

	case dhtpb.Message_PUT_VALUE:
		// TODO before merging: verifyRecord(req.GetRecord())
		putRoutingRecord(s.routingBackend, util.Key(req.GetKey()), req.GetRecord())
		return p, req // TODO before merging: verify that we should return record

	case dhtpb.Message_FIND_NODE:
		p := s.peerstore.PeerInfo(peer.ID(req.GetKey()))
		pri := []dhtpb.PeerRoutingInfo{
			dhtpb.PeerRoutingInfo{
				PeerInfo: p,
				// Connectedness: TODO
			},
		}
		response.CloserPeers = dhtpb.PeerRoutingInfosToPBPeers(pri)
		return p.ID, response

	case dhtpb.Message_ADD_PROVIDER:
		// FIXME(btc): do we want to store these locally? I think the
		// storeProvidersToPeerstore behavior is straight from the DHT message
		// handler.
		storeProvidersToPeerstore(s.peerstore, p, req.GetProviderPeers())

		if err := putRoutingProviders(s.routingBackend, util.Key(req.GetKey()), req.GetProviderPeers()); err != nil {
			return "", nil
		}
		return "", nil

	case dhtpb.Message_GET_PROVIDERS:
		providers, err := getRoutingProviders(s.routingBackend, util.Key(req.GetKey()))
		if err != nil {
			return "", nil
		}
		response.ProviderPeers = providers
		return p, response

	case dhtpb.Message_PING:
		return p, req
	default:
	}
	return "", nil
}

var _ proxy.RequestHandler = &Server{}
var _ proxy.Proxy = &Server{}

func getRoutingRecord(ds datastore.Datastore, k util.Key) (*dhtpb.Record, error) {
	dskey := k.DsKey()
	val, err := ds.Get(dskey)
	if err != nil {
		return nil, errors.Wrap(err)
	}
	recordBytes, ok := val.([]byte)
	if !ok {
		return nil, fmt.Errorf("datastore had non byte-slice value for %v", dskey)
	}
	var record dhtpb.Record
	if err := proto.Unmarshal(recordBytes, &record); err != nil {
		return nil, errors.New("failed to unmarshal dht record from datastore")
	}
	return &record, nil
}

func putRoutingRecord(ds datastore.Datastore, k util.Key, value *dhtpb.Record) error {
	data, err := proto.Marshal(value)
	if err != nil {
		return err
	}
	dskey := k.DsKey()
	// TODO namespace
	if err := ds.Put(dskey, data); err != nil {
		return err
	}
	return nil
}

func putRoutingProviders(ds datastore.Datastore, k util.Key, newRecords []*dhtpb.Message_Peer) error {
	log.Event(context.Background(), "putRoutingProviders", &k)
	oldRecords, err := getRoutingProviders(ds, k)
	if err != nil {
		return err
	}
	mergedRecords := make(map[string]*dhtpb.Message_Peer)
	for _, provider := range oldRecords {
		mergedRecords[provider.GetId()] = provider // add original records
	}
	for _, provider := range newRecords {
		mergedRecords[provider.GetId()] = provider // overwrite old record if new exists
	}
	var protomsg dhtpb.Message
	protomsg.ProviderPeers = make([]*dhtpb.Message_Peer, 0, len(mergedRecords))
	for _, provider := range mergedRecords {
		protomsg.ProviderPeers = append(protomsg.ProviderPeers, provider)
	}
	data, err := proto.Marshal(&protomsg)
	if err != nil {
		return err
	}
	pkey := datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()})
	return ds.Put(pkey, data)
}

func storeProvidersToPeerstore(ps peer.Peerstore, p peer.ID, providers []*dhtpb.Message_Peer) {
	for _, provider := range providers {
		providerID := peer.ID(provider.GetId())
		if providerID != p {
			log.Errorf("provider message came from third-party %s", p)
			continue
		}
		for _, maddr := range provider.Addresses() {
			// as a router, we want to store addresses for peers who have provided
			ps.AddAddr(p, maddr, peer.AddressTTL)
		}
	}
}

func getRoutingProviders(ds datastore.Datastore, k util.Key) ([]*dhtpb.Message_Peer, error) {
	e := log.EventBegin(context.Background(), "getProviders", &k)
	defer e.Done()
	var providers []*dhtpb.Message_Peer
	pkey := datastore.KeyWithNamespaces([]string{"routing", "providers", k.String()}) // TODO key fmt
	if v, err := ds.Get(pkey); err == nil {
		if data, ok := v.([]byte); ok {
			var msg dhtpb.Message
			if err := proto.Unmarshal(data, &msg); err != nil {
				return nil, err
			}
			providers = append(providers, msg.GetProviderPeers()...)
		}
	}
	return providers, nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"./file"
	"fmt"
	"os"
)

func main() {
	hello := []byte{'h', 'e', 'l', 'l', 'o', ',', ' ', 'w', 'o', 'r', 'l', 'd', '\n'}
	file.Stdout.Write(hello)
	file, err := file.Open("/does/not/exist",  0,  0)
	if file == nil {
		fmt.Printf("can't open file; err=%s\n",  err.String())
		os.Exit(1)
	}
}
<commit_msg>doc: Use byte("...") over byte{...}<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"./file"
	"fmt"
	"os"
)

func main() {
	hello := []byte("hello, world\n")
	file.Stdout.Write(hello)
	file, err := file.Open("/does/not/exist",  0,  0)
	if file == nil {
		fmt.Printf("can't open file; err=%s\n",  err.String())
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

// Sample graphviz-web is a Cloud Run service which provides a CLI tool over HTTP.
package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

func main() {
	// Verify the the dot utility is available at startup
	// instead of waiting for a first request.
	fileInfo, err := os.Stat("/usr/bin/dot")
	if err != nil {
		log.Fatalf("graphviz-web: %v", err)
	}
	if fileInfo.Mode()&0111 == 0 {
		log.Fatal(`graphviz-web: ("/usr/bin/dot") not executable`)
	}

	http.HandleFunc("/diagram.png", diagramHandler)

	// Determine port for HTTP service.
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
		log.Printf("Defaulting to port %s", port)
	}

	// Start HTTP server.
	log.Printf("Listening on port %s", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}

// [START run_system_package_handler]

// diagramHandler renders a diagram using HTTP request parameters and the dot command.
func diagramHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		log.Printf("method not allowed: %s", r.Method)
		http.Error(w, fmt.Sprintf("HTTP Method %s Not Allowed", r.Method), http.StatusMethodNotAllowed)
		return
	}

	q := r.URL.Query()
	dot := q.Get("dot")
	if dot == "" {
		log.Print("no graphviz definition provided")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	// Cache header must be set before writing a response.
	w.Header().Set("Cache-Control", "public, max-age=86400")

	input := strings.NewReader(dot)
	if err := createDiagram(w, input); err != nil {
		log.Printf("createDiagram: %v", err)
		// Do not cache error responses.
		w.Header().Del("Cache-Control")
		if strings.Contains(err.Error(), "syntax") {
			http.Error(w, "Bad Request: DOT syntax error", http.StatusBadRequest)
		} else {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		}
	}
}

// [END run_system_package_handler]

// [START run_system_package_exec]

// createDiagram generates a diagram image from the provided io.Reader written to the io.Writer.
func createDiagram(w io.Writer, r io.Reader) error {
	stderr := new(bytes.Buffer)
	args := []string{
		"-Glabel=Made on Cloud Run",
		"-Gfontsize=10",
		"-Glabeljust=right",
		"-Glabelloc=bottom",
		"-Gfontcolor=gray",
		"-Tpng",
	}
	cmd := exec.Command("/usr/bin/dot", args...)
	cmd.Stdin = r
	cmd.Stdout = w
	cmd.Stderr = stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("exec(%s) failed (%v): %s", cmd.Path, err, stderr.String())
	}

	return nil
}

// [END run_system_package_exec]
<commit_msg>fix(run/system_package): correct minor typos of comment in code (#1710)<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://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.

// Sample graphviz-web is a Cloud Run service which provides a CLI tool over HTTP.
package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"strings"
)

func main() {
	// Verify the dot utility is available at startup
	// instead of waiting for a first request.
	fileInfo, err := os.Stat("/usr/bin/dot")
	if err != nil {
		log.Fatalf("graphviz-web: %v", err)
	}
	if fileInfo.Mode()&0111 == 0 {
		log.Fatal(`graphviz-web: ("/usr/bin/dot") not executable`)
	}

	http.HandleFunc("/diagram.png", diagramHandler)

	// Determine port for HTTP service.
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
		log.Printf("Defaulting to port %s", port)
	}

	// Start HTTP server.
	log.Printf("Listening on port %s", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatal(err)
	}
}

// [START run_system_package_handler]

// diagramHandler renders a diagram using HTTP request parameters and the dot command.
func diagramHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		log.Printf("method not allowed: %s", r.Method)
		http.Error(w, fmt.Sprintf("HTTP Method %s Not Allowed", r.Method), http.StatusMethodNotAllowed)
		return
	}

	q := r.URL.Query()
	dot := q.Get("dot")
	if dot == "" {
		log.Print("no graphviz definition provided")
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	// Cache header must be set before writing a response.
	w.Header().Set("Cache-Control", "public, max-age=86400")

	input := strings.NewReader(dot)
	if err := createDiagram(w, input); err != nil {
		log.Printf("createDiagram: %v", err)
		// Do not cache error responses.
		w.Header().Del("Cache-Control")
		if strings.Contains(err.Error(), "syntax") {
			http.Error(w, "Bad Request: DOT syntax error", http.StatusBadRequest)
		} else {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		}
	}
}

// [END run_system_package_handler]

// [START run_system_package_exec]

// createDiagram generates a diagram image from the provided io.Reader written to the io.Writer.
func createDiagram(w io.Writer, r io.Reader) error {
	stderr := new(bytes.Buffer)
	args := []string{
		"-Glabel=Made on Cloud Run",
		"-Gfontsize=10",
		"-Glabeljust=right",
		"-Glabelloc=bottom",
		"-Gfontcolor=gray",
		"-Tpng",
	}
	cmd := exec.Command("/usr/bin/dot", args...)
	cmd.Stdin = r
	cmd.Stdout = w
	cmd.Stderr = stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("exec(%s) failed (%v): %s", cmd.Path, err, stderr.String())
	}

	return nil
}

// [END run_system_package_exec]
<|endoftext|>


<commit_before>package checklog

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
)

// Format implements fmt.Formatter.
func (s *state) Format(f fmt.State, c rune) {
	if s == nil {
		fmt.Fprintf(f, "<nil>")
		return
	}
	fmt.Printf("%"+string(c), *s)
}

func TestLoadStateIfFileNotExist(t *testing.T) {
	file := "testdata/file_not_found"
	s, err := loadState(file)
	if err != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

func TestLoadStateIfAccessDenied(t *testing.T) {
	var file string
	switch runtime.GOOS {
	case "windows":
		file = `C:\pagefile.sys`
	default:
		file = "testdata/plain.txt/file"
	}
	s, err := loadState(file)
	if err == nil {
		t.Errorf("loadState(%q) = %v; want an error", file, s)
	}
}

func TestLoadStateUnexpectedJson(t *testing.T) {
	file := "testdata/unexpected_json.json"
	s, err := loadState(file)
	if err != errStateFileCorrupted {
		t.Errorf("loadState(%q) = %v; want errStateFileCorrupted", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

func TestLoadStateBrokenJson(t *testing.T) {
	file := "testdata/invalid_json.json"
	s, err := loadState(file)
	if err != errStateFileCorrupted {
		t.Errorf("loadState(%q) = %v; want errStateFileCorrupted", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

// TODO(lufia): We might be better to test a condition too that file is exist but loadState can't read it.

func TestSaveStateIfFileNotExist(t *testing.T) {
	file := "testdata/file_will_create"
	defer func() {
		err := os.Remove(file)
		if err != nil && !os.IsNotExist(err) {
			t.Fatal(err)
		}
	}()

	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	testSaveLoadState(t, file, s)
}

func TestSaveStateOverwrittenIfFileExist(t *testing.T) {
	file := "testdata/state_overwritten"
	defer func() {
		err := os.Remove(file)
		if err != nil && !os.IsNotExist(err) {
			t.Fatal(err)
		}
	}()

	err := ioutil.WriteFile(file, []byte(`{"skip_bytes": 10, "inode": 100}`), 0644)
	if err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	testSaveLoadState(t, file, s)
}

func testSaveLoadState(t *testing.T, file string, s *state) {
	t.Helper()

	if err := saveState(file, s); err != nil {
		t.Errorf("saveState(%v) = %v; want nil", file, *s)
		return
	}
	s1, err := loadState(file)
	if err != nil {
		t.Errorf("loadState: %v", err)
		return
	}
	if !reflect.DeepEqual(s, s1) {
		t.Errorf("saveState(%v) -> loadState() = %v", s, s1)
	}
}

func TestSaveStateIfAccessDenied(t *testing.T) {
	switch runtime.GOOS {
	case "windows": // Almost os.Chmod operations will ignore on Windows.
		t.Skip()
		return
	}
	file := "testdata/readonly/state"
	dir := filepath.Dir(file)
	defer func() {
		if err := os.Chmod(dir, 0755); err != nil {
			t.Fatalf("Chmod: %v", err)
		}
		err := os.RemoveAll(dir)
		if err != nil && !os.IsNotExist(err) {
			t.Fatalf("RemoveAll: %v", err)
		}
	}()

	if err := os.MkdirAll(dir, 0700); err != nil {
		t.Errorf("MkdirAll: %v", err)
		return
	}
	data := []byte(`{"skip_bytes": 10, "inode": 100}`)
	if err := ioutil.WriteFile(file, data, 0644); err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	if err := os.Chmod(dir, 0500); err != nil {
		t.Errorf("Chmod: %v", err)
		return
	}
	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	saveState(file, s) // an error can be ignored in this case.
	data1, err := ioutil.ReadFile(file)
	if err != nil {
		t.Errorf("ReadFile: %v", err)
		return
	}
	if !bytes.Equal(data1, data) {
		t.Errorf("saveState into readonly directory should keep original contents: result = %s", data1)
	}
}

func TestGetBytesToSkipOld(t *testing.T) {
	file := "testdata/old-state.txt"
	n, err := getBytesToSkipOld(file)
	if err != nil {
		t.Fatal(err)
	}
	if want := int64(20); n != want {
		t.Errorf("getBytesToSkip(%s) = %d; want %d", file, n, want)
	}
}

func TestGetBytesToSkipOldIfFileNotExist(t *testing.T) {
	file := "testdata/file_not_found"
	n, err := getBytesToSkipOld(file)
	if err != errValidStateFileNotFound {
		t.Fatal(err)
	}
	if want := int64(0); n != want {
		t.Errorf("getBytesToSkip(%s) = %d; want %d", file, n, want)
	}
}

func TestGetBytesToSkipOldErr(t *testing.T) {
	var file string
	switch runtime.GOOS {
	case "windows":
		file = `C:\pagefile.sys`
	default:
		file = "testdata/plain.txt/file"
	}
	n, err := getBytesToSkipOld(file)
	if err == nil {
		t.Errorf("getBytesToSkip(%s) = %d; but want an error", file, n)
	}
}
<commit_msg>[check-log] skip few tests on Windows<commit_after>package checklog

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"runtime"
	"testing"
)

// Format implements fmt.Formatter.
func (s *state) Format(f fmt.State, c rune) {
	if s == nil {
		fmt.Fprintf(f, "<nil>")
		return
	}
	fmt.Printf("%"+string(c), *s)
}

func TestLoadStateIfFileNotExist(t *testing.T) {
	file := "testdata/file_not_found"
	s, err := loadState(file)
	if err != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

func TestLoadStateIfAccessDenied(t *testing.T) {
	switch runtime.GOOS {
	case "windows":
		t.Skip()
	}
	file := "testdata/plain.txt/file"
	s, err := loadState(file)
	if err == nil {
		t.Errorf("loadState(%q) = %v; want an error", file, s)
	}
}

func TestLoadStateUnexpectedJson(t *testing.T) {
	file := "testdata/unexpected_json.json"
	s, err := loadState(file)
	if err != errStateFileCorrupted {
		t.Errorf("loadState(%q) = %v; want errStateFileCorrupted", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

func TestLoadStateBrokenJson(t *testing.T) {
	file := "testdata/invalid_json.json"
	s, err := loadState(file)
	if err != errStateFileCorrupted {
		t.Errorf("loadState(%q) = %v; want errStateFileCorrupted", file, err)
	}
	if s != nil {
		t.Errorf("loadState(%q) = %v; want nil", file, *s)
	}
}

// TODO(lufia): We might be better to test a condition too that file is exist but loadState can't read it.

func TestSaveStateIfFileNotExist(t *testing.T) {
	file := "testdata/file_will_create"
	defer func() {
		err := os.Remove(file)
		if err != nil && !os.IsNotExist(err) {
			t.Fatal(err)
		}
	}()

	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	testSaveLoadState(t, file, s)
}

func TestSaveStateOverwrittenIfFileExist(t *testing.T) {
	file := "testdata/state_overwritten"
	defer func() {
		err := os.Remove(file)
		if err != nil && !os.IsNotExist(err) {
			t.Fatal(err)
		}
	}()

	err := ioutil.WriteFile(file, []byte(`{"skip_bytes": 10, "inode": 100}`), 0644)
	if err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	testSaveLoadState(t, file, s)
}

func testSaveLoadState(t *testing.T, file string, s *state) {
	t.Helper()

	if err := saveState(file, s); err != nil {
		t.Errorf("saveState(%v) = %v; want nil", file, *s)
		return
	}
	s1, err := loadState(file)
	if err != nil {
		t.Errorf("loadState: %v", err)
		return
	}
	if !reflect.DeepEqual(s, s1) {
		t.Errorf("saveState(%v) -> loadState() = %v", s, s1)
	}
}

func TestSaveStateIfAccessDenied(t *testing.T) {
	switch runtime.GOOS {
	case "windows": // Almost os.Chmod operations will ignore on Windows.
		t.Skip()
		return
	}
	file := "testdata/readonly/state"
	dir := filepath.Dir(file)
	defer func() {
		if err := os.Chmod(dir, 0755); err != nil {
			t.Fatalf("Chmod: %v", err)
		}
		err := os.RemoveAll(dir)
		if err != nil && !os.IsNotExist(err) {
			t.Fatalf("RemoveAll: %v", err)
		}
	}()

	if err := os.MkdirAll(dir, 0700); err != nil {
		t.Errorf("MkdirAll: %v", err)
		return
	}
	data := []byte(`{"skip_bytes": 10, "inode": 100}`)
	if err := ioutil.WriteFile(file, data, 0644); err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	if err := os.Chmod(dir, 0500); err != nil {
		t.Errorf("Chmod: %v", err)
		return
	}
	s := &state{
		SkipBytes: 15,
		Inode:     150,
	}
	saveState(file, s) // an error can be ignored in this case.
	data1, err := ioutil.ReadFile(file)
	if err != nil {
		t.Errorf("ReadFile: %v", err)
		return
	}
	if !bytes.Equal(data1, data) {
		t.Errorf("saveState into readonly directory should keep original contents: result = %s", data1)
	}
}

func TestGetBytesToSkipOld(t *testing.T) {
	file := "testdata/old-state.txt"
	n, err := getBytesToSkipOld(file)
	if err != nil {
		t.Fatal(err)
	}
	if want := int64(20); n != want {
		t.Errorf("getBytesToSkip(%s) = %d; want %d", file, n, want)
	}
}

func TestGetBytesToSkipOldIfFileNotExist(t *testing.T) {
	file := "testdata/file_not_found"
	n, err := getBytesToSkipOld(file)
	if err != errValidStateFileNotFound {
		t.Fatal(err)
	}
	if want := int64(0); n != want {
		t.Errorf("getBytesToSkip(%s) = %d; want %d", file, n, want)
	}
}

func TestGetBytesToSkipOldErr(t *testing.T) {
	switch runtime.GOOS {
	case "windows":
		t.Skip()
	}
	file := "testdata/plain.txt/file"
	n, err := getBytesToSkipOld(file)
	if err == nil {
		t.Errorf("getBytesToSkip(%s) = %d; but want an error", file, n)
	}
}
<|endoftext|>


<commit_before>package gincrud

import (
	"encoding/json"
	"fmt"
	"github.com/fatih/structs"
	"github.com/gin-gonic/gin"
	"github.com/osiloke/gostore"
	"strconv"
)

type GetKey func(interface{}, *gin.Context) string

//TODO: Extract core logic from each crud function i.e make doGetAll, doGet, ... they return data, err
func GetAll(bucket string, store gostore.Store, c *gin.Context) {
	var results []map[string]interface{}
	count := 10
	q := c.Request.URL.Query()
	if val, ok := q["_perPage"]; ok {
		count, _ = strconv.Atoi(val[0])
	}
	data, _ := store.GetAll(count, bucket)
	for _, element := range data {
		var result map[string]interface{}
		_ = json.Unmarshal(element[1], &result)
		result["key"] = string(element[0])
		results = append(results, result)
	}
	if len(results) == 0 {
		c.JSON(200, []string{})
	} else {
		stats, _ := store.Stats(bucket)
		c.Writer.Header().Set("X-Total-Count", fmt.Sprintf("%d", stats["KeyN"]))
		c.JSON(200, results)
	}
}

func Get(key, bucket string, store gostore.Store, c *gin.Context, record interface{}) {
	data, err := store.Get([]byte(key), bucket)
	if err != nil {
		//TODO: Does not exist error for store
		fmt.Println("Could not retrieve", key, "from", bucket)
		c.JSON(404, gin.H{"msg": fmt.Sprintf("Not found")})
	} else {
		_ = json.Unmarshal(data[1], record)
		m := structs.Map(record)
		m["key"] = string(data[0])
		c.JSON(200, m)
	}
}

func Post(bucket string, store gostore.Store, c *gin.Context, record interface{}, fn GetKey) {
	if b := c.Bind(record); b != false {
		m := structs.Map(record)
		data, err := json.Marshal(&record)
		key := fn(m, c)
		if err != nil {
			c.JSON(500, gin.H{"msg": "An error occured and this item could not be saved"})
		} else {
			store.Save([]byte([]byte(key)), data, bucket)
			m["key"] = key
			c.JSON(200, m)
		}
	} else {
		c.JSON(400, gin.H{"msg": "Seems like the data submitted is not formatted properly"})
	}
}

func Put(key, bucket string, store gostore.Store, c *gin.Context, record interface{}) {
	if b := c.Bind(record); b != false {
		data, err := json.Marshal(&record)
		if err != nil {
			c.JSON(500, gin.H{"msg": "An error occured and this item could not be saved"})
		} else {
			store.Save([]byte([]byte(key)), data, bucket)
			c.JSON(200, structs.Map(record))
		}
	} else {
		c.JSON(400, gin.H{"msg": "Seems like the data submitted is not formatted properly"})
	}
}

func Delete(key, bucket string, store gostore.Store, c *gin.Context) {
	err := store.Delete([]byte(key), bucket)
	if err != nil {
		c.JSON(500, gin.H{"msg": "The item [" + key + "] was not deleted"})
	} else {
		c.JSON(200, gin.H{"msg": "The item [" + key + "] was deleted"})
	}
}
<commit_msg>refactor to use pagination<commit_after>package gincrud

import (
	"encoding/json"
	"fmt"
	"github.com/fatih/structs"
	"github.com/gin-gonic/gin"
	"github.com/osiloke/gostore"
	"strconv"
)

type GetKey func(interface{}, *gin.Context) string

type Results struct {
	Data []map[string]interface{} `json:"data"`
	Next string                   `json:"next,omitempty"`
}

//TODO: Extract core logic from each crud function i.e make doGetAll, doGet, ... they return data, err
func GetAll(bucket string, store gostore.Store, c *gin.Context) {
	var results []map[string]interface{}
	count := 10
	q := c.Request.URL.Query()
	if val, ok := q["_perPage"]; ok {
		count, _ = strconv.Atoi(val[0])
	}
	var data [][][]byte
	var next [][]byte
	if val, ok := q["afterKey"]; ok {
		data, next, _ = store.GetAllAfter([]byte(val[0]), count, bucket)
	} else if val, ok := q["beforeKey"]; ok {
		data, next, _ = store.GetAllBefore([]byte(val[0]), count, bucket)
	} else {
		data, next, _ = store.GetAll(count, bucket)
	}

	for _, element := range data {
		var result map[string]interface{}
		_ = json.Unmarshal(element[1], &result)
		result["key"] = string(element[0])
		results = append(results, result)
	}
	if len(results) == 0 {
		c.JSON(200, []string{})
	} else {
		stats, _ := store.Stats(bucket)
		c.Writer.Header().Set("X-Total-Count", fmt.Sprintf("%d", stats["KeyN"]))
		c.JSON(200, Results{results, string(next[0])})
	}
}

func Get(key, bucket string, store gostore.Store, c *gin.Context, record interface{}) {
	data, err := store.Get([]byte(key), bucket)
	if err != nil {
		//TODO: Does not exist error for store
		fmt.Println("Could not retrieve", key, "from", bucket)
		c.JSON(404, gin.H{"msg": fmt.Sprintf("Not found")})
	} else {
		_ = json.Unmarshal(data[1], record)
		m := structs.Map(record)
		m["key"] = string(data[0])
		c.JSON(200, m)
	}
}

func Post(bucket string, store gostore.Store, c *gin.Context, record interface{}, fn GetKey) {
	if b := c.Bind(record); b != false {
		m := structs.Map(record)
		data, err := json.Marshal(&record)
		key := fn(m, c)
		if err != nil {
			c.JSON(500, gin.H{"msg": "An error occured and this item could not be saved"})
		} else {
			store.Save([]byte([]byte(key)), data, bucket)
			m["key"] = key
			c.JSON(200, m)
		}
	} else {
		c.JSON(400, gin.H{"msg": "Seems like the data submitted is not formatted properly"})
	}
}

func Put(key, bucket string, store gostore.Store, c *gin.Context, record interface{}) {
	if b := c.Bind(record); b != false {
		data, err := json.Marshal(&record)
		if err != nil {
			c.JSON(500, gin.H{"msg": "An error occured and this item could not be saved"})
		} else {
			store.Save([]byte([]byte(key)), data, bucket)
			c.JSON(200, structs.Map(record))
		}
	} else {
		c.JSON(400, gin.H{"msg": "Seems like the data submitted is not formatted properly"})
	}
}

func Delete(key, bucket string, store gostore.Store, c *gin.Context) {
	err := store.Delete([]byte(key), bucket)
	if err != nil {
		c.JSON(500, gin.H{"msg": "The item [" + key + "] was not deleted"})
	} else {
		c.JSON(200, gin.H{"msg": "The item [" + key + "] was deleted"})
	}
}
<|endoftext|>


<commit_before>package pwshandler

import (
	"net/http"

	"github.com/gorilla/websocket"
)

// Environment is a common data for connections in one pool
type Environment interface{}

// PoolManager is common interface for a structures which merge
// websocket connections in a pools with access to common environment
// data such as context for starting other goroutines or something
// else
type PoolManager interface {
	// AddConn creates connection to a pool and returns environment
	// data
	AddConn(*websocket.Conn) (Environment, error)
	// DelConn removes passed connection from a pool if exists
	DelConn(*websocket.Conn)
}

// ConnManager contains methods for processing websocket connections
// with passed environment data
type ConnManager interface {
	// Handle handles connection using passed environment data
	Handle(*websocket.Conn, Environment) error
	// HandleError processes an errors
	HandleError(http.ResponseWriter, *http.Request, error)
}

// RequestVerifier verifies requests
type RequestVerifier interface {
	Verify(*http.Request) error
}

// PoolHandler is handler which receives websocket requests and joins
// connections in a pools with common data in order with passed pool
// manager
type PoolHandler struct {
	poolMgr  PoolManager
	connMgr  ConnManager
	verifier RequestVerifier
	upgrader *websocket.Upgrader
}

// NewPoolHandler creates new PoolHandler with passed pool manager,
// connection manager, verifier and upgrader
func NewPoolHandler(poolMgr PoolManager, connMgr ConnManager,
	verifier RequestVerifier, upgrader *websocket.Upgrader,
) http.Handler {
	return &PoolHandler{poolMgr, connMgr, verifier, upgrader}
}

// Implementing http.Handler interface
func (ph *PoolHandler) ServeHTTP(w http.ResponseWriter,
	r *http.Request) {

	var err error
	// Verify request
	if ph.verifier != nil {
		if err = ph.verifier.Verify(r); err != nil {
			ph.connMgr.HandleError(w, r, err)
			return
		}
	}

	// Upgrade connection to websocket
	var conn *websocket.Conn
	if ph.upgrader != nil {
		conn, err = ph.upgrader.Upgrade(w, r, nil)
	} else {
		conn, err = (&websocket.Upgrader{}).Upgrade(w, r, nil)
	}
	if err != nil {
		ph.connMgr.HandleError(w, r, err)
		return
	}

	// Create connection to a pool and get envitonment data
	var env Environment
	if env, err = ph.poolMgr.AddConn(conn); err != nil {
		ph.connMgr.HandleError(w, r, err)
		return
	}

	// Handle connection
	if err = ph.connMgr.Handle(conn, env); err != nil {
		ph.connMgr.HandleError(w, r, err)
	}

	// Delete connection from a pool
	ph.poolMgr.DelConn(conn)

	// Close connection
	if err = conn.Close(); err != nil {
		ph.connMgr.HandleError(nil, nil, err)
	}
}
<commit_msg>pass req<commit_after>package pwshandler

import (
	"net/http"

	"github.com/gorilla/websocket"
)

// Environment is a common data for connections in one pool
type Environment interface{}

// PoolManager is common interface for a structures which merge
// websocket connections in a pools with access to common environment
// data such as context for starting other goroutines or something
// else
type PoolManager interface {
	// AddConn creates connection to a pool and returns environment
	// data
	AddConn(*websocket.Conn) (Environment, error)
	// DelConn removes passed connection from a pool if exists
	DelConn(*websocket.Conn)
}

// ConnManager contains methods for processing websocket connections
// with passed environment data
type ConnManager interface {
	// Handle handles connection using passed environment data
	Handle(*websocket.Conn, Environment) error
	// HandleError processes an errors
	HandleError(http.ResponseWriter, *http.Request, error)
}

// RequestVerifier verifies requests
type RequestVerifier interface {
	Verify(*http.Request) error
}

// PoolHandler is handler which receives websocket requests and joins
// connections in a pools with common data in order with passed pool
// manager
type PoolHandler struct {
	poolMgr  PoolManager
	connMgr  ConnManager
	verifier RequestVerifier
	upgrader *websocket.Upgrader
}

// NewPoolHandler creates new PoolHandler with passed pool manager,
// connection manager, verifier and upgrader
func NewPoolHandler(poolMgr PoolManager, connMgr ConnManager,
	verifier RequestVerifier, upgrader *websocket.Upgrader,
) http.Handler {
	return &PoolHandler{poolMgr, connMgr, verifier, upgrader}
}

// Implementing http.Handler interface
func (ph *PoolHandler) ServeHTTP(w http.ResponseWriter,
	r *http.Request) {

	var err error
	// Verify request
	if ph.verifier != nil {
		if err = ph.verifier.Verify(r); err != nil {
			ph.connMgr.HandleError(w, r, err)
			return
		}
	}

	// Upgrade connection to websocket
	var conn *websocket.Conn
	if ph.upgrader != nil {
		conn, err = ph.upgrader.Upgrade(w, r, nil)
	} else {
		conn, err = (&websocket.Upgrader{}).Upgrade(w, r, nil)
	}
	if err != nil {
		ph.connMgr.HandleError(w, r, err)
		return
	}

	// Create connection to a pool and get envitonment data
	var env Environment
	if env, err = ph.poolMgr.AddConn(conn); err != nil {
		ph.connMgr.HandleError(w, r, err)
		return
	}

	// Handle connection
	if err = ph.connMgr.Handle(conn, env); err != nil {
		ph.connMgr.HandleError(w, r, err)
	}

	// Delete connection from a pool
	ph.poolMgr.DelConn(conn)

	// Close connection
	if err = conn.Close(); err != nil {
		ph.connMgr.HandleError(nil, r, err)
	}
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"github.com/toolkits/core"
	"github.com/toolkits/nux"
	"net/http"
)

func configDfRoutes() {
	http.HandleFunc("/page/df", func(w http.ResponseWriter, r *http.Request) {
		mountPoints, err := nux.ListMountPoint()
		if err != nil {
			RenderMsgJson(w, err.Error())
			return
		}

		var ret [][]interface{} = make([][]interface{}, 0)
		for idx := range mountPoints {
			var du *nux.DeviceUsage
			du, err = nux.BuildDeviceUsage(mountPoints[idx][0], mountPoints[idx][1], mountPoints[idx][2])
			if err == nil {
				ret = append(ret,
					[]interface{}{
						du.FsSpec,
						core.ReadableSize(float64(du.BlocksAll)),
						core.ReadableSize(float64(du.BlocksUsed)),
						core.ReadableSize(float64(du.BlocksFree)),
						fmt.Sprintf("%.1f%%", du.BlocksUsedPercent),
						du.FsFile,
						core.ReadableSize(float64(du.InodesAll)),
						core.ReadableSize(float64(du.InodesUsed)),
						core.ReadableSize(float64(du.InodesFree)),
						fmt.Sprintf("%.1f%%", du.InodesUsedPercent),
						du.FsVfstype,
					})
			}
		}

		RenderDataJson(w, ret)
	})
}
<commit_msg>add disk usage for http<commit_after>package http

import (
	"fmt"
	"github.com/toolkits/core"
	"net/http"
	"log"
	"../tools/disk"
)

func configDfRoutes() {
	http.HandleFunc("/page/df", func(w http.ResponseWriter, r *http.Request) {
		var diskTotal uint64 = 0
		var diskUsed uint64 = 0

		var mountPoints, err = disk.Partitions(false)
		if err != nil {
			RenderMsgJson(w, err.Error())
			return
		}

		var ret [][]interface{} = make([][]interface{}, 0)
		for _, ds := range mountPoints {

			var dusage, err = disk.Usage(ds.Mountpoint)
			if err != nil {
				log.Println(err)
			}

			diskTotal += dusage.Total
			diskUsed += dusage.Used

			if err == nil {
				ret = append(ret,
					[]interface{}{
						dusage.Fstype,
						core.ReadableSize(float64(dusage.Total)),
						core.ReadableSize(float64(dusage.Used)),
						core.ReadableSize(float64(dusage.Free)),
						fmt.Sprintf("%.1f%%", dusage.UsedPercent),
						ds.Mountpoint,
						core.ReadableSize(float64(dusage.InodesTotal)),
						core.ReadableSize(float64(dusage.InodesUsed)),
						core.ReadableSize(float64(dusage.InodesFree)),
						fmt.Sprintf("%.1f%%", dusage.InodesUsedPercent),
						ds.Fstype,
					})
			}
		}

		RenderDataJson(w, ret)
	})
}
<|endoftext|>