<commit_before>// +build fvtests

// Copyright (c) 2019 Tigera, Inc. 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 fv_test

import (
	"fmt"
	"strconv"

	"github.com/davecgh/go-spew/spew"
	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"

	"github.com/projectcalico/libcalico-go/lib/apiconfig"

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

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/infrastructure"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/felix/fv/workload"
	api "github.com/projectcalico/libcalico-go/lib/apis/v3"
	client "github.com/projectcalico/libcalico-go/lib/clientv3"
)

var _ = infrastructure.DatastoreDescribe("_BPF-NAT_ _BPF-SAFE_ BPF NAT tests", []apiconfig.DatastoreType{apiconfig.Kubernetes}, func(getInfra infrastructure.InfraFactory) {

	var (
		infra          infrastructure.DatastoreInfra
		felixes        []*infrastructure.Felix
		w              [2][2]*workload.Workload // 1st workload on each host
		hostW          [2]*workload.Workload
		client         client.Interface
		cc             *workload.ConnectivityChecker
		externalClient *containers.Container
	)

	BeforeEach(func() {
		var err error
		infra = getInfra()

		options := infrastructure.DefaultTopologyOptions()
		options.FelixLogSeverity = "debug"
		felixes, client = infrastructure.StartNNodeTopology(2, options, infra)
		cc = &workload.ConnectivityChecker{}

		// Start a host networked workload on each host for connectivity checks.
		for ii, felix := range felixes {
			// We tell each host-networked workload to open:
			// TODO: Copied from another test
			// - its normal (uninteresting) port, 8055
			// - port 2379, which is both an inbound and an outbound failsafe port
			// - port 22, which is an inbound failsafe port.
			// This allows us to test the interaction between do-not-track policy and failsafe
			// ports.
			const portsToOpen = "8055,2379,22"
			hostW[ii] = workload.Run(
				felixes[ii],
				fmt.Sprintf("host%d", ii),
				"default",
				felixes[ii].IP, // Same IP as felix means "run in the host's namespace"
				portsToOpen,
				"tcp")

			// Two workloads on each host so we can check the same host and other host cases.
			iiStr := strconv.Itoa(ii)
			w[ii][0] = workload.Run(felix, "w"+iiStr+"0", "default", "10.65."+iiStr+".2", "8055", "tcp")
			w[ii][0].WorkloadEndpoint.Labels = map[string]string{"name": w[ii][0].Name}
			w[ii][0].ConfigureInDatastore(infra)
			w[ii][1] = workload.Run(felix, "w"+iiStr+"1", "default", "10.65."+iiStr+".3", "8056", "tcp")
			w[ii][1].WorkloadEndpoint.Labels = map[string]string{"name": w[ii][1].Name}
			w[ii][1].ConfigureInDatastore(infra)
		}

		// We will use this container to model an external client trying to connect into
		// workloads on a host.  Create a route in the container for the workload CIDR.
		// TODO: Copied from another test
		externalClient = containers.Run("external-client",
			containers.RunOpts{AutoRemove: true},
			"--privileged", // So that we can add routes inside the container.
			utils.Config.BusyboxImage,
			"/bin/sh", "-c", "sleep 1000")
		_ = externalClient

		err = infra.AddDefaultDeny()
		Expect(err).To(BeNil())
	})

	JustAfterEach(func() {
		if CurrentGinkgoTestDescription().Failed {
			for _, felix := range felixes {
				felix.Exec("iptables-save", "-c")
				felix.Exec("ip", "r")
				felix.Exec("calico-bpf", "ipsets", "dump")
			}
		}
	})

	AfterEach(func() {
		for _, f := range felixes {
			f.Stop()
		}
		infra.Stop()
		externalClient.Stop()
	})

	It("should deny all by default", func() {
		// Same host, other workload.
		cc.ExpectNone(w[0][0], w[0][1])
		cc.ExpectNone(w[0][1], w[0][0])
		// Other host.
		cc.ExpectNone(w[0][0], w[1][0])
		cc.ExpectNone(w[1][0], w[0][0])
		cc.CheckConnectivity()
	})

	createPolicy := func(policy *api.GlobalNetworkPolicy) *api.GlobalNetworkPolicy {
		log.WithField("policy", dumpResource(policy)).Info("Creating policy")
		policy, err := client.GlobalNetworkPolicies().Create(utils.Ctx, policy, utils.NoOptions)
		Expect(err).NotTo(HaveOccurred())
		return policy
	}

	updatePolicy := func(policy *api.GlobalNetworkPolicy) *api.GlobalNetworkPolicy {
		log.WithField("policy", dumpResource(policy)).Info("Updating policy")
		policy, err := client.GlobalNetworkPolicies().Update(utils.Ctx, policy, utils.NoOptions)
		Expect(err).NotTo(HaveOccurred())
		return policy
	}
	_ = updatePolicy

	Context("with a policy allowing ingress to w[0][0]", func() {
		var (
			pol       *api.GlobalNetworkPolicy
			k8sClient *kubernetes.Clientset
		)

		BeforeEach(func() {
			pol = api.NewGlobalNetworkPolicy()
			pol.Namespace = "fv"
			pol.Name = "policy-1"
			pol.Spec.Ingress = []api.Rule{
				{
					Action: "Allow",
					Source: api.EntityRule{
						Selector: "all()",
					},
				},
			}
			pol.Spec.Egress = []api.Rule{
				{
					Action: "Allow",
					Source: api.EntityRule{
						Selector: "all()",
					},
				},
			}
			pol.Spec.Selector = "all()"

			pol = createPolicy(pol)

			k8sClient = infra.(*infrastructure.K8sDatastoreInfra).K8sClient
			_ = k8sClient
		})

		It("connectivity from all workloads via workload 0's main IP", func() {
			cc.ExpectSome(w[0][1], w[0][0])
			cc.ExpectSome(w[1][0], w[0][0])
			cc.ExpectSome(w[1][1], w[0][0])
			cc.CheckConnectivity()
		})

		Context("with nat configured 10.96.0.1:80 -> 10.65.0.1:8055 (w[0][0])", func() {
			BeforeEach(func() {
				testSvc := &v1.Service{
					TypeMeta:   typeMetaV1("Service"),
					ObjectMeta: objectMetaV1("test-service"),
					Spec: v1.ServiceSpec{
						ClusterIP: "10.101.0.10",
						Type:      v1.ServiceTypeClusterIP,
						Selector: map[string]string{
							"name": w[0][0].Name,
						},
						Ports: []v1.ServicePort{
							{
								Protocol:   v1.ProtocolTCP,
								Port:       80,
								Name:       "port-8055",
								TargetPort: intstr.FromInt(8055),
							},
						},
					},
				}

				_, err := k8sClient.CoreV1().Services("default").Create(testSvc)
				Expect(err).NotTo(HaveOccurred())

				getEndpointSubsets := func() []v1.EndpointSubset {
					ep, _ := k8sClient.CoreV1().Endpoints("default").Get("test-service", metav1.GetOptions{})
					log.WithField("endpoints", spew.Sprint(ep)).Info("Got endpoints for test-service")
					return ep.Subsets
				}
				Eventually(getEndpointSubsets, "10s").Should(HaveLen(1),
					"Service endpoints didn't get created? Is controller-manager happy?")
			})

			It("connectivity from all workloads via workload 0's NAT", func() {
				cc.ExpectSome(w[0][1], workload.IP("10.101.0.10"), 80)
				cc.ExpectSome(w[1][0], workload.IP("10.101.0.10"), 80)
				cc.ExpectSome(w[1][1], workload.IP("10.101.0.10"), 80)
				cc.CheckConnectivity()
			})
		})
	})
})

func typeMetaV1(kind string) metav1.TypeMeta {
	return metav1.TypeMeta{
		Kind:       kind,
		APIVersion: "v1",
	}
}

func objectMetaV1(name string) metav1.ObjectMeta {
	return metav1.ObjectMeta{
		Name:      name,
		Namespace: "default",
	}
}
<commit_msg>fv/bpf_test can read and compare NAT maps<commit_after>// +build fvtests

// Copyright (c) 2019 Tigera, Inc. 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 fv_test

import (
	"fmt"
	"strconv"

	"github.com/davecgh/go-spew/spew"
	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/kubernetes"

	"github.com/projectcalico/libcalico-go/lib/apiconfig"

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

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/projectcalico/felix/bpf"
	bpfm "github.com/projectcalico/felix/bpf/proxy/maps"
	"github.com/projectcalico/felix/fv/containers"
	"github.com/projectcalico/felix/fv/infrastructure"
	"github.com/projectcalico/felix/fv/utils"
	"github.com/projectcalico/felix/fv/workload"
	api "github.com/projectcalico/libcalico-go/lib/apis/v3"
	client "github.com/projectcalico/libcalico-go/lib/clientv3"
)

var _ = infrastructure.DatastoreDescribe("_BPF-NAT_ _BPF-SAFE_ BPF NAT tests", []apiconfig.DatastoreType{apiconfig.Kubernetes}, func(getInfra infrastructure.InfraFactory) {

	var (
		infra          infrastructure.DatastoreInfra
		felixes        []*infrastructure.Felix
		w              [2][2]*workload.Workload // 1st workload on each host
		hostW          [2]*workload.Workload
		client         client.Interface
		cc             *workload.ConnectivityChecker
		externalClient *containers.Container
	)

	BeforeEach(func() {
		var err error
		infra = getInfra()

		options := infrastructure.DefaultTopologyOptions()
		options.FelixLogSeverity = "debug"
		felixes, client = infrastructure.StartNNodeTopology(2, options, infra)
		cc = &workload.ConnectivityChecker{}

		// Start a host networked workload on each host for connectivity checks.
		for ii, felix := range felixes {
			// We tell each host-networked workload to open:
			// TODO: Copied from another test
			// - its normal (uninteresting) port, 8055
			// - port 2379, which is both an inbound and an outbound failsafe port
			// - port 22, which is an inbound failsafe port.
			// This allows us to test the interaction between do-not-track policy and failsafe
			// ports.
			const portsToOpen = "8055,2379,22"
			hostW[ii] = workload.Run(
				felixes[ii],
				fmt.Sprintf("host%d", ii),
				"default",
				felixes[ii].IP, // Same IP as felix means "run in the host's namespace"
				portsToOpen,
				"tcp")

			// Two workloads on each host so we can check the same host and other host cases.
			iiStr := strconv.Itoa(ii)
			w[ii][0] = workload.Run(felix, "w"+iiStr+"0", "default", "10.65."+iiStr+".2", "8055", "tcp")
			w[ii][0].WorkloadEndpoint.Labels = map[string]string{"name": w[ii][0].Name}
			w[ii][0].ConfigureInDatastore(infra)
			w[ii][1] = workload.Run(felix, "w"+iiStr+"1", "default", "10.65."+iiStr+".3", "8056", "tcp")
			w[ii][1].WorkloadEndpoint.Labels = map[string]string{"name": w[ii][1].Name}
			w[ii][1].ConfigureInDatastore(infra)
		}

		// We will use this container to model an external client trying to connect into
		// workloads on a host.  Create a route in the container for the workload CIDR.
		// TODO: Copied from another test
		externalClient = containers.Run("external-client",
			containers.RunOpts{AutoRemove: true},
			"--privileged", // So that we can add routes inside the container.
			utils.Config.BusyboxImage,
			"/bin/sh", "-c", "sleep 1000")
		_ = externalClient

		err = infra.AddDefaultDeny()
		Expect(err).To(BeNil())
	})

	JustAfterEach(func() {
		if CurrentGinkgoTestDescription().Failed {
			currBpfsvcs, currBpfeps := dumpNATmaps(felixes)

			for i, felix := range felixes {
				felix.Exec("iptables-save", "-c")
				felix.Exec("ip", "r")
				felix.Exec("calico-bpf", "ipsets", "dump")
				log.Infof("[%d]NATMap: %+v", i, currBpfsvcs[i])
				log.Infof("[%d]NATBackend: %+v", i, currBpfeps[i])
			}
		}
	})

	AfterEach(func() {
		for _, f := range felixes {
			f.Stop()
		}
		infra.Stop()
		externalClient.Stop()
	})

	It("should deny all by default", func() {
		// Same host, other workload.
		cc.ExpectNone(w[0][0], w[0][1])
		cc.ExpectNone(w[0][1], w[0][0])
		// Other host.
		cc.ExpectNone(w[0][0], w[1][0])
		cc.ExpectNone(w[1][0], w[0][0])
		cc.CheckConnectivity()
	})

	createPolicy := func(policy *api.GlobalNetworkPolicy) *api.GlobalNetworkPolicy {
		log.WithField("policy", dumpResource(policy)).Info("Creating policy")
		policy, err := client.GlobalNetworkPolicies().Create(utils.Ctx, policy, utils.NoOptions)
		Expect(err).NotTo(HaveOccurred())
		return policy
	}

	updatePolicy := func(policy *api.GlobalNetworkPolicy) *api.GlobalNetworkPolicy {
		log.WithField("policy", dumpResource(policy)).Info("Updating policy")
		policy, err := client.GlobalNetworkPolicies().Update(utils.Ctx, policy, utils.NoOptions)
		Expect(err).NotTo(HaveOccurred())
		return policy
	}
	_ = updatePolicy

	Context("with a policy allowing ingress to w[0][0]", func() {
		var (
			pol       *api.GlobalNetworkPolicy
			k8sClient *kubernetes.Clientset
		)

		BeforeEach(func() {
			pol = api.NewGlobalNetworkPolicy()
			pol.Namespace = "fv"
			pol.Name = "policy-1"
			pol.Spec.Ingress = []api.Rule{
				{
					Action: "Allow",
					Source: api.EntityRule{
						Selector: "all()",
					},
				},
			}
			pol.Spec.Egress = []api.Rule{
				{
					Action: "Allow",
					Source: api.EntityRule{
						Selector: "all()",
					},
				},
			}
			pol.Spec.Selector = "all()"

			pol = createPolicy(pol)

			k8sClient = infra.(*infrastructure.K8sDatastoreInfra).K8sClient
			_ = k8sClient
		})

		It("connectivity from all workloads via workload 0's main IP", func() {
			cc.ExpectSome(w[0][1], w[0][0])
			cc.ExpectSome(w[1][0], w[0][0])
			cc.ExpectSome(w[1][1], w[0][0])
			cc.CheckConnectivity()
		})

		var testSvc *v1.Service
		var getTestSvcEps func() []v1.EndpointSubset

		BeforeEach(func() {
			testSvc = &v1.Service{
				TypeMeta:   typeMetaV1("Service"),
				ObjectMeta: objectMetaV1("test-service"),
				Spec: v1.ServiceSpec{
					ClusterIP: "10.101.0.10",
					Type:      v1.ServiceTypeClusterIP,
					Selector: map[string]string{
						"name": w[0][0].Name,
					},
					Ports: []v1.ServicePort{
						{
							Protocol:   v1.ProtocolTCP,
							Port:       80,
							Name:       "port-8055",
							TargetPort: intstr.FromInt(8055),
						},
					},
				},
			}

			getTestSvcEps = func() []v1.EndpointSubset {
				ep, _ := k8sClient.CoreV1().
					Endpoints(testSvc.ObjectMeta.Namespace).
					Get(testSvc.ObjectMeta.Name, metav1.GetOptions{})
				log.WithField("endpoints",
					spew.Sprint(ep)).Infof("Got endpoints for %s", testSvc.ObjectMeta.Name)
				return ep.Subsets
			}
		})

		Context("with test-service configured 10.101.0.10:80 -> w[0][0].IP:8055", func() {
			BeforeEach(func() {
				_, err := k8sClient.CoreV1().Services(testSvc.ObjectMeta.Namespace).Create(testSvc)
				Expect(err).NotTo(HaveOccurred())

				getEndpointSubsets := func() []v1.EndpointSubset {
					ep, _ := k8sClient.CoreV1().Endpoints("default").Get("test-service", metav1.GetOptions{})
					log.WithField("endpoints", spew.Sprint(ep)).Info("Got endpoints for test-service")
					return ep.Subsets
				}
				Eventually(getEndpointSubsets, "10s").Should(HaveLen(1),
					"Service endpoints didn't get created? Is controller-manager happy?")
			})

			It("should have connectivity from all workloads via a service to workload 0", func() {
				ip := testSvc.Spec.ClusterIP
				port := uint16(testSvc.Spec.Ports[0].Port)

				cc.ExpectSome(w[0][1], workload.IP(ip), port)
				cc.ExpectSome(w[1][0], workload.IP(ip), port)
				cc.ExpectSome(w[1][1], workload.IP(ip), port)
				cc.CheckConnectivity()

			})

			Context("with test-service removed", func() {
				var (
					prevBpfsvcs []bpfm.NATMapMem
					prevBpfeps  []bpfm.NATBackendMapMem
				)

				BeforeEach(func() {
					prevBpfsvcs, prevBpfeps = dumpNATmaps(felixes)
					err := k8sClient.CoreV1().
						Services(testSvc.ObjectMeta.Namespace).
						Delete(testSvc.ObjectMeta.Name, &metav1.DeleteOptions{})
					Expect(err).NotTo(HaveOccurred())
					Eventually(getTestSvcEps, "10s").Should(HaveLen(0))
				})

				It("should not have connectivity from workloads via a service to workload 0", func() {
					ip := testSvc.Spec.ClusterIP
					port := uint16(testSvc.Spec.Ports[0].Port)

					cc.ExpectNone(w[0][1], workload.IP(ip), port)
					cc.ExpectNone(w[1][0], workload.IP(ip), port)
					cc.ExpectNone(w[1][1], workload.IP(ip), port)
					cc.CheckConnectivity()

					natmaps, natbacks := dumpNATmaps(felixes)
					for i := range felixes {
						Expect(natmaps[i]).To(HaveLen(len(prevBpfsvcs[i]) - 1))
						Expect(natbacks[i]).To(HaveLen(len(prevBpfeps[i]) - 1))
					}
				})
			})
		})
	})
})

func typeMetaV1(kind string) metav1.TypeMeta {
	return metav1.TypeMeta{
		Kind:       kind,
		APIVersion: "v1",
	}
}

func objectMetaV1(name string) metav1.ObjectMeta {
	return metav1.ObjectMeta{
		Name:      name,
		Namespace: "default",
	}
}

func dumpNATmaps(felixes []*infrastructure.Felix) ([]bpfm.NATMapMem, []bpfm.NATBackendMapMem) {
	bpfsvcs := make([]bpfm.NATMapMem, len(felixes))
	bpfeps := make([]bpfm.NATBackendMapMem, len(felixes))

	for i, felix := range felixes {
		// just one map to get the file names and cmds, it is the same for all felixes
		bm := bpfm.NATMap()
		cmd, err := bpf.DumpMapCmd(bm)
		Expect(err).NotTo(HaveOccurred())
		bpfsvcs[i] = make(bpfm.NATMapMem)
		out, err := felix.ExecOutput(cmd...)
		Expect(err).NotTo(HaveOccurred())
		err = bpf.IterMapCmdOutput([]byte(out), bpfm.NATMapMemIter(bpfsvcs[i]))
		Expect(err).NotTo(HaveOccurred())

		bb := bpfm.BackendMap()
		cmd, err = bpf.DumpMapCmd(bb)
		Expect(err).NotTo(HaveOccurred())
		bpfeps[i] = make(bpfm.NATBackendMapMem)
		out, err = felix.ExecOutput(cmd...)
		Expect(err).NotTo(HaveOccurred())
		err = bpf.IterMapCmdOutput([]byte(out), bpfm.NATBackendMapMemIter(bpfeps[i]))
		Expect(err).NotTo(HaveOccurred())
	}

	return bpfsvcs, bpfeps
}
<|endoftext|>


<commit_before>package s3utils

import (
	"log"
	"path"
	"strings"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

// An s3 input looks like: s3://bucket/dir
// Where dir can be a path

// getBucket extracts the bucket from an s3 input
func GetBucket(input string) (string, error) {
	return strings.Split(strings.TrimPrefix(input, "s3://"), "/")[0], nil
}

// getPath extracts the path from an s3 input
func GetPath(input string) (string, error) {
	return path.Join(strings.Split(strings.TrimPrefix(input, "s3://"), "/")[1:]...), nil
}

func NewBucket(bucketName string) (*s3.Bucket, error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		log.Print(err)
		return nil, err
	}
	client := s3.New(auth, aws.USWest)

	return client.Bucket(bucketName), nil
}
<commit_msg>Adds a PutMulti function which wraps goamz multiputs.<commit_after>package s3utils

import (
	"bytes"
	"io"
	"log"
	"path"
	"strings"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

const (
	// AWS says that parts must be at least 5MB, it's unclear if that means 5 *
	// 10^6 or 5 2^10 so we went with the larger.
	minPart = 5242880 // 5MB
)

// An s3 input looks like: s3://bucket/dir
// Where dir can be a path

// getBucket extracts the bucket from an s3 input
func GetBucket(input string) (string, error) {
	return strings.Split(strings.TrimPrefix(input, "s3://"), "/")[0], nil
}

// getPath extracts the path from an s3 input
func GetPath(input string) (string, error) {
	return path.Join(strings.Split(strings.TrimPrefix(input, "s3://"), "/")[1:]...), nil
}

func NewBucket(bucketName string) (*s3.Bucket, error) {
	auth, err := aws.EnvAuth()
	if err != nil {
		log.Print(err)
		return nil, err
	}
	client := s3.New(auth, aws.USWest)

	return client.Bucket(bucketName), nil
}

// PutMulti is like a smart bucket.Put in that it will automatically do a
// multiput if the input reader has enough data that it makes sense to do so.
func PutMulti(bucket *s3.Bucket, path string, r io.Reader, contType string, perm s3.ACL) error {
	// A pointer to a Multi transaction, if this is non nil it means we're
	// putting the data in parts.
	var multi *s3.Multi = nil
	var parts []s3.Part
	for i := 0; ; i++ {
		var data []byte
		n, err := io.ReadAtLeast(r, data, minPart)
		if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
			// EOF means that r was empty
			// ErrUnexpectedEOF means that r had less than minPart bytes in it
			// Note that we can hit these errors after having read several 5MB
			// chunks from r already
			return err
		}
		if n >= minPart && multi == nil {
			// We have enough data to do a MultiPut but we haven't started on yet.
			// That means it's time to start one
			multi, err = bucket.Multi(path, contType, perm)
			if err != nil {
				return err
			}
		}
		// Now we upload the actual data
		if multi == nil {
			// We're not doing a multi transaction,
			return bucket.Put(path, data, contType, perm)
		} else {
			part, err := multi.PutPart(i, bytes.NewReader(data))
			if err != nil {
				return err
			}
			parts = append(parts, part)
		}
		if n < minPart {
			// That means this was the last batch of data, time to break out of
			// this loop
			break
		}
	}
	if multi != nil {
		if err := multi.Complete(parts); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Code for genetic algorithms
package genetic

import (
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/kevinburke/rct/tracks"
)

// constants which may be altered to affect the ride runtime
const PARENTS = 2
const FIT_PERCENTAGE = 0.2
const MUTATION_RATE = 0.05

// crossover with a probability of 0.6 (taken from the book & De Jong 1975)
const CROSSOVER_PROBABILITY = 0.6
const POOL_SIZE = 500
const ITERATIONS = 10
const PRINT_RESULTS_EVERY = 1

// create a directory and ignore "directory exists" errors
func mkdir(name string) error {
	err := os.Mkdir(name, 0755)
	// ugh
	if err != nil && err.(*os.PathError).Err.Error() != "file exists" {
		return err
	}
	return nil
}

func encode(name string, v interface{}) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(v)
}

type ExperimentMetadata struct {
	Hash  string
	Date  time.Time
	Notes string
}

func Run(outputDirectory string, packageRoot string) error {
	expDir := path.Join(outputDirectory, "experiments")
	err := mkdir(expDir)
	if err != nil {
		return err
	}
	id := fmt.Sprintf("exp_%s", uuid.New())
	expIdDir := path.Join(expDir, id)
	err = mkdir(expIdDir)
	if err != nil {
		return err
	}
	iterationsDir := path.Join(expIdDir, "iterations")
	err = mkdir(iterationsDir)
	if err != nil {
		return err
	}
	cmd := exec.Command("git", "rev-parse", "HEAD")
	cmd.Dir = packageRoot
	hashb, err := cmd.Output()
	mtd := ExperimentMetadata{
		Hash:  strings.TrimSpace(string(hashb)),
		Date:  time.Now().UTC(),
		Notes: "(none)", // XXX
	}
	metadataPath := path.Join(expIdDir, "meta.json")
	err = encode(metadataPath, mtd)
	if err != nil {
		return err
	}
	fmt.Printf("Experiment %s\n======================================\n", id)
	pool := SeedPool(POOL_SIZE)
	pool.Id = id
	for i := 0; i < ITERATIONS; i++ {
		pool = pool.Crossover()
		pool.Mutate(MUTATION_RATE)
		pool.Evaluate()
		iterationDir := path.Join(iterationsDir, strconv.Itoa(i))
		err = mkdir(iterationDir)
		if err != nil {
			return err
		}
		pool.Statistics(i, outputDirectory)
	}
	return nil
}

type Pool struct {
	Id      string
	Members []*Member
}

type Member struct {
	Id    string
	Track []tracks.Element
	Score int64
	// Advantage or disadvantage in reproducing
	Fitness float64
}

func (p *Pool) Statistics(iteration int, outputDirectory string) {
	if iteration%PRINT_RESULTS_EVERY == 0 {
		var highestScore int64 = -1
		var worstScore int64 = 7000
		bestMember := new(Member)
		for i := 0; i < len(p.Members); i++ {
			if p.Members[i].Score > highestScore {
				highestScore = p.Members[i].Score
				bestMember = p.Members[i]
			}
			if p.Members[i].Score < worstScore {
				worstScore = p.Members[i].Score
			}
		}
		fmt.Printf("Iteration %d: %d members, best member has score %d, worst has score %d\n", iteration, len(p.Members), bestMember.Score, worstScore)
	}
	// XXX, move offline to a goroutine
	for i := 0; i < len(p.Members); i++ {
		pth := path.Join(outputDirectory, "experiments", p.Id, "iterations", strconv.Itoa(iteration), p.Members[i].Id)
		err := encode(pth, p.Members[i])
		if err != nil {
			log.Print(err)
		}
	}
}

// Create an initial pool
func SeedPool(size int) *Pool {
	// 1. Create a station of length 10
	// 2. For the start station piece, generate a list of possible pieces.
	// 3. Choose one at random. Advance a pointer one forward.
	// 4. Repeat for 50 pieces (Woodchip is length 108. Mischief is 123)
	members := make([]*Member, POOL_SIZE)
	for i := 0; i < POOL_SIZE; i++ {
		track := CreateStation()
		idx := STATION_LENGTH - 1
		for j := 0; j < INITIAL_TRACK_LENGTH-STATION_LENGTH; j++ {
			poss := track[idx].Possibilities()
			track = append(track, poss[rand.Intn(len(poss))])
		}
		score := GetScore(track)
		members[i] = &Member{
			Id:    fmt.Sprintf("iter_%s", uuid.New()),
			Track: track,
			Score: score,
		}
	}
	return &Pool{Members: members}
}

func (p *Pool) Mutate(rate float64) {
	// for each ride:
	// for each position in the ride:
	// add in a possibility of mutation - addition
	// if addition:
	//	find a piece that is compatible with both ends.
	//  if no piece:
	//	  advance to the next track piece and try again
}

// Assign scores for every member of the pool
func (p *Pool) Evaluate() {
	for i := 0; i < POOL_SIZE; i++ {
		p.Members[i].Score = GetScore(p.Members[i].Track)
	}

	// Assign fitness for every member. For now, every member gets a fitness of
	// 1. In the future, consider sorting/giving higher score members a better
	// chance of reproducing.
	for i := 0; i < POOL_SIZE; i++ {
		p.Members[i].Fitness = float64(p.Members[i].Score)
	}
}

// Select chooses a member of the population at random
func (p *Pool) Select() (int, *Member) {
	// Stupid dumb version of this taken from here:
	// http://eli.thegreenplace.net/2010/01/22/weighted-random-generation-in-python
	// If it's a bottleneck, rewrite it.
	var weightedTotal float64 = 0
	totals := make([]float64, len(p.Members))
	for i := 0; i < len(p.Members); i++ {
		weightedTotal += p.Members[i].Fitness
		totals[i] = weightedTotal
	}
	rnd := rand.Float64() * weightedTotal
	for index, element := range totals {
		if rnd < element {
			return index, p.Members[index]
		}
	}
	return -1, &Member{}
}

func min(a int, b int) int {
	if a < b {
		return a
	}
	return b
}

func crossoverOne(parent1 *Member, parent2 *Member) (*Member, *Member) {
	//	choose a random point between the beginning and the end
	minval := min(len(parent1.Track), len(parent2.Track))
	crossPoint1 := rand.Intn(minval)
	crossPoint2 := crossPoint1
	foundMatch := false
	for {
		if tracks.Compatible(parent1.Track[crossPoint1], parent2.Track[crossPoint2]) {
			foundMatch = true
			break
		}
		crossPoint1++
		if crossPoint1 >= len(parent1.Track) {
			break
		}
		if tracks.Compatible(parent1.Track[crossPoint1], parent2.Track[crossPoint2]) {
			foundMatch = true
			break
		}
		crossPoint2++
		if crossPoint2 >= len(parent2.Track) {
			break
		}
	}
	//	swap the track pieces at the chosen point on track A and track B
	if foundMatch {
		return Swap(parent1, parent2, crossPoint1, crossPoint2)
	}
	return parent1, parent2
}

// Crossover chooses two members of a pool and joins them at random.
func (p *Pool) Crossover() *Pool {
	halfLen := len(p.Members) / 2
	for i := 0; i < halfLen; i++ {
		// select 2 parents at random
		idx1, parent1 := p.Select()
		idx2, parent2 := p.Select()
		if idx1 == -1 || idx2 == -1 {
			continue
		}
		if rand.Float64() < CROSSOVER_PROBABILITY {
			// XXX delete parents
			child1, child2 := crossoverOne(parent1, parent2)
			p.Members[idx1] = child1
			p.Members[idx2] = child2
		}
	}
	return p
}

// Swap creates two children out of the parents, by crossing over the tracks at
// the given cross points. The sum of the two track lengths may be the same,
// but the tracks themselves will change.
func Swap(parent1 *Member, parent2 *Member, crossPoint1 int, crossPoint2 int) (*Member, *Member) {
	child1len := crossPoint1 + (len(parent2.Track) - crossPoint1)
	child2len := crossPoint2 + (len(parent1.Track) - crossPoint2)
	// XXX, probably something fancy you can do with xor, or a temporary array.
	child1track := make([]tracks.Element, child1len)
	child2track := make([]tracks.Element, child2len)
	copy(child1track[:crossPoint1], parent1.Track[:crossPoint1])
	copy(child1track[crossPoint1:], parent2.Track[crossPoint1:])
	copy(child2track[:crossPoint2], parent2.Track[:crossPoint2])
	copy(child2track[crossPoint2:], parent1.Track[crossPoint2:])
	child1 := &Member{
		Id:    fmt.Sprintf("iter_%s", uuid.New()),
		Track: child1track,
	}
	child2 := &Member{
		Id:    fmt.Sprintf("iter_%s", uuid.New()),
		Track: child2track,
	}
	return child1, child2
}

func (p *Pool) Spawn(numParents int) *Pool {
	return nil
}
<commit_msg>slight formatting<commit_after>// Code for genetic algorithms
package genetic

import (
	"encoding/json"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/kevinburke/rct/tracks"
)

// constants which may be altered to affect the ride runtime
const PARENTS = 2
const FIT_PERCENTAGE = 0.2
const MUTATION_RATE = 0.05

// crossover with a probability of 0.6 (taken from the book & De Jong 1975)
const CROSSOVER_PROBABILITY = 0.6
const POOL_SIZE = 500
const ITERATIONS = 10
const PRINT_RESULTS_EVERY = 1

// create a directory and ignore "directory exists" errors
func mkdir(name string) error {
	err := os.Mkdir(name, 0755)
	// ugh
	if err != nil && err.(*os.PathError).Err.Error() != "file exists" {
		return err
	}
	return nil
}

func encode(name string, v interface{}) error {
	f, err := os.OpenFile(name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	enc := json.NewEncoder(f)
	return enc.Encode(v)
}

type ExperimentMetadata struct {
	Hash    string
	Date    time.Time
	Notes   string
	Runtime time.Duration
}

func Run(outputDirectory string, packageRoot string) error {
	expDir := path.Join(outputDirectory, "experiments")
	err := mkdir(expDir)
	if err != nil {
		return err
	}
	id := fmt.Sprintf("exp_%s", uuid.New())
	expIdDir := path.Join(expDir, id)
	err = mkdir(expIdDir)
	if err != nil {
		return err
	}
	iterationsDir := path.Join(expIdDir, "iterations")
	err = mkdir(iterationsDir)
	if err != nil {
		return err
	}
	cmd := exec.Command("git", "rev-parse", "HEAD")
	cmd.Dir = packageRoot
	hashb, err := cmd.Output()
	mtd := ExperimentMetadata{
		Hash:  strings.TrimSpace(string(hashb)),
		Date:  time.Now().UTC(),
		Notes: "(none)", // XXX
	}
	metadataPath := path.Join(expIdDir, "meta.json")
	err = encode(metadataPath, mtd)
	if err != nil {
		return err
	}
	fmt.Printf("Experiment %s\n======================================\n", id)
	pool := SeedPool(POOL_SIZE)
	pool.Id = id
	for i := 0; i < ITERATIONS; i++ {
		pool = pool.Crossover()
		pool.Mutate(MUTATION_RATE)
		pool.Evaluate()
		iterationDir := path.Join(iterationsDir, strconv.Itoa(i))
		err = mkdir(iterationDir)
		if err != nil {
			return err
		}
		pool.Statistics(i, outputDirectory)
	}
	return nil
}

type Pool struct {
	Id      string
	Members []*Member
}

type Member struct {
	Id    string
	Score int64
	// Advantage or disadvantage in reproducing
	Fitness float64
	Runtime time.Duration
	Track   []tracks.Element
}

func (p *Pool) Statistics(iteration int, outputDirectory string) {
	if iteration%PRINT_RESULTS_EVERY == 0 {
		var highestScore int64 = -1
		var worstScore int64 = 7000
		bestMember := new(Member)
		for i := 0; i < len(p.Members); i++ {
			if p.Members[i].Score > highestScore {
				highestScore = p.Members[i].Score
				bestMember = p.Members[i]
			}
			if p.Members[i].Score < worstScore {
				worstScore = p.Members[i].Score
			}
		}
		fmt.Printf("Iteration %d: %d members, best member has score %d, worst has score %d\n", iteration, len(p.Members), bestMember.Score, worstScore)
	}
	// XXX, move offline to a goroutine
	for i := 0; i < len(p.Members); i++ {
		pth := path.Join(outputDirectory, "experiments", p.Id,
			"iterations", strconv.Itoa(iteration), p.Members[i].Id, ".json")
		err := encode(pth, p.Members[i])
		if err != nil {
			log.Print(err)
		}
	}
}

// Create an initial pool
func SeedPool(size int) *Pool {
	// 1. Create a station of length 10
	// 2. For the start station piece, generate a list of possible pieces.
	// 3. Choose one at random. Advance a pointer one forward.
	// 4. Repeat for 50 pieces (Woodchip is length 108. Mischief is 123)
	members := make([]*Member, POOL_SIZE)
	for i := 0; i < POOL_SIZE; i++ {
		track := CreateStation()
		idx := STATION_LENGTH - 1
		for j := 0; j < INITIAL_TRACK_LENGTH-STATION_LENGTH; j++ {
			poss := track[idx].Possibilities()
			track = append(track, poss[rand.Intn(len(poss))])
		}
		score := GetScore(track)
		members[i] = &Member{
			Id:    fmt.Sprintf("iter_%s", uuid.New()),
			Track: track,
			Score: score,
		}
	}
	return &Pool{Members: members}
}

func (p *Pool) Mutate(rate float64) {
	// for each ride:
	// for each position in the ride:
	// add in a possibility of mutation - addition
	// if addition:
	//	find a piece that is compatible with both ends.
	//  if no piece:
	//	  advance to the next track piece and try again
}

// Assign scores for every member of the pool
func (p *Pool) Evaluate() {
	for i := 0; i < POOL_SIZE; i++ {
		p.Members[i].Score = GetScore(p.Members[i].Track)
	}

	// Assign fitness for every member. For now, every member gets a fitness of
	// 1. In the future, consider sorting/giving higher score members a better
	// chance of reproducing.
	for i := 0; i < POOL_SIZE; i++ {
		p.Members[i].Fitness = float64(p.Members[i].Score)
	}
}

// Select chooses a member of the population at random
func (p *Pool) Select() (int, *Member) {
	// Stupid dumb version of this taken from here:
	// http://eli.thegreenplace.net/2010/01/22/weighted-random-generation-in-python
	// If it's a bottleneck, rewrite it.
	var weightedTotal float64 = 0
	totals := make([]float64, len(p.Members))
	for i := 0; i < len(p.Members); i++ {
		weightedTotal += p.Members[i].Fitness
		totals[i] = weightedTotal
	}
	rnd := rand.Float64() * weightedTotal
	for index, element := range totals {
		if rnd < element {
			return index, p.Members[index]
		}
	}
	return -1, &Member{}
}

func min(a int, b int) int {
	if a < b {
		return a
	}
	return b
}

func crossoverOne(parent1 *Member, parent2 *Member) (*Member, *Member) {
	//	choose a random point between the beginning and the end
	minval := min(len(parent1.Track), len(parent2.Track))
	crossPoint1 := rand.Intn(minval)
	crossPoint2 := crossPoint1
	foundMatch := false
	for {
		if tracks.Compatible(parent1.Track[crossPoint1], parent2.Track[crossPoint2]) {
			foundMatch = true
			break
		}
		crossPoint1++
		if crossPoint1 >= len(parent1.Track) {
			break
		}
		if tracks.Compatible(parent1.Track[crossPoint1], parent2.Track[crossPoint2]) {
			foundMatch = true
			break
		}
		crossPoint2++
		if crossPoint2 >= len(parent2.Track) {
			break
		}
	}
	//	swap the track pieces at the chosen point on track A and track B
	if foundMatch {
		return Swap(parent1, parent2, crossPoint1, crossPoint2)
	}
	return parent1, parent2
}

// Crossover chooses two members of a pool and joins them at random.
func (p *Pool) Crossover() *Pool {
	halfLen := len(p.Members) / 2
	for i := 0; i < halfLen; i++ {
		// select 2 parents at random
		idx1, parent1 := p.Select()
		idx2, parent2 := p.Select()
		if idx1 == -1 || idx2 == -1 {
			continue
		}
		if rand.Float64() < CROSSOVER_PROBABILITY {
			// XXX delete parents
			child1, child2 := crossoverOne(parent1, parent2)
			p.Members[idx1] = child1
			p.Members[idx2] = child2
		}
	}
	return p
}

// Swap creates two children out of the parents, by crossing over the tracks at
// the given cross points. The sum of the two track lengths may be the same,
// but the tracks themselves will change.
func Swap(parent1 *Member, parent2 *Member, crossPoint1 int, crossPoint2 int) (*Member, *Member) {
	child1len := crossPoint1 + (len(parent2.Track) - crossPoint1)
	child2len := crossPoint2 + (len(parent1.Track) - crossPoint2)
	// XXX, probably something fancy you can do with xor, or a temporary array.
	child1track := make([]tracks.Element, child1len)
	child2track := make([]tracks.Element, child2len)
	copy(child1track[:crossPoint1], parent1.Track[:crossPoint1])
	copy(child1track[crossPoint1:], parent2.Track[crossPoint1:])
	copy(child2track[:crossPoint2], parent2.Track[:crossPoint2])
	copy(child2track[crossPoint2:], parent1.Track[crossPoint2:])
	child1 := &Member{
		Id:    fmt.Sprintf("iter_%s", uuid.New()),
		Track: child1track,
	}
	child2 := &Member{
		Id:    fmt.Sprintf("iter_%s", uuid.New()),
		Track: child2track,
	}
	return child1, child2
}

func (p *Pool) Spawn(numParents int) *Pool {
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"math"
	"os"
	"path"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/mattn/go-gtk/gdk"
	"github.com/mattn/go-gtk/gtk"
	"github.com/naoina/toml"
)

//Conf structure contains configuration
type Conf struct {
	window   *gtk.Window `toml:",omitempty"`
	filename string      `toml:",omitempty"`

	UI struct {
		MenuBarVisible   bool `toml:"menubar-visible"`
		StatusBarVisible bool `toml:"statusbar-visible"`
	}
	TextView struct {
		Font           string `toml:"font"`
		LineHightlight bool   `toml:"line-hightlight"`
		LineNumbers    bool   `toml:"line-numbers"`
		WordWrap       bool   `toml:"word-wrap"`
		IndentSpace    bool   `toml:"indent-space"`
		IndentWidth    int    `toml:"indent-width"`
		StyleScheme    string `toml:"style-scheme"`
	}
	Tabs struct {
		Homogeneous bool  `toml:"homogeneous"`
		CloseBtns   bool  `toml:"close-buttons"`
		Height      int   `toml:"height"`
		FGNormal    []int `toml:"fg-normal"`
		FGModified  []int `toml:"fg-modified"`
		FGNew       []int `toml:"fg-new"`
	}
	Search struct {
		MaxItems int `toml:"max-items"`
	}
	Hex struct {
		BytesInLine int `toml:"bytes-in-line"`
	}
}

//NewConf set default values for configuration and parse config file
func NewConf() *Conf {
	// default values
	c := new(Conf)
	c.UI.MenuBarVisible = true
	c.UI.StatusBarVisible = false

	c.TextView.Font = "Liberation Mono 8"
	c.TextView.LineHightlight = true
	c.TextView.LineNumbers = true
	c.TextView.WordWrap = true
	c.TextView.IndentSpace = false
	c.TextView.IndentWidth = 2

	c.Tabs.Homogeneous = true
	c.Tabs.CloseBtns = true
	c.Tabs.Height = 24
	c.Tabs.FGNormal = []int{200, 200, 200}
	c.Tabs.FGModified = []int{220, 20, 20}
	c.Tabs.FGNew = []int{250, 200, 10}

	c.Search.MaxItems = 1024

	c.Hex.BytesInLine = 16

	//parse config files
	for _, configfile := range []string{
		path.Join(os.Getenv("XDG_CONFIG_HOME"), "goatee", "goatee.conf"),
		path.Join(os.Getenv("HOME"), ".config", "goatee", "goatee.conf"),
		"goatee.conf",
	} {
		data, err := ioutil.ReadFile(configfile)
		if err != nil {
			continue
		}

		err = toml.Unmarshal(data, &c)
		if err != nil {
			log.Printf("failed decode config file '%s', reason: %s", configfile, err)
			continue
		}

		c.filename = configfile
		break
	}

	return c
}

func (c *Conf) Write() {
	c.CreateDirConfig()

	log.Println("write", c.filename)

	if c.filename == "" {
		return
	}

	data, err := toml.Marshal(&c)
	if err != nil {
		log.Println(err)
		return
	}

	err = ioutil.WriteFile(c.filename, data, 0644)
	if err != nil {
		log.Println(err)
		return
	}
}

func (c *Conf) CreateDirConfig() {
	//if file already created nothing do
	if len(c.filename) > 0 {
		return
	}

	dirConf := os.Getenv("XDG_CONFIG_HOME")
	dirHome := os.Getenv("HOME")

	if len(dirConf) == 0 && len(dirHome) != 0 {
		dirConf = path.Join(dirHome, ".config")
		stat, err := os.Stat(dirConf)
		if err != nil || stat.IsDir() {
			return
		}
	}

	if len(dirConf) == 0 {
		return
	}

	dirGoateeConf := path.Join(dirConf, "goatee")
	err := os.MkdirAll(dirGoateeConf, 0700)
	if err != nil {
		log.Println("failed create directory for save configuration,reason:", err)
		return
	}
	c.filename = path.Join(dirGoateeConf, "goatee.conf")
}

//OpenWindow open window configuration
func (c *Conf) OpenWindow() {
	if c.window == nil {
		c.CreateWindow()
	}
	c.window.ShowAll()
}

//CreateWindow create window configuration
func (c *Conf) CreateWindow() {
	c.window = gtk.NewWindow(gtk.WINDOW_TOPLEVEL)
	c.window.SetName("Preferences")
	c.window.SetTypeHint(gdk.WINDOW_TYPE_HINT_DIALOG)
	c.window.SetDefaultSize(300, 300)
	c.window.SetSizeRequest(300, 300)

	vbox := gtk.NewVBox(false, 0)

	notebook := gtk.NewNotebook()

	rc := reflect.TypeOf(*conf)
	rv := reflect.ValueOf(conf).Elem()
	for i := 0; i < rc.NumField(); i++ {
		if rc.Field(i).Type.Kind() != reflect.Struct {
			continue
		}

		fvbox := gtk.NewVBox(false, 0)
		notebook.AppendPage(fvbox, gtk.NewLabel(rc.Field(i).Name))

		confStruct := rc.Field(i).Type
		for j := 0; j < confStruct.NumField(); j++ {
			field := confStruct.Field(j)
			val := rv.Field(i).Field(j)

			//prepare name field
			name := strings.Split(field.Tag.Get("toml"), ",")[0]
			if len(name) == 0 {
				name = field.Name
			}
			name = c.FormatName(name)

			hbox := gtk.NewHBox(false, 0)

			label := gtk.NewLabel(name)
			hbox.PackStart(label, false, false, 5)

			switch val.Kind() {
			case reflect.Bool:
				widget := gtk.NewCheckButton()
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetActive(val.Bool())

				w := &ConfWidget{chkbtn: widget, Field: val, conf: c}
				widget.Connect("clicked", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.String:
				widget := gtk.NewEntry()
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetText(val.String())

				w := &ConfWidget{entry: widget, Field: val, conf: c}
				widget.Connect("changed", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.Int:
				widget := gtk.NewSpinButtonWithRange(-1, 2048, 1)
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetValue(float64(val.Int()))

				w := &ConfWidget{spnbtn: widget, Field: val, conf: c}
				widget.Connect("changed", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.TypeOf([]int{}).Kind(): //color
				color := convertColor(val.Interface().([]int))
				widget := gtk.NewColorButtonWithColor(color)
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)

				w := &ConfWidget{colbtn: widget, Field: val, conf: c}
				widget.Connect("color-set", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			}

			fvbox.PackStart(hbox, false, false, 5)
		}
	}

	closebtn := gtk.NewButtonFromStock(gtk.STOCK_CLOSE)
	closebtn.Clicked(c.CloseWindow)
	hbox := gtk.NewHBox(false, 0)
	hbox.PackEnd(closebtn, false, false, 5)

	vbox.Add(notebook)
	vbox.PackEnd(hbox, false, false, 5)

	c.window.Add(vbox)
}

func (c *Conf) CloseWindow() {
	c.Write()

	c.window.Hide()
}

type ConfWidget struct {
	Field reflect.Value

	conf *Conf

	chkbtn *gtk.CheckButton
	entry  *gtk.Entry
	spnbtn *gtk.SpinButton
	colbtn *gtk.ColorButton
}

func (w *ConfWidget) UpdateValue() {
	switch {
	case w.chkbtn != nil:
		w.Field.SetBool(w.chkbtn.GetActive())
	case w.entry != nil:
		w.Field.SetString(w.entry.GetText())
	case w.spnbtn != nil:
		n, _ := strconv.Atoi(w.spnbtn.Entry.GetText())
		w.Field.SetInt(int64(n))
	case w.colbtn != nil:
		col := w.colbtn.GetColor()
		r := int(math.Sqrt(float64(col.Red())))
		g := int(math.Sqrt(float64(col.Green())))
		b := int(math.Sqrt(float64(col.Blue())))

		w.Field.Set(reflect.ValueOf([3]int{r, g, b}))
	}
}

func (c *Conf) FormatName(name string) string {
	name = strings.Replace(name, "-", " ", -1)
	name = regexp.MustCompile("([A-Z])").ReplaceAllString(name, " $1")
	name = strings.TrimSpace(name)
	name = strings.Title(name)
	return name
}
<commit_msg>some fixes to config<commit_after>package main

import (
	"io/ioutil"
	"log"
	"math"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"github.com/mattn/go-gtk/gdk"
	"github.com/mattn/go-gtk/gtk"
	"github.com/naoina/toml"
)

//Conf structure contains configuration
type Conf struct {
	window   *gtk.Window `toml:",omitempty"`
	filename string      `toml:",omitempty"`

	UI struct {
		MenuBarVisible   bool `toml:"menubar-visible"`
		StatusBarVisible bool `toml:"statusbar-visible"`
	}
	TextView struct {
		Font           string `toml:"font"`
		LineHightlight bool   `toml:"line-hightlight"`
		LineNumbers    bool   `toml:"line-numbers"`
		WordWrap       bool   `toml:"word-wrap"`
		IndentSpace    bool   `toml:"indent-space"`
		IndentWidth    int    `toml:"indent-width"`
		StyleScheme    string `toml:"style-scheme"`
	}
	Tabs struct {
		Homogeneous bool  `toml:"homogeneous"`
		CloseBtns   bool  `toml:"close-buttons"`
		Height      int   `toml:"height"`
		FGNormal    []int `toml:"fg-normal"`
		FGModified  []int `toml:"fg-modified"`
		FGNew       []int `toml:"fg-new"`
	}
	Search struct {
		MaxItems int `toml:"max-items"`
	}
	Hex struct {
		BytesInLine int `toml:"bytes-in-line"`
	}
}

//NewConf set default values for configuration and parse config file
func NewConf() *Conf {
	confdir := os.Getenv("XDG_CONFIG_HOME")
	if confdir == "" {
		confdir = path.Join(os.Getenv("HOME"), ".config")
	}
	confdir = path.Join(confdir, "goatee")

	// default values
	c := new(Conf)
	c.filename = path.Join(confdir, "goatee.conf")
	c.UI.MenuBarVisible = true
	c.UI.StatusBarVisible = false

	c.TextView.Font = "Liberation Mono 8"
	c.TextView.LineHightlight = true
	c.TextView.LineNumbers = true
	c.TextView.WordWrap = true
	c.TextView.IndentSpace = false
	c.TextView.IndentWidth = 2

	c.Tabs.Homogeneous = true
	c.Tabs.CloseBtns = true
	c.Tabs.Height = 16
	c.Tabs.FGNormal = []int{200, 200, 200}
	c.Tabs.FGModified = []int{220, 20, 20}
	c.Tabs.FGNew = []int{250, 200, 10}

	c.Search.MaxItems = 1024

	c.Hex.BytesInLine = 16

	//parse config files
	for _, configfile := range []string{
		path.Join(confdir, "goatee.conf"),
		"goatee.conf",
	} {
		data, err := ioutil.ReadFile(configfile)
		if err != nil {
			continue
		}

		err = toml.Unmarshal(data, &c)
		if err != nil {
			log.Printf("failed decode config file '%s', reason: %s", configfile, err)
			continue
		}

		c.filename = configfile
		break
	}

	return c
}

func (c *Conf) Write() {
	os.MkdirAll(filepath.Dir(c.filename), 0755)

	log.Println("write config to:", c.filename)

	if c.filename == "" {
		return
	}

	data, err := toml.Marshal(&c)
	if err != nil {
		log.Println(err)
		return
	}

	err = ioutil.WriteFile(c.filename, data, 0644)
	if err != nil {
		log.Println(err)
		return
	}
}

//OpenWindow open window configuration
func (c *Conf) OpenWindow() {
	if c.window == nil {
		c.CreateWindow()
	}
	c.window.ShowAll()
}

//CreateWindow create window configuration
func (c *Conf) CreateWindow() {
	c.window = gtk.NewWindow(gtk.WINDOW_TOPLEVEL)
	c.window.SetName("Preferences")
	c.window.SetTypeHint(gdk.WINDOW_TYPE_HINT_DIALOG)
	c.window.SetDefaultSize(300, 300)
	c.window.SetSizeRequest(300, 300)

	vbox := gtk.NewVBox(false, 0)

	notebook := gtk.NewNotebook()

	rc := reflect.TypeOf(*conf)
	rv := reflect.ValueOf(conf).Elem()
	for i := 0; i < rc.NumField(); i++ {
		if rc.Field(i).Type.Kind() != reflect.Struct {
			continue
		}

		fvbox := gtk.NewVBox(false, 0)
		notebook.AppendPage(fvbox, gtk.NewLabel(rc.Field(i).Name))

		confStruct := rc.Field(i).Type
		for j := 0; j < confStruct.NumField(); j++ {
			field := confStruct.Field(j)
			val := rv.Field(i).Field(j)

			//prepare name field
			name := strings.Split(field.Tag.Get("toml"), ",")[0]
			if len(name) == 0 {
				name = field.Name
			}
			name = c.FormatName(name)

			hbox := gtk.NewHBox(false, 0)

			label := gtk.NewLabel(name)
			hbox.PackStart(label, false, false, 5)

			switch val.Kind() {
			case reflect.Bool:
				widget := gtk.NewCheckButton()
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetActive(val.Bool())

				w := &ConfWidget{chkbtn: widget, Field: val, conf: c}
				widget.Connect("clicked", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.String:
				widget := gtk.NewEntry()
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetText(val.String())

				w := &ConfWidget{entry: widget, Field: val, conf: c}
				widget.Connect("changed", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.Int:
				widget := gtk.NewSpinButtonWithRange(-1, 2048, 1)
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)
				widget.SetValue(float64(val.Int()))

				w := &ConfWidget{spnbtn: widget, Field: val, conf: c}
				widget.Connect("changed", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			case reflect.TypeOf([]int{}).Kind(): //color
				color := convertColor(val.Interface().([]int))
				widget := gtk.NewColorButtonWithColor(color)
				widget.SetName(field.Name)
				widget.SetSizeRequest(120, 20)

				w := &ConfWidget{colbtn: widget, Field: val, conf: c}
				widget.Connect("color-set", w.UpdateValue)

				hbox.PackEnd(widget, false, false, 5)
			}

			fvbox.PackStart(hbox, false, false, 5)
		}
	}

	closebtn := gtk.NewButtonFromStock(gtk.STOCK_CLOSE)
	closebtn.Clicked(c.CloseWindow)
	hbox := gtk.NewHBox(false, 0)
	hbox.PackEnd(closebtn, false, false, 5)

	vbox.Add(notebook)
	vbox.PackEnd(hbox, false, false, 5)

	c.window.Add(vbox)
}

func (c *Conf) CloseWindow() {
	c.Write()

	c.window.Hide()
}

type ConfWidget struct {
	Field reflect.Value

	conf *Conf

	chkbtn *gtk.CheckButton
	entry  *gtk.Entry
	spnbtn *gtk.SpinButton
	colbtn *gtk.ColorButton
}

func (w *ConfWidget) UpdateValue() {
	switch {
	case w.chkbtn != nil:
		w.Field.SetBool(w.chkbtn.GetActive())
	case w.entry != nil:
		w.Field.SetString(w.entry.GetText())
	case w.spnbtn != nil:
		n, _ := strconv.Atoi(w.spnbtn.Entry.GetText())
		w.Field.SetInt(int64(n))
	case w.colbtn != nil:
		col := w.colbtn.GetColor()
		r := int(math.Sqrt(float64(col.Red())))
		g := int(math.Sqrt(float64(col.Green())))
		b := int(math.Sqrt(float64(col.Blue())))

		w.Field.Set(reflect.ValueOf([]int{r, g, b}))
	}
}

func (c *Conf) FormatName(name string) string {
	name = strings.Replace(name, "-", " ", -1)
	name = regexp.MustCompile("([A-Z])").ReplaceAllString(name, " $1")
	name = strings.TrimSpace(name)
	name = strings.Title(name)
	return name
}
<|endoftext|>


<commit_before>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 (
	log "github.com/Sirupsen/logrus"
	"github.com/Sirupsen/logrus/hooks/syslog"
	"github.com/jessevdk/go-flags"
	p "github.com/kr/pretty"
	"github.com/osrg/gobgp/config"
	ops "github.com/osrg/gobgp/openswitch"
	"github.com/osrg/gobgp/server"
	"io/ioutil"
	"log/syslog"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
)

func main() {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGHUP, syscall.SIGTERM)

	var opts struct {
		ConfigFile      string `short:"f" long:"config-file" description:"specifying a config file"`
		ConfigType      string `short:"t" long:"config-type" description:"specifying config type (toml, yaml, json)" default:"toml"`
		LogLevel        string `short:"l" long:"log-level" description:"specifying log level"`
		LogPlain        bool   `short:"p" long:"log-plain" description:"use plain format for logging (json by default)"`
		UseSyslog       string `short:"s" long:"syslog" description:"use syslogd"`
		Facility        string `long:"syslog-facility" description:"specify syslog facility"`
		DisableStdlog   bool   `long:"disable-stdlog" description:"disable standard logging"`
		CPUs            int    `long:"cpus" description:"specify the number of CPUs to be used"`
		Ops             bool   `long:"openswitch" description:"openswitch mode"`
		GrpcPort        int    `short:"g" long:"grpc-port" description:"grpc port" default:"50051"`
		GracefulRestart bool   `short:"r" long:"graceful-restart" description:"flag restart-state in graceful-restart capability"`
		Dry             bool   `short:"d" long:"dry-run" description:"check configuration"`
	}
	_, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}

	if opts.CPUs == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		if runtime.NumCPU() < opts.CPUs {
			log.Errorf("Only %d CPUs are available but %d is specified", runtime.NumCPU(), opts.CPUs)
			os.Exit(1)
		}
		runtime.GOMAXPROCS(opts.CPUs)
	}

	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()

	switch opts.LogLevel {
	case "debug":
		log.SetLevel(log.DebugLevel)
	case "info":
		log.SetLevel(log.InfoLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	if opts.DisableStdlog == true {
		log.SetOutput(ioutil.Discard)
	} else {
		log.SetOutput(os.Stdout)
	}

	if opts.UseSyslog != "" {
		dst := strings.SplitN(opts.UseSyslog, ":", 2)
		network := ""
		addr := ""
		if len(dst) == 2 {
			network = dst[0]
			addr = dst[1]
		}

		facility := syslog.Priority(0)
		switch opts.Facility {
		case "kern":
			facility = syslog.LOG_KERN
		case "user":
			facility = syslog.LOG_USER
		case "mail":
			facility = syslog.LOG_MAIL
		case "daemon":
			facility = syslog.LOG_DAEMON
		case "auth":
			facility = syslog.LOG_AUTH
		case "syslog":
			facility = syslog.LOG_SYSLOG
		case "lpr":
			facility = syslog.LOG_LPR
		case "news":
			facility = syslog.LOG_NEWS
		case "uucp":
			facility = syslog.LOG_UUCP
		case "cron":
			facility = syslog.LOG_CRON
		case "authpriv":
			facility = syslog.LOG_AUTHPRIV
		case "ftp":
			facility = syslog.LOG_FTP
		case "local0":
			facility = syslog.LOG_LOCAL0
		case "local1":
			facility = syslog.LOG_LOCAL1
		case "local2":
			facility = syslog.LOG_LOCAL2
		case "local3":
			facility = syslog.LOG_LOCAL3
		case "local4":
			facility = syslog.LOG_LOCAL4
		case "local5":
			facility = syslog.LOG_LOCAL5
		case "local6":
			facility = syslog.LOG_LOCAL6
		case "local7":
			facility = syslog.LOG_LOCAL7
		}

		hook, err := logrus_syslog.NewSyslogHook(network, addr, syslog.LOG_INFO|facility, "bgpd")
		if err != nil {
			log.Error("Unable to connect to syslog daemon, ", opts.UseSyslog)
			os.Exit(1)
		} else {
			log.AddHook(hook)
		}
	}

	if opts.LogPlain == false {
		log.SetFormatter(&log.JSONFormatter{})
	}

	configCh := make(chan config.BgpConfigSet)
	reloadCh := make(chan bool)
	if opts.Dry {
		go config.ReadConfigfileServe(opts.ConfigFile, opts.ConfigType, configCh, reloadCh)
		reloadCh <- true
		c := <-configCh
		if opts.LogLevel == "debug" {
			p.Println(c)
		}
		os.Exit(0)
	}

	log.Info("gobgpd started")

	bgpServer := server.NewBgpServer()
	if opts.Ops {
		m, err := ops.NewOpsManager(bgpServer.GrpcReqCh)
		if err != nil {
			log.Errorf("Failed to start ops config manager: %s", err)
			os.Exit(1)
		}
		log.Info("Coordination with OpenSwitch")
		m.Serve()
	} else if opts.ConfigFile != "" {
		go config.ReadConfigfileServe(opts.ConfigFile, opts.ConfigType, configCh, reloadCh)
		reloadCh <- true
	}
	go bgpServer.Serve()

	// start grpc Server
	grpcServer := server.NewGrpcServer(opts.GrpcPort, bgpServer.GrpcReqCh)
	go func() {
		if err := grpcServer.Serve(); err != nil {
			log.Fatalf("failed to listen grpc port: %s", err)
		}
	}()

	var bgpConfig *config.Bgp = nil
	var policyConfig *config.RoutingPolicy = nil
	for {
		select {
		case newConfig := <-configCh:
			var added, deleted, updated []config.Neighbor

			if bgpConfig == nil {
				bgpServer.SetGlobalType(newConfig.Bgp.Global)
				bgpConfig = &newConfig.Bgp
				bgpServer.SetRpkiConfig(newConfig.Bgp.RpkiServers)
				added = newConfig.Bgp.Neighbors
				if opts.GracefulRestart {
					for i, n := range added {
						if n.GracefulRestart.Config.Enabled {
							added[i].GracefulRestart.State.LocalRestarting = true
						}
					}
				}
				deleted = []config.Neighbor{}
				updated = []config.Neighbor{}
			} else {
				bgpConfig, added, deleted, updated = config.UpdateConfig(bgpConfig, &newConfig.Bgp)
			}

			if policyConfig == nil {
				policyConfig = &newConfig.Policy
				// FIXME: Currently the following code
				// is safe because the above
				// SetRpkiConfig will be blocked
				// because the length of rpkiConfigCh
				// is zero. So server.GlobalRib is
				// allocated before the above
				// SetPolicy. But this should be
				// handled more cleanly.
				if err := bgpServer.SetRoutingPolicy(newConfig.Policy); err != nil {
					log.Fatal(err)
				}
			} else {
				if config.CheckPolicyDifference(policyConfig, &newConfig.Policy) {
					log.Info("Policy config is updated")
					bgpServer.UpdatePolicy(newConfig.Policy)
				}
			}

			for _, p := range added {
				log.Infof("Peer %v is added", p.Config.NeighborAddress)
				bgpServer.PeerAdd(p)
			}
			for _, p := range deleted {
				log.Infof("Peer %v is deleted", p.Config.NeighborAddress)
				bgpServer.PeerDelete(p)
			}
			for _, p := range updated {
				log.Infof("Peer %v is updated", p.Config.NeighborAddress)
				bgpServer.PeerUpdate(p)
			}
		case sig := <-sigCh:
			switch sig {
			case syscall.SIGHUP:
				log.Info("reload the config file")
				reloadCh <- true
			case syscall.SIGKILL, syscall.SIGTERM:
				bgpServer.Shutdown()
			}
		}
	}
}
<commit_msg>gobgpd: stop coloring when stdout log is disabled<commit_after>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// 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 (
	log "github.com/Sirupsen/logrus"
	"github.com/Sirupsen/logrus/hooks/syslog"
	"github.com/jessevdk/go-flags"
	p "github.com/kr/pretty"
	"github.com/osrg/gobgp/config"
	ops "github.com/osrg/gobgp/openswitch"
	"github.com/osrg/gobgp/server"
	"io/ioutil"
	"log/syslog"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
)

func main() {
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGHUP, syscall.SIGTERM)

	var opts struct {
		ConfigFile      string `short:"f" long:"config-file" description:"specifying a config file"`
		ConfigType      string `short:"t" long:"config-type" description:"specifying config type (toml, yaml, json)" default:"toml"`
		LogLevel        string `short:"l" long:"log-level" description:"specifying log level"`
		LogPlain        bool   `short:"p" long:"log-plain" description:"use plain format for logging (json by default)"`
		UseSyslog       string `short:"s" long:"syslog" description:"use syslogd"`
		Facility        string `long:"syslog-facility" description:"specify syslog facility"`
		DisableStdlog   bool   `long:"disable-stdlog" description:"disable standard logging"`
		CPUs            int    `long:"cpus" description:"specify the number of CPUs to be used"`
		Ops             bool   `long:"openswitch" description:"openswitch mode"`
		GrpcPort        int    `short:"g" long:"grpc-port" description:"grpc port" default:"50051"`
		GracefulRestart bool   `short:"r" long:"graceful-restart" description:"flag restart-state in graceful-restart capability"`
		Dry             bool   `short:"d" long:"dry-run" description:"check configuration"`
	}
	_, err := flags.Parse(&opts)
	if err != nil {
		os.Exit(1)
	}

	if opts.CPUs == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		if runtime.NumCPU() < opts.CPUs {
			log.Errorf("Only %d CPUs are available but %d is specified", runtime.NumCPU(), opts.CPUs)
			os.Exit(1)
		}
		runtime.GOMAXPROCS(opts.CPUs)
	}

	go func() {
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()

	switch opts.LogLevel {
	case "debug":
		log.SetLevel(log.DebugLevel)
	case "info":
		log.SetLevel(log.InfoLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	if opts.DisableStdlog == true {
		log.SetOutput(ioutil.Discard)
	} else {
		log.SetOutput(os.Stdout)
	}

	if opts.UseSyslog != "" {
		dst := strings.SplitN(opts.UseSyslog, ":", 2)
		network := ""
		addr := ""
		if len(dst) == 2 {
			network = dst[0]
			addr = dst[1]
		}

		facility := syslog.Priority(0)
		switch opts.Facility {
		case "kern":
			facility = syslog.LOG_KERN
		case "user":
			facility = syslog.LOG_USER
		case "mail":
			facility = syslog.LOG_MAIL
		case "daemon":
			facility = syslog.LOG_DAEMON
		case "auth":
			facility = syslog.LOG_AUTH
		case "syslog":
			facility = syslog.LOG_SYSLOG
		case "lpr":
			facility = syslog.LOG_LPR
		case "news":
			facility = syslog.LOG_NEWS
		case "uucp":
			facility = syslog.LOG_UUCP
		case "cron":
			facility = syslog.LOG_CRON
		case "authpriv":
			facility = syslog.LOG_AUTHPRIV
		case "ftp":
			facility = syslog.LOG_FTP
		case "local0":
			facility = syslog.LOG_LOCAL0
		case "local1":
			facility = syslog.LOG_LOCAL1
		case "local2":
			facility = syslog.LOG_LOCAL2
		case "local3":
			facility = syslog.LOG_LOCAL3
		case "local4":
			facility = syslog.LOG_LOCAL4
		case "local5":
			facility = syslog.LOG_LOCAL5
		case "local6":
			facility = syslog.LOG_LOCAL6
		case "local7":
			facility = syslog.LOG_LOCAL7
		}

		hook, err := logrus_syslog.NewSyslogHook(network, addr, syslog.LOG_INFO|facility, "bgpd")
		if err != nil {
			log.Error("Unable to connect to syslog daemon, ", opts.UseSyslog)
			os.Exit(1)
		} else {
			log.AddHook(hook)
		}
	}

	if opts.LogPlain {
		if opts.DisableStdlog {
			log.SetFormatter(&log.TextFormatter{
				DisableColors: true,
			})
		}
	} else {
		log.SetFormatter(&log.JSONFormatter{})
	}

	configCh := make(chan config.BgpConfigSet)
	reloadCh := make(chan bool)
	if opts.Dry {
		go config.ReadConfigfileServe(opts.ConfigFile, opts.ConfigType, configCh, reloadCh)
		reloadCh <- true
		c := <-configCh
		if opts.LogLevel == "debug" {
			p.Println(c)
		}
		os.Exit(0)
	}

	log.Info("gobgpd started")

	bgpServer := server.NewBgpServer()
	if opts.Ops {
		m, err := ops.NewOpsManager(bgpServer.GrpcReqCh)
		if err != nil {
			log.Errorf("Failed to start ops config manager: %s", err)
			os.Exit(1)
		}
		log.Info("Coordination with OpenSwitch")
		m.Serve()
	} else if opts.ConfigFile != "" {
		go config.ReadConfigfileServe(opts.ConfigFile, opts.ConfigType, configCh, reloadCh)
		reloadCh <- true
	}
	go bgpServer.Serve()

	// start grpc Server
	grpcServer := server.NewGrpcServer(opts.GrpcPort, bgpServer.GrpcReqCh)
	go func() {
		if err := grpcServer.Serve(); err != nil {
			log.Fatalf("failed to listen grpc port: %s", err)
		}
	}()

	var bgpConfig *config.Bgp = nil
	var policyConfig *config.RoutingPolicy = nil
	for {
		select {
		case newConfig := <-configCh:
			var added, deleted, updated []config.Neighbor

			if bgpConfig == nil {
				bgpServer.SetGlobalType(newConfig.Bgp.Global)
				bgpConfig = &newConfig.Bgp
				bgpServer.SetRpkiConfig(newConfig.Bgp.RpkiServers)
				added = newConfig.Bgp.Neighbors
				if opts.GracefulRestart {
					for i, n := range added {
						if n.GracefulRestart.Config.Enabled {
							added[i].GracefulRestart.State.LocalRestarting = true
						}
					}
				}
				deleted = []config.Neighbor{}
				updated = []config.Neighbor{}
			} else {
				bgpConfig, added, deleted, updated = config.UpdateConfig(bgpConfig, &newConfig.Bgp)
			}

			if policyConfig == nil {
				policyConfig = &newConfig.Policy
				// FIXME: Currently the following code
				// is safe because the above
				// SetRpkiConfig will be blocked
				// because the length of rpkiConfigCh
				// is zero. So server.GlobalRib is
				// allocated before the above
				// SetPolicy. But this should be
				// handled more cleanly.
				if err := bgpServer.SetRoutingPolicy(newConfig.Policy); err != nil {
					log.Fatal(err)
				}
			} else {
				if config.CheckPolicyDifference(policyConfig, &newConfig.Policy) {
					log.Info("Policy config is updated")
					bgpServer.UpdatePolicy(newConfig.Policy)
				}
			}

			for _, p := range added {
				log.Infof("Peer %v is added", p.Config.NeighborAddress)
				bgpServer.PeerAdd(p)
			}
			for _, p := range deleted {
				log.Infof("Peer %v is deleted", p.Config.NeighborAddress)
				bgpServer.PeerDelete(p)
			}
			for _, p := range updated {
				log.Infof("Peer %v is updated", p.Config.NeighborAddress)
				bgpServer.PeerUpdate(p)
			}
		case sig := <-sigCh:
			switch sig {
			case syscall.SIGHUP:
				log.Info("reload the config file")
				reloadCh <- true
			case syscall.SIGKILL, syscall.SIGTERM:
				bgpServer.Shutdown()
			}
		}
	}
}
<|endoftext|>


<commit_before>package gojpegoptim

// #cgo CFLAGS: -I/opt/local/include -I/usr/local/include
// #cgo LDFLAGS: -L/opt/local/lib -L/usr/local/lib -ljpeg
// #include <stdlib.h>
// extern int optimizeJPEG(unsigned char *inputbuffer, unsigned long inputsize, unsigned char **outputbuffer, unsigned long *outputsize, int quality);
// extern int encodeJPEG(unsigned char *inputbuffer, int width, int height, unsigned char **outputbuffer, unsigned long *outputsize, int quality);
import "C"

import (
	"errors"
	"image"
	"image/draw"
	"io"
	"unsafe"
)

// Options are the encoding parameters.
// Quality ranges from -1 to 100 inclusive, higher is better.
type Options struct {
	Quality int
}

// Encode writes the Image m to w in JPEG format with the given options.
func Encode(w io.Writer, m image.Image, o *Options) error {
	// Get an image.RGBA if it is one
	rgba, ok := m.(*image.RGBA)
	if !ok {
		b := m.Bounds()
		tm := image.NewRGBA(b)
		draw.Draw(tm, b, m, b.Min, draw.Src)
		m = tm
	}
	var coutimg *C.uchar
	var coutsize C.ulong
	input := (*C.uchar)(unsafe.Pointer(&rgba.Pix[0]))
	code := C.encodeJPEG(input, C.int(m.Bounds().Size().X), C.int(m.Bounds().Size().Y), &coutimg, &coutsize, C.int(o.Quality))
	if code != 0 || coutsize == 0 {
		return errors.New("Encoding error")
	}
	outBytes := C.GoBytes(unsafe.Pointer(coutimg), C.int(coutsize))
	w.Write(outBytes)
	C.free(unsafe.Pointer(coutimg))
	return nil
}

// Optimize a JPEG bytes array if quality is -1, Optimize & Recompress if quality is between [0 - 100]
func EncodeBytesOptimized(srcBytes []byte, o *Options) (outBytes []byte, err error) {
	if len(srcBytes) == 0 {
		err = errors.New("Image source is empty")
		return
	}
	// Clip quality to [-1, 100].
	if o != nil {
		quality := o.Quality
		if quality < -1 {
			quality = -1
		} else if quality > 100 {
			quality = 100
		}
	}

	csrcimg := (*C.uchar)(unsafe.Pointer(&srcBytes[0]))
	cinputsize := C.ulong(len(srcBytes))
	var coutimg *C.uchar
	var coutsize C.ulong
	code := C.optimizeJPEG(csrcimg, cinputsize, &coutimg, &coutsize, C.int(o.Quality))
	if code != 0 || coutsize == 0 {
		err = errors.New("Optimize failed")
		return
	}
	outBytes = C.GoBytes(unsafe.Pointer(coutimg), C.int(coutsize))
	C.free(unsafe.Pointer(coutimg))
	return
}
<commit_msg>better os detection avoid warning ldflags<commit_after>package gojpegoptim

// #cgo LDFLAGS: -ljpeg
// #cgo darwin LDFLAGS: -L/opt/local/lib
// #cgo darwin CFLAGS: -I/opt/local/include
// #cgo freebsd LDFLAGS: -L/usr/local/lib
// #cgo freebsd CFLAGS: -I/usr/local/include
// #include <stdlib.h>
// extern int optimizeJPEG(unsigned char *inputbuffer, unsigned long inputsize, unsigned char **outputbuffer, unsigned long *outputsize, int quality);
// extern int encodeJPEG(unsigned char *inputbuffer, int width, int height, unsigned char **outputbuffer, unsigned long *outputsize, int quality);
import "C"

import (
	"errors"
	"image"
	"image/draw"
	"io"
	"unsafe"
)

// Options are the encoding parameters.
// Quality ranges from -1 to 100 inclusive, higher is better.
type Options struct {
	Quality int
}

// Encode writes the Image m to w in JPEG format with the given options.
func Encode(w io.Writer, m image.Image, o *Options) error {
	// Get an image.RGBA if it is one
	rgba, ok := m.(*image.RGBA)
	if !ok {
		b := m.Bounds()
		tm := image.NewRGBA(b)
		draw.Draw(tm, b, m, b.Min, draw.Src)
		m = tm
	}
	var coutimg *C.uchar
	var coutsize C.ulong
	input := (*C.uchar)(unsafe.Pointer(&rgba.Pix[0]))
	code := C.encodeJPEG(input, C.int(m.Bounds().Size().X), C.int(m.Bounds().Size().Y), &coutimg, &coutsize, C.int(o.Quality))
	if code != 0 || coutsize == 0 {
		return errors.New("Encoding error")
	}
	outBytes := C.GoBytes(unsafe.Pointer(coutimg), C.int(coutsize))
	w.Write(outBytes)
	C.free(unsafe.Pointer(coutimg))
	return nil
}

// Optimize a JPEG bytes array if quality is -1, Optimize & Recompress if quality is between [0 - 100]
func EncodeBytesOptimized(srcBytes []byte, o *Options) (outBytes []byte, err error) {
	if len(srcBytes) == 0 {
		err = errors.New("Image source is empty")
		return
	}
	// Clip quality to [-1, 100].
	if o != nil {
		quality := o.Quality
		if quality < -1 {
			quality = -1
		} else if quality > 100 {
			quality = 100
		}
	}

	csrcimg := (*C.uchar)(unsafe.Pointer(&srcBytes[0]))
	cinputsize := C.ulong(len(srcBytes))
	var coutimg *C.uchar
	var coutsize C.ulong
	code := C.optimizeJPEG(csrcimg, cinputsize, &coutimg, &coutsize, C.int(o.Quality))
	if code != 0 || coutsize == 0 {
		err = errors.New("Optimize failed")
		return
	}
	outBytes = C.GoBytes(unsafe.Pointer(coutimg), C.int(coutsize))
	C.free(unsafe.Pointer(coutimg))
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Google Inc. 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 main

import (
	"io/ioutil"
	"path"
	"reflect"
	"testing"
	"time"

	"github.com/google/googet/client"
	"github.com/google/googet/goolib"
	"github.com/google/googet/oswrap"
)

func TestRepoList(t *testing.T) {
	testRepo := "https://foo.com/googet/bar"

	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	f, err := oswrap.Create(path.Join(tempDir, "test.repo"))
	if err != nil {
		t.Fatalf("error creating repo file: %v", err)
	}
	if _, err := f.Write([]byte("url: " + testRepo)); err != nil {
		t.Fatalf("error writing repo: %v", err)
	}
	if err := f.Close(); err != nil {
		t.Fatalf("error closing repo file: %v", err)
	}

	got, err := repoList(tempDir)
	if err != nil {
		t.Fatal(err)
	}

	if !reflect.DeepEqual(got, []string{testRepo}) {
		t.Errorf("returned repo does not match expected repo: got %v, want %v", got, testRepo)
	}
}

func TestInstalledPackages(t *testing.T) {
	state := []client.PackageState{
		{
			PackageSpec: &goolib.PkgSpec{
				Name:    "foo",
				Version: "1.2.3@4",
				Arch:    "noarch",
			},
		},
		{
			PackageSpec: &goolib.PkgSpec{
				Name:    "bar",
				Version: "0.1.0@1",
				Arch:    "noarch",
			},
		},
	}

	want := packageMap{"foo.noarch": "1.2.3@4", "bar.noarch": "0.1.0@1"}
	got := installedPackages(state)

	if !reflect.DeepEqual(got, want) {
		t.Errorf("returned map does not match expected map: got %v, want %v", got, want)
	}
}

func TestReadConf(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	confPath := path.Join(tempDir, "test.conf")
	f, err := oswrap.Create(confPath)
	if err != nil {
		t.Fatalf("error creating conf file: %v", err)
	}

	content := []byte("archs: [noarch, x86_64]\ncachelife: 10m")
	if _, err := f.Write(content); err != nil {
		t.Fatalf("error writing conf file: %v", err)
	}
	if err := f.Close(); err != nil {
		t.Fatalf("error closing conf file: %v", err)
	}

	readConf(confPath)

	ea := []string{"noarch", "x86_64"}
	if !reflect.DeepEqual(archs, ea) {
		t.Errorf("readConf did not create expected arch list, want: %s, got: %s", ea, archs)
	}

	ecl := time.Duration(10 * time.Minute)
	if cacheLife != ecl {
		t.Errorf("readConf did not create expected cacheLife, want: %s, got: %s", ecl, cacheLife)
	}
}

func TestRotateLog(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	table := []struct {
		name    string
		size    int64
		rotated bool
	}{
		{"test1.log", 10 * 1024, true},
		{"test2.log", 9 * 1024, false},
	}

	for _, tt := range table {
		logPath := path.Join(tempDir, tt.name)
		f, err := oswrap.Create(logPath)
		if err != nil {
			t.Fatalf("error creating log file: %v", err)
		}

		if err := f.Truncate(tt.size); err != nil {
			t.Fatalf("error truncating log file: %v", err)
		}

		if err := f.Close(); err != nil {
			t.Fatalf("error closing log file: %v", err)
		}

		if err := rotateLog(logPath, 10*1024); err != nil {
			t.Errorf("error running rotateLog: %v", err)
		}

		switch tt.rotated {
		case true:
			if _, err := oswrap.Stat(logPath); err == nil {
				t.Error("rotateLog did not rotate log as expected, old log file still exists")
			}
			if _, err := oswrap.Stat(logPath + ".old"); err != nil {
				t.Error("rotateLog did not rotate log as expected, .old file does not exist")
			}
		case false:
			if _, err := oswrap.Stat(logPath); err != nil {
				t.Error("rotateLog rotated a log we didn't expect")
			}
		}
	}
}

func TestWriteReadState(t *testing.T) {
	want := &client.GooGetState{
		client.PackageState{PackageSpec: &goolib.PkgSpec{Name: "test"}},
	}

	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	sf := path.Join(tempDir, "test.state")

	if err := writeState(want, sf); err != nil {
		t.Errorf("error running writeState: %v", err)
	}

	got, err := readState(sf)
	if err != nil {
		t.Errorf("error running readState: %v", err)
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("did not get expected state, got: %+v, want %+v", got, want)
	}
}
<commit_msg>Add tests<commit_after>/*
Copyright 2016 Google Inc. 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 main

import (
	"io/ioutil"
	"path"
	"reflect"
	"testing"
	"time"

	"github.com/google/googet/client"
	"github.com/google/googet/goolib"
	"github.com/google/googet/oswrap"
)

func TestRepoList(t *testing.T) {
	testRepo := "https://foo.com/googet/bar"

	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	testFile := path.Join(tempDir, "test.repo")

	repoTests := []struct {
		content []byte
		result  []string
	}{
		{[]byte("url: " + testRepo), []string{testRepo}},
		{[]byte("- url: " + testRepo), []string{testRepo}},
		{[]byte("- url: " + testRepo + "\n\n- url: " + testRepo), []string{testRepo, testRepo}},
	}

	for _, tt := range repoTests {
		if err := ioutil.WriteFile(testFile, tt.content, 0660); err != nil {
			t.Fatalf("error writing repo: %v", err)
		}
		got, err := repoList(tempDir)
		if err != nil {
			t.Fatal(err)
		}
		if !reflect.DeepEqual(got, tt.result) {
			t.Errorf("returned repo does not match expected repo: got %v, want %v", got, testRepo)
		}
	}
}

func TestInstalledPackages(t *testing.T) {
	state := []client.PackageState{
		{
			PackageSpec: &goolib.PkgSpec{
				Name:    "foo",
				Version: "1.2.3@4",
				Arch:    "noarch",
			},
		},
		{
			PackageSpec: &goolib.PkgSpec{
				Name:    "bar",
				Version: "0.1.0@1",
				Arch:    "noarch",
			},
		},
	}

	want := packageMap{"foo.noarch": "1.2.3@4", "bar.noarch": "0.1.0@1"}
	got := installedPackages(state)

	if !reflect.DeepEqual(got, want) {
		t.Errorf("returned map does not match expected map: got %v, want %v", got, want)
	}
}

func TestReadConf(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	confPath := path.Join(tempDir, "test.conf")
	f, err := oswrap.Create(confPath)
	if err != nil {
		t.Fatalf("error creating conf file: %v", err)
	}

	content := []byte("archs: [noarch, x86_64]\ncachelife: 10m")
	if _, err := f.Write(content); err != nil {
		t.Fatalf("error writing conf file: %v", err)
	}
	if err := f.Close(); err != nil {
		t.Fatalf("error closing conf file: %v", err)
	}

	readConf(confPath)

	ea := []string{"noarch", "x86_64"}
	if !reflect.DeepEqual(archs, ea) {
		t.Errorf("readConf did not create expected arch list, want: %s, got: %s", ea, archs)
	}

	ecl := time.Duration(10 * time.Minute)
	if cacheLife != ecl {
		t.Errorf("readConf did not create expected cacheLife, want: %s, got: %s", ecl, cacheLife)
	}
}

func TestRotateLog(t *testing.T) {
	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	table := []struct {
		name    string
		size    int64
		rotated bool
	}{
		{"test1.log", 10 * 1024, true},
		{"test2.log", 9 * 1024, false},
	}

	for _, tt := range table {
		logPath := path.Join(tempDir, tt.name)
		f, err := oswrap.Create(logPath)
		if err != nil {
			t.Fatalf("error creating log file: %v", err)
		}

		if err := f.Truncate(tt.size); err != nil {
			t.Fatalf("error truncating log file: %v", err)
		}

		if err := f.Close(); err != nil {
			t.Fatalf("error closing log file: %v", err)
		}

		if err := rotateLog(logPath, 10*1024); err != nil {
			t.Errorf("error running rotateLog: %v", err)
		}

		switch tt.rotated {
		case true:
			if _, err := oswrap.Stat(logPath); err == nil {
				t.Error("rotateLog did not rotate log as expected, old log file still exists")
			}
			if _, err := oswrap.Stat(logPath + ".old"); err != nil {
				t.Error("rotateLog did not rotate log as expected, .old file does not exist")
			}
		case false:
			if _, err := oswrap.Stat(logPath); err != nil {
				t.Error("rotateLog rotated a log we didn't expect")
			}
		}
	}
}

func TestWriteReadState(t *testing.T) {
	want := &client.GooGetState{
		client.PackageState{PackageSpec: &goolib.PkgSpec{Name: "test"}},
	}

	tempDir, err := ioutil.TempDir("", "")
	if err != nil {
		t.Fatalf("error creating temp directory: %v", err)
	}
	defer oswrap.RemoveAll(tempDir)

	sf := path.Join(tempDir, "test.state")

	if err := writeState(want, sf); err != nil {
		t.Errorf("error running writeState: %v", err)
	}

	got, err := readState(sf)
	if err != nil {
		t.Errorf("error running readState: %v", err)
	}

	if !reflect.DeepEqual(got, want) {
		t.Errorf("did not get expected state, got: %+v, want %+v", got, want)
	}
}
<|endoftext|>


<commit_before>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"
	"sync"

	"github.com/Sirupsen/logrus"
)

type NetworkGraph struct {
	// nodeName -> Node
	NodesMap  map[string]*NetworkNode `json:"nodes"`
	NodesLock *sync.RWMutex           `json:"-"`

	// nodeName,nodeName -> NetworkLink
	LinksMap  map[string]*NetworkLink `json:"edges"`
	LinksLock *sync.RWMutex           `json:"-"`

	RoutesMap  map[string]*NetworkRoute `json:"routes"`
	RoutesLock *sync.RWMutex            `json:"-"`

	// event stuff
	eventChannels     map[chan *Event]bool
	eventRegistration chan chan *Event
	internalEvents    chan *Event
}

func Create() *NetworkGraph {
	g := &NetworkGraph{
		NodesMap:   make(map[string]*NetworkNode),
		NodesLock:  &sync.RWMutex{},
		LinksMap:   make(map[string]*NetworkLink),
		LinksLock:  &sync.RWMutex{},
		RoutesMap:  make(map[string]*NetworkRoute),
		RoutesLock: &sync.RWMutex{},

		eventChannels:     make(map[chan *Event]bool),
		eventRegistration: make(chan chan *Event),
		internalEvents:    make(chan *Event),
	}

	go g.publisher()

	return g
}

// TODO: buffer messages?
// goroutine target to do all the publishing of events
func (g *NetworkGraph) publisher() {
	for {
		select {
		case newChannel := <-g.eventRegistration:
			g.eventChannels[newChannel] = true
		case newEvent := <-g.internalEvents:
			//logrus.Infof("got a new event! %v", newEvent)
			for subscriberChannel := range g.eventChannels {
				select {
				case subscriberChannel <- newEvent:
					//logrus.Infof("sent that event to a channel")
				default:
					//logrus.Infof("Unable to send event to that subscriber, killing")
					delete(g.eventChannels, subscriberChannel)
					close(subscriberChannel)
				}
			}
		}
	}
}

// TODO: locking here
// Dump everything in the NetworkGraph into a channel
func (g *NetworkGraph) EventDumpChannel() chan *Event {
	// TODO: buffer?
	c := make(chan *Event)
	go func(c chan *Event) {
		// dump nodes
		for _, n := range g.NodesMap {
			c <- &Event{
				E:    addEvent,
				Item: n,
			}
		}
		// dump links
		for _, l := range g.LinksMap {
			c <- &Event{
				E:    addEvent,
				Item: l,
			}
		}
		// dump routes
		for _, route := range g.RoutesMap {
			c <- &Event{
				E:    addEvent,
				Item: route,
			}
		}

		close(c)
	}(c)
	return c
}

// add subscriber to our events
func (g *NetworkGraph) Subscribe(c chan *Event) {
	g.eventRegistration <- c
}

func (g *NetworkGraph) IncrNode(name string, newNode *NetworkNode) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		if newNode == nil {
			n = NewNetworkNode(name, g.internalEvents)
		} else {
			n = newNode
			n.updateChan = g.internalEvents
		}
		g.NodesMap[name] = n

		// Now that there is a new thing we fire an addEvent.
		// Note: if the background DNS lookup was initiated an updateEvent
		// will fire as soon as the lookup completes
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: n,
		}
	}
	n.refCount++
	return n, !ok
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return nil, false
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: n,
		}
		return n, true
	}
	return nil, false
}

func (g *NetworkGraph) IncrLink(src, dst string, newLink *NetworkLink) (*NetworkLink, bool) {
	key := src + ";" + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		if newLink == nil {
			srcNode, _ := g.IncrNode(src, nil)
			dstNode, _ := g.IncrNode(dst, nil)
			l = &NetworkLink{
				SrcName: srcNode.Name,
				srcNode: srcNode,
				DstName: dstNode.Name,
				dstNode: dstNode,
			}
		} else {
			srcNode, _ := g.IncrNode(src, newLink.srcNode)
			dstNode, _ := g.IncrNode(dst, newLink.dstNode)
			// update child pointers
			newLink.srcNode = srcNode
			newLink.dstNode = dstNode
			l = newLink
		}
		g.LinksMap[key] = l
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: l,
		}
	}
	l.refCount++
	return l, !ok
}

func (g *NetworkGraph) GetLink(key string) *NetworkLink {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) (*NetworkLink, bool) {
	key := src + ";" + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return nil, false
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: l,
		}
		return l, true
	}
	return nil, false
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string, newRoute *NetworkRoute) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Debugf("New Route: key=%s %v", key, hops)
		if newRoute == nil {
			path := make([]*NetworkNode, 0, len(hops))
			for i, hop := range hops {
				hopNode, _ := g.IncrNode(hop, nil)
				path = append(path, hopNode)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(hops[i-1], hop, nil)
				}
			}
			route = &NetworkRoute{
				Path:       hops,
				path:       path,
				State:      Up,
				metricRing: ring.New(100), // TODO: config
				mLock:      &sync.RWMutex{},
				updateChan: g.internalEvents,
			}
		} else {
			newRoute.path = make([]*NetworkNode, len(newRoute.Path))
			for i, nodeName := range newRoute.Path {
				// Increment the node (this will convert the name to a pointer)
				node, _ := g.IncrNode(nodeName, nil)
				// set the pointer in our `path`
				newRoute.path[i] = node

				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(newRoute.path[i-1].Name, nodeName, nil)
				}
			}
			route = newRoute
			route.updateChan = g.internalEvents
		}

		g.RoutesMap[key] = route

		g.internalEvents <- &Event{
			E:    addEvent,
			Item: route,
		}
	}

	// increment route's refcount
	route.refCount++

	return route, !ok
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return nil, false
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, nodeName := range r.Path {
			g.DecrNode(nodeName)
			if i-1 >= 0 {
				g.DecrLink(r.Path[i-1], nodeName)
			}
		}

		delete(g.RoutesMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: r,
		}
		return r, true
	}
	return r, false
}
<commit_msg>Node/Link Decr should return the respective object if they had it (regardless of whether we delete it)<commit_after>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"
	"sync"

	"github.com/Sirupsen/logrus"
)

type NetworkGraph struct {
	// nodeName -> Node
	NodesMap  map[string]*NetworkNode `json:"nodes"`
	NodesLock *sync.RWMutex           `json:"-"`

	// nodeName,nodeName -> NetworkLink
	LinksMap  map[string]*NetworkLink `json:"edges"`
	LinksLock *sync.RWMutex           `json:"-"`

	RoutesMap  map[string]*NetworkRoute `json:"routes"`
	RoutesLock *sync.RWMutex            `json:"-"`

	// event stuff
	eventChannels     map[chan *Event]bool
	eventRegistration chan chan *Event
	internalEvents    chan *Event
}

func Create() *NetworkGraph {
	g := &NetworkGraph{
		NodesMap:   make(map[string]*NetworkNode),
		NodesLock:  &sync.RWMutex{},
		LinksMap:   make(map[string]*NetworkLink),
		LinksLock:  &sync.RWMutex{},
		RoutesMap:  make(map[string]*NetworkRoute),
		RoutesLock: &sync.RWMutex{},

		eventChannels:     make(map[chan *Event]bool),
		eventRegistration: make(chan chan *Event),
		internalEvents:    make(chan *Event),
	}

	go g.publisher()

	return g
}

// TODO: buffer messages?
// goroutine target to do all the publishing of events
func (g *NetworkGraph) publisher() {
	for {
		select {
		case newChannel := <-g.eventRegistration:
			g.eventChannels[newChannel] = true
		case newEvent := <-g.internalEvents:
			//logrus.Infof("got a new event! %v", newEvent)
			for subscriberChannel := range g.eventChannels {
				select {
				case subscriberChannel <- newEvent:
					//logrus.Infof("sent that event to a channel")
				default:
					//logrus.Infof("Unable to send event to that subscriber, killing")
					delete(g.eventChannels, subscriberChannel)
					close(subscriberChannel)
				}
			}
		}
	}
}

// TODO: locking here
// Dump everything in the NetworkGraph into a channel
func (g *NetworkGraph) EventDumpChannel() chan *Event {
	// TODO: buffer?
	c := make(chan *Event)
	go func(c chan *Event) {
		// dump nodes
		for _, n := range g.NodesMap {
			c <- &Event{
				E:    addEvent,
				Item: n,
			}
		}
		// dump links
		for _, l := range g.LinksMap {
			c <- &Event{
				E:    addEvent,
				Item: l,
			}
		}
		// dump routes
		for _, route := range g.RoutesMap {
			c <- &Event{
				E:    addEvent,
				Item: route,
			}
		}

		close(c)
	}(c)
	return c
}

// add subscriber to our events
func (g *NetworkGraph) Subscribe(c chan *Event) {
	g.eventRegistration <- c
}

func (g *NetworkGraph) IncrNode(name string, newNode *NetworkNode) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		if newNode == nil {
			n = NewNetworkNode(name, g.internalEvents)
		} else {
			n = newNode
			n.updateChan = g.internalEvents
		}
		g.NodesMap[name] = n

		// Now that there is a new thing we fire an addEvent.
		// Note: if the background DNS lookup was initiated an updateEvent
		// will fire as soon as the lookup completes
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: n,
		}
	}
	n.refCount++
	return n, !ok
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return nil, false
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: n,
		}
		return n, true
	}
	return n, false
}

func (g *NetworkGraph) IncrLink(src, dst string, newLink *NetworkLink) (*NetworkLink, bool) {
	key := src + ";" + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		if newLink == nil {
			srcNode, _ := g.IncrNode(src, nil)
			dstNode, _ := g.IncrNode(dst, nil)
			l = &NetworkLink{
				SrcName: srcNode.Name,
				srcNode: srcNode,
				DstName: dstNode.Name,
				dstNode: dstNode,
			}
		} else {
			srcNode, _ := g.IncrNode(src, newLink.srcNode)
			dstNode, _ := g.IncrNode(dst, newLink.dstNode)
			// update child pointers
			newLink.srcNode = srcNode
			newLink.dstNode = dstNode
			l = newLink
		}
		g.LinksMap[key] = l
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: l,
		}
	}
	l.refCount++
	return l, !ok
}

func (g *NetworkGraph) GetLink(key string) *NetworkLink {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) (*NetworkLink, bool) {
	key := src + ";" + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return nil, false
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: l,
		}
		return l, true
	}
	return l, false
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string, newRoute *NetworkRoute) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Debugf("New Route: key=%s %v", key, hops)
		if newRoute == nil {
			path := make([]*NetworkNode, 0, len(hops))
			for i, hop := range hops {
				hopNode, _ := g.IncrNode(hop, nil)
				path = append(path, hopNode)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(hops[i-1], hop, nil)
				}
			}
			route = &NetworkRoute{
				Path:       hops,
				path:       path,
				State:      Up,
				metricRing: ring.New(100), // TODO: config
				mLock:      &sync.RWMutex{},
				updateChan: g.internalEvents,
			}
		} else {
			newRoute.path = make([]*NetworkNode, len(newRoute.Path))
			for i, nodeName := range newRoute.Path {
				// Increment the node (this will convert the name to a pointer)
				node, _ := g.IncrNode(nodeName, nil)
				// set the pointer in our `path`
				newRoute.path[i] = node

				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(newRoute.path[i-1].Name, nodeName, nil)
				}
			}
			route = newRoute
			route.updateChan = g.internalEvents
		}

		g.RoutesMap[key] = route

		g.internalEvents <- &Event{
			E:    addEvent,
			Item: route,
		}
	}

	// increment route's refcount
	route.refCount++

	return route, !ok
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return nil, false
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, nodeName := range r.Path {
			g.DecrNode(nodeName)
			if i-1 >= 0 {
				g.DecrLink(r.Path[i-1], nodeName)
			}
		}

		delete(g.RoutesMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: r,
		}
		return r, true
	}
	return r, false
}
<|endoftext|>


<commit_before>// +build ignore

package main

import (
	"github.com/klauspost/shutdown"
	"log"
	"net/http"
	"os"
	"sync"
	"syscall"
)

// This example shows a server that has logging to a file
//
// When the webserver is closed, it will close the file when all requests have
// been finished.
//
// In a real world, you would not want multiple goroutines writing to the same file
//
// To execute, use 'go run simple-func.go'

var logFile *os.File

func main() {
	// Make shutdown catch Ctrl+c and system terminate
	shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM)

	// Create a log file
	logFile, _ = os.Create("log.txt")

	// When shutdown is initiated, close the file
	shutdown.FirstFunc(func(interface{}) {
		log.Println("Closing log...")
		logFile.Close()
	}, nil)

	// Start a webserver
	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		// Get a lock, and write to the file if we get it.
		// While we have the lock the file will not be closed.
		if shutdown.Lock() {
			_, _ = logFile.WriteString(req.URL.String() + "\n")
			shutdown.Unlock()
		}
	})
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>Make example more clear.<commit_after>// +build ignore

package main

import (
	"github.com/klauspost/shutdown"
	"log"
	"net/http"
	"os"
	"syscall"
)

// This example shows a server that has logging to a file
//
// When the webserver is closed, it will close the file when all requests have
// been finished.
//
// In a real world, you would not want multiple goroutines writing to the same file
//
// To execute, use 'go run simple-func.go'

// This is the function we would like to execute at shutdown.
func closeFile(i interface{}) {
	f := i.(*os.File)
	log.Println("Closing", f.Name()+"...")
	f.Close()
}

func main() {
	// Make shutdown catch Ctrl+c and system terminate
	shutdown.OnSignal(0, os.Interrupt, syscall.SIGTERM)

	// Create a log file
	var logFile *os.File
	logFile, _ = os.Create("log.txt")

	// When shutdown is initiated, close the file
	shutdown.FirstFunc(closeFile, logFile)

	// Start a webserver
	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		// Get a lock, and write to the file if we get it.
		// While we have the lock the file will not be closed.
		if shutdown.Lock() {
			_, _ = logFile.WriteString(req.URL.String() + "\n")
			shutdown.Unlock()
		}
	})
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>