<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/openconfig/models-ci/commonci"
	"gopkg.in/yaml.v3"
)

var (
	// flags
	modelRoot string
	repoSlug  string
	commitSHA string
	prNumber  int
	extraPV   string

	// derived flags
	owner string
	repo  string

	// disabledModelPaths are the paths whose models should not undergo CI.
	// These should be temporary.
	disabledModelPaths = map[string]bool{
		"wifi:access-points": true,
		"wifi:ap-manager":    true,
		"wifi:mac":           true,
		"wifi:phy":           true,
	}
)

func init() {
	flag.StringVar(&modelRoot, "modelRoot", "", "root directory to OpenConfig models")
	flag.StringVar(&repoSlug, "repo-slug", "openconfig/public", "repo where CI is run")
	flag.StringVar(&commitSHA, "commit-sha", "", "commit SHA of the PR")
	flag.IntVar(&prNumber, "pr-number", 0, "PR number")
	flag.StringVar(&extraPV, "extra-pyang-versions", "", "Extra pyang versions to run")
}

type ModelInfo struct {
	Name       string
	DocFiles   []string `yaml:"docs"`
	BuildFiles []string `yaml:"build"`
	RunCi      bool     `yaml:"run-ci"`
}

type OpenConfigModelMap struct {
	ModelRoot    string
	ModelInfoMap map[string][]ModelInfo
}

func parseModels(modelRoot string) (OpenConfigModelMap, error) {
	modelInfoMap := map[string][]ModelInfo{}
	err := filepath.Walk(modelRoot, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("prevent panic by handling failure accessing a path %q: %v\n", path, err)
		}
		if !info.IsDir() && info.Name() == ".spec.yml" {
			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("failed to open spec file at path %q: %v\n", path, err)
			}
			m := []ModelInfo{}
			if err := yaml.NewDecoder(file).Decode(&m); err != nil {
				return fmt.Errorf("error while unmarshalling spec file at path %q: %v\n", path, err)
			}

			// Change the build paths to the absolute correct paths.
			for _, info := range m {
				for i, fileName := range info.BuildFiles {
					info.BuildFiles[i] = filepath.Join(modelRoot, strings.TrimPrefix(fileName, "yang/"))
				}
			}

			parentPath := filepath.Dir(path)
			relPath, err := filepath.Rel(modelRoot, parentPath)
			if err != nil {
				return fmt.Errorf("failed to calculate relpath at path %q (modelRoot %q): %v\n", path, modelRoot, err)
			}
			// Allow nested model directories to be used later on as a partial file name.
			relPath = strings.ReplaceAll(relPath, "/", ":")
			modelInfoMap[relPath] = m
		}
		return nil
	})

	return OpenConfigModelMap{ModelRoot: modelRoot, ModelInfoMap: modelInfoMap}, err
}

func createFmtStr(validatorId string) string {
	switch validatorId {
	case "pyang":
		return `if ! $@ -p %s -p %s/third_party/ietf %s &> %s; then
  mv %s %s
fi
`
	case "oc-pyang":
		return `if ! $@ -p %s -p %s/third_party/ietf --openconfig --ignore-error=OC_RELATIVE_PATH %s &> %s; then
  mv %s %s
fi
`
	case "pyangbind":
		return `if ! $@ -p %s -p %s/third_party/ietf -f pybind -o binding.py %s &> %s; then
  mv %s %s
fi
`
	case "goyang-ygot":
		return `if ! go run /go/src/github.com/openconfig/ygot/generator/generator.go \
-path=%s,%s/third_party/ietf \
-output_file=/go/src/github.com/openconfig/ygot/exampleoc/oc.go \
-package_name=exampleoc -generate_fakeroot -fakeroot_name=device -compress_paths=true \
-exclude_modules=ietf-interfaces -generate_rename -generate_append -generate_getters \
-generate_leaf_getters -generate_delete -annotations \
%s &> %s; then
  mv %s %s
fi
`
	case "yanglint":
		return `if ! yanglint -p %s -p %s/third_party/ietf %s &> %s; then
  mv %s %s
fi
`
	}
	return ""
}

func genOpenConfigLinterCmd(g *commonci.GithubRequestHandler, validatorId, folderPath string, modelMap OpenConfigModelMap) string {
	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("#!/bin/bash\nmkdir -p %s\n", folderPath))

	modelDirNames := make([]string, 0, len(modelMap.ModelInfoMap))
	for modelDirName := range modelMap.ModelInfoMap {
		modelDirNames = append(modelDirNames, modelDirName)
	}
	sort.Strings(modelDirNames)

	fmtStr := createFmtStr(validatorId)

	for _, modelDirName := range modelDirNames {
		for _, modelInfo := range modelMap.ModelInfoMap[modelDirName] {
			// First check whether to skip CI.
			if !modelInfo.RunCi || len(modelInfo.BuildFiles) == 0 {
				continue
			} else if disabledModelPaths[modelDirName] {
				log.Infof("skipping disabled model directory %s", modelDirName)
				g.PostLabel("skipped: "+modelDirName, commonci.LabelColors["orange"], owner, repo, prNumber)
				continue
			}
			outputFile := filepath.Join(folderPath, fmt.Sprintf("%s==%s==pass", modelDirName, modelInfo.Name))
			failFile := filepath.Join(folderPath, fmt.Sprintf("%s==%s==fail", modelDirName, modelInfo.Name))
			builder.WriteString(fmt.Sprintf(fmtStr, modelMap.ModelRoot, commonci.RootDir, strings.Join(modelInfo.BuildFiles, " "), outputFile, outputFile, failFile))
		}
	}

	return builder.String()
}

func postPendingStatuses(g *commonci.GithubRequestHandler, validatorId string, versions []string, prApproved bool) []error {
	var errs []error
	validator := commonci.Validators[validatorId]
	for _, version := range versions {
		validatorName := validator.Name + version
		// Update the status to pending so that the user can see that we have received
		// this request and are ready to run the CI.
		update := &commonci.GithubPRUpdate{
			Owner:       owner,
			Repo:        repo,
			Ref:         commitSHA,
			Description: validatorName + " Running",
			NewStatus:   "pending",
			Context:     validatorName,
		}
		if !prApproved && !validator.RunBeforeApproval {
			update.Description = validatorName + " Skipped (PR not approved)"
			update.NewStatus = "error"
		}

		if err := g.UpdatePRStatus(update); err != nil {
			log.Printf("error: couldn't update PR: %s", err)
			errs = append(errs, err)
		}
	}
	return errs
}

func main() {
	flag.Parse()
	repoSplit := strings.Split(repoSlug, "/")
	owner = repoSplit[0]
	repo = repoSplit[1]
	if commitSHA == "" {
		log.Fatalf("no commit SHA")
	}
	if prNumber == 0 {
		log.Fatalf("no PR number")
	}

	modelMap, err := parseModels(modelRoot)
	if err != nil {
		log.Fatalf("CI flow failed due to error encountered while parsing spec files, parseModels: %v", err)
	}

	h := commonci.NewGitHubRequestHandler()

	prApproved, err := h.IsPRApproved(owner, repo, prNumber)
	if err != nil {
		log.Fatalf("warning: Could not check PR approved status, running all checks: %v", err)
		prApproved = true
	}

	versionsToRun := map[string][]string{}
	for validatorId, validator := range commonci.Validators {
		switch validatorId {
		case "pyang":
			versionsToRun["pyang"] = append([]string{""}, strings.Split(extraPV, ",")...)
		default:
			// Empty string is the "head" version, which is always run.
			versionsToRun[validatorId] = []string{""}
		}

		if errs := postPendingStatuses(h, validatorId, versionsToRun[validatorId], prApproved); errs != nil {
			log.Fatal(errs)
		}

		if !prApproved && !validator.RunBeforeApproval {
			// We don't run less important and long tests until PR is approved.
			continue
		}
		// Generate validation commands for the validator.
		for _, version := range versionsToRun[validatorId] {
			validatorResultPath := filepath.Join(commonci.ResultsDir, validatorId+version)
			if err := os.MkdirAll(validatorResultPath, 0644); err != nil {
				log.Fatalf("error while creating directory %q: %v", validatorResultPath, err)
			}
			log.Printf("Created results directory %q", validatorResultPath)

			scriptStr := genOpenConfigLinterCmd(h, validatorId, validatorResultPath, modelMap)
			scriptPath := filepath.Join(validatorResultPath, commonci.ScriptFileName)
			if err := ioutil.WriteFile(scriptPath, []byte(scriptStr), 0744); err != nil {
				log.Printf("error while writing script to path %q: %v", scriptPath, err)
			}
		}
	}

	for validatorId, versions := range versionsToRun {
		extraVersionFile := filepath.Join(commonci.ResultsDir, fmt.Sprintf("extra-%s-versions.txt", validatorId))
		if err := ioutil.WriteFile(extraVersionFile, []byte(strings.Join(versions, " ")), 0444); err != nil {
			log.Fatalf("error while writing extra versions file %q", extraVersionFile)
		}
	}
}
<commit_msg>style improvements<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/openconfig/models-ci/commonci"
	"gopkg.in/yaml.v3"
)

var (
	// flags
	modelRoot string
	repoSlug  string
	commitSHA string
	prNumber  int
	extraPV   string

	// derived flags
	owner string
	repo  string

	// disabledModelPaths are the paths whose models should not undergo CI.
	// These should be temporary.
	disabledModelPaths = map[string]bool{
		"wifi:access-points": true,
		"wifi:ap-manager":    true,
		"wifi:mac":           true,
		"wifi:phy":           true,
	}
)

func init() {
	flag.StringVar(&modelRoot, "modelRoot", "", "root directory to OpenConfig models")
	flag.StringVar(&repoSlug, "repo-slug", "openconfig/public", "repo where CI is run")
	flag.StringVar(&commitSHA, "commit-sha", "", "commit SHA of the PR")
	flag.IntVar(&prNumber, "pr-number", 0, "PR number")
	flag.StringVar(&extraPV, "extra-pyang-versions", "", "Extra pyang versions to run")
}

type ModelInfo struct {
	Name       string
	DocFiles   []string `yaml:"docs"`
	BuildFiles []string `yaml:"build"`
	RunCi      bool     `yaml:"run-ci"`
}

type OpenConfigModelMap struct {
	ModelRoot    string
	ModelInfoMap map[string][]ModelInfo
}

func parseModels(modelRoot string) (OpenConfigModelMap, error) {
	modelInfoMap := map[string][]ModelInfo{}
	err := filepath.Walk(modelRoot, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return fmt.Errorf("prevent panic by handling failure accessing a path %q: %v\n", path, err)
		}
		if !info.IsDir() && info.Name() == ".spec.yml" {
			file, err := os.Open(path)
			if err != nil {
				return fmt.Errorf("failed to open spec file at path %q: %v\n", path, err)
			}
			m := []ModelInfo{}
			if err := yaml.NewDecoder(file).Decode(&m); err != nil {
				return fmt.Errorf("error while unmarshalling spec file at path %q: %v\n", path, err)
			}

			// Change the build paths to the absolute correct paths.
			for _, info := range m {
				for i, fileName := range info.BuildFiles {
					info.BuildFiles[i] = filepath.Join(modelRoot, strings.TrimPrefix(fileName, "yang/"))
				}
			}

			parentPath := filepath.Dir(path)
			relPath, err := filepath.Rel(modelRoot, parentPath)
			if err != nil {
				return fmt.Errorf("failed to calculate relpath at path %q (modelRoot %q): %v\n", path, modelRoot, err)
			}
			// Allow nested model directories to be used later on as a partial file name.
			relPath = strings.ReplaceAll(relPath, "/", ":")
			modelInfoMap[relPath] = m
		}
		return nil
	})

	return OpenConfigModelMap{ModelRoot: modelRoot, ModelInfoMap: modelInfoMap}, err
}

func createFmtStr(validatorId string) string {
	switch validatorId {
	case "pyang":
		return `if ! $@ -p %s -p %s/third_party/ietf %s &> %s; then
  mv %s %s
fi
`
	case "oc-pyang":
		return `if ! $@ -p %s -p %s/third_party/ietf --openconfig --ignore-error=OC_RELATIVE_PATH %s &> %s; then
  mv %s %s
fi
`
	case "pyangbind":
		return `if ! $@ -p %s -p %s/third_party/ietf -f pybind -o binding.py %s &> %s; then
  mv %s %s
fi
`
	case "goyang-ygot":
		return `if ! go run /go/src/github.com/openconfig/ygot/generator/generator.go \
-path=%s,%s/third_party/ietf \
-output_file=/go/src/github.com/openconfig/ygot/exampleoc/oc.go \
-package_name=exampleoc -generate_fakeroot -fakeroot_name=device -compress_paths=true \
-exclude_modules=ietf-interfaces -generate_rename -generate_append -generate_getters \
-generate_leaf_getters -generate_delete -annotations \
%s &> %s; then
  mv %s %s
fi
`
	case "yanglint":
		return `if ! yanglint -p %s -p %s/third_party/ietf %s &> %s; then
  mv %s %s
fi
`
	}
	return ""
}

func genOpenConfigLinterCmd(g *commonci.GithubRequestHandler, validatorId, folderPath string, modelMap OpenConfigModelMap) string {
	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("#!/bin/bash\nmkdir -p %s\n", folderPath))

	modelDirNames := make([]string, 0, len(modelMap.ModelInfoMap))
	for modelDirName := range modelMap.ModelInfoMap {
		modelDirNames = append(modelDirNames, modelDirName)
	}
	sort.Strings(modelDirNames)

	fmtStr := createFmtStr(validatorId)

	for _, modelDirName := range modelDirNames {
		for _, modelInfo := range modelMap.ModelInfoMap[modelDirName] {
			// First check whether to skip CI.
			if !modelInfo.RunCi || len(modelInfo.BuildFiles) == 0 {
				continue
			} else if disabledModelPaths[modelDirName] {
				log.Printf("skipping disabled model directory %s", modelDirName)
				g.PostLabel("skipped: "+modelDirName, commonci.LabelColors["orange"], owner, repo, prNumber)
				continue
			}
			outputFile := filepath.Join(folderPath, fmt.Sprintf("%s==%s==pass", modelDirName, modelInfo.Name))
			failFile := filepath.Join(folderPath, fmt.Sprintf("%s==%s==fail", modelDirName, modelInfo.Name))
			builder.WriteString(fmt.Sprintf(fmtStr, modelMap.ModelRoot, commonci.RootDir, strings.Join(modelInfo.BuildFiles, " "), outputFile, outputFile, failFile))
		}
	}

	return builder.String()
}

func postPendingStatuses(g *commonci.GithubRequestHandler, validatorId string, versions []string, prApproved bool) []error {
	var errs []error
	validator := commonci.Validators[validatorId]
	for _, version := range versions {
		validatorName := validator.Name + version
		// Update the status to pending so that the user can see that we have received
		// this request and are ready to run the CI.
		update := &commonci.GithubPRUpdate{
			Owner:       owner,
			Repo:        repo,
			Ref:         commitSHA,
			Description: validatorName + " Running",
			NewStatus:   "pending",
			Context:     validatorName,
		}
		if !prApproved && !validator.RunBeforeApproval {
			update.Description = validatorName + " Skipped (PR not approved)"
			update.NewStatus = "error"
		}

		if err := g.UpdatePRStatus(update); err != nil {
			log.Printf("error: couldn't update PR: %s", err)
			errs = append(errs, err)
		}
	}
	return errs
}

func main() {
	flag.Parse()
	repoSplit := strings.Split(repoSlug, "/")
	owner = repoSplit[0]
	repo = repoSplit[1]
	if commitSHA == "" {
		log.Fatalf("no commit SHA")
	}
	if prNumber == 0 {
		log.Fatalf("no PR number")
	}

	modelMap, err := parseModels(modelRoot)
	if err != nil {
		log.Fatalf("CI flow failed due to error encountered while parsing spec files, parseModels: %v", err)
	}

	h := commonci.NewGitHubRequestHandler()

	prApproved, err := h.IsPRApproved(owner, repo, prNumber)
	if err != nil {
		log.Fatalf("warning: Could not check PR approved status, running all checks: %v", err)
		prApproved = true
	}

	versionsToRun := map[string][]string{}
	for validatorId, validator := range commonci.Validators {
		switch validatorId {
		case "pyang":
			versionsToRun["pyang"] = append([]string{""}, strings.Split(extraPV, ",")...)
		default:
			// Empty string is the "head" version, which is always run.
			versionsToRun[validatorId] = []string{""}
		}

		if errs := postPendingStatuses(h, validatorId, versionsToRun[validatorId], prApproved); errs != nil {
			log.Fatal(errs)
		}

		if !prApproved && !validator.RunBeforeApproval {
			// We don't run less important and long tests until PR is approved.
			continue
		}
		// Generate validation commands for the validator.
		for _, version := range versionsToRun[validatorId] {
			validatorResultPath := filepath.Join(commonci.ResultsDir, validatorId+version)
			if err := os.MkdirAll(validatorResultPath, 0644); err != nil {
				log.Fatalf("error while creating directory %q: %v", validatorResultPath, err)
			}
			log.Printf("Created results directory %q", validatorResultPath)

			scriptStr := genOpenConfigLinterCmd(h, validatorId, validatorResultPath, modelMap)
			scriptPath := filepath.Join(validatorResultPath, commonci.ScriptFileName)
			if err := ioutil.WriteFile(scriptPath, []byte(scriptStr), 0744); err != nil {
				log.Printf("error while writing script to path %q: %v", scriptPath, err)
			}
		}
	}

	for validatorId, versions := range versionsToRun {
		extraVersionFile := filepath.Join(commonci.ResultsDir, fmt.Sprintf("extra-%s-versions.txt", validatorId))
		if err := ioutil.WriteFile(extraVersionFile, []byte(strings.Join(versions, " ")), 0444); err != nil {
			log.Fatalf("error while writing extra versions file %q", extraVersionFile)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"./src/gdssh"
	"fmt"
	"log"
)

type sfTask struct {
	local  string
	remote string
}

func parseSfOptions(opt GdshOptions) *sfTask {
	task := sfTask{}

	if len(opt.Args) == 2 {
		task.local = opt.Args[0]
		task.remote = opt.Args[1]
	} else if len(opt.Args) == 4 {
		skip := false
		for i, arg := range opt.Args {
			if skip {
				skip = false
				continue
			} else if arg == "--local" {
				task.local = opt.Args[i+1]
				skip = true
			} else if arg == "--remote" {
				task.remote = opt.Args[i+1]
				skip = true
			}
		}
	} else {
		log.Fatal("Wrong number of arguments.")
	}

	return &task
}

func (task *sfTask) Run(conn *gdssh.Conn) error {
	// stat localFile, check for errors
	// copy localfile
	//conn.Scp(task.script.Bytes(), "0555", task.filename)
	fmt.Printf("scp %s %s:%s\n", task.local, conn.Host, task.remote)
	return nil
}

func cmdSendfile(opt GdshOptions) int {
	//stat, err := os.Stat(opt.)
	pool := sshPool(opt)

	/*	f, err := os.Open(opt.Script)
		if err != nil {
			log.Fatal("Could not read script file '", opt.Script, "': ", err)
		}
		io.Copy(run.script, f)
		f.Close()
	} */

	task := parseSfOptions(opt)
	pool.All(task)
	pool.Close()

	return 1
}

// vim: ts=4 sw=4 noet tw=120 softtabstop=4
<commit_msg>Working sendfile.<commit_after>package main

import (
	"./src/gdssh"
	"fmt"
	"log"
	"os"
	"strings"
)

type sfTask struct {
	local  string
	remote string
}

func parseSfOptions(opt GdshOptions) *sfTask {
	task := sfTask{}

	// bare argument style, e.g. gdsh sf /etc/hosts /etc/hosts
	if len(opt.Args) == 2 {
		if strings.HasPrefix(opt.Args[0], "-") {
			log.Fatal("Malformed command? Leading dashes are not allowed with bare arguments.")
		}

		if strings.HasPrefix(opt.Args[1], "-") {
			log.Fatal("Malformed command? Leading dashes are not allowed with bare arguments.")
		}

		_, err := os.Stat(opt.Args[0])
		if err != nil {
			log.Fatal(fmt.Sprintf("Not enough arguments. '%s' is not a file: %s", opt.Args[0], err))
		}

		task.local = opt.Args[0]
		task.remote = opt.Args[1]

		// fully-specified style, e.g. gdsh scp --local /etc/hosts --remote /etc/hosts
	} else if len(opt.Args) == 4 {
		skip := false
		for i, arg := range opt.Args {
			if skip {
				skip = false
				continue
			} else if arg == "--local" || arg == "-L" {
				_, err := os.Stat(opt.Args[i+1])
				if err != nil {
					log.Fatal(fmt.Sprintf("Invalid local file '%s': %s", opt.Args[i+1], err))
				}
				task.local = opt.Args[i+1]
				skip = true
			} else if arg == "--remote" || arg == "-R" {
				task.remote = opt.Args[i+1]
				skip = true
			}
		}
	} else {
		log.Fatal("Wrong number of arguments.")
	}

	return &task
}

// the file will be opened for each remote host, but that's fine since
// the reads will end up shared on modern operating systems
func (task *sfTask) Run(conn *gdssh.Conn) error {
	conn.Scp(task.local, task.remote)
	return nil
}

func cmdSendfile(opt GdshOptions) int {
	pool := sshPool(opt)
	task := parseSfOptions(opt)
	pool.All(task)
	pool.Close()

	return 1
}

// vim: ts=4 sw=4 noet tw=120 softtabstop=4
<|endoftext|>


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

// Cat concatenates files and prints to stdout.
//
// Synopsis:
//     cat [-u] [FILES]...
//
// Description:
//     If no files are specified, read from stdin.
//
// Options:
//     -u: ignored flag
package main

import (
	"flag"
	"io"
	"log"
	"os"
)

var (
	_ = flag.Bool("u", false, "ignored")
)

func catFile(w io.Writer, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = io.Copy(w, f)
	return err
}

func cat(w io.Writer, files []string) error {
	for _, name := range files {
		if err := catFile(w, name); err != nil {
			return err
		}
	}
	return nil
}

func main() {
	flag.Parse()

	if len(os.Args) == 1 {
		if _, err := io.Copy(os.Stdout, os.Stdin); err != nil {
			log.Fatalf("error concatenating stdin to stdout: %v", err)
		}
	}

	if err := cat(os.Stdout, flag.Args()); err != nil {
		log.Fatalf("cat: %v", err)
	}
}
<commit_msg>cat: Fix arg counting<commit_after>// Copyright 2012-2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Cat concatenates files and prints to stdout.
//
// Synopsis:
//     cat [-u] [FILES]...
//
// Description:
//     If no files are specified, read from stdin.
//
// Options:
//     -u: ignored flag
package main

import (
	"flag"
	"io"
	"log"
	"os"
)

var (
	_ = flag.Bool("u", false, "ignored")
)

func catFile(w io.Writer, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = io.Copy(w, f)
	return err
}

func cat(w io.Writer, files []string) error {
	for _, name := range files {
		if err := catFile(w, name); err != nil {
			return err
		}
	}
	return nil
}

func main() {
	flag.Parse()

	if flag.NArg() == 0 {
		if _, err := io.Copy(os.Stdout, os.Stdin); err != nil {
			log.Fatalf("error concatenating stdin to stdout: %v", err)
		}
	}

	if err := cat(os.Stdout, flag.Args()); err != nil {
		log.Fatalf("cat: %v", err)
	}
}
<|endoftext|>


<commit_before>// Package cmds provides commands to a bot
package cmds

import (
    "fmt"
    "os"
    "time"

    "encoding/json"
    "net/http"

    "github.com/weirdtales/senor-rosado/slack"
)

// partial struct for a google map api hit - so gross
type weatherLocation struct {
    Results []struct {
        FormattedAddress string `json:"formatted_address"`
        Geometry struct {
            Location struct {
                Lat float64 `json:"lat"`
                Lng float64 `json:"lng"`
            } `json:"location"`
        } `json:"geometry"`
        Status string `json:"status"`
    } `json:"results"`
}

// partial struct for a darksky api hit - so pretty
type weatherInfo struct {
    Latitude float64 `json:"latitude"`
    Longitude float64 `json:"longitude"`
    Timezone string `json:"timezone"`
    Currently struct {
        Time int64 `json:"time"`
        Summary string `json:"summary"`
        Icon string `json:"icon"`
        Temperature float64 `json:"temperature"`
        ApparentTemperature float64 `json:"apparentTemperature"`
        DewPoint float64 `json:"dewPoint"`
        Humidity float64 `json:"humidity"`
        WindSpeed float64 `json:"windSpeed"`
        Visibility float64 `json:"visibility"`
        CloudCover float64 `json:"cloudCover"`
        Pressure float64 `json:"pressure"`
        Ozone float64 `json:"ozone"`
    } `json:"currently"`
    Daily struct {
        Summary string `json:"summary"`
    } `json:"daily"`
    Hourly struct {
        Summary string `json:"summary"`
    } `json:"hourly"`

    GoogleName string // use the google maps FormattedAddress value - it's good
}

var httpClient = &http.Client{Timeout: 10 * time.Second}

// Weather reports the current forecast for a location
func Weather(m slack.Message, c slack.Conn) {
    // TODO needs to be cached for a minute or so
    reply := getReply(m)
    reply.Channel = m.Channel
    c.Send(m, reply)
}

func getReply(m slack.Message) (r slack.Reply) {
    token := os.Getenv("DARKSKY_API")
    if token == "" {
        r.Text = "I got no DARKSKY_API token in my env. Sort it."
        return
    }

    if m.Subcommand == "" {
        r.Text = "Specify a location, fool, then I can help."
        return
    }

    loc, err := getLocation(m.Subcommand)
    if err != nil {
        r.Text = fmt.Sprintf("Problem getting location: %s", err)
        return
    }

    w, err := getWeather(token, loc)
    if err != nil {
        r.Text = fmt.Sprintf("Problem getting weather: %s", err)
        return
    }

    report := "*%s* :point_right: %s %s, %.1f°C. %s %s"
    t := time.Unix(w.Currently.Time, 0).Format("15:04")
    r.Text = fmt.Sprintf(report, w.GoogleName, t, w.Currently.Summary, w.Currently.Temperature, w.Hourly.Summary, w.Daily.Summary)
    return
}

// returns a weatherLocation struct from google maps api
func getLocation(q string) (loc weatherLocation, err error) {
    url := "http://maps.googleapis.com/maps/api/geocode/json?address=%s&sensor=false"
    url = fmt.Sprintf(url, q) // TODO does `q` need to be sanitized?

    err = getJson(url, &loc)
    if err != nil {
        return
    }

    if len(loc.Results) < 1 {
        err = fmt.Errorf("google maps failed to convert the query into a location")
    }
    return
}

// returns a weatherInfo struct from darksky
func getWeather(token string, loc weatherLocation) (w weatherInfo, err error) {
    w.GoogleName = loc.Results[0].FormattedAddress
    url := "https://api.darksky.net/forecast/%s/%f,%f?units=si"
    url = fmt.Sprintf(url, token, loc.Results[0].Geometry.Location.Lat, loc.Results[0].Geometry.Location.Lng)

    err = getJson(url, &w)
    if err != nil {
        return
    }

    if w.Timezone == "" {
        err = fmt.Errorf("darksky doesn't have weather data for the location")
    }
    return
}

// pointless scaffolding
func getJson(url string, target interface{}) error {
    r, err := httpClient.Get(url)
    if err != nil {
        return err
    }
    defer r.Body.Close()
    if r.StatusCode != 200 {
        return err
    }

    return json.NewDecoder(r.Body).Decode(target)
}
<commit_msg>clean up<commit_after>// Package cmds provides commands to a bot.
package cmds

import (
    "fmt"
    "os"
    "time"

    "encoding/json"
    "net/http"

    "github.com/weirdtales/senor-rosado/slack"
)

// partial struct for a google map api hit - so gross.
type weatherLocation struct {
    Results []struct {
        FormattedAddress string `json:"formatted_address"`
        Geometry struct {
            Location struct {
                Lat float64 `json:"lat"`
                Lng float64 `json:"lng"`
            } `json:"location"`
        } `json:"geometry"`
        Status string `json:"status"`
    } `json:"results"`
}

// partial struct for a darksky api hit - so pretty.
type weatherInfo struct {
    Latitude float64 `json:"latitude"`
    Longitude float64 `json:"longitude"`
    Timezone string `json:"timezone"`
    Currently struct {
        Time int64 `json:"time"`
        Summary string `json:"summary"`
        Icon string `json:"icon"`
        Temperature float64 `json:"temperature"`
        ApparentTemperature float64 `json:"apparentTemperature"`
        DewPoint float64 `json:"dewPoint"`
        Humidity float64 `json:"humidity"`
        WindSpeed float64 `json:"windSpeed"`
        Visibility float64 `json:"visibility"`
        CloudCover float64 `json:"cloudCover"`
        Pressure float64 `json:"pressure"`
        Ozone float64 `json:"ozone"`
    } `json:"currently"`
    Daily struct {
        Summary string `json:"summary"`
    } `json:"daily"`
    Hourly struct {
        Summary string `json:"summary"`
    } `json:"hourly"`

    GoogleName string // use the google maps FormattedAddress value - it's good
}

var httpClient = &http.Client{Timeout: 10 * time.Second}

// Weather reports the current forecast for a location.
func Weather(m slack.Message, c slack.Conn) {
    // TODO needs to be cached for a minute or so
    reply := getReply(m)
    reply.Channel = m.Channel
    c.Send(m, reply)
}

// eats a slack.Message and returns a slack.Reply with some text set.
func getReply(m slack.Message) (r slack.Reply) {
    token := os.Getenv("DARKSKY_API")
    if token == "" {
        r.Text = "I got no DARKSKY_API token in my env. Sort it."
        return
    }

    if m.Subcommand == "" {
        r.Text = "Specify a location, fool, then I can help."
        return
    }

    loc, err := getLocation(m.Subcommand)
    if err != nil {
        r.Text = fmt.Sprintf("Problem getting location: %s", err)
        return
    }

    w, err := getWeather(token, loc)
    if err != nil {
        r.Text = fmt.Sprintf("Problem getting weather: %s", err)
        return
    }

    report := "*%s* :point_right: %s, %.1f°C. %s %s"
    r.Text = fmt.Sprintf(report, w.GoogleName, w.Currently.Summary, w.Currently.Temperature, w.Hourly.Summary, w.Daily.Summary)
    return
}

// returns a weatherLocation struct from google maps api.
func getLocation(q string) (loc weatherLocation, err error) {
    url := "http://maps.googleapis.com/maps/api/geocode/json?address=%s&sensor=false"
    url = fmt.Sprintf(url, q) // TODO does `q` need to be sanitized?

    err = getJson(url, &loc)
    if err != nil {
        return
    }

    if len(loc.Results) < 1 {
        err = fmt.Errorf("google maps failed to convert the query into a location")
    }
    return
}

// returns a weatherInfo struct from darksky.
func getWeather(token string, loc weatherLocation) (w weatherInfo, err error) {
    w.GoogleName = loc.Results[0].FormattedAddress
    url := "https://api.darksky.net/forecast/%s/%f,%f?units=si"
    url = fmt.Sprintf(url, token, loc.Results[0].Geometry.Location.Lat, loc.Results[0].Geometry.Location.Lng)

    err = getJson(url, &w)
    if err != nil {
        return
    }

    if w.Timezone == "" {
        err = fmt.Errorf("darksky doesn't have weather data for the location")
    }
    return
}

// pointless scaffolding.
func getJson(url string, target interface{}) error {
    r, err := httpClient.Get(url)
    if err != nil {
        return err
    }
    defer r.Body.Close()
    if r.StatusCode != 200 {
        return err
    }

    return json.NewDecoder(r.Body).Decode(target)
}
<|endoftext|>


<commit_before>// Package cmdtmpl provides the templated command system for mcdev
package cmdtmpl

import (
	"bytes"
	"os"
	"os/exec"
	"text/template"
)

type Command struct {
	Cmd  string
	Args []*template.Template
}

func NewCommand(args []string) (*Command, error) {
	result := new(Command)
	result.Cmd = args[0]
	result.Args = make([]*template.Template, len(args)-1)

	for i, arg := range args[1:] {
		t, err := template.New("arg").Parse(arg)
		if err != nil {
			return nil, err
		}
		result.Args[i] = t
	}
	return result, nil
}

func (cmd *Command) Run(ctx interface{}) error {
	args := make([]string, len(cmd.Args))

	for i, t := range cmd.Args {
		var buf bytes.Buffer
		err := t.Execute(&buf, ctx)
		if err != nil {
			return err
		}
		args[i] = buf.String()
	}

	proc := exec.Command(cmd.Cmd, args...)
	proc.Stdout = os.Stdout
	proc.Stderr = os.Stderr

	return proc.Run()
}
<commit_msg>Add Make to cmdtpl.Command<commit_after>// Package cmdtmpl provides the templated command system for mcdev
package cmdtmpl

import (
	"bytes"
	"os"
	"os/exec"
	"text/template"
)

type Command struct {
	Cmd  string
	Args []*template.Template
}

func NewCommand(args []string) (*Command, error) {
	result := new(Command)
	result.Cmd = args[0]
	result.Args = make([]*template.Template, len(args)-1)

	for i, arg := range args[1:] {
		t, err := template.New("arg").Parse(arg)
		if err != nil {
			return nil, err
		}
		result.Args[i] = t
	}
	return result, nil
}

func (cmd *Command) Make(ctx interface{}) (*exec.Cmd, error) {
	args := make([]string, len(cmd.Args))

	for i, t := range cmd.Args {
		var buf bytes.Buffer
		err := t.Execute(&buf, ctx)
		if err != nil {
			return nil, err
		}
		args[i] = buf.String()
	}

	proc := exec.Command(cmd.Cmd, args...)
	proc.Stdout = os.Stdout
	proc.Stderr = os.Stderr
	return proc, nil
}

func (cmd *Command) Run(ctx interface{}) error {
	proc, err := cmd.Make(ctx)
	if err != nil {
		return err
	}
	return proc.Run()
}
<|endoftext|>


<commit_before>package coal

import (
	"context"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// Tracer is used by a traced collection to push tracing spans for database
// queries.
type Tracer interface {
	Push(name string)
	Tag(key string, value interface{})
	Log(key string, value interface{})
	Pop()
}

// TracedCollection wraps a collection to automatically push tracing spans for
// run queries.
type TracedCollection struct {
	coll   *mongo.Collection
	tracer Tracer
}

// AggregateAll wraps the native Aggregate collection method and decodes all
// documents to the provided slice.
func (c *TracedCollection) AggregateAll(ctx context.Context, slicePtr interface{}, pipeline interface{}, opts ...*options.AggregateOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Aggregate")
	c.tracer.Tag("pipeline", pipeline)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Aggregate(ctx, pipeline, opts...)
	if err != nil {
		return err
	}

	// decode all documents
	err = csr.All(ctx, slicePtr)
	if err != nil {
		return err
	}

	return nil
}

// BulkWrite wraps the native BulkWrite collection method.
func (c *TracedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.BulkWrite")
	defer c.tracer.Pop()

	// run query
	return c.coll.BulkWrite(ctx, models, opts...)
}

// CountDocuments wraps the native CountDocuments collection method.
func (c *TracedCollection) CountDocuments(ctx context.Context, filter interface{}, opts ...*options.CountOptions) (int64, error) {
	// push span
	c.tracer.Push("mongo/Collection.CountDocuments")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.CountDocuments(ctx, filter, opts...)
}

// DeleteMany wraps the native DeleteMany collection method.
func (c *TracedCollection) DeleteMany(ctx context.Context, filter interface{}, opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.DeleteMany")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.DeleteMany(ctx, filter, opts...)
}

// DeleteOne wraps the native DeleteOne collection method.
func (c *TracedCollection) DeleteOne(ctx context.Context, filter interface{}, opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.DeleteOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.DeleteOne(ctx, filter, opts...)
}

// Distinct wraps the native Distinct collection method.
func (c *TracedCollection) Distinct(ctx context.Context, fieldName string, filter interface{}, opts ...*options.DistinctOptions) ([]interface{}, error) {
	// push span
	c.tracer.Push("mongo/Collection.Distinct")
	c.tracer.Tag("fieldName", fieldName)
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.Distinct(ctx, fieldName, filter, opts...)
}

// EstimatedDocumentCount wraps the native EstimatedDocumentCount collection method.
func (c *TracedCollection) EstimatedDocumentCount(ctx context.Context, opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
	// push span
	c.tracer.Push("mongo/Collection.EstimatedDocumentCount")
	defer c.tracer.Pop()

	// run query
	return c.coll.EstimatedDocumentCount(ctx, opts...)
}

// FindAll wraps the native Find collection method and decodes all documents to
// the provided slice.
func (c *TracedCollection) FindAll(ctx context.Context, slicePtr interface{}, filter interface{}, opts ...*options.FindOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Find")
	c.tracer.Tag("filter", filter)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Find(ctx, filter, opts...)
	if err != nil {
		return err
	}

	// decode all documents
	err = csr.All(ctx, slicePtr)
	if err != nil {
		return err
	}

	return nil
}

// FindOne wraps the native FindOne collection method.
func (c *TracedCollection) FindOne(ctx context.Context, filter interface{}, opts ...*options.FindOneOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOne(ctx, filter, opts...)
}

// FindOneAndDelete wraps the native FindOneAndDelete collection method.
func (c *TracedCollection) FindOneAndDelete(ctx context.Context, filter interface{}, opts ...*options.FindOneAndDeleteOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndDelete")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndDelete(ctx, filter, opts...)
}

// FindOneAndReplace wraps the native FindOneAndReplace collection method.
func (c *TracedCollection) FindOneAndReplace(ctx context.Context, filter interface{}, replacement interface{}, opts ...*options.FindOneAndReplaceOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndReplace")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndReplace(ctx, filter, replacement, opts...)
}

// FindOneAndUpdate wraps the native FindOneAndUpdate collection method.
func (c *TracedCollection) FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{}, opts ...*options.FindOneAndUpdateOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndUpdate")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndUpdate(ctx, filter, update, opts...)
}

// InsertMany wraps the native InsertMany collection method.
func (c *TracedCollection) InsertMany(ctx context.Context, documents []interface{}, opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.InsertMany")
	defer c.tracer.Pop()

	// run query
	return c.coll.InsertMany(ctx, documents, opts...)
}

// InsertOne wraps the native InsertOne collection method.
func (c *TracedCollection) InsertOne(ctx context.Context, document interface{}, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.InsertOne")
	defer c.tracer.Pop()

	// run query
	return c.coll.InsertOne(ctx, document, opts...)
}

// ReplaceOne wraps the native ReplaceOne collection method.
func (c *TracedCollection) ReplaceOne(ctx context.Context, filter interface{}, replacement interface{}, opts ...*options.ReplaceOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.ReplaceOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.ReplaceOne(ctx, filter, replacement, opts...)
}

// UpdateMany wraps the native UpdateMany collection method.
func (c *TracedCollection) UpdateMany(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.UpdateMany")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.UpdateMany(ctx, filter, update, opts...)
}

// UpdateOne wraps the native UpdateOne collection method.
func (c *TracedCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.UpdateOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.UpdateOne(ctx, filter, update, opts...)
}
<commit_msg>added traced itera variants<commit_after>package coal

import (
	"context"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// Tracer is used by a traced collection to push tracing spans for database
// queries.
type Tracer interface {
	Push(name string)
	Tag(key string, value interface{})
	Log(key string, value interface{})
	Pop()
}

// TracedCollection wraps a collection to automatically push tracing spans for
// run queries.
type TracedCollection struct {
	coll   *mongo.Collection
	tracer Tracer
}

// AggregateAll wraps the native Aggregate collection method and decodes all
// documents to the provided slice.
func (c *TracedCollection) AggregateAll(ctx context.Context, slicePtr interface{}, pipeline interface{}, opts ...*options.AggregateOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Aggregate")
	c.tracer.Tag("pipeline", pipeline)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Aggregate(ctx, pipeline, opts...)
	if err != nil {
		return err
	}

	// decode all documents
	err = csr.All(ctx, slicePtr)
	if err != nil {
		return err
	}

	return nil
}

// AggregateIter wraps the native Aggregate collection method and calls the
// provided callback with the decode method until an error is returned or the
// cursor has been exhausted.
func (c *TracedCollection) AggregateIter(ctx context.Context, pipeline interface{}, fn func(func(interface{}) error) error, opts ...*options.AggregateOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Aggregate")
	c.tracer.Tag("pipeline", pipeline)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Aggregate(ctx, pipeline, opts...)
	if err != nil {
		return err
	}

	// ensure cursor is closed
	defer csr.Close(ctx)

	// iterate over all documents
	for csr.Next(ctx) {
		err = fn(csr.Decode)
		if err != nil {
			return err
		}
	}

	// close cursor
	err = csr.Close(nil)
	if err != nil {
		return err
	}

	return nil
}

// BulkWrite wraps the native BulkWrite collection method.
func (c *TracedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.BulkWrite")
	defer c.tracer.Pop()

	// run query
	return c.coll.BulkWrite(ctx, models, opts...)
}

// CountDocuments wraps the native CountDocuments collection method.
func (c *TracedCollection) CountDocuments(ctx context.Context, filter interface{}, opts ...*options.CountOptions) (int64, error) {
	// push span
	c.tracer.Push("mongo/Collection.CountDocuments")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.CountDocuments(ctx, filter, opts...)
}

// DeleteMany wraps the native DeleteMany collection method.
func (c *TracedCollection) DeleteMany(ctx context.Context, filter interface{}, opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.DeleteMany")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.DeleteMany(ctx, filter, opts...)
}

// DeleteOne wraps the native DeleteOne collection method.
func (c *TracedCollection) DeleteOne(ctx context.Context, filter interface{}, opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.DeleteOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.DeleteOne(ctx, filter, opts...)
}

// Distinct wraps the native Distinct collection method.
func (c *TracedCollection) Distinct(ctx context.Context, fieldName string, filter interface{}, opts ...*options.DistinctOptions) ([]interface{}, error) {
	// push span
	c.tracer.Push("mongo/Collection.Distinct")
	c.tracer.Tag("fieldName", fieldName)
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.Distinct(ctx, fieldName, filter, opts...)
}

// EstimatedDocumentCount wraps the native EstimatedDocumentCount collection method.
func (c *TracedCollection) EstimatedDocumentCount(ctx context.Context, opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
	// push span
	c.tracer.Push("mongo/Collection.EstimatedDocumentCount")
	defer c.tracer.Pop()

	// run query
	return c.coll.EstimatedDocumentCount(ctx, opts...)
}

// FindAll wraps the native Find collection method and decodes all documents to
// the provided slice.
func (c *TracedCollection) FindAll(ctx context.Context, slicePtr interface{}, filter interface{}, opts ...*options.FindOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Find")
	c.tracer.Tag("filter", filter)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Find(ctx, filter, opts...)
	if err != nil {
		return err
	}

	// decode all documents
	err = csr.All(ctx, slicePtr)
	if err != nil {
		return err
	}

	return nil
}

// FindIter wraps the native Find collection method and calls the provided
// callback with the decode method until an error is returned or the cursor has
// been exhausted.
func (c *TracedCollection) FindIter(ctx context.Context, filter interface{}, fn func(func(interface{}) error) error, opts ...*options.FindOptions) error {
	// push span
	c.tracer.Push("mongo/Collection.Find")
	c.tracer.Tag("filter", filter)
	defer c.tracer.Pop()

	// run query
	csr, err := c.coll.Find(ctx, filter, opts...)
	if err != nil {
		return err
	}

	// ensure cursor is closed
	defer csr.Close(ctx)

	// iterate over all documents
	for csr.Next(ctx) {
		err = fn(csr.Decode)
		if err != nil {
			return err
		}
	}

	// close cursor
	err = csr.Close(nil)
	if err != nil {
		return err
	}

	return nil
}

// FindOne wraps the native FindOne collection method.
func (c *TracedCollection) FindOne(ctx context.Context, filter interface{}, opts ...*options.FindOneOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOne(ctx, filter, opts...)
}

// FindOneAndDelete wraps the native FindOneAndDelete collection method.
func (c *TracedCollection) FindOneAndDelete(ctx context.Context, filter interface{}, opts ...*options.FindOneAndDeleteOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndDelete")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndDelete(ctx, filter, opts...)
}

// FindOneAndReplace wraps the native FindOneAndReplace collection method.
func (c *TracedCollection) FindOneAndReplace(ctx context.Context, filter interface{}, replacement interface{}, opts ...*options.FindOneAndReplaceOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndReplace")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndReplace(ctx, filter, replacement, opts...)
}

// FindOneAndUpdate wraps the native FindOneAndUpdate collection method.
func (c *TracedCollection) FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{}, opts ...*options.FindOneAndUpdateOptions) *mongo.SingleResult {
	// push span
	c.tracer.Push("mongo/Collection.FindOneAndUpdate")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.FindOneAndUpdate(ctx, filter, update, opts...)
}

// InsertMany wraps the native InsertMany collection method.
func (c *TracedCollection) InsertMany(ctx context.Context, documents []interface{}, opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.InsertMany")
	defer c.tracer.Pop()

	// run query
	return c.coll.InsertMany(ctx, documents, opts...)
}

// InsertOne wraps the native InsertOne collection method.
func (c *TracedCollection) InsertOne(ctx context.Context, document interface{}, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.InsertOne")
	defer c.tracer.Pop()

	// run query
	return c.coll.InsertOne(ctx, document, opts...)
}

// ReplaceOne wraps the native ReplaceOne collection method.
func (c *TracedCollection) ReplaceOne(ctx context.Context, filter interface{}, replacement interface{}, opts ...*options.ReplaceOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.ReplaceOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.ReplaceOne(ctx, filter, replacement, opts...)
}

// UpdateMany wraps the native UpdateMany collection method.
func (c *TracedCollection) UpdateMany(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.UpdateMany")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.UpdateMany(ctx, filter, update, opts...)
}

// UpdateOne wraps the native UpdateOne collection method.
func (c *TracedCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	// push span
	c.tracer.Push("mongo/Collection.UpdateOne")
	c.tracer.Log("filter", filter)
	defer c.tracer.Pop()

	// run query
	return c.coll.UpdateOne(ctx, filter, update, opts...)
}
<|endoftext|>


<commit_before>package wkhtmltopdf

import (
	"bytes"
	"io/ioutil"
	"os"
	"reflect"
	"runtime"
	"strings"
	"testing"
)

func newTestPDFGenerator(tb testing.TB) *PDFGenerator {

	pdfg, err := NewPDFGenerator()
	if err != nil {
		tb.Fatal(err)
	}

	pdfg.Dpi.Set(600)
	pdfg.NoCollate.Set(false)
	pdfg.PageSize.Set(PageSizeA4)
	pdfg.MarginBottom.Set(40)
	pdfg.MarginLeft.Set(0)

	page1 := NewPage("https://www.google.com")

	page1.DisableSmartShrinking.Set(true)
	page1.HeaderSpacing.Set(10.01)
	page1.Allow.Set("/usr/local/html")
	page1.Allow.Set("/usr/local/images")
	page1.CustomHeader.Set("X-AppKey", "abcdef")
	page1.ViewportSize.Set("3840x2160")

	if runtime.GOOS == "darwin" {
		page1.LoadErrorHandling.Set("ignore")
	}

	pdfg.AddPage(page1)

	pdfg.Cover.Input = "https://wkhtmltopdf.org/index.html"
	pdfg.Cover.Zoom.Set(0.75)

	pdfg.TOC.Include = true
	pdfg.TOC.DisableDottedLines.Set(true)

	return pdfg
}

func wantArgString() string {
	return "--dpi 600 --margin-bottom 40 --margin-left 0 --page-size A4 cover https://wkhtmltopdf.org/index.html --zoom 0.750 toc --disable-dotted-lines page https://www.google.com --allow /usr/local/html --allow /usr/local/images --custom-header X-AppKey abcdef --disable-smart-shrinking --viewport-size 3840x2160 --header-spacing 10.010 -"
}

func TestArgString(t *testing.T) {
	pdfg := newTestPDFGenerator(t)
	want := wantArgString()
	if pdfg.ArgString() != want {
		t.Errorf("Want argstring:\n%s\nHave:\n%s", want, pdfg.ArgString())
	}
	pdfg.SetPages(pdfg.pages)
	if pdfg.ArgString() != want {
		t.Errorf("Want argstring:\n%s\nHave:\n%s", want, pdfg.ArgString())
	}
}

func TestVersion(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	pdfg.Version.Set(true)
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
}

func TestNoInput(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.Create()
	if err == nil {
		t.Fatal("Want an error when there is no input, have no error")
	}
	//TODO temp error check because older versions of wkhtmltopdf return a different error :(
	wantErrNew := "You need to specify at least one input file, and exactly one output file"
	wantErrOld := "You need to specify atleast one input file, and exactly one output file"
	if strings.HasPrefix(err.Error(), wantErrNew) == false && strings.HasPrefix(err.Error(), wantErrOld) == false {
		t.Errorf("Want error prefix %s or %s, have %s", wantErrNew, wantErrOld, err.Error())
	}
}

func TestGeneratePDF(t *testing.T) {
	pdfg := newTestPDFGenerator(t)
	err := pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePDF.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
}

func TestGeneratePdfFromStdinSimple(t *testing.T) {
	//Use a new blank PDF generator
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	htmlfile, err := ioutil.ReadFile("./testfiles/htmlsimple.html")
	if err != nil {
		t.Fatal(err)
	}
	pdfg.AddPage(NewPageReader(bytes.NewReader(htmlfile)))
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePdfFromStdinSimple.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
	if pdfg.Buffer().Len() != len(pdfg.Bytes()) {
		t.Errorf("Buffersize not equal")
	}
}

func TestPDFGeneratorOutputFile(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	htmlfile, err := os.Open("./testfiles/htmlsimple.html")
	if err != nil {
		t.Fatal(err)
	}
	defer htmlfile.Close()

	pdfg.OutputFile = "./testfiles/TestPDFGeneratorOutputFile.pdf"

	pdfg.AddPage(NewPageReader(htmlfile))
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}

	pdfFile, err := os.Open("./testfiles/TestPDFGeneratorOutputFile.pdf")
	if err != nil {
		t.Fatal(err)
	}
	defer pdfFile.Close()

	stat, err := pdfFile.Stat()
	if err != nil {
		t.Fatal(err)
	}
	if stat.Size() < 100 {
		t.Errorf("generated PDF is size under 100 bytes")
	}
}

func TestGeneratePdfFromStdinHtml5(t *testing.T) {
	//Use newTestPDFGenerator and append to page1 and TOC
	pdfg := newTestPDFGenerator(t)
	htmlfile, err := ioutil.ReadFile("./testfiles/html5.html")
	if err != nil {
		t.Fatal(err)
	}
	page2 := NewPageReader(bytes.NewReader(htmlfile))
	pdfg.AddPage(page2)
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePdfFromStdinHtml5.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
}

func TestPath(t *testing.T) {
	path := "/usr/wkhtmltopdf/wkhtmltopdf"
	SetPath(path)
	defer SetPath("")
	if GetPath() != path {
		t.Errorf("Have path %q, want %q", GetPath(), path)
	}
}

func BenchmarkArgs(b *testing.B) {
	pdfg := newTestPDFGenerator(b)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pdfg.Args()
	}
}

func TestStringOption(t *testing.T) {
	opt := stringOption{
		option: "stringopt",
	}
	opt.Set("value123")

	want := []string{"--stringopt", "value123"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestSliceOption(t *testing.T) {
	opt := sliceOption{
		option: "sliceopt",
	}
	opt.Set("string15183")
	opt.Set("foo")
	opt.Set("bar")

	want := []string{"--sliceopt", "string15183", "--sliceopt", "foo", "--sliceopt", "bar"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestMapOption(t *testing.T) {
	opt := mapOption{
		option: "mapopt",
	}

	opt.Set("key1", "foo")
	opt.Set("key2", "bar")
	opt.Set("key3", "Hello")

	result := strings.Join(opt.Parse(), " ")
	if !strings.Contains(result, "--mapopt key1 foo") {
		t.Error("missing map option key1")
	}
	if !strings.Contains(result, "--mapopt key2 bar") {
		t.Error("missing map option key2")
	}
	if !strings.Contains(result, "--mapopt key3 Hello") {
		t.Error("missing map option key3")
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestUIntOption(t *testing.T) {
	opt := uintOption{
		option: "uintopt",
	}
	opt.Set(14860)

	want := []string{"--uintopt", "14860"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestFloatOption(t *testing.T) {
	opt := floatOption{
		option: "flopt",
	}
	opt.Set(239.75)

	want := []string{"--flopt", "239.750"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestBoolOption(t *testing.T) {
	opt := boolOption{
		option: "boolopt",
	}
	opt.Set(true)

	want := []string{"--boolopt"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}
<commit_msg>remove OSX flag in tests<commit_after>package wkhtmltopdf

import (
	"bytes"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"testing"
)

func newTestPDFGenerator(tb testing.TB) *PDFGenerator {

	pdfg, err := NewPDFGenerator()
	if err != nil {
		tb.Fatal(err)
	}

	pdfg.Dpi.Set(600)
	pdfg.NoCollate.Set(false)
	pdfg.PageSize.Set(PageSizeA4)
	pdfg.MarginBottom.Set(40)
	pdfg.MarginLeft.Set(0)

	page1 := NewPage("https://www.google.com")

	page1.DisableSmartShrinking.Set(true)
	page1.HeaderSpacing.Set(10.01)
	page1.Allow.Set("/usr/local/html")
	page1.Allow.Set("/usr/local/images")
	page1.CustomHeader.Set("X-AppKey", "abcdef")
	page1.ViewportSize.Set("3840x2160")

	pdfg.AddPage(page1)

	pdfg.Cover.Input = "https://wkhtmltopdf.org/index.html"
	pdfg.Cover.Zoom.Set(0.75)

	pdfg.TOC.Include = true
	pdfg.TOC.DisableDottedLines.Set(true)

	return pdfg
}

func wantArgString() string {
	return "--dpi 600 --margin-bottom 40 --margin-left 0 --page-size A4 cover https://wkhtmltopdf.org/index.html --zoom 0.750 toc --disable-dotted-lines page https://www.google.com --allow /usr/local/html --allow /usr/local/images --custom-header X-AppKey abcdef --disable-smart-shrinking --viewport-size 3840x2160 --header-spacing 10.010 -"
}

func TestArgString(t *testing.T) {
	pdfg := newTestPDFGenerator(t)
	want := wantArgString()
	if pdfg.ArgString() != want {
		t.Errorf("Want argstring:\n%s\nHave:\n%s", want, pdfg.ArgString())
	}
	pdfg.SetPages(pdfg.pages)
	if pdfg.ArgString() != want {
		t.Errorf("Want argstring:\n%s\nHave:\n%s", want, pdfg.ArgString())
	}
}

func TestVersion(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	pdfg.Version.Set(true)
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
}

func TestNoInput(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.Create()
	if err == nil {
		t.Fatal("Want an error when there is no input, have no error")
	}
	//TODO temp error check because older versions of wkhtmltopdf return a different error :(
	wantErrNew := "You need to specify at least one input file, and exactly one output file"
	wantErrOld := "You need to specify atleast one input file, and exactly one output file"
	if strings.HasPrefix(err.Error(), wantErrNew) == false && strings.HasPrefix(err.Error(), wantErrOld) == false {
		t.Errorf("Want error prefix %s or %s, have %s", wantErrNew, wantErrOld, err.Error())
	}
}

func TestGeneratePDF(t *testing.T) {
	pdfg := newTestPDFGenerator(t)
	err := pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePDF.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
}

func TestGeneratePdfFromStdinSimple(t *testing.T) {
	//Use a new blank PDF generator
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	htmlfile, err := ioutil.ReadFile("./testfiles/htmlsimple.html")
	if err != nil {
		t.Fatal(err)
	}
	pdfg.AddPage(NewPageReader(bytes.NewReader(htmlfile)))
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePdfFromStdinSimple.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
	if pdfg.Buffer().Len() != len(pdfg.Bytes()) {
		t.Errorf("Buffersize not equal")
	}
}

func TestPDFGeneratorOutputFile(t *testing.T) {
	pdfg, err := NewPDFGenerator()
	if err != nil {
		t.Fatal(err)
	}
	htmlfile, err := os.Open("./testfiles/htmlsimple.html")
	if err != nil {
		t.Fatal(err)
	}
	defer htmlfile.Close()

	pdfg.OutputFile = "./testfiles/TestPDFGeneratorOutputFile.pdf"

	pdfg.AddPage(NewPageReader(htmlfile))
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}

	pdfFile, err := os.Open("./testfiles/TestPDFGeneratorOutputFile.pdf")
	if err != nil {
		t.Fatal(err)
	}
	defer pdfFile.Close()

	stat, err := pdfFile.Stat()
	if err != nil {
		t.Fatal(err)
	}
	if stat.Size() < 100 {
		t.Errorf("generated PDF is size under 100 bytes")
	}
}

func TestGeneratePdfFromStdinHtml5(t *testing.T) {
	//Use newTestPDFGenerator and append to page1 and TOC
	pdfg := newTestPDFGenerator(t)
	htmlfile, err := ioutil.ReadFile("./testfiles/html5.html")
	if err != nil {
		t.Fatal(err)
	}
	page2 := NewPageReader(bytes.NewReader(htmlfile))
	pdfg.AddPage(page2)
	err = pdfg.Create()
	if err != nil {
		t.Fatal(err)
	}
	err = pdfg.WriteFile("./testfiles/TestGeneratePdfFromStdinHtml5.pdf")
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("PDF size %vkB", len(pdfg.Bytes())/1024)
}

func TestPath(t *testing.T) {
	path := "/usr/wkhtmltopdf/wkhtmltopdf"
	SetPath(path)
	defer SetPath("")
	if GetPath() != path {
		t.Errorf("Have path %q, want %q", GetPath(), path)
	}
}

func BenchmarkArgs(b *testing.B) {
	pdfg := newTestPDFGenerator(b)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pdfg.Args()
	}
}

func TestStringOption(t *testing.T) {
	opt := stringOption{
		option: "stringopt",
	}
	opt.Set("value123")

	want := []string{"--stringopt", "value123"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestSliceOption(t *testing.T) {
	opt := sliceOption{
		option: "sliceopt",
	}
	opt.Set("string15183")
	opt.Set("foo")
	opt.Set("bar")

	want := []string{"--sliceopt", "string15183", "--sliceopt", "foo", "--sliceopt", "bar"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestMapOption(t *testing.T) {
	opt := mapOption{
		option: "mapopt",
	}

	opt.Set("key1", "foo")
	opt.Set("key2", "bar")
	opt.Set("key3", "Hello")

	result := strings.Join(opt.Parse(), " ")
	if !strings.Contains(result, "--mapopt key1 foo") {
		t.Error("missing map option key1")
	}
	if !strings.Contains(result, "--mapopt key2 bar") {
		t.Error("missing map option key2")
	}
	if !strings.Contains(result, "--mapopt key3 Hello") {
		t.Error("missing map option key3")
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestUIntOption(t *testing.T) {
	opt := uintOption{
		option: "uintopt",
	}
	opt.Set(14860)

	want := []string{"--uintopt", "14860"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestFloatOption(t *testing.T) {
	opt := floatOption{
		option: "flopt",
	}
	opt.Set(239.75)

	want := []string{"--flopt", "239.750"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}

func TestBoolOption(t *testing.T) {
	opt := boolOption{
		option: "boolopt",
	}
	opt.Set(true)

	want := []string{"--boolopt"}

	if !reflect.DeepEqual(opt.Parse(), want) {
		t.Errorf("expected %v, have %v", want, opt.Parse())
	}

	opt.Unset()
	if !reflect.DeepEqual(opt.Parse(), []string{}) {
		t.Errorf("not empty after unset")
	}
}
<|endoftext|>


<commit_before>package wsmux

import (
	"bytes"
	"crypto/sha256"
	"io"
	"io/ioutil"
	"net/http"
	"testing"

	"github.com/gorilla/websocket"
)

// utils
var upgrader websocket.Upgrader = websocket.Upgrader{
	ReadBufferSize:  64 * 1024,
	WriteBufferSize: 64 * 1024,
}

func genTransferHandler(b *testing.B) http.Handler {
	tr := func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
			return
		}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			b.Fatal(err)
		}

		server := Server(conn, Config{Log: genLogger("transfer-bench-test"), StreamBufferSize: 64 * 1024})
		hash := sha256.New()

		str, err := server.Accept()
		if err != nil {
			b.Fatal(err)
		}
		// copy does not report EOF as error
		_, err = io.Copy(hash, str)
		if err != nil {
			b.Fatal(err)
		}
		_, err = str.Write(hash.Sum(nil))
		if err != nil {
			b.Fatal(err)
		}
		err = str.Close()
		if err != nil {
			b.Fatal(err)
		}
	}
	return http.HandlerFunc(tr)
}

// test large transfer
func BenchmarkTransfer(b *testing.B) {
	server := &http.Server{
		Addr:    ":9999",
		Handler: genTransferHandler(b),
	}
	defer func() {
		_ = server.Close()
	}()
	go func() {
		_ = server.ListenAndServe()
	}()
	conn, _, err := (&websocket.Dialer{}).Dial("ws://127.0.0.1:9999", nil)
	if err != nil {
		b.Fatal(err)
	}
	client := Client(conn, Config{Log: genLogger("transfer-client-bench-test"), StreamBufferSize: 64 * 1024})
	str, err := client.Open()
	if err != nil {
		b.Fatal(err)
	}

	hash := sha256.New()
	// write 4MB of data
	// how: generate a fixed array of size 2048
	// write it to hash and stream 2048 times

	// generate array
	buf := make([]byte, 2048)
	for i := 0; i < 2048; i++ {
		buf[i] = byte(i % 127)
	}

	b.ResetTimer()
	// write this array 4096 times
	for i := 0; i < 4096; i++ {
		_, _ = hash.Write(buf)
		_, err := str.Write(buf)
		if err != nil {
			b.Fatal(err)
		}
	}
	err = str.Close()

	// read hash generated by remote and compare
	buf, err = ioutil.ReadAll(str)
	if err != nil {
		b.Fatal(err)
	}

	if !bytes.Equal(hash.Sum(nil), buf) {
		b.Fatalf("bad message")
	}
}
<commit_msg>bench test slight refactor<commit_after>package wsmux

import (
	"bytes"
	"crypto/sha256"
	"io"
	"io/ioutil"
	"net/http"
	"testing"

	"github.com/gorilla/websocket"
)

// utils
var upgrader websocket.Upgrader = websocket.Upgrader{
	ReadBufferSize:  64 * 1024,
	WriteBufferSize: 64 * 1024,
}

func genTransferHandler(b *testing.B) http.Handler {
	tr := func(w http.ResponseWriter, r *http.Request) {
		if !websocket.IsWebSocketUpgrade(r) {
			http.NotFound(w, r)
			return
		}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			b.Fatal(err)
		}

		server := Server(conn, Config{StreamBufferSize: 64 * 1024})
		hash := sha256.New()

		str, err := server.Accept()
		if err != nil {
			b.Fatal(err)
		}
		// copy does not report EOF as error
		_, err = io.Copy(hash, str)
		if err != nil {
			b.Fatal(err)
		}
		_, err = str.Write(hash.Sum(nil))
		if err != nil {
			b.Fatal(err)
		}
		err = str.Close()
		if err != nil {
			b.Fatal(err)
		}
	}
	return http.HandlerFunc(tr)
}

// test large transfer
func BenchmarkTransfer(b *testing.B) {
	server := &http.Server{
		Addr:    ":9999",
		Handler: genTransferHandler(b),
	}
	defer func() {
		_ = server.Close()
	}()
	go func() {
		_ = server.ListenAndServe()
	}()
	conn, _, err := (&websocket.Dialer{}).Dial("ws://127.0.0.1:9999", nil)
	if err != nil {
		b.Fatal(err)
	}
	client := Client(conn, Config{StreamBufferSize: 64 * 1024})
	str, err := client.Open()
	if err != nil {
		b.Fatal(err)
	}

	hash := sha256.New()

	// write 64MB of data
	size := 64 * 1024 * 1024
	buf := make([]byte, size)
	for i := 0; i < size; i++ {
		buf[i] = byte(i % 127)
	}

	_, _ = hash.Write(buf)
	b.ResetTimer()
	_, err = str.Write(buf)
	if err != nil {
		b.Fatal(err)
	}
	err = str.Close()
	if err != nil {
		b.Fatal(err)
	}

	// read hash generated by remote and compare
	buf, err = ioutil.ReadAll(str)
	if err != nil {
		b.Fatal(err)
	}

	if !bytes.Equal(hash.Sum(nil), buf) {
		b.Fatalf("bad message")
	}
}
<|endoftext|>


<commit_before>/*

agorn is a wrapper around gorename for use with Acme.
It renames the entity under the cursor.

Usage:
	agorn [name]

Example:
	'agorn Foo' renames the entity under the cursor with 'Foo'.

gorename must be installed:
	% go get code.google.com/p/go.tools/cmd/gorename
*/
package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"code.google.com/p/goplan9/plan9/acme"
)

type window struct {
	win    *acme.Win
	offset int
	name   string
}

func currentAcmeWin() (*acme.Win, error) {
	winid := os.Getenv("winid")
	if winid == "" {
		return nil, fmt.Errorf("$winid not set - not running inside acme?")
	}
	id, err := strconv.Atoi(winid)
	if err != nil {
		return nil, fmt.Errorf("invalid $winid %q", winid)
	}
	win, err := acme.Open(id, nil)
	if err != nil {
		return nil, fmt.Errorf("cannot open acme window: %v", err)
	}
	return win, nil
}

func currentWindow() (*window, error) {
	win, err := currentAcmeWin()
	if err != nil {
		return nil, err
	}
	_, _, err = win.ReadAddr()
	if err != nil {
		return nil, fmt.Errorf("cannot read address: %v", err)
	}
	err = win.Ctl("addr=dot")
	if err != nil {
		return nil, fmt.Errorf("cannot set addr=dot: %v", err)
	}
	q0, _, err := win.ReadAddr()
	if err != nil {
		return nil, fmt.Errorf("cannot read address: %v", err)
	}
	b, err := win.ReadAll("tag")
	if err != nil {
		return nil, fmt.Errorf("cannot read tag: %v", err)
	}
	tag := string(b)
	i := strings.Index(tag, " ")
	if i == -1 {
		return nil, fmt.Errorf("tag with no spaces")
	}
	body, err := readBody(win)
	if err != nil {
		return nil, fmt.Errorf("cannot read body: %v", err)
	}
	return &window{win: win, name: tag[0:i], offset: toByteOffset(body, q0)}, nil
}

// We would use win.ReadAll except for a bug in acme
// where it crashes when reading trying to read more
// than the negotiated 9P message size.
func readBody(win *acme.Win) ([]byte, error) {
	var body []byte
	buf := make([]byte, 8000)
	for {
		n, err := win.Read("body", buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		body = append(body, buf[0:n]...)
	}
	return body, nil
}

func toByteOffset(b []byte, off int) int {
	r := 0
	for i := range string(b) {
		if r == off {
			return i
		}
		r++
	}
	return len(b)
}

func (w *window) showAddr(addr string) {
	w.win.Fprintf("addr", addr)
	w.win.Ctl("dot=addr")
	w.win.Ctl("show")
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "usage: agorn NAME\nReplaces the name under the cursor with NAME\n")
		return
	}
	to := os.Args[1]

	win, err := currentWindow()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		return
	}

	err = exec.Command("gorename", "-offset", fmt.Sprintf("%s:#%d", win.name, win.offset), "-to", to).Run()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		return
	}
	// fmt.Printf("gorename -offset %s:#%d -to %s", filename, offset, to)

	win.win.Ctl("get")
	win.showAddr(fmt.Sprintf("#%d", win.offset))
}
<commit_msg>Remove useless comment.<commit_after>/*

agorn is a wrapper around gorename for use with Acme.
It renames the entity under the cursor.

Usage:
	agorn [name]

Example:
	'agorn Foo' renames the entity under the cursor with 'Foo'.

gorename must be installed:
	% go get code.google.com/p/go.tools/cmd/gorename
*/
package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"code.google.com/p/goplan9/plan9/acme"
)

type window struct {
	win    *acme.Win
	offset int
	name   string
}

func currentAcmeWin() (*acme.Win, error) {
	winid := os.Getenv("winid")
	if winid == "" {
		return nil, fmt.Errorf("$winid not set - not running inside acme?")
	}
	id, err := strconv.Atoi(winid)
	if err != nil {
		return nil, fmt.Errorf("invalid $winid %q", winid)
	}
	win, err := acme.Open(id, nil)
	if err != nil {
		return nil, fmt.Errorf("cannot open acme window: %v", err)
	}
	return win, nil
}

func currentWindow() (*window, error) {
	win, err := currentAcmeWin()
	if err != nil {
		return nil, err
	}
	_, _, err = win.ReadAddr()
	if err != nil {
		return nil, fmt.Errorf("cannot read address: %v", err)
	}
	err = win.Ctl("addr=dot")
	if err != nil {
		return nil, fmt.Errorf("cannot set addr=dot: %v", err)
	}
	q0, _, err := win.ReadAddr()
	if err != nil {
		return nil, fmt.Errorf("cannot read address: %v", err)
	}
	b, err := win.ReadAll("tag")
	if err != nil {
		return nil, fmt.Errorf("cannot read tag: %v", err)
	}
	tag := string(b)
	i := strings.Index(tag, " ")
	if i == -1 {
		return nil, fmt.Errorf("tag with no spaces")
	}
	body, err := readBody(win)
	if err != nil {
		return nil, fmt.Errorf("cannot read body: %v", err)
	}
	return &window{win: win, name: tag[0:i], offset: toByteOffset(body, q0)}, nil
}

// We would use win.ReadAll except for a bug in acme
// where it crashes when reading trying to read more
// than the negotiated 9P message size.
func readBody(win *acme.Win) ([]byte, error) {
	var body []byte
	buf := make([]byte, 8000)
	for {
		n, err := win.Read("body", buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, err
		}
		body = append(body, buf[0:n]...)
	}
	return body, nil
}

func toByteOffset(b []byte, off int) int {
	r := 0
	for i := range string(b) {
		if r == off {
			return i
		}
		r++
	}
	return len(b)
}

func (w *window) showAddr(addr string) {
	w.win.Fprintf("addr", addr)
	w.win.Ctl("dot=addr")
	w.win.Ctl("show")
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "usage: agorn NAME\nReplaces the name under the cursor with NAME\n")
		return
	}
	to := os.Args[1]

	win, err := currentWindow()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		return
	}

	err = exec.Command("gorename", "-offset", fmt.Sprintf("%s:#%d", win.name, win.offset), "-to", to).Run()
	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		return
	}

	win.win.Ctl("get")
	win.showAddr(fmt.Sprintf("#%d", win.offset))
}
<|endoftext|>


<commit_before>package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync/atomic"
	"time"
)

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) (pb *ProgressBar) {
	pb = &ProgressBar{
		Total:        int64(total),
		RefreshRate:  DEFAULT_REFRESH_RATE,
		ShowPercent:  true,
		ShowCounters: true,
		ShowBar:      true,
		ShowTimeLeft: true,
	}
	pb.Format(FORMAT)
	return
}

// Create new object and start
func StartNew(total int) (pb *ProgressBar) {
	pb = New(total)
	pb.Start()
	return
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            int

	isFinish  bool
	startTime time.Time

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string
}

// Start print
func (pb *ProgressBar) Start() {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	go pb.writer()
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(atomic.AddInt64(&pb.current, int64(add)))
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
func (pb *ProgressBar) Format(format string) (bar *ProgressBar) {
	bar = pb
	formatEntries := strings.Split(format, "")
	if len(formatEntries) != 5 {
		return
	}
	pb.BarStart = formatEntries[0]
	pb.BarEnd = formatEntries[4]
	pb.Empty = formatEntries[3]
	pb.Current = formatEntries[1]
	pb.CurrentN = formatEntries[2]
	return
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) (bar *ProgressBar) {
	bar = pb
	pb.RefreshRate = rate
	return
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units int) (bar *ProgressBar) {
	bar = pb
	switch units {
	case U_NO, U_BYTES:
		pb.Units = units
	}
	return
}

// End print
func (pb *ProgressBar) Finish() {
	pb.isFinish = true
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	fmt.Println(str)
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// Create new proxy reader over bar
func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
	return &Reader{r, pb}
}

func (pb *ProgressBar) write(current int64) {
	width, _ := terminalWidth()
	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	if pb.ShowTimeLeft && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		perEntry := fromStart / time.Duration(current)
		left := time.Duration(pb.Total-current) * perEntry
		left = (left / time.Second) * time.Second
		if left > 0 {
			timeLeftBox = left.String()
		}
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+pb.BarStart+pb.BarEnd+percentBox+timeLeftBox+speedBox)
		if size > 0 {
			curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
			emptCount := size - curCount
			barBox = pb.BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
			barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
		}
	}

	// check len
	out = countersBox + barBox + percentBox + speedBox + timeLeftBox
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	out = countersBox + barBox + percentBox + speedBox + timeLeftBox

	// and print!
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func (pb *ProgressBar) writer() {
	var c, oc int64
	oc = -1
	for {
		if pb.isFinish {
			break
		}
		c = atomic.LoadInt64(&pb.current)
		if c != oc {
			pb.write(c)
			oc = c
		}
		time.Sleep(pb.RefreshRate)
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<commit_msg>going back<commit_after>package pb

import (
	"fmt"
	"io"
	"math"
	"strings"
	"sync/atomic"
	"time"
)

const (
	// Default refresh rate - 200ms
	DEFAULT_REFRESH_RATE = time.Millisecond * 200
	FORMAT               = "[=>-]"
)

// DEPRECATED
// variables for backward compatibility, from now do not work
// use pb.Format and pb.SetRefreshRate
var (
	DefaultRefreshRate                         = DEFAULT_REFRESH_RATE
	BarStart, BarEnd, Empty, Current, CurrentN string
)

// Create new progress bar object
func New(total int) (pb *ProgressBar) {
	pb = &ProgressBar{
		Total:        int64(total),
		RefreshRate:  DEFAULT_REFRESH_RATE,
		ShowPercent:  true,
		ShowCounters: true,
		ShowBar:      true,
		ShowTimeLeft: true,
	}
	pb.Format(FORMAT)
	return
}

// Create new object and start
func StartNew(total int) (pb *ProgressBar) {
	pb = New(total)
	pb.Start()
	return
}

// Callback for custom output
// For example:
// bar.Callback = func(s string) {
//     mySuperPrint(s)
// }
//
type Callback func(out string)

type ProgressBar struct {
	current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)

	Total                            int64
	RefreshRate                      time.Duration
	ShowPercent, ShowCounters        bool
	ShowSpeed, ShowTimeLeft, ShowBar bool
	Output                           io.Writer
	Callback                         Callback
	NotPrint                         bool
	Units                            int

	isFinish  bool
	startTime time.Time

	BarStart string
	BarEnd   string
	Empty    string
	Current  string
	CurrentN string
}

// Start print
func (pb *ProgressBar) Start() {
	pb.startTime = time.Now()
	if pb.Total == 0 {
		pb.ShowBar = false
		pb.ShowTimeLeft = false
		pb.ShowPercent = false
	}
	go pb.writer()
}

// Increment current value
func (pb *ProgressBar) Increment() int {
	return pb.Add(1)
}

// Set current value
func (pb *ProgressBar) Set(current int) {
	atomic.StoreInt64(&pb.current, int64(current))
}

// Add to current value
func (pb *ProgressBar) Add(add int) int {
	return int(atomic.AddInt64(&pb.current, int64(add)))
}

// Set custom format for bar
// Example: bar.Format("[=>_]")
func (pb *ProgressBar) Format(format string) (bar *ProgressBar) {
	bar = pb
	formatEntries := strings.Split(format, "")
	if len(formatEntries) != 5 {
		return
	}
	pb.BarStart = formatEntries[0]
	pb.BarEnd = formatEntries[4]
	pb.Empty = formatEntries[3]
	pb.Current = formatEntries[1]
	pb.CurrentN = formatEntries[2]
	return
}

// Set bar refresh rate
func (pb *ProgressBar) SetRefreshRate(rate time.Duration) (bar *ProgressBar) {
	bar = pb
	pb.RefreshRate = rate
	return
}

// Set units
// bar.SetUnits(U_NO) - by default
// bar.SetUnits(U_BYTES) - for Mb, Kb, etc
func (pb *ProgressBar) SetUnits(units int) (bar *ProgressBar) {
	bar = pb
	switch units {
	case U_NO, U_BYTES:
		pb.Units = units
	}
	return
}

// End print
func (pb *ProgressBar) Finish() {
	pb.isFinish = true
	pb.write(atomic.LoadInt64(&pb.current))
	if !pb.NotPrint {
		fmt.Println()
	}
}

// End print and write string 'str'
func (pb *ProgressBar) FinishPrint(str string) {
	pb.Finish()
	fmt.Println(str)
}

// implement io.Writer
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

// implement io.Reader
func (pb *ProgressBar) Read(p []byte) (n int, err error) {
	n = len(p)
	pb.Add(n)
	return
}

func (pb *ProgressBar) write(current int64) {
	width, _ := terminalWidth()
	var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string

	// percents
	if pb.ShowPercent {
		percent := float64(current) / (float64(pb.Total) / float64(100))
		percentBox = fmt.Sprintf(" %#.02f %% ", percent)
	}

	// counters
	if pb.ShowCounters {
		if pb.Total > 0 {
			countersBox = fmt.Sprintf("%s / %s ", Format(current, pb.Units), Format(pb.Total, pb.Units))
		} else {
			countersBox = Format(current, pb.Units) + " "
		}
	}

	// time left
	if pb.ShowTimeLeft && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		perEntry := fromStart / time.Duration(current)
		left := time.Duration(pb.Total-current) * perEntry
		left = (left / time.Second) * time.Second
		if left > 0 {
			timeLeftBox = left.String()
		}
	}

	// speed
	if pb.ShowSpeed && current > 0 {
		fromStart := time.Now().Sub(pb.startTime)
		speed := float64(current) / (float64(fromStart) / float64(time.Second))
		speedBox = Format(int64(speed), pb.Units) + "/s "
	}

	// bar
	if pb.ShowBar {
		size := width - len(countersBox+pb.BarStart+pb.BarEnd+percentBox+timeLeftBox+speedBox)
		if size > 0 {
			curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
			emptCount := size - curCount
			barBox = pb.BarStart
			if emptCount < 0 {
				emptCount = 0
			}
			if curCount > size {
				curCount = size
			}
			barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
			barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
		}
	}

	// check len
	out = countersBox + barBox + percentBox + speedBox + timeLeftBox
	if len(out) < width {
		end = strings.Repeat(" ", width-len(out))
	}

	out = countersBox + barBox + percentBox + speedBox + timeLeftBox

	// and print!
	switch {
	case pb.Output != nil:
		fmt.Fprint(pb.Output, out+end)
	case pb.Callback != nil:
		pb.Callback(out + end)
	case !pb.NotPrint:
		fmt.Print("\r" + out + end)
	}
}

func (pb *ProgressBar) writer() {
	var c, oc int64
	oc = -1
	for {
		if pb.isFinish {
			break
		}
		c = atomic.LoadInt64(&pb.current)
		if c != oc {
			pb.write(c)
			oc = c
		}
		time.Sleep(pb.RefreshRate)
	}
}

type window struct {
	Row    uint16
	Col    uint16
	Xpixel uint16
	Ypixel uint16
}
<|endoftext|>


<commit_before>package qq

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"
)

var (
	LogFile = "qq.log"
	mu      sync.Mutex
)

func Println(a ...interface{}) {
	f := filepath.Join(os.TempDir(), LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		// TODO: don't panic. people will forget and leave qq.Print() calls in
		// their code, which will end up in prod. we don't want to crash the
		// server because we don't have permissions to write to /tmp.
		panic(err)
	}
	defer fd.Close()

	pc, file, line, ok := runtime.Caller(1)
	if ok {
		names, err := argNames(file, line)
		if err == nil {
			a = formatArgs(names, a)
		}

		p := []interface{}{prefix(pc, file, line)}
		a = append(p, a...)
	}

	a = append(a, "\n")
	mu.Lock()
	_, err = fmt.Fprintln(fd, a...)
	mu.Unlock()

	if err != nil {
		panic(err) // TODO: don't panic
	}
}

func Printf(format string, a ...interface{}) {
	f := filepath.Join(os.TempDir(), LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err) // TODO: don't panic
	}
	defer fd.Close()

	pc, file, line, ok := runtime.Caller(1)
	if !ok {
		mu.Lock()
		_, err = fmt.Fprintf(fd, format, a...)
		mu.Unlock()
		return
	}

	p := prefix(pc, file, line)
	mu.Lock()
	_, err = fmt.Fprintf(fd, p+" "+format, a...)
	mu.Unlock()

	if err != nil {
		panic(err) // TODO: don't panic
	}
}

func prefix(pc uintptr, file string, line int) string {
	t := time.Now().Format("15:04:05")
	shortFile := filepath.Base(file)
	callerName := runtime.FuncForPC(pc).Name()

	return fmt.Sprintf("[%s %s:%d %s]", t, shortFile, line, callerName)
}

// formatArgs turns a slice of arguments into pretty-printed strings. If the
// argument variable name is present in names, it will be returned as a
// name=value string, e.g. "port=443".
func formatArgs(names []string, values []interface{}) []interface{} {
	for i := 0; i < len(values); i++ {
		if names[i] == "" {
			values[i] = fmt.Sprintf("%#v", values[i])
		} else {
			values[i] = fmt.Sprintf("%s=%#v", names[i], values[i])
		}
	}
	return values
}

// argNames returns the names of all the variable arguments for the qq.Print*()
// call at the given file and line number. If the argument is not a variable,
// the slice will contain an empty string at the index position for that
// argument. For example, qq.Print(a, 123) will result in []string{"a", ""}
// for arg names, because 123 is not a variable name.
func argNames(file string, line int) ([]string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, file, nil, 0)
	if err != nil {
		return nil, err
	}

	var names []string
	ast.Inspect(f, func(n ast.Node) bool {
		call, is := n.(*ast.CallExpr)
		if !is {
			return true
		}

		if fset.Position(call.End()).Line != line {
			return true
		}

		if !qqCall(call) {
			return true
		}

		for _, arg := range call.Args {
			names = append(names, argName(arg))
		}
		return true
	})

	return names, nil
}

// qqCall returns true if the given function call expression is for a function
// in the qq package, e.g. qq.Printf().
func qqCall(n *ast.CallExpr) bool {
	sel, is := n.Fun.(*ast.SelectorExpr)
	if !is {
		return false
	}

	ident, is := sel.X.(*ast.Ident)
	if !is {
		return false
	}

	return ident.Name == "qq"
}

// argName returns the name of the given argument if it's a variable. If the
// argument is something else, like a literal or a function call, argName
// returns an empty string.
func argName(arg ast.Expr) string {
	ident, is := arg.(*ast.Ident)
	if !is {
		return ""
	}

	if ident.Obj.Kind != ast.Var {
		return ""
	}

	return ident.Obj.Name
}
<commit_msg>Remove don't panic comments<commit_after>package qq

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"time"
)

var (
	LogFile = "qq.log"
	mu      sync.Mutex
)

func Println(a ...interface{}) {
	f := filepath.Join(os.TempDir(), LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	pc, file, line, ok := runtime.Caller(1)
	if ok {
		names, err := argNames(file, line)
		if err == nil {
			a = formatArgs(names, a)
		}

		p := []interface{}{prefix(pc, file, line)}
		a = append(p, a...)
	}

	a = append(a, "\n")
	mu.Lock()
	_, err = fmt.Fprintln(fd, a...)
	mu.Unlock()

	if err != nil {
		panic(err)
	}
}

func Printf(format string, a ...interface{}) {
	f := filepath.Join(os.TempDir(), LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	pc, file, line, ok := runtime.Caller(1)
	if !ok {
		mu.Lock()
		_, err = fmt.Fprintf(fd, format, a...)
		mu.Unlock()
		return
	}

	p := prefix(pc, file, line)
	mu.Lock()
	_, err = fmt.Fprintf(fd, p+" "+format, a...)
	mu.Unlock()

	if err != nil {
		panic(err)
	}
}

func prefix(pc uintptr, file string, line int) string {
	t := time.Now().Format("15:04:05")
	shortFile := filepath.Base(file)
	callerName := runtime.FuncForPC(pc).Name()

	return fmt.Sprintf("[%s %s:%d %s]", t, shortFile, line, callerName)
}

// formatArgs turns a slice of arguments into pretty-printed strings. If the
// argument variable name is present in names, it will be returned as a
// name=value string, e.g. "port=443".
func formatArgs(names []string, values []interface{}) []interface{} {
	for i := 0; i < len(values); i++ {
		if names[i] == "" {
			values[i] = fmt.Sprintf("%#v", values[i])
		} else {
			values[i] = fmt.Sprintf("%s=%#v", names[i], values[i])
		}
	}
	return values
}

// argNames returns the names of all the variable arguments for the qq.Print*()
// call at the given file and line number. If the argument is not a variable,
// the slice will contain an empty string at the index position for that
// argument. For example, qq.Print(a, 123) will result in []string{"a", ""}
// for arg names, because 123 is not a variable name.
func argNames(file string, line int) ([]string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, file, nil, 0)
	if err != nil {
		return nil, err
	}

	var names []string
	ast.Inspect(f, func(n ast.Node) bool {
		call, is := n.(*ast.CallExpr)
		if !is {
			return true
		}

		if fset.Position(call.End()).Line != line {
			return true
		}

		if !qqCall(call) {
			return true
		}

		for _, arg := range call.Args {
			names = append(names, argName(arg))
		}
		return true
	})

	return names, nil
}

// qqCall returns true if the given function call expression is for a function
// in the qq package, e.g. qq.Printf().
func qqCall(n *ast.CallExpr) bool {
	sel, is := n.Fun.(*ast.SelectorExpr)
	if !is {
		return false
	}

	ident, is := sel.X.(*ast.Ident)
	if !is {
		return false
	}

	return ident.Name == "qq"
}

// argName returns the name of the given argument if it's a variable. If the
// argument is something else, like a literal or a function call, argName
// returns an empty string.
func argName(arg ast.Expr) string {
	ident, is := arg.(*ast.Ident)
	if !is {
		return ""
	}

	if ident.Obj.Kind != ast.Var {
		return ""
	}

	return ident.Obj.Name
}
<|endoftext|>


<commit_before>package dnsr

import (
	"strings"

	"github.com/miekg/dns"
)

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// RRs represents a slice of DNS resource records.
type RRs []RR

// emptyRRs is an empty, non-nil slice of RRs.
// It is used to save allocations at runtime.
var emptyRRs = RRs{}

// ICANN specifies that DNS servers should return the special value 127.0.53.53
// for A record queries of TLDs that have recently entered the root zone,
// that have a high likelyhood of colliding with private DNS names.
// The record returned is a notice to network administrators to adjust their
// DNS configuration.
// https://www.icann.org/resources/pages/name-collision-2013-12-06-en#127.0.53.53
const NameCollision = "127.0.53.53"

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

// convertRR converts a dns.RR to an RR.
// If the RR is not a type that this package uses,
// It will attempt to translate this if there are enough parameters
// Should all maners of translation fail, it returns an undefined RR and false.
func convertRR(drr dns.RR) (RR, bool) {
	switch t := drr.(type) {
	case *dns.SOA:
		return RR{toLowerFQDN(t.Hdr.Name), "SOA", toLowerFQDN(t.Ns)}, true
	case *dns.NS:
		return RR{toLowerFQDN(t.Hdr.Name), "NS", toLowerFQDN(t.Ns)}, true
	case *dns.CNAME:
		return RR{toLowerFQDN(t.Hdr.Name), "CNAME", toLowerFQDN(t.Target)}, true
	case *dns.A:
		return RR{toLowerFQDN(t.Hdr.Name), "A", t.A.String()}, true
	case *dns.AAAA:
		return RR{toLowerFQDN(t.Hdr.Name), "AAAA", t.AAAA.String()}, true
	case *dns.TXT:
		return RR{toLowerFQDN(t.Hdr.Name), "TXT", strings.Join(t.Txt, "\t")}, true
	default:
		fields := strings.Fields(drr.String())
		if len(fields) >= 4 {
			return RR{toLowerFQDN(fields[0]), fields[3], strings.Join(fields[4:len(fields)], "\t")}, true
		}
	}
	return RR{}, false
}
<commit_msg>adjsut comments and remove len<commit_after>package dnsr

import (
	"strings"

	"github.com/miekg/dns"
)

// RR represents a DNS resource record.
type RR struct {
	Name  string
	Type  string
	Value string
}

// RRs represents a slice of DNS resource records.
type RRs []RR

// emptyRRs is an empty, non-nil slice of RRs.
// It is used to save allocations at runtime.
var emptyRRs = RRs{}

// ICANN specifies that DNS servers should return the special value 127.0.53.53
// for A record queries of TLDs that have recently entered the root zone,
// that have a high likelyhood of colliding with private DNS names.
// The record returned is a notice to network administrators to adjust their
// DNS configuration.
// https://www.icann.org/resources/pages/name-collision-2013-12-06-en#127.0.53.53
const NameCollision = "127.0.53.53"

// String returns a string representation of an RR in zone-file format.
func (rr *RR) String() string {
	return rr.Name + "\t      3600\tIN\t" + rr.Type + "\t" + rr.Value
}

// convertRR converts a dns.RR to an RR.
// If the RR is not a type that this package uses,
// It will attempt to translate this if there are enough parameters
// Should all translation fail, it returns an undefined RR and false.
func convertRR(drr dns.RR) (RR, bool) {
	switch t := drr.(type) {
	case *dns.SOA:
		return RR{toLowerFQDN(t.Hdr.Name), "SOA", toLowerFQDN(t.Ns)}, true
	case *dns.NS:
		return RR{toLowerFQDN(t.Hdr.Name), "NS", toLowerFQDN(t.Ns)}, true
	case *dns.CNAME:
		return RR{toLowerFQDN(t.Hdr.Name), "CNAME", toLowerFQDN(t.Target)}, true
	case *dns.A:
		return RR{toLowerFQDN(t.Hdr.Name), "A", t.A.String()}, true
	case *dns.AAAA:
		return RR{toLowerFQDN(t.Hdr.Name), "AAAA", t.AAAA.String()}, true
	case *dns.TXT:
		return RR{toLowerFQDN(t.Hdr.Name), "TXT", strings.Join(t.Txt, "\t")}, true
	default:
		fields := strings.Fields(drr.String())
		if len(fields) >= 4 {
			return RR{toLowerFQDN(fields[0]), fields[3], strings.Join(fields[4:], "\t")}, true
		}
	}
	return RR{}, false
}
<|endoftext|>