<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-etcd/etcd"
	"github.com/ryanuber/columnize"

	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	ETCD_ENTRY_ALREADY_EXISTS = 105
	ETCD_ENTRY_NOT_EXISTS     = 100
)

var config struct {
	Host       string `toml:"host"`
	PrivateKey string `toml:"private_key"`
}

func ensureEtcClient(c *cli.Context) *etcd.Client {
	machines := strings.Split(c.GlobalString("etcd"), ",")
	ectdClient := etcd.NewClient(machines)
	return ectdClient
}

func ensureAppParam(c *cli.Context, command string) string {
	app := c.Args().First()
	if app == "" {
		println("ERROR: app name missing")
		cli.ShowCommandHelp(c, command)
		os.Exit(1)
	}
	return app
}

func appList(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	path := "/" + c.GlobalString("env") + "/" + c.GlobalString("pool")

	entries, err := etcdClient.Get(path, false, false)
	if err != nil && err.(*etcd.EtcdError).ErrorCode == ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Environment (%s) or pool (%s) does not exist.\n",
			c.GlobalString("env"), c.GlobalString("pool"))
		os.Exit(1)
	}

	if err != nil {
		fmt.Printf("ERROR: Could not find registered apps: %s\n", err)
		os.Exit(1)
	}

	columns := []string{"NAME | CONFIGURED | VERSION"}
	for _, entry := range entries.Node.Nodes {
		name := filepath.Base(entry.Key)
		// skip runtime host entry
		if name == "hosts" {
			continue
		}

		environmentConfigured := false
		_, err := etcdClient.Get(filepath.Join(path, name, "environment"), false, false)
		if err == nil {
			environmentConfigured = true
		}

		versionDeployed := ""
		version, err := etcdClient.Get(filepath.Join(path, name, "version"), false, false)
		if err == nil {
			versionDeployed = version.Node.Value
		}

		columns = append(columns, strings.Join([]string{
			name, strconv.FormatBool(environmentConfigured),
			versionDeployed}, " | "))
	}
	output, _ := columnize.SimpleFormat(columns)
	fmt.Println(output)
}

func appDeploy(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "app:deploy")

	version := c.Args().Tail()[0]
	if version == "" {
		println("ERROR: app name missing")
		cli.ShowCommandHelp(c, "config")
		os.Exit(1)
	}

	_, err := etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/version", version, 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store version: %s\n", err)
		os.Exit(1)
	}
}

func configList(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config")

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	if err != nil && err.(*etcd.EtcdError).ErrorCode == ETCD_ENTRY_NOT_EXISTS {
		return
	}

	var env map[string]string
	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for k, v := range env {
		fmt.Printf("%s=%s\n", k, v)
	}
}

func configSet(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:set")

	var env map[string]string

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	if err == nil || err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		err = json.Unmarshal([]byte(resp.Node.Value), &env)
		if err != nil {
			fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
			os.Exit(1)
		}
	}

	for _, arg := range c.Args().Tail() {
		if !strings.Contains(arg, "=") {
			fmt.Printf("ERROR: bad config variable format: %s\n", arg)
			cli.ShowCommandHelp(c, "config")
			os.Exit(1)

		}
		values := strings.Split(arg, "=")
		env[strings.ToUpper(values[0])] = values[1]
	}

	serialized, err := json.Marshal(env)
	if err != nil {
		fmt.Printf("ERROR: Could not marshall config: %s\n", err)
		os.Exit(1)
	}

	resp, err = etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", string(serialized), 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store config: %s\n", err)
		os.Exit(1)
	}
}

func configUnset(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:unset")

	env := map[string]string{}

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for _, arg := range c.Args().Tail() {
		delete(env, strings.ToUpper(arg))
	}

	serialized, err := json.Marshal(env)
	if err != nil {
		fmt.Printf("ERROR: Could not marshall config: %s\n", err)
		os.Exit(1)
	}

	resp, err = etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", string(serialized), 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store config: %s\n", err)
		os.Exit(1)
	}
}

func configGet(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:get")

	env := map[string]string{}

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for _, arg := range c.Args().Tail() {
		fmt.Printf("%s=%s\n", strings.ToUpper(arg), env[strings.ToUpper(arg)])
	}
}

func login(c *cli.Context) {

	if c.Args().First() == "" {
		println("ERROR: host missing")
		cli.ShowCommandHelp(c, "login")
		os.Exit(1)
	}

	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}

	configDir := filepath.Join(currentUser.HomeDir, ".galaxy")
	_, err = os.Stat(configDir)
	if err != nil && os.IsNotExist(err) {
		os.Mkdir(configDir, 0700)
	}
	availableKeys := findSshKeys(currentUser.HomeDir)

	if len(availableKeys) == 0 {
		fmt.Printf("ERROR: No SSH private keys found.  Create one first.\n")
		os.Exit(1)
	}

	for i, key := range availableKeys {
		fmt.Printf("%d) %s\n", i, key)
	}

	fmt.Printf("Select private key to use [0]: ")
	var i int
	fmt.Scanf("%d", &i)

	if i < 0 || i > len(availableKeys) {
		i = 0
	}
	fmt.Printf("Using %s\n", availableKeys[i])

	config.Host = c.Args().First()
	config.PrivateKey = availableKeys[i]

	configFile, err := os.Create(filepath.Join(configDir, "galaxy.toml"))
	if err != nil {
		fmt.Printf("ERROR: Unable to create config file: %s\n", err)
		os.Exit(1)
	}
	defer configFile.Close()

	encoder := toml.NewEncoder(configFile)
	encoder.Encode(config)
	configFile.WriteString("\n")
	fmt.Printf("Login sucessful")
}

func logout(c *cli.Context) {
	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}
	configFile := filepath.Join(currentUser.HomeDir, ".galaxy", "galaxy.toml")

	_, err = os.Stat(configFile)
	if err == nil {
		err = os.Remove(configFile)
		if err != nil {
			fmt.Printf("ERROR: Unable to logout: %s\n", err)
			os.Exit(1)
		}
	}
	fmt.Printf("Logout sucessful\n")
}

func runRemote() {
	Sshcmd(config.Host, "galaxy "+strings.Join(os.Args[1:], " "), false, false)
}

func loadConfig() {

	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}
	configFile := filepath.Join(currentUser.HomeDir, ".galaxy", "galaxy.toml")

	_, err = os.Stat(configFile)
	if err == nil {
		if _, err := toml.DecodeFile(configFile, &config); err != nil {
			fmt.Printf("ERROR: Unable to logout: %s\n", err)
			os.Exit(1)
		}
	}

}
func main() {

	loadConfig()
	if config.Host != "" {
		runRemote()
		return
	}

	app := cli.NewApp()
	app.Name = "galaxy"
	app.Usage = "galaxy cli"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "etcd", Value: "http://127.0.0.1:4001", Usage: "host:port[,host:port,..]"},
		cli.StringFlag{Name: "env", Value: "dev", Usage: "environment (dev, test, prod, etc.)"},
		cli.StringFlag{Name: "pool", Value: "web", Usage: "pool (web, worker, etc.)"},
	}

	app.Commands = []cli.Command{
		{
			Name:        "login",
			Usage:       "login to a controller",
			Action:      login,
			Description: "login host[:port]",
		},
		{
			Name:        "logout",
			Usage:       "logout off a controller",
			Action:      logout,
			Description: "logout",
		},
		{
			Name:        "app",
			Usage:       "list the apps currently created",
			Action:      appList,
			Description: "app",
		},
		{
			Name:        "app:deploy",
			Usage:       "deploy a new version of an app",
			Action:      appDeploy,
			Description: "config <app> <version>",
		},
		{
			Name:        "config",
			Usage:       "list the config values for an app",
			Action:      configList,
			Description: "config <app>",
		},
		{
			Name:        "config:set",
			Usage:       "set one or more configuration variables",
			Action:      configSet,
			Description: "config <app> KEY=VALUE[,KEY=VALUE,..]",
		},
		{
			Name:        "config:unset",
			Usage:       "unset one or more configuration variables",
			Action:      configUnset,
			Description: "config <app> KEY[ KEY, etc..]",
		},
		{
			Name:   "config:get",
			Usage:  "display the config value for an app",
			Action: configGet,
		},
	}
	app.Run(os.Args)
}
<commit_msg>Add create app command<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/codegangsta/cli"
	"github.com/coreos/go-etcd/etcd"
	"github.com/ryanuber/columnize"

	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
)

const (
	ETCD_ENTRY_ALREADY_EXISTS = 105
	ETCD_ENTRY_NOT_EXISTS     = 100
)

var config struct {
	Host       string `toml:"host"`
	PrivateKey string `toml:"private_key"`
}

func ensureEtcClient(c *cli.Context) *etcd.Client {
	machines := strings.Split(c.GlobalString("etcd"), ",")
	ectdClient := etcd.NewClient(machines)
	return ectdClient
}

func ensureAppParam(c *cli.Context, command string) string {
	app := c.Args().First()
	if app == "" {
		println("ERROR: app name missing")
		cli.ShowCommandHelp(c, command)
		os.Exit(1)
	}
	return app
}

func appList(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	path := "/" + c.GlobalString("env") + "/" + c.GlobalString("pool")

	entries, err := etcdClient.Get(path, false, false)
	if err != nil && err.(*etcd.EtcdError).ErrorCode == ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Environment (%s) or pool (%s) does not exist.\n",
			c.GlobalString("env"), c.GlobalString("pool"))
		os.Exit(1)
	}

	if err != nil {
		fmt.Printf("ERROR: Could not find registered apps: %s\n", err)
		os.Exit(1)
	}

	columns := []string{"NAME | CONFIGURED | VERSION"}
	for _, entry := range entries.Node.Nodes {
		name := filepath.Base(entry.Key)
		// skip runtime host entry
		if name == "hosts" {
			continue
		}

		environmentConfigured := false
		_, err := etcdClient.Get(filepath.Join(path, name, "environment"), false, false)
		if err == nil {
			environmentConfigured = true
		}

		versionDeployed := ""
		version, err := etcdClient.Get(filepath.Join(path, name, "version"), false, false)
		if err == nil {
			versionDeployed = version.Node.Value
		}

		columns = append(columns, strings.Join([]string{
			name, strconv.FormatBool(environmentConfigured),
			versionDeployed}, " | "))
	}
	output, _ := columnize.SimpleFormat(columns)
	fmt.Println(output)
}

func appCreate(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "app:create")

	_, err := etcdClient.CreateDir("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app, 0)
	if err != nil {
		fmt.Printf("ERROR: Could not create app: %s\n", err)
		os.Exit(1)
	}
}

func appDeploy(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "app:deploy")

	version := ""
	if len(c.Args().Tail()) == 1 {
		version = c.Args().Tail()[0]
	}

	if version == "" {
		println("ERROR: version missing")
		cli.ShowCommandHelp(c, "app:deploy")
		os.Exit(1)
	}

	_, err := etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/version", version, 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store version: %s\n", err)
		os.Exit(1)
	}
}

func configList(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config")

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	if err != nil && err.(*etcd.EtcdError).ErrorCode == ETCD_ENTRY_NOT_EXISTS {
		return
	}

	var env map[string]string
	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for k, v := range env {
		fmt.Printf("%s=%s\n", k, v)
	}
}

func configSet(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:set")

	var env map[string]string

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	if err == nil || err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		err = json.Unmarshal([]byte(resp.Node.Value), &env)
		if err != nil {
			fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
			os.Exit(1)
		}
	}

	for _, arg := range c.Args().Tail() {
		if !strings.Contains(arg, "=") {
			fmt.Printf("ERROR: bad config variable format: %s\n", arg)
			cli.ShowCommandHelp(c, "config")
			os.Exit(1)

		}
		values := strings.Split(arg, "=")
		env[strings.ToUpper(values[0])] = values[1]
	}

	serialized, err := json.Marshal(env)
	if err != nil {
		fmt.Printf("ERROR: Could not marshall config: %s\n", err)
		os.Exit(1)
	}

	resp, err = etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", string(serialized), 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store config: %s\n", err)
		os.Exit(1)
	}
}

func configUnset(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:unset")

	env := map[string]string{}

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for _, arg := range c.Args().Tail() {
		delete(env, strings.ToUpper(arg))
	}

	serialized, err := json.Marshal(env)
	if err != nil {
		fmt.Printf("ERROR: Could not marshall config: %s\n", err)
		os.Exit(1)
	}

	resp, err = etcdClient.Set("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", string(serialized), 0)
	if err != nil {
		fmt.Printf("ERROR: Could not store config: %s\n", err)
		os.Exit(1)
	}
}

func configGet(c *cli.Context) {

	etcdClient := ensureEtcClient(c)
	app := ensureAppParam(c, "config:get")

	env := map[string]string{}

	resp, err := etcdClient.Get("/"+c.GlobalString("env")+"/"+c.GlobalString("pool")+"/"+app+"/environment", true, true)
	if err != nil && err.(*etcd.EtcdError).ErrorCode != ETCD_ENTRY_NOT_EXISTS {
		fmt.Printf("ERROR: Could not connect to etcd: %s\n", err)
		os.Exit(1)
	}

	err = json.Unmarshal([]byte(resp.Node.Value), &env)
	if err != nil {
		fmt.Printf("ERROR: Could not unmarshall config: %s\n", err)
		os.Exit(1)
	}

	for _, arg := range c.Args().Tail() {
		fmt.Printf("%s=%s\n", strings.ToUpper(arg), env[strings.ToUpper(arg)])
	}
}

func login(c *cli.Context) {

	if c.Args().First() == "" {
		println("ERROR: host missing")
		cli.ShowCommandHelp(c, "login")
		os.Exit(1)
	}

	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}

	configDir := filepath.Join(currentUser.HomeDir, ".galaxy")
	_, err = os.Stat(configDir)
	if err != nil && os.IsNotExist(err) {
		os.Mkdir(configDir, 0700)
	}
	availableKeys := findSshKeys(currentUser.HomeDir)

	if len(availableKeys) == 0 {
		fmt.Printf("ERROR: No SSH private keys found.  Create one first.\n")
		os.Exit(1)
	}

	for i, key := range availableKeys {
		fmt.Printf("%d) %s\n", i, key)
	}

	fmt.Printf("Select private key to use [0]: ")
	var i int
	fmt.Scanf("%d", &i)

	if i < 0 || i > len(availableKeys) {
		i = 0
	}
	fmt.Printf("Using %s\n", availableKeys[i])

	config.Host = c.Args().First()
	config.PrivateKey = availableKeys[i]

	configFile, err := os.Create(filepath.Join(configDir, "galaxy.toml"))
	if err != nil {
		fmt.Printf("ERROR: Unable to create config file: %s\n", err)
		os.Exit(1)
	}
	defer configFile.Close()

	encoder := toml.NewEncoder(configFile)
	encoder.Encode(config)
	configFile.WriteString("\n")
	fmt.Printf("Login sucessful")
}

func logout(c *cli.Context) {
	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}
	configFile := filepath.Join(currentUser.HomeDir, ".galaxy", "galaxy.toml")

	_, err = os.Stat(configFile)
	if err == nil {
		err = os.Remove(configFile)
		if err != nil {
			fmt.Printf("ERROR: Unable to logout: %s\n", err)
			os.Exit(1)
		}
	}
	fmt.Printf("Logout sucessful\n")
}

func runRemote() {
	Sshcmd(config.Host, "galaxy "+strings.Join(os.Args[1:], " "), false, false)
}

func loadConfig() {

	currentUser, err := user.Current()
	if err != nil {
		fmt.Printf("ERROR: Unable to determine current user: %s\n", err)
		os.Exit(1)
	}
	configFile := filepath.Join(currentUser.HomeDir, ".galaxy", "galaxy.toml")

	_, err = os.Stat(configFile)
	if err == nil {
		if _, err := toml.DecodeFile(configFile, &config); err != nil {
			fmt.Printf("ERROR: Unable to logout: %s\n", err)
			os.Exit(1)
		}
	}

}
func main() {

	loadConfig()
	if config.Host != "" {
		runRemote()
		return
	}

	app := cli.NewApp()
	app.Name = "galaxy"
	app.Usage = "galaxy cli"
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "etcd", Value: "http://127.0.0.1:4001", Usage: "host:port[,host:port,..]"},
		cli.StringFlag{Name: "env", Value: "dev", Usage: "environment (dev, test, prod, etc.)"},
		cli.StringFlag{Name: "pool", Value: "web", Usage: "pool (web, worker, etc.)"},
	}

	app.Commands = []cli.Command{
		{
			Name:        "login",
			Usage:       "login to a controller",
			Action:      login,
			Description: "login host[:port]",
		},
		{
			Name:        "logout",
			Usage:       "logout off a controller",
			Action:      logout,
			Description: "logout",
		},
		{
			Name:        "app",
			Usage:       "list the apps currently created",
			Action:      appList,
			Description: "app",
		},
		{
			Name:        "app:create",
			Usage:       "create a new app",
			Action:      appCreate,
			Description: "app:create",
		},
		{
			Name:        "app:deploy",
			Usage:       "deploy a new version of an app",
			Action:      appDeploy,
			Description: "config <app> <version>",
		},
		{
			Name:        "config",
			Usage:       "list the config values for an app",
			Action:      configList,
			Description: "config <app>",
		},
		{
			Name:        "config:set",
			Usage:       "set one or more configuration variables",
			Action:      configSet,
			Description: "config <app> KEY=VALUE[,KEY=VALUE,..]",
		},
		{
			Name:        "config:unset",
			Usage:       "unset one or more configuration variables",
			Action:      configUnset,
			Description: "config <app> KEY[ KEY, etc..]",
		},
		{
			Name:   "config:get",
			Usage:  "display the config value for an app",
			Action: configGet,
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gc

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/goharbor/harbor/src/common"
	common_http "github.com/goharbor/harbor/src/common/http"
	"github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/registryctl"
	"github.com/goharbor/harbor/src/common/utils"
	reg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/jobservice/env"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/registryctl/client"
)

const (
	dialConnectionTimeout = 30 * time.Second
	dialReadTimeout       = time.Minute + 10*time.Second
	dialWriteTimeout      = 10 * time.Second
	blobPrefix            = "blobs::*"
	repoPrefix            = "repository::*"
)

// GarbageCollector is the struct to run registry's garbage collection
type GarbageCollector struct {
	registryCtlClient client.Client
	logger            logger.Interface
	coreclient        *common_http.Client
	CoreURL           string
	insecure          bool
	redisURL          string
}

// MaxFails implements the interface in job/Interface
func (gc *GarbageCollector) MaxFails() uint {
	return 1
}

// ShouldRetry implements the interface in job/Interface
func (gc *GarbageCollector) ShouldRetry() bool {
	return false
}

// Validate implements the interface in job/Interface
func (gc *GarbageCollector) Validate(params map[string]interface{}) error {
	return nil
}

// Run implements the interface in job/Interface
func (gc *GarbageCollector) Run(ctx env.JobContext, params map[string]interface{}) error {
	if err := gc.init(ctx, params); err != nil {
		return err
	}
	readOnlyCur, err := gc.getReadOnly()
	if err != nil {
		return err
	}
	if readOnlyCur != true {
		if err := gc.setReadOnly(true); err != nil {
			return err
		}
	}
	defer gc.setReadOnly(readOnlyCur)
	if err := gc.registryCtlClient.Health(); err != nil {
		gc.logger.Errorf("failed to start gc as registry controller is unreachable: %v", err)
		return err
	}
	gc.logger.Infof("start to run gc in job.")
	gcr, err := gc.registryCtlClient.StartGC()
	if err != nil {
		gc.logger.Errorf("failed to get gc result: %v", err)
		return err
	}
	if err := gc.cleanCache(); err != nil {
		return err
	}
	gc.logger.Infof("GC results: status: %t, message: %s, start: %s, end: %s.", gcr.Status, gcr.Msg, gcr.StartTime, gcr.EndTime)
	gc.logger.Infof("success to run gc in job.")
	return nil
}

func (gc *GarbageCollector) init(ctx env.JobContext, params map[string]interface{}) error {
	registryctl.Init()
	gc.registryCtlClient = registryctl.RegistryCtlClient
	gc.logger = ctx.GetLogger()
	cred := auth.NewSecretAuthorizer(os.Getenv("JOBSERVICE_SECRET"))
	gc.insecure = false
	gc.coreclient = common_http.NewClient(&http.Client{
		Transport: reg.GetHTTPTransport(gc.insecure),
	}, cred)
	errTpl := "Failed to get required property: %s"
	if v, ok := ctx.Get(common.CoreURL); ok && len(v.(string)) > 0 {
		gc.CoreURL = v.(string)
	} else {
		return fmt.Errorf(errTpl, common.CoreURL)
	}
	gc.redisURL = params["redis_url_reg"].(string)
	return nil
}

func (gc *GarbageCollector) getReadOnly() (bool, error) {
	cfgs := map[string]interface{}{}
	if err := gc.coreclient.Get(fmt.Sprintf("%s/api/configs", gc.CoreURL), &cfgs); err != nil {
		return false, err
	}
	return utils.SafeCastBool(cfgs[common.ReadOnly]), nil
}

func (gc *GarbageCollector) setReadOnly(switcher bool) error {
	if err := gc.coreclient.Put(fmt.Sprintf("%s/api/configurations", gc.CoreURL), struct {
		ReadOnly bool `json:"read_only"`
	}{
		ReadOnly: switcher,
	}); err != nil {
		gc.logger.Errorf("failed to send readonly request to %s: %v", gc.CoreURL, err)
		return err
	}
	gc.logger.Info("the readonly request has been sent successfully")
	return nil
}

// cleanCache is to clean the registry cache for GC.
// To do this is because the issue https://github.com/docker/distribution/issues/2094
func (gc *GarbageCollector) cleanCache() error {

	con, err := redis.DialURL(
		gc.redisURL,
		redis.DialConnectTimeout(dialConnectionTimeout),
		redis.DialReadTimeout(dialReadTimeout),
		redis.DialWriteTimeout(dialWriteTimeout),
	)

	if err != nil {
		gc.logger.Errorf("failed to connect to redis %v", err)
		return err
	}
	defer con.Close()

	// clean all keys in registry redis DB.

	// sample of keys in registry redis:
	// 1) "blobs::sha256:1a6fd470b9ce10849be79e99529a88371dff60c60aab424c077007f6979b4812"
	// 2) "repository::library/hello-world::blobs::sha256:4ab4c602aa5eed5528a6620ff18a1dc4faef0e1ab3a5eddeddb410714478c67f"
	err = delKeys(con, blobPrefix)
	if err != nil {
		gc.logger.Errorf("failed to clean registry cache %v, pattern blobs::*", err)
		return err
	}
	err = delKeys(con, repoPrefix)
	if err != nil {
		gc.logger.Errorf("failed to clean registry cache %v, pattern repository::*", err)
		return err
	}

	return nil
}

func delKeys(con redis.Conn, pattern string) error {
	iter := 0
	keys := []string{}
	for {
		arr, err := redis.Values(con.Do("SCAN", iter, "MATCH", pattern))
		if err != nil {
			return fmt.Errorf("error retrieving '%s' keys", pattern)
		}
		iter, err := redis.Int(arr[0], nil)
		if err != nil {
			return fmt.Errorf("unexpected type for Int, got type %T", err)
		}
		k, err := redis.Strings(arr[1], nil)
		if err != nil {
			return fmt.Errorf("converts an array command reply to a []string %v", err)
		}
		keys = append(keys, k...)

		if iter == 0 {
			break
		}
	}
	for _, key := range keys {
		_, err := con.Do("DEL", key)
		if err != nil {
			return fmt.Errorf("failed to clean registry cache %v", err)
		}
	}
	return nil
}
<commit_msg>Fix per comments by wk<commit_after>// Copyright Project Harbor Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gc

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/goharbor/harbor/src/common"
	common_http "github.com/goharbor/harbor/src/common/http"
	"github.com/goharbor/harbor/src/common/http/modifier/auth"
	"github.com/goharbor/harbor/src/common/registryctl"
	"github.com/goharbor/harbor/src/common/utils"
	reg "github.com/goharbor/harbor/src/common/utils/registry"
	"github.com/goharbor/harbor/src/jobservice/env"
	"github.com/goharbor/harbor/src/jobservice/logger"
	"github.com/goharbor/harbor/src/registryctl/client"
)

const (
	dialConnectionTimeout = 30 * time.Second
	dialReadTimeout       = time.Minute + 10*time.Second
	dialWriteTimeout      = 10 * time.Second
	blobPrefix            = "blobs::*"
	repoPrefix            = "repository::*"
)

// GarbageCollector is the struct to run registry's garbage collection
type GarbageCollector struct {
	registryCtlClient client.Client
	logger            logger.Interface
	coreclient        *common_http.Client
	CoreURL           string
	insecure          bool
	redisURL          string
}

// MaxFails implements the interface in job/Interface
func (gc *GarbageCollector) MaxFails() uint {
	return 1
}

// ShouldRetry implements the interface in job/Interface
func (gc *GarbageCollector) ShouldRetry() bool {
	return false
}

// Validate implements the interface in job/Interface
func (gc *GarbageCollector) Validate(params map[string]interface{}) error {
	return nil
}

// Run implements the interface in job/Interface
func (gc *GarbageCollector) Run(ctx env.JobContext, params map[string]interface{}) error {
	if err := gc.init(ctx, params); err != nil {
		return err
	}
	readOnlyCur, err := gc.getReadOnly()
	if err != nil {
		return err
	}
	if readOnlyCur != true {
		if err := gc.setReadOnly(true); err != nil {
			return err
		}
		defer gc.setReadOnly(readOnlyCur)
	}
	if err := gc.registryCtlClient.Health(); err != nil {
		gc.logger.Errorf("failed to start gc as registry controller is unreachable: %v", err)
		return err
	}
	gc.logger.Infof("start to run gc in job.")
	gcr, err := gc.registryCtlClient.StartGC()
	if err != nil {
		gc.logger.Errorf("failed to get gc result: %v", err)
		return err
	}
	if err := gc.cleanCache(); err != nil {
		return err
	}
	gc.logger.Infof("GC results: status: %t, message: %s, start: %s, end: %s.", gcr.Status, gcr.Msg, gcr.StartTime, gcr.EndTime)
	gc.logger.Infof("success to run gc in job.")
	return nil
}

func (gc *GarbageCollector) init(ctx env.JobContext, params map[string]interface{}) error {
	registryctl.Init()
	gc.registryCtlClient = registryctl.RegistryCtlClient
	gc.logger = ctx.GetLogger()
	cred := auth.NewSecretAuthorizer(os.Getenv("JOBSERVICE_SECRET"))
	gc.insecure = false
	gc.coreclient = common_http.NewClient(&http.Client{
		Transport: reg.GetHTTPTransport(gc.insecure),
	}, cred)
	errTpl := "Failed to get required property: %s"
	if v, ok := ctx.Get(common.CoreURL); ok && len(v.(string)) > 0 {
		gc.CoreURL = v.(string)
	} else {
		return fmt.Errorf(errTpl, common.CoreURL)
	}
	gc.redisURL = params["redis_url_reg"].(string)
	return nil
}

func (gc *GarbageCollector) getReadOnly() (bool, error) {
	cfgs := map[string]interface{}{}
	if err := gc.coreclient.Get(fmt.Sprintf("%s/api/configs", gc.CoreURL), &cfgs); err != nil {
		return false, err
	}
	return utils.SafeCastBool(cfgs[common.ReadOnly]), nil
}

func (gc *GarbageCollector) setReadOnly(switcher bool) error {
	if err := gc.coreclient.Put(fmt.Sprintf("%s/api/configurations", gc.CoreURL), struct {
		ReadOnly bool `json:"read_only"`
	}{
		ReadOnly: switcher,
	}); err != nil {
		gc.logger.Errorf("failed to send readonly request to %s: %v", gc.CoreURL, err)
		return err
	}
	gc.logger.Info("the readonly request has been sent successfully")
	return nil
}

// cleanCache is to clean the registry cache for GC.
// To do this is because the issue https://github.com/docker/distribution/issues/2094
func (gc *GarbageCollector) cleanCache() error {

	con, err := redis.DialURL(
		gc.redisURL,
		redis.DialConnectTimeout(dialConnectionTimeout),
		redis.DialReadTimeout(dialReadTimeout),
		redis.DialWriteTimeout(dialWriteTimeout),
	)

	if err != nil {
		gc.logger.Errorf("failed to connect to redis %v", err)
		return err
	}
	defer con.Close()

	// clean all keys in registry redis DB.

	// sample of keys in registry redis:
	// 1) "blobs::sha256:1a6fd470b9ce10849be79e99529a88371dff60c60aab424c077007f6979b4812"
	// 2) "repository::library/hello-world::blobs::sha256:4ab4c602aa5eed5528a6620ff18a1dc4faef0e1ab3a5eddeddb410714478c67f"
	err = delKeys(con, blobPrefix)
	if err != nil {
		gc.logger.Errorf("failed to clean registry cache %v, pattern blobs::*", err)
		return err
	}
	err = delKeys(con, repoPrefix)
	if err != nil {
		gc.logger.Errorf("failed to clean registry cache %v, pattern repository::*", err)
		return err
	}

	return nil
}

func delKeys(con redis.Conn, pattern string) error {
	iter := 0
	keys := []string{}
	for {
		arr, err := redis.Values(con.Do("SCAN", iter, "MATCH", pattern))
		if err != nil {
			return fmt.Errorf("error retrieving '%s' keys", pattern)
		}
		iter, err = redis.Int(arr[0], nil)
		if err != nil {
			return fmt.Errorf("unexpected type for Int, got type %T", err)
		}
		k, err := redis.Strings(arr[1], nil)
		if err != nil {
			return fmt.Errorf("converts an array command reply to a []string %v", err)
		}
		keys = append(keys, k...)

		if iter == 0 {
			break
		}
	}
	for _, key := range keys {
		_, err := con.Do("DEL", key)
		if err != nil {
			return fmt.Errorf("failed to clean registry cache %v", err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package branch

import (
	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pps/store"
)

type Brancher interface {
	GetOutputCommitID(
		inputRepositoryName string,
		inputCommitID string,
		outputRepositoryName string,
	) (string, error)
	CommitOutstanding() error
	// TODO(pedge)
	//DeleteOutstanding() error
}

func NewBrancher(
	pfsAPIClient pfs.ApiClient,
	storeClient store.Client,
) Brancher {
	return newBrancher(
		pfsAPIClient,
		storeClient,
	)
}
<commit_msg>build fix<commit_after>package branch

import (
	"github.com/pachyderm/pachyderm/src/pfs"
	"github.com/pachyderm/pachyderm/src/pkg/timing"
	"github.com/pachyderm/pachyderm/src/pps/store"
)

type Brancher interface {
	GetOutputCommitID(
		inputRepositoryName string,
		inputCommitID string,
		outputRepositoryName string,
	) (string, error)
	CommitOutstanding() error
	// TODO(pedge)
	//DeleteOutstanding() error
}

func NewBrancher(
	pfsAPIClient pfs.ApiClient,
	storeClient store.Client,
	timer timing.Timer,
) Brancher {
	return newBrancher(
		pfsAPIClient,
		storeClient,
		timer,
	)
}
<|endoftext|>


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

package compilergraph

//go:generate stringer -type=GraphLayerKind

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/serulian/compiler/compilerutil"

	"github.com/cayleygraph/cayley"
	"github.com/cayleygraph/cayley/graph/memstore"
	"github.com/cayleygraph/cayley/quad"
)

// GraphLayer represents a single layer in the overall project graph.
type GraphLayer struct {
	id                string         // Unique ID for the layer.
	prefix            string         // The predicate prefix
	cayleyStore       *cayley.Handle // Handle to the cayley store.
	nodeKindPredicate Predicate      // Name of the predicate for representing the kind of a node in this layer.
	nodeKindEnum      TaggedValue    // Tagged value type that is the enum of possible node kinds.
	isFrozen          bool           // Whether the layer is frozen. Once frozen, a layer cannot be modified.
}

// NewGraphLayer returns a new graph layer of the given kind.
func (sg *SerulianGraph) NewGraphLayer(uniqueID string, nodeKindEnum TaggedValue) *GraphLayer {
	return &GraphLayer{
		id:                compilerutil.NewUniqueId(),
		prefix:            uniqueID,
		cayleyStore:       sg.cayleyStore,
		nodeKindPredicate: "node-kind",
		nodeKindEnum:      nodeKindEnum,
	}
}

// NewModifier returns a new layer modifier for modifying the graph.
func (gl *GraphLayer) NewModifier() GraphLayerModifier {
	if gl.isFrozen {
		panic("Cannot modify a frozen graph layer")
	}

	return gl.createNewModifier()
}

// Freeze freezes the layer, preventing any further modification.
func (gl *GraphLayer) Freeze() {
	gl.isFrozen = true
}

// Unfreeze unfreezes the layer, allowing for additional modification.
func (gl *GraphLayer) Unfreeze() {
	gl.isFrozen = false
}

// GetNode returns a node found in the graph layer.
func (gl *GraphLayer) GetNode(nodeID GraphNodeId) GraphNode {
	result, found := gl.TryGetNode(nodeID)
	if !found {
		panic(fmt.Sprintf("Unknown node %s in layer %s (%s)", nodeID, gl.prefix, gl.id))
	}
	return result
}

// TryGetNode tries to return a node found in the graph layer.
func (gl *GraphLayer) TryGetNode(nodeID GraphNodeId) (GraphNode, bool) {
	// Note: For efficiency reasons related to the overhead of constructing Cayley iterators,
	// we instead perform the lookup of the node directly off of the memstore's QuadIterator.
	// This code was originally:
	//	return gl.StartQuery(nodeID).TryGetNode()

	// Lookup an iterator of all quads with the node's ID as a subject.
	subjectValue := gl.cayleyStore.ValueOf(nodeIdToValue(nodeID))
	if it, ok := gl.cayleyStore.QuadIterator(quad.Subject, subjectValue).(*memstore.Iterator); ok {
		// Find a node with a predicate matching the prefixed "kind" predicate for the layer, which
		// indicates this is a node in this layer.
		fullKindPredicate := gl.getPrefixedPredicate(gl.nodeKindPredicate)

		for it.Next() {
			quad := gl.cayleyStore.Quad(it.Result())
			if quad.Predicate == fullKindPredicate {
				return GraphNode{GraphNodeId(nodeID), quad.Object, gl}, true
			}
		}
	}

	return GraphNode{}, false
}

// WalkResult is a result for each step of a walk.
type WalkResult struct {
	ParentNode        *GraphNode        // The parent node that led to this node in the walk. May be nil.
	IncomingPredicate string            // The predicate followed from the parent node to this node.
	Node              GraphNode         // The current node.
	Predicates        map[string]string // The list of outgoing predicates on this node.
}

// WalkCallback is a callback invoked for each step of a walk. If the callback returns false, the
// walk is terminated immediately.
type WalkCallback func(result *WalkResult) bool

// WalkOutward walks the graph layer outward, starting from the specified nodes, and hitting each
// node found from the outgoing predicates in the layer. Note that this method can be quite slow,
// so it should only be used for testing.
func (gl *GraphLayer) WalkOutward(startingNodes []GraphNode, callback WalkCallback) {
	encountered := map[GraphNodeId]bool{}
	var workList = make([]*WalkResult, len(startingNodes))

	// Start with walk results at the roots.
	for index, startNode := range startingNodes {
		workList[index] = &WalkResult{nil, "", startNode, map[string]string{}}
	}

outer:
	for {
		if len(workList) == 0 {
			break
		}

		// Trim the work list.
		currentResult := workList[0]
		workList = workList[1:]

		// Skip this node if we have seen it already. This prevents cycles from infinitely looping.
		currentID := currentResult.Node.NodeId
		if _, ok := encountered[currentID]; ok {
			continue
		}
		encountered[currentID] = true

		// Lookup all quads in the system from the current node, outward.
		subjectValue := gl.cayleyStore.ValueOf(nodeIdToValue(currentID))
		it := gl.cayleyStore.QuadIterator(quad.Subject, subjectValue)

		var nextWorkList = make([]*WalkResult, 0)

		for it.Next() {
			currentQuad := gl.cayleyStore.Quad(it.Result())

			// Note: We skip any predicates that are not part of this graph layer.
			predicate := valueToPredicateString(currentQuad.Predicate)
			if !strings.HasPrefix(predicate, gl.prefix+"-") {
				continue
			}

			// Try to retrieve the object as a node. If found, then we have another step in the walk.
			// Otherwise, we have a string predicate value.
			_, isPossibleNodeID := currentQuad.Object.(quad.Raw)
			found := false
			targetNode := GraphNode{}

			if isPossibleNodeID {
				targetNode, found = gl.TryGetNode(valueToNodeId(currentQuad.Object))
			}

			if isPossibleNodeID && found {
				nextWorkList = append(nextWorkList, &WalkResult{&currentResult.Node, predicate, targetNode, map[string]string{}})
			} else {
				// This is a value predicate.
				switch objectValue := currentQuad.Object.(type) {
				case quad.String:
					currentResult.Predicates[predicate] = string(objectValue)

				case quad.Raw:
					currentResult.Predicates[predicate] = string(objectValue)

				case quad.Int:
					currentResult.Predicates[predicate] = strconv.Itoa(int(objectValue))

				default:
					panic("Unknown object value type")
				}
			}
		}

		if !callback(currentResult) {
			continue outer
		}

		for _, result := range nextWorkList {
			workList = append(workList, result)
		}
	}
}

// getTaggedKey returns a unique Quad value representing the tagged name and associated value, such
// that it doesn't conflict with other tagged values in the system with the same data.
func (gl *GraphLayer) getTaggedKey(value TaggedValue) quad.Value {
	return taggedValueDataToValue(value.Value() + "|" + value.Name() + "|" + gl.prefix)
}

// parseTaggedKey parses an tagged value key (as returned by getTaggedKey) and returns the underlying value.
func (gl *GraphLayer) parseTaggedKey(value quad.Value, example TaggedValue) interface{} {
	strValue := valueToTaggedValueData(value)
	pieces := strings.SplitN(strValue, "|", 3)
	if len(pieces) != 3 {
		panic(fmt.Sprintf("Expected 3 pieces in tagged key, found: %v for value '%s'", pieces, strValue))
	}

	if pieces[2] != gl.prefix {
		panic(fmt.Sprintf("Expected tagged suffix %s, found: %s", gl.prefix, pieces[2]))
	}

	if pieces[1] != example.Name() {
		panic(fmt.Sprintf("Expected tagged key %s, found: %s", example.Name(), pieces[1]))
	}

	return example.Build(pieces[0])
}

// getPrefixedPredicate returns the given predicate prefixed with the layer prefix.
func (gl *GraphLayer) getPrefixedPredicate(predicate Predicate) quad.Value {
	return predicateToValue(Predicate(gl.prefix + "-" + string(predicate)))
}

// getPrefixedPredicates returns the given predicates prefixed with the layer prefix.
func (gl *GraphLayer) getPrefixedPredicates(predicates ...Predicate) []interface{} {
	adjusted := make([]interface{}, len(predicates))
	for index, predicate := range predicates {
		fullPredicate := gl.getPrefixedPredicate(predicate)
		adjusted[index] = fullPredicate
	}
	return adjusted
}

// getPredicatesListForDebugging returns a developer-friendly set of predicate description strings
// for all the predicates on a node.
func (gl *GraphLayer) getPredicatesListForDebugging(graphNode GraphNode) []string {
	var predicates = make([]string, 0)

	nodeIDValue := gl.cayleyStore.ValueOf(nodeIdToValue(graphNode.NodeId))
	iit := gl.cayleyStore.QuadIterator(quad.Subject, nodeIDValue)
	for iit.Next() {
		quad := gl.cayleyStore.Quad(iit.Result())
		predicates = append(predicates, fmt.Sprintf("Outgoing predicate: %v => %v", quad.Predicate, quad.Object))
	}

	oit := gl.cayleyStore.QuadIterator(quad.Object, nodeIDValue)
	for oit.Next() {
		quad := gl.cayleyStore.Quad(oit.Result())
		predicates = append(predicates, fmt.Sprintf("Incoming predicate: %v <= %v", quad.Predicate, quad.Subject))
	}

	return predicates
}
<commit_msg>Small optimizations on tagged keys<commit_after>// Copyright 2015 The Serulian Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package compilergraph

//go:generate stringer -type=GraphLayerKind

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/serulian/compiler/compilerutil"

	"github.com/cayleygraph/cayley"
	"github.com/cayleygraph/cayley/graph/memstore"
	"github.com/cayleygraph/cayley/quad"
)

// GraphLayer represents a single layer in the overall project graph.
type GraphLayer struct {
	id                string         // Unique ID for the layer.
	prefix            string         // The predicate prefix
	cayleyStore       *cayley.Handle // Handle to the cayley store.
	nodeKindPredicate Predicate      // Name of the predicate for representing the kind of a node in this layer.
	nodeKindEnum      TaggedValue    // Tagged value type that is the enum of possible node kinds.
	isFrozen          bool           // Whether the layer is frozen. Once frozen, a layer cannot be modified.
}

// NewGraphLayer returns a new graph layer of the given kind.
func (sg *SerulianGraph) NewGraphLayer(uniqueID string, nodeKindEnum TaggedValue) *GraphLayer {
	return &GraphLayer{
		id:                compilerutil.NewUniqueId(),
		prefix:            uniqueID,
		cayleyStore:       sg.cayleyStore,
		nodeKindPredicate: "node-kind",
		nodeKindEnum:      nodeKindEnum,
	}
}

// NewModifier returns a new layer modifier for modifying the graph.
func (gl *GraphLayer) NewModifier() GraphLayerModifier {
	if gl.isFrozen {
		panic("Cannot modify a frozen graph layer")
	}

	return gl.createNewModifier()
}

// Freeze freezes the layer, preventing any further modification.
func (gl *GraphLayer) Freeze() {
	gl.isFrozen = true
}

// Unfreeze unfreezes the layer, allowing for additional modification.
func (gl *GraphLayer) Unfreeze() {
	gl.isFrozen = false
}

// GetNode returns a node found in the graph layer.
func (gl *GraphLayer) GetNode(nodeID GraphNodeId) GraphNode {
	result, found := gl.TryGetNode(nodeID)
	if !found {
		panic(fmt.Sprintf("Unknown node %s in layer %s (%s)", nodeID, gl.prefix, gl.id))
	}
	return result
}

// TryGetNode tries to return a node found in the graph layer.
func (gl *GraphLayer) TryGetNode(nodeID GraphNodeId) (GraphNode, bool) {
	// Note: For efficiency reasons related to the overhead of constructing Cayley iterators,
	// we instead perform the lookup of the node directly off of the memstore's QuadIterator.
	// This code was originally:
	//	return gl.StartQuery(nodeID).TryGetNode()

	// Lookup an iterator of all quads with the node's ID as a subject.
	subjectValue := gl.cayleyStore.ValueOf(nodeIdToValue(nodeID))
	if it, ok := gl.cayleyStore.QuadIterator(quad.Subject, subjectValue).(*memstore.Iterator); ok {
		// Find a node with a predicate matching the prefixed "kind" predicate for the layer, which
		// indicates this is a node in this layer.
		fullKindPredicate := gl.getPrefixedPredicate(gl.nodeKindPredicate)

		for it.Next() {
			quad := gl.cayleyStore.Quad(it.Result())
			if quad.Predicate == fullKindPredicate {
				return GraphNode{GraphNodeId(nodeID), quad.Object, gl}, true
			}
		}
	}

	return GraphNode{}, false
}

// WalkResult is a result for each step of a walk.
type WalkResult struct {
	ParentNode        *GraphNode        // The parent node that led to this node in the walk. May be nil.
	IncomingPredicate string            // The predicate followed from the parent node to this node.
	Node              GraphNode         // The current node.
	Predicates        map[string]string // The list of outgoing predicates on this node.
}

// WalkCallback is a callback invoked for each step of a walk. If the callback returns false, the
// walk is terminated immediately.
type WalkCallback func(result *WalkResult) bool

// WalkOutward walks the graph layer outward, starting from the specified nodes, and hitting each
// node found from the outgoing predicates in the layer. Note that this method can be quite slow,
// so it should only be used for testing.
func (gl *GraphLayer) WalkOutward(startingNodes []GraphNode, callback WalkCallback) {
	encountered := map[GraphNodeId]bool{}
	var workList = make([]*WalkResult, len(startingNodes))

	// Start with walk results at the roots.
	for index, startNode := range startingNodes {
		workList[index] = &WalkResult{nil, "", startNode, map[string]string{}}
	}

outer:
	for {
		if len(workList) == 0 {
			break
		}

		// Trim the work list.
		currentResult := workList[0]
		workList = workList[1:]

		// Skip this node if we have seen it already. This prevents cycles from infinitely looping.
		currentID := currentResult.Node.NodeId
		if _, ok := encountered[currentID]; ok {
			continue
		}
		encountered[currentID] = true

		// Lookup all quads in the system from the current node, outward.
		subjectValue := gl.cayleyStore.ValueOf(nodeIdToValue(currentID))
		it := gl.cayleyStore.QuadIterator(quad.Subject, subjectValue)

		var nextWorkList = make([]*WalkResult, 0)

		for it.Next() {
			currentQuad := gl.cayleyStore.Quad(it.Result())

			// Note: We skip any predicates that are not part of this graph layer.
			predicate := valueToPredicateString(currentQuad.Predicate)
			if !strings.HasPrefix(predicate, gl.prefix+"-") {
				continue
			}

			// Try to retrieve the object as a node. If found, then we have another step in the walk.
			// Otherwise, we have a string predicate value.
			_, isPossibleNodeID := currentQuad.Object.(quad.Raw)
			found := false
			targetNode := GraphNode{}

			if isPossibleNodeID {
				targetNode, found = gl.TryGetNode(valueToNodeId(currentQuad.Object))
			}

			if isPossibleNodeID && found {
				nextWorkList = append(nextWorkList, &WalkResult{&currentResult.Node, predicate, targetNode, map[string]string{}})
			} else {
				// This is a value predicate.
				switch objectValue := currentQuad.Object.(type) {
				case quad.String:
					currentResult.Predicates[predicate] = string(objectValue)

				case quad.Raw:
					currentResult.Predicates[predicate] = string(objectValue)

				case quad.Int:
					currentResult.Predicates[predicate] = strconv.Itoa(int(objectValue))

				default:
					panic("Unknown object value type")
				}
			}
		}

		if !callback(currentResult) {
			continue outer
		}

		for _, result := range nextWorkList {
			workList = append(workList, result)
		}
	}
}

const taggedDelimeter = '|'

// getTaggedKey returns a unique Quad value representing the tagged name and associated value, such
// that it doesn't conflict with other tagged values in the system with the same data.
func (gl *GraphLayer) getTaggedKey(value TaggedValue) quad.Value {
	return taggedValueDataToValue(value.Value() + string(taggedDelimeter) + value.Name() + string(taggedDelimeter) + gl.prefix)
}

// parseTaggedKey parses an tagged value key (as returned by getTaggedKey) and returns the underlying value.
func (gl *GraphLayer) parseTaggedKey(value quad.Value, example TaggedValue) interface{} {
	strValue := valueToTaggedValueData(value)
	endIndex := strings.IndexByte(strValue, taggedDelimeter)
	return example.Build(strValue[0:endIndex])
}

// getPrefixedPredicate returns the given predicate prefixed with the layer prefix.
func (gl *GraphLayer) getPrefixedPredicate(predicate Predicate) quad.Value {
	return predicateToValue(Predicate(gl.prefix + "-" + string(predicate)))
}

// getPrefixedPredicates returns the given predicates prefixed with the layer prefix.
func (gl *GraphLayer) getPrefixedPredicates(predicates ...Predicate) []interface{} {
	adjusted := make([]interface{}, len(predicates))
	for index, predicate := range predicates {
		fullPredicate := gl.getPrefixedPredicate(predicate)
		adjusted[index] = fullPredicate
	}
	return adjusted
}

// getPredicatesListForDebugging returns a developer-friendly set of predicate description strings
// for all the predicates on a node.
func (gl *GraphLayer) getPredicatesListForDebugging(graphNode GraphNode) []string {
	var predicates = make([]string, 0)

	nodeIDValue := gl.cayleyStore.ValueOf(nodeIdToValue(graphNode.NodeId))
	iit := gl.cayleyStore.QuadIterator(quad.Subject, nodeIDValue)
	for iit.Next() {
		quad := gl.cayleyStore.Quad(iit.Result())
		predicates = append(predicates, fmt.Sprintf("Outgoing predicate: %v => %v", quad.Predicate, quad.Object))
	}

	oit := gl.cayleyStore.QuadIterator(quad.Object, nodeIDValue)
	for oit.Next() {
		quad := gl.cayleyStore.Quad(oit.Result())
		predicates = append(predicates, fmt.Sprintf("Incoming predicate: %v <= %v", quad.Predicate, quad.Subject))
	}

	return predicates
}
<|endoftext|>


<commit_before>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package component

import (
	"context"

	"github.com/spf13/viper"

	"go.opentelemetry.io/collector/config/configmodels"
)

// Component is either a receiver, exporter, processor or extension.
type Component interface {
	// Start tells the component to start. Host parameter can be used for communicating
	// with the host after Start() has already returned. If error is returned by
	// Start() then the collector startup will be aborted.
	// If this is an exporter component it may prepare for exporting
	// by connecting to the endpoint.
	//
	// If the component needs to perform a long-running starting operation then it is recommended
	// that Start() returns quickly and the long-running operation is performed in background.
	// In that case make sure that the long-running operation does not use the context passed
	// to Start() function since that context will be cancelled soon and can abort the long-running operation.
	// Create a new context from the context.Background() for long-running operations.
	Start(ctx context.Context, host Host) error

	// Shutdown is invoked during service shutdown.
	//
	// If there are any background operations running by the component they must be aborted as soon as possible.
	// Remember that if you started any long-running background operation from the Start() method that operation
	// must be also cancelled. If there are any buffer in the component, it should be cleared and the data sent
	// immediately to the next component.
	Shutdown(ctx context.Context) error
}

// Kind specified one of the 4 components kinds, see consts below.
type Kind int

const (
	_ Kind = iota // skip 0, start types from 1.
	KindReceiver
	KindProcessor
	KindExporter
	KindExtension
)

// Host represents the entity that is hosting a Component. It is used to allow communication
// between the Component and its host (normally the service.Application is the host).
type Host interface {
	// ReportFatalError is used to report to the host that the extension
	// encountered a fatal error (i.e.: an error that the instance can't recover
	// from) after its start function had already returned.
	ReportFatalError(err error)

	// GetFactory of the specified kind. Returns the factory for a component type.
	// This allows components to create other components. For example:
	//   func (r MyReceiver) Start(host component.Host) error {
	//     apacheFactory := host.GetFactory(KindReceiver,"apache").(component.ReceiverFactory)
	//     receiver, err := apacheFactory.CreateMetricsReceiver(...)
	//     ...
	//   }
	// GetFactory can be called by the component anytime after Start() begins and
	// until Shutdown() is called. Note that the component is responsible for destroying
	// other components that it creates.
	GetFactory(kind Kind, componentType configmodels.Type) Factory

	// Return map of extensions. Only enabled and created extensions will be returned.
	// Typically is used to find an extension by type or by full config name. Both cases
	// can be done by iterating the returned map. There are typically very few extensions
	// so there there is no performance implications due to iteration.
	GetExtensions() map[configmodels.Extension]ServiceExtension

	// Return map of exporters. Only enabled and created exporters will be returned.
	// Typically is used to find exporters by type or by full config name. Both cases
	// can be done by iterating the returned map. There are typically very few exporters
	// so there there is no performance implications due to iteration.
	// This returns a map by DataType of maps by exporter configs to the exporter instance.
	// Note that an exporter with the same name may be attached to multiple pipelines and
	// thus we may have an instance of the exporter for multiple data types.
	// This is an experimental function that may change or even be removed completely.
	GetExporters() map[configmodels.DataType]map[configmodels.Exporter]Exporter
}

// Factory interface must be implemented by all component factories.
type Factory interface {
	// Type gets the type of the component created by this factory.
	Type() configmodels.Type
}

// ConfigUnmarshaler interface is an optional interface that if implemented by a Factory,
// the configuration loading system will use to unmarshal the config.
type ConfigUnmarshaler interface {
	// Unmarshal is a function that un-marshals a viper data into a config struct in a custom way.
	// componentViperSection *viper.Viper
	//   The config for this specific component. May be nil or empty if no config available.
	// intoCfg interface{}
	//   An empty interface wrapping a pointer to the config struct to unmarshal into.
	Unmarshal(componentViperSection *viper.Viper, intoCfg interface{}) error
}

// CustomUnmarshaler is a function that un-marshals a viper data into a config struct
// in a custom way.
// componentViperSection *viper.Viper
//   The config for this specific component. May be nil or empty if no config available.
// intoCfg interface{}
//   An empty interface wrapping a pointer to the config struct to unmarshal into.
type CustomUnmarshaler func(componentViperSection *viper.Viper, intoCfg interface{}) error

// ApplicationStartInfo is the information that is logged at the application start and
// passed into each component. This information can be overridden in custom builds.
type ApplicationStartInfo struct {
	// Executable file name, e.g. "otelcol".
	ExeName string

	// Long name, used e.g. in the logs.
	LongName string

	// Version string.
	Version string

	// Git hash of the source code.
	GitHash string
}
<commit_msg>docs(component): improves the documentation to clarify that components should not accept anymore data on shutdown (#2481). (#2504)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package component

import (
	"context"

	"github.com/spf13/viper"

	"go.opentelemetry.io/collector/config/configmodels"
)

// Component is either a receiver, exporter, processor or extension.
type Component interface {
	// Start tells the component to start. Host parameter can be used for communicating
	// with the host after Start() has already returned. If error is returned by
	// Start() then the collector startup will be aborted.
	// If this is an exporter component it may prepare for exporting
	// by connecting to the endpoint.
	//
	// If the component needs to perform a long-running starting operation then it is recommended
	// that Start() returns quickly and the long-running operation is performed in background.
	// In that case make sure that the long-running operation does not use the context passed
	// to Start() function since that context will be cancelled soon and can abort the long-running operation.
	// Create a new context from the context.Background() for long-running operations.
	Start(ctx context.Context, host Host) error

	// Shutdown is invoked during service shutdown. After Shutdown() is called, if the component accept data in
	// any way, it should not accept it anymore.
	//
	// If there are any background operations running by the component they must be aborted as soon as possible.
	// Remember that if you started any long-running background operation from the Start() method that operation
	// must be also cancelled. If there are any buffer in the component, it should be cleared and the data sent
	// immediately to the next component.
	Shutdown(ctx context.Context) error
}

// Kind specified one of the 4 components kinds, see consts below.
type Kind int

const (
	_ Kind = iota // skip 0, start types from 1.
	KindReceiver
	KindProcessor
	KindExporter
	KindExtension
)

// Host represents the entity that is hosting a Component. It is used to allow communication
// between the Component and its host (normally the service.Application is the host).
type Host interface {
	// ReportFatalError is used to report to the host that the extension
	// encountered a fatal error (i.e.: an error that the instance can't recover
	// from) after its start function had already returned.
	ReportFatalError(err error)

	// GetFactory of the specified kind. Returns the factory for a component type.
	// This allows components to create other components. For example:
	//   func (r MyReceiver) Start(host component.Host) error {
	//     apacheFactory := host.GetFactory(KindReceiver,"apache").(component.ReceiverFactory)
	//     receiver, err := apacheFactory.CreateMetricsReceiver(...)
	//     ...
	//   }
	// GetFactory can be called by the component anytime after Start() begins and
	// until Shutdown() is called. Note that the component is responsible for destroying
	// other components that it creates.
	GetFactory(kind Kind, componentType configmodels.Type) Factory

	// Return map of extensions. Only enabled and created extensions will be returned.
	// Typically is used to find an extension by type or by full config name. Both cases
	// can be done by iterating the returned map. There are typically very few extensions
	// so there there is no performance implications due to iteration.
	GetExtensions() map[configmodels.Extension]ServiceExtension

	// Return map of exporters. Only enabled and created exporters will be returned.
	// Typically is used to find exporters by type or by full config name. Both cases
	// can be done by iterating the returned map. There are typically very few exporters
	// so there there is no performance implications due to iteration.
	// This returns a map by DataType of maps by exporter configs to the exporter instance.
	// Note that an exporter with the same name may be attached to multiple pipelines and
	// thus we may have an instance of the exporter for multiple data types.
	// This is an experimental function that may change or even be removed completely.
	GetExporters() map[configmodels.DataType]map[configmodels.Exporter]Exporter
}

// Factory interface must be implemented by all component factories.
type Factory interface {
	// Type gets the type of the component created by this factory.
	Type() configmodels.Type
}

// ConfigUnmarshaler interface is an optional interface that if implemented by a Factory,
// the configuration loading system will use to unmarshal the config.
type ConfigUnmarshaler interface {
	// Unmarshal is a function that un-marshals a viper data into a config struct in a custom way.
	// componentViperSection *viper.Viper
	//   The config for this specific component. May be nil or empty if no config available.
	// intoCfg interface{}
	//   An empty interface wrapping a pointer to the config struct to unmarshal into.
	Unmarshal(componentViperSection *viper.Viper, intoCfg interface{}) error
}

// CustomUnmarshaler is a function that un-marshals a viper data into a config struct
// in a custom way.
// componentViperSection *viper.Viper
//   The config for this specific component. May be nil or empty if no config available.
// intoCfg interface{}
//   An empty interface wrapping a pointer to the config struct to unmarshal into.
type CustomUnmarshaler func(componentViperSection *viper.Viper, intoCfg interface{}) error

// ApplicationStartInfo is the information that is logged at the application start and
// passed into each component. This information can be overridden in custom builds.
type ApplicationStartInfo struct {
	// Executable file name, e.g. "otelcol".
	ExeName string

	// Long name, used e.g. in the logs.
	LongName string

	// Version string.
	Version string

	// Git hash of the source code.
	GitHash string
}
<|endoftext|>


<commit_before>package s3gof3r

import (
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"sync"
	"syscall"
	"time"
)

const (
	qWaitMax = 2
)

type getter struct {
	url   url.URL
	b     *Bucket
	bufsz int64
	err   error

	chunkID    int
	rChunk     *chunk
	contentLen int64
	bytesRead  int64
	chunkTotal int

	readCh   chan *chunk
	getCh    chan *chunk
	quit     chan struct{}
	qWait    map[int]*chunk
	qWaitLen uint
	cond     sync.Cond

	sp *bp

	closed bool
	c      *Config

	md5  hash.Hash
	cIdx int64
}

type chunk struct {
	id     int
	header http.Header
	start  int64
	size   int64
	b      []byte
}

func newGetter(getURL url.URL, c *Config, b *Bucket) (io.ReadCloser, http.Header, error) {
	g := new(getter)
	g.url = getURL
	g.c, g.b = new(Config), new(Bucket)
	*g.c, *g.b = *c, *b
	g.bufsz = max64(c.PartSize, 1)
	g.c.NTry = max(c.NTry, 1)
	g.c.Concurrency = max(c.Concurrency, 1)

	g.getCh = make(chan *chunk)
	g.readCh = make(chan *chunk)
	g.quit = make(chan struct{})
	g.qWait = make(map[int]*chunk)
	g.b = b
	g.md5 = md5.New()
	g.cond = sync.Cond{L: &sync.Mutex{}}

	// use get instead of head for error messaging
	resp, err := g.retryRequest("GET", g.url.String(), nil)
	if err != nil {
		return nil, nil, err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return nil, nil, newRespError(resp)
	}

	// Golang changes content-length to -1 when chunked transfer encoding / EOF close response detected
	if resp.ContentLength == -1 {
		return nil, nil, fmt.Errorf("Retrieving objects with undefined content-length " +
			" responses (chunked transfer encoding / EOF close) is not supported")
	}

	g.contentLen = resp.ContentLength
	g.chunkTotal = int((g.contentLen + g.bufsz - 1) / g.bufsz) // round up, integer division
	logger.debugPrintf("object size: %3.2g MB", float64(g.contentLen)/float64((1*mb)))

	g.sp = bufferPool(g.bufsz)

	for i := 0; i < g.c.Concurrency; i++ {
		go g.worker()
	}
	go g.initChunks()
	return g, resp.Header, nil
}

func (g *getter) retryRequest(method, urlStr string, body io.ReadSeeker) (resp *http.Response, err error) {
	for i := 0; i < g.c.NTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		g.b.Sign(req)
		resp, err = g.c.Client.Do(req)
		if err == nil {
			return
		}
		logger.debugPrintln(err)
		if body != nil {
			if _, err = body.Seek(0, 0); err != nil {
				return
			}
		}
	}
	return
}

func (g *getter) initChunks() {
	id := 0
	for i := int64(0); i < g.contentLen; {
		size := min64(g.bufsz, g.contentLen-i)
		c := &chunk{
			id: id,
			header: http.Header{
				"Range": {fmt.Sprintf("bytes=%d-%d",
					i, i+size-1)},
			},
			start: i,
			size:  size,
			b:     nil,
		}
		i += size
		id++
		g.getCh <- c
	}
	close(g.getCh)
}

func (g *getter) worker() {
	for c := range g.getCh {
		g.retryGetChunk(c)
	}

}

func (g *getter) retryGetChunk(c *chunk) {
	var err error
	c.b = <-g.sp.get
	for i := 0; i < g.c.NTry; i++ {
		time.Sleep(time.Duration(math.Exp2(float64(i))) * 100 * time.Millisecond) // exponential back-off
		err = g.getChunk(c)
		if err == nil {
			return
		}
		logger.debugPrintf("error on attempt %d: retrying chunk: %v, error: %s", i, c.id, err)
	}
	g.err = err
	close(g.quit) // out of tries, ensure quit by closing channel
}

func (g *getter) getChunk(c *chunk) error {
	// ensure buffer is empty

	r, err := http.NewRequest("GET", g.url.String(), nil)
	if err != nil {
		return err
	}
	r.Header = c.header
	g.b.Sign(r)
	resp, err := g.c.Client.Do(r)
	if err != nil {
		return err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 206 {
		return newRespError(resp)
	}
	n, err := io.ReadAtLeast(resp.Body, c.b, int(c.size))
	if err != nil {
		return err
	}
	if err := resp.Body.Close(); err != nil {
		return err
	}
	if int64(n) != c.size {
		return fmt.Errorf("chunk %d: Expected %d bytes, received %d",
			c.id, c.size, n)
	}
	g.readCh <- c

	// wait for qWait to drain before starting next chunk
	g.cond.L.Lock()
	for g.qWaitLen >= qWaitMax {
		g.cond.Wait()
	}
	g.cond.L.Unlock()
	return nil
}

func (g *getter) Read(p []byte) (int, error) {
	var err error
	if g.closed {
		return 0, syscall.EINVAL
	}
	if g.err != nil {
		return 0, g.err
	}
	nw := 0
	for nw < len(p) {
		if g.bytesRead == g.contentLen {
			return nw, io.EOF
		} else if g.bytesRead > g.contentLen {
			// Here for robustness / completeness
			// Should not occur as golang uses LimitedReader up to content-length
			return nw, fmt.Errorf("Expected %d bytes, received %d (too many bytes)",
				g.contentLen, g.bytesRead)
		}

		// If for some reason no more chunks to be read and bytes are off, error, incomplete result
		if g.chunkID >= g.chunkTotal {
			return nw, fmt.Errorf("Expected %d bytes, received %d and chunkID %d >= chunkTotal %d (no more chunks remaining)",
				g.contentLen, g.bytesRead, g.chunkID, g.chunkTotal)
		}

		if g.rChunk == nil {
			g.rChunk, err = g.nextChunk()
			if err != nil {
				return 0, err
			}
			g.cIdx = 0
		}

		n := copy(p[nw:], g.rChunk.b[g.cIdx:g.rChunk.size])
		g.cIdx += int64(n)
		nw += n
		g.bytesRead += int64(n)

		if g.cIdx >= g.rChunk.size { // chunk complete
			g.sp.give <- g.rChunk.b
			g.chunkID++
			g.rChunk = nil
		}
	}
	return nw, nil

}

func (g *getter) nextChunk() (*chunk, error) {
	for {

		// first check qWait
		c := g.qWait[g.chunkID]
		if c != nil {
			delete(g.qWait, g.chunkID)
			g.cond.L.Lock()
			g.qWaitLen--
			g.cond.L.Unlock()
			g.cond.Signal() // wake up waiting worker goroutine
			if g.c.Md5Check {
				if _, err := g.md5.Write(c.b[:c.size]); err != nil {
					return nil, err
				}
			}
			return c, nil
		}
		// if next chunk not in qWait, read from channel
		select {
		case c := <-g.readCh:
			g.qWait[c.id] = c
			g.cond.L.Lock()
			g.qWaitLen++
			g.cond.L.Unlock()
		case <-g.quit:
			return nil, g.err // fatal error, quit.
		}
	}
}

func (g *getter) Close() error {
	if g.closed {
		return syscall.EINVAL
	}
	g.closed = true
	close(g.sp.quit)
	if g.err != nil {
		return g.err
	}
	if g.bytesRead != g.contentLen {
		return fmt.Errorf("read error: %d bytes read. expected: %d", g.bytesRead, g.contentLen)
	}
	if g.c.Md5Check {
		if err := g.checkMd5(); err != nil {
			return err
		}
	}
	return nil
}

func (g *getter) checkMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", g.md5.Sum(nil))
	md5Path := fmt.Sprint(".md5", g.url.Path, ".md5")
	md5Url, err := g.b.url(md5Path, g.c)
	if err != nil {
		return err
	}

	logger.debugPrintln("md5: ", calcMd5)
	logger.debugPrintln("md5Path: ", md5Path)
	resp, err := g.retryRequest("GET", md5Url.String(), nil)
	if err != nil {
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return fmt.Errorf("MD5 check failed: %s not found: %s", md5Url.String(), newRespError(resp))
	}
	givenMd5, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	if calcMd5 != string(givenMd5) {
		return fmt.Errorf("MD5 mismatch. given:%s calculated:%s", givenMd5, calcMd5)
	}
	return
}
<commit_msg>avoid multiple workers closing quit channel<commit_after>package s3gof3r

import (
	"crypto/md5"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"sync"
	"syscall"
	"time"
)

const (
	qWaitMax = 2
)

type getter struct {
	url   url.URL
	b     *Bucket
	bufsz int64
	err   error

	chunkID    int
	rChunk     *chunk
	contentLen int64
	bytesRead  int64
	chunkTotal int

	readCh   chan *chunk
	getCh    chan *chunk
	quit     chan struct{}
	qWait    map[int]*chunk
	qWaitLen uint
	cond     sync.Cond

	sp *bp

	closed bool
	c      *Config

	md5  hash.Hash
	cIdx int64
}

type chunk struct {
	id     int
	header http.Header
	start  int64
	size   int64
	b      []byte
}

func newGetter(getURL url.URL, c *Config, b *Bucket) (io.ReadCloser, http.Header, error) {
	g := new(getter)
	g.url = getURL
	g.c, g.b = new(Config), new(Bucket)
	*g.c, *g.b = *c, *b
	g.bufsz = max64(c.PartSize, 1)
	g.c.NTry = max(c.NTry, 1)
	g.c.Concurrency = max(c.Concurrency, 1)

	g.getCh = make(chan *chunk)
	g.readCh = make(chan *chunk)
	g.quit = make(chan struct{})
	g.qWait = make(map[int]*chunk)
	g.b = b
	g.md5 = md5.New()
	g.cond = sync.Cond{L: &sync.Mutex{}}

	// use get instead of head for error messaging
	resp, err := g.retryRequest("GET", g.url.String(), nil)
	if err != nil {
		return nil, nil, err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return nil, nil, newRespError(resp)
	}

	// Golang changes content-length to -1 when chunked transfer encoding / EOF close response detected
	if resp.ContentLength == -1 {
		return nil, nil, fmt.Errorf("Retrieving objects with undefined content-length " +
			" responses (chunked transfer encoding / EOF close) is not supported")
	}

	g.contentLen = resp.ContentLength
	g.chunkTotal = int((g.contentLen + g.bufsz - 1) / g.bufsz) // round up, integer division
	logger.debugPrintf("object size: %3.2g MB", float64(g.contentLen)/float64((1*mb)))

	g.sp = bufferPool(g.bufsz)

	for i := 0; i < g.c.Concurrency; i++ {
		go g.worker()
	}
	go g.initChunks()
	return g, resp.Header, nil
}

func (g *getter) retryRequest(method, urlStr string, body io.ReadSeeker) (resp *http.Response, err error) {
	for i := 0; i < g.c.NTry; i++ {
		var req *http.Request
		req, err = http.NewRequest(method, urlStr, body)
		if err != nil {
			return
		}
		g.b.Sign(req)
		resp, err = g.c.Client.Do(req)
		if err == nil {
			return
		}
		logger.debugPrintln(err)
		if body != nil {
			if _, err = body.Seek(0, 0); err != nil {
				return
			}
		}
	}
	return
}

func (g *getter) initChunks() {
	id := 0
	for i := int64(0); i < g.contentLen; {
		size := min64(g.bufsz, g.contentLen-i)
		c := &chunk{
			id: id,
			header: http.Header{
				"Range": {fmt.Sprintf("bytes=%d-%d",
					i, i+size-1)},
			},
			start: i,
			size:  size,
			b:     nil,
		}
		i += size
		id++
		g.getCh <- c
	}
	close(g.getCh)
}

func (g *getter) worker() {
	for c := range g.getCh {
		g.retryGetChunk(c)
	}

}

func (g *getter) retryGetChunk(c *chunk) {
	var err error
	c.b = <-g.sp.get
	for i := 0; i < g.c.NTry; i++ {
		time.Sleep(time.Duration(math.Exp2(float64(i))) * 100 * time.Millisecond) // exponential back-off
		err = g.getChunk(c)
		if err == nil {
			return
		}
		logger.debugPrintf("error on attempt %d: retrying chunk: %v, error: %s", i, c.id, err)
	}
	select {
	case <-g.quit: // check for closed quit channel before closing
		return
	default:
		close(g.quit) // out of tries, ensure quit by closing channel
		g.err = err
	}
}

func (g *getter) getChunk(c *chunk) error {
	// ensure buffer is empty

	r, err := http.NewRequest("GET", g.url.String(), nil)
	if err != nil {
		return err
	}
	r.Header = c.header
	g.b.Sign(r)
	resp, err := g.c.Client.Do(r)
	if err != nil {
		return err
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 206 {
		return newRespError(resp)
	}
	n, err := io.ReadAtLeast(resp.Body, c.b, int(c.size))
	if err != nil {
		return err
	}
	if err := resp.Body.Close(); err != nil {
		return err
	}
	if int64(n) != c.size {
		return fmt.Errorf("chunk %d: Expected %d bytes, received %d",
			c.id, c.size, n)
	}
	g.readCh <- c

	// wait for qWait to drain before starting next chunk
	g.cond.L.Lock()
	for g.qWaitLen >= qWaitMax {
		g.cond.Wait()
	}
	g.cond.L.Unlock()
	return nil
}

func (g *getter) Read(p []byte) (int, error) {
	var err error
	if g.closed {
		return 0, syscall.EINVAL
	}
	if g.err != nil {
		return 0, g.err
	}
	nw := 0
	for nw < len(p) {
		if g.bytesRead == g.contentLen {
			return nw, io.EOF
		} else if g.bytesRead > g.contentLen {
			// Here for robustness / completeness
			// Should not occur as golang uses LimitedReader up to content-length
			return nw, fmt.Errorf("Expected %d bytes, received %d (too many bytes)",
				g.contentLen, g.bytesRead)
		}

		// If for some reason no more chunks to be read and bytes are off, error, incomplete result
		if g.chunkID >= g.chunkTotal {
			return nw, fmt.Errorf("Expected %d bytes, received %d and chunkID %d >= chunkTotal %d (no more chunks remaining)",
				g.contentLen, g.bytesRead, g.chunkID, g.chunkTotal)
		}

		if g.rChunk == nil {
			g.rChunk, err = g.nextChunk()
			if err != nil {
				return 0, err
			}
			g.cIdx = 0
		}

		n := copy(p[nw:], g.rChunk.b[g.cIdx:g.rChunk.size])
		g.cIdx += int64(n)
		nw += n
		g.bytesRead += int64(n)

		if g.cIdx >= g.rChunk.size { // chunk complete
			g.sp.give <- g.rChunk.b
			g.chunkID++
			g.rChunk = nil
		}
	}
	return nw, nil

}

func (g *getter) nextChunk() (*chunk, error) {
	for {

		// first check qWait
		c := g.qWait[g.chunkID]
		if c != nil {
			delete(g.qWait, g.chunkID)
			g.cond.L.Lock()
			g.qWaitLen--
			g.cond.L.Unlock()
			g.cond.Signal() // wake up waiting worker goroutine
			if g.c.Md5Check {
				if _, err := g.md5.Write(c.b[:c.size]); err != nil {
					return nil, err
				}
			}
			return c, nil
		}
		// if next chunk not in qWait, read from channel
		select {
		case c := <-g.readCh:
			g.qWait[c.id] = c
			g.cond.L.Lock()
			g.qWaitLen++
			g.cond.L.Unlock()
		case <-g.quit:
			return nil, g.err // fatal error, quit.
		}
	}
}

func (g *getter) Close() error {
	if g.closed {
		return syscall.EINVAL
	}
	g.closed = true
	close(g.sp.quit)
	if g.err != nil {
		return g.err
	}
	if g.bytesRead != g.contentLen {
		return fmt.Errorf("read error: %d bytes read. expected: %d", g.bytesRead, g.contentLen)
	}
	if g.c.Md5Check {
		if err := g.checkMd5(); err != nil {
			return err
		}
	}
	return nil
}

func (g *getter) checkMd5() (err error) {
	calcMd5 := fmt.Sprintf("%x", g.md5.Sum(nil))
	md5Path := fmt.Sprint(".md5", g.url.Path, ".md5")
	md5Url, err := g.b.url(md5Path, g.c)
	if err != nil {
		return err
	}

	logger.debugPrintln("md5: ", calcMd5)
	logger.debugPrintln("md5Path: ", md5Path)
	resp, err := g.retryRequest("GET", md5Url.String(), nil)
	if err != nil {
		return
	}
	defer checkClose(resp.Body, &err)
	if resp.StatusCode != 200 {
		return fmt.Errorf("MD5 check failed: %s not found: %s", md5Url.String(), newRespError(resp))
	}
	givenMd5, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	if calcMd5 != string(givenMd5) {
		return fmt.Errorf("MD5 mismatch. given:%s calculated:%s", givenMd5, calcMd5)
	}
	return
}
<|endoftext|>


<commit_before>package waku

import (
	"errors"
	"fmt"
	"io"
	"math"
	"reflect"
	"strconv"
	"strings"

	"github.com/ethereum/go-ethereum/rlp"
)

// statusOptionKey is a current type used in statusOptions as a key.
type statusOptionKey uint

// statusOptionKeyType is a type of a statusOptions key used for a particular instance of statusOptions struct.
type statusOptionKeyType uint

const (
	sOKTS statusOptionKeyType = iota + 1 // Status Option Key Type String
	sOKTU                                // Status Option Key Type Uint
)

var (
	defaultMinPoW = math.Float64bits(0.001)
	idxFieldKey   = make(map[int]statusOptionKey)
	keyFieldIdx   = make(map[statusOptionKey]int)
)

// statusOptions defines additional information shared between peers
// during the handshake.
// There might be more options provided then fields in statusOptions
// and they should be ignored during deserialization to stay forward compatible.
// In the case of RLP, options should be serialized to an array of tuples
// where the first item is a field name and the second is a RLP-serialized value.
type statusOptions struct {
	PoWRequirement       *uint64     `rlp:"key=0"` // RLP does not support float64 natively
	BloomFilter          []byte      `rlp:"key=1"`
	LightNodeEnabled     *bool       `rlp:"key=2"`
	ConfirmationsEnabled *bool       `rlp:"key=3"`
	RateLimits           *RateLimits `rlp:"key=4"`
	TopicInterest        []TopicType `rlp:"key=5"`
	keyType              statusOptionKeyType
}

// initFLPKeyFields initialises the values of `idxFieldKey` and `keyFieldIdx`
func initRLPKeyFields() error {
	o := statusOptions{}
	v := reflect.ValueOf(o)

	for i := 0; i < v.NumField(); i++ {
		// skip unexported fields
		if !v.Field(i).CanInterface() {
			continue
		}
		rlpTag := v.Type().Field(i).Tag.Get("rlp")
		// skip fields without rlp field tag
		if rlpTag == "" {
			continue
		}

		keys := strings.Split(rlpTag, "=")

		if len(keys) != 2 || keys[0] != "key" {
			panic("invalid value of \"rlp\" tag, expected \"key=N\" where N is uint")
		}

		// parse keys[1] as an uint
		key, err := strconv.ParseUint(keys[1], 10, 64)
		if err != nil {
			return fmt.Errorf("malformed rlp tag '%s', expected \"key=N\" where N is uint: %v", rlpTag, err)
		}

		// typecast key to be of statusOptionKey type
		keyFieldIdx[statusOptionKey(key)] = i
		idxFieldKey[i] = statusOptionKey(key)
	}

	return nil
}

// WithDefaults adds the default values for a given peer.
// This are not the host default values, but the default values that ought to
// be used when receiving from an update from a peer.
func (o statusOptions) WithDefaults() statusOptions {
	if o.PoWRequirement == nil {
		o.PoWRequirement = &defaultMinPoW
	}

	if o.LightNodeEnabled == nil {
		lightNodeEnabled := false
		o.LightNodeEnabled = &lightNodeEnabled
	}

	if o.ConfirmationsEnabled == nil {
		confirmationsEnabled := false
		o.ConfirmationsEnabled = &confirmationsEnabled
	}

	if o.RateLimits == nil {
		o.RateLimits = &RateLimits{}
	}

	if o.BloomFilter == nil {
		o.BloomFilter = MakeFullNodeBloom()
	}

	return o
}

func (o statusOptions) PoWRequirementF() *float64 {
	if o.PoWRequirement == nil {
		return nil
	}
	result := math.Float64frombits(*o.PoWRequirement)
	return &result
}

func (o *statusOptions) SetPoWRequirementFromF(val float64) {
	requirement := math.Float64bits(val)
	o.PoWRequirement = &requirement
}

func (o statusOptions) EncodeRLP(w io.Writer) error {
	v := reflect.ValueOf(o)
	var optionsList []interface{}
	for i := 0; i < v.NumField(); i++ {
		field := v.Field(i)

		// skip unexported fields
		if !field.CanInterface() {
			continue
		}

		if field.IsNil() {
			continue
		}

		value := field.Interface()
		key, ok := idxFieldKey[i]
		if !ok {
			continue
		}

		if value != nil {
			optionsList = append(optionsList, []interface{}{o.encodeKey(key), value})
		}
	}
	return rlp.Encode(w, optionsList)
}

func (o *statusOptions) DecodeRLP(s *rlp.Stream) error {
	_, err := s.List()
	if err != nil {
		return fmt.Errorf("expected an outer list: %v", err)
	}
	v := reflect.ValueOf(o)

loop:
	for {
		_, err := s.List()
		switch err {
		case nil:
			// continue to decode a key
		case rlp.EOL:
			break loop
		default:
			return fmt.Errorf("expected an inner list: %v", err)
		}

		key, keyType, err := o.decodeKey(s)
		o.setKeyType(keyType)

		// Skip processing if a key does not exist.
		// It might happen when there is a new peer
		// which supports a new option with
		// a higher index.
		idx, ok := keyFieldIdx[key]
		if !ok {
			// Read the rest of the list items and dump them.
			_, err := s.Raw()
			if err != nil {
				return fmt.Errorf("failed to read the value of key %d: %v", key, err)
			}
			continue
		}

		if err := s.Decode(v.Elem().Field(idx).Addr().Interface()); err != nil {
			return fmt.Errorf("failed to decode an option %d: %v", key, err)
		}
		if err := s.ListEnd(); err != nil {
			return err
		}
	}

	return s.ListEnd()
}

func (o statusOptions) decodeKey(s *rlp.Stream) (statusOptionKey, statusOptionKeyType, error) {
	var key statusOptionKey

	// If statusOptionKey (uint) can be decoded return it
	// Ignore the first error and attempt string decoding
	if err := s.Decode(&key); err == nil {
		return key, sOKTU, nil
	}

	// Attempt decoding into a string
	var sKey string
	if err := s.Decode(&sKey); err != nil {
		return key, 0, err
	}

	// Parse string into uint
	uKey, err := strconv.ParseUint(sKey, 10, 64)
	if err != nil {
		return key, 0, err
	}

	key = statusOptionKey(uKey)
	return key, sOKTS, nil
}

// setKeyType sets a statusOptions' keyType if it hasn't previously been set
func (o *statusOptions) setKeyType(t statusOptionKeyType) {
	if o.keyType == 0 {
		o.keyType = t
	}
}

func (o statusOptions) encodeKey(key statusOptionKey) interface{} {
	if o.keyType == sOKTS {
		return fmt.Sprint(key)
	}

	return key
}

func (o statusOptions) Validate() error {
	if len(o.TopicInterest) > 1000 {
		return errors.New("topic interest is limited by 1000 items")
	}
	return nil
}
<commit_msg>Added error handling for waku handshake DecodeRLP<commit_after>package waku

import (
	"errors"
	"fmt"
	"io"
	"math"
	"reflect"
	"strconv"
	"strings"

	"github.com/ethereum/go-ethereum/rlp"
)

// statusOptionKey is a current type used in statusOptions as a key.
type statusOptionKey uint

// statusOptionKeyType is a type of a statusOptions key used for a particular instance of statusOptions struct.
type statusOptionKeyType uint

const (
	sOKTS statusOptionKeyType = iota + 1 // Status Option Key Type String
	sOKTU                                // Status Option Key Type Uint
)

var (
	defaultMinPoW = math.Float64bits(0.001)
	idxFieldKey   = make(map[int]statusOptionKey)
	keyFieldIdx   = make(map[statusOptionKey]int)
)

// statusOptions defines additional information shared between peers
// during the handshake.
// There might be more options provided then fields in statusOptions
// and they should be ignored during deserialization to stay forward compatible.
// In the case of RLP, options should be serialized to an array of tuples
// where the first item is a field name and the second is a RLP-serialized value.
type statusOptions struct {
	PoWRequirement       *uint64     `rlp:"key=0"` // RLP does not support float64 natively
	BloomFilter          []byte      `rlp:"key=1"`
	LightNodeEnabled     *bool       `rlp:"key=2"`
	ConfirmationsEnabled *bool       `rlp:"key=3"`
	RateLimits           *RateLimits `rlp:"key=4"`
	TopicInterest        []TopicType `rlp:"key=5"`
	keyType              statusOptionKeyType
}

// initFLPKeyFields initialises the values of `idxFieldKey` and `keyFieldIdx`
func initRLPKeyFields() error {
	o := statusOptions{}
	v := reflect.ValueOf(o)

	for i := 0; i < v.NumField(); i++ {
		// skip unexported fields
		if !v.Field(i).CanInterface() {
			continue
		}
		rlpTag := v.Type().Field(i).Tag.Get("rlp")
		// skip fields without rlp field tag
		if rlpTag == "" {
			continue
		}

		keys := strings.Split(rlpTag, "=")

		if len(keys) != 2 || keys[0] != "key" {
			panic("invalid value of \"rlp\" tag, expected \"key=N\" where N is uint")
		}

		// parse keys[1] as an uint
		key, err := strconv.ParseUint(keys[1], 10, 64)
		if err != nil {
			return fmt.Errorf("malformed rlp tag '%s', expected \"key=N\" where N is uint: %v", rlpTag, err)
		}

		// typecast key to be of statusOptionKey type
		keyFieldIdx[statusOptionKey(key)] = i
		idxFieldKey[i] = statusOptionKey(key)
	}

	return nil
}

// WithDefaults adds the default values for a given peer.
// This are not the host default values, but the default values that ought to
// be used when receiving from an update from a peer.
func (o statusOptions) WithDefaults() statusOptions {
	if o.PoWRequirement == nil {
		o.PoWRequirement = &defaultMinPoW
	}

	if o.LightNodeEnabled == nil {
		lightNodeEnabled := false
		o.LightNodeEnabled = &lightNodeEnabled
	}

	if o.ConfirmationsEnabled == nil {
		confirmationsEnabled := false
		o.ConfirmationsEnabled = &confirmationsEnabled
	}

	if o.RateLimits == nil {
		o.RateLimits = &RateLimits{}
	}

	if o.BloomFilter == nil {
		o.BloomFilter = MakeFullNodeBloom()
	}

	return o
}

func (o statusOptions) PoWRequirementF() *float64 {
	if o.PoWRequirement == nil {
		return nil
	}
	result := math.Float64frombits(*o.PoWRequirement)
	return &result
}

func (o *statusOptions) SetPoWRequirementFromF(val float64) {
	requirement := math.Float64bits(val)
	o.PoWRequirement = &requirement
}

func (o statusOptions) EncodeRLP(w io.Writer) error {
	v := reflect.ValueOf(o)
	var optionsList []interface{}
	for i := 0; i < v.NumField(); i++ {
		field := v.Field(i)

		// skip unexported fields
		if !field.CanInterface() {
			continue
		}

		if field.IsNil() {
			continue
		}

		value := field.Interface()
		key, ok := idxFieldKey[i]
		if !ok {
			continue
		}

		if value != nil {
			optionsList = append(optionsList, []interface{}{o.encodeKey(key), value})
		}
	}
	return rlp.Encode(w, optionsList)
}

func (o *statusOptions) DecodeRLP(s *rlp.Stream) error {
	_, err := s.List()
	if err != nil {
		return fmt.Errorf("expected an outer list: %v", err)
	}
	v := reflect.ValueOf(o)

loop:
	for {
		_, err := s.List()
		switch err {
		case nil:
			// continue to decode a key
		case rlp.EOL:
			break loop
		default:
			return fmt.Errorf("expected an inner list: %v", err)
		}

		key, keyType, err := o.decodeKey(s)
		if err != nil {
			return fmt.Errorf("key decode failure: %v", err)
		}
		o.setKeyType(keyType)

		// Skip processing if a key does not exist.
		// It might happen when there is a new peer
		// which supports a new option with
		// a higher index.
		idx, ok := keyFieldIdx[key]
		if !ok {
			// Read the rest of the list items and dump them.
			_, err := s.Raw()
			if err != nil {
				return fmt.Errorf("failed to read the value of key %d: %v", key, err)
			}
			continue
		}

		if err := s.Decode(v.Elem().Field(idx).Addr().Interface()); err != nil {
			return fmt.Errorf("failed to decode an option %d: %v", key, err)
		}
		if err := s.ListEnd(); err != nil {
			return err
		}
	}

	return s.ListEnd()
}

func (o statusOptions) decodeKey(s *rlp.Stream) (statusOptionKey, statusOptionKeyType, error) {
	var key statusOptionKey

	// If statusOptionKey (uint) can be decoded return it
	// Ignore the first error and attempt string decoding
	if err := s.Decode(&key); err == nil {
		return key, sOKTU, nil
	}

	// Attempt decoding into a string
	var sKey string
	if err := s.Decode(&sKey); err != nil {
		return key, 0, err
	}

	// Parse string into uint
	uKey, err := strconv.ParseUint(sKey, 10, 64)
	if err != nil {
		return key, 0, err
	}

	key = statusOptionKey(uKey)
	return key, sOKTS, nil
}

// setKeyType sets a statusOptions' keyType if it hasn't previously been set
func (o *statusOptions) setKeyType(t statusOptionKeyType) {
	if o.keyType == 0 {
		o.keyType = t
	}
}

func (o statusOptions) encodeKey(key statusOptionKey) interface{} {
	if o.keyType == sOKTS {
		return fmt.Sprint(key)
	}

	return key
}

func (o statusOptions) Validate() error {
	if len(o.TopicInterest) > 1000 {
		return errors.New("topic interest is limited by 1000 items")
	}
	return nil
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	slib "servicelib"
)

// Returns data via legacy endpoints for consumption by tools which require
// map data formatted in certain way
func serviceVulnAuto(rw http.ResponseWriter, req *http.Request) {
	op := opContext{}
	op.newContext(dbconn, false, req.RemoteAddr)

	var (
		ret                slib.VulnAutoList
		srchm              string
		destoper, destteam string
		destv2boverride    sql.NullString
	)

	// Extract ownership information from the interlink table to
	// build the response
	rows, err := op.Query(`SELECT srchostmatch,
		destoperatormatch, destteammatch, destv2boverride
		FROM interlinks
		WHERE ruletype = $1 ORDER BY srchostmatch`, HOST_OWNERSHIP)
	for rows.Next() {
		err = rows.Scan(&srchm, &destoper, &destteam, &destv2boverride)
		if err != nil {
			rows.Close()
			op.logf(err.Error())
			http.Error(rw, err.Error(), 500)
			return
		}
		nent := slib.VulnAuto{
			Team:     destteam,
			Operator: destoper,
			Match:    srchm,
			V2BKey:   destoper + "-" + destteam,
		}
		if destv2boverride.Valid && destv2boverride.String != "" {
			nent.V2BKey = destv2boverride.String
		}
		ret.VulnAuto = append(ret.VulnAuto, nent)
	}
	err = rows.Err()
	if err != nil {
		op.logf(err.Error())
		http.Error(rw, err.Error(), 500)
		return
	}

	buf, err := json.Marshal(&ret)
	if err != nil {
		op.logf(err.Error())
		http.Error(rw, err.Error(), 500)
		return
	}
	fmt.Fprintf(rw, string(buf))
}
<commit_msg>order vulnauto reply by ruleid<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Contributor:
// - Aaron Meihm ameihm@mozilla.com

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	slib "servicelib"
)

// Returns data via legacy endpoints for consumption by tools which require
// map data formatted in certain way
func serviceVulnAuto(rw http.ResponseWriter, req *http.Request) {
	op := opContext{}
	op.newContext(dbconn, false, req.RemoteAddr)

	var (
		ret                slib.VulnAutoList
		srchm              string
		destoper, destteam string
		destv2boverride    sql.NullString
	)

	// Extract ownership information from the interlink table to
	// build the response
	rows, err := op.Query(`SELECT srchostmatch,
		destoperatormatch, destteammatch, destv2boverride
		FROM interlinks
		WHERE ruletype = $1 ORDER BY ruleid`, HOST_OWNERSHIP)
	for rows.Next() {
		err = rows.Scan(&srchm, &destoper, &destteam, &destv2boverride)
		if err != nil {
			rows.Close()
			op.logf(err.Error())
			http.Error(rw, err.Error(), 500)
			return
		}
		nent := slib.VulnAuto{
			Team:     destteam,
			Operator: destoper,
			Match:    srchm,
			V2BKey:   destoper + "-" + destteam,
		}
		if destv2boverride.Valid && destv2boverride.String != "" {
			nent.V2BKey = destv2boverride.String
		}
		ret.VulnAuto = append(ret.VulnAuto, nent)
	}
	err = rows.Err()
	if err != nil {
		op.logf(err.Error())
		http.Error(rw, err.Error(), 500)
		return
	}

	buf, err := json.Marshal(&ret)
	if err != nil {
		op.logf(err.Error())
		http.Error(rw, err.Error(), 500)
		return
	}
	fmt.Fprintf(rw, string(buf))
}
<|endoftext|>


<commit_before>package main

import (
	"net/http"

	"github.com/koffeinsource/go-klogger"
)

// Config is the struct used to configure go-URLextract
type config struct {
	Client *http.Client
	Logger *klogger.KLogger
}

func (c *config) HTTPClient() *http.Client {
	return c.Client
}
func (c *config) Log() *klogger.KLogger {
	return c.Logger
}
<commit_msg>Removed old file.<commit_after><|endoftext|>


<commit_before>package sketches

import (
	"strconv"
	"testing"

	"datamodel"
	pb "datamodel/protobuf"
	"testutils"
	"utils"
)

func TestAddCML(t *testing.T) {
	testutils.SetupTests()
	defer testutils.TearDownTests()

	info := datamodel.NewEmptyInfo()
	info.Properties.MaxUniqueItems = utils.Int64p(1000000)
	info.Name = utils.Stringp("marvel")
	typ := pb.SketchType_FREQ
	info.Type = &typ
	sketch, err := NewCMLSketch(info)

	if err != nil {
		t.Error("expected avengers to have no error, got", err)
	}

	values := [][]byte{
		[]byte("sabertooth"),
		[]byte("thunderbolt"),
		[]byte("havoc"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("havoc")}

	if _, err := sketch.Add(values); err != nil {
		t.Error("expected no errors, got", err)
	}

	if res, err := sketch.Get([][]byte{[]byte("cyclops")}); err != nil {
		t.Error("expected no errors, got", err)
	} else if res.(*pb.FrequencyResult).Frequencies[0].GetCount() != 3 {
		t.Error("expected 'cyclops' count == 3, got", res.(*pb.FrequencyResult).Frequencies[0].GetCount())
	}
}

func BenchmarkCML(b *testing.B) {
	values := make([][]byte, 10)
	for i := 0; i < 1024; i++ {
		avenger := "avenger" + strconv.Itoa(i)
		values = append(values, []byte(avenger))
	}

	for n := 0; n < b.N; n++ {
		info := datamodel.NewEmptyInfo()
		info.Properties.MaxUniqueItems = utils.Int64p(1000)
		info.Name = utils.Stringp("marvel2")
		sketch, err := NewCMLSketch(info)
		if err != nil {
			b.Error("expected no errors, got", err)
		}
		for i := 0; i < 1000; i++ {
			if _, err := sketch.Add(values); err != nil {
				b.Error("expected no errors, got", err)
			}
		}
	}
}
<commit_msg>Add unit test for CML threshold<commit_after>package sketches

import (
	"fmt"
	"strconv"
	"testing"

	"datamodel"
	pb "datamodel/protobuf"
	"math/rand"
	"testutils"
	"utils"
)

func TestAddCML(t *testing.T) {
	testutils.SetupTests()
	defer testutils.TearDownTests()

	info := datamodel.NewEmptyInfo()
	info.Properties.MaxUniqueItems = utils.Int64p(1000000)
	info.Name = utils.Stringp("marvel")
	typ := pb.SketchType_FREQ
	info.Type = &typ
	sketch, err := NewCMLSketch(info)

	if err != nil {
		t.Error("expected avengers to have no error, got", err)
	}

	values := [][]byte{
		[]byte("sabertooth"),
		[]byte("thunderbolt"),
		[]byte("havoc"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("cyclops"),
		[]byte("havoc")}

	if _, err := sketch.Add(values); err != nil {
		t.Error("expected no errors, got", err)
	}

	if res, err := sketch.Get([][]byte{[]byte("cyclops")}); err != nil {
		t.Error("expected no errors, got", err)
	} else if res.(*pb.FrequencyResult).Frequencies[0].GetCount() != 3 {
		t.Error("expected 'cyclops' count == 3, got", res.(*pb.FrequencyResult).Frequencies[0].GetCount())
	}
}

func TestAddCMLThreshold(t *testing.T) {
	testutils.SetupTests()
	defer testutils.TearDownTests()

	info := datamodel.NewEmptyInfo()
	info.Properties.MaxUniqueItems = utils.Int64p(1024)
	info.Name = utils.Stringp("marvel")
	typ := pb.SketchType_FREQ
	info.Type = &typ
	sketch, err := NewCMLSketch(info)

	if err != nil {
		t.Error("expected avengers to have no error, got", err)
	}

	rValues := make(map[string]uint64)
	var sValues [][]byte
	thresholdSize := int64(sketch.threshold.size)
	for i := int64(0); i < info.GetProperties().GetMaxUniqueItems()/10; i++ {
		value := fmt.Sprintf("value-%d", i)
		freq := uint64(rand.Int63()) % 100

		values := make([][]byte, freq, freq)
		for i := range values {
			values[i] = []byte(value)
		}
		if _, err := sketch.Add(values); err != nil {
			t.Error("expected no errors, got", err)
		}
		rValues[value] = freq
		sValues = append(sValues, []byte(value))
		// Threshold should be nil once more than 10% is filled
		if sketch.threshold != nil && i >= thresholdSize {
			t.Error("expected threshold == nil for i ==", i)
		}

		if res, err := sketch.Get(sValues); err != nil {
			t.Error("expected no errors, got", err)
		} else {
			tmp := res.(*pb.FrequencyResult)
			mres := tmp.GetFrequencies()
			for i := 0; i < len(mres); i++ {
				if key := mres[i].GetValue(); rValues[key] != uint64(mres[i].GetCount()) {
					t.Fatalf("expected %s: %d, got %d", mres[i].GetValue(), rValues[key], uint64(mres[i].GetCount()))
				}
			}
		}
	}
}

func BenchmarkCML(b *testing.B) {
	values := make([][]byte, 10)
	for i := 0; i < 1024; i++ {
		avenger := "avenger" + strconv.Itoa(i)
		values = append(values, []byte(avenger))
	}

	for n := 0; n < b.N; n++ {
		info := datamodel.NewEmptyInfo()
		info.Properties.MaxUniqueItems = utils.Int64p(1000)
		info.Name = utils.Stringp("marvel2")
		sketch, err := NewCMLSketch(info)
		if err != nil {
			b.Error("expected no errors, got", err)
		}
		for i := 0; i < 1000; i++ {
			if _, err := sketch.Add(values); err != nil {
				b.Error("expected no errors, got", err)
			}
		}
	}
}
<|endoftext|>