<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 util

import (
	"fmt"
	"github.com/golang/glog"
	federation_v1alpha1 "k8s.io/kubernetes/federation/apis/federation/v1alpha1"
	"k8s.io/kubernetes/pkg/client/restclient"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
	utilnet "k8s.io/kubernetes/pkg/util/net"
	"net"
	"os"
)

const (
	KubeAPIQPS              = 20.0
	KubeAPIBurst            = 30
	KubeconfigSecretDataKey = "kubeconfig"
)

func BuildClusterConfig(c *federation_v1alpha1.Cluster) (*restclient.Config, error) {
	var serverAddress string
	var clusterConfig *restclient.Config
	hostIP, err := utilnet.ChooseHostInterface()
	if err != nil {
		return nil, err
	}

	for _, item := range c.Spec.ServerAddressByClientCIDRs {
		_, cidrnet, err := net.ParseCIDR(item.ClientCIDR)
		if err != nil {
			return nil, err
		}
		myaddr := net.ParseIP(hostIP.String())
		if cidrnet.Contains(myaddr) == true {
			serverAddress = item.ServerAddress
			break
		}
	}
	if serverAddress != "" {
		if c.Spec.SecretRef == nil {
			glog.Infof("didnt find secretRef for cluster %s. Trying insecure access", c.Name)
			clusterConfig, err = clientcmd.BuildConfigFromFlags(serverAddress, "")
		} else {
			kubeconfigGetter := KubeconfigGetterForCluster(c)
			clusterConfig, err = clientcmd.BuildConfigFromKubeconfigGetter(serverAddress, kubeconfigGetter)
		}
		if err != nil {
			return nil, err
		}
		clusterConfig.QPS = KubeAPIQPS
		clusterConfig.Burst = KubeAPIBurst
	}
	return clusterConfig, nil
}

// This is to inject a different kubeconfigGetter in tests.
// We dont use the standard one which calls NewInCluster in tests to avoid having to setup service accounts and mount files with secret tokens.
var KubeconfigGetterForCluster = func(c *federation_v1alpha1.Cluster) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		secretRefName := ""
		if c.Spec.SecretRef != nil {
			secretRefName = c.Spec.SecretRef.Name
		} else {
			glog.Infof("didnt find secretRef for cluster %s. Trying insecure access", c.Name)
		}
		return KubeconfigGetterForSecret(secretRefName)()
	}
}

// KubeconfigGettterForSecret is used to get the kubeconfig from the given secret.
var KubeconfigGetterForSecret = func(secretName string) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		var data []byte
		if secretName != "" {
			// Get the namespace this is running in from the env variable.
			namespace := os.Getenv("POD_NAMESPACE")
			if namespace == "" {
				return nil, fmt.Errorf("unexpected: POD_NAMESPACE env var returned empty string")
			}
			// Get a client to talk to the k8s apiserver, to fetch secrets from it.
			client, err := client.NewInCluster()
			if err != nil {
				return nil, fmt.Errorf("error in creating in-cluster client: %s", err)
			}
			data = []byte{}
			secret, err := client.Secrets(namespace).Get(secretName)
			if err != nil {
				return nil, fmt.Errorf("error in fetching secret: %s", err)
			}
			ok := false
			data, ok = secret.Data[KubeconfigSecretDataKey]
			if !ok {
				return nil, fmt.Errorf("secret does not have data with key: %s", KubeconfigSecretDataKey)
			}
		}
		return clientcmd.Load(data)
	}
}
<commit_msg>Adding retries to fetching secret in controller manager<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

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 util

import (
	"fmt"
	"net"
	"os"
	"time"

	"github.com/golang/glog"
	federation_v1alpha1 "k8s.io/kubernetes/federation/apis/federation/v1alpha1"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/client/restclient"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/client/unversioned/clientcmd"
	clientcmdapi "k8s.io/kubernetes/pkg/client/unversioned/clientcmd/api"
	utilnet "k8s.io/kubernetes/pkg/util/net"
	"k8s.io/kubernetes/pkg/util/wait"
)

const (
	KubeAPIQPS              = 20.0
	KubeAPIBurst            = 30
	KubeconfigSecretDataKey = "kubeconfig"
	getSecretTimeout        = 1 * time.Minute
)

func BuildClusterConfig(c *federation_v1alpha1.Cluster) (*restclient.Config, error) {
	var serverAddress string
	var clusterConfig *restclient.Config
	hostIP, err := utilnet.ChooseHostInterface()
	if err != nil {
		return nil, err
	}

	for _, item := range c.Spec.ServerAddressByClientCIDRs {
		_, cidrnet, err := net.ParseCIDR(item.ClientCIDR)
		if err != nil {
			return nil, err
		}
		myaddr := net.ParseIP(hostIP.String())
		if cidrnet.Contains(myaddr) == true {
			serverAddress = item.ServerAddress
			break
		}
	}
	if serverAddress != "" {
		if c.Spec.SecretRef == nil {
			glog.Infof("didnt find secretRef for cluster %s. Trying insecure access", c.Name)
			clusterConfig, err = clientcmd.BuildConfigFromFlags(serverAddress, "")
		} else {
			kubeconfigGetter := KubeconfigGetterForCluster(c)
			clusterConfig, err = clientcmd.BuildConfigFromKubeconfigGetter(serverAddress, kubeconfigGetter)
		}
		if err != nil {
			return nil, err
		}
		clusterConfig.QPS = KubeAPIQPS
		clusterConfig.Burst = KubeAPIBurst
	}
	return clusterConfig, nil
}

// This is to inject a different kubeconfigGetter in tests.
// We dont use the standard one which calls NewInCluster in tests to avoid having to setup service accounts and mount files with secret tokens.
var KubeconfigGetterForCluster = func(c *federation_v1alpha1.Cluster) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		secretRefName := ""
		if c.Spec.SecretRef != nil {
			secretRefName = c.Spec.SecretRef.Name
		} else {
			glog.Infof("didnt find secretRef for cluster %s. Trying insecure access", c.Name)
		}
		return KubeconfigGetterForSecret(secretRefName)()
	}
}

// KubeconfigGettterForSecret is used to get the kubeconfig from the given secret.
var KubeconfigGetterForSecret = func(secretName string) clientcmd.KubeconfigGetter {
	return func() (*clientcmdapi.Config, error) {
		var data []byte
		if secretName != "" {
			// Get the namespace this is running in from the env variable.
			namespace := os.Getenv("POD_NAMESPACE")
			if namespace == "" {
				return nil, fmt.Errorf("unexpected: POD_NAMESPACE env var returned empty string")
			}
			// Get a client to talk to the k8s apiserver, to fetch secrets from it.
			client, err := client.NewInCluster()
			if err != nil {
				return nil, fmt.Errorf("error in creating in-cluster client: %s", err)
			}
			data = []byte{}
			var secret *api.Secret
			err = wait.PollImmediate(1*time.Second, getSecretTimeout, func() (bool, error) {
				secret, err = client.Secrets(namespace).Get(secretName)
				if err == nil {
					return true, nil
				}
				glog.Warningf("error in fetching secret: %s", err)
				return false, nil
			})
			if err != nil {
				return nil, fmt.Errorf("timed out waiting for secret: %s", err)
			}
			if secret == nil {
				return nil, fmt.Errorf("unexpected: received null secret %s", secretName)
			}
			ok := false
			data, ok = secret.Data[KubeconfigSecretDataKey]
			if !ok {
				return nil, fmt.Errorf("secret does not have data with key: %s", KubeconfigSecretDataKey)
			}
		}
		return clientcmd.Load(data)
	}
}
<|endoftext|>


<commit_before>package shezmu

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/localhots/shezmu/stats"
)

// Shezmu is the master daemon.
type Shezmu struct {
	Subscriber  Subscriber
	Publisher   Publisher
	DaemonStats stats.Publisher
	Logger      Logger
	NumWorkers  int

	daemons      []Daemon
	queue        chan *task
	runtimeStats stats.Manager

	wgWorkers       sync.WaitGroup
	wgSystem        sync.WaitGroup
	shutdownWorkers chan struct{}
	shutdownSystem  chan struct{}
}

// Actor is a function that could be executed by daemon workers.
type Actor func()

// Subscriber is the interface that is used by daemons to subscribe to messages.
type Subscriber interface {
	Subscribe(consumer, topic string) Streamer
}

// Streamer is the interface that wraps message consumers. Error handling
// should be provided by the implementation. Feel free to panic.
type Streamer interface {
	Messages() <-chan []byte
	Close()
}

// Publisher is the interface that wraps message publishers. Error handling
// should be provided by the implementation. Feel free to panic.
type Publisher interface {
	Publish(msg []byte)
	Close()
}

// Logger is the interface that implements minimal logging functions.
type Logger interface {
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

type task struct {
	daemon    Daemon
	actor     Actor
	createdAt time.Time
	system    bool
	name      string
}

const (
	// DefaultNumWorkers is the default number of workers that would process
	// tasks.
	DefaultNumWorkers = 100
)

// Summon creates a new instance of Shezmu.
func Summon() *Shezmu {
	return &Shezmu{
		DaemonStats:     &stats.Void{},
		Logger:          log.New(os.Stdout, "[daemons] ", log.LstdFlags),
		NumWorkers:      DefaultNumWorkers,
		queue:           make(chan *task),
		runtimeStats:    stats.NewBasicStats(),
		shutdownWorkers: make(chan struct{}),
		shutdownSystem:  make(chan struct{}),
	}
}

// AddDaemon adds a new daemon.
func (s *Shezmu) AddDaemon(d Daemon) {
	base := d.base()
	base.self = d
	base.subscriber = s.Subscriber
	base.publisher = s.Publisher
	base.queue = s.queue
	base.logger = s.Logger
	base.shutdown = s.shutdownSystem

	go d.Startup()
	s.daemons = append(s.daemons, d)
}

// StartDaemons starts all registered daemons.
func (s *Shezmu) StartDaemons() {
	s.Logger.Printf("Starting %d workers", s.NumWorkers)
	for i := 0; i < s.NumWorkers; i++ {
		go s.runWorker()
	}
}

// StopDaemons stops all running daemons.
func (s *Shezmu) StopDaemons() {
	close(s.shutdownSystem)
	for _, d := range s.daemons {
		d.Shutdown()
	}

	s.wgSystem.Wait()
	close(s.shutdownWorkers)
	s.wgWorkers.Wait()
	close(s.queue)

	fmt.Println(s.runtimeStats.Fetch(stats.Latency))
}

func (s *Shezmu) runWorker() {
	s.wgWorkers.Add(1)
	defer s.wgWorkers.Done()
	defer func() {
		if err := recover(); err != nil {
			s.Logger.Printf("Worker crashed. Error: %v\n", err)
			debug.PrintStack()
			go s.runWorker() // Restarting worker
		}
	}()

	for {
		select {
		case t := <-s.queue:
			s.processTask(t)
		case <-s.shutdownWorkers:
			return
		}
	}
}

func (s *Shezmu) processTask(t *task) {
	dur := time.Now().Sub(t.createdAt)
	s.runtimeStats.Add(stats.Latency, dur)

	if t.system {
		s.processSystemTask(t)
	} else {
		s.processGeneralTask(t)
	}
}

func (s *Shezmu) processSystemTask(t *task) {
	// Abort starting a system task if shutdown was already called. Otherwise
	// incrementing a wait group counter will cause a panic. This should be an
	// extremely rare scenario when a system task crashes and tries to restart
	// after a shutdown call.
	select {
	case <-s.shutdownSystem:
		return
	default:
	}

	s.wgSystem.Add(1)
	defer s.wgSystem.Done()
	defer func() {
		if err := recover(); err != nil {
			s.Logger.Printf("System task %s recovered from a panic\nError: %v\n", t, err)
			debug.PrintStack()

			t.createdAt = time.Now()
			s.queue <- t // Restarting task
		} else {
			s.Logger.Printf("System task %s has stopped\n", t)
		}
	}()

	s.Logger.Printf("Starting system task %s\n", t)
	t.actor() // <--- ACTION STARTS HERE
}

func (s *Shezmu) processGeneralTask(t *task) {
	defer func() {
		if err := recover(); err != nil {
			s.DaemonStats.Error(t.daemon.String())
			t.daemon.base().handlePanic(err)
			s.Logger.Printf("Daemon %s recovered from a panic\nError: %v\n", t.daemon, err)
			debug.PrintStack()
		}
	}()
	defer func(start time.Time) {
		dur := time.Now().Sub(start)
		s.DaemonStats.Add(t.daemon.String(), dur)
	}(time.Now())

	t.actor() // <--- ACTION STARTS HERE
}

func (t *task) String() string {
	if t.name == "" {
		return fmt.Sprintf("[unnamed %s process]", t.daemon)
	}

	return fmt.Sprintf("%s[%s]", t.daemon, t.name)
}
<commit_msg>Add a function to clear daemons list<commit_after>package shezmu

import (
	"fmt"
	"log"
	"os"
	"runtime/debug"
	"sync"
	"time"

	"github.com/localhots/shezmu/stats"
)

// Shezmu is the master daemon.
type Shezmu struct {
	Subscriber  Subscriber
	Publisher   Publisher
	DaemonStats stats.Publisher
	Logger      Logger
	NumWorkers  int

	daemons      []Daemon
	queue        chan *task
	runtimeStats stats.Manager

	wgWorkers       sync.WaitGroup
	wgSystem        sync.WaitGroup
	shutdownWorkers chan struct{}
	shutdownSystem  chan struct{}
}

// Actor is a function that could be executed by daemon workers.
type Actor func()

// Subscriber is the interface that is used by daemons to subscribe to messages.
type Subscriber interface {
	Subscribe(consumer, topic string) Streamer
}

// Streamer is the interface that wraps message consumers. Error handling
// should be provided by the implementation. Feel free to panic.
type Streamer interface {
	Messages() <-chan []byte
	Close()
}

// Publisher is the interface that wraps message publishers. Error handling
// should be provided by the implementation. Feel free to panic.
type Publisher interface {
	Publish(msg []byte)
	Close()
}

// Logger is the interface that implements minimal logging functions.
type Logger interface {
	Printf(format string, v ...interface{})
	Println(v ...interface{})
}

type task struct {
	daemon    Daemon
	actor     Actor
	createdAt time.Time
	system    bool
	name      string
}

const (
	// DefaultNumWorkers is the default number of workers that would process
	// tasks.
	DefaultNumWorkers = 100
)

// Summon creates a new instance of Shezmu.
func Summon() *Shezmu {
	return &Shezmu{
		DaemonStats:     &stats.Void{},
		Logger:          log.New(os.Stdout, "[daemons] ", log.LstdFlags),
		NumWorkers:      DefaultNumWorkers,
		queue:           make(chan *task),
		runtimeStats:    stats.NewBasicStats(),
		shutdownWorkers: make(chan struct{}),
		shutdownSystem:  make(chan struct{}),
	}
}

// AddDaemon adds a new daemon.
func (s *Shezmu) AddDaemon(d Daemon) {
	base := d.base()
	base.self = d
	base.subscriber = s.Subscriber
	base.publisher = s.Publisher
	base.queue = s.queue
	base.logger = s.Logger
	base.shutdown = s.shutdownSystem

	go d.Startup()
	s.daemons = append(s.daemons, d)
}

// ClearDaemons clears the list of added daemons. StopDaemons() function MUST be
// called before calling ClearDaemons().
func (s *Shezmu) ClearDaemons() {
	s.daemons = []Daemon{}
}

// StartDaemons starts all registered daemons.
func (s *Shezmu) StartDaemons() {
	s.Logger.Printf("Starting %d workers", s.NumWorkers)
	for i := 0; i < s.NumWorkers; i++ {
		go s.runWorker()
	}
}

// StopDaemons stops all running daemons.
func (s *Shezmu) StopDaemons() {
	close(s.shutdownSystem)
	for _, d := range s.daemons {
		d.Shutdown()
	}

	s.wgSystem.Wait()
	close(s.shutdownWorkers)
	s.wgWorkers.Wait()
	close(s.queue)

	fmt.Println(s.runtimeStats.Fetch(stats.Latency))
}

func (s *Shezmu) runWorker() {
	s.wgWorkers.Add(1)
	defer s.wgWorkers.Done()
	defer func() {
		if err := recover(); err != nil {
			s.Logger.Printf("Worker crashed. Error: %v\n", err)
			debug.PrintStack()
			go s.runWorker() // Restarting worker
		}
	}()

	for {
		select {
		case t := <-s.queue:
			s.processTask(t)
		case <-s.shutdownWorkers:
			return
		}
	}
}

func (s *Shezmu) processTask(t *task) {
	dur := time.Now().Sub(t.createdAt)
	s.runtimeStats.Add(stats.Latency, dur)

	if t.system {
		s.processSystemTask(t)
	} else {
		s.processGeneralTask(t)
	}
}

func (s *Shezmu) processSystemTask(t *task) {
	// Abort starting a system task if shutdown was already called. Otherwise
	// incrementing a wait group counter will cause a panic. This should be an
	// extremely rare scenario when a system task crashes and tries to restart
	// after a shutdown call.
	select {
	case <-s.shutdownSystem:
		return
	default:
	}

	s.wgSystem.Add(1)
	defer s.wgSystem.Done()
	defer func() {
		if err := recover(); err != nil {
			s.Logger.Printf("System task %s recovered from a panic\nError: %v\n", t, err)
			debug.PrintStack()

			t.createdAt = time.Now()
			s.queue <- t // Restarting task
		} else {
			s.Logger.Printf("System task %s has stopped\n", t)
		}
	}()

	s.Logger.Printf("Starting system task %s\n", t)
	t.actor() // <--- ACTION STARTS HERE
}

func (s *Shezmu) processGeneralTask(t *task) {
	defer func() {
		if err := recover(); err != nil {
			s.DaemonStats.Error(t.daemon.String())
			t.daemon.base().handlePanic(err)
			s.Logger.Printf("Daemon %s recovered from a panic\nError: %v\n", t.daemon, err)
			debug.PrintStack()
		}
	}()
	defer func(start time.Time) {
		dur := time.Now().Sub(start)
		s.DaemonStats.Add(t.daemon.String(), dur)
	}(time.Now())

	t.actor() // <--- ACTION STARTS HERE
}

func (t *task) String() string {
	if t.name == "" {
		return fmt.Sprintf("[unnamed %s process]", t.daemon)
	}

	return fmt.Sprintf("%s[%s]", t.daemon, t.name)
}
<|endoftext|>


<commit_before>/*
© Copyright IBM Corporation 2017

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 main

import (
	"archive/tar"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/volume"
	"github.com/docker/docker/client"
	"github.com/docker/docker/pkg/stdcopy"
	"github.com/moby/moby/pkg/jsonmessage"
)

func imageName() string {
	image, ok := os.LookupEnv("TEST_IMAGE")
	if !ok {
		image = "mq-devserver:latest-x86-64"
	}
	return image
}

func coverage() bool {
	cover := os.Getenv("TEST_COVER")
	if cover == "true" || cover == "1" {
		return true
	}
	return false
}

// coverageDir returns the host directory to use for code coverage data
func coverageDir(t *testing.T) string {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	return filepath.Join(dir, "coverage")
}

// coverageBind returns a string to use to add a bind-mounted directory for code coverage data
func coverageBind(t *testing.T) string {
	return coverageDir(t) + ":/var/coverage"
}

func cleanContainer(t *testing.T, cli *client.Client, ID string) {
	i, err := cli.ContainerInspect(context.Background(), ID)
	if err == nil {
		// Log the results and continue
		t.Logf("Inspected container %v: %#v", ID, i)
		s, err := json.MarshalIndent(i, "", "    ")
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Inspected container %v: %v", ID, string(s))
	}
	t.Logf("Stopping container: %v", ID)
	timeout := 10 * time.Second
	// Stop the container.  This allows the coverage output to be generated.
	err = cli.ContainerStop(context.Background(), ID, &timeout)
	if err != nil {
		// Just log the error and continue
		t.Log(err)
	}
	t.Log("Container stopped")
	// If a code coverage file has been generated, then rename it to match the test name
	os.Rename(filepath.Join(coverageDir(t), "container.cov"), filepath.Join(coverageDir(t), t.Name()+".cov"))
	// Log the container output for any container we're about to delete
	t.Logf("Console log from container %v:\n%v", ID, inspectLogs(t, cli, ID))

	t.Logf("Removing container: %s", ID)
	opts := types.ContainerRemoveOptions{
		RemoveVolumes: true,
		Force:         true,
	}
	err = cli.ContainerRemove(context.Background(), ID, opts)
	if err != nil {
		t.Error(err)
	}
}

// runContainer creates and starts a container.  If no image is specified in
// the container config, then the image name is retrieved from the TEST_IMAGE
// environment variable.
func runContainer(t *testing.T, cli *client.Client, containerConfig *container.Config) string {
	if containerConfig.Image == "" {
		containerConfig.Image = imageName()
	}
	// if coverage
	containerConfig.Env = append(containerConfig.Env, "COVERAGE_FILE="+t.Name()+".cov")
	hostConfig := container.HostConfig{
		// PortBindings: nat.PortMap{
		// 	"1414/tcp": []nat.PortBinding{
		// 		{
		// 			HostIP:   "0.0.0.0",
		// 			HostPort: "1414",
		// 		},
		// 	},
		// },
		Binds: []string{
			coverageBind(t),
		},
	}
	networkingConfig := network.NetworkingConfig{}
	t.Logf("Running container (%s)", containerConfig.Image)
	ctr, err := cli.ContainerCreate(context.Background(), containerConfig, &hostConfig, &networkingConfig, t.Name())
	if err != nil {
		t.Fatal(err)
	}
	startContainer(t, cli, ctr.ID)
	return ctr.ID
}

func runContainerOneShot(t *testing.T, cli *client.Client, command ...string) (int64, string) {
	containerConfig := container.Config{
		Entrypoint: command,
	}
	id := runContainer(t, cli, &containerConfig)
	defer cleanContainer(t, cli, id)
	return waitForContainer(t, cli, id, 10), inspectLogs(t, cli, id)
}

func startContainer(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Starting container: %v", ID)
	startOptions := types.ContainerStartOptions{}
	err := cli.ContainerStart(context.Background(), ID, startOptions)
	if err != nil {
		t.Fatal(err)
	}
}

func stopContainer(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Stopping container: %v", ID)
	timeout := 10 * time.Second
	err := cli.ContainerStop(context.Background(), ID, &timeout) //Duration(20)*time.Second)
	if err != nil {
		t.Fatal(err)
	}
}

func getCoverageExitCode(t *testing.T, orig int64) int64 {
	f := filepath.Join(coverageDir(t), "exitCode")
	_, err := os.Stat(f)
	if err != nil {
		t.Log(err)
		return orig
	}
	// Remove the file, ready for the next test
	defer os.Remove(f)
	buf, err := ioutil.ReadFile(f)
	if err != nil {
		t.Log(err)
		return orig
	}
	rc, err := strconv.Atoi(string(buf))
	if err != nil {
		t.Log(err)
		return orig
	}
	t.Logf("Retrieved exit code %v from file", rc)
	return int64(rc)
}

// waitForContainer waits until a container has exited
func waitForContainer(t *testing.T, cli *client.Client, ID string, timeout int64) int64 {
	//ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	//defer cancel()
	rc, err := cli.ContainerWait(context.Background(), ID)

	if coverage() {
		// COVERAGE: When running coverage, the exit code is written to a file,
		// to allow the coverage to be generated (which doesn't happen for non-zero
		// exit codes)
		rc = getCoverageExitCode(t, rc)
	}

	//	err := <-errC
	if err != nil {
		t.Fatal(err)
	}
	//	wait := <-waitC
	return rc
}

// execContainerWithExitCode runs a command in a running container, and returns the exit code
// Note: due to a bug in Docker/Moby code, you always get an exit code of 0 if you attach to the
// container to get output.  This is why these are two separate commands.
func execContainerWithExitCode(t *testing.T, cli *client.Client, ID string, user string, cmd []string) int {
	config := types.ExecConfig{
		User:        user,
		Privileged:  false,
		Tty:         false,
		AttachStdin: false,
		// Note that you still need to attach stdout/stderr, even though they're not wanted
		AttachStdout: true,
		AttachStderr: true,
		Detach:       false,
		Cmd:          cmd,
	}
	resp, err := cli.ContainerExecCreate(context.Background(), ID, config)
	if err != nil {
		t.Fatal(err)
	}
	cli.ContainerExecStart(context.Background(), resp.ID, types.ExecStartCheck{
		Detach: false,
		Tty:    false,
	})
	if err != nil {
		t.Fatal(err)
	}
	inspect, err := cli.ContainerExecInspect(context.Background(), resp.ID)
	if err != nil {
		t.Fatal(err)
	}
	return inspect.ExitCode
}

// execContainerWithOutput runs a command in a running container, and returns the output from stdout/stderr
// Note: due to a bug in Docker/Moby code, you always get an exit code of 0 if you attach to the
// container to get output.  This is why these are two separate commands.
func execContainerWithOutput(t *testing.T, cli *client.Client, ID string, user string, cmd []string) string {
	config := types.ExecConfig{
		User:         user,
		Privileged:   false,
		Tty:          false,
		AttachStdin:  false,
		AttachStdout: true,
		AttachStderr: true,
		Detach:       false,
		Cmd:          cmd,
	}
	resp, err := cli.ContainerExecCreate(context.Background(), ID, config)
	if err != nil {
		t.Fatal(err)
	}
	hijack, err := cli.ContainerExecAttach(context.Background(), resp.ID, config)
	if err != nil {
		t.Fatal(err)
	}
	cli.ContainerExecStart(context.Background(), resp.ID, types.ExecStartCheck{
		Detach: false,
		Tty:    false,
	})
	if err != nil {
		t.Fatal(err)
	}
	buf := new(bytes.Buffer)
	// Each output line has a header, which needs to be removed
	_, err = stdcopy.StdCopy(buf, buf, hijack.Reader)
	if err != nil {
		log.Fatal(err)
	}
	return strings.TrimSpace(buf.String())
}

func waitForReady(t *testing.T, cli *client.Client, ID string) {
	for {
		rc := execContainerWithExitCode(t, cli, ID, "mqm", []string{"chkmqready"})
		if rc == 0 {
			t.Log("MQ is ready")
			return
		}
	}
}

func getIPAddress(t *testing.T, cli *client.Client, ID string) string {
	ctr, err := cli.ContainerInspect(context.Background(), ID)
	if err != nil {
		t.Fatal(err)
	}
	return ctr.NetworkSettings.IPAddress
}

func createNetwork(t *testing.T, cli *client.Client) string {
	name := "test"
	t.Logf("Creating network: %v", name)
	opts := types.NetworkCreate{}
	net, err := cli.NetworkCreate(context.Background(), name, opts)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Created network %v with ID %v", name, net.ID)
	return net.ID
}

func removeNetwork(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Removing network ID: %v", ID)
	err := cli.NetworkRemove(context.Background(), ID)
	if err != nil {
		t.Fatal(err)
	}
}

func createVolume(t *testing.T, cli *client.Client) types.Volume {
	v, err := cli.VolumeCreate(context.Background(), volume.VolumesCreateBody{
		Driver:     "local",
		DriverOpts: map[string]string{},
		Labels:     map[string]string{},
		Name:       t.Name(),
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Created volume %v", t.Name())
	return v
}

func removeVolume(t *testing.T, cli *client.Client, name string) {
	t.Logf("Removing volume %v", name)
	err := cli.VolumeRemove(context.Background(), name, true)
	if err != nil {
		t.Fatal(err)
	}
}

func inspectLogs(t *testing.T, cli *client.Client, ID string) string {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	reader, err := cli.ContainerLogs(ctx, ID, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
	})
	if err != nil {
		log.Fatal(err)
	}
	buf := new(bytes.Buffer)
	// Each output line has a header, which needs to be removed
	_, err = stdcopy.StdCopy(buf, buf, reader)
	if err != nil {
		log.Fatal(err)
	}
	return buf.String()
}

// generateTAR creates a TAR-formatted []byte, with the specified files included.
func generateTAR(t *testing.T, files []struct{ Name, Body string }) []byte {
	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)
	for _, file := range files {
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0600,
			Size: int64(len(file.Body)),
		}
		err := tw.WriteHeader(hdr)
		if err != nil {
			t.Fatal(err)
		}
		_, err = tw.Write([]byte(file.Body))
		if err != nil {
			t.Fatal(err)
		}
	}
	err := tw.Close()
	if err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

// createImage creates a new Docker image with the specified files included.
func createImage(t *testing.T, cli *client.Client, files []struct{ Name, Body string }) string {
	r := bytes.NewReader(generateTAR(t, files))
	tag := strings.ToLower(t.Name())
	buildOptions := types.ImageBuildOptions{
		Context: r,
		Tags:    []string{tag},
	}
	resp, err := cli.ImageBuild(context.Background(), r, buildOptions)
	if err != nil {
		t.Fatal(err)
	}
	// resp (ImageBuildResponse) contains a series of JSON messages
	dec := json.NewDecoder(resp.Body)
	for {
		m := jsonmessage.JSONMessage{}
		err := dec.Decode(&m)
		if m.Error != nil {
			t.Fatal(m.ErrorMessage)
		}
		t.Log(strings.TrimSpace(m.Stream))
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
	}
	return tag
}

// deleteImage deletes a Docker image
func deleteImage(t *testing.T, cli *client.Client, id string) {
	cli.ImageRemove(context.Background(), id, types.ImageRemoveOptions{
		Force: true,
	})
}
<commit_msg>Switch from Moby to Docker import<commit_after>/*
© Copyright IBM Corporation 2017

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 main

import (
	"archive/tar"
	"bytes"
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/volume"
	"github.com/docker/docker/client"
	"github.com/docker/docker/pkg/jsonmessage"
	"github.com/docker/docker/pkg/stdcopy"
)

func imageName() string {
	image, ok := os.LookupEnv("TEST_IMAGE")
	if !ok {
		image = "mq-devserver:latest-x86-64"
	}
	return image
}

func coverage() bool {
	cover := os.Getenv("TEST_COVER")
	if cover == "true" || cover == "1" {
		return true
	}
	return false
}

// coverageDir returns the host directory to use for code coverage data
func coverageDir(t *testing.T) string {
	dir, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	return filepath.Join(dir, "coverage")
}

// coverageBind returns a string to use to add a bind-mounted directory for code coverage data
func coverageBind(t *testing.T) string {
	return coverageDir(t) + ":/var/coverage"
}

func cleanContainer(t *testing.T, cli *client.Client, ID string) {
	i, err := cli.ContainerInspect(context.Background(), ID)
	if err == nil {
		// Log the results and continue
		t.Logf("Inspected container %v: %#v", ID, i)
		s, err := json.MarshalIndent(i, "", "    ")
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Inspected container %v: %v", ID, string(s))
	}
	t.Logf("Stopping container: %v", ID)
	timeout := 10 * time.Second
	// Stop the container.  This allows the coverage output to be generated.
	err = cli.ContainerStop(context.Background(), ID, &timeout)
	if err != nil {
		// Just log the error and continue
		t.Log(err)
	}
	t.Log("Container stopped")
	// If a code coverage file has been generated, then rename it to match the test name
	os.Rename(filepath.Join(coverageDir(t), "container.cov"), filepath.Join(coverageDir(t), t.Name()+".cov"))
	// Log the container output for any container we're about to delete
	t.Logf("Console log from container %v:\n%v", ID, inspectLogs(t, cli, ID))

	t.Logf("Removing container: %s", ID)
	opts := types.ContainerRemoveOptions{
		RemoveVolumes: true,
		Force:         true,
	}
	err = cli.ContainerRemove(context.Background(), ID, opts)
	if err != nil {
		t.Error(err)
	}
}

// runContainer creates and starts a container.  If no image is specified in
// the container config, then the image name is retrieved from the TEST_IMAGE
// environment variable.
func runContainer(t *testing.T, cli *client.Client, containerConfig *container.Config) string {
	if containerConfig.Image == "" {
		containerConfig.Image = imageName()
	}
	// if coverage
	containerConfig.Env = append(containerConfig.Env, "COVERAGE_FILE="+t.Name()+".cov")
	hostConfig := container.HostConfig{
		// PortBindings: nat.PortMap{
		// 	"1414/tcp": []nat.PortBinding{
		// 		{
		// 			HostIP:   "0.0.0.0",
		// 			HostPort: "1414",
		// 		},
		// 	},
		// },
		Binds: []string{
			coverageBind(t),
		},
	}
	networkingConfig := network.NetworkingConfig{}
	t.Logf("Running container (%s)", containerConfig.Image)
	ctr, err := cli.ContainerCreate(context.Background(), containerConfig, &hostConfig, &networkingConfig, t.Name())
	if err != nil {
		t.Fatal(err)
	}
	startContainer(t, cli, ctr.ID)
	return ctr.ID
}

func runContainerOneShot(t *testing.T, cli *client.Client, command ...string) (int64, string) {
	containerConfig := container.Config{
		Entrypoint: command,
	}
	id := runContainer(t, cli, &containerConfig)
	defer cleanContainer(t, cli, id)
	return waitForContainer(t, cli, id, 10), inspectLogs(t, cli, id)
}

func startContainer(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Starting container: %v", ID)
	startOptions := types.ContainerStartOptions{}
	err := cli.ContainerStart(context.Background(), ID, startOptions)
	if err != nil {
		t.Fatal(err)
	}
}

func stopContainer(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Stopping container: %v", ID)
	timeout := 10 * time.Second
	err := cli.ContainerStop(context.Background(), ID, &timeout) //Duration(20)*time.Second)
	if err != nil {
		t.Fatal(err)
	}
}

func getCoverageExitCode(t *testing.T, orig int64) int64 {
	f := filepath.Join(coverageDir(t), "exitCode")
	_, err := os.Stat(f)
	if err != nil {
		t.Log(err)
		return orig
	}
	// Remove the file, ready for the next test
	defer os.Remove(f)
	buf, err := ioutil.ReadFile(f)
	if err != nil {
		t.Log(err)
		return orig
	}
	rc, err := strconv.Atoi(string(buf))
	if err != nil {
		t.Log(err)
		return orig
	}
	t.Logf("Retrieved exit code %v from file", rc)
	return int64(rc)
}

// waitForContainer waits until a container has exited
func waitForContainer(t *testing.T, cli *client.Client, ID string, timeout int64) int64 {
	//ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
	//defer cancel()
	rc, err := cli.ContainerWait(context.Background(), ID)

	if coverage() {
		// COVERAGE: When running coverage, the exit code is written to a file,
		// to allow the coverage to be generated (which doesn't happen for non-zero
		// exit codes)
		rc = getCoverageExitCode(t, rc)
	}

	//	err := <-errC
	if err != nil {
		t.Fatal(err)
	}
	//	wait := <-waitC
	return rc
}

// execContainerWithExitCode runs a command in a running container, and returns the exit code
// Note: due to a bug in Docker/Moby code, you always get an exit code of 0 if you attach to the
// container to get output.  This is why these are two separate commands.
func execContainerWithExitCode(t *testing.T, cli *client.Client, ID string, user string, cmd []string) int {
	config := types.ExecConfig{
		User:        user,
		Privileged:  false,
		Tty:         false,
		AttachStdin: false,
		// Note that you still need to attach stdout/stderr, even though they're not wanted
		AttachStdout: true,
		AttachStderr: true,
		Detach:       false,
		Cmd:          cmd,
	}
	resp, err := cli.ContainerExecCreate(context.Background(), ID, config)
	if err != nil {
		t.Fatal(err)
	}
	cli.ContainerExecStart(context.Background(), resp.ID, types.ExecStartCheck{
		Detach: false,
		Tty:    false,
	})
	if err != nil {
		t.Fatal(err)
	}
	inspect, err := cli.ContainerExecInspect(context.Background(), resp.ID)
	if err != nil {
		t.Fatal(err)
	}
	return inspect.ExitCode
}

// execContainerWithOutput runs a command in a running container, and returns the output from stdout/stderr
// Note: due to a bug in Docker/Moby code, you always get an exit code of 0 if you attach to the
// container to get output.  This is why these are two separate commands.
func execContainerWithOutput(t *testing.T, cli *client.Client, ID string, user string, cmd []string) string {
	config := types.ExecConfig{
		User:         user,
		Privileged:   false,
		Tty:          false,
		AttachStdin:  false,
		AttachStdout: true,
		AttachStderr: true,
		Detach:       false,
		Cmd:          cmd,
	}
	resp, err := cli.ContainerExecCreate(context.Background(), ID, config)
	if err != nil {
		t.Fatal(err)
	}
	hijack, err := cli.ContainerExecAttach(context.Background(), resp.ID, config)
	if err != nil {
		t.Fatal(err)
	}
	cli.ContainerExecStart(context.Background(), resp.ID, types.ExecStartCheck{
		Detach: false,
		Tty:    false,
	})
	if err != nil {
		t.Fatal(err)
	}
	buf := new(bytes.Buffer)
	// Each output line has a header, which needs to be removed
	_, err = stdcopy.StdCopy(buf, buf, hijack.Reader)
	if err != nil {
		log.Fatal(err)
	}
	return strings.TrimSpace(buf.String())
}

func waitForReady(t *testing.T, cli *client.Client, ID string) {
	for {
		rc := execContainerWithExitCode(t, cli, ID, "mqm", []string{"chkmqready"})
		if rc == 0 {
			t.Log("MQ is ready")
			return
		}
	}
}

func getIPAddress(t *testing.T, cli *client.Client, ID string) string {
	ctr, err := cli.ContainerInspect(context.Background(), ID)
	if err != nil {
		t.Fatal(err)
	}
	return ctr.NetworkSettings.IPAddress
}

func createNetwork(t *testing.T, cli *client.Client) string {
	name := "test"
	t.Logf("Creating network: %v", name)
	opts := types.NetworkCreate{}
	net, err := cli.NetworkCreate(context.Background(), name, opts)
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Created network %v with ID %v", name, net.ID)
	return net.ID
}

func removeNetwork(t *testing.T, cli *client.Client, ID string) {
	t.Logf("Removing network ID: %v", ID)
	err := cli.NetworkRemove(context.Background(), ID)
	if err != nil {
		t.Fatal(err)
	}
}

func createVolume(t *testing.T, cli *client.Client) types.Volume {
	v, err := cli.VolumeCreate(context.Background(), volume.VolumesCreateBody{
		Driver:     "local",
		DriverOpts: map[string]string{},
		Labels:     map[string]string{},
		Name:       t.Name(),
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Logf("Created volume %v", t.Name())
	return v
}

func removeVolume(t *testing.T, cli *client.Client, name string) {
	t.Logf("Removing volume %v", name)
	err := cli.VolumeRemove(context.Background(), name, true)
	if err != nil {
		t.Fatal(err)
	}
}

func inspectLogs(t *testing.T, cli *client.Client, ID string) string {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	reader, err := cli.ContainerLogs(ctx, ID, types.ContainerLogsOptions{
		ShowStdout: true,
		ShowStderr: true,
	})
	if err != nil {
		log.Fatal(err)
	}
	buf := new(bytes.Buffer)
	// Each output line has a header, which needs to be removed
	_, err = stdcopy.StdCopy(buf, buf, reader)
	if err != nil {
		log.Fatal(err)
	}
	return buf.String()
}

// generateTAR creates a TAR-formatted []byte, with the specified files included.
func generateTAR(t *testing.T, files []struct{ Name, Body string }) []byte {
	buf := new(bytes.Buffer)
	tw := tar.NewWriter(buf)
	for _, file := range files {
		hdr := &tar.Header{
			Name: file.Name,
			Mode: 0600,
			Size: int64(len(file.Body)),
		}
		err := tw.WriteHeader(hdr)
		if err != nil {
			t.Fatal(err)
		}
		_, err = tw.Write([]byte(file.Body))
		if err != nil {
			t.Fatal(err)
		}
	}
	err := tw.Close()
	if err != nil {
		t.Fatal(err)
	}
	return buf.Bytes()
}

// createImage creates a new Docker image with the specified files included.
func createImage(t *testing.T, cli *client.Client, files []struct{ Name, Body string }) string {
	r := bytes.NewReader(generateTAR(t, files))
	tag := strings.ToLower(t.Name())
	buildOptions := types.ImageBuildOptions{
		Context: r,
		Tags:    []string{tag},
	}
	resp, err := cli.ImageBuild(context.Background(), r, buildOptions)
	if err != nil {
		t.Fatal(err)
	}
	// resp (ImageBuildResponse) contains a series of JSON messages
	dec := json.NewDecoder(resp.Body)
	for {
		m := jsonmessage.JSONMessage{}
		err := dec.Decode(&m)
		if m.Error != nil {
			t.Fatal(m.ErrorMessage)
		}
		t.Log(strings.TrimSpace(m.Stream))
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
	}
	return tag
}

// deleteImage deletes a Docker image
func deleteImage(t *testing.T, cli *client.Client, id string) {
	cli.ImageRemove(context.Background(), id, types.ImageRemoveOptions{
		Force: true,
	})
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package wrappers

import (
	"bytes"
	"testing"
)

func TestPackerByte(t *testing.T) {
	p := Packer{MaxSize: 1}

	p.PackByte(0x01)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 1 {
		t.Fatalf("Packer.PackByte wrote %d byte(s) but expected %d byte(s)", size, 1)
	}

	expected := []byte{0x01}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackByte wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerShort(t *testing.T) {
	p := Packer{MaxSize: 2}

	p.PackShort(0x0102)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 2 {
		t.Fatalf("Packer.PackShort wrote %d byte(s) but expected %d byte(s)", size, 2)
	}

	expected := []byte{0x01, 0x02}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackShort wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerInt(t *testing.T) {
	p := Packer{MaxSize: 4}

	p.PackInt(0x01020304)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 4 {
		t.Fatalf("Packer.PackInt wrote %d byte(s) but expected %d byte(s)", size, 4)
	}

	expected := []byte{0x01, 0x02, 0x03, 0x04}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackInt wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerLong(t *testing.T) {
	p := Packer{MaxSize: 8}

	p.PackLong(0x0102030405060708)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 8 {
		t.Fatalf("Packer.PackLong wrote %d byte(s) but expected %d byte(s)", size, 8)
	}

	expected := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackLong wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerString(t *testing.T) {
	p := Packer{MaxSize: 5}

	p.PackStr("Ava")

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 5 {
		t.Fatalf("Packer.PackStr wrote %d byte(s) but expected %d byte(s)", size, 5)
	}

	expected := []byte{0x00, 0x03, 0x41, 0x76, 0x61}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackStr wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPacker(t *testing.T) {
	packer := Packer{
		MaxSize: 3,
	}

	if packer.Errored() {
		t.Fatalf("Packer has error %s", packer.Err)
	}

	packer.PackShort(17)
	if len(packer.Bytes) != 2 {
		t.Fatalf("Wrong byte length")
	}

	packer.PackShort(1)
	if !packer.Errored() {
		t.Fatalf("Packer should have error")
	}

	newPacker := Packer{
		Bytes: packer.Bytes,
	}

	if newPacker.UnpackShort() != 17 {
		t.Fatalf("Unpacked wrong value")
	}
}

func TestPackBool(t *testing.T) {
	p := Packer{MaxSize: 3}
	p.PackBool(false)
	p.PackBool(true)
	p.PackBool(false)
	if p.Errored() {
		t.Fatal("should have been able to pack 3 bools")
	}

	p2 := Packer{Bytes: p.Bytes}
	bool1, bool2, bool3 := p2.UnpackBool(), p2.UnpackBool(), p2.UnpackBool()

	if p.Errored() {
		t.Fatalf("errors while unpacking bools: %v", p.Errs)
	}

	if bool1 || !bool2 || bool3 {
		t.Fatal("got back wrong values")
	}
}
<commit_msg>utils: Add test for Packer.UnpackByte<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package wrappers

import (
	"bytes"
	"testing"
)

const (
	ByteSentinal  = 0
)

func TestPackerPackByte(t *testing.T) {
	p := Packer{MaxSize: 1}

	p.PackByte(0x01)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 1 {
		t.Fatalf("Packer.PackByte wrote %d byte(s) but expected %d byte(s)", size, 1)
	}

	expected := []byte{0x01}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackByte wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}

	p.PackByte(0x02)
	if !p.Errored() {
		t.Fatal("Packer.PackByte did not fail when attempt was beyond p.MaxSize")
	}
}

func TestPackerUnpackByte(t *testing.T) {
	var (
		p                = Packer{Bytes: []byte{0x01}, Offset: 0}
		actual           = p.UnpackByte()
		expected    byte = 1
		expectedLen      = ByteLen
	)
	if p.Errored() {
		t.Fatalf("Packer.UnpackByte unexpectedly raised %s", p.Err)
	} else if actual != expected {
		t.Fatalf("Packer.UnpackByte returned %d, but expected %d", actual, expected)
	} else if p.Offset != expectedLen {
		t.Fatalf("Packer.UnpackByte left Offset %d, expected %d", p.Offset, expectedLen)
	}

	actual = p.UnpackByte()
	if !p.Errored() {
		t.Fatalf("Packer.UnpackByte should have set error, due to attempted out of bounds read")
	} else if actual != ByteSentinal {
		t.Fatalf("Packer.UnpackByte returned %d, expected sentinal value %d", actual, ByteSentinal)
	}
}

func TestPackerShort(t *testing.T) {
	p := Packer{MaxSize: 2}

	p.PackShort(0x0102)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 2 {
		t.Fatalf("Packer.PackShort wrote %d byte(s) but expected %d byte(s)", size, 2)
	}

	expected := []byte{0x01, 0x02}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackShort wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerInt(t *testing.T) {
	p := Packer{MaxSize: 4}

	p.PackInt(0x01020304)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 4 {
		t.Fatalf("Packer.PackInt wrote %d byte(s) but expected %d byte(s)", size, 4)
	}

	expected := []byte{0x01, 0x02, 0x03, 0x04}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackInt wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerLong(t *testing.T) {
	p := Packer{MaxSize: 8}

	p.PackLong(0x0102030405060708)

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 8 {
		t.Fatalf("Packer.PackLong wrote %d byte(s) but expected %d byte(s)", size, 8)
	}

	expected := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackLong wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPackerString(t *testing.T) {
	p := Packer{MaxSize: 5}

	p.PackStr("Ava")

	if p.Errored() {
		t.Fatal(p.Err)
	}

	if size := len(p.Bytes); size != 5 {
		t.Fatalf("Packer.PackStr wrote %d byte(s) but expected %d byte(s)", size, 5)
	}

	expected := []byte{0x00, 0x03, 0x41, 0x76, 0x61}
	if !bytes.Equal(p.Bytes, expected) {
		t.Fatalf("Packer.PackStr wrote:\n%v\nExpected:\n%v", p.Bytes, expected)
	}
}

func TestPacker(t *testing.T) {
	packer := Packer{
		MaxSize: 3,
	}

	if packer.Errored() {
		t.Fatalf("Packer has error %s", packer.Err)
	}

	packer.PackShort(17)
	if len(packer.Bytes) != 2 {
		t.Fatalf("Wrong byte length")
	}

	packer.PackShort(1)
	if !packer.Errored() {
		t.Fatalf("Packer should have error")
	}

	newPacker := Packer{
		Bytes: packer.Bytes,
	}

	if newPacker.UnpackShort() != 17 {
		t.Fatalf("Unpacked wrong value")
	}
}

func TestPackBool(t *testing.T) {
	p := Packer{MaxSize: 3}
	p.PackBool(false)
	p.PackBool(true)
	p.PackBool(false)
	if p.Errored() {
		t.Fatal("should have been able to pack 3 bools")
	}

	p2 := Packer{Bytes: p.Bytes}
	bool1, bool2, bool3 := p2.UnpackBool(), p2.UnpackBool(), p2.UnpackBool()

	if p.Errored() {
		t.Fatalf("errors while unpacking bools: %v", p.Errs)
	}

	if bool1 || !bool2 || bool3 {
		t.Fatal("got back wrong values")
	}
}
<|endoftext|>


<commit_before>package testing

import (
	"golang.org/x/tools/godoc/vfs"
	"sourcegraph.com/sourcegraph/go-vcs/vcs"
)

type MockRepository struct {
	ResolveRevision_ func(spec string) (vcs.CommitID, error)
	ResolveTag_      func(name string) (vcs.CommitID, error)
	ResolveBranch_   func(name string) (vcs.CommitID, error)

	Branches_ func() ([]*vcs.Branch, error)
	Tags_     func() ([]*vcs.Tag, error)

	GetCommit_ func(vcs.CommitID) (*vcs.Commit, error)
	Commits_   func(vcs.CommitsOptions) ([]*vcs.Commit, uint, error)

	FileSystem_ func(at vcs.CommitID) (vfs.FileSystem, error)
}

var _ vcs.Repository = MockRepository{}

func (r MockRepository) ResolveRevision(spec string) (vcs.CommitID, error) {
	return r.ResolveRevision_(spec)
}

func (r MockRepository) ResolveTag(name string) (vcs.CommitID, error) {
	return r.ResolveTag_(name)
}

func (r MockRepository) ResolveBranch(name string) (vcs.CommitID, error) {
	return r.ResolveBranch_(name)
}

func (r MockRepository) Branches() ([]*vcs.Branch, error) {
	return r.Branches_()
}

func (r MockRepository) Tags() ([]*vcs.Tag, error) {
	return r.Tags_()
}

func (r MockRepository) GetCommit(id vcs.CommitID) (*vcs.Commit, error) {
	return r.GetCommit_(id)
}

func (r MockRepository) Commits(opt vcs.CommitsOptions) ([]*vcs.Commit, uint, error) {
	return r.Commits_(opt)
}

func (r MockRepository) FileSystem(at vcs.CommitID) (vfs.FileSystem, error) {
	return r.FileSystem_(at)
}
<commit_msg>add Blame mock<commit_after>package testing

import (
	"golang.org/x/tools/godoc/vfs"
	"sourcegraph.com/sourcegraph/go-vcs/vcs"
)

type MockRepository struct {
	ResolveRevision_ func(spec string) (vcs.CommitID, error)
	ResolveTag_      func(name string) (vcs.CommitID, error)
	ResolveBranch_   func(name string) (vcs.CommitID, error)

	Branches_ func() ([]*vcs.Branch, error)
	Tags_     func() ([]*vcs.Tag, error)

	GetCommit_ func(vcs.CommitID) (*vcs.Commit, error)
	Commits_   func(vcs.CommitsOptions) ([]*vcs.Commit, uint, error)

	BlameFile_ func(path string, opt *vcs.BlameOptions) ([]*vcs.Hunk, error)

	FileSystem_ func(at vcs.CommitID) (vfs.FileSystem, error)
}

var (
	_ interface {
		vcs.Repository
		vcs.Blamer
	} = MockRepository{}
)

func (r MockRepository) ResolveRevision(spec string) (vcs.CommitID, error) {
	return r.ResolveRevision_(spec)
}

func (r MockRepository) ResolveTag(name string) (vcs.CommitID, error) {
	return r.ResolveTag_(name)
}

func (r MockRepository) ResolveBranch(name string) (vcs.CommitID, error) {
	return r.ResolveBranch_(name)
}

func (r MockRepository) Branches() ([]*vcs.Branch, error) {
	return r.Branches_()
}

func (r MockRepository) Tags() ([]*vcs.Tag, error) {
	return r.Tags_()
}

func (r MockRepository) GetCommit(id vcs.CommitID) (*vcs.Commit, error) {
	return r.GetCommit_(id)
}

func (r MockRepository) Commits(opt vcs.CommitsOptions) ([]*vcs.Commit, uint, error) {
	return r.Commits_(opt)
}

func (r MockRepository) BlameFile(path string, opt *vcs.BlameOptions) ([]*vcs.Hunk, error) {
	return r.BlameFile_(path, opt)
}

func (r MockRepository) FileSystem(at vcs.CommitID) (vfs.FileSystem, error) {
	return r.FileSystem_(at)
}
<|endoftext|>


<commit_before>// The ultimateq bot framework.
package main

import (
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/aarondl/query"
	"github.com/aarondl/quotes"
	"github.com/aarondl/ultimateq/bot"
	"github.com/aarondl/ultimateq/data"
	"github.com/aarondl/ultimateq/dispatch/cmd"
	"github.com/aarondl/ultimateq/irc"
)

var (
	sanitizeNewline = strings.NewReplacer("\r\n", " ", "\n", " ")
	rgxSpace        = regexp.MustCompile(`\s{2,}`)
	queryConf       query.Config
)

const (
	dateFormat = "January 02, 2006 at 3:04pm MST"
)

/* =====================
 Helper methods.
===================== */
func sanitize(str string) string {
	return rgxSpace.ReplaceAllString(sanitizeNewline.Replace(str), " ")
}

type Quoter struct {
	db *quotes.QuoteDB
}

type Queryer struct {
}

type Handler struct {
	b *bot.Bot
}

// Let reflection hook up the commands, instead of doing it here.
func (_ *Quoter) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

func (_ *Queryer) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

func (_ *Handler) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

/* =====================
 Quoter methods.
===================== */

func (q *Quoter) Addquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	quote := ev.GetArg("quote")
	if len(quote) == 0 {
		return nil
	}

	ev.Close()

	err := q.db.AddQuote(nick, quote)
	if err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else {
		w.Notice(nick, "\x02Quote:\x02 Added.")
	}
	return nil
}

func (q *Quoter) Delquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}
	if did, err := q.db.DelQuote(int(id)); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else if !did {
		w.Notice(nick, "\x02Quote:\x02 Could not find quote %d.", id)
	} else {
		w.Noticef(nick, "\x02Quote:\x02 Quote %d deleted.", id)
	}
	return nil
}

func (q *Quoter) Editquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	quote := ev.GetArg("quote")
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if len(quote) == 0 {
		return nil
	}

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}
	if did, err := q.db.EditQuote(int(id), quote); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else if !did {
		w.Notice(nick, "\x02Quote:\x02 Could not find quote %d.", id)
	} else {
		w.Noticef(nick, "\x02Quote:\x02 Quote %d updated.", id)
	}
	return nil
}

func (q *Quoter) Quote(w irc.Writer, ev *cmd.Event) error {
	strid := ev.GetArg("id")
	nick := ev.Nick()
	ev.Close()

	var quote string
	var id int
	var err error
	if len(strid) > 0 {
		getid, err := strconv.Atoi(strid)
		id = int(getid)
		if err != nil {
			w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
			return nil
		}
		quote, err = q.db.GetQuote(id)
	} else {
		id, quote, err = q.db.RandomQuote()
	}
	if err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
		return nil
	}

	if len(quote) == 0 {
		w.Notify(ev.Event, nick, "\x02Quote:\x02 Does not exist.")
	} else {
		w.Notifyf(ev.Event, nick, "\x02Quote (\x02#%d\x02):\x02 %s",
			id, quote)
	}
	return nil
}

func (q *Quoter) Quotes(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	ev.Close()

	w.Notifyf(ev.Event, nick, "\x02Quote:\x02 %d quote(s) in database.",
		q.db.NQuotes())
	return nil
}

func (q *Quoter) Details(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}

	if date, author, err := q.db.GetDetails(int(id)); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else {
		w.Notifyf(ev.Event, nick,
			"\x02Quote (\x02#%d\x02):\x02 Created on %s by %s",
			id, time.Unix(date, 0).UTC().Format(dateFormat), author)
	}

	return nil
}

/* =====================
 Queryer methods.
===================== */

func (_ *Queryer) PrivmsgChannel(w irc.Writer, ev *irc.Event) {
	if out, err := query.YouTube(ev.Message()); len(out) != 0 {
		w.Privmsg(ev.Target(), out)
	} else if err != nil {
		nick := ev.Nick()
		w.Notice(nick, err.Error())
	}
}

func (_ *Queryer) Calc(w irc.Writer, ev *cmd.Event) error {
	q := ev.GetArg("query")
	nick := ev.Nick()
	ev.Close()

	if out, err := query.Wolfram(q, &queryConf); len(out) != 0 {
		out = sanitize(out)

		// Ensure two lines only
		// ircmaxlen - maxhostsize - PRIVMSG - targetsize - spacing - colons
		maxlen := 2 * (510 - 62 - 7 - len(ev.Target()) - 3 - 2)
		if len(out) > maxlen {
			out = out[:maxlen-3]
			out += "..."
		}

		w.Notify(ev.Event, nick, out)
	} else if err != nil {
		w.Notice(nick, err.Error())
	}

	return nil
}

func (_ *Queryer) Google(w irc.Writer, ev *cmd.Event) error {
	q := ev.GetArg("query")
	nick := ev.Nick()
	ev.Close()

	if out, err := query.Google(q); len(out) != 0 {
		out = sanitize(out)
		w.Notify(ev.Event, nick, out)
	} else if err != nil {
		w.Notice(nick, err.Error())
	}

	return nil
}

/* =====================
 Handler methods.
===================== */

func (h *Handler) Up(w irc.Writer, ev *cmd.Event) error {
	user := ev.UserAccess
	ch := ev.TargetChannel
	if ch == nil {
		return fmt.Errorf("Must be a channel that the bot is on.")
	}
	chname := ch.Name()

	if !putPeopleUp(ev.Event, chname, user, w) {
		return cmd.MakeFlagsError("ov")
	}
	return nil
}

func (h *Handler) HandleRaw(w irc.Writer, ev *irc.Event) {
	if ev.Name == irc.JOIN {
		h.b.UsingStore(func(s *data.Store) {
			a := s.GetAuthedUser(ev.NetworkID, ev.Sender)
			ch := ev.Target()
			putPeopleUp(ev, ch, a, w)
		})
	}
}

func putPeopleUp(ev *irc.Event, ch string,
	a *data.UserAccess, w irc.Writer) (did bool) {
	if a != nil {
		nick := ev.Nick()
		if a.HasFlag(ev.NetworkID, ch, 'o') {
			w.Sendf("MODE %s +o :%s", ch, nick)
			did = true
		} else if a.HasFlag(ev.NetworkID, ch, 'v') {
			w.Sendf("MODE %s +v :%s", ch, nick)
			did = true
		}
	}
	return
}

func (h *Handler) PrivmsgUser(w irc.Writer, ev *irc.Event) {
	flds := strings.Fields(ev.Message())
	if ev.Nick() == "Aaron" && flds[0] == "do" {
		w.Send(strings.Join(flds[1:], " "))
	}
}

func main() {

	var queryer Queryer
	if conf := query.NewConfig("wolfid.toml"); conf != nil {
		queryConf = *conf
	} else {
		log.Println("Error loading wolfram configuration.")
	}
	qdb, err := quotes.OpenDB("quotes.sqlite3")
	if err != nil {
		log.Fatalln("Error opening quotes db:", err)
	}
	defer qdb.Close()
	var quoter = Quoter{qdb}

	err = bot.Run(func(b *bot.Bot) {
		// Quote commands
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Retrieves a quote. Randomly selects a quote if no id is provided.",
			"quote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "[id]",
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Shows the number of quotes in the database.",
			"quotes",
			&quoter,
			cmd.PRIVMSG, cmd.ALL,
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Gets the details for a specific quote.",
			"details",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "id",
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Adds a quote to the database.",
			"addquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "quote...",
		))
		b.RegisterCmd(cmd.MkAuthCmd(
			"quote",
			"Removes a quote from the database.",
			"delquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, 0, "Q", "id",
		))
		b.RegisterCmd(cmd.MkAuthCmd(
			"quote",
			"Edits an existing quote.",
			"editquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, 0, "Q", "id", "quote...",
		))

		// Queryer commands
		b.Register(irc.PRIVMSG, &queryer)
		b.RegisterCmd(cmd.MkCmd(
			"query",
			"Submits a query to Google.",
			"google",
			&queryer,
			cmd.PRIVMSG, cmd.ALL, "query...",
		))
		b.RegisterCmd(cmd.MkCmd(
			"query",
			"Submits a query to Wolfram Alpha.",
			"calc",
			&queryer,
			cmd.PRIVMSG, cmd.ALL, "query...",
		))

		// Handler commands
		handler := Handler{b}
		b.Register(irc.PRIVMSG, &handler)
		b.Register(irc.JOIN, &handler)
		b.RegisterCmd(cmd.MkAuthCmd(
			"simple",
			"Gives the user ops or voice if they have o or v flags respectively.",
			"up",
			&handler,
			cmd.PRIVMSG, cmd.ALL, 0, "", "#chan",
		))
	})

	if err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>Added weather command to get weather report from yr.no<commit_after>// The ultimateq bot framework.
package main

import (
	"fmt"
	"log"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/aarondl/query"
	"github.com/aarondl/quotes"
	"github.com/aarondl/ultimateq/bot"
	"github.com/aarondl/ultimateq/data"
	"github.com/aarondl/ultimateq/dispatch/cmd"
	"github.com/aarondl/ultimateq/irc"
)

var (
	sanitizeNewline = strings.NewReplacer("\r\n", " ", "\n", " ")
	rgxSpace        = regexp.MustCompile(`\s{2,}`)
	queryConf       query.Config
)

const (
	dateFormat = "January 02, 2006 at 3:04pm MST"
)

/* =====================
 Helper methods.
===================== */
func sanitize(str string) string {
	return rgxSpace.ReplaceAllString(sanitizeNewline.Replace(str), " ")
}

type Quoter struct {
	db *quotes.QuoteDB
}

type Queryer struct {
}

type Handler struct {
	b *bot.Bot
}

// Let reflection hook up the commands, instead of doing it here.
func (_ *Quoter) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

func (_ *Queryer) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

func (_ *Handler) Cmd(_ string, _ irc.Writer, _ *cmd.Event) error {
	return nil
}

/* =====================
 Quoter methods.
===================== */

func (q *Quoter) Addquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	quote := ev.GetArg("quote")
	if len(quote) == 0 {
		return nil
	}

	ev.Close()

	err := q.db.AddQuote(nick, quote)
	if err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else {
		w.Notice(nick, "\x02Quote:\x02 Added.")
	}
	return nil
}

func (q *Quoter) Delquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}
	if did, err := q.db.DelQuote(int(id)); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else if !did {
		w.Notice(nick, "\x02Quote:\x02 Could not find quote %d.", id)
	} else {
		w.Noticef(nick, "\x02Quote:\x02 Quote %d deleted.", id)
	}
	return nil
}

func (q *Quoter) Editquote(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	quote := ev.GetArg("quote")
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if len(quote) == 0 {
		return nil
	}

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}
	if did, err := q.db.EditQuote(int(id), quote); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else if !did {
		w.Notice(nick, "\x02Quote:\x02 Could not find quote %d.", id)
	} else {
		w.Noticef(nick, "\x02Quote:\x02 Quote %d updated.", id)
	}
	return nil
}

func (q *Quoter) Quote(w irc.Writer, ev *cmd.Event) error {
	strid := ev.GetArg("id")
	nick := ev.Nick()
	ev.Close()

	var quote string
	var id int
	var err error
	if len(strid) > 0 {
		getid, err := strconv.Atoi(strid)
		id = int(getid)
		if err != nil {
			w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
			return nil
		}
		quote, err = q.db.GetQuote(id)
	} else {
		id, quote, err = q.db.RandomQuote()
	}
	if err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
		return nil
	}

	if len(quote) == 0 {
		w.Notify(ev.Event, nick, "\x02Quote:\x02 Does not exist.")
	} else {
		w.Notifyf(ev.Event, nick, "\x02Quote (\x02#%d\x02):\x02 %s",
			id, quote)
	}
	return nil
}

func (q *Quoter) Quotes(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	ev.Close()

	w.Notifyf(ev.Event, nick, "\x02Quote:\x02 %d quote(s) in database.",
		q.db.NQuotes())
	return nil
}

func (q *Quoter) Details(w irc.Writer, ev *cmd.Event) error {
	nick := ev.Nick()
	id, err := strconv.Atoi(ev.GetArg("id"))
	ev.Close()

	if err != nil {
		w.Notice(nick, "\x02Quote:\x02 Not a valid id.")
		return nil
	}

	if date, author, err := q.db.GetDetails(int(id)); err != nil {
		w.Noticef(nick, "\x02Quote:\x02 %v", err)
	} else {
		w.Notifyf(ev.Event, nick,
			"\x02Quote (\x02#%d\x02):\x02 Created on %s by %s",
			id, time.Unix(date, 0).UTC().Format(dateFormat), author)
	}

	return nil
}

/* =====================
 Queryer methods.
===================== */

func (_ *Queryer) PrivmsgChannel(w irc.Writer, ev *irc.Event) {
	if out, err := query.YouTube(ev.Message()); len(out) != 0 {
		w.Privmsg(ev.Target(), out)
	} else if err != nil {
		nick := ev.Nick()
		w.Notice(nick, err.Error())
	}
}

func (_ *Queryer) Calc(w irc.Writer, ev *cmd.Event) error {
	q := ev.GetArg("query")
	nick := ev.Nick()
	ev.Close()

	if out, err := query.Wolfram(q, &queryConf); len(out) != 0 {
		out = sanitize(out)

		// Ensure two lines only
		// ircmaxlen - maxhostsize - PRIVMSG - targetsize - spacing - colons
		maxlen := 2 * (510 - 62 - 7 - len(ev.Target()) - 3 - 2)
		if len(out) > maxlen {
			out = out[:maxlen-3]
			out += "..."
		}

		w.Notify(ev.Event, nick, out)
	} else if err != nil {
		w.Notice(nick, err.Error())
	}

	return nil
}

func (_ *Queryer) Google(w irc.Writer, ev *cmd.Event) error {
	q := ev.GetArg("query")
	nick := ev.Nick()
	ev.Close()

	if out, err := query.Google(q); len(out) != 0 {
		out = sanitize(out)
		w.Notify(ev.Event, nick, out)
	} else if err != nil {
		w.Notice(nick, err.Error())
	}

	return nil
}

func (_ *Queryer) Weather(m *irc.Message, e *data.DataEndpoint,
	c *commander.CommandData) error {

	c.Close()

	q := c.GetArg("query")
	nick := m.Nick()
	if out, err := query.Weather(q); len(out) != 0 {
		out = sanitize(out)
		if targ := m.Target(); isNick(targ) {
			e.Notice(nick, out)
		} else {
			e.Privmsg(targ, out)
		}
	} else if err != nil {
		e.Notice(nick, err.Error())
	}

	return nil
}

/* =====================
 Handler methods.
===================== */

func (h *Handler) Up(w irc.Writer, ev *cmd.Event) error {
	user := ev.UserAccess
	ch := ev.TargetChannel
	if ch == nil {
		return fmt.Errorf("Must be a channel that the bot is on.")
	}
	chname := ch.Name()

	if !putPeopleUp(ev.Event, chname, user, w) {
		return cmd.MakeFlagsError("ov")
	}
	return nil
}

func (h *Handler) HandleRaw(w irc.Writer, ev *irc.Event) {
	if ev.Name == irc.JOIN {
		h.b.UsingStore(func(s *data.Store) {
			a := s.GetAuthedUser(ev.NetworkID, ev.Sender)
			ch := ev.Target()
			putPeopleUp(ev, ch, a, w)
		})
	}
}

func putPeopleUp(ev *irc.Event, ch string,
	a *data.UserAccess, w irc.Writer) (did bool) {
	if a != nil {
		nick := ev.Nick()
		if a.HasFlag(ev.NetworkID, ch, 'o') {
			w.Sendf("MODE %s +o :%s", ch, nick)
			did = true
		} else if a.HasFlag(ev.NetworkID, ch, 'v') {
			w.Sendf("MODE %s +v :%s", ch, nick)
			did = true
		}
	}
	return
}

func (h *Handler) PrivmsgUser(w irc.Writer, ev *irc.Event) {
	flds := strings.Fields(ev.Message())
	if ev.Nick() == "Aaron" && flds[0] == "do" {
		w.Send(strings.Join(flds[1:], " "))
	}
}

func main() {

	var queryer Queryer
	if conf := query.NewConfig("wolfid.toml"); conf != nil {
		queryConf = *conf
	} else {
		log.Println("Error loading wolfram configuration.")
	}
	qdb, err := quotes.OpenDB("quotes.sqlite3")
	if err != nil {
		log.Fatalln("Error opening quotes db:", err)
	}
	defer qdb.Close()
	var quoter = Quoter{qdb}

	err = bot.Run(func(b *bot.Bot) {
		// Quote commands
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Retrieves a quote. Randomly selects a quote if no id is provided.",
			"quote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "[id]",
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Shows the number of quotes in the database.",
			"quotes",
			&quoter,
			cmd.PRIVMSG, cmd.ALL,
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Gets the details for a specific quote.",
			"details",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "id",
		))
		b.RegisterCmd(cmd.MkCmd(
			"quote",
			"Adds a quote to the database.",
			"addquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, "quote...",
		))
		b.RegisterCmd(cmd.MkAuthCmd(
			"quote",
			"Removes a quote from the database.",
			"delquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, 0, "Q", "id",
		))
		b.RegisterCmd(cmd.MkAuthCmd(
			"quote",
			"Edits an existing quote.",
			"editquote",
			&quoter,
			cmd.PRIVMSG, cmd.ALL, 0, "Q", "id", "quote...",
		))

		// Queryer commands
		b.Register(irc.PRIVMSG, &queryer)
		b.RegisterCmd(cmd.MkCmd(
			"query",
			"Submits a query to Google.",
			"google",
			&queryer,
			cmd.PRIVMSG, cmd.ALL, "query...",
		))
		b.RegisterCmd(cmd.MkCmd(
			"query",
			"Submits a query to Wolfram Alpha.",
			"calc",
			&queryer,
			cmd.PRIVMSG, cmd.ALL, "query...",
		))
		b.RegisterCommand(commander.MkCmd(
			"query",
			"Fetches a weather report from yr.no.",
			"weather",
			&queryer,
			commander.PRIVMSG, commander.ALL, "query...",
		))

		// Handler commands
		handler := Handler{b}
		b.Register(irc.PRIVMSG, &handler)
		b.Register(irc.JOIN, &handler)
		b.RegisterCmd(cmd.MkAuthCmd(
			"simple",
			"Gives the user ops or voice if they have o or v flags respectively.",
			"up",
			&handler,
			cmd.PRIVMSG, cmd.ALL, 0, "", "#chan",
		))
	})

	if err != nil {
		log.Fatalln(err)
	}
}
<|endoftext|>


<commit_before>package testutil

import (
	"path/filepath"

	"veyron.io/tools/lib/collect"
	"veyron.io/tools/lib/util"
)

// runJSTest is a harness for executing javascript tests.
func (t *testEnv) runJSTest(ctx *util.Context, testName, testDir, target string, cleanFn func() error, env map[string]string) (_ *TestResult, e error) {
	// Initialize the test.
	cleanup, err := t.initTest(ctx, testName, []string{"web"})
	if err != nil {
		return nil, err
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	// Navigate to the target directory.
	if err := ctx.Run().Chdir(testDir); err != nil {
		return nil, err
	}

	// Clean up after previous instances of the test.
	opts := t.setTestEnv(ctx.Run().Opts())
	if err := ctx.Run().CommandWithOpts(opts, "make", "clean"); err != nil {
		return nil, err
	}
	if cleanFn != nil {
		if err := cleanFn(); err != nil {
			return nil, err
		}
	}

	// Run the test target.
	if err := ctx.Run().CommandWithOpts(opts, "make", target); err != nil {
		return &TestResult{Status: TestFailed}, nil
	}

	return &TestResult{Status: TestPassed}, nil
}

// veyronJSBuildExtension tests the veyron javascript build extension.
func (t *testEnv) veyronJSBuildExtension(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "extension/veyron.crx"
	return t.runJSTest(ctx, testName, testDir, target, nil, nil)
}

// veyronJSDoc (re)generates the content of the veyron javascript
// documentation server.
func (t *testEnv) veyronJSDoc(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "docs"
	webDir, jsDocDir := "/usr/share/nginx/www/jsdoc", filepath.Join(testDir, "docs")
	cleanFn := func() error {
		if err := ctx.Run().RemoveAll(webDir); err != nil {
			return err
		}
		return nil
	}
	result, err := t.runJSTest(ctx, testName, testDir, target, cleanFn, nil)
	if err != nil {
		return nil, err
	}
	// Move generated js documentation to the web server directory.
	if err := ctx.Run().Rename(jsDocDir, webDir); err != nil {
		return nil, err
	}
	return result, nil
}

// veyronJSBrowserIntegrationTest runs the veyron javascript integration test in a browser environment using nacl plugin.
func (t *testEnv) veyronJSBrowserIntegrationTest(ctx *util.Context, testName string) (*TestResult, error) {
	// TODO(aghassemi): Re-enable the test when it is fixed.
	return &TestResult{Status: TestPassed}, nil

	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-integration-browser"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["BROWSER_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSNodeIntegrationTest runs the veyron javascript integration test in NodeJS environment using wspr.
func (t *testEnv) veyronJSNodeIntegrationTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-integration-node"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSUnitTest runs the veyron javascript unit test.
func (t *testEnv) veyronJSUnitTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-unit"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSVdlTest runs the veyron javascript vdl test.
func (t *testEnv) veyronJSVdlTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-vdl"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSVomTest runs the veyron javascript vom test.
func (t *testEnv) veyronJSVomTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron", "javascript", "vom")
	target := "test"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}
<commit_msg>TBR: tools/testutil/javascript: fix env processing.<commit_after>package testutil

import (
	"path/filepath"

	"veyron.io/tools/lib/collect"
	"veyron.io/tools/lib/util"
)

// runJSTest is a harness for executing javascript tests.
func (t *testEnv) runJSTest(ctx *util.Context, testName, testDir, target string, cleanFn func() error, env map[string]string) (_ *TestResult, e error) {
	// Initialize the test.
	cleanup, err := t.initTest(ctx, testName, []string{"web"})
	if err != nil {
		return nil, err
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	// Navigate to the target directory.
	if err := ctx.Run().Chdir(testDir); err != nil {
		return nil, err
	}

	// Clean up after previous instances of the test.
	opts := t.setTestEnv(ctx.Run().Opts())
	for key, value := range env {
		opts.Env[key] = value
	}
	if err := ctx.Run().CommandWithOpts(opts, "make", "clean"); err != nil {
		return nil, err
	}
	if cleanFn != nil {
		if err := cleanFn(); err != nil {
			return nil, err
		}
	}

	// Run the test target.
	if err := ctx.Run().CommandWithOpts(opts, "make", target); err != nil {
		return &TestResult{Status: TestFailed}, nil
	}

	return &TestResult{Status: TestPassed}, nil
}

// veyronJSBuildExtension tests the veyron javascript build extension.
func (t *testEnv) veyronJSBuildExtension(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "extension/veyron.crx"
	return t.runJSTest(ctx, testName, testDir, target, nil, nil)
}

// veyronJSDoc (re)generates the content of the veyron javascript
// documentation server.
func (t *testEnv) veyronJSDoc(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "docs"
	webDir, jsDocDir := "/usr/share/nginx/www/jsdoc", filepath.Join(testDir, "docs")
	cleanFn := func() error {
		if err := ctx.Run().RemoveAll(webDir); err != nil {
			return err
		}
		return nil
	}
	result, err := t.runJSTest(ctx, testName, testDir, target, cleanFn, nil)
	if err != nil {
		return nil, err
	}
	// Move generated js documentation to the web server directory.
	if err := ctx.Run().Rename(jsDocDir, webDir); err != nil {
		return nil, err
	}
	return result, nil
}

// veyronJSBrowserIntegrationTest runs the veyron javascript integration test in a browser environment using nacl plugin.
func (t *testEnv) veyronJSBrowserIntegrationTest(ctx *util.Context, testName string) (*TestResult, error) {
	// TODO(aghassemi): Re-enable the test when it is fixed.
	return &TestResult{Status: TestPassed}, nil

	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-integration-browser"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["BROWSER_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSNodeIntegrationTest runs the veyron javascript integration test in NodeJS environment using wspr.
func (t *testEnv) veyronJSNodeIntegrationTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-integration-node"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSUnitTest runs the veyron javascript unit test.
func (t *testEnv) veyronJSUnitTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-unit"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSVdlTest runs the veyron javascript vdl test.
func (t *testEnv) veyronJSVdlTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron.js")
	target := "test-vdl"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}

// veyronJSVomTest runs the veyron javascript vom test.
func (t *testEnv) veyronJSVomTest(ctx *util.Context, testName string) (*TestResult, error) {
	root, err := util.VeyronRoot()
	if err != nil {
		return nil, err
	}
	testDir := filepath.Join(root, "veyron", "javascript", "vom")
	target := "test"
	env := map[string]string{}
	env["XUNIT"] = "true"
	env["NODE_OUTPUT"] = XUnitReportPath(testName)
	return t.runJSTest(ctx, testName, testDir, target, nil, env)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"

	log "github.com/Sirupsen/logrus"
)

type healthCheckHandler struct{}

func (_ *healthCheckHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func mustServeSnitch(addr string) {
	mux := http.NewServeMux()
	hc := &healthCheckHandler{}

	mux.Handle("/health_check", hc)
	log.WithFields(log.Fields{"address": addr}).Info("Serving snitch.")
	log.Fatal(http.ListenAndServe(addr, mux))
}
<commit_msg>Actually exposes the Prometheus metrics.<commit_after>package main

import (
	"fmt"
	"net/http"

	log "github.com/Sirupsen/logrus"
	"github.com/prometheus/client_golang/prometheus"
)

type healthCheckHandler struct{}

func (_ *healthCheckHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "OK")
}

func mustServeSnitch(addr string) {
	mux := http.NewServeMux()
	hc := &healthCheckHandler{}

	mux.Handle("/metrics", prometheus.Handler())
	mux.Handle("/health_check", hc)
	log.WithFields(log.Fields{"address": addr}).Info("Serving snitch.")
	log.Fatal(http.ListenAndServe(addr, mux))
}
<|endoftext|>


<commit_before>// +build linux

package ipvs

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"unsafe"

	"github.com/Sirupsen/logrus"
	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
)

var (
	native     = nl.NativeEndian()
	ipvsFamily int
	ipvsOnce   sync.Once
)

type genlMsgHdr struct {
	cmd      uint8
	version  uint8
	reserved uint16
}

type ipvsFlags struct {
	flags uint32
	mask  uint32
}

func deserializeGenlMsg(b []byte) (hdr *genlMsgHdr) {
	return (*genlMsgHdr)(unsafe.Pointer(&b[0:unsafe.Sizeof(*hdr)][0]))
}

func (hdr *genlMsgHdr) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*hdr)]byte)(unsafe.Pointer(hdr)))[:]
}

func (hdr *genlMsgHdr) Len() int {
	return int(unsafe.Sizeof(*hdr))
}

func (f *ipvsFlags) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*f)]byte)(unsafe.Pointer(f)))[:]
}

func (f *ipvsFlags) Len() int {
	return int(unsafe.Sizeof(*f))
}

func setup() {
	ipvsOnce.Do(func() {
		var err error
		if out, err := exec.Command("modprobe", "-va", "ip_vs").CombinedOutput(); err != nil {
			logrus.Warnf("Running modprobe ip_vs failed with message: `%s`, error: %v", strings.TrimSpace(string(out)), err)
		}

		ipvsFamily, err = getIPVSFamily()
		if err != nil {
			logrus.Error("Could not get ipvs family information from the kernel. It is possible that ipvs is not enabled in your kernel. Native loadbalancing will not work until this is fixed.")
		}
	})
}

func fillService(s *Service) nl.NetlinkRequestData {
	cmdAttr := nl.NewRtAttr(ipvsCmdAttrService, nil)
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrAddressFamily, nl.Uint16Attr(s.AddressFamily))
	if s.FWMark != 0 {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrFWMark, nl.Uint32Attr(s.FWMark))
	} else {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrProtocol, nl.Uint16Attr(s.Protocol))
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrAddress, rawIPData(s.Address))

		// Port needs to be in network byte order.
		portBuf := new(bytes.Buffer)
		binary.Write(portBuf, binary.BigEndian, s.Port)
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrPort, portBuf.Bytes())
	}

	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrSchedName, nl.ZeroTerminated(s.SchedName))
	if s.PEName != "" {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrPEName, nl.ZeroTerminated(s.PEName))
	}

	f := &ipvsFlags{
		flags: s.Flags,
		mask:  0xFFFFFFFF,
	}
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrFlags, f.Serialize())
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrTimeout, nl.Uint32Attr(s.Timeout))
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrNetmask, nl.Uint32Attr(s.Netmask))
	return cmdAttr
}

func fillDestinaton(d *Destination) nl.NetlinkRequestData {
	cmdAttr := nl.NewRtAttr(ipvsCmdAttrDest, nil)

	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrAddress, rawIPData(d.Address))
	// Port needs to be in network byte order.
	portBuf := new(bytes.Buffer)
	binary.Write(portBuf, binary.BigEndian, d.Port)
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrPort, portBuf.Bytes())

	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrForwardingMethod, nl.Uint32Attr(d.ConnectionFlags&ConnectionFlagFwdMask))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrWeight, nl.Uint32Attr(uint32(d.Weight)))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrUpperThreshold, nl.Uint32Attr(d.UpperThreshold))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrLowerThreshold, nl.Uint32Attr(d.LowerThreshold))

	return cmdAttr
}

func (i *Handle) doCmd(s *Service, d *Destination, cmd uint8) error {
	req := newIPVSRequest(cmd)
	req.Seq = atomic.AddUint32(&i.seq, 1)
	req.AddData(fillService(s))

	if d != nil {
		req.AddData(fillDestinaton(d))
	}

	if _, err := execute(i.sock, req, 0); err != nil {
		return err
	}

	return nil
}

func getIPVSFamily() (int, error) {
	sock, err := nl.GetNetlinkSocketAt(netns.None(), netns.None(), syscall.NETLINK_GENERIC)
	if err != nil {
		return 0, err
	}

	req := newGenlRequest(genlCtrlID, genlCtrlCmdGetFamily)
	req.AddData(nl.NewRtAttr(genlCtrlAttrFamilyName, nl.ZeroTerminated("IPVS")))

	msgs, err := execute(sock, req, 0)
	if err != nil {
		return 0, err
	}

	for _, m := range msgs {
		hdr := deserializeGenlMsg(m)
		attrs, err := nl.ParseRouteAttr(m[hdr.Len():])
		if err != nil {
			return 0, err
		}

		for _, attr := range attrs {
			switch int(attr.Attr.Type) {
			case genlCtrlAttrFamilyID:
				return int(native.Uint16(attr.Value[0:2])), nil
			}
		}
	}

	return 0, fmt.Errorf("no family id in the netlink response")
}

func rawIPData(ip net.IP) []byte {
	family := nl.GetIPFamily(ip)
	if family == nl.FAMILY_V4 {
		return ip.To4()
	}

	return ip
}

func newIPVSRequest(cmd uint8) *nl.NetlinkRequest {
	return newGenlRequest(ipvsFamily, cmd)
}

func newGenlRequest(familyID int, cmd uint8) *nl.NetlinkRequest {
	req := nl.NewNetlinkRequest(familyID, syscall.NLM_F_ACK)
	req.AddData(&genlMsgHdr{cmd: cmd, version: 1})
	return req
}

func execute(s *nl.NetlinkSocket, req *nl.NetlinkRequest, resType uint16) ([][]byte, error) {
	var (
		err error
	)

	if err := s.Send(req); err != nil {
		return nil, err
	}

	pid, err := s.GetPid()
	if err != nil {
		return nil, err
	}

	var res [][]byte

done:
	for {
		msgs, err := s.Receive()
		if err != nil {
			return nil, err
		}
		for _, m := range msgs {
			if m.Header.Seq != req.Seq {
				continue
			}
			if m.Header.Pid != pid {
				return nil, fmt.Errorf("Wrong pid %d, expected %d", m.Header.Pid, pid)
			}
			if m.Header.Type == syscall.NLMSG_DONE {
				break done
			}
			if m.Header.Type == syscall.NLMSG_ERROR {
				error := int32(native.Uint32(m.Data[0:4]))
				if error == 0 {
					break done
				}
				return nil, syscall.Errno(-error)
			}
			if resType != 0 && m.Header.Type != resType {
				continue
			}
			res = append(res, m.Data)
			if m.Header.Flags&syscall.NLM_F_MULTI == 0 {
				break done
			}
		}
	}
	return res, nil
}
<commit_msg>Do not leak ipvs netlink socket<commit_after>// +build linux

package ipvs

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"os/exec"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"unsafe"

	"github.com/Sirupsen/logrus"
	"github.com/vishvananda/netlink/nl"
	"github.com/vishvananda/netns"
)

var (
	native     = nl.NativeEndian()
	ipvsFamily int
	ipvsOnce   sync.Once
)

type genlMsgHdr struct {
	cmd      uint8
	version  uint8
	reserved uint16
}

type ipvsFlags struct {
	flags uint32
	mask  uint32
}

func deserializeGenlMsg(b []byte) (hdr *genlMsgHdr) {
	return (*genlMsgHdr)(unsafe.Pointer(&b[0:unsafe.Sizeof(*hdr)][0]))
}

func (hdr *genlMsgHdr) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*hdr)]byte)(unsafe.Pointer(hdr)))[:]
}

func (hdr *genlMsgHdr) Len() int {
	return int(unsafe.Sizeof(*hdr))
}

func (f *ipvsFlags) Serialize() []byte {
	return (*(*[unsafe.Sizeof(*f)]byte)(unsafe.Pointer(f)))[:]
}

func (f *ipvsFlags) Len() int {
	return int(unsafe.Sizeof(*f))
}

func setup() {
	ipvsOnce.Do(func() {
		var err error
		if out, err := exec.Command("modprobe", "-va", "ip_vs").CombinedOutput(); err != nil {
			logrus.Warnf("Running modprobe ip_vs failed with message: `%s`, error: %v", strings.TrimSpace(string(out)), err)
		}

		ipvsFamily, err = getIPVSFamily()
		if err != nil {
			logrus.Error("Could not get ipvs family information from the kernel. It is possible that ipvs is not enabled in your kernel. Native loadbalancing will not work until this is fixed.")
		}
	})
}

func fillService(s *Service) nl.NetlinkRequestData {
	cmdAttr := nl.NewRtAttr(ipvsCmdAttrService, nil)
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrAddressFamily, nl.Uint16Attr(s.AddressFamily))
	if s.FWMark != 0 {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrFWMark, nl.Uint32Attr(s.FWMark))
	} else {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrProtocol, nl.Uint16Attr(s.Protocol))
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrAddress, rawIPData(s.Address))

		// Port needs to be in network byte order.
		portBuf := new(bytes.Buffer)
		binary.Write(portBuf, binary.BigEndian, s.Port)
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrPort, portBuf.Bytes())
	}

	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrSchedName, nl.ZeroTerminated(s.SchedName))
	if s.PEName != "" {
		nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrPEName, nl.ZeroTerminated(s.PEName))
	}

	f := &ipvsFlags{
		flags: s.Flags,
		mask:  0xFFFFFFFF,
	}
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrFlags, f.Serialize())
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrTimeout, nl.Uint32Attr(s.Timeout))
	nl.NewRtAttrChild(cmdAttr, ipvsSvcAttrNetmask, nl.Uint32Attr(s.Netmask))
	return cmdAttr
}

func fillDestinaton(d *Destination) nl.NetlinkRequestData {
	cmdAttr := nl.NewRtAttr(ipvsCmdAttrDest, nil)

	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrAddress, rawIPData(d.Address))
	// Port needs to be in network byte order.
	portBuf := new(bytes.Buffer)
	binary.Write(portBuf, binary.BigEndian, d.Port)
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrPort, portBuf.Bytes())

	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrForwardingMethod, nl.Uint32Attr(d.ConnectionFlags&ConnectionFlagFwdMask))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrWeight, nl.Uint32Attr(uint32(d.Weight)))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrUpperThreshold, nl.Uint32Attr(d.UpperThreshold))
	nl.NewRtAttrChild(cmdAttr, ipvsDestAttrLowerThreshold, nl.Uint32Attr(d.LowerThreshold))

	return cmdAttr
}

func (i *Handle) doCmd(s *Service, d *Destination, cmd uint8) error {
	req := newIPVSRequest(cmd)
	req.Seq = atomic.AddUint32(&i.seq, 1)
	req.AddData(fillService(s))

	if d != nil {
		req.AddData(fillDestinaton(d))
	}

	if _, err := execute(i.sock, req, 0); err != nil {
		return err
	}

	return nil
}

func getIPVSFamily() (int, error) {
	sock, err := nl.GetNetlinkSocketAt(netns.None(), netns.None(), syscall.NETLINK_GENERIC)
	if err != nil {
		return 0, err
	}
	defer sock.Close()

	req := newGenlRequest(genlCtrlID, genlCtrlCmdGetFamily)
	req.AddData(nl.NewRtAttr(genlCtrlAttrFamilyName, nl.ZeroTerminated("IPVS")))

	msgs, err := execute(sock, req, 0)
	if err != nil {
		return 0, err
	}

	for _, m := range msgs {
		hdr := deserializeGenlMsg(m)
		attrs, err := nl.ParseRouteAttr(m[hdr.Len():])
		if err != nil {
			return 0, err
		}

		for _, attr := range attrs {
			switch int(attr.Attr.Type) {
			case genlCtrlAttrFamilyID:
				return int(native.Uint16(attr.Value[0:2])), nil
			}
		}
	}

	return 0, fmt.Errorf("no family id in the netlink response")
}

func rawIPData(ip net.IP) []byte {
	family := nl.GetIPFamily(ip)
	if family == nl.FAMILY_V4 {
		return ip.To4()
	}

	return ip
}

func newIPVSRequest(cmd uint8) *nl.NetlinkRequest {
	return newGenlRequest(ipvsFamily, cmd)
}

func newGenlRequest(familyID int, cmd uint8) *nl.NetlinkRequest {
	req := nl.NewNetlinkRequest(familyID, syscall.NLM_F_ACK)
	req.AddData(&genlMsgHdr{cmd: cmd, version: 1})
	return req
}

func execute(s *nl.NetlinkSocket, req *nl.NetlinkRequest, resType uint16) ([][]byte, error) {
	var (
		err error
	)

	if err := s.Send(req); err != nil {
		return nil, err
	}

	pid, err := s.GetPid()
	if err != nil {
		return nil, err
	}

	var res [][]byte

done:
	for {
		msgs, err := s.Receive()
		if err != nil {
			return nil, err
		}
		for _, m := range msgs {
			if m.Header.Seq != req.Seq {
				continue
			}
			if m.Header.Pid != pid {
				return nil, fmt.Errorf("Wrong pid %d, expected %d", m.Header.Pid, pid)
			}
			if m.Header.Type == syscall.NLMSG_DONE {
				break done
			}
			if m.Header.Type == syscall.NLMSG_ERROR {
				error := int32(native.Uint32(m.Data[0:4]))
				if error == 0 {
					break done
				}
				return nil, syscall.Errno(-error)
			}
			if resType != 0 && m.Header.Type != resType {
				continue
			}
			res = append(res, m.Data)
			if m.Header.Flags&syscall.NLM_F_MULTI == 0 {
				break done
			}
		}
	}
	return res, nil
}
<|endoftext|>