<commit_before>package main

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func newProjectInDir(path string) error {
	// set path to be nested inside $GOPATH/src
	gopath := os.Getenv("GOPATH")
	path = filepath.Join(gopath, "src", path)

	// check if anything exists at the path, ask if it should be overwritten
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		fmt.Println("Path exists, overwrite contents? (y/N):")

		answer, err := getAnswer()
		if err != nil {
			return err
		}

		switch answer {
		case "n", "no", "\r\n", "\n", "":
			fmt.Println("")

		case "y", "yes":
			err := os.RemoveAll(path)
			if err != nil {
				return fmt.Errorf("Failed to overwrite %s. \n%s", path, err)
			}

			return createProjectInDir(path)

		default:
			fmt.Println("Input not recognized. No files overwritten. Answer as 'y' or 'n' only.")
		}

		return nil
	}

	return createProjectInDir(path)
}

var ponzuRepo = []string{"github.com", "ponzu-cms", "ponzu"}

func getAnswer() (string, error) {
	var answer string
	_, err := fmt.Scanf("%s\n", &answer)
	if err != nil {
		if err.Error() == "unexpected newline" {
			answer = ""
		} else {
			return "", err
		}
	}

	answer = strings.ToLower(answer)

	return answer, nil
}

func createProjectInDir(path string) error {
	gopath := os.Getenv("GOPATH")
	repo := ponzuRepo
	local := filepath.Join(gopath, "src", filepath.Join(repo...))
	network := "https://" + strings.Join(repo, "/") + ".git"
	if !strings.HasPrefix(path, gopath) {
		path = filepath.Join(gopath, path)
	}

	// create the directory or overwrite it
	err := os.MkdirAll(path, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	if dev {
		if fork != "" {
			local = filepath.Join(gopath, "src", fork)
		}

		devClone := exec.Command("git", "clone", local, "--branch", "ponzu-dev", "--single-branch", path)
		devClone.Stdout = os.Stdout
		devClone.Stderr = os.Stderr

		err = devClone.Start()
		if err != nil {
			return err
		}

		err = devClone.Wait()
		if err != nil {
			return err
		}

		err = vendorCorePackages(path)
		if err != nil {
			return err
		}

		fmt.Println("Dev build cloned from " + local + ":ponzu-dev")
		return nil
	}

	// try to git clone the repository from the local machine's $GOPATH
	localClone := exec.Command("git", "clone", local, path)
	localClone.Stdout = os.Stdout
	localClone.Stderr = os.Stderr

	err = localClone.Start()
	if err != nil {
		return err
	}
	err = localClone.Wait()
	if err != nil {
		fmt.Println("Couldn't clone from", local, "- trying network...")

		// try to git clone the repository over the network
		networkClone := exec.Command("git", "clone", network, path)
		networkClone.Stdout = os.Stdout
		networkClone.Stderr = os.Stderr

		err = networkClone.Start()
		if err != nil {
			fmt.Println("Network clone failed to start. Try again and make sure you have a network connection.")
			return err
		}
		err = networkClone.Wait()
		if err != nil {
			fmt.Println("Network clone failure.")
			// failed
			return fmt.Errorf("Failed to clone files from local machine [%s] and over the network [%s].\n%s", local, network, err)
		}
	}

	// create an internal vendor directory in ./cmd/ponzu and move content,
	// management and system packages into it
	err = vendorCorePackages(path)
	if err != nil {
		return err
	}

	gitDir := filepath.Join(path, ".git")
	err = os.RemoveAll(gitDir)
	if err != nil {
		fmt.Println("Failed to remove .git directory from your project path. Consider removing it manually.")
	}

	fmt.Println("New ponzu project created at", path)
	return nil
}

func vendorCorePackages(path string) error {
	vendorPath := filepath.Join(path, "cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu")
	err := os.MkdirAll(vendorPath, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	dirs := []string{"content", "management", "system"}
	for _, dir := range dirs {
		err = os.Rename(filepath.Join(path, dir), filepath.Join(vendorPath, dir))
		if err != nil {
			return err
		}
	}

	// create a user content directory at project root
	contentPath := filepath.Join(path, "content")
	err = os.Mkdir(contentPath, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	return nil
}

func copyFileNoRoot(src, dst string) error {
	noRoot := strings.Split(src, string(filepath.Separator))[1:]
	path := filepath.Join(noRoot...)
	dstFile, err := os.Create(filepath.Join(dst, path))
	defer dstFile.Close()
	if err != nil {
		return err
	}

	srcFile, err := os.Open(src)
	defer srcFile.Close()
	if err != nil {
		return err
	}

	_, err = io.Copy(dstFile, srcFile)
	if err != nil {
		return err
	}

	return nil
}

func copyFilesWarnConflicts(srcDir, dstDir string, conflicts []string) error {
	err := filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		for _, conflict := range conflicts {
			if info.Name() == conflict {
				fmt.Println("Ponzu couldn't fully build your project:")
				fmt.Println("You must rename the following file, as it conflicts with Ponzu core:")
				fmt.Println(path)
				fmt.Println("")
				fmt.Println("Once the files above have been renamed, run '$ ponzu build' to retry.")
				return errors.New("Ponzu has very few internal conflicts, sorry for the inconvenience.")
			}
		}

		if info.IsDir() {
			// don't copy root directory
			if path == srcDir {
				return nil
			}

			if len(path) > len(srcDir) {
				path = path[len(srcDir)+1:]
			}
			dir := filepath.Join(dstDir, path)
			err := os.MkdirAll(dir, os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}

			return nil
		}

		err = copyFileNoRoot(path, dstDir)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}

func emptyDir(path string) error {
	d, err := os.Open(path)
	if err != nil {
		return err
	}
	defer d.Close()

	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(path, name))
		if err != nil {
			return err
		}
	}

	return nil
}

func buildPonzuServer(args []string) error {
	pwd, err := os.Getwd()
	if err != nil {
		return err
	}

	// copy all ./content files to internal vendor directory
	src := "content"
	dst := filepath.Join("cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu", "content")
	err = emptyDir(dst)
	if err != nil {
		return err
	}
	err = copyFilesWarnConflicts(src, dst, []string{"doc.go"})
	if err != nil {
		return err
	}

	// copy all ./addons files & dirs to internal vendor directory
	src = "addons"
	dst = filepath.Join("cmd", "ponzu", "vendor")
	err = copyFilesWarnConflicts(src, dst, nil)
	if err != nil {
		return err
	}

	// execute go build -o ponzu-cms cmd/ponzu/*.go
	buildOptions := []string{"build", "-o", buildOutputName()}
	cmdBuildFiles := []string{
		"main.go", "options.go", "generate.go",
		"usage.go", "paths.go",
	}
	var cmdBuildFilePaths []string
	for _, file := range cmdBuildFiles {
		p := filepath.Join(pwd, "cmd", "ponzu", file)
		cmdBuildFilePaths = append(cmdBuildFilePaths, p)
	}

	build := exec.Command(gocmd, append(buildOptions, cmdBuildFilePaths...)...)
	build.Stderr = os.Stderr
	build.Stdout = os.Stdout

	err = build.Start()
	if err != nil {
		return errors.New("Ponzu build step failed. Please try again. " + "\n" + err.Error())

	}
	err = build.Wait()
	if err != nil {
		return errors.New("Ponzu build step failed. Please try again. " + "\n" + err.Error())

	}

	return nil
}

func copyAll(src, dst string) error {
	err := filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		sep := string(filepath.Separator)

		// base == the ponzu project dir + string(filepath.Separator)
		parts := strings.Split(src, sep)
		base := strings.Join(parts[:len(parts)-1], sep)
		base += sep

		target := filepath.Join(dst, path[len(base):])

		// if its a directory, make dir in dst
		if info.IsDir() {
			err := os.MkdirAll(target, os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}
		} else {
			// if its a file, move file to dir of dst
			err = os.Rename(path, target)
			if err != nil {
				return err
			}
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}

func upgradePonzuProjectDir(path string) error {
	core := []string{
		".gitattributes",
		"LICENSE",
		"ponzu-banner.png",
		"README.md",
		"cmd",
		"deployment",
		"management",
		"system",
	}

	stamp := fmt.Sprintf("ponzu-%d.bak", time.Now().Unix())
	temp := filepath.Join(os.TempDir(), stamp)
	err := os.Mkdir(temp, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	// track non-Ponzu core items (added by user)
	var user []os.FileInfo
	list, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	for _, item := range list {
		// check if in core
		var isCore bool
		for _, name := range core {
			if item.Name() == name {
				isCore = true
				break
			}
		}

		if !isCore {
			user = append(user, item)
		}
	}

	// move non-Ponzu files to temp location
	fmt.Println("Preserving files to be restored after upgrade...")
	for _, item := range user {
		src := filepath.Join(path, item.Name())
		if item.IsDir() {
			err := os.Mkdir(filepath.Join(temp, item.Name()), os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}
		}

		err := copyAll(src, temp)
		if err != nil {
			return err
		}

		fmt.Println(" [-]", item.Name())

	}

	// remove all files in path
	for _, item := range list {
		err := os.RemoveAll(filepath.Join(path, item.Name()))
		if err != nil {
			return fmt.Errorf("Failed to remove old Ponzu files.\n%s", err)
		}
	}

	err = createProjectInDir(path)
	if err != nil {
		fmt.Println("")
		fmt.Println("Upgrade failed...")
		fmt.Println("Your code is backed up at the following location:")
		fmt.Println(temp)
		fmt.Println("")
		fmt.Println("Manually create a new Ponzu project here and copy those files within it to fully restore.")
		fmt.Println("")
		return err
	}

	// move non-Ponzu files from temp location backed
	restore, err := ioutil.ReadDir(temp)
	if err != nil {
		return err
	}

	fmt.Println("Restoring files preserved before upgrade...")
	for _, r := range restore {
		p := filepath.Join(temp, r.Name())
		err = copyAll(p, path)
		if err != nil {
			fmt.Println("Couldn't merge your previous project files with upgraded one.")
			fmt.Println("Manually copy your files from the following directory:")
			fmt.Println(temp)
			return err
		}

		fmt.Println(" [+]", r.Name())
	}

	// clean-up
	backups := []string{filepath.Join(path, stamp), temp}
	for _, bak := range backups {
		err := os.RemoveAll(bak)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>added add.go to cmdBuildFiles in buidPonzuServer()<commit_after>package main

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

func newProjectInDir(path string) error {
	// set path to be nested inside $GOPATH/src
	gopath := os.Getenv("GOPATH")
	path = filepath.Join(gopath, "src", path)

	// check if anything exists at the path, ask if it should be overwritten
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		fmt.Println("Path exists, overwrite contents? (y/N):")

		answer, err := getAnswer()
		if err != nil {
			return err
		}

		switch answer {
		case "n", "no", "\r\n", "\n", "":
			fmt.Println("")

		case "y", "yes":
			err := os.RemoveAll(path)
			if err != nil {
				return fmt.Errorf("Failed to overwrite %s. \n%s", path, err)
			}

			return createProjectInDir(path)

		default:
			fmt.Println("Input not recognized. No files overwritten. Answer as 'y' or 'n' only.")
		}

		return nil
	}

	return createProjectInDir(path)
}

var ponzuRepo = []string{"github.com", "ponzu-cms", "ponzu"}

func getAnswer() (string, error) {
	var answer string
	_, err := fmt.Scanf("%s\n", &answer)
	if err != nil {
		if err.Error() == "unexpected newline" {
			answer = ""
		} else {
			return "", err
		}
	}

	answer = strings.ToLower(answer)

	return answer, nil
}

func createProjectInDir(path string) error {
	gopath := os.Getenv("GOPATH")
	repo := ponzuRepo
	local := filepath.Join(gopath, "src", filepath.Join(repo...))
	network := "https://" + strings.Join(repo, "/") + ".git"
	if !strings.HasPrefix(path, gopath) {
		path = filepath.Join(gopath, path)
	}

	// create the directory or overwrite it
	err := os.MkdirAll(path, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	if dev {
		if fork != "" {
			local = filepath.Join(gopath, "src", fork)
		}

		devClone := exec.Command("git", "clone", local, "--branch", "ponzu-dev", "--single-branch", path)
		devClone.Stdout = os.Stdout
		devClone.Stderr = os.Stderr

		err = devClone.Start()
		if err != nil {
			return err
		}

		err = devClone.Wait()
		if err != nil {
			return err
		}

		err = vendorCorePackages(path)
		if err != nil {
			return err
		}

		fmt.Println("Dev build cloned from " + local + ":ponzu-dev")
		return nil
	}

	// try to git clone the repository from the local machine's $GOPATH
	localClone := exec.Command("git", "clone", local, path)
	localClone.Stdout = os.Stdout
	localClone.Stderr = os.Stderr

	err = localClone.Start()
	if err != nil {
		return err
	}
	err = localClone.Wait()
	if err != nil {
		fmt.Println("Couldn't clone from", local, "- trying network...")

		// try to git clone the repository over the network
		networkClone := exec.Command("git", "clone", network, path)
		networkClone.Stdout = os.Stdout
		networkClone.Stderr = os.Stderr

		err = networkClone.Start()
		if err != nil {
			fmt.Println("Network clone failed to start. Try again and make sure you have a network connection.")
			return err
		}
		err = networkClone.Wait()
		if err != nil {
			fmt.Println("Network clone failure.")
			// failed
			return fmt.Errorf("Failed to clone files from local machine [%s] and over the network [%s].\n%s", local, network, err)
		}
	}

	// create an internal vendor directory in ./cmd/ponzu and move content,
	// management and system packages into it
	err = vendorCorePackages(path)
	if err != nil {
		return err
	}

	gitDir := filepath.Join(path, ".git")
	err = os.RemoveAll(gitDir)
	if err != nil {
		fmt.Println("Failed to remove .git directory from your project path. Consider removing it manually.")
	}

	fmt.Println("New ponzu project created at", path)
	return nil
}

func vendorCorePackages(path string) error {
	vendorPath := filepath.Join(path, "cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu")
	err := os.MkdirAll(vendorPath, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	dirs := []string{"content", "management", "system"}
	for _, dir := range dirs {
		err = os.Rename(filepath.Join(path, dir), filepath.Join(vendorPath, dir))
		if err != nil {
			return err
		}
	}

	// create a user content directory at project root
	contentPath := filepath.Join(path, "content")
	err = os.Mkdir(contentPath, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	return nil
}

func copyFileNoRoot(src, dst string) error {
	noRoot := strings.Split(src, string(filepath.Separator))[1:]
	path := filepath.Join(noRoot...)
	dstFile, err := os.Create(filepath.Join(dst, path))
	defer dstFile.Close()
	if err != nil {
		return err
	}

	srcFile, err := os.Open(src)
	defer srcFile.Close()
	if err != nil {
		return err
	}

	_, err = io.Copy(dstFile, srcFile)
	if err != nil {
		return err
	}

	return nil
}

func copyFilesWarnConflicts(srcDir, dstDir string, conflicts []string) error {
	err := filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		for _, conflict := range conflicts {
			if info.Name() == conflict {
				fmt.Println("Ponzu couldn't fully build your project:")
				fmt.Println("You must rename the following file, as it conflicts with Ponzu core:")
				fmt.Println(path)
				fmt.Println("")
				fmt.Println("Once the files above have been renamed, run '$ ponzu build' to retry.")
				return errors.New("Ponzu has very few internal conflicts, sorry for the inconvenience.")
			}
		}

		if info.IsDir() {
			// don't copy root directory
			if path == srcDir {
				return nil
			}

			if len(path) > len(srcDir) {
				path = path[len(srcDir)+1:]
			}
			dir := filepath.Join(dstDir, path)
			err := os.MkdirAll(dir, os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}

			return nil
		}

		err = copyFileNoRoot(path, dstDir)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}

func emptyDir(path string) error {
	d, err := os.Open(path)
	if err != nil {
		return err
	}
	defer d.Close()

	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(path, name))
		if err != nil {
			return err
		}
	}

	return nil
}

func buildPonzuServer(args []string) error {
	pwd, err := os.Getwd()
	if err != nil {
		return err
	}

	// copy all ./content files to internal vendor directory
	src := "content"
	dst := filepath.Join("cmd", "ponzu", "vendor", "github.com", "ponzu-cms", "ponzu", "content")
	err = emptyDir(dst)
	if err != nil {
		return err
	}
	err = copyFilesWarnConflicts(src, dst, []string{"doc.go"})
	if err != nil {
		return err
	}

	// copy all ./addons files & dirs to internal vendor directory
	src = "addons"
	dst = filepath.Join("cmd", "ponzu", "vendor")
	err = copyFilesWarnConflicts(src, dst, nil)
	if err != nil {
		return err
	}

	// execute go build -o ponzu-cms cmd/ponzu/*.go
	buildOptions := []string{"build", "-o", buildOutputName()}
	cmdBuildFiles := []string{
		"main.go", "options.go", "generate.go",
		"usage.go", "paths.go", "add.go",
	}
	var cmdBuildFilePaths []string
	for _, file := range cmdBuildFiles {
		p := filepath.Join(pwd, "cmd", "ponzu", file)
		cmdBuildFilePaths = append(cmdBuildFilePaths, p)
	}

	build := exec.Command(gocmd, append(buildOptions, cmdBuildFilePaths...)...)
	build.Stderr = os.Stderr
	build.Stdout = os.Stdout

	err = build.Start()
	if err != nil {
		return errors.New("Ponzu build step failed. Please try again. " + "\n" + err.Error())

	}
	err = build.Wait()
	if err != nil {
		return errors.New("Ponzu build step failed. Please try again. " + "\n" + err.Error())

	}

	return nil
}

func copyAll(src, dst string) error {
	err := filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		sep := string(filepath.Separator)

		// base == the ponzu project dir + string(filepath.Separator)
		parts := strings.Split(src, sep)
		base := strings.Join(parts[:len(parts)-1], sep)
		base += sep

		target := filepath.Join(dst, path[len(base):])

		// if its a directory, make dir in dst
		if info.IsDir() {
			err := os.MkdirAll(target, os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}
		} else {
			// if its a file, move file to dir of dst
			err = os.Rename(path, target)
			if err != nil {
				return err
			}
		}

		return nil
	})
	if err != nil {
		return err
	}

	return nil
}

func upgradePonzuProjectDir(path string) error {
	core := []string{
		".gitattributes",
		"LICENSE",
		"ponzu-banner.png",
		"README.md",
		"cmd",
		"deployment",
		"management",
		"system",
	}

	stamp := fmt.Sprintf("ponzu-%d.bak", time.Now().Unix())
	temp := filepath.Join(os.TempDir(), stamp)
	err := os.Mkdir(temp, os.ModeDir|os.ModePerm)
	if err != nil {
		return err
	}

	// track non-Ponzu core items (added by user)
	var user []os.FileInfo
	list, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	for _, item := range list {
		// check if in core
		var isCore bool
		for _, name := range core {
			if item.Name() == name {
				isCore = true
				break
			}
		}

		if !isCore {
			user = append(user, item)
		}
	}

	// move non-Ponzu files to temp location
	fmt.Println("Preserving files to be restored after upgrade...")
	for _, item := range user {
		src := filepath.Join(path, item.Name())
		if item.IsDir() {
			err := os.Mkdir(filepath.Join(temp, item.Name()), os.ModeDir|os.ModePerm)
			if err != nil {
				return err
			}
		}

		err := copyAll(src, temp)
		if err != nil {
			return err
		}

		fmt.Println(" [-]", item.Name())

	}

	// remove all files in path
	for _, item := range list {
		err := os.RemoveAll(filepath.Join(path, item.Name()))
		if err != nil {
			return fmt.Errorf("Failed to remove old Ponzu files.\n%s", err)
		}
	}

	err = createProjectInDir(path)
	if err != nil {
		fmt.Println("")
		fmt.Println("Upgrade failed...")
		fmt.Println("Your code is backed up at the following location:")
		fmt.Println(temp)
		fmt.Println("")
		fmt.Println("Manually create a new Ponzu project here and copy those files within it to fully restore.")
		fmt.Println("")
		return err
	}

	// move non-Ponzu files from temp location backed
	restore, err := ioutil.ReadDir(temp)
	if err != nil {
		return err
	}

	fmt.Println("Restoring files preserved before upgrade...")
	for _, r := range restore {
		p := filepath.Join(temp, r.Name())
		err = copyAll(p, path)
		if err != nil {
			fmt.Println("Couldn't merge your previous project files with upgraded one.")
			fmt.Println("Manually copy your files from the following directory:")
			fmt.Println(temp)
			return err
		}

		fmt.Println(" [+]", r.Name())
	}

	// clean-up
	backups := []string{filepath.Join(path, stamp), temp}
	for _, bak := range backups {
		err := os.RemoveAll(bak)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2020, Ben Morgan. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/cassava/repoctl/pacman/aur"
	"github.com/goulash/pr"
	"github.com/spf13/cobra"
)

func init() {
	MainCmd.AddCommand(queryCmd)
}

var queryCmd = &cobra.Command{
	Use:   "query [pkgname...]",
	Short: "query package information from AUR",
	Long: `Query package information from AUR.

  This command queries AUR for the specified packages and returns as much
  information on these packages as AUR gives us. The results are combined and
  sorted alphabetically.

  Note that this command is very similar to the results given from "search -i"
  command, but it uses a different AUR request. This command shows the
  following additional metadata:

	- Groups
	- Dependencies
	- Make Dependencies
	- Optional Dependencies
	- Conflicts
	- Provides
	- Replaces
	- Keywords

  Metadata properties that are empty are not shown.
`,
	Example: `  repoctl query firefox56 flirc-bin`,
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		// Prevent errors that we print being printed a second time by cobra.
		cmd.SilenceErrors = true
		cmd.SilenceUsage = true

		return nil
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		pkgs, err := aur.ReadAll(args)
		if err != nil {
			nfe, ok := err.(*aur.NotFoundError)
			if !ok {
				return err
			}
			for _, n := range nfe.Names {
				fmt.Fprintf(os.Stderr, "warning: unknown package %s\n", n)
			}
		}

		// Get the terminal width and fallback to a massive value if it's not
		// available. This prevents wrapping and lets us for example grep the
		// output better.
		terminalWidth := pr.StdoutTerminalWidth()
		if terminalWidth <= 0 {
			// FIXME: This is a hack
			terminalWidth = 1024
		}

		// Print the list
		pkgset := make(map[string]bool)
		for _, p := range pkgs {
			// Only add unique names to the list of packages
			if pkgset[p.Name] {
				continue
			}
			pkgset[p.Name] = true

			col.Printf("@{!m}aur/@{!w}%s @{!g}%s @{r}(%d)\n@|", p.Name, p.Version, p.NumVotes)
			col.Printf("@.%s\n", formatAURPackageInfo(p, terminalWidth))
		}

		return nil
	},
}

func formatAURPackageInfo(p *aur.Package, hspace int) string {
	// We want formatList to give us something like this:
	//		Depends: package package package package package
	//				 package package package package
	wrap := func(xs []string, prefixLen int) string {
		var buf strings.Builder

		n := prefixLen
		for i := 0; i < len(xs); i++ {
			x := xs[i]
			k := len(x)
			if n+k+1 > hspace && n != prefixLen {
				// If n == prefixLen, then that means we are at the beginning of the line,
				// and we still don't have enough space. We'll just have to deal with it.
				// A possible optimization here would be to try to reduce prefixLen and
				// see if it would fit then. But that can be done some other day.

				// Add a newline and the prefix
				buf.WriteRune('\n')
				buf.WriteString(strings.Repeat(" ", prefixLen))
				n = prefixLen
			}
			if n != prefixLen {
				buf.WriteRune(' ')
			}
			buf.WriteString(x)
			n += k + 1
		}

		return buf.String()
	}

	var buf strings.Builder
	fmt.Fprintf(&buf, "    Name: %s\n", p.Name)
	if p.PackageBase != p.Name {
		fmt.Fprintf(&buf, "    Base Name: %s\n", p.PackageBase)
	}
	fmt.Fprintf(&buf, "    Version: %s\n", p.Version)
	fmt.Fprintf(&buf, "    Description: %s\n", wrap(strings.Split(p.Description, " "), 17))
	if len(p.URL) != 0 {
		fmt.Fprintf(&buf, "    URL: %s\n", p.URL)
	}

	// The following is not available from the information we get when using the
	// search API, but might be useful in the future.
	//
	if len(p.License) > 0 {
		fmt.Fprintf(&buf, "    Licenses: %s\n", wrap(p.License, 14))
	}
	if len(p.Groups) > 0 {
		fmt.Fprintf(&buf, "    Groups: %s\n", wrap(p.Groups, 14))
	}
	if len(p.Provides) > 0 {
		fmt.Fprintf(&buf, "    Provides: %s\n", wrap(p.Provides, 14))
	}
	if len(p.Conflicts) > 0 {
		fmt.Fprintf(&buf, "    Conflicts: %s\n", wrap(p.Conflicts, 15))
	}
	if len(p.Replaces) > 0 {
		fmt.Fprintf(&buf, "    Replaces: %s\n", wrap(p.Replaces, 14))
	}
	if len(p.Depends) > 0 {
		fmt.Fprintf(&buf, "    Dependencies: %s\n", wrap(p.Depends, 16))
	}
	if len(p.OptDepends) > 0 {
		fmt.Fprintf(&buf, "    Optional Dependencies:\n")
		for _, d := range p.OptDepends {
			fmt.Fprintf(&buf, "        %s\n", d)
		}
	}
	if len(p.MakeDepends) > 0 {
		fmt.Fprintf(&buf, "    Build Dependencies: %s\n", wrap(p.MakeDepends, 15))
	}
	if len(p.Keywords) > 0 {
		fmt.Fprintf(&buf, "    Keywords: %s\n", wrap(p.Keywords, 14))
	}

	fmt.Fprintf(&buf, "    Snapshot URL: %s\n", p.URLPath)
	fmt.Fprintf(&buf, "    Maintainer: %s\n", p.Maintainer)
	fmt.Fprintf(&buf, "    Votes: %d\n", p.NumVotes)
	fmt.Fprintf(&buf, "    Popularity: %f\n", p.Popularity)
	fmt.Fprintf(&buf, "    First Submitted: %s\n", time.Unix(int64(p.FirstSubmitted), 0))
	fmt.Fprintf(&buf, "    Last Updated: %s\n", time.Unix(int64(p.LastModified), 0))
	fmt.Fprintf(&buf, "    Out-Of-Date: %v", p.OutOfDate != 0)
	return buf.String()
}
<commit_msg>Replace tab characters in usage messages<commit_after>// Copyright (c) 2020, Ben Morgan. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

package main

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/cassava/repoctl/pacman/aur"
	"github.com/goulash/pr"
	"github.com/spf13/cobra"
)

func init() {
	MainCmd.AddCommand(queryCmd)
}

var queryCmd = &cobra.Command{
	Use:   "query [pkgname...]",
	Short: "query package information from AUR",
	Long: `Query package information from AUR.

  This command queries AUR for the specified packages and returns as much
  information on these packages as AUR gives us. The results are combined and
  sorted alphabetically.

  Note that this command is very similar to the results given from "search -i"
  command, but it uses a different AUR request. This command shows the
  following additional metadata:

    - Groups
    - Dependencies
    - Make Dependencies
    - Optional Dependencies
    - Conflicts
    - Provides
    - Replaces
    - Keywords

  Metadata properties that are empty are not shown.
`,
	Example: `  repoctl query firefox56 flirc-bin`,
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
		// Prevent errors that we print being printed a second time by cobra.
		cmd.SilenceErrors = true
		cmd.SilenceUsage = true

		return nil
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		pkgs, err := aur.ReadAll(args)
		if err != nil {
			nfe, ok := err.(*aur.NotFoundError)
			if !ok {
				return err
			}
			for _, n := range nfe.Names {
				fmt.Fprintf(os.Stderr, "warning: unknown package %s\n", n)
			}
		}

		// Get the terminal width and fallback to a massive value if it's not
		// available. This prevents wrapping and lets us for example grep the
		// output better.
		terminalWidth := pr.StdoutTerminalWidth()
		if terminalWidth <= 0 {
			// FIXME: This is a hack
			terminalWidth = 1024
		}

		// Print the list
		pkgset := make(map[string]bool)
		for _, p := range pkgs {
			// Only add unique names to the list of packages
			if pkgset[p.Name] {
				continue
			}
			pkgset[p.Name] = true

			col.Printf("@{!m}aur/@{!w}%s @{!g}%s @{r}(%d)\n@|", p.Name, p.Version, p.NumVotes)
			col.Printf("@.%s\n", formatAURPackageInfo(p, terminalWidth))
		}

		return nil
	},
}

func formatAURPackageInfo(p *aur.Package, hspace int) string {
	// We want formatList to give us something like this:
	//		Depends: package package package package package
	//				 package package package package
	wrap := func(xs []string, prefixLen int) string {
		var buf strings.Builder

		n := prefixLen
		for i := 0; i < len(xs); i++ {
			x := xs[i]
			k := len(x)
			if n+k+1 > hspace && n != prefixLen {
				// If n == prefixLen, then that means we are at the beginning of the line,
				// and we still don't have enough space. We'll just have to deal with it.
				// A possible optimization here would be to try to reduce prefixLen and
				// see if it would fit then. But that can be done some other day.

				// Add a newline and the prefix
				buf.WriteRune('\n')
				buf.WriteString(strings.Repeat(" ", prefixLen))
				n = prefixLen
			}
			if n != prefixLen {
				buf.WriteRune(' ')
			}
			buf.WriteString(x)
			n += k + 1
		}

		return buf.String()
	}

	var buf strings.Builder
	fmt.Fprintf(&buf, "    Name: %s\n", p.Name)
	if p.PackageBase != p.Name {
		fmt.Fprintf(&buf, "    Base Name: %s\n", p.PackageBase)
	}
	fmt.Fprintf(&buf, "    Version: %s\n", p.Version)
	fmt.Fprintf(&buf, "    Description: %s\n", wrap(strings.Split(p.Description, " "), 17))
	if len(p.URL) != 0 {
		fmt.Fprintf(&buf, "    URL: %s\n", p.URL)
	}

	// The following is not available from the information we get when using the
	// search API, but might be useful in the future.
	//
	if len(p.License) > 0 {
		fmt.Fprintf(&buf, "    Licenses: %s\n", wrap(p.License, 14))
	}
	if len(p.Groups) > 0 {
		fmt.Fprintf(&buf, "    Groups: %s\n", wrap(p.Groups, 14))
	}
	if len(p.Provides) > 0 {
		fmt.Fprintf(&buf, "    Provides: %s\n", wrap(p.Provides, 14))
	}
	if len(p.Conflicts) > 0 {
		fmt.Fprintf(&buf, "    Conflicts: %s\n", wrap(p.Conflicts, 15))
	}
	if len(p.Replaces) > 0 {
		fmt.Fprintf(&buf, "    Replaces: %s\n", wrap(p.Replaces, 14))
	}
	if len(p.Depends) > 0 {
		fmt.Fprintf(&buf, "    Dependencies: %s\n", wrap(p.Depends, 16))
	}
	if len(p.OptDepends) > 0 {
		fmt.Fprintf(&buf, "    Optional Dependencies:\n")
		for _, d := range p.OptDepends {
			fmt.Fprintf(&buf, "        %s\n", d)
		}
	}
	if len(p.MakeDepends) > 0 {
		fmt.Fprintf(&buf, "    Build Dependencies: %s\n", wrap(p.MakeDepends, 15))
	}
	if len(p.Keywords) > 0 {
		fmt.Fprintf(&buf, "    Keywords: %s\n", wrap(p.Keywords, 14))
	}

	fmt.Fprintf(&buf, "    Snapshot URL: %s\n", p.URLPath)
	fmt.Fprintf(&buf, "    Maintainer: %s\n", p.Maintainer)
	fmt.Fprintf(&buf, "    Votes: %d\n", p.NumVotes)
	fmt.Fprintf(&buf, "    Popularity: %f\n", p.Popularity)
	fmt.Fprintf(&buf, "    First Submitted: %s\n", time.Unix(int64(p.FirstSubmitted), 0))
	fmt.Fprintf(&buf, "    Last Updated: %s\n", time.Unix(int64(p.LastModified), 0))
	fmt.Fprintf(&buf, "    Out-Of-Date: %v", p.OutOfDate != 0)
	return buf.String()
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io"
	"math/big"
	mr "math/rand"
	"net"
	"os"
	"time"
)

const (
	tlsRSABits           = 3072
	tlsDefaultCommonName = "syncthing"
)

func newCertificate(certFile, keyFile, name string) (tls.Certificate, error) {
	l.Infof("Generating RSA key and certificate for %s...", name)

	priv, err := rsa.GenerateKey(rand.Reader, tlsRSABits)
	if err != nil {
		l.Fatalln("generate key:", err)
	}

	notBefore := time.Now()
	notAfter := time.Date(2049, 12, 31, 23, 59, 59, 0, time.UTC)

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(mr.Int63()),
		Subject: pkix.Name{
			CommonName: name,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		l.Fatalln("create cert:", err)
	}

	certOut, err := os.Create(certFile)
	if err != nil {
		l.Fatalln("save cert:", err)
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	if err != nil {
		l.Fatalln("save cert:", err)
	}
	err = certOut.Close()
	if err != nil {
		l.Fatalln("save cert:", err)
	}

	keyOut, err := os.OpenFile(keyFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		l.Fatalln("save key:", err)
	}
	err = pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	if err != nil {
		l.Fatalln("save key:", err)
	}
	err = keyOut.Close()
	if err != nil {
		l.Fatalln("save key:", err)
	}

	return tls.LoadX509KeyPair(certFile, keyFile)
}

type DowngradingListener struct {
	net.Listener
	TLSConfig *tls.Config
}

type WrappedConnection struct {
	io.Reader
	net.Conn
}

func (l *DowngradingListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	br := bufio.NewReader(conn)
	bs, err := br.Peek(1)
	if err != nil {
		// We hit a read error here, but the Accept() call succeeded so we must not return an error.
		// We return the connection as is and let whoever tries to use it deal with the error.
		return conn, nil
	}

	wrapper := &WrappedConnection{br, conn}

	// 0x16 is the first byte of a TLS handshake
	if bs[0] == 0x16 {
		return tls.Server(wrapper, l.TLSConfig), nil
	}

	return wrapper, nil
}

func (c *WrappedConnection) Read(b []byte) (n int, err error) {
	return c.Reader.Read(b)
}
<commit_msg>Add timeout for peek (fixes #1035)<commit_after>// Copyright (C) 2014 The Syncthing Authors.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package main

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io"
	"math/big"
	mr "math/rand"
	"net"
	"os"
	"time"
)

const (
	tlsRSABits           = 3072
	tlsDefaultCommonName = "syncthing"
)

func newCertificate(certFile, keyFile, name string) (tls.Certificate, error) {
	l.Infof("Generating RSA key and certificate for %s...", name)

	priv, err := rsa.GenerateKey(rand.Reader, tlsRSABits)
	if err != nil {
		l.Fatalln("generate key:", err)
	}

	notBefore := time.Now()
	notAfter := time.Date(2049, 12, 31, 23, 59, 59, 0, time.UTC)

	template := x509.Certificate{
		SerialNumber: new(big.Int).SetInt64(mr.Int63()),
		Subject: pkix.Name{
			CommonName: name,
		},
		NotBefore: notBefore,
		NotAfter:  notAfter,

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		l.Fatalln("create cert:", err)
	}

	certOut, err := os.Create(certFile)
	if err != nil {
		l.Fatalln("save cert:", err)
	}
	err = pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	if err != nil {
		l.Fatalln("save cert:", err)
	}
	err = certOut.Close()
	if err != nil {
		l.Fatalln("save cert:", err)
	}

	keyOut, err := os.OpenFile(keyFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		l.Fatalln("save key:", err)
	}
	err = pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	if err != nil {
		l.Fatalln("save key:", err)
	}
	err = keyOut.Close()
	if err != nil {
		l.Fatalln("save key:", err)
	}

	return tls.LoadX509KeyPair(certFile, keyFile)
}

type DowngradingListener struct {
	net.Listener
	TLSConfig *tls.Config
}

type WrappedConnection struct {
	io.Reader
	net.Conn
}

func (l *DowngradingListener) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}

	br := bufio.NewReader(conn)
	conn.SetReadDeadline(time.Now().Add(1 * time.Second))
	bs, err := br.Peek(1)
	conn.SetReadDeadline(time.Time{})
	if err != nil {
		// We hit a read error here, but the Accept() call succeeded so we must not return an error.
		// We return the connection as is and let whoever tries to use it deal with the error.
		return conn, nil
	}

	wrapper := &WrappedConnection{br, conn}

	// 0x16 is the first byte of a TLS handshake
	if bs[0] == 0x16 {
		return tls.Server(wrapper, l.TLSConfig), nil
	}

	return wrapper, nil
}

func (c *WrappedConnection) Read(b []byte) (n int, err error) {
	return c.Reader.Read(b)
}
<|endoftext|>


<commit_before>// Copyright 2015 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.

// The upload command writes a file to Google Cloud Storage. It's used
// exclusively by the Makefiles in the Go project repos. Think of it
// as a very light version of gsutil or gcloud, but with some
// Go-specific configuration knowledge baked in.
package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"golang.org/x/build/auth"
	"golang.org/x/build/envutil"
	"golang.org/x/net/context"
	"google.golang.org/api/option"
)

var (
	public    = flag.Bool("public", false, "object should be world-readable")
	cacheable = flag.Bool("cacheable", true, "object should be cacheable")
	file      = flag.String("file", "-", "Filename to read object from, or '-' for stdin. If it begins with 'go:' then the rest is considered to be a Go target to install first, and then upload.")
	verbose   = flag.Bool("verbose", false, "verbose logging")
	osarch    = flag.String("osarch", "", "Optional 'GOOS-GOARCH' value to cross-compile; used only if --file begins with 'go:'. As a special case, if the value contains a '.' byte, anything up to and including that period is discarded.")
	project   = flag.String("project", "", "GCE Project. If blank, it's automatically inferred from the bucket name for the common Go buckets.")
	tags      = flag.String("tags", "", "tags to pass to go list, go install, etc. Only applicable if the --file value begins with 'go:'")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: upload [--public] [--file=...] <bucket/object>\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	args := strings.SplitN(flag.Arg(0), "/", 2)
	if len(args) != 2 {
		flag.Usage()
		os.Exit(1)
	}
	if strings.HasPrefix(*file, "go:") {
		buildGoTarget()
	}
	bucket, object := args[0], args[1]

	proj := *project
	if proj == "" {
		proj, _ = bucketProject[bucket]
		if proj == "" {
			log.Fatalf("bucket %q doesn't have an associated project in upload.go", bucket)
		}
	}

	ts, err := auth.ProjectTokenSource(proj, storage.ScopeReadWrite)
	if err != nil {
		log.Fatalf("Failed to get an OAuth2 token source: %v", err)
	}

	ctx := context.Background()
	storageClient, err := storage.NewClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		log.Fatalf("storage.NewClient: %v", err)
	}

	if alreadyUploaded(storageClient, bucket, object) {
		if *verbose {
			log.Printf("Already uploaded.")
		}
		return
	}

	w := storageClient.Bucket(bucket).Object(object).NewWriter(ctx)
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public, so
	// won't render the "Shared Publicly" link. So we do that, even
	// though it's dumb and unnecessary otherwise:
	w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.ACLEntity("project-owners-" + proj), Role: storage.RoleOwner})
	if *public {
		w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.AllUsers, Role: storage.RoleReader})
		if !*cacheable {
			w.CacheControl = "no-cache"
		}
	}
	var content io.Reader
	if *file == "-" {
		content = os.Stdin
	} else {
		content, err = os.Open(*file)
		if err != nil {
			log.Fatal(err)
		}
	}

	const maxSlurp = 1 << 20
	var buf bytes.Buffer
	n, err := io.CopyN(&buf, content, maxSlurp)
	if err != nil && err != io.EOF {
		log.Fatalf("Error reading from stdin: %v, %v", n, err)
	}
	w.ContentType = http.DetectContentType(buf.Bytes())

	_, err = io.Copy(w, io.MultiReader(&buf, content))
	if cerr := w.Close(); cerr != nil && err == nil {
		err = cerr
	}
	if err != nil {
		log.Fatalf("Write error: %v", err)
	}
	if *verbose {
		log.Printf("Wrote %v", object)
	}
	os.Exit(0)
}

var bucketProject = map[string]string{
	"dev-gccgo-builder-data": "gccgo-dashboard-dev",
	"dev-go-builder-data":    "go-dashboard-dev",
	"gccgo-builder-data":     "gccgo-dashboard-builders",
	"go-builder-data":        "symbolic-datum-552",
	"go-build-log":           "symbolic-datum-552",
	"http2-demo-server-tls":  "symbolic-datum-552",
	"winstrap":               "999119582588",
	"gobuilder":              "999119582588", // deprecated
}

func buildGoTarget() {
	target := strings.TrimPrefix(*file, "go:")
	var goos, goarch string
	if *osarch != "" {
		*osarch = (*osarch)[strings.LastIndex(*osarch, ".")+1:]
		v := strings.Split(*osarch, "-")
		if len(v) == 3 {
			v = v[:2] // support e.g. "linux-arm-scaleway" as GOOS=linux, GOARCH=arm
		}
		if len(v) != 2 || v[0] == "" || v[1] == "" {
			log.Fatalf("invalid -osarch value %q", *osarch)
		}
		goos, goarch = v[0], v[1]
	}

	env := envutil.Dedup(runtime.GOOS == "windows", append(os.Environ(), "GOOS="+goos, "GOARCH="+goarch))
	cmd := exec.Command("go", "list", "--tags="+*tags, "-f", "{{.Target}}", target)
	cmd.Env = env
	out, err := cmd.Output()
	if err != nil {
		log.Fatalf("go list: %v", err)
	}
	outFile := string(bytes.TrimSpace(out))
	fi0, err := os.Stat(outFile)
	if os.IsNotExist(err) {
		if *verbose {
			log.Printf("File %s doesn't exist; building...", outFile)
		}
	}

	version := os.Getenv("USER") + "-" + time.Now().Format(time.RFC3339)
	cmd = exec.Command("go", "install", "--tags="+*tags, "-x", "--ldflags=-X main.Version="+version, target)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if *verbose {
		cmd.Stdout = os.Stdout
	}
	cmd.Env = env
	if err := cmd.Run(); err != nil {
		log.Fatalf("go install %s: %v, %s", target, err, stderr.Bytes())
	}

	fi1, err := os.Stat(outFile)
	if err != nil {
		log.Fatalf("Expected output file %s stat failure after go install %v: %v", outFile, target, err)
	}
	if !os.SameFile(fi0, fi1) {
		if *verbose {
			log.Printf("File %s rebuilt.", outFile)
		}
	}
	*file = outFile
}

// alreadyUploaded reports whether *file has already been uploaded and the correct contents
// are on cloud storage already.
func alreadyUploaded(storageClient *storage.Client, bucket, object string) bool {
	if *file == "-" {
		return false // don't know.
	}
	o, err := storageClient.Bucket(bucket).Object(object).Attrs(context.Background())
	if err == storage.ErrObjectNotExist {
		return false
	}
	if err != nil {
		log.Printf("Warning: stat failure: %v", err)
		return false
	}
	m5 := md5.New()
	fi, err := os.Stat(*file)
	if err != nil {
		log.Fatal(err)
	}
	if fi.Size() != o.Size {
		return false
	}
	f, err := os.Open(*file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	n, err := io.Copy(m5, f)
	if err != nil {
		log.Fatal(err)
	}
	if n != fi.Size() {
		log.Printf("Warning: file size of %v changed", *file)
	}
	return bytes.Equal(m5.Sum(nil), o.MD5)
}
<commit_msg>cmd/upload: add --gzip flag<commit_after>// Copyright 2015 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.

// The upload command writes a file to Google Cloud Storage. It's used
// exclusively by the Makefiles in the Go project repos. Think of it
// as a very light version of gsutil or gcloud, but with some
// Go-specific configuration knowledge baked in.
package main

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"golang.org/x/build/auth"
	"golang.org/x/build/envutil"
	"golang.org/x/net/context"
	"google.golang.org/api/option"
)

var (
	public    = flag.Bool("public", false, "object should be world-readable")
	cacheable = flag.Bool("cacheable", true, "object should be cacheable")
	file      = flag.String("file", "-", "Filename to read object from, or '-' for stdin. If it begins with 'go:' then the rest is considered to be a Go target to install first, and then upload.")
	verbose   = flag.Bool("verbose", false, "verbose logging")
	osarch    = flag.String("osarch", "", "Optional 'GOOS-GOARCH' value to cross-compile; used only if --file begins with 'go:'. As a special case, if the value contains a '.' byte, anything up to and including that period is discarded.")
	project   = flag.String("project", "", "GCE Project. If blank, it's automatically inferred from the bucket name for the common Go buckets.")
	tags      = flag.String("tags", "", "tags to pass to go list, go install, etc. Only applicable if the --file value begins with 'go:'")
	doGzip    = flag.Bool("gzip", false, "gzip the stored contents (not the upload's Content-Encoding); this forces the Content-Type to be application/octet-stream. To prevent misuse, the object name must also end in '.gz'")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: upload [--public] [--file=...] <bucket/object>\n")
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	args := strings.SplitN(flag.Arg(0), "/", 2)
	if len(args) != 2 {
		flag.Usage()
		os.Exit(1)
	}
	if strings.HasPrefix(*file, "go:") {
		buildGoTarget()
	}
	bucket, object := args[0], args[1]

	if *doGzip && !strings.HasSuffix(object, ".gz") {
		log.Fatalf("--gzip flag requires object ending in .gz")
	}

	proj := *project
	if proj == "" {
		proj, _ = bucketProject[bucket]
		if proj == "" {
			log.Fatalf("bucket %q doesn't have an associated project in upload.go", bucket)
		}
	}

	ts, err := auth.ProjectTokenSource(proj, storage.ScopeReadWrite)
	if err != nil {
		log.Fatalf("Failed to get an OAuth2 token source: %v", err)
	}

	ctx := context.Background()
	storageClient, err := storage.NewClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		log.Fatalf("storage.NewClient: %v", err)
	}

	if alreadyUploaded(storageClient, bucket, object) {
		if *verbose {
			log.Printf("Already uploaded.")
		}
		return
	}

	w := storageClient.Bucket(bucket).Object(object).NewWriter(ctx)
	// If you don't give the owners access, the web UI seems to
	// have a bug and doesn't have access to see that it's public, so
	// won't render the "Shared Publicly" link. So we do that, even
	// though it's dumb and unnecessary otherwise:
	w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.ACLEntity("project-owners-" + proj), Role: storage.RoleOwner})
	if *public {
		w.ACL = append(w.ACL, storage.ACLRule{Entity: storage.AllUsers, Role: storage.RoleReader})
		if !*cacheable {
			w.CacheControl = "no-cache"
		}
	}
	var content io.Reader
	if *file == "-" {
		content = os.Stdin
	} else {
		content, err = os.Open(*file)
		if err != nil {
			log.Fatal(err)
		}
	}
	if *doGzip {
		var zbuf bytes.Buffer
		zw := gzip.NewWriter(&zbuf)
		if _, err := io.Copy(zw, content); err != nil {
			log.Fatalf("compressing content: %v", err)
		}
		if err := zw.Close(); err != nil {
			log.Fatalf("gzip.Close: %v", err)
		}
		content = &zbuf
	}

	const maxSlurp = 1 << 20
	var buf bytes.Buffer
	n, err := io.CopyN(&buf, content, maxSlurp)
	if err != nil && err != io.EOF {
		log.Fatalf("Error reading from stdin: %v, %v", n, err)
	}
	if *doGzip {
		w.ContentType = "application/octet-stream"
	} else {
		w.ContentType = http.DetectContentType(buf.Bytes())
	}

	_, err = io.Copy(w, io.MultiReader(&buf, content))
	if cerr := w.Close(); cerr != nil && err == nil {
		err = cerr
	}
	if err != nil {
		log.Fatalf("Write error: %v", err)
	}
	if *verbose {
		log.Printf("Wrote %v", object)
	}
	os.Exit(0)
}

var bucketProject = map[string]string{
	"dev-gccgo-builder-data": "gccgo-dashboard-dev",
	"dev-go-builder-data":    "go-dashboard-dev",
	"gccgo-builder-data":     "gccgo-dashboard-builders",
	"go-builder-data":        "symbolic-datum-552",
	"go-build-log":           "symbolic-datum-552",
	"http2-demo-server-tls":  "symbolic-datum-552",
	"winstrap":               "999119582588",
	"gobuilder":              "999119582588", // deprecated
}

func buildGoTarget() {
	target := strings.TrimPrefix(*file, "go:")
	var goos, goarch string
	if *osarch != "" {
		*osarch = strings.TrimSuffix(*osarch, ".gz")
		*osarch = (*osarch)[strings.LastIndex(*osarch, ".")+1:]
		v := strings.Split(*osarch, "-")
		if len(v) == 3 {
			v = v[:2] // support e.g. "linux-arm-scaleway" as GOOS=linux, GOARCH=arm
		}
		if len(v) != 2 || v[0] == "" || v[1] == "" {
			log.Fatalf("invalid -osarch value %q", *osarch)
		}
		goos, goarch = v[0], v[1]
	}

	env := envutil.Dedup(runtime.GOOS == "windows", append(os.Environ(), "GOOS="+goos, "GOARCH="+goarch))
	cmd := exec.Command("go", "list", "--tags="+*tags, "-f", "{{.Target}}", target)
	cmd.Env = env
	out, err := cmd.Output()
	if err != nil {
		log.Fatalf("go list: %v", err)
	}
	outFile := string(bytes.TrimSpace(out))
	fi0, err := os.Stat(outFile)
	if os.IsNotExist(err) {
		if *verbose {
			log.Printf("File %s doesn't exist; building...", outFile)
		}
	}

	version := os.Getenv("USER") + "-" + time.Now().Format(time.RFC3339)
	cmd = exec.Command("go", "install", "--tags="+*tags, "-x", "--ldflags=-X main.Version="+version, target)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	if *verbose {
		cmd.Stdout = os.Stdout
	}
	cmd.Env = env
	if err := cmd.Run(); err != nil {
		log.Fatalf("go install %s: %v, %s", target, err, stderr.Bytes())
	}

	fi1, err := os.Stat(outFile)
	if err != nil {
		log.Fatalf("Expected output file %s stat failure after go install %v: %v", outFile, target, err)
	}
	if !os.SameFile(fi0, fi1) {
		if *verbose {
			log.Printf("File %s rebuilt.", outFile)
		}
	}
	*file = outFile
}

// alreadyUploaded reports whether *file has already been uploaded and the correct contents
// are on cloud storage already.
func alreadyUploaded(storageClient *storage.Client, bucket, object string) bool {
	if *file == "-" {
		return false // don't know.
	}
	o, err := storageClient.Bucket(bucket).Object(object).Attrs(context.Background())
	if err == storage.ErrObjectNotExist {
		return false
	}
	if err != nil {
		log.Printf("Warning: stat failure: %v", err)
		return false
	}
	m5 := md5.New()
	fi, err := os.Stat(*file)
	if err != nil {
		log.Fatal(err)
	}
	if fi.Size() != o.Size {
		return false
	}
	f, err := os.Open(*file)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	n, err := io.Copy(m5, f)
	if err != nil {
		log.Fatal(err)
	}
	if n != fi.Size() {
		log.Printf("Warning: file size of %v changed", *file)
	}
	return bytes.Equal(m5.Sum(nil), o.MD5)
}
<|endoftext|>


<commit_before>// Package bugsnagrevel adds Bugsnag to revel.
// It lets you pass *revel.Controller into bugsnag.Notify(),
// and provides a Filter to catch errors.
package bugsnagrevel

import (
	"strings"
	"sync"
	"strconv"
	"net/http"
	"github.com/bugsnag/bugsnag-go"
	"github.com/revel/revel"
)

var once sync.Once

const FrameworkName string = "Revel"

var errorHandlingState = bugsnag.HandledState{
	bugsnag.SeverityReasonUnhandledMiddlewareError,
	bugsnag.SeverityError,
	true,
	FrameworkName,
}

// Filter should be added to the filter chain just after the PanicFilter.
// It sends errors to Bugsnag automatically. Configuration is read out of
// conf/app.conf, you should set bugsnag.apikey, and can also set
// bugsnag.endpoint, bugsnag.releasestage, bugsnag.apptype, bugsnag.appversion,
// bugsnag.projectroot, bugsnag.projectpackages if needed.
func Filter(c *revel.Controller, fc []revel.Filter) {
	defer bugsnag.AutoNotify(c, errorHandlingState)
	fc[0](c, fc[1:])
}

// Add support to bugsnag for reading data out of *revel.Controllers
func middleware(event *bugsnag.Event, config *bugsnag.Configuration) error {
	for _, datum := range event.RawData {
		if controller, ok := datum.(*revel.Controller); ok {
			// make the request visible to the builtin HttpMiddleware
			if version("0.18.0") {
				event.RawData = append(event.RawData, controller.Request)
			} else {
				req := struct{*http.Request}{}
				event.RawData = append(event.RawData, req.Request)
			}
			event.RawData = append(event.RawData, controller.Request)
			event.Context = controller.Action
			event.MetaData.AddStruct("Session", controller.Session)
		}
	}

	return nil
}

func init() {
	revel.OnAppStart(func() {
		bugsnag.OnBeforeNotify(middleware)

		var projectPackages []string
		if packages, ok := revel.Config.String("bugsnag.projectpackages"); ok {
			projectPackages = strings.Split(packages, ",")
		} else {
			projectPackages = []string{revel.ImportPath + "/app/*", revel.ImportPath + "/app"}
		}

		bugsnag.Configure(bugsnag.Configuration{
			APIKey:          revel.Config.StringDefault("bugsnag.apikey", ""),
			Endpoint:        revel.Config.StringDefault("bugsnag.endpoint", ""),
			AppType:         revel.Config.StringDefault("bugsnag.apptype", ""),
			AppVersion:      revel.Config.StringDefault("bugsnag.appversion", ""),
			ReleaseStage:    revel.Config.StringDefault("bugsnag.releasestage", revel.RunMode),
			ProjectPackages: projectPackages,
			Logger:          revel.ERROR,
		})
	})
}

// Very basic semantic versioning.
// Returns true if given version matches or is above revel.Version
func version(reqVersion string) bool{
	req := strings.Split(reqVersion, ".")
	cur := strings.Split(revel.Version, ".")
	for i:=0;i<2;i++{
		rV,_ := strconv.Atoi(req[i])
		cV,_ := strconv.Atoi(cur[i])
		if (rV<cV && i==0) || (rV<cV && i==1){
			return true
		}
	}
	return false
}
<commit_msg>chore: gofmt revel integration<commit_after>// Package bugsnagrevel adds Bugsnag to revel.
// It lets you pass *revel.Controller into bugsnag.Notify(),
// and provides a Filter to catch errors.
package bugsnagrevel

import (
	"github.com/bugsnag/bugsnag-go"
	"github.com/revel/revel"
	"net/http"
	"strconv"
	"strings"
	"sync"
)

var once sync.Once

const FrameworkName string = "Revel"

var errorHandlingState = bugsnag.HandledState{
	bugsnag.SeverityReasonUnhandledMiddlewareError,
	bugsnag.SeverityError,
	true,
	FrameworkName,
}

// Filter should be added to the filter chain just after the PanicFilter.
// It sends errors to Bugsnag automatically. Configuration is read out of
// conf/app.conf, you should set bugsnag.apikey, and can also set
// bugsnag.endpoint, bugsnag.releasestage, bugsnag.apptype, bugsnag.appversion,
// bugsnag.projectroot, bugsnag.projectpackages if needed.
func Filter(c *revel.Controller, fc []revel.Filter) {
	defer bugsnag.AutoNotify(c, errorHandlingState)
	fc[0](c, fc[1:])
}

// Add support to bugsnag for reading data out of *revel.Controllers
func middleware(event *bugsnag.Event, config *bugsnag.Configuration) error {
	for _, datum := range event.RawData {
		if controller, ok := datum.(*revel.Controller); ok {
			// make the request visible to the builtin HttpMiddleware
			if version("0.18.0") {
				event.RawData = append(event.RawData, controller.Request)
			} else {
				req := struct{ *http.Request }{}
				event.RawData = append(event.RawData, req.Request)
			}
			event.RawData = append(event.RawData, controller.Request)
			event.Context = controller.Action
			event.MetaData.AddStruct("Session", controller.Session)
		}
	}

	return nil
}

func init() {
	revel.OnAppStart(func() {
		bugsnag.OnBeforeNotify(middleware)

		var projectPackages []string
		if packages, ok := revel.Config.String("bugsnag.projectpackages"); ok {
			projectPackages = strings.Split(packages, ",")
		} else {
			projectPackages = []string{revel.ImportPath + "/app/*", revel.ImportPath + "/app"}
		}

		bugsnag.Configure(bugsnag.Configuration{
			APIKey:          revel.Config.StringDefault("bugsnag.apikey", ""),
			Endpoint:        revel.Config.StringDefault("bugsnag.endpoint", ""),
			AppType:         revel.Config.StringDefault("bugsnag.apptype", ""),
			AppVersion:      revel.Config.StringDefault("bugsnag.appversion", ""),
			ReleaseStage:    revel.Config.StringDefault("bugsnag.releasestage", revel.RunMode),
			ProjectPackages: projectPackages,
			Logger:          revel.ERROR,
		})
	})
}

// Very basic semantic versioning.
// Returns true if given version matches or is above revel.Version
func version(reqVersion string) bool {
	req := strings.Split(reqVersion, ".")
	cur := strings.Split(revel.Version, ".")
	for i := 0; i < 2; i++ {
		rV, _ := strconv.Atoi(req[i])
		cV, _ := strconv.Atoi(cur[i])
		if (rV < cV && i == 0) || (rV < cV && i == 1) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
    "container/heap"
    "flag"
    "github.com/bitly/go-simplejson"
    //"github.com/bitly/nsq/nsq"
    "github.com/bitly/nsq/tree/v0.2.21/nsq"
    "log"
    "strconv"
    "time"
    "net/http"
    "bytes"
    "io/ioutil"
)

var (
    // for input
    topic            = flag.String("topic", "", "nsq topic")
    channel          = flag.String("channel", "", "nsq topic")
    maxInFlight      = flag.Int("max-in-flight", 10, "max number of messages to allow in flight")
    lookupdHTTPAddrs = flag.String("lookupd-http-address", "127.0.0.1:4161", "lookupd HTTP address")
    // for output
    outNsqTCPAddrs   = flag.String("out-nsqd-tcp-address", "127.0.0.1:4151", "out nsqd TCP address")
    outTopic         = flag.String("out-topic", "", "nsq topic")
    outChannel       = flag.String("out-channel", "", "nsq channel")

    lag_time         = flag.Int("lag", 10, "lag before emitting in seconds")
    timeKey          = flag.String("key","","key that holds time")
)

type WriteMessage struct {
    val          []byte
    t            time.Time
    responseChan chan bool
}

type PQMessage struct {
    val          []byte
    t            time.Time
    index        int
    killChan     chan bool
    responseChan chan bool
}

// PRIORITY QUEUE
// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*PQMessage

func (pq PriorityQueue) Len() int {
    return len(pq)
}

func (pq PriorityQueue) Less(i, j int) bool {
    // We want Pop to give us the highest, not lowest, priority so we use greater than here.
    return pq[i].t.Before(pq[j].t)
}

func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
    pq[i].index = i
    pq[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
    n := len(*pq)
    item := x.(*PQMessage)
    item.index = n
    *pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    item.index = -1 // for safety
    *pq = old[0 : n-1]
    return item
}

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *PQMessage, val []byte, time time.Time) {
    heap.Remove(pq, item.index)
    item.val = val
    item.t = time
    heap.Push(pq, item)
}

func store(writeChan chan WriteMessage, out chan []byte, pq *PriorityQueue, lag time.Duration) {

    var emit_time time.Time
    nextMsg := &PQMessage{
        t: time.Now(),
    }

    getNext := make(chan bool)

    emitter := time.AfterFunc(24 * 365 * time.Hour, func(){
        log.Println("...")
    })

    const layout = "2006-01-02 15:04:05 -0700"

    count := 0
    heapCount := 0
    errorCount := 0
    for {
        select {
            case inMsg := <-writeChan:

                outMsg := &PQMessage{
                    val:          inMsg.val,
                    t:            inMsg.t,
                }

                outTime := outMsg.t.Add(lag)
                outDur := outTime.Sub(time.Now())

                if outDur > time.Duration(0 * time.Second) {
                    heap.Push(pq, outMsg)

                    heapCount ++

                    if heapCount % 500 == 0{
                        log.Println( "HEAP: " + strconv.Itoa(pq.Len()))
                    }

                    if outMsg.t.Before(nextMsg.t) {
                        heap.Push(pq, nextMsg)
                        nextMsg = heap.Pop(pq).(*PQMessage)
                        emit_time = nextMsg.t.Add(lag)
                        duration := emit_time.Sub(time.Now())

                        emitter.Stop()

                        emitter = time.AfterFunc(duration, func() {
                            out<-nextMsg.val
                            count = count + 1
                            if count % 250 == 0 {
                                diff := nextMsg.t.Sub( time.Now() )
                                log.Println("POP: " + diff.String() + "IN QUEUE:" + strconv.Itoa(pq.Len()) )
                            }
                            getNext<- true
                        })
                    }
                } else {
                    errorCount++
                    if errorCount % 250 == 0 {
                        log.Println("error: " + outDur.String() + " message reads: " + outMsg.t.Format(layout) )
                    }

                }

                //inMsg.responseChan <- true

            case <-getNext:
                if pq.Len() > 0 {
                    nextMsg = heap.Pop(pq).(*PQMessage)
                    emit_time = nextMsg.t.Add(lag)
                    duration := emit_time.Sub(time.Now())

                    emitter = time.AfterFunc(duration, func() {
                        out<-nextMsg.val
                        count = count + 1
                        if count % 250 == 0 {
                            diff := nextMsg.t.Sub( time.Now() )
                            log.Println("POP: " + diff.String() + "IN QUEUE:" + strconv.Itoa(pq.Len()) )
                        }
                        getNext<- true
                    })
                }
        }
    }
}

func emitter(tcpAddr string, topic string, out chan []byte){
    outCount := 0

    client := &http.Client{}

    for{
        select{
        case msg := <- out:
            outCount ++
            if outCount % 250 == 0{
                log.Println("OUT: " + strconv.Itoa(outCount) )
            }
            test := bytes.NewReader(msg)
            resp, err := client.Post("http://" + tcpAddr + "/put?topic=" + topic,"data/multi-part", test)
            if err != nil {
                log.Println(err.Error())
            }
            body, err := ioutil.ReadAll(resp.Body)

            if string(body) != "OK" {
                log.Println(body)
            }

            resp.Body.Close()
        }
    }
}

type SyncHandler struct{
    writeChan chan WriteMessage
    timeKey string
}

func (self *SyncHandler) HandleMessage(m *nsq.Message) error {

    reject := false

    blob, err := simplejson.NewJson(m.Body)

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    msg_time, err := blob.Get(self.timeKey).Int64()

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    // milliseconds
    t := time.Unix(0, msg_time * 1000 * 1000)
    mblob, err := blob.MarshalJSON()

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    responseChan := make(chan bool)

    msg := WriteMessage{
        t:            t,
        val:          mblob,
        responseChan: responseChan,
    }

    if !reject {
        self.writeChan <- msg
    }

    return nil
}

func main() {

    flag.Parse()

    stop := make(chan bool)
    wc := make(chan WriteMessage, 500000)
    oc := make(chan []byte)
    pq := &PriorityQueue{}
    heap.Init(pq)

    lag := time.Duration(time.Duration(*lag_time) * time.Second)

    go store(wc, oc, pq, lag)
    go emitter(*outNsqTCPAddrs, *outTopic, oc)

    for j := 0; j < 10; j++ {
        r, _ := nsq.NewReader(*topic, *channel)
        r.SetMaxInFlight(*maxInFlight)

        for i := 0; i < 5; i++ {
            sh := SyncHandler{
                writeChan: wc,
                timeKey: *timeKey,
            }
            r.AddHandler(&sh)
        }

        _ = r.ConnectToLookupd(*lookupdHTTPAddrs)
    }

    <-stop

}<commit_msg>still buggy on reads<commit_after>package main

import (
    "container/heap"
    "flag"
    "github.com/bitly/go-simplejson"
    //"github.com/bitly/nsq/nsq"
    "github.com/bitly/nsq/tree/v0.2.21/nsq"
    "log"
    "strconv"
    "time"
    "net/http"
    "bytes"
    "io/ioutil"
)

var (
    // for input
    topic            = flag.String("topic", "", "nsq topic")
    channel          = flag.String("channel", "", "nsq topic")
    maxInFlight      = flag.Int("max-in-flight", 10, "max number of messages to allow in flight")
    lookupdHTTPAddrs = flag.String("lookupd-http-address", "127.0.0.1:4161", "lookupd HTTP address")
    // for output
    outNsqTCPAddrs   = flag.String("out-nsqd-tcp-address", "127.0.0.1:4151", "out nsqd TCP address")
    outTopic         = flag.String("out-topic", "", "nsq topic")
    outChannel       = flag.String("out-channel", "", "nsq channel")

    lag_time         = flag.Int("lag", 10, "lag before emitting in seconds")
    timeKey          = flag.String("key","","key that holds time")
)

type WriteMessage struct {
    val          []byte
    t            time.Time
    responseChan chan bool
}

type PQMessage struct {
    val          []byte
    t            time.Time
    index        int
    killChan     chan bool
    responseChan chan bool
}

// PRIORITY QUEUE
// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*PQMessage

func (pq PriorityQueue) Len() int {
    return len(pq)
}

func (pq PriorityQueue) Less(i, j int) bool {
    // We want Pop to give us the highest, not lowest, priority so we use greater than here.
    return pq[i].t.Before(pq[j].t)
}

func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
    pq[i].index = i
    pq[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
    n := len(*pq)
    item := x.(*PQMessage)
    item.index = n
    *pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    item.index = -1 // for safety
    *pq = old[0 : n-1]
    return item
}

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *PQMessage, val []byte, time time.Time) {
    heap.Remove(pq, item.index)
    item.val = val
    item.t = time
    heap.Push(pq, item)
}

func store(writeChan chan WriteMessage, out chan []byte, pq *PriorityQueue, lag time.Duration) {

    var emit_time time.Time
    nextMsg := &PQMessage{
        t: time.Now(),
    }

    getNext := make(chan bool)

    emitter := time.AfterFunc(24 * 365 * time.Hour, func(){
        log.Println("...")
    })

    const layout = "2006-01-02 15:04:05 -0700"

    count := 0
    heapCount := 0
    errorCount := 0
    for {
        select {
            case inMsg := <-writeChan:

                outMsg := &PQMessage{
                    val:          inMsg.val,
                    t:            inMsg.t,
                }

                outTime := outMsg.t.Add(lag)
                outDur := outTime.Sub(time.Now())

                if outDur > time.Duration(0 * time.Second) {
                    heap.Push(pq, outMsg)

                    heapCount ++

                    if heapCount % 500 == 0{
                        log.Println( "HEAP: " + strconv.Itoa(pq.Len()))
                    }

                    if outMsg.t.Before(nextMsg.t) {
                        heap.Push(pq, nextMsg)
                        nextMsg = heap.Pop(pq).(*PQMessage)
                        emit_time = nextMsg.t.Add(lag)
                        duration := emit_time.Sub(time.Now())

                        emitter.Stop()

                        emitter = time.AfterFunc(duration, func() {
                            out<-nextMsg.val
                            count = count + 1
                            if count % 250 == 0 {
                                diff := nextMsg.t.Sub( time.Now() )
                                log.Println("POP: " + diff.String() + "IN QUEUE:" + strconv.Itoa(pq.Len()) )
                            }
                            getNext<- true
                        })
                    }
                } else {
                    errorCount++
                    if errorCount % 250 == 0 {
                        log.Println("error: " + outDur.String() + " message reads: " + outMsg.t.Format(layout) )
                    }

                }

                //inMsg.responseChan <- true

            case <-getNext:
                if pq.Len() > 0 {
                    nextMsg = heap.Pop(pq).(*PQMessage)
                    emit_time = nextMsg.t.Add(lag)
                    duration := emit_time.Sub(time.Now())

                    emitter = time.AfterFunc(duration, func() {
                        out<-nextMsg.val
                        count = count + 1
                        if count % 250 == 0 {
                            diff := nextMsg.t.Sub( time.Now() )
                            log.Println("POP: " + diff.String() + "IN QUEUE:" + strconv.Itoa(pq.Len()) )
                        }
                        getNext<- true
                    })
                }
        }
    }
}

func emitter(tcpAddr string, topic string, out chan []byte){
    outCount := 0

    client := &http.Client{}

    for{
        select{
        case msg := <- out:
            outCount ++
            if outCount % 250 == 0{
                log.Println("OUT: " + strconv.Itoa(outCount) )
            }
            test := bytes.NewReader(msg)
            resp, err := client.Post("http://" + tcpAddr + "/put?topic=" + topic,"data/multi-part", test)
            if err != nil {
                log.Println(err.Error())
            }
            body, err := ioutil.ReadAll(resp.Body)

            if string(body) != "OK" {
                log.Println(body)
            }

            resp.Body.Close()
        }
    }
}

type SyncHandler struct{
    writeChan chan WriteMessage
    timeKey string
}

func (self *SyncHandler) HandleMessage(m *nsq.Message) error {

    reject := false

    blob, err := simplejson.NewJson(m.Body)

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    msg_time, err := blob.Get(self.timeKey).Int64()

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    // milliseconds
    t := time.Unix(0, msg_time * 1000 * 1000)
    mblob, err := blob.MarshalJSON()

    if err != nil {
        reject = true
        log.Println(err.Error())
    }

    responseChan := make(chan bool)

    msg := WriteMessage{
        t:            t,
        val:          mblob,
        responseChan: responseChan,
    }

    if !reject {
        self.writeChan <- msg
    }

    return nil
}

func main() {

    flag.Parse()

    stop := make(chan bool)
    wc := make(chan WriteMessage, 500000)
    oc := make(chan []byte)
    pq := &PriorityQueue{}
    heap.Init(pq)

    lag := time.Duration(time.Duration(*lag_time) * time.Second)

    go store(wc, oc, pq, lag)
    go emitter(*outNsqTCPAddrs, *outTopic, oc)

    for j := 0; j < 10; j++ {
        go func(){
            r, _ := nsq.NewReader(*topic, *channel)
            r.SetMaxInFlight(*maxInFlight)

            for i := 0; i < 5; i++ {
                sh := SyncHandler{
                    writeChan: wc,
                    timeKey: *timeKey,
                }
                r.AddHandler(&sh)
            }

            _ = r.ConnectToLookupd(*lookupdHTTPAddrs)
        }()
    }

    <-stop

}<|endoftext|>


<commit_before>package graphs

import "testing"

func TestDFS(t *testing.T) {
	graph := NewDigraph()

	graph.AddEdge(1, 2, 0)
	graph.AddEdge(2, 3, 0)
	graph.AddEdge(3, 4, 0)
	graph.AddEdge(1, 5, 0)
	graph.AddEdge(5, 6, 0)
	graph.AddEdge(6, 3, 0)
	graph.AddEdge(1, 7, 0)

	walks := 0
	DFS(graph, 1, func(v Vertex, stop *bool) {
		walks++
	})

	if walks != 7 {
		t.Errorf("should visit 7 vertices; visited %d", walks)
	}

	walks = 0
	DFS(graph, 1, func(v Vertex, stop *bool) {
		walks++
		if v == 5 {
			*stop = true
		}
	})

	if walks != 5 {
		t.Errorf("should visit 5 vertices; visited %d", walks)
	}
}
<commit_msg>DFS: Fix flaky test<commit_after>package graphs

import "testing"

func TestDFS(t *testing.T) {
	graph := NewDigraph()

	graph.AddEdge(1, 2, 0)
	graph.AddEdge(2, 3, 0)
	graph.AddEdge(3, 4, 0)
	graph.AddEdge(1, 5, 0)
	graph.AddEdge(5, 6, 0)
	graph.AddEdge(6, 3, 0)
	graph.AddEdge(1, 7, 0)

	walks := 0
	DFS(graph, 1, func(v Vertex, stop *bool) {
		walks++
	})

	if walks != 7 {
		t.Errorf("should visit 7 vertices; visited %d", walks)
	}

	visited := make(map[Vertex]bool)
	DFS(graph, 1, func(v Vertex, stop *bool) {
		visited[v] = true
		if v == 5 {
			*stop = true
		}
	})
	if visited6 := visited[Vertex(6)]; visited6 {
		t.Errorf("visited vertex 6 vertices, but should not")
	}
}
<|endoftext|>


<commit_before>package kite

import (
	"crypto/tls"
	"flag"
	"fmt"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/op/go-logging"
)

func init() {
	// Debugging helper: Prints stacktrace on SIGUSR1.
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGUSR1)
	go func() {
		for {
			s := <-c
			fmt.Println("Got signal:", s)
			buf := make([]byte, 1<<16)
			runtime.Stack(buf, true)
			fmt.Println(string(buf))
			fmt.Print("Number of goroutines:", runtime.NumGoroutine())
			m := new(runtime.MemStats)
			runtime.GC()
			runtime.ReadMemStats(m)
			fmt.Printf(", Memory allocated: %+v\n", m.Alloc)
		}
	}()
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Is this Kite Public or Private? Default is Private.
	Visibility protocol.Visibility

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// Use Koding.com's reverse-proxy server for incoming connections.
	// Instead of the Kite's address, address of the Proxy Kite will be
	// registered to Kontrol.
	proxyEnabled bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Should handlers run concurrently? Default is true.
	concurrent bool

	// Dnode rpc server
	server *rpc.Server

	listener net.Listener

	// Handlers to call when a Kite opens a connection to this Kite.
	onConnectHandlers []func(*RemoteKite)

	// Handlers to call when a client has disconnected.
	onDisconnectHandlers []func(*RemoteKite)

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*Request) error

	// Should kite disable authenticators for incoming requests? Disabled by default
	disableAuthenticate bool

	// Kontrol keys to trust. Kontrol will issue access tokens for kites
	// that are signed with the private counterpart of these keys.
	// Key data must be PEM encoded.
	trustedKontrolKeys map[string][]byte

	// Trusted root certificates for TLS connections (wss://).
	// Certificate data must be PEM encoded.
	tlsCertificates [][]byte

	// Used to signal if the kite is ready to start and make calls to
	// other kites.
	ready chan bool
	end   chan bool

	// Prints logging messages to stderr and syslog.
	Log *logging.Logger
}

// New creates, initialize and then returns a new Kite instance. It accepts
// a single options argument that is a config struct that needs to be filled
// with several informations like Name, Port, IP and so on.
func New(options *Options) *Kite {
	var err error
	if options == nil {
		options, err = ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	options.validate() // exits if validating fails

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	// Enable authentication. options.DisableAuthentication is false by
	// default due to Go's varible initialization.
	var kodingKey string
	if !options.DisableAuthentication {
		kodingKey, err = utils.GetKodingKey()
		if err != nil {
			log.Fatal("Couldn't find koding.key. Please run 'kd register'.")
		}
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Environment: options.Environment,
			Region:      options.Region,
			Visibility:  options.Visibility,
			URL: protocol.KiteURL{
				&url.URL{
					Scheme: "ws",
					Host:   net.JoinHostPort(options.PublicIP, options.Port),
					Path:   "/dnode",
				},
			},
		},
		KodingKey:           kodingKey,
		server:              rpc.NewServer(),
		concurrent:          true,
		KontrolEnabled:      true,
		RegisterToKontrol:   true,
		trustedKontrolKeys:  make(map[string][]byte),
		Authenticators:      make(map[string]func(*Request) error),
		disableAuthenticate: options.DisableAuthentication,
		handlers:            make(map[string]HandlerFunc),
		ready:               make(chan bool),
		end:                 make(chan bool, 1),
	}

	k.TrustKontrolKey("koding.com", kodingKontrolPub)
	k.AddRootCertificate(kontrol_pem())

	k.server.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback, onError)
	k.server.Properties()["localKite"] = k

	k.Log = newLogger(k.Name, k.hasDebugFlag())
	k.Kontrol = k.NewKontrol(options.KontrolURL)

	// Call registered handlers when a client has disconnected.
	k.server.OnDisconnect(func(c *rpc.Client) {
		if r, ok := c.Properties()["remoteKite"]; ok {
			// Run OnDisconnect handlers.
			k.notifyRemoteKiteDisconnected(r.(*RemoteKite))
		}
	})

	k.server.OnConnect(func(c *rpc.Client) {
		k.Log.Info("Client is connected: %s", c.Conn.Request().RemoteAddr)
	})

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("systemInfo", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) DisableConcurrency() {
	k.server.SetConcurrent(false)
}

func (k *Kite) EnableTLS(certFile, keyFile string) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		k.Log.Fatal(err.Error())
	}

	k.server.TlsConfig = &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	k.Kite.URL.Scheme = "wss"
}

func (k *Kite) EnableProxy() {
	k.proxyEnabled = true
}

func (k *Kite) TrustKontrolKey(issuer string, key []byte) {
	k.trustedKontrolKeys[issuer] = key
}

func (k *Kite) AddRootCertificate(cert []byte) {
	k.tlsCertificates = append(k.tlsCertificates, cert)
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously.
func (k *Kite) Run() {
	k.Start()
	<-k.end
	k.Log.Notice("Kite server is closed.")
}

// Start is like Run(), but does not wait for it to complete. It's nonblocking.
func (k *Kite) Start() {
	k.parseVersionFlag()

	go func() {
		err := k.listenAndServe()
		if err != nil {
			k.Log.Fatal(err)
		}
	}()

	<-k.ready // wait until we are ready
}

// Close stops the server.
func (k *Kite) Close() {
	k.Log.Notice("Closing server...")
	k.listener.Close()
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args := r.Args.MustSliceOfLength(2)
	seconds := args[0].MustFloat64()
	ping := args[1].MustFunction()

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	msg := r.Args.One().MustString()
	k.Log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, msg))
	return nil, nil
}

func init() {
	// These logging related stuff needs to be called once because stupid
	// logging library uses global variables and resets the backends every time.
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))
	stderrBackend := logging.NewLogBackend(os.Stderr, "", log.LstdFlags)
	stderrBackend.Color = true
	syslogBackend, _ := logging.NewSyslogBackend("")
	logging.SetBackend(stderrBackend, syslogBackend)
}

// newLogger returns a new logger object for desired name and level.
func newLogger(name string, debug bool) *logging.Logger {
	logger := logging.MustGetLogger(name)

	level := logging.INFO
	if debug {
		level = logging.DEBUG
	}

	logging.SetLevel(level, name)
	return logger
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}

	// We can't use flags when running "go test" command.
	// This is another way to print debug logs.
	if os.Getenv("DEBUG") != "" {
		return true
	}

	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() (err error) {
	k.listener, err = net.Listen("tcp4", k.Kite.URL.Host)
	if err != nil {
		return err
	}

	k.Log.Notice("Listening: %s", k.listener.Addr().String())

	// Enable TLS
	if k.server.TlsConfig != nil {
		k.listener = tls.NewListener(k.listener, k.server.TlsConfig)
	}

	// Port is known here if "0" is used as port number
	host, _, _ := net.SplitHostPort(k.Kite.URL.Host)
	_, port, _ := net.SplitHostPort(k.listener.Addr().String())
	k.Kite.URL.Host = net.JoinHostPort(host, port)

	registerURLs := make(chan *url.URL, 1)

	if k.proxyEnabled {
		// Register to Proxy Kite and stay connected.
		// Fill the channel with registered Proxy URLs.
		go k.keepRegisteredToProxyKite(registerURLs)
	} else {
		// Register with Kite's own URL.
		registerURLs <- k.URL.URL
	}

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if err = k.Kontrol.DialForever(); err != nil {
			return
		}

		if k.RegisterToKontrol {
			go k.keepRegisteredToKontrol(registerURLs)
		}
	}

	k.ready <- true // listener is ready, unblock Start().

	// An error string equivalent to net.errClosing for using with http.Serve()
	// during a graceful exit. Needed to declare here again because it is not
	// exported by "net" package.
	const errClosing = "use of closed network connection"

	err = http.Serve(k.listener, k.server)
	if strings.Contains(err.Error(), errClosing) {
		// The server is closed by Close() method
		err = nil
	}

	k.end <- true // Serving is finished.

	return err
}

// OnConnect registers a function to run when a Kite connects to this Kite.
func (k *Kite) OnConnect(handler func(*RemoteKite)) {
	k.onConnectHandlers = append(k.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run when a connected Kite is disconnected.
func (k *Kite) OnDisconnect(handler func(*RemoteKite)) {
	k.onDisconnectHandlers = append(k.onDisconnectHandlers, handler)
}

// notifyRemoteKiteConnected runs the registered handlers with OnConnect().
func (k *Kite) notifyRemoteKiteConnected(r *RemoteKite) {
	k.Log.Info("Client '%s' is identified as '%s'",
		r.client.Conn.Request().RemoteAddr, r.Name)

	for _, handler := range k.onConnectHandlers {
		go handler(r)
	}
}

func (k *Kite) notifyRemoteKiteDisconnected(r *RemoteKite) {
	k.Log.Info("Client has disconnected: %s", r.Name)

	for _, handler := range k.onDisconnectHandlers {
		go handler(r)
	}
}
<commit_msg>kite: don't panic if there is no key available.<commit_after>package kite

import (
	"crypto/tls"
	"flag"
	"fmt"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"koding/newkite/utils"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/op/go-logging"
)

func init() {
	// Debugging helper: Prints stacktrace on SIGUSR1.
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGUSR1)
	go func() {
		for {
			s := <-c
			fmt.Println("Got signal:", s)
			buf := make([]byte, 1<<16)
			runtime.Stack(buf, true)
			fmt.Println(string(buf))
			fmt.Print("Number of goroutines:", runtime.NumGoroutine())
			m := new(runtime.MemStats)
			runtime.GC()
			runtime.ReadMemStats(m)
			fmt.Printf(", Memory allocated: %+v\n", m.Alloc)
		}
	}()
}

// Kite defines a single process that enables distributed service messaging
// amongst the peers it is connected. A Kite process acts as a Client and as a
// Server. That means it can receive request, process them, but it also can
// make request to other kites. A Kite can be anything. It can be simple Image
// processing kite (which would process data), it could be a Chat kite that
// enables peer-to-peer chat. For examples we have FileSystem kite that expose
// the file system to a client, which in order build the filetree.
type Kite struct {
	protocol.Kite

	// KodingKey is used for authenticate to Kontrol.
	KodingKey string

	// Is this Kite Public or Private? Default is Private.
	Visibility protocol.Visibility

	// Points to the Kontrol instance if enabled
	Kontrol *Kontrol

	// Wheter we want to connect to Kontrol on startup, true by default.
	KontrolEnabled bool

	// Wheter we want to register our Kite to Kontrol, true by default.
	RegisterToKontrol bool

	// Use Koding.com's reverse-proxy server for incoming connections.
	// Instead of the Kite's address, address of the Proxy Kite will be
	// registered to Kontrol.
	proxyEnabled bool

	// method map for exported methods
	handlers map[string]HandlerFunc

	// Should handlers run concurrently? Default is true.
	concurrent bool

	// Dnode rpc server
	server *rpc.Server

	listener net.Listener

	// Handlers to call when a Kite opens a connection to this Kite.
	onConnectHandlers []func(*RemoteKite)

	// Handlers to call when a client has disconnected.
	onDisconnectHandlers []func(*RemoteKite)

	// Contains different functions for authenticating user from request.
	// Keys are the authentication types (options.authentication.type).
	Authenticators map[string]func(*Request) error

	// Should kite disable authenticators for incoming requests? Disabled by default
	disableAuthenticate bool

	// Kontrol keys to trust. Kontrol will issue access tokens for kites
	// that are signed with the private counterpart of these keys.
	// Key data must be PEM encoded.
	trustedKontrolKeys map[string][]byte

	// Trusted root certificates for TLS connections (wss://).
	// Certificate data must be PEM encoded.
	tlsCertificates [][]byte

	// Used to signal if the kite is ready to start and make calls to
	// other kites.
	ready chan bool
	end   chan bool

	// Prints logging messages to stderr and syslog.
	Log *logging.Logger
}

// New creates, initialize and then returns a new Kite instance. It accepts
// a single options argument that is a config struct that needs to be filled
// with several informations like Name, Port, IP and so on.
func New(options *Options) *Kite {
	var err error
	if options == nil {
		options, err = ReadKiteOptions("manifest.json")
		if err != nil {
			log.Fatal("error: could not read config file", err)
		}
	}

	options.validate() // exits if validating fails

	hostname, _ := os.Hostname()
	kiteID := utils.GenerateUUID()

	// Enable authentication. options.DisableAuthentication is false by
	// default due to Go's varible initialization.
	var kodingKey string
	if !options.DisableAuthentication {
		kodingKey, err = utils.GetKodingKey()
		if err != nil {
			// don't fatal until we find a better way to integrate kite into other applications
			log.Println("Couldn't find koding.key. Please run 'kd register'.")
		}
	}

	k := &Kite{
		Kite: protocol.Kite{
			Name:        options.Kitename,
			Username:    options.Username,
			ID:          kiteID,
			Version:     options.Version,
			Hostname:    hostname,
			Environment: options.Environment,
			Region:      options.Region,
			Visibility:  options.Visibility,
			URL: protocol.KiteURL{
				&url.URL{
					Scheme: "ws",
					Host:   net.JoinHostPort(options.PublicIP, options.Port),
					Path:   "/dnode",
				},
			},
		},
		KodingKey:           kodingKey,
		server:              rpc.NewServer(),
		concurrent:          true,
		KontrolEnabled:      true,
		RegisterToKontrol:   true,
		trustedKontrolKeys:  make(map[string][]byte),
		Authenticators:      make(map[string]func(*Request) error),
		disableAuthenticate: options.DisableAuthentication,
		handlers:            make(map[string]HandlerFunc),
		ready:               make(chan bool),
		end:                 make(chan bool, 1),
	}

	k.TrustKontrolKey("koding.com", kodingKontrolPub)
	k.AddRootCertificate(kontrol_pem())

	k.server.SetWrappers(wrapMethodArgs, wrapCallbackArgs, runMethod, runCallback, onError)
	k.server.Properties()["localKite"] = k

	k.Log = newLogger(k.Name, k.hasDebugFlag())
	k.Kontrol = k.NewKontrol(options.KontrolURL)

	// Call registered handlers when a client has disconnected.
	k.server.OnDisconnect(func(c *rpc.Client) {
		if r, ok := c.Properties()["remoteKite"]; ok {
			// Run OnDisconnect handlers.
			k.notifyRemoteKiteDisconnected(r.(*RemoteKite))
		}
	})

	k.server.OnConnect(func(c *rpc.Client) {
		k.Log.Info("Client is connected: %s", c.Conn.Request().RemoteAddr)
	})

	// Every kite should be able to authenticate the user from token.
	k.Authenticators["token"] = k.AuthenticateFromToken
	// A kite accepts requests from Kontrol.
	k.Authenticators["kodingKey"] = k.AuthenticateFromKodingKey

	// Register our internal methods
	k.HandleFunc("systemInfo", new(Status).Info)
	k.HandleFunc("heartbeat", k.handleHeartbeat)
	k.HandleFunc("log", k.handleLog)

	return k
}

func (k *Kite) DisableConcurrency() {
	k.server.SetConcurrent(false)
}

func (k *Kite) EnableTLS(certFile, keyFile string) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		k.Log.Fatal(err.Error())
	}

	k.server.TlsConfig = &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	k.Kite.URL.Scheme = "wss"
}

func (k *Kite) EnableProxy() {
	k.proxyEnabled = true
}

func (k *Kite) TrustKontrolKey(issuer string, key []byte) {
	k.trustedKontrolKeys[issuer] = key
}

func (k *Kite) AddRootCertificate(cert []byte) {
	k.tlsCertificates = append(k.tlsCertificates, cert)
}

// Run is a blocking method. It runs the kite server and then accepts requests
// asynchronously.
func (k *Kite) Run() {
	k.Start()
	<-k.end
	k.Log.Notice("Kite server is closed.")
}

// Start is like Run(), but does not wait for it to complete. It's nonblocking.
func (k *Kite) Start() {
	k.parseVersionFlag()

	go func() {
		err := k.listenAndServe()
		if err != nil {
			k.Log.Fatal(err)
		}
	}()

	<-k.ready // wait until we are ready
}

// Close stops the server.
func (k *Kite) Close() {
	k.Log.Notice("Closing server...")
	k.listener.Close()
}

func (k *Kite) handleHeartbeat(r *Request) (interface{}, error) {
	args := r.Args.MustSliceOfLength(2)
	seconds := args[0].MustFloat64()
	ping := args[1].MustFunction()

	go func() {
		for {
			time.Sleep(time.Duration(seconds) * time.Second)
			if ping() != nil {
				return
			}
		}
	}()

	return nil, nil
}

// handleLog prints a log message to stdout.
func (k *Kite) handleLog(r *Request) (interface{}, error) {
	msg := r.Args.One().MustString()
	k.Log.Info(fmt.Sprintf("%s: %s", r.RemoteKite.Name, msg))
	return nil, nil
}

func init() {
	// These logging related stuff needs to be called once because stupid
	// logging library uses global variables and resets the backends every time.
	logging.SetFormatter(logging.MustStringFormatter("%{level:-8s} ▶ %{message}"))
	stderrBackend := logging.NewLogBackend(os.Stderr, "", log.LstdFlags)
	stderrBackend.Color = true
	syslogBackend, _ := logging.NewSyslogBackend("")
	logging.SetBackend(stderrBackend, syslogBackend)
}

// newLogger returns a new logger object for desired name and level.
func newLogger(name string, debug bool) *logging.Logger {
	logger := logging.MustGetLogger(name)

	level := logging.INFO
	if debug {
		level = logging.DEBUG
	}

	logging.SetLevel(level, name)
	return logger
}

// If the user wants to call flag.Parse() the flag must be defined in advance.
var _ = flag.Bool("version", false, "show version")
var _ = flag.Bool("debug", false, "print debug logs")

// parseVersionFlag prints the version number of the kite and exits with 0
// if "-version" flag is enabled.
// We did not use the "flag" package because it causes trouble if the user
// also calls "flag.Parse()" in his code. flag.Parse() can be called only once.
func (k *Kite) parseVersionFlag() {
	for _, flag := range os.Args {
		if flag == "-version" {
			fmt.Println(k.Version)
			os.Exit(0)
		}
	}
}

// hasDebugFlag returns true if -debug flag is present in os.Args.
func (k *Kite) hasDebugFlag() bool {
	for _, flag := range os.Args {
		if flag == "-debug" {
			return true
		}
	}

	// We can't use flags when running "go test" command.
	// This is another way to print debug logs.
	if os.Getenv("DEBUG") != "" {
		return true
	}

	return false
}

// listenAndServe starts our rpc server with the given addr.
func (k *Kite) listenAndServe() (err error) {
	k.listener, err = net.Listen("tcp4", k.Kite.URL.Host)
	if err != nil {
		return err
	}

	k.Log.Notice("Listening: %s", k.listener.Addr().String())

	// Enable TLS
	if k.server.TlsConfig != nil {
		k.listener = tls.NewListener(k.listener, k.server.TlsConfig)
	}

	// Port is known here if "0" is used as port number
	host, _, _ := net.SplitHostPort(k.Kite.URL.Host)
	_, port, _ := net.SplitHostPort(k.listener.Addr().String())
	k.Kite.URL.Host = net.JoinHostPort(host, port)

	registerURLs := make(chan *url.URL, 1)

	if k.proxyEnabled {
		// Register to Proxy Kite and stay connected.
		// Fill the channel with registered Proxy URLs.
		go k.keepRegisteredToProxyKite(registerURLs)
	} else {
		// Register with Kite's own URL.
		registerURLs <- k.URL.URL
	}

	// We must connect to Kontrol after starting to listen on port
	if k.KontrolEnabled {
		if err = k.Kontrol.DialForever(); err != nil {
			return
		}

		if k.RegisterToKontrol {
			go k.keepRegisteredToKontrol(registerURLs)
		}
	}

	k.ready <- true // listener is ready, unblock Start().

	// An error string equivalent to net.errClosing for using with http.Serve()
	// during a graceful exit. Needed to declare here again because it is not
	// exported by "net" package.
	const errClosing = "use of closed network connection"

	err = http.Serve(k.listener, k.server)
	if strings.Contains(err.Error(), errClosing) {
		// The server is closed by Close() method
		err = nil
	}

	k.end <- true // Serving is finished.

	return err
}

// OnConnect registers a function to run when a Kite connects to this Kite.
func (k *Kite) OnConnect(handler func(*RemoteKite)) {
	k.onConnectHandlers = append(k.onConnectHandlers, handler)
}

// OnDisconnect registers a function to run when a connected Kite is disconnected.
func (k *Kite) OnDisconnect(handler func(*RemoteKite)) {
	k.onDisconnectHandlers = append(k.onDisconnectHandlers, handler)
}

// notifyRemoteKiteConnected runs the registered handlers with OnConnect().
func (k *Kite) notifyRemoteKiteConnected(r *RemoteKite) {
	k.Log.Info("Client '%s' is identified as '%s'",
		r.client.Conn.Request().RemoteAddr, r.Name)

	for _, handler := range k.onConnectHandlers {
		go handler(r)
	}
}

func (k *Kite) notifyRemoteKiteDisconnected(r *RemoteKite) {
	k.Log.Info("Client has disconnected: %s", r.Name)

	for _, handler := range k.onDisconnectHandlers {
		go handler(r)
	}
}
<|endoftext|>