<commit_before>// Based on: https://github.com/btcsuite/btcwallet/blob/master/walletsetup.go
/*
* Copyright (c) 2014-2015 The btcsuite developers
*
* Permission to use, copy, modify, and distribute this software for any
* purpose with or without fee is hereby granted, provided that the above
* copyright notice and this permission notice appear in all copies.
*
* THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
* WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
* MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
* ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
* WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
* ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
* OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package wallet

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
	"github.com/btcsuite/btcutil/hdkeychain"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/wallet"
	"github.com/btcsuite/btcwallet/walletdb"
)

var (
	// TODO(roasbeef): lnwallet config file
	lnwalletHomeDir    = btcutil.AppDataDir("lnwallet", false)
	defaultDataDir     = lnwalletHomeDir
	defaultLogFilename = "lnwallet.log"

	defaultLogDirname = "logs"

	// defaultPubPassphrase is the default public wallet passphrase which is
	// used when the user indicates they do not want additional protection
	// provided by having all public data in the wallet encrypted by a
	// passphrase only known to them.
	defaultPubPassphrase = []byte("public")

	defaultLogDir = filepath.Join(lnwalletHomeDir, defaultLogDirname)

	walletDbName = "lnwallet.db"
)

// filesExists reports whether the named file or directory exists.
func fileExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// networkDir returns the directory name of a network directory to hold wallet
// files.
func networkDir(dataDir string, chainParams *chaincfg.Params) string {
	netname := chainParams.Name

	// For now, we must always name the testnet data directory as "testnet"
	// and not "testnet3" or any other version, as the chaincfg testnet3
	// paramaters will likely be switched to being named "testnet3" in the
	// future.  This is done to future proof that change, and an upgrade
	// plan to move the testnet3 data directory can be worked out later.
	if chainParams.Net == wire.TestNet3 {
		netname = "testnet"
	}

	return filepath.Join(dataDir, netname)
}

// checkCreateDir checks that the path exists and is a directory.
// If path does not exist, it is created.
func checkCreateDir(path string) error {
	if fi, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(path, 0700); err != nil {
				return fmt.Errorf("cannot create directory: %s", err)
			}
		} else {
			return fmt.Errorf("error checking directory: %s", err)
		}
	} else {
		if !fi.IsDir() {
			return fmt.Errorf("path '%s' is not a directory", path)
		}
	}

	return nil
}

// createWallet generates a new wallet. The new wallet will reside at the
// provided path.
// TODO(roasbeef): maybe pass in config after all for testing purposes?
func createWallet(privPass, pubPass, userSeed []byte,
	dbPath string) error {
	// TODO(roasbeef): replace with tadge's seed format?
	hdSeed := userSeed
	var seedErr error
	if userSeed == nil {
		hdSeed, seedErr = hdkeychain.GenerateSeed(hdkeychain.RecommendedSeedLen)
		if seedErr != nil {
			return seedErr
		}
	}

	// Create the wallet.
	fmt.Println("Creating the wallet...")

	// Create the wallet database backed by bolt db.
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return err
	}

	// Create the address manager.
	namespace, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return err
	}
	manager, err := waddrmgr.Create(namespace, hdSeed, []byte(pubPass),
		[]byte(privPass), ActiveNetParams, nil)
	if err != nil {
		return err
	}

	manager.Close()
	fmt.Println("The lnwallet has been created successfully.")
	return nil
}

// openDb opens and returns a walletdb.DB (boltdb here) given the directory and
// dbname
func openDb(directory string, dbname string) (walletdb.DB, error) {
	dbPath := filepath.Join(directory, dbname)

	// Ensure that the network directory exists.
	if err := checkCreateDir(directory); err != nil {
		return nil, err
	}

	// Open the database using the boltdb backend.
	return walletdb.Open("bdb", dbPath)
}

// promptSeed is used to prompt for the wallet seed which maybe required during
// upgrades.
func promptSeed() ([]byte, error) {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("Enter existing wallet seed: ")
		seedStr, err := reader.ReadString('\n')
		if err != nil {
			return nil, err
		}
		seedStr = strings.TrimSpace(strings.ToLower(seedStr))

		seed, err := hex.DecodeString(seedStr)
		if err != nil || len(seed) < hdkeychain.MinSeedBytes ||
			len(seed) > hdkeychain.MaxSeedBytes {

			fmt.Printf("Invalid seed specified.  Must be a "+
				"hexadecimal value that is at least %d bits and "+
				"at most %d bits\n", hdkeychain.MinSeedBytes*8,
				hdkeychain.MaxSeedBytes*8)
			continue
		}

		return seed, nil
	}
}

// promptPrivPassPhrase is used to prompt for the private passphrase which maybe
// required during upgrades.
func promptPrivPassPhrase() ([]byte, error) {
	prompt := "Enter the private passphrase of your wallet: "
	for {
		fmt.Print(prompt)
		pass, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			return nil, err
		}
		fmt.Print("\n")
		pass = bytes.TrimSpace(pass)
		if len(pass) == 0 {
			continue
		}

		return pass, nil
	}
}

// openWallet returns a wallet. The function handles opening an existing wallet
// database, the address manager and the transaction store and uses the values
// to open a wallet.Wallet
func openWallet(pubPass []byte, dbDir string) (*wallet.Wallet, walletdb.DB, error) {
	db, err := openDb(dbDir, walletDbName)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to open database: %v", err)
	}

	addrMgrNS, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, nil, err
	}
	txMgrNS, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, nil, err
	}

	// TODO(roasbeef): pass these in as funcs instead, priv pass already
	// loaded into memory, use tadge's format to read HD seed.
	cbs := &waddrmgr.OpenCallbacks{
		ObtainSeed:        promptSeed,
		ObtainPrivatePass: promptPrivPassPhrase,
	}
	w, err := wallet.Open(pubPass, ActiveNetParams, db, addrMgrNS, txMgrNS,
		cbs)
	return w, db, err
}
<commit_msg>import walletdb/bdb for side effects (registers driver)<commit_after>// Based on: https://github.com/btcsuite/btcwallet/blob/master/walletsetup.go
/*
* Copyright (c) 2014-2015 The btcsuite developers
*
* Permission to use, copy, modify, and distribute this software for any
* purpose with or without fee is hereby granted, provided that the above
* copyright notice and this permission notice appear in all copies.
*
* THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
* WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
* MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
* ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
* WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
* ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
* OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package wallet

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/wire"
	"github.com/btcsuite/btcutil"
	"github.com/btcsuite/btcutil/hdkeychain"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/wallet"
	"github.com/btcsuite/btcwallet/walletdb"
	_ "github.com/btcsuite/btcwallet/walletdb/bdb"
)

var (
	// TODO(roasbeef): lnwallet config file
	lnwalletHomeDir    = btcutil.AppDataDir("lnwallet", false)
	defaultDataDir     = lnwalletHomeDir
	defaultLogFilename = "lnwallet.log"

	defaultLogDirname = "logs"

	// defaultPubPassphrase is the default public wallet passphrase which is
	// used when the user indicates they do not want additional protection
	// provided by having all public data in the wallet encrypted by a
	// passphrase only known to them.
	defaultPubPassphrase = []byte("public")

	defaultLogDir = filepath.Join(lnwalletHomeDir, defaultLogDirname)

	walletDbName = "lnwallet.db"
)

// filesExists reports whether the named file or directory exists.
func fileExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// networkDir returns the directory name of a network directory to hold wallet
// files.
func networkDir(dataDir string, chainParams *chaincfg.Params) string {
	netname := chainParams.Name

	// For now, we must always name the testnet data directory as "testnet"
	// and not "testnet3" or any other version, as the chaincfg testnet3
	// paramaters will likely be switched to being named "testnet3" in the
	// future.  This is done to future proof that change, and an upgrade
	// plan to move the testnet3 data directory can be worked out later.
	if chainParams.Net == wire.TestNet3 {
		netname = "testnet"
	}

	return filepath.Join(dataDir, netname)
}

// checkCreateDir checks that the path exists and is a directory.
// If path does not exist, it is created.
func checkCreateDir(path string) error {
	if fi, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			// Attempt data directory creation
			if err = os.MkdirAll(path, 0700); err != nil {
				return fmt.Errorf("cannot create directory: %s", err)
			}
		} else {
			return fmt.Errorf("error checking directory: %s", err)
		}
	} else {
		if !fi.IsDir() {
			return fmt.Errorf("path '%s' is not a directory", path)
		}
	}

	return nil
}

// createWallet generates a new wallet. The new wallet will reside at the
// provided path.
// TODO(roasbeef): maybe pass in config after all for testing purposes?
func createWallet(privPass, pubPass, userSeed []byte,
	dbPath string) error {
	// TODO(roasbeef): replace with tadge's seed format?
	hdSeed := userSeed
	var seedErr error
	if userSeed == nil {
		hdSeed, seedErr = hdkeychain.GenerateSeed(hdkeychain.RecommendedSeedLen)
		if seedErr != nil {
			return seedErr
		}
	}

	// Create the wallet.
	fmt.Println("Creating the wallet...")

	// Create the wallet database backed by bolt db.
	db, err := walletdb.Create("bdb", dbPath)
	if err != nil {
		return err
	}

	// Create the address manager.
	namespace, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return err
	}
	manager, err := waddrmgr.Create(namespace, hdSeed, []byte(pubPass),
		[]byte(privPass), ActiveNetParams, nil)
	if err != nil {
		return err
	}

	manager.Close()
	fmt.Println("The lnwallet has been created successfully.")
	return nil
}

// openDb opens and returns a walletdb.DB (boltdb here) given the directory and
// dbname
func openDb(directory string, dbname string) (walletdb.DB, error) {
	dbPath := filepath.Join(directory, dbname)

	// Ensure that the network directory exists.
	if err := checkCreateDir(directory); err != nil {
		return nil, err
	}

	// Open the database using the boltdb backend.
	return walletdb.Open("bdb", dbPath)
}

// promptSeed is used to prompt for the wallet seed which maybe required during
// upgrades.
func promptSeed() ([]byte, error) {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("Enter existing wallet seed: ")
		seedStr, err := reader.ReadString('\n')
		if err != nil {
			return nil, err
		}
		seedStr = strings.TrimSpace(strings.ToLower(seedStr))

		seed, err := hex.DecodeString(seedStr)
		if err != nil || len(seed) < hdkeychain.MinSeedBytes ||
			len(seed) > hdkeychain.MaxSeedBytes {

			fmt.Printf("Invalid seed specified.  Must be a "+
				"hexadecimal value that is at least %d bits and "+
				"at most %d bits\n", hdkeychain.MinSeedBytes*8,
				hdkeychain.MaxSeedBytes*8)
			continue
		}

		return seed, nil
	}
}

// promptPrivPassPhrase is used to prompt for the private passphrase which maybe
// required during upgrades.
func promptPrivPassPhrase() ([]byte, error) {
	prompt := "Enter the private passphrase of your wallet: "
	for {
		fmt.Print(prompt)
		pass, err := terminal.ReadPassword(int(os.Stdin.Fd()))
		if err != nil {
			return nil, err
		}
		fmt.Print("\n")
		pass = bytes.TrimSpace(pass)
		if len(pass) == 0 {
			continue
		}

		return pass, nil
	}
}

// openWallet returns a wallet. The function handles opening an existing wallet
// database, the address manager and the transaction store and uses the values
// to open a wallet.Wallet
func openWallet(pubPass []byte, dbDir string) (*wallet.Wallet, walletdb.DB, error) {
	db, err := openDb(dbDir, walletDbName)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to open database: %v", err)
	}

	addrMgrNS, err := db.Namespace(waddrmgrNamespaceKey)
	if err != nil {
		return nil, nil, err
	}
	txMgrNS, err := db.Namespace(wtxmgrNamespaceKey)
	if err != nil {
		return nil, nil, err
	}

	// TODO(roasbeef): pass these in as funcs instead, priv pass already
	// loaded into memory, use tadge's format to read HD seed.
	cbs := &waddrmgr.OpenCallbacks{
		ObtainSeed:        promptSeed,
		ObtainPrivatePass: promptPrivPassPhrase,
	}
	w, err := wallet.Open(pubPass, ActiveNetParams, db, addrMgrNS, txMgrNS,
		cbs)
	return w, db, err
}
<|endoftext|>


<commit_before>// The unexport command unexports exported identifiers which are not imported
// by any other Go code.
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"unicode"
	"unicode/utf8"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
	"golang.org/x/tools/refactor/importgraph"
)

func main() {
	var (
		flagPackage    = flag.String("package", "", "package import path to be unexported")
		flagIdentifier = flag.String("identifier", "", "comma-separated list of identifiers names; if empty all identifiers are unexported")
		flagDryRun     = flag.Bool("dryrun", false, "show the change, but do not apply")
		flagVerbose    = flag.Bool("verbose", false, "show more information. Useful for debugging.")
	)

	flag.Var((*buildutil.TagsFlag)(&build.Default.BuildTags), "tags", buildutil.TagsFlagDoc)

	flag.Parse()
	log.SetPrefix("unexport:")

	identifiers := []string{}
	if *flagIdentifier != "" {
		identifiers = strings.Split(*flagIdentifier, ",")
	}

	if err := runMain(&config{
		importPath:   *flagPackage,
		identifiers:  identifiers,
		buildContext: &build.Default,
		dryRun:       *flagDryRun,
		verbose:      *flagVerbose,
	}); err != nil {
		log.Println(err)
		os.Exit(1)
	}
}

// config is used to define how unexport should be work
type config struct {
	// importPath defines the package defined with the importpath
	importPath string

	// identifiers is used to limit the changes of unexporting to certain identifiers
	identifiers []string

	// build context
	buildContext *build.Context

	// logging/development ...
	dryRun  bool
	verbose bool
}

// runMain runs the actual command. It's an helper function so we can easily
// calls defers or return errors.
func runMain(conf *config) error {
	if conf.importPath == "" {
		return errors.New("import path of the package must be given")
	}

	path := conf.importPath

	prog, err := loadProgram(conf.buildContext, map[string]bool{path: true})
	if err != nil {
		return err
	}

	_, rev, errors := importgraph.Build(conf.buildContext)
	if len(errors) > 0 {
		// With a large GOPATH tree, errors are inevitable.
		// Report them but proceed.
		log.Printf("while scanning Go workspace:\n")
		for path, err := range errors {
			log.Printf("Package %q: %s.\n", path, err)
		}
	}

	// Enumerate the set of potentially affected packages.
	possiblePackages := make(map[string]bool)
	for _, obj := range findExportedObjects(prog, path) {
		for path := range rev.Search(obj.Pkg().Path()) {
			possiblePackages[path] = true
		}
	}

	if conf.verbose {
		log.Println("Possible affected packages:")
		for pkg := range possiblePackages {
			log.Println("\t", pkg)
		}
	}

	// reload the program with all possible packages to fetch the packageinfo's
	globalProg, err := loadProgram(conf.buildContext, possiblePackages)
	if err != nil {
		return err
	}

	objsToUpdate := make(map[types.Object]bool, 0)
	objects := findExportedObjects(globalProg, path)

	if conf.verbose {
		log.Println("Exported identififers are:")
		for _, obj := range objects {
			log.Println("\t", obj)
		}
	}

	// filter out identifiers which can't be renamed due any collision or if
	// they are part of any other third party package (because renaming would
	// cause to break other packages)
	for _, info := range globalProg.Imported {
		safeObjects := filterObjects(info, objects, conf.identifiers)
		for _, obj := range safeObjects {
			// we don't care about other packages anymore
			if info.Pkg.Path() != path {
				continue
			}

			// don't include collisions
			newName := toLowerCase(obj.Name())
			if info.Pkg.Path() == obj.Pkg().Path() && hasObject(info, newName) {
				log.Printf("WARNING! can't unexport %q due collision. Identifier %q already exists.\n",
					obj.Name(), newName)
				continue
			}

			objsToUpdate[obj] = true
		}
	}

	if conf.verbose {
		log.Println("Safe to unexport identifiers are:")
		for obj := range objsToUpdate {
			log.Println("\t", obj)
		}
	}

	// first create the files that needs an update and modify the fileset
	var nidents int
	var filesToUpdate = make(map[*token.File]bool)
	for _, info := range globalProg.Imported {
		for id, obj := range info.Defs {
			if objsToUpdate[obj] {
				nidents++
				id.Name = toLowerCase(obj.Name())
				filesToUpdate[globalProg.Fset.File(id.Pos())] = true
			}
		}
		for id, obj := range info.Uses {
			if objsToUpdate[obj] {
				nidents++
				id.Name = toLowerCase(obj.Name())
				filesToUpdate[globalProg.Fset.File(id.Pos())] = true
			}
		}
	}

	// now start to rewrite the files
	var nerrs, npkgs int
	for _, info := range globalProg.Imported {
		first := true
		for _, f := range info.Files {
			tokenFile := globalProg.Fset.File(f.Pos())
			if filesToUpdate[tokenFile] {
				if first {
					npkgs++
					first = false
				}

				if conf.dryRun {
					continue
				}

				if err := rewriteFile(globalProg.Fset, f, tokenFile.Name()); err != nil {
					log.Println(err)
					nerrs++
				}
			}
		}
	}

	log.Printf("Unexported %d identifier%s in %d file%s in %d package%s.\n", nidents, plural(nidents),
		len(filesToUpdate), plural(len(filesToUpdate)),
		npkgs, plural(npkgs))
	if conf.verbose {
		log.Println("Identifiers changed:")
		for obj := range objsToUpdate {
			log.Println("\t", obj)
		}
		log.Println("Files changed:")
		for f := range filesToUpdate {
			log.Println("\t", f.Name())
		}

	}

	if nerrs > 0 {
		return fmt.Errorf("failed to rewrite %d file%s", nerrs, plural(nerrs))
	}

	return nil
}

func toLowerCase(s string) string {
	if s == "" {
		return ""
	}

	r, n := utf8.DecodeRuneInString(s)
	return string(unicode.ToLower(r)) + s[n:]
}

// plural, copied from golang.org/x/tools/refactor/rename
func plural(n int) string {
	if n != 1 {
		return "s"
	}
	return ""
}

func rewriteFile(fset *token.FileSet, f *ast.File, filename string) error {
	var buf bytes.Buffer
	if err := format.Node(&buf, fset, f); err != nil {
		return fmt.Errorf("failed to pretty-print syntax tree: %v", err)
	}
	return ioutil.WriteFile(filename, buf.Bytes(), 0644)
}

// filterObjects filters the given objects and returns objects which are not in use by the given info package.
func filterObjects(info *loader.PackageInfo, exported map[*ast.Ident]types.Object, allowed []string) map[*ast.Ident]types.Object {
	isAllowed := func(id string) bool {
		for _, i := range allowed {
			if i == id {
				return true
			}
		}
		return false
	}

	if len(allowed) == 0 {
		isAllowed = func(id string) bool { return true }
	}

	filtered := make(map[*ast.Ident]types.Object, 0)
	for id, ex := range exported {
		// TODO(arslan): not sure if we should allow test functions?
		if !hasUse(info, ex) && isAllowed(ex.Name()) {
			filtered[id] = ex
		}
	}

	return filtered
}

// hasUse returns true if the given obj is part of the use in info
func hasUse(info *loader.PackageInfo, obj types.Object) bool {
	for _, o := range info.Uses {
		if o == obj {
			return true
		}
	}
	return false
}

// hasObjects returns true if the given name exists in the definition of the package info
func hasObject(info *loader.PackageInfo, name string) bool {
	for _, obj := range info.Defs {
		if obj == nil {
			continue
		}

		if obj.Name() == name {
			return true
		}
	}

	return false
}

// exportedObjects returns objects which are exported only
func exportedObjects(info *loader.PackageInfo) map[*ast.Ident]types.Object {
	objects := make(map[*ast.Ident]types.Object, 0)
	for id, obj := range info.Defs {
		if obj == nil {
			continue
		}

		if obj.Exported() {
			objects[id] = obj
		}
	}

	return objects
}

// findExportedObjects returns a map of exported paths which are part of the
// given import path
func findExportedObjects(prog *loader.Program, path string) map[*ast.Ident]types.Object {
	var pkgObj *types.Package
	for pkg := range prog.AllPackages {
		if pkg.Path() == path {
			pkgObj = pkg
			break
		}
	}

	info := prog.AllPackages[pkgObj]
	return exportedObjects(info)
}

func loadProgram(ctxt *build.Context, pkgs map[string]bool) (*loader.Program, error) {
	conf := loader.Config{
		Build:       ctxt,
		ParserMode:  parser.ParseComments,
		AllowErrors: false,
	}

	for pkg := range pkgs {
		conf.ImportWithTests(pkg)
	}
	return conf.Load()
}
<commit_msg>unexport: just add some notes<commit_after>// The unexport command unexports exported identifiers which are not imported
// by any other Go code.
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"unicode"
	"unicode/utf8"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types"
	"golang.org/x/tools/refactor/importgraph"
)

func main() {
	var (
		flagPackage    = flag.String("package", "", "package import path to be unexported")
		flagIdentifier = flag.String("identifier", "", "comma-separated list of identifiers names; if empty all identifiers are unexported")
		flagDryRun     = flag.Bool("dryrun", false, "show the change, but do not apply")
		flagVerbose    = flag.Bool("verbose", false, "show more information. Useful for debugging.")
	)

	flag.Var((*buildutil.TagsFlag)(&build.Default.BuildTags), "tags", buildutil.TagsFlagDoc)

	flag.Parse()
	log.SetPrefix("unexport:")

	identifiers := []string{}
	if *flagIdentifier != "" {
		identifiers = strings.Split(*flagIdentifier, ",")
	}

	if err := runMain(&config{
		importPath:   *flagPackage,
		identifiers:  identifiers,
		buildContext: &build.Default,
		dryRun:       *flagDryRun,
		verbose:      *flagVerbose,
	}); err != nil {
		log.Println(err)
		os.Exit(1)
	}
}

// config is used to define how unexport should be work
type config struct {
	// importPath defines the package defined with the importpath
	importPath string

	// identifiers is used to limit the changes of unexporting to certain identifiers
	identifiers []string

	// build context
	buildContext *build.Context

	// logging/development ...
	dryRun  bool
	verbose bool
}

// runMain runs the actual command. It's an helper function so we can easily
// calls defers or return errors.  most of the functionality can be seen in
// `gorename` code source. Actually unexport is something that probably
// gorename can do for us.
func runMain(conf *config) error {
	if conf.importPath == "" {
		return errors.New("import path of the package must be given")
	}

	path := conf.importPath

	prog, err := loadProgram(conf.buildContext, map[string]bool{path: true})
	if err != nil {
		return err
	}

	_, rev, errors := importgraph.Build(conf.buildContext)
	if len(errors) > 0 {
		// With a large GOPATH tree, errors are inevitable.
		// Report them but proceed.
		log.Printf("while scanning Go workspace:\n")
		for path, err := range errors {
			log.Printf("Package %q: %s.\n", path, err)
		}
	}

	// Enumerate the set of potentially affected packages.
	possiblePackages := make(map[string]bool)
	for _, obj := range findExportedObjects(prog, path) {
		for path := range rev.Search(obj.Pkg().Path()) {
			possiblePackages[path] = true
		}
	}

	if conf.verbose {
		log.Println("Possible affected packages:")
		for pkg := range possiblePackages {
			log.Println("\t", pkg)
		}
	}

	// reload the program with all possible packages to fetch the packageinfo's
	globalProg, err := loadProgram(conf.buildContext, possiblePackages)
	if err != nil {
		return err
	}

	objsToUpdate := make(map[types.Object]bool, 0)
	objects := findExportedObjects(globalProg, path)

	if conf.verbose {
		log.Println("Exported identififers are:")
		for _, obj := range objects {
			log.Println("\t", obj)
		}
	}

	// filter out identifiers which can't be renamed due any collision or if
	// they are part of any other third party package (because renaming would
	// cause to break other packages)
	for _, info := range globalProg.Imported {
		safeObjects := filterObjects(info, objects, conf.identifiers)
		for _, obj := range safeObjects {
			// we don't care about other packages anymore
			if info.Pkg.Path() != path {
				continue
			}

			// don't include collisions
			newName := toLowerCase(obj.Name())
			if info.Pkg.Path() == obj.Pkg().Path() && hasObject(info, newName) {
				log.Printf("WARNING! can't unexport %q due collision. Identifier %q already exists.\n",
					obj.Name(), newName)
				continue
			}

			objsToUpdate[obj] = true
		}
	}

	if conf.verbose {
		log.Println("Safe to unexport identifiers are:")
		for obj := range objsToUpdate {
			log.Println("\t", obj)
		}
	}

	// first create the files that needs an update and modify the fileset
	var nidents int
	var filesToUpdate = make(map[*token.File]bool)
	for _, info := range globalProg.Imported {
		for id, obj := range info.Defs {
			if objsToUpdate[obj] {
				nidents++
				id.Name = toLowerCase(obj.Name())
				filesToUpdate[globalProg.Fset.File(id.Pos())] = true
			}
		}
		for id, obj := range info.Uses {
			if objsToUpdate[obj] {
				nidents++
				id.Name = toLowerCase(obj.Name())
				filesToUpdate[globalProg.Fset.File(id.Pos())] = true
			}
		}
	}

	// now start to rewrite the files
	var nerrs, npkgs int
	for _, info := range globalProg.Imported {
		first := true
		for _, f := range info.Files {
			tokenFile := globalProg.Fset.File(f.Pos())
			if filesToUpdate[tokenFile] {
				if first {
					npkgs++
					first = false
				}

				if conf.dryRun {
					continue
				}

				if err := rewriteFile(globalProg.Fset, f, tokenFile.Name()); err != nil {
					log.Println(err)
					nerrs++
				}
			}
		}
	}

	log.Printf("Unexported %d identifier%s in %d file%s in %d package%s.\n", nidents, plural(nidents),
		len(filesToUpdate), plural(len(filesToUpdate)),
		npkgs, plural(npkgs))
	if conf.verbose {
		log.Println("Identifiers changed:")
		for obj := range objsToUpdate {
			log.Println("\t", obj)
		}
		log.Println("Files changed:")
		for f := range filesToUpdate {
			log.Println("\t", f.Name())
		}

	}

	if nerrs > 0 {
		return fmt.Errorf("failed to rewrite %d file%s", nerrs, plural(nerrs))
	}

	return nil
}

func toLowerCase(s string) string {
	if s == "" {
		return ""
	}

	r, n := utf8.DecodeRuneInString(s)
	return string(unicode.ToLower(r)) + s[n:]
}

// plural, copied from golang.org/x/tools/refactor/rename
func plural(n int) string {
	if n != 1 {
		return "s"
	}
	return ""
}

func rewriteFile(fset *token.FileSet, f *ast.File, filename string) error {
	var buf bytes.Buffer
	if err := format.Node(&buf, fset, f); err != nil {
		return fmt.Errorf("failed to pretty-print syntax tree: %v", err)
	}
	return ioutil.WriteFile(filename, buf.Bytes(), 0644)
}

// filterObjects filters the given objects and returns objects which are not in use by the given info package.
func filterObjects(info *loader.PackageInfo, exported map[*ast.Ident]types.Object, allowed []string) map[*ast.Ident]types.Object {
	isAllowed := func(id string) bool {
		for _, i := range allowed {
			if i == id {
				return true
			}
		}
		return false
	}

	if len(allowed) == 0 {
		isAllowed = func(id string) bool { return true }
	}

	filtered := make(map[*ast.Ident]types.Object, 0)
	for id, ex := range exported {
		// TODO(arslan): not sure if we should allow test functions?
		if !hasUse(info, ex) && isAllowed(ex.Name()) {
			filtered[id] = ex
		}
	}

	return filtered
}

// hasUse returns true if the given obj is part of the use in info
func hasUse(info *loader.PackageInfo, obj types.Object) bool {
	for _, o := range info.Uses {
		if o == obj {
			return true
		}
	}
	return false
}

// hasObjects returns true if the given name exists in the definition of the package info
func hasObject(info *loader.PackageInfo, name string) bool {
	for _, obj := range info.Defs {
		if obj == nil {
			continue
		}

		if obj.Name() == name {
			return true
		}
	}

	return false
}

// exportedObjects returns objects which are exported only
func exportedObjects(info *loader.PackageInfo) map[*ast.Ident]types.Object {
	objects := make(map[*ast.Ident]types.Object, 0)
	for id, obj := range info.Defs {
		if obj == nil {
			continue
		}

		if obj.Exported() {
			objects[id] = obj
		}
	}

	return objects
}

// findExportedObjects returns a map of exported paths which are part of the
// given import path
func findExportedObjects(prog *loader.Program, path string) map[*ast.Ident]types.Object {
	var pkgObj *types.Package
	for pkg := range prog.AllPackages {
		if pkg.Path() == path {
			pkgObj = pkg
			break
		}
	}

	info := prog.AllPackages[pkgObj]
	return exportedObjects(info)
}

func loadProgram(ctxt *build.Context, pkgs map[string]bool) (*loader.Program, error) {
	conf := loader.Config{
		Build:       ctxt,
		ParserMode:  parser.ParseComments,
		AllowErrors: false,
	}

	for pkg := range pkgs {
		conf.ImportWithTests(pkg)
	}
	return conf.Load()
}
<|endoftext|>


<commit_before>package checkload

import (
	"os/exec"
	"strconv"
	"strings"
)

func getloadavg() (loadavgs [3]float64, err error) {
	outputBytes, err := exec.Command("sysctl", "-n", "vm.loadavg").Output()
	if err != nil {
		return loadavgs, err
	}

	output := string(outputBytes)

	// fields will be "<loadavg1>, <loadavg5>, <loadavg15>"
	fields := strings.Fields(output)
	if len(fields) != 3 {
		return loadavgs, err
	}

	for i := 0; i < 3; i++ {
		loadavg, err := strconv.ParseFloat(fields[i], 64)
		if err != nil {
			return loadavgs, err
		}
		loadavgs[i] = loadavg
	}
  return loadavgs, nil
}
<commit_msg>updated check-load for netbsd, using go-osstat instead of sysctl command<commit_after>package checkload

import (
  "github.com/mackerelio/go-osstat/loadavg"
)

func getloadavg() (loadavgs [3]float64, err error) {
  output, err := loadavg.Get()
  loadavgs = [3]float64{output.Loadavg1,output.Loadavg5,output.Loadavg15}

  return loadavgs, nil
}
<|endoftext|>


<commit_before>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(codecContext)

	return ist.Index(), ost.Index(), nil
}

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	streamMap := make(map[int]int, 0)
	var lastDelta int64

	// create input context
	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	// create output context
	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	// add video stream to streamMap
	srcVideoStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)
	if err != nil {
		return err
	}
	videoCodec := getVideoCodec(job)

	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		return err
	}
	streamMap[i] = o

	// add audio stream to streamMap
	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}
	audioCodec := getAudioCodec(job)

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	streamMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}

	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())
	framesCount := float64(0)

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(streamMap[ist.Index()])
		if err != nil {
			return err
		}

		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.SetNbSamples(ost.CodecCtx().FrameSize())
				frame.SetFormat(ost.CodecCtx().SampleFmt())
				frame.SetChannelLayout(ost.CodecCtx().ChannelLayout())
				frame.SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}
		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(streamMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getVideoCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
		"aac":    "aac",
	}

	if codec, ok := codecs[job.Preset.Video.Codec]; ok {
		return codec
	}
	return "libx264"
}

func getAudioCodec(job types.Job) string {
	codecs := map[string]string{
		"aac":    "aac",
		"vorbis": "vorbis",
	}
	if codec, ok := codecs[job.Preset.Audio.Codec]; ok {
		return codec
	}
	return "aac"
}

func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudioCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<commit_msg>encoders/ffmpeg: move addStream() to bottom<commit_after>package encoders

import (
	"errors"
	"strconv"

	"code.cloudfoundry.org/lager"

	"github.com/3d0c/gmf"
	"github.com/snickers/snickers/db"
	"github.com/snickers/snickers/types"
)

// FFMPEGEncode function is responsible for encoding the file
func FFMPEGEncode(logger lager.Logger, dbInstance db.Storage, jobID string) error {
	log := logger.Session("ffmpeg-encode")
	log.Info("started", lager.Data{"job": jobID})
	defer log.Info("finished")

	gmf.LogSetLevel(gmf.AV_LOG_FATAL)
	job, _ := dbInstance.RetrieveJob(jobID)
	streamMap := make(map[int]int, 0)
	var lastDelta int64

	// create input context
	inputCtx, err := gmf.NewInputCtx(job.LocalSource)
	if err != nil {
		log.Error("input-failed", err)
		return err
	}
	defer inputCtx.CloseInputAndRelease()

	// create output context
	outputCtx, err := gmf.NewOutputCtx(job.LocalDestination)
	if err != nil {
		log.Error("output-failed", err)
		return err
	}
	defer outputCtx.CloseOutputAndRelease()

	job.Status = types.JobEncoding
	job.Details = "0%"
	dbInstance.UpdateJob(job.ID, job)

	// add video stream to streamMap
	srcVideoStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_VIDEO)
	if err != nil {
		return err
	}
	videoCodec := getVideoCodec(job)

	i, o, err := addStream(job, videoCodec, outputCtx, srcVideoStream)
	if err != nil {
		return err
	}
	streamMap[i] = o

	// add audio stream to streamMap
	srcAudioStream, err := inputCtx.GetBestStream(gmf.AVMEDIA_TYPE_AUDIO)
	if err != nil {
		return err
	}
	audioCodec := getAudioCodec(job)

	i, o, err = addStream(job, audioCodec, outputCtx, srcAudioStream)
	if err != nil {
		return err
	}
	streamMap[i] = o

	if err := outputCtx.WriteHeader(); err != nil {
		return err
	}

	totalFrames := float64(srcVideoStream.NbFrames() + srcAudioStream.NbFrames())

	for packet := range inputCtx.GetNewPackets() {
		ist, err := inputCtx.GetStream(packet.StreamIndex())
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(streamMap[ist.Index()])
		if err != nil {
			return err
		}

		framesCount := float64(0)
		for frame := range packet.Frames(ist.CodecCtx()) {
			if ost.IsAudio() {
				fsTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}
				outTb := gmf.AVR{Num: 1, Den: ist.CodecCtx().SampleRate()}

				frame.SetPts(packet.Pts())

				pts := gmf.RescaleDelta(ist.TimeBase(), frame.Pts(), fsTb.AVRational(), frame.NbSamples(), &lastDelta, outTb.AVRational())

				frame.SetNbSamples(ost.CodecCtx().FrameSize())
				frame.SetFormat(ost.CodecCtx().SampleFmt())
				frame.SetChannelLayout(ost.CodecCtx().ChannelLayout())
				frame.SetPts(pts)
			} else {
				frame.SetPts(ost.Pts)
			}

			if p, ready, _ := frame.EncodeNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			}

			ost.Pts++
			framesCount++
			percentage := string(strconv.FormatInt(int64(framesCount/totalFrames*100), 10) + "%")
			if percentage != job.Details {
				job.Details = percentage
				dbInstance.UpdateJob(job.ID, job)
			}
		}

		gmf.Release(packet)
	}

	for i := 0; i < outputCtx.StreamsCnt(); i++ {
		ist, err := inputCtx.GetStream(0)
		if err != nil {
			return err
		}
		ost, err := outputCtx.GetStream(streamMap[ist.Index()])
		if err != nil {
			return err
		}

		frame := gmf.NewFrame()

		for {
			if p, ready, _ := frame.FlushNewPacket(ost.CodecCtx()); ready {
				if p.Pts() != gmf.AV_NOPTS_VALUE {
					p.SetPts(gmf.RescaleQ(p.Pts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				if p.Dts() != gmf.AV_NOPTS_VALUE {
					p.SetDts(gmf.RescaleQ(p.Dts(), ost.CodecCtx().TimeBase(), ost.TimeBase()))
				}

				p.SetStreamIndex(ost.Index())

				if err := outputCtx.WritePacket(p); err != nil {
					return err
				}
				gmf.Release(p)
			} else {
				gmf.Release(p)
				break
			}

			ost.Pts++
		}

		gmf.Release(frame)
	}
	if job.Details != "100%" {
		job.Details = "100%"
		dbInstance.UpdateJob(job.ID, job)
	}

	return nil
}

func addStream(job types.Job, codecName string, oc *gmf.FmtCtx, ist *gmf.Stream) (int, int, error) {
	var codecContext *gmf.CodecCtx
	var ost *gmf.Stream

	codec, err := gmf.FindEncoder(codecName)
	if err != nil {
		return 0, 0, err
	}

	if ost = oc.NewStream(codec); ost == nil {
		return 0, 0, errors.New("unable to create stream in output context")
	}
	defer gmf.Release(ost)

	if codecContext = gmf.NewCodecCtx(codec); codecContext == nil {
		return 0, 0, errors.New("unable to create codec context")
	}
	defer gmf.Release(codecContext)

	// https://ffmpeg.org/pipermail/ffmpeg-devel/2008-January/046900.html
	if oc.IsGlobalHeader() {
		codecContext.SetFlag(gmf.CODEC_FLAG_GLOBAL_HEADER)
	}

	if codec.IsExperimental() {
		codecContext.SetStrictCompliance(gmf.FF_COMPLIANCE_EXPERIMENTAL)
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_AUDIO {
		err := setAudioCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if codecContext.Type() == gmf.AVMEDIA_TYPE_VIDEO {
		err := setVideoCtxParams(codecContext, ist, job)
		if err != nil {
			return 0, 0, err
		}
	}

	if err := codecContext.Open(nil); err != nil {
		return 0, 0, err
	}

	ost.SetCodecCtx(codecContext)

	return ist.Index(), ost.Index(), nil
}

func getProfile(job types.Job) int {
	profiles := map[string]int{
		"baseline": gmf.FF_PROFILE_H264_BASELINE,
		"main":     gmf.FF_PROFILE_H264_MAIN,
		"high":     gmf.FF_PROFILE_H264_HIGH,
	}

	if job.Preset.Video.Profile != "" {
		return profiles[job.Preset.Video.Profile]
	}
	return gmf.FF_PROFILE_H264_MAIN
}

func getVideoCodec(job types.Job) string {
	codecs := map[string]string{
		"h264":   "libx264",
		"vp8":    "libvpx",
		"vp9":    "libvpx-vp9",
		"theora": "libtheora",
		"aac":    "aac",
	}

	if codec, ok := codecs[job.Preset.Video.Codec]; ok {
		return codec
	}
	return "libx264"
}

func getAudioCodec(job types.Job) string {
	codecs := map[string]string{
		"aac":    "aac",
		"vorbis": "vorbis",
	}
	if codec, ok := codecs[job.Preset.Audio.Codec]; ok {
		return codec
	}
	return "aac"
}

func GetResolution(job types.Job, inputWidth int, inputHeight int) (int, int) {
	var width, height int
	if job.Preset.Video.Width == "" && job.Preset.Video.Height == "" {
		return inputWidth, inputHeight
	} else if job.Preset.Video.Width == "" {
		height, _ = strconv.Atoi(job.Preset.Video.Height)
		width = (inputWidth * height) / inputHeight
	} else if job.Preset.Video.Height == "" {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height = (inputHeight * width) / inputWidth
	} else {
		width, _ = strconv.Atoi(job.Preset.Video.Width)
		height, _ = strconv.Atoi(job.Preset.Video.Height)
	}
	return width, height
}

func setAudioCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	bitrate, err := strconv.Atoi(job.Preset.Audio.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetBitRate(bitrate)
	codecContext.SetSampleFmt(ist.CodecCtx().SampleFmt())
	codecContext.SetSampleRate(ist.CodecCtx().SampleRate())
	codecContext.SetChannels(ist.CodecCtx().Channels())
	codecContext.SelectChannelLayout()
	codecContext.SelectSampleRate()
	return nil
}

func setVideoCtxParams(codecContext *gmf.CodecCtx, ist *gmf.Stream, job types.Job) error {
	codecContext.SetTimeBase(gmf.AVR{Num: 1, Den: 25}) // what is this

	if job.Preset.Video.Codec == "h264" {
		profile := getProfile(job)
		codecContext.SetProfile(profile)
	}

	gop, err := strconv.Atoi(job.Preset.Video.GopSize)
	if err != nil {
		return err
	}

	width, height := GetResolution(job, ist.CodecCtx().Width(), ist.CodecCtx().Height())

	bitrate, err := strconv.Atoi(job.Preset.Video.Bitrate)
	if err != nil {
		return err
	}

	codecContext.SetDimension(width, height)
	codecContext.SetGopSize(gop)
	codecContext.SetBitRate(bitrate)
	codecContext.SetPixFmt(ist.CodecCtx().PixFmt())

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"

	"github.com/docker/libnetwork/ipamapi"
	weaveapi "github.com/weaveworks/weave/api"
	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/common/docker"
	weavenet "github.com/weaveworks/weave/net"
	ipamplugin "github.com/weaveworks/weave/plugin/ipam"
	netplugin "github.com/weaveworks/weave/plugin/net"
	"github.com/weaveworks/weave/plugin/skel"
)

var version = "unreleased"

var Log = common.Log

func main() {
	var (
		justVersion      bool
		address          string
		meshAddress      string
		logLevel         string
		noMulticastRoute bool
	)

	flag.BoolVar(&justVersion, "version", false, "print version and exit")
	flag.StringVar(&logLevel, "log-level", "info", "logging level (debug, info, warning, error)")
	flag.StringVar(&address, "socket", "/run/docker/plugins/weave.sock", "socket on which to listen")
	flag.StringVar(&meshAddress, "meshsocket", "/run/docker/plugins/weavemesh.sock", "socket on which to listen in mesh mode")
	flag.BoolVar(&noMulticastRoute, "no-multicast-route", false, "deprecated (this is now the default)")

	flag.Parse()

	if justVersion {
		fmt.Printf("weave plugin %s\n", version)
		os.Exit(0)
	}

	common.SetLogLevel(logLevel)

	weave := weaveapi.NewClient(os.Getenv("WEAVE_HTTP_ADDR"), Log)

	// API 1.21 is the first version that supports docker network commands
	dockerClient, err := docker.NewVersionedClientFromEnv("1.21")
	if err != nil {
		Log.Fatalf("unable to connect to docker: %s", err)
	}

	Log.Println("Weave plugin", version, "Command line options:", os.Args[1:])
	if noMulticastRoute {
		Log.Warning("--no-multicast-route option has been removed; multicast is off by default")
	}
	Log.Info(dockerClient.Info())

	err = run(dockerClient, weave, address, meshAddress)
	if err != nil {
		Log.Fatal(err)
	}
}

func run(dockerClient *docker.Client, weave *weaveapi.Client, address, meshAddress string) error {
	endChan := make(chan error, 1)
	if address != "" {
		globalListener, err := listenAndServe(dockerClient, weave, address, endChan, "global", false)
		if err != nil {
			return err
		}
		defer os.Remove(address)
		defer globalListener.Close()
	}
	if meshAddress != "" {
		meshListener, err := listenAndServe(dockerClient, weave, meshAddress, endChan, "local", true)
		if err != nil {
			return err
		}
		defer os.Remove(meshAddress)
		defer meshListener.Close()
	}

	statusListener, err := weavenet.ListenUnixSocket("/home/weave/status.sock")
	if err != nil {
		return err
	}
	go serveStatus(statusListener)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	select {
	case sig := <-sigChan:
		Log.Debugf("Caught signal %s; shutting down", sig)
		return nil
	case err := <-endChan:
		return err
	}
}

func listenAndServe(dockerClient *docker.Client, weave *weaveapi.Client, address string, endChan chan<- error, scope string, withIpam bool) (net.Listener, error) {
	name := strings.TrimSuffix(path.Base(address), ".sock")
	d, err := netplugin.New(dockerClient, weave, name, scope)
	if err != nil {
		return nil, err
	}

	var i ipamapi.Ipam
	if withIpam {
		i = ipamplugin.NewIpam(weave)
	}

	listener, err := weavenet.ListenUnixSocket(address)
	if err != nil {
		return nil, err
	}
	Log.Printf("Listening on %s for %s scope", address, scope)

	go func() {
		endChan <- skel.Listen(listener, d, i)
	}()

	return listener, nil
}

func serveStatus(listener net.Listener) {
	server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "ok")
	})}
	if err := server.Serve(listener); err != nil {
		Log.Fatalf("ListenAndServeStatus failed: %s", err)
	}
}
<commit_msg>Allow docker API location to be specified or blanked<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"

	"github.com/docker/libnetwork/ipamapi"
	weaveapi "github.com/weaveworks/weave/api"
	"github.com/weaveworks/weave/common"
	"github.com/weaveworks/weave/common/docker"
	weavenet "github.com/weaveworks/weave/net"
	ipamplugin "github.com/weaveworks/weave/plugin/ipam"
	netplugin "github.com/weaveworks/weave/plugin/net"
	"github.com/weaveworks/weave/plugin/skel"
)

var version = "unreleased"

var Log = common.Log

func main() {
	var (
		justVersion      bool
		dockerAPI        string
		address          string
		meshAddress      string
		logLevel         string
		noMulticastRoute bool

		defaultDockerHost = "unix:///var/run/docker.sock"
	)

	if val := os.Getenv("DOCKER_HOST"); val != "" {
		defaultDockerHost = val
	}

	flag.BoolVar(&justVersion, "version", false, "print version and exit")
	flag.StringVar(&logLevel, "log-level", "info", "logging level (debug, info, warning, error)")
	flag.StringVar(&dockerAPI, "docker-api", defaultDockerHost, "Docker API endpoint")
	flag.StringVar(&address, "socket", "/run/docker/plugins/weave.sock", "socket on which to listen")
	flag.StringVar(&meshAddress, "meshsocket", "/run/docker/plugins/weavemesh.sock", "socket on which to listen in mesh mode")
	flag.BoolVar(&noMulticastRoute, "no-multicast-route", false, "deprecated (this is now the default)")

	flag.Parse()

	if justVersion {
		fmt.Printf("weave plugin %s\n", version)
		os.Exit(0)
	}

	common.SetLogLevel(logLevel)

	weave := weaveapi.NewClient(os.Getenv("WEAVE_HTTP_ADDR"), Log)

	var dockerClient *docker.Client
	var err error
	if dockerAPI != "" {
		// API 1.21 is the first version that supports docker network commands
		dockerClient, err = docker.NewVersionedClient(dockerAPI, "1.21")
		if err != nil {
			Log.Fatalf("unable to connect to docker: %s", err)
		}
	}

	Log.Println("Weave plugin", version, "Command line options:", os.Args[1:])
	if noMulticastRoute {
		Log.Warning("--no-multicast-route option has been removed; multicast is off by default")
	}
	if dockerClient == nil {
		Log.Info("Running without Docker API connection")
	} else {
		Log.Info(dockerClient.Info())
	}

	err = run(dockerClient, weave, address, meshAddress)
	if err != nil {
		Log.Fatal(err)
	}
}

func run(dockerClient *docker.Client, weave *weaveapi.Client, address, meshAddress string) error {
	endChan := make(chan error, 1)
	if address != "" {
		globalListener, err := listenAndServe(dockerClient, weave, address, endChan, "global", false)
		if err != nil {
			return err
		}
		defer os.Remove(address)
		defer globalListener.Close()
	}
	if meshAddress != "" {
		meshListener, err := listenAndServe(dockerClient, weave, meshAddress, endChan, "local", true)
		if err != nil {
			return err
		}
		defer os.Remove(meshAddress)
		defer meshListener.Close()
	}

	statusListener, err := weavenet.ListenUnixSocket("/home/weave/status.sock")
	if err != nil {
		return err
	}
	go serveStatus(statusListener)
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	select {
	case sig := <-sigChan:
		Log.Debugf("Caught signal %s; shutting down", sig)
		return nil
	case err := <-endChan:
		return err
	}
}

func listenAndServe(dockerClient *docker.Client, weave *weaveapi.Client, address string, endChan chan<- error, scope string, withIpam bool) (net.Listener, error) {
	name := strings.TrimSuffix(path.Base(address), ".sock")
	d, err := netplugin.New(dockerClient, weave, name, scope)
	if err != nil {
		return nil, err
	}

	var i ipamapi.Ipam
	if withIpam {
		i = ipamplugin.NewIpam(weave)
	}

	listener, err := weavenet.ListenUnixSocket(address)
	if err != nil {
		return nil, err
	}
	Log.Printf("Listening on %s for %s scope", address, scope)

	go func() {
		endChan <- skel.Listen(listener, d, i)
	}()

	return listener, nil
}

func serveStatus(listener net.Listener) {
	server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "ok")
	})}
	if err := server.Serve(listener); err != nil {
		Log.Fatalf("ListenAndServeStatus failed: %s", err)
	}
}
<|endoftext|>


<commit_before>package ersatzjson

import (
	"testing"
)

type Animal struct {
	Name  string
	Order string
}

func TestUnmarshal(t *testing.T) {
	tests := []struct {
		jsonBlob    []byte
		expected    []Animal
		expectedErr string
	}{
		{jsonBlob: []byte(""), expected: []Animal{}, expectedErr: "unexpected end of JSON input"},
		// from http://golang.org/pkg/encoding/json/#example_Unmarshal
		{jsonBlob: []byte(`[
      {"Name": "Platypus", "Order": "Monotremata"},
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
		{jsonBlob: []byte(`// comments
      [
      {"Name": "Platypus", "Order": "Monotremata"},
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
		{jsonBlob: []byte(`[ /* block comment */
      {"Name": "Platypus", "Order": "Monotremata"},
      # another one
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
	}

	for i, test := range tests {
		var animals []Animal
		err := Unmarshal(test.jsonBlob, &animals)
		if err != nil {
			if err.Error() != test.expectedErr {
				t.Errorf("%d: expected err %q got %q", i, test.expectedErr, err)
			}
			continue
		}
		if len(test.expected) != len(animals) {
			t.Errorf("%d: expected %d results, go %d", len(test.expected), len(animals))
			continue
		}
		for i, v := range animals {
			j, found := findAnimal(v, test.expected)
			if !found {
				t.Errorf("%d: unexpected result: %q", i, v.Name)
				continue
			}
			if test.expected[j].Order != v.Order {
				t.Errorf("%d expected the order for %q to be %q, got %q", v.Name, test.expected[j].Order, v.Order)
			}
		}
	}
}

func findAnimal(val Animal, animals []Animal) (int, bool) {
	for i, v := range animals {
		if val.Name == v.Name {
			return i, true
		}
	}
	return 0, false
}
<commit_msg>fix testing error messages<commit_after>package ersatzjson

import (
	"testing"
)

type Animal struct {
	Name  string
	Order string
}

func TestUnmarshal(t *testing.T) {
	tests := []struct {
		jsonBlob    []byte
		expected    []Animal
		expectedErr string
	}{
		{jsonBlob: []byte(""), expected: []Animal{}, expectedErr: "unexpected end of JSON input"},
		// from http://golang.org/pkg/encoding/json/#example_Unmarshal
		{jsonBlob: []byte(`[
      {"Name": "Platypus", "Order": "Monotremata"},
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
		{jsonBlob: []byte(`// comments
      [
      {"Name": "Platypus", "Order": "Monotremata"},
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
		{jsonBlob: []byte(`[ /* block comment */
      {"Name": "Platypus", "Order": "Monotremata"},
      # another one
      {"Name": "Quoll", "Order": "Dasyuromorphia"}
    ]`), expected: []Animal{
			Animal{"Platypus", "Monotremata"},
			Animal{"Quoll", "Dasyuromorphia"},
		}, expectedErr: ""},
	}

	for i, test := range tests {
		var animals []Animal
		err := Unmarshal(test.jsonBlob, &animals)
		if err != nil {
			if err.Error() != test.expectedErr {
				t.Errorf("%d: expected err %q got %q", i, test.expectedErr, err)
			}
			continue
		}
		if len(test.expected) != len(animals) {
			t.Errorf("%d: expected %d results, go %d", i, len(test.expected), len(animals))
			continue
		}
		for i, v := range animals {
			j, found := findAnimal(v, test.expected)
			if !found {
				t.Errorf("%d: unexpected result: %q", i, v.Name)
				continue
			}
			if test.expected[j].Order != v.Order {
				t.Errorf("%d expected the order for %q to be %q, got %q", i, v.Name, test.expected[j].Order, v.Order)
			}
		}
	}
}

func findAnimal(val Animal, animals []Animal) (int, bool) {
	for i, v := range animals {
		if val.Name == v.Name {
			return i, true
		}
	}
	return 0, false
}
<|endoftext|>


<commit_before>package csrapprover

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"net"
	"net/url"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	certv1beta1 "k8s.io/api/certificates/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	e2e "k8s.io/kubernetes/test/e2e/framework"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"

	kubeclient "k8s.io/client-go/kubernetes"
	certclientv1beta1 "k8s.io/client-go/kubernetes/typed/certificates/v1beta1"
	restclient "k8s.io/client-go/rest"

	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/ibmcloud"
	"github.com/openshift/origin/test/extended/util/image"
)

var _ = g.Describe("[sig-cluster-lifecycle]", func() {
	oc := exutil.NewCLI("cluster-client-cert")
	defer g.GinkgoRecover()

	g.It("Pods cannot access the /config/master API endpoint", func() {
		// https://issues.redhat.com/browse/CO-760
		if e2e.TestContext.Provider == ibmcloud.ProviderName {
			e2eskipper.Skipf("IBM ROKS clusters do not expose machine configuration externally because they don't use RHCOS workers. " +
				"Remove this skip when https://issues.redhat.com/browse/CO-760 (RHCOS support) is implemented")
		}

		// the /config/master API port+endpoint is only visible from inside the cluster
		// (-> we need to create a pod to try to reach it) and contains the token
		// of the node-bootstrapper SA, so no random pods should be able to see it
		pod, err := exutil.NewPodExecutor(oc, "get-bootstrap-creds", image.ShellImage())
		o.Expect(err).NotTo(o.HaveOccurred())

		// get the API server URL, mutate to internal API (use infra.Status.APIServerURLInternal) once API is bumped
		infra, err := oc.AdminConfigClient().ConfigV1().Infrastructures().Get(context.Background(), "cluster", metav1.GetOptions{})
		o.Expect(err).NotTo(o.HaveOccurred())

		internalAPI, err := url.Parse(infra.Status.APIServerURL)
		o.Expect(err).NotTo(o.HaveOccurred())
		internalAPI.Host = strings.Replace(internalAPI.Host, "api.", "api-int.", 1)

		host, _, err := net.SplitHostPort(internalAPI.Host)
		o.Expect(err).ToNot(o.HaveOccurred())

		internalAPI.Host = net.JoinHostPort(host, "22623")
		internalAPI.Path = "/config/master"

		// we should not be able to reach the endpoint
		curlOutput, err := pod.Exec(fmt.Sprintf("curl -k %s", internalAPI.String()))
		o.Expect(err).To(o.HaveOccurred())
		o.Expect(curlOutput).To(o.ContainSubstring("Connection refused"))
	})

	g.It("CSRs from machines that are not recognized by the cloud provider are not approved", func() {
		if e2e.TestContext.Provider == ibmcloud.ProviderName {
			e2eskipper.Skipf("IBM ROKS clusters do not handle node bootstrapping in the cluster. The openshift-machine-config-operator/node-bootstrapper service account does not exist")
		}

		// we somehow were able to get the node-approver token, make sure we can't
		// create node certs with client auth with it
		priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		o.Expect(err).NotTo(o.HaveOccurred())

		certRequestTemplate := &x509.CertificateRequest{
			SignatureAlgorithm: x509.ECDSAWithSHA256,
			Subject: pkix.Name{
				CommonName:   "system:node:hacking-node.ec2.internal",
				Organization: []string{"system:nodes"},
			},
			PublicKey: priv.PublicKey,
		}

		csrbytes, err := x509.CreateCertificateRequest(rand.Reader, certRequestTemplate, priv)
		o.Expect(err).NotTo(o.HaveOccurred())

		// get the token of the node-bootstrapper and use it to build a client for it
		bootstrapperToken, err := oc.AsAdmin().Run("sa").Args("get-token", "node-bootstrapper", "-n", "openshift-machine-config-operator").Output()
		o.Expect(err).NotTo(o.HaveOccurred())

		saClientConfig := restclient.AnonymousClientConfig(oc.AdminConfig())
		saClientConfig.BearerToken = bootstrapperToken

		bootstrapperClient := kubeclient.NewForConfigOrDie(saClientConfig)

		csrName := "node-client-csr"
		bootstrapperClient.CertificatesV1beta1().CertificateSigningRequests().Create(context.Background(), &certv1beta1.CertificateSigningRequest{
			ObjectMeta: metav1.ObjectMeta{
				Name: csrName,
			},
			Spec: certv1beta1.CertificateSigningRequestSpec{
				Request: pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrbytes}),
				Usages: []certv1beta1.KeyUsage{
					certv1beta1.UsageDigitalSignature,
					certv1beta1.UsageKeyEncipherment,
					certv1beta1.UsageClientAuth,
				},
			},
		}, metav1.CreateOptions{})

		csrClient := oc.AdminKubeClient().CertificatesV1beta1().CertificateSigningRequests()
		defer cleanupCSR(csrClient, csrName)

		err = waitCSRStatus(csrClient, csrName)
		// if status did not change in 30 sec, the CSR is still in pending
		// which is fine as the machine-approver does not deny
		if err != wait.ErrWaitTimeout {
			o.Expect(err).NotTo(o.HaveOccurred())
		}
	})
})

// waits for the CSR object to change status, checks that it did not get approved
func waitCSRStatus(csrAdminClient certclientv1beta1.CertificateSigningRequestInterface, csrName string) error {
	return wait.Poll(1*time.Second, 30*time.Second, func() (bool, error) {
		csr, err := csrAdminClient.Get(context.Background(), csrName, metav1.GetOptions{})
		if err != nil {
			return false, err
		}
		if len(csr.Status.Conditions) > 0 {
			for _, c := range csr.Status.Conditions {
				if c.Type == certv1beta1.CertificateApproved {
					return true, fmt.Errorf("CSR for unknown node should not be approved")
				}
			}
		}
		return false, nil
	})
}

func cleanupCSR(csrAdminClient certclientv1beta1.CertificateSigningRequestInterface, name string) {
	csrAdminClient.Delete(context.Background(), name, metav1.DeleteOptions{})
}
<commit_msg>csrapprover: fix deprecated CSR api client<commit_after>package csrapprover

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"net"
	"net/url"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	e2e "k8s.io/kubernetes/test/e2e/framework"
	e2eskipper "k8s.io/kubernetes/test/e2e/framework/skipper"

	certv1 "k8s.io/api/certificates/v1"
	kubeclient "k8s.io/client-go/kubernetes"
	certclientv1 "k8s.io/client-go/kubernetes/typed/certificates/v1"
	restclient "k8s.io/client-go/rest"

	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/ibmcloud"
	"github.com/openshift/origin/test/extended/util/image"
)

var _ = g.Describe("[sig-cluster-lifecycle]", func() {
	oc := exutil.NewCLI("cluster-client-cert")
	defer g.GinkgoRecover()

	g.It("Pods cannot access the /config/master API endpoint", func() {
		// https://issues.redhat.com/browse/CO-760
		if e2e.TestContext.Provider == ibmcloud.ProviderName {
			e2eskipper.Skipf("IBM ROKS clusters do not expose machine configuration externally because they don't use RHCOS workers. " +
				"Remove this skip when https://issues.redhat.com/browse/CO-760 (RHCOS support) is implemented")
		}

		// the /config/master API port+endpoint is only visible from inside the cluster
		// (-> we need to create a pod to try to reach it) and contains the token
		// of the node-bootstrapper SA, so no random pods should be able to see it
		pod, err := exutil.NewPodExecutor(oc, "get-bootstrap-creds", image.ShellImage())
		o.Expect(err).NotTo(o.HaveOccurred())

		// get the API server URL, mutate to internal API (use infra.Status.APIServerURLInternal) once API is bumped
		infra, err := oc.AdminConfigClient().ConfigV1().Infrastructures().Get(context.Background(), "cluster", metav1.GetOptions{})
		o.Expect(err).NotTo(o.HaveOccurred())

		internalAPI, err := url.Parse(infra.Status.APIServerURL)
		o.Expect(err).NotTo(o.HaveOccurred())
		internalAPI.Host = strings.Replace(internalAPI.Host, "api.", "api-int.", 1)

		host, _, err := net.SplitHostPort(internalAPI.Host)
		o.Expect(err).ToNot(o.HaveOccurred())

		internalAPI.Host = net.JoinHostPort(host, "22623")
		internalAPI.Path = "/config/master"

		// we should not be able to reach the endpoint
		curlOutput, err := pod.Exec(fmt.Sprintf("curl -k %s", internalAPI.String()))
		o.Expect(err).To(o.HaveOccurred())
		o.Expect(curlOutput).To(o.ContainSubstring("Connection refused"))
	})

	g.It("CSRs from machines that are not recognized by the cloud provider are not approved", func() {
		if e2e.TestContext.Provider == ibmcloud.ProviderName {
			e2eskipper.Skipf("IBM ROKS clusters do not handle node bootstrapping in the cluster. The openshift-machine-config-operator/node-bootstrapper service account does not exist")
		}

		// we somehow were able to get the node-approver token, make sure we can't
		// create node certs with client auth with it
		priv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		o.Expect(err).NotTo(o.HaveOccurred())

		certRequestTemplate := &x509.CertificateRequest{
			SignatureAlgorithm: x509.ECDSAWithSHA256,
			Subject: pkix.Name{
				CommonName:   "system:node:hacking-node.ec2.internal",
				Organization: []string{"system:nodes"},
			},
			PublicKey: priv.PublicKey,
		}

		csrbytes, err := x509.CreateCertificateRequest(rand.Reader, certRequestTemplate, priv)
		o.Expect(err).NotTo(o.HaveOccurred())

		// get the token of the node-bootstrapper and use it to build a client for it
		bootstrapperToken, err := oc.AsAdmin().Run("sa").Args("get-token", "node-bootstrapper", "-n", "openshift-machine-config-operator").Output()
		o.Expect(err).NotTo(o.HaveOccurred())

		saClientConfig := restclient.AnonymousClientConfig(oc.AdminConfig())
		saClientConfig.BearerToken = bootstrapperToken

		bootstrapperClient := kubeclient.NewForConfigOrDie(saClientConfig)

		csrName := "node-client-csr"
		_, err = bootstrapperClient.CertificatesV1().CertificateSigningRequests().Create(context.Background(), &certv1.CertificateSigningRequest{
			ObjectMeta: metav1.ObjectMeta{
				Name: csrName,
			},
			Spec: certv1.CertificateSigningRequestSpec{
				SignerName: "kubernetes.io/kubelet-serving",
				Request:    pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrbytes}),
				Usages: []certv1.KeyUsage{
					certv1.UsageDigitalSignature,
					certv1.UsageKeyEncipherment,
					certv1.UsageClientAuth,
				},
			},
		}, metav1.CreateOptions{})
		o.Expect(err).NotTo(o.HaveOccurred())

		csrClient := oc.AdminKubeClient().CertificatesV1().CertificateSigningRequests()
		defer cleanupCSR(csrClient, csrName)

		err = waitCSRStatus(csrClient, csrName)
		// if status did not change in 30 sec, the CSR is still in pending
		// which is fine as the machine-approver does not deny
		if err != wait.ErrWaitTimeout {
			o.Expect(err).NotTo(o.HaveOccurred())
		}
	})
})

// waits for the CSR object to change status, checks that it did not get approved
func waitCSRStatus(csrAdminClient certclientv1.CertificateSigningRequestInterface, csrName string) error {
	return wait.Poll(1*time.Second, 30*time.Second, func() (bool, error) {
		csr, err := csrAdminClient.Get(context.Background(), csrName, metav1.GetOptions{})
		if err != nil {
			return false, err
		}
		if len(csr.Status.Conditions) > 0 {
			for _, c := range csr.Status.Conditions {
				if c.Type == certv1.CertificateApproved {
					return true, fmt.Errorf("CSR for unknown node should not be approved")
				}
			}
		}
		return false, nil
	})
}

func cleanupCSR(csrAdminClient certclientv1.CertificateSigningRequestInterface, name string) {
	csrAdminClient.Delete(context.Background(), name, metav1.DeleteOptions{})
}
<|endoftext|>


<commit_before>package registry

import (
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/micro/go-micro/v3/logger"
	"github.com/micro/go-micro/v3/router"
)

// table is an in-memory routing table
type table struct {
	sync.RWMutex
	// lookup for a service
	lookup func(string) ([]router.Route, error)
	// routes stores service routes
	routes map[string]map[uint64]*route
	// watchers stores table watchers
	watchers map[string]*tableWatcher
}

type route struct {
	route   router.Route
	updated time.Time
}

// newtable creates a new routing table and returns it
func newTable(lookup func(string) ([]router.Route, error), opts ...router.Option) *table {
	return &table{
		lookup:   lookup,
		routes:   make(map[string]map[uint64]*route),
		watchers: make(map[string]*tableWatcher),
	}
}

// pruneRoutes will prune routes older than the time specified
func (t *table) pruneRoutes(olderThan time.Duration) {
	var routes []router.Route

	t.Lock()

	// search for all the routes
	for _, routeList := range t.routes {
		for _, r := range routeList {
			// if any route is older than
			if time.Since(r.updated).Seconds() > olderThan.Seconds() {
				routes = append(routes, r.route)
			}
		}
	}

	t.Unlock()

	// delete the routes we've found
	for _, route := range routes {
		t.Delete(route)
	}
}

// deleteService removes the entire service
func (t *table) deleteService(service, network string) {
	t.Lock()

	routes, ok := t.routes[service]
	if !ok {
		return
	}

	// delete the routes for the service
	for hash, rt := range routes {
		// TODO: check if this causes a problem
		// with * in the network if that is a thing
		// or blank strings
		if rt.route.Network != network {
			continue
		}
		delete(routes, hash)
	}

	// delete the map for the service if its empty
	if len(routes) == 0 {
		delete(t.routes, service)
		return
	}

	// save the routes
	t.routes[service] = routes

	t.Unlock()
}

// saveRoutes completely replaces the routes for a service
func (t *table) saveRoutes(service string, routes []router.Route) {
	t.Lock()
	defer t.Unlock()

	// delete old routes
	delete(t.routes, service)
	// make new map
	t.routes[service] = make(map[uint64]*route)

	// iterate through new routes and save
	for _, rt := range routes {
		// save the new route
		t.routes[service][rt.Hash()] = &route{rt, time.Now()}
	}
}

// sendEvent sends events to all subscribed watchers
func (t *table) sendEvent(e *router.Event) {
	t.RLock()
	defer t.RUnlock()

	if len(e.Id) == 0 {
		e.Id = uuid.New().String()
	}

	for _, w := range t.watchers {
		select {
		case w.resChan <- e:
		case <-w.done:
		// don't block forever
		case <-time.After(time.Second):
		}
	}
}

// Create creates new route in the routing table
func (t *table) Create(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	// check if there are any routes in the table for the route destination
	if _, ok := t.routes[service]; !ok {
		t.routes[service] = make(map[uint64]*route)
	}

	// add new route to the table for the route destination
	if _, ok := t.routes[service][sum]; ok {
		return router.ErrDuplicateRoute
	}

	// create the route
	t.routes[service][sum] = &route{r, time.Now()}

	if logger.V(logger.DebugLevel, logger.DefaultLogger) {
		logger.Debugf("Router emitting %s for route: %s", router.Create, r.Address)
	}

	// send a route created event
	go t.sendEvent(&router.Event{Type: router.Create, Timestamp: time.Now(), Route: r})

	return nil
}

// Delete deletes the route from the routing table
func (t *table) Delete(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	if _, ok := t.routes[service]; !ok {
		return router.ErrRouteNotFound
	}

	if _, ok := t.routes[service][sum]; !ok {
		return router.ErrRouteNotFound
	}

	// delete the route from the service
	delete(t.routes[service], sum)

	// delete the whole map if there are no routes left
	if len(t.routes[service]) == 0 {
		delete(t.routes, service)
	}

	if logger.V(logger.DebugLevel, logger.DefaultLogger) {
		logger.Debugf("Router emitting %s for route: %s", router.Delete, r.Address)
	}
	go t.sendEvent(&router.Event{Type: router.Delete, Timestamp: time.Now(), Route: r})

	return nil
}

// Update updates routing table with the new route
func (t *table) Update(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	// check if the route destination has any routes in the table
	if _, ok := t.routes[service]; !ok {
		t.routes[service] = make(map[uint64]*route)
	}

	if _, ok := t.routes[service][sum]; !ok {
		// update the route
		t.routes[service][sum] = &route{r, time.Now()}

		if logger.V(logger.DebugLevel, logger.DefaultLogger) {
			logger.Debugf("Router emitting %s for route: %s", router.Update, r.Address)
		}
		go t.sendEvent(&router.Event{Type: router.Update, Timestamp: time.Now(), Route: r})
		return nil
	}

	// just update the route, but dont emit Update event
	t.routes[service][sum] = &route{r, time.Now()}

	return nil
}

// List returns a list of all routes in the table
func (t *table) List() ([]router.Route, error) {
	t.RLock()
	defer t.RUnlock()

	var routes []router.Route
	for _, rmap := range t.routes {
		for _, route := range rmap {
			routes = append(routes, route.route)
		}
	}

	return routes, nil
}

// isMatch checks if the route matches given query options
func isMatch(route router.Route, address, gateway, network, rtr string, strategy router.Strategy) bool {
	// matches the values provided
	match := func(a, b string) bool {
		if a == "*" || b == "*" || a == b {
			return true
		}
		return false
	}

	// a simple struct to hold our values
	type compare struct {
		a string
		b string
	}

	// by default assume we are querying all routes
	link := "*"
	// if AdvertiseLocal change the link query accordingly
	if strategy == router.AdvertiseLocal {
		link = "local"
	}

	// compare the following values
	values := []compare{
		{gateway, route.Gateway},
		{network, route.Network},
		{rtr, route.Router},
		{address, route.Address},
		{link, route.Link},
	}

	for _, v := range values {
		// attempt to match each value
		if !match(v.a, v.b) {
			return false
		}
	}

	return true
}

// filterRoutes finds all the routes for given network and router and returns them
func filterRoutes(routes map[uint64]*route, opts router.QueryOptions) []router.Route {
	address := opts.Address
	gateway := opts.Gateway
	network := opts.Network
	rtr := opts.Router
	strategy := opts.Strategy

	// routeMap stores the routes we're going to advertise
	routeMap := make(map[string][]router.Route)

	for _, rt := range routes {
		// get the actual route
		route := rt.route

		if isMatch(route, address, gateway, network, rtr, strategy) {
			// add matchihg route to the routeMap
			routeKey := route.Service + "@" + route.Network
			// append the first found route to routeMap
			_, ok := routeMap[routeKey]
			if !ok {
				routeMap[routeKey] = append(routeMap[routeKey], route)
				continue
			}

			// if AdvertiseAll, keep appending
			if strategy == router.AdvertiseAll || strategy == router.AdvertiseLocal {
				routeMap[routeKey] = append(routeMap[routeKey], route)
				continue
			}

			// now we're going to find the best routes
			if strategy == router.AdvertiseBest {
				// if the current optimal route metric is higher than routing table route, replace it
				if len(routeMap[routeKey]) > 0 {
					// NOTE: we know that when AdvertiseBest is set, we only ever have one item in current
					if routeMap[routeKey][0].Metric > route.Metric {
						routeMap[routeKey][0] = route
						continue
					}
				}
			}
		}
	}

	var results []router.Route
	for _, route := range routeMap {
		results = append(results, route...)
	}

	return results
}

// Lookup queries routing table and returns all routes that match the lookup query
func (t *table) Query(q ...router.QueryOption) ([]router.Route, error) {
	// create new query options
	opts := router.NewQuery(q...)

	// create a cwslicelist of query results
	results := make([]router.Route, 0, len(t.routes))

	// if No routes are queried, return early
	if opts.Strategy == router.AdvertiseNone {
		return results, nil
	}

	// readAndFilter routes for this service under read lock.
	readAndFilter := func(q router.QueryOptions) ([]router.Route, bool) {
		t.RLock()
		defer t.RUnlock()

		routes, ok := t.routes[q.Service]
		if !ok || len(routes) == 0 {
			return nil, false
		}

		return filterRoutes(routes, q), true
	}

	if opts.Service != "*" {
		// try and load services from the cache
		if routes, ok := readAndFilter(opts); ok {
			return routes, nil
		}

		// lookup the route and try again
		// TODO: move this logic out of the hot path
		// being hammered on queries will require multiple lookups
		routes, err := t.lookup(opts.Service)
		if err != nil {
			return nil, err
		}

		// cache the routes
		t.saveRoutes(opts.Service, routes)

		// try again
		if routes, ok := readAndFilter(opts); ok {
			return routes, nil
		}

		return nil, router.ErrRouteNotFound
	}

	// search through all destinations
	t.RLock()

	for _, routes := range t.routes {
		// filter the routes
		found := filterRoutes(routes, opts)
		// ensure we don't append zero length routes
		if len(found) == 0 {
			continue
		}
		results = append(results, found...)
	}

	t.RUnlock()

	return results, nil
}

// Watch returns routing table entry watcher
func (t *table) Watch(opts ...router.WatchOption) (router.Watcher, error) {
	// by default watch everything
	wopts := router.WatchOptions{
		Service: "*",
	}

	for _, o := range opts {
		o(&wopts)
	}

	w := &tableWatcher{
		id:      uuid.New().String(),
		opts:    wopts,
		resChan: make(chan *router.Event, 10),
		done:    make(chan struct{}),
	}

	// when the watcher is stopped delete it
	go func() {
		<-w.done
		t.Lock()
		delete(t.watchers, w.id)
		t.Unlock()
	}()

	// save the watcher
	t.Lock()
	t.watchers[w.id] = w
	t.Unlock()

	return w, nil
}
<commit_msg>fix deadlock bug<commit_after>package registry

import (
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/micro/go-micro/v3/logger"
	"github.com/micro/go-micro/v3/router"
)

// table is an in-memory routing table
type table struct {
	sync.RWMutex
	// lookup for a service
	lookup func(string) ([]router.Route, error)
	// routes stores service routes
	routes map[string]map[uint64]*route
	// watchers stores table watchers
	watchers map[string]*tableWatcher
}

type route struct {
	route   router.Route
	updated time.Time
}

// newtable creates a new routing table and returns it
func newTable(lookup func(string) ([]router.Route, error), opts ...router.Option) *table {
	return &table{
		lookup:   lookup,
		routes:   make(map[string]map[uint64]*route),
		watchers: make(map[string]*tableWatcher),
	}
}

// pruneRoutes will prune routes older than the time specified
func (t *table) pruneRoutes(olderThan time.Duration) {
	var routes []router.Route

	t.Lock()

	// search for all the routes
	for _, routeList := range t.routes {
		for _, r := range routeList {
			// if any route is older than
			if time.Since(r.updated).Seconds() > olderThan.Seconds() {
				routes = append(routes, r.route)
			}
		}
	}

	t.Unlock()

	// delete the routes we've found
	for _, route := range routes {
		t.Delete(route)
	}
}

// deleteService removes the entire service
func (t *table) deleteService(service, network string) {
	t.Lock()
	defer t.Unlock()

	routes, ok := t.routes[service]
	if !ok {
		return
	}

	// delete the routes for the service
	for hash, rt := range routes {
		// TODO: check if this causes a problem
		// with * in the network if that is a thing
		// or blank strings
		if rt.route.Network != network {
			continue
		}
		delete(routes, hash)
	}

	// delete the map for the service if its empty
	if len(routes) == 0 {
		delete(t.routes, service)
		return
	}

	// save the routes
	t.routes[service] = routes
}

// saveRoutes completely replaces the routes for a service
func (t *table) saveRoutes(service string, routes []router.Route) {
	t.Lock()
	defer t.Unlock()

	// delete old routes
	delete(t.routes, service)
	// make new map
	t.routes[service] = make(map[uint64]*route)

	// iterate through new routes and save
	for _, rt := range routes {
		// save the new route
		t.routes[service][rt.Hash()] = &route{rt, time.Now()}
	}
}

// sendEvent sends events to all subscribed watchers
func (t *table) sendEvent(e *router.Event) {
	t.RLock()
	defer t.RUnlock()

	if len(e.Id) == 0 {
		e.Id = uuid.New().String()
	}

	for _, w := range t.watchers {
		select {
		case w.resChan <- e:
		case <-w.done:
		// don't block forever
		case <-time.After(time.Second):
		}
	}
}

// Create creates new route in the routing table
func (t *table) Create(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	// check if there are any routes in the table for the route destination
	if _, ok := t.routes[service]; !ok {
		t.routes[service] = make(map[uint64]*route)
	}

	// add new route to the table for the route destination
	if _, ok := t.routes[service][sum]; ok {
		return router.ErrDuplicateRoute
	}

	// create the route
	t.routes[service][sum] = &route{r, time.Now()}

	if logger.V(logger.DebugLevel, logger.DefaultLogger) {
		logger.Debugf("Router emitting %s for route: %s", router.Create, r.Address)
	}

	// send a route created event
	go t.sendEvent(&router.Event{Type: router.Create, Timestamp: time.Now(), Route: r})

	return nil
}

// Delete deletes the route from the routing table
func (t *table) Delete(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	if _, ok := t.routes[service]; !ok {
		return router.ErrRouteNotFound
	}

	if _, ok := t.routes[service][sum]; !ok {
		return router.ErrRouteNotFound
	}

	// delete the route from the service
	delete(t.routes[service], sum)

	// delete the whole map if there are no routes left
	if len(t.routes[service]) == 0 {
		delete(t.routes, service)
	}

	if logger.V(logger.DebugLevel, logger.DefaultLogger) {
		logger.Debugf("Router emitting %s for route: %s", router.Delete, r.Address)
	}
	go t.sendEvent(&router.Event{Type: router.Delete, Timestamp: time.Now(), Route: r})

	return nil
}

// Update updates routing table with the new route
func (t *table) Update(r router.Route) error {
	service := r.Service
	sum := r.Hash()

	t.Lock()
	defer t.Unlock()

	// check if the route destination has any routes in the table
	if _, ok := t.routes[service]; !ok {
		t.routes[service] = make(map[uint64]*route)
	}

	if _, ok := t.routes[service][sum]; !ok {
		// update the route
		t.routes[service][sum] = &route{r, time.Now()}

		if logger.V(logger.DebugLevel, logger.DefaultLogger) {
			logger.Debugf("Router emitting %s for route: %s", router.Update, r.Address)
		}
		go t.sendEvent(&router.Event{Type: router.Update, Timestamp: time.Now(), Route: r})
		return nil
	}

	// just update the route, but dont emit Update event
	t.routes[service][sum] = &route{r, time.Now()}

	return nil
}

// List returns a list of all routes in the table
func (t *table) List() ([]router.Route, error) {
	t.RLock()
	defer t.RUnlock()

	var routes []router.Route
	for _, rmap := range t.routes {
		for _, route := range rmap {
			routes = append(routes, route.route)
		}
	}

	return routes, nil
}

// isMatch checks if the route matches given query options
func isMatch(route router.Route, address, gateway, network, rtr string, strategy router.Strategy) bool {
	// matches the values provided
	match := func(a, b string) bool {
		if a == "*" || b == "*" || a == b {
			return true
		}
		return false
	}

	// a simple struct to hold our values
	type compare struct {
		a string
		b string
	}

	// by default assume we are querying all routes
	link := "*"
	// if AdvertiseLocal change the link query accordingly
	if strategy == router.AdvertiseLocal {
		link = "local"
	}

	// compare the following values
	values := []compare{
		{gateway, route.Gateway},
		{network, route.Network},
		{rtr, route.Router},
		{address, route.Address},
		{link, route.Link},
	}

	for _, v := range values {
		// attempt to match each value
		if !match(v.a, v.b) {
			return false
		}
	}

	return true
}

// filterRoutes finds all the routes for given network and router and returns them
func filterRoutes(routes map[uint64]*route, opts router.QueryOptions) []router.Route {
	address := opts.Address
	gateway := opts.Gateway
	network := opts.Network
	rtr := opts.Router
	strategy := opts.Strategy

	// routeMap stores the routes we're going to advertise
	routeMap := make(map[string][]router.Route)

	for _, rt := range routes {
		// get the actual route
		route := rt.route

		if isMatch(route, address, gateway, network, rtr, strategy) {
			// add matchihg route to the routeMap
			routeKey := route.Service + "@" + route.Network
			// append the first found route to routeMap
			_, ok := routeMap[routeKey]
			if !ok {
				routeMap[routeKey] = append(routeMap[routeKey], route)
				continue
			}

			// if AdvertiseAll, keep appending
			if strategy == router.AdvertiseAll || strategy == router.AdvertiseLocal {
				routeMap[routeKey] = append(routeMap[routeKey], route)
				continue
			}

			// now we're going to find the best routes
			if strategy == router.AdvertiseBest {
				// if the current optimal route metric is higher than routing table route, replace it
				if len(routeMap[routeKey]) > 0 {
					// NOTE: we know that when AdvertiseBest is set, we only ever have one item in current
					if routeMap[routeKey][0].Metric > route.Metric {
						routeMap[routeKey][0] = route
						continue
					}
				}
			}
		}
	}

	var results []router.Route
	for _, route := range routeMap {
		results = append(results, route...)
	}

	return results
}

// Lookup queries routing table and returns all routes that match the lookup query
func (t *table) Query(q ...router.QueryOption) ([]router.Route, error) {
	// create new query options
	opts := router.NewQuery(q...)

	// create a cwslicelist of query results
	results := make([]router.Route, 0, len(t.routes))

	// if No routes are queried, return early
	if opts.Strategy == router.AdvertiseNone {
		return results, nil
	}

	// readAndFilter routes for this service under read lock.
	readAndFilter := func(q router.QueryOptions) ([]router.Route, bool) {
		t.RLock()
		defer t.RUnlock()

		routes, ok := t.routes[q.Service]
		if !ok || len(routes) == 0 {
			return nil, false
		}

		return filterRoutes(routes, q), true
	}

	if opts.Service != "*" {
		// try and load services from the cache
		if routes, ok := readAndFilter(opts); ok {
			return routes, nil
		}

		// lookup the route and try again
		// TODO: move this logic out of the hot path
		// being hammered on queries will require multiple lookups
		routes, err := t.lookup(opts.Service)
		if err != nil {
			return nil, err
		}

		// cache the routes
		t.saveRoutes(opts.Service, routes)

		// try again
		if routes, ok := readAndFilter(opts); ok {
			return routes, nil
		}

		return nil, router.ErrRouteNotFound
	}

	// search through all destinations
	t.RLock()

	for _, routes := range t.routes {
		// filter the routes
		found := filterRoutes(routes, opts)
		// ensure we don't append zero length routes
		if len(found) == 0 {
			continue
		}
		results = append(results, found...)
	}

	t.RUnlock()

	return results, nil
}

// Watch returns routing table entry watcher
func (t *table) Watch(opts ...router.WatchOption) (router.Watcher, error) {
	// by default watch everything
	wopts := router.WatchOptions{
		Service: "*",
	}

	for _, o := range opts {
		o(&wopts)
	}

	w := &tableWatcher{
		id:      uuid.New().String(),
		opts:    wopts,
		resChan: make(chan *router.Event, 10),
		done:    make(chan struct{}),
	}

	// when the watcher is stopped delete it
	go func() {
		<-w.done
		t.Lock()
		delete(t.watchers, w.id)
		t.Unlock()
	}()

	// save the watcher
	t.Lock()
	t.watchers[w.id] = w
	t.Unlock()

	return w, nil
}
<|endoftext|>


<commit_before>package project_test

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

	"github.com/getantibody/antibody/project"
	"github.com/stretchr/testify/assert"
)

func TestDownloadAllKinds(t *testing.T) {
	assert := assert.New(t)
	urls := []string{
		"caarlos0/ports",
		"http://github.com/caarlos0/ports",
		"http://github.com/caarlos0/ports.git",
		"https://github.com/caarlos0/ports",
		"https://github.com/caarlos0/ports.git",
		"git://github.com/caarlos0/ports.git",
		"ssh://git@github.com/caarlos0/ports.git",
		"git@github.com:caarlos0/ports.git",
	}
	for _, url := range urls {
		home := home()
		defer os.RemoveAll(home)
		assert.NoError(
			project.NewGit(home, url).Download(),
			"Repo "+url+" failed to download",
		)
	}
}

func TestDownloadAnotherBranch(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	assert.NoError(project.NewGit(home, "caarlos0/jvm branch:gh-pages").Download())
}

func TestDownloadAndUpdate(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.NoError(repo.Download())
	assert.NoError(repo.Update())
}

func TestUpdateNonExistentLocalRepo(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.Error(repo.Update())
}

func TestDownloadNonExistenRepo(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "doesn-not-exist-really branch:also-nope")
	assert.Error(repo.Download())
}

func TestDownloadMultipleTimes(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.NoError(repo.Download())
	assert.NoError(repo.Download())
}

func TestDownloadFolderNaming(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.Equal(
		home+"/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-ports",
		repo.Folder(),
	)
}

func home() string {
	home, err := ioutil.TempDir(os.TempDir(), "antibody")
	if err != nil {
		panic(err.Error())
	}
	return home
}
<commit_msg>commenting out a line, temp<commit_after>package project_test

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

	"github.com/getantibody/antibody/project"
	"github.com/stretchr/testify/assert"
)

func TestDownloadAllKinds(t *testing.T) {
	assert := assert.New(t)
	urls := []string{
		"caarlos0/ports",
		"http://github.com/caarlos0/ports",
		"http://github.com/caarlos0/ports.git",
		"https://github.com/caarlos0/ports",
		"https://github.com/caarlos0/ports.git",
		"git://github.com/caarlos0/ports.git",
		// "ssh://git@github.com/caarlos0/ports.git", TODO
		"git@github.com:caarlos0/ports.git",
	}
	for _, url := range urls {
		home := home()
		defer os.RemoveAll(home)
		assert.NoError(
			project.NewGit(home, url).Download(),
			"Repo "+url+" failed to download",
		)
	}
}

func TestDownloadAnotherBranch(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	assert.NoError(project.NewGit(home, "caarlos0/jvm branch:gh-pages").Download())
}

func TestDownloadAndUpdate(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.NoError(repo.Download())
	assert.NoError(repo.Update())
}

func TestUpdateNonExistentLocalRepo(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.Error(repo.Update())
}

func TestDownloadNonExistenRepo(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "doesn-not-exist-really branch:also-nope")
	assert.Error(repo.Download())
}

func TestDownloadMultipleTimes(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.NoError(repo.Download())
	assert.NoError(repo.Download())
}

func TestDownloadFolderNaming(t *testing.T) {
	assert := assert.New(t)
	home := home()
	defer os.RemoveAll(home)
	repo := project.NewGit(home, "caarlos0/ports")
	assert.Equal(
		home+"/https-COLON--SLASH--SLASH-github.com-SLASH-caarlos0-SLASH-ports",
		repo.Folder(),
	)
}

func home() string {
	home, err := ioutil.TempDir(os.TempDir(), "antibody")
	if err != nil {
		panic(err.Error())
	}
	return home
}
<|endoftext|>