<commit_before>package testutil

import (
	"net/http"
	"testing"

	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/testutil/assert"
)

// RequestTestFunc is a function which will be run by RequestTestSpec.Run.
// This function should generally call some request method (such as those built-in
// to the lib.Client interface or external ones, usually in the lib/requests/* packages.
//
// Any validation of the results of that request method should be done as part of the
// RequestTestFunc
type RequestTestFunc func(lib.Client)

// RequestTestSpec is used to build up httptest servers for a test, then run a
// function and check it makes the correct request.
type RequestTestSpec struct {
	// MuxHandlers will be used if defined - this allows for the test to support
	// multiple endpoints, URLs, methods, etc. while still keeping as DRY as
	// possible. Otherwise, set the Method, Endpoint, URL, AssertRequest,
	// Response and StatusCode.
	MuxHandlers *MuxHandlers

	// Method is used to assert that the request was given the correct type
	// it is only used if MuxHandlers is nil
	Method string
	// Endpoint is used to build a MuxHandlers from when MuxHandlers is nil
	Endpoint lib.Endpoint
	// URL is used to build a MuxHandlers when MuxHandlers is nil
	URL string
	// Response will be JSON marshalled
	// to use a raw string (i.e. if you don't want to use JSON) cast it to
	// encoding/json.RawMessage - this will be reproduced verbatim
	Response interface{}
	// StatusCode is the status code that will be returned. If unset, will
	// default to whatever http.ResponseWriter.Write defaults to.
	// Only used if MuxHandlers is nil
	StatusCode int
	// AssertRequest is an optional func which will get called to check the
	// request object further - for example to check the URL has particular
	// query string keys
	AssertRequest func(t *testing.T, testName string, r *http.Request)

	// Auth is used to determine whether authentication should be performed - if
	// it is set to true it will be. Auth is automatically set by Run and only
	// used
	Auth bool
	// NoVerify is used to disable visit count verification in case the test is
	// expected not to get as far as calling out to this API call.
	NoVerify bool
	// visits is how many times the generated handler has been called
	visits int
}

// Run sets up fake servers, creates a client that talks to them, then passes the client to fn.
// fn should run some request method using the client & test the results of that function.
func (rts *RequestTestSpec) Run(t *testing.T, testName string, auth bool, fn RequestTestFunc) {
	client, servers, err := NewClientAndServers(t, rts)
	defer servers.Close()
	if err != nil {
		t.Fatalf("%s NewClientAndServers failed - %s", testName, err)
	}
	if auth {
		err = client.AuthWithCredentials(map[string]string{})
		if err != nil {
			t.Fatalf("%s AuthWithCredentials failed - %s", testName, err)
		}
	}

	fn(client)
	rts.Verify(t)
}

// MakeServers creates a Servers object for this RequestTestSpec. If MuxHandlers
// is set, that is used as the entirety of the basis for the Servers. Otherwise,
// makes a HandlerMap and calls MakeServers on that.
func (rts *RequestTestSpec) MakeServers(t *testing.T) Servers {
	if rts.MuxHandlers != nil {
		return rts.MuxHandlers.MakeServers(t)
	}
	return HandlerMap{
		rts.Endpoint: rts.Handler(t),
	}.MakeServers(t)
}

func (rts *RequestTestSpec) Verify(t *testing.T) {
	if rts.NoVerify {
		return
	}
	if rts.MuxHandlers == nil {
		if rts.visits == 0 {
			t.Error("never called the HTTP endpoint")
		}
	}
}

// Handler returns an http.Handler for the request(s) expected by this
// RequestTestSpec.
func (rts *RequestTestSpec) Handler(t *testing.T) http.Handler {
	return Mux{
		rts.URL: func(wr http.ResponseWriter, r *http.Request) {
			assert.Method(rts.Method)(t, "", r)
			rts.visits++
			if rts.Auth {
				assert.Auth(lib.TokenType(rts.Endpoint))(t, "", r)
			}
			if rts.AssertRequest != nil {
				rts.AssertRequest(t, "", r)
			}
			if rts.StatusCode != 0 {
				wr.WriteHeader(rts.StatusCode)
			}
			WriteJSON(t, wr, rts.Response)
		},
	}.Handler()
}
<commit_msg>add documentation comment for RequestTestSpec.Verify<commit_after>package testutil

import (
	"net/http"
	"testing"

	"github.com/BytemarkHosting/bytemark-client/lib"
	"github.com/BytemarkHosting/bytemark-client/lib/testutil/assert"
)

// RequestTestFunc is a function which will be run by RequestTestSpec.Run.
// This function should generally call some request method (such as those built-in
// to the lib.Client interface or external ones, usually in the lib/requests/* packages.
//
// Any validation of the results of that request method should be done as part of the
// RequestTestFunc
type RequestTestFunc func(lib.Client)

// RequestTestSpec is used to build up httptest servers for a test, then run a
// function and check it makes the correct request.
type RequestTestSpec struct {
	// MuxHandlers will be used if defined - this allows for the test to support
	// multiple endpoints, URLs, methods, etc. while still keeping as DRY as
	// possible. Otherwise, set the Method, Endpoint, URL, AssertRequest,
	// Response and StatusCode.
	MuxHandlers *MuxHandlers

	// Method is used to assert that the request was given the correct type
	// it is only used if MuxHandlers is nil
	Method string
	// Endpoint is used to build a MuxHandlers from when MuxHandlers is nil
	Endpoint lib.Endpoint
	// URL is used to build a MuxHandlers when MuxHandlers is nil
	URL string
	// Response will be JSON marshalled
	// to use a raw string (i.e. if you don't want to use JSON) cast it to
	// encoding/json.RawMessage - this will be reproduced verbatim
	Response interface{}
	// StatusCode is the status code that will be returned. If unset, will
	// default to whatever http.ResponseWriter.Write defaults to.
	// Only used if MuxHandlers is nil
	StatusCode int
	// AssertRequest is an optional func which will get called to check the
	// request object further - for example to check the URL has particular
	// query string keys
	AssertRequest func(t *testing.T, testName string, r *http.Request)

	// Auth is used to determine whether authentication should be performed - if
	// it is set to true it will be. Auth is automatically set by Run and only
	// used
	Auth bool
	// NoVerify is used to disable visit count verification in case the test is
	// expected not to get as far as calling out to this API call.
	NoVerify bool
	// visits is how many times the generated handler has been called
	visits int
}

// Run sets up fake servers, creates a client that talks to them, then passes the client to fn.
// fn should run some request method using the client & test the results of that function.
func (rts *RequestTestSpec) Run(t *testing.T, testName string, auth bool, fn RequestTestFunc) {
	client, servers, err := NewClientAndServers(t, rts)
	defer servers.Close()
	if err != nil {
		t.Fatalf("%s NewClientAndServers failed - %s", testName, err)
	}
	if auth {
		err = client.AuthWithCredentials(map[string]string{})
		if err != nil {
			t.Fatalf("%s AuthWithCredentials failed - %s", testName, err)
		}
	}

	fn(client)
	rts.Verify(t)
}

// MakeServers creates a Servers object for this RequestTestSpec. If MuxHandlers
// is set, that is used as the entirety of the basis for the Servers. Otherwise,
// makes a HandlerMap and calls MakeServers on that.
func (rts *RequestTestSpec) MakeServers(t *testing.T) Servers {
	if rts.MuxHandlers != nil {
		return rts.MuxHandlers.MakeServers(t)
	}
	return HandlerMap{
		rts.Endpoint: rts.Handler(t),
	}.MakeServers(t)
}

// Verify ensures that the request was visited at least once, as long as
// NoVerify is false and this RequestTestSpec is not the MuxHandlers style.
func (rts *RequestTestSpec) Verify(t *testing.T) {
	if rts.NoVerify {
		return
	}
	if rts.MuxHandlers == nil {
		if rts.visits == 0 {
			t.Error("never called the HTTP endpoint")
		}
	}
}

// Handler returns an http.Handler for the request(s) expected by this
// RequestTestSpec.
func (rts *RequestTestSpec) Handler(t *testing.T) http.Handler {
	return Mux{
		rts.URL: func(wr http.ResponseWriter, r *http.Request) {
			assert.Method(rts.Method)(t, "", r)
			rts.visits++
			if rts.Auth {
				assert.Auth(lib.TokenType(rts.Endpoint))(t, "", r)
			}
			if rts.AssertRequest != nil {
				rts.AssertRequest(t, "", r)
			}
			if rts.StatusCode != 0 {
				wr.WriteHeader(rts.StatusCode)
			}
			WriteJSON(t, wr, rts.Response)
		},
	}.Handler()
}
<|endoftext|>


<commit_before>package drouter

import (
	"net"
	"strconv"
	"errors"
	"strings"
	//"os/exec"
	//"fmt"
	"time"
	"os"
	"bufio"
	//"os/signal"
	//"syscall"
	//"bytes"
	//"io/ioutil"

	log "github.com/Sirupsen/logrus"
	//"github.com/samalba/dockerclient"
	dockerclient "github.com/docker/engine-api/client"
	dockertypes "github.com/docker/engine-api/types"
	dockerfilters "github.com/docker/engine-api/types/filters"
	dockernetworks "github.com/docker/engine-api/types/network"
	dockerevents "github.com/docker/engine-api/types/events"
	"golang.org/x/net/context"
	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netns"
	"github.com/ziutek/utils/netaddr"
	"github.com/llimllib/ipaddress"
	"github.com/vdemeester/docker-events"
)

var (
	docker                *dockerclient.Client
	self_container        dockertypes.ContainerJSON
	networks              = make(map[string]bool)
	host_ns_h             *netlink.Handle
	self_ns_h             *netlink.Handle
	host_route_link_index int
	host_route_gw		    net.IP
	my_pid                = os.Getpid()
)

func init() {
	var err error

	if my_pid == 1 {
		log.Fatal("Running as Pid 1. drouter must be run with --pid=host")
	}

	defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
	docker, err = dockerclient.NewClient("unix:///var/run/docker.sock", "v1.23", nil, defaultHeaders)
	if err != nil {
		log.Error("Error connecting to docker socket")
		log.Fatal(err)
	}
	self_container, err = getSelf()
	if err != nil {
		log.Error("Error getting self container. Is this processs running in a container? Is the docker socket passed through?")
		log.Fatal(err)
	}

	// Prepopulate networks that this container is a member of
	for _, settings := range self_container.NetworkSettings.Networks {
		networks[settings.NetworkID] = true
	}

	self_ns, err := netns.Get()
	if err != nil {
		log.Error("Error getting self namespace.")
		log.Fatal(err)
	}
	self_ns_h, err = netlink.NewHandleAt(self_ns)
	if err != nil {
		log.Error("Error getting handle at self namespace.")
		log.Fatal(err)
	}
	host_ns, err := netns.GetFromPid(1)
	if err != nil {
		log.Error("Error getting host namespace. Is this container running in priveleged mode?")
		log.Fatal(err)
	}
	host_ns_h, err = netlink.NewHandleAt(host_ns)
	if err != nil {
		log.Error("Error getting handle at host namespace.")
		log.Fatal(err)
	}
}

// Loop to watch for new networks created and create interfaces when needed
func WatchNetworks(IPOffset int) {
	log.Info("Watching Networks")
	for {
		nets, err := docker.NetworkList(context.Background(), dockertypes.NetworkListOptions{ Filters: dockerfilters.NewArgs(), })
		if err != nil {
			log.Error("Error getting network list")
			log.Error(err)
		}
		for i := range nets {
			drouter_str := nets[i].Options["drouter"]
			drouter := false
			if drouter_str != "" {
				drouter, err = strconv.ParseBool(drouter_str)
				if err != nil {
					log.Errorf("Error parsing drouter option: %v", drouter_str)
					log.Error(err)
				}
			}

			if drouter && !networks[nets[i].ID] {
				log.Debugf("Joining Net: %+v", nets[i])
				err := joinNet(&nets[i], IPOffset)
				if err != nil {
					log.Errorf("Error joining network: %v", nets[i])
					log.Error(err)
				}
			} else if !drouter && networks[nets[i].ID] {
				log.Debugf("Leaving Net: %+v", nets[i])
				err := leaveNet(&nets[i])
				if err != nil {
					log.Errorf("Error leaving network: %v", nets[i])
					log.Error(err)
				}
			}
		}
		time.Sleep(5 * time.Second)
	}
}

func WatchEvents() {
	errChan := events.Monitor(context.Background(), docker, dockertypes.EventsOptions{}, func(event dockerevents.Message) {
		log.Debugf("Event: %v", event)
		log.Debugf("Event.Status: %v", event.Status)
		log.Debugf("Event.ID: %v", event.ID)
		log.Debugf("Event.From: %v", event.From)
		log.Debugf("Event.Type: %v", event.Type)
		log.Debugf("Event.Action: %v", event.Action)
		log.Debugf("Event.Actor: %v", event.Actor)
})
	if err := <-errChan; err != nil {
		log.Error(err)
	}
}

func joinNet(n *dockertypes.NetworkResource, IPOffset int) error {
	endpointSettings := &dockernetworks.EndpointSettings{}
	if IPOffset != 0 {
		for i := range n.IPAM.Config {
			ipamconfig := n.IPAM.Config[i]
			log.Debugf("ip-offset configured")
			_, subnet, err := net.ParseCIDR(ipamconfig.Subnet)
			if err != nil {
				return err
			}
			var ip net.IP
			if IPOffset > 0 {
				ip = netaddr.IPAdd(subnet.IP, IPOffset)
			} else {
				last := ipaddress.LastAddress(subnet)
				ip = netaddr.IPAdd(last, IPOffset)
			}
			log.Debugf("Setting IP to %v", ip)
			if endpointSettings.IPAddress == "" {
				endpointSettings.IPAddress = ip.String()
				endpointSettings.IPAMConfig =&dockernetworks.EndpointIPAMConfig{
					IPv4Address: ip.String(),
				}
			} else {
				endpointSettings.Aliases = append(endpointSettings.Aliases, ip.String())
			}
		}
	}

	err := docker.NetworkConnect(context.Background(), n.ID, self_container.ID, endpointSettings)
	if err != nil {
		return err
	}
	networks[n.ID] = true
	for i := range n.IPAM.Config {
		ipamconfig := n.IPAM.Config[i]
		_, dst, err := net.ParseCIDR(ipamconfig.Subnet)
		if err != nil {
			return err
		}
		route := &netlink.Route{
			LinkIndex: host_route_link_index,
			Gw: host_route_gw,
			Dst: dst,
		}
		err = host_ns_h.RouteAdd(route)
		if err != nil {
			return err
		}
	}
	return nil
}

func leaveNet(n *dockertypes.NetworkResource) error {
	err := docker.NetworkDisconnect(context.Background(), n.ID, self_container.ID, true)
	if err != nil {
		return err
	}
	networks[n.ID] = false
	return nil
}

func getSelf() (dockertypes.ContainerJSON, error) {
	cgroup, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return dockertypes.ContainerJSON{}, err
	}
	defer cgroup.Close()

	scanner := bufio.NewScanner(cgroup)
	for scanner.Scan() {
		line := strings.Split(scanner.Text(), "/")
		id := line[len(line) - 1]
		containerInfo, err := docker.ContainerInspect(context.Background(), id)
		if err != nil {
			log.Warnf("Error inspecting container: %v", id)
			log.Warn(err)
			continue
		}
		return containerInfo, nil
	}
	return dockertypes.ContainerJSON{}, errors.New("Container not found")
}

func MakeP2PLink(p2p_addr string) error {
	host_link_veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "drouter_veth0"},
		PeerName:  "drouter_veth1",
	}
	err := host_ns_h.LinkAdd(host_link_veth)
	if err != nil {
		return err
	}
	host_link, err := host_ns_h.LinkByName("drouter_veth0")
	if err != nil {
		return err
	}
	host_route_link_index = host_link.Attrs().Index

	int_link, err := host_ns_h.LinkByName("drouter_veth1")
	if err != nil {
		return err
	}
	err = host_ns_h.LinkSetNsPid(int_link, my_pid)
	if err != nil {
		return err
	}
	int_link, err = self_ns_h.LinkByName("drouter_veth1")
	if err != nil {
		return err
	}

	_, p2p_net, err := net.ParseCIDR(p2p_addr)
	if err != nil {
		return err
	}

	host_addr := *p2p_net
	host_addr.IP = netaddr.IPAdd(host_addr.IP, 1)
	host_netlink_addr := &netlink.Addr{
		IPNet: &host_addr,
		Label: "",
	}
	err = host_ns_h.AddrAdd(host_link, host_netlink_addr)
	if err != nil {
		return err
	}

	int_addr := *p2p_net
	int_addr.IP = netaddr.IPAdd(int_addr.IP, 2)
	int_netlink_addr := &netlink.Addr{
		IPNet: &int_addr,
		Label: "",
	}
	err = self_ns_h.AddrAdd(int_link, int_netlink_addr)
	if err != nil {
		return err
	}

	host_route_gw = int_addr.IP

	err = self_ns_h.LinkSetUp(int_link)
	if err != nil {
		return err
	}

	err = host_ns_h.LinkSetUp(host_link)
	if err != nil {
		return err
	}

	return nil
}

func Cleanup() error {
	log.Info("Cleaning Up")
	return removeP2PLink()
}

func removeP2PLink() error {
	host_link, err := host_ns_h.LinkByName("drouter_veth0")
	if err != nil {
		return err
	}
	return host_ns_h.LinkDel(host_link)
}

<commit_msg>need to store gateway in networks[] so that it can be quickly added to the container. In progress<commit_after>package drouter

import (
	"net"
	"strconv"
	"errors"
	"strings"
	//"os/exec"
	//"fmt"
	"time"
	"os"
	"bufio"
	//"os/signal"
	//"syscall"
	//"bytes"
	//"io/ioutil"

	log "github.com/Sirupsen/logrus"
	//"github.com/samalba/dockerclient"
	dockerclient "github.com/docker/engine-api/client"
	dockertypes "github.com/docker/engine-api/types"
	dockerfilters "github.com/docker/engine-api/types/filters"
	dockernetworks "github.com/docker/engine-api/types/network"
	dockerevents "github.com/docker/engine-api/types/events"
	"golang.org/x/net/context"
	"github.com/vishvananda/netlink"
	"github.com/vishvananda/netns"
	"github.com/ziutek/utils/netaddr"
	"github.com/llimllib/ipaddress"
	"github.com/vdemeester/docker-events"
)

var (
	docker                *dockerclient.Client
	self_container        dockertypes.ContainerJSON
	networks              = make(map[string]net.IP)
	host_ns_h             *netlink.Handle
	self_ns_h             *netlink.Handle
	host_route_link_index int
	host_route_gw		    net.IP
	my_pid                = os.Getpid()
)

func init() {
	var err error

	if my_pid == 1 {
		log.Fatal("Running as Pid 1. drouter must be run with --pid=host")
	}

	defaultHeaders := map[string]string{"User-Agent": "engine-api-cli-1.0"}
	docker, err = dockerclient.NewClient("unix:///var/run/docker.sock", "v1.23", nil, defaultHeaders)
	if err != nil {
		log.Error("Error connecting to docker socket")
		log.Fatal(err)
	}
	self_container, err = getSelf()
	if err != nil {
		log.Error("Error getting self container. Is this processs running in a container? Is the docker socket passed through?")
		log.Fatal(err)
	}

	// Prepopulate networks that this container is a member of
	for _, settings := range self_container.NetworkSettings.Networks {
		networks[settings.NetworkID] = net.ParseIP(settings.IPAMConfig[0].IPv4Address)
	}

	self_ns, err := netns.Get()
	if err != nil {
		log.Error("Error getting self namespace.")
		log.Fatal(err)
	}
	self_ns_h, err = netlink.NewHandleAt(self_ns)
	if err != nil {
		log.Error("Error getting handle at self namespace.")
		log.Fatal(err)
	}
	host_ns, err := netns.GetFromPid(1)
	if err != nil {
		log.Error("Error getting host namespace. Is this container running in priveleged mode?")
		log.Fatal(err)
	}
	host_ns_h, err = netlink.NewHandleAt(host_ns)
	if err != nil {
		log.Error("Error getting handle at host namespace.")
		log.Fatal(err)
	}
}

// Loop to watch for new networks created and create interfaces when needed
func WatchNetworks(IPOffset int) {
	log.Info("Watching Networks")
	for {
		nets, err := docker.NetworkList(context.Background(), dockertypes.NetworkListOptions{ Filters: dockerfilters.NewArgs(), })
		if err != nil {
			log.Error("Error getting network list")
			log.Error(err)
		}
		for i := range nets {
			drouter_str := nets[i].Options["drouter"]
			drouter := false
			if drouter_str != "" {
				drouter, err = strconv.ParseBool(drouter_str)
				if err != nil {
					log.Errorf("Error parsing drouter option: %v", drouter_str)
					log.Error(err)
				}
			}

			if drouter && !networks[nets[i].ID] {
				log.Debugf("Joining Net: %+v", nets[i])
				err := joinNet(&nets[i], IPOffset)
				if err != nil {
					log.Errorf("Error joining network: %v", nets[i])
					log.Error(err)
				}
			} else if !drouter && networks[nets[i].ID] {
				log.Debugf("Leaving Net: %+v", nets[i])
				err := leaveNet(&nets[i])
				if err != nil {
					log.Errorf("Error leaving network: %v", nets[i])
					log.Error(err)
				}
			}
		}
		time.Sleep(5 * time.Second)
	}
}

func WatchEvents() {
	errChan := events.Monitor(context.Background(), docker, dockertypes.EventsOptions{}, func(event dockerevents.Message) {
		if event.Type != "network" { return }
		if event.Action != "connect" { return }
		// don't run on self events
		if event.Actor.Attributes["container"] == self_container.ID { return }
		// don't run if this network is not being managed
		if !networks[event.Actor.ID] { return }

		log.Debugf("Event.Actor: %v", event.Actor)
	})
	if err := <-errChan; err != nil {
		log.Error(err)
	}
}

func joinNet(n *dockertypes.NetworkResource, IPOffset int) error {
	endpointSettings := &dockernetworks.EndpointSettings{}
	if IPOffset != 0 {
		for i := range n.IPAM.Config {
			ipamconfig := n.IPAM.Config[i]
			log.Debugf("ip-offset configured")
			_, subnet, err := net.ParseCIDR(ipamconfig.Subnet)
			if err != nil {
				return err
			}
			var ip net.IP
			if IPOffset > 0 {
				ip = netaddr.IPAdd(subnet.IP, IPOffset)
			} else {
				last := ipaddress.LastAddress(subnet)
				ip = netaddr.IPAdd(last, IPOffset)
			}
			log.Debugf("Setting IP to %v", ip)
			if endpointSettings.IPAddress == "" {
				endpointSettings.IPAddress = ip.String()
				endpointSettings.IPAMConfig =&dockernetworks.EndpointIPAMConfig{
					IPv4Address: ip.String(),
				}
			} else {
				endpointSettings.Aliases = append(endpointSettings.Aliases, ip.String())
			}
		}
	}

	err := docker.NetworkConnect(context.Background(), n.ID, self_container.ID, endpointSettings)
	if err != nil {
		return err
	}
	networks[n.ID] = true
	for i := range n.IPAM.Config {
		ipamconfig := n.IPAM.Config[i]
		_, dst, err := net.ParseCIDR(ipamconfig.Subnet)
		if err != nil {
			return err
		}
		route := &netlink.Route{
			LinkIndex: host_route_link_index,
			Gw: host_route_gw,
			Dst: dst,
		}
		err = host_ns_h.RouteAdd(route)
		if err != nil {
			return err
		}
	}
	return nil
}

func leaveNet(n *dockertypes.NetworkResource) error {
	err := docker.NetworkDisconnect(context.Background(), n.ID, self_container.ID, true)
	if err != nil {
		return err
	}
	networks[n.ID] = false
	return nil
}

func getSelf() (dockertypes.ContainerJSON, error) {
	cgroup, err := os.Open("/proc/self/cgroup")
	if err != nil {
		return dockertypes.ContainerJSON{}, err
	}
	defer cgroup.Close()

	scanner := bufio.NewScanner(cgroup)
	for scanner.Scan() {
		line := strings.Split(scanner.Text(), "/")
		id := line[len(line) - 1]
		containerInfo, err := docker.ContainerInspect(context.Background(), id)
		if err != nil {
			log.Warnf("Error inspecting container: %v", id)
			log.Warn(err)
			continue
		}
		return containerInfo, nil
	}
	return dockertypes.ContainerJSON{}, errors.New("Container not found")
}

func MakeP2PLink(p2p_addr string) error {
	host_link_veth := &netlink.Veth{
		LinkAttrs: netlink.LinkAttrs{Name: "drouter_veth0"},
		PeerName:  "drouter_veth1",
	}
	err := host_ns_h.LinkAdd(host_link_veth)
	if err != nil {
		return err
	}
	host_link, err := host_ns_h.LinkByName("drouter_veth0")
	if err != nil {
		return err
	}
	host_route_link_index = host_link.Attrs().Index

	int_link, err := host_ns_h.LinkByName("drouter_veth1")
	if err != nil {
		return err
	}
	err = host_ns_h.LinkSetNsPid(int_link, my_pid)
	if err != nil {
		return err
	}
	int_link, err = self_ns_h.LinkByName("drouter_veth1")
	if err != nil {
		return err
	}

	_, p2p_net, err := net.ParseCIDR(p2p_addr)
	if err != nil {
		return err
	}

	host_addr := *p2p_net
	host_addr.IP = netaddr.IPAdd(host_addr.IP, 1)
	host_netlink_addr := &netlink.Addr{
		IPNet: &host_addr,
		Label: "",
	}
	err = host_ns_h.AddrAdd(host_link, host_netlink_addr)
	if err != nil {
		return err
	}

	int_addr := *p2p_net
	int_addr.IP = netaddr.IPAdd(int_addr.IP, 2)
	int_netlink_addr := &netlink.Addr{
		IPNet: &int_addr,
		Label: "",
	}
	err = self_ns_h.AddrAdd(int_link, int_netlink_addr)
	if err != nil {
		return err
	}

	host_route_gw = int_addr.IP

	err = self_ns_h.LinkSetUp(int_link)
	if err != nil {
		return err
	}

	err = host_ns_h.LinkSetUp(host_link)
	if err != nil {
		return err
	}

	return nil
}

func Cleanup() error {
	log.Info("Cleaning Up")
	return removeP2PLink()
}

func removeP2PLink() error {
	host_link, err := host_ns_h.LinkByName("drouter_veth0")
	if err != nil {
		return err
	}
	return host_ns_h.LinkDel(host_link)
}

<|endoftext|>


<commit_before>package main

import (
	"os"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/authelia/authelia/internal/utils"
)

func buildAutheliaBinary() {
	cmd := utils.CommandWithStdout("go", "build", "-o", "../../"+OutputDir+"/authelia")
	cmd.Dir = "cmd/authelia"

	cmd.Env = append(os.Environ(),
		"GOOS=linux", "GOARCH=amd64", "CGO_ENABLED=1")

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func buildFrontend() {
	cmd := utils.CommandWithStdout("yarn", "install")
	cmd.Dir = webDirectory

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("yarn", "build")
	cmd.Dir = webDirectory

	cmd.Env = append(os.Environ(), "INLINE_RUNTIME_CHUNK=false")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("rm", "-rf", "internal/server/public_html")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	err = os.Rename("web/build", "internal/server/public_html")
	if err != nil {
		log.Fatal(err)
	}
}

func buildSwagger() {
	swaggerVer := "3.48.0"
	cmd := utils.CommandWithStdout("bash", "-c", "wget -q https://github.com/swagger-api/swagger-ui/archive/v"+swaggerVer+".tar.gz -O ./v"+swaggerVer+".tar.gz")

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("cp", "-r", "api", "internal/server/public_html")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("tar", "-C", "internal/server/public_html/api", "--exclude=index.html", "--strip-components=2", "-xf", "v"+swaggerVer+".tar.gz", "swagger-ui-"+swaggerVer+"/dist")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("rm", "./v"+swaggerVer+".tar.gz")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func cleanAssets() {
	if err := os.Rename("internal/server/public_html", OutputDir+"/public_html"); err != nil {
		log.Fatal(err)
	}

	cmd := utils.CommandWithStdout("mkdir", "-p", "internal/server/public_html/api")

	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("bash", "-c", "touch internal/server/public_html/{index.html,api/index.html,api/openapi.yml}")

	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
}

// Build build Authelia.
func Build(cobraCmd *cobra.Command, args []string) {
	log.Info("Building Authelia...")

	Clean(cobraCmd, args)

	log.Debug("Creating `" + OutputDir + "` directory")
	err := os.MkdirAll(OutputDir, os.ModePerm)

	if err != nil {
		log.Fatal(err)
	}

	log.Debug("Building Authelia frontend...")
	buildFrontend()

	log.Debug("Building swagger-ui frontend...")
	buildSwagger()

	log.Debug("Building Authelia Go binary...")
	buildAutheliaBinary()
	cleanAssets()
}
<commit_msg>build(deps): update swagger-ui to v3.50.0 (#2070)<commit_after>package main

import (
	"os"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"

	"github.com/authelia/authelia/internal/utils"
)

func buildAutheliaBinary() {
	cmd := utils.CommandWithStdout("go", "build", "-o", "../../"+OutputDir+"/authelia")
	cmd.Dir = "cmd/authelia"

	cmd.Env = append(os.Environ(),
		"GOOS=linux", "GOARCH=amd64", "CGO_ENABLED=1")

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func buildFrontend() {
	cmd := utils.CommandWithStdout("yarn", "install")
	cmd.Dir = webDirectory

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("yarn", "build")
	cmd.Dir = webDirectory

	cmd.Env = append(os.Environ(), "INLINE_RUNTIME_CHUNK=false")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("rm", "-rf", "internal/server/public_html")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	err = os.Rename("web/build", "internal/server/public_html")
	if err != nil {
		log.Fatal(err)
	}
}

func buildSwagger() {
	swaggerVer := "3.50.0"
	cmd := utils.CommandWithStdout("bash", "-c", "wget -q https://github.com/swagger-api/swagger-ui/archive/v"+swaggerVer+".tar.gz -O ./v"+swaggerVer+".tar.gz")

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("cp", "-r", "api", "internal/server/public_html")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("tar", "-C", "internal/server/public_html/api", "--exclude=index.html", "--strip-components=2", "-xf", "v"+swaggerVer+".tar.gz", "swagger-ui-"+swaggerVer+"/dist")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("rm", "./v"+swaggerVer+".tar.gz")

	err = cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
}

func cleanAssets() {
	if err := os.Rename("internal/server/public_html", OutputDir+"/public_html"); err != nil {
		log.Fatal(err)
	}

	cmd := utils.CommandWithStdout("mkdir", "-p", "internal/server/public_html/api")

	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}

	cmd = utils.CommandWithStdout("bash", "-c", "touch internal/server/public_html/{index.html,api/index.html,api/openapi.yml}")

	if err := cmd.Run(); err != nil {
		log.Fatal(err)
	}
}

// Build build Authelia.
func Build(cobraCmd *cobra.Command, args []string) {
	log.Info("Building Authelia...")

	Clean(cobraCmd, args)

	log.Debug("Creating `" + OutputDir + "` directory")
	err := os.MkdirAll(OutputDir, os.ModePerm)

	if err != nil {
		log.Fatal(err)
	}

	log.Debug("Building Authelia frontend...")
	buildFrontend()

	log.Debug("Building swagger-ui frontend...")
	buildSwagger()

	log.Debug("Building Authelia Go binary...")
	buildAutheliaBinary()
	cleanAssets()
}
<|endoftext|>


<commit_before>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	"github.com/funkygao/golib/sync2"
	"github.com/funkygao/golib/timewheel"
	log "github.com/funkygao/log4go"
	"github.com/samuel/go-zookeeper/zk"
)

type subServer struct {
	*webServer

	idleConnsWg   sync2.WaitGroupTimeout // wait for all inflight http connections done
	closedConnCh  chan string            // channel of remote addr
	idleConns     map[net.Conn]struct{}
	idleConnsLock sync.Mutex

	auditor log.Logger

	timer *timewheel.TimeWheel

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics *subMetrics

	throttleBadGroup *ratelimiter.LeakyBuckets
	goodGroupClients map[string]struct{} // key is remote addr(port inclusive)
	goodGroupLock    sync.RWMutex
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:        newWebServer("sub_server", httpAddr, httpsAddr, maxClients, Options.HttpReadTimeout, gw),
		closedConnCh:     make(chan string, 1<<10),
		idleConns:        make(map[net.Conn]struct{}, 200),
		wsReadLimit:      8 << 10,
		wsPongWait:       time.Minute,
		timer:            timewheel.NewTimeWheel(time.Second, 120),
		throttleBadGroup: ratelimiter.NewLeakyBuckets(3, time.Minute),
		goodGroupClients: make(map[string]struct{}, 100),
		ackShutdown:      0,
		ackCh:            make(chan ackOffsets, 100),
		ackedOffsets:     make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)
		atomic.AddInt32(&this.activeConnN, 1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

			this.idleConnsLock.Lock()
			delete(this.idleConns, c)
			this.idleConnsLock.Unlock()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c] = struct{}{}
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

		atomic.AddInt32(&this.activeConnN, -1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		this.closedConnCh <- c.RemoteAddr().String()
		this.idleConnsWg.Done()
		atomic.AddInt32(&this.activeConnN, -1)

		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()
	}
}

func (this *subServer) waitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	if this.idleConnsWg.WaitTimeout(Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()
	this.timer.Stop()

	this.gw.wg.Done()
	close(this.closed)
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	var n int64
	const flushPerN int64 = 4000
	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}

				n++
				if n%flushPerN == 0 {
					// TODO reset the tick timer?
					this.commitOffsets()
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("cluster[%s] group[%s] commit offset {T:%s/%d O:%d}", cluster, group, topic, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group, strconv.Itoa(partition), offset); err != nil {
						log.Error("cluster[%s] group[%s] commit offset {T:%s/%d O:%d} %v", cluster, group, topic, partition, offset, err)

						if err == zk.ErrNoNode {
							// invalid offset commit request, will not retry
							this.ackedOffsets[cluster][topic][group][partition] = -1
						}
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<commit_msg>BUG FIX: a potential DDOS bug<commit_after>package gateway

import (
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/funkygao/gafka/cmd/kateway/meta"
	"github.com/funkygao/golib/ratelimiter"
	"github.com/funkygao/golib/sync2"
	"github.com/funkygao/golib/timewheel"
	log "github.com/funkygao/log4go"
	"github.com/samuel/go-zookeeper/zk"
)

type subServer struct {
	*webServer

	idleConnsWg   sync2.WaitGroupTimeout // wait for all inflight http connections done
	closedConnCh  chan string            // channel of remote addr
	idleConns     map[net.Conn]struct{}
	idleConnsLock sync.Mutex

	auditor log.Logger

	timer *timewheel.TimeWheel

	// websocket heartbeat configuration
	wsReadLimit int64
	wsPongWait  time.Duration

	shutdownOnce sync.Once
	ackShutdown  int32                                          // sync shutdown with ack handlers goroutines
	ackCh        chan ackOffsets                                // client ack'ed offsets
	ackedOffsets map[string]map[string]map[string]map[int]int64 // [cluster][topic][group][partition]: offset

	subMetrics *subMetrics

	throttleBadGroup *ratelimiter.LeakyBuckets
	goodGroupClients map[string]struct{} // key is remote addr(port inclusive)
	goodGroupLock    sync.RWMutex
}

func newSubServer(httpAddr, httpsAddr string, maxClients int, gw *Gateway) *subServer {
	this := &subServer{
		webServer:        newWebServer("sub_server", httpAddr, httpsAddr, maxClients, Options.HttpReadTimeout, gw),
		closedConnCh:     make(chan string, 1<<10),
		idleConns:        make(map[net.Conn]struct{}, 200),
		wsReadLimit:      8 << 10,
		wsPongWait:       time.Minute,
		timer:            timewheel.NewTimeWheel(time.Second, 120),
		throttleBadGroup: ratelimiter.NewLeakyBuckets(3, time.Minute),
		goodGroupClients: make(map[string]struct{}, 100),
		ackShutdown:      0,
		ackCh:            make(chan ackOffsets, 100),
		ackedOffsets:     make(map[string]map[string]map[string]map[int]int64),
	}
	this.subMetrics = NewSubMetrics(this.gw)
	this.waitExitFunc = this.waitExit
	this.connStateFunc = this.connStateHandler

	if this.httpsServer != nil {
		this.httpsServer.ConnState = this.connStateFunc
	}

	if this.httpServer != nil {
		this.httpServer.ConnState = this.connStateFunc
	}

	this.auditor = log.NewDefaultLogger(log.TRACE)
	this.auditor.DeleteFilter("stdout")

	_ = os.Mkdir("audit", os.ModePerm)
	rotateEnabled, discardWhenDiskFull := true, false
	filer := log.NewFileLogWriter("audit/sub_audit.log", rotateEnabled, discardWhenDiskFull, 0644)
	if filer == nil {
		panic("failed to open sub audit log")
	}
	filer.SetFormat("[%d %T] [%L] (%S) %M")
	if Options.LogRotateSize > 0 {
		filer.SetRotateSize(Options.LogRotateSize)
	}
	filer.SetRotateLines(0)
	filer.SetRotateDaily(true)
	this.auditor.AddFilter("file", logLevel, filer)

	return this
}

func (this *subServer) Start() {
	this.gw.wg.Add(1)
	go this.ackCommitter()

	this.subMetrics.Load()
	this.webServer.Start()
}

func (this *subServer) connStateHandler(c net.Conn, cs http.ConnState) {
	switch cs {
	case http.StateNew:
		// Connections begin at StateNew and then
		// transition to either StateActive or StateClosed
		this.idleConnsWg.Add(1)
		atomic.AddInt32(&this.activeConnN, 1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Inc(1)
		}

	case http.StateActive:
		// StateActive fires before the request has entered a handler
		// and doesn't fire again until the request has been
		// handled.
		// After the request is handled, the state
		// transitions to StateClosed, StateHijacked, or StateIdle.
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

	case http.StateIdle:
		// StateIdle represents a connection that has finished
		// handling a request and is in the keep-alive state, waiting
		// for a new request. Connections transition from StateIdle
		// to either StateActive or StateClosed.
		select {
		case <-this.gw.shutdownCh:
			// actively close the client safely because IO is all done
			c.Close()

			this.idleConnsLock.Lock()
			delete(this.idleConns, c)
			this.idleConnsLock.Unlock()

		default:
			this.idleConnsLock.Lock()
			this.idleConns[c] = struct{}{}
			this.idleConnsLock.Unlock()
		}

	case http.StateHijacked:
		// websocket steals the socket
		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()

		atomic.AddInt32(&this.activeConnN, -1)

		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)

			this.gw.svrMetrics.ConcurrentSubWs.Inc(1)
		}

	case http.StateClosed:
		if this.gw != nil && !Options.DisableMetrics {
			this.gw.svrMetrics.ConcurrentSub.Dec(1)
		}

		remoteAddr := c.RemoteAddr().String()
		log.Debug("closed from %s", remoteAddr)

		this.goodGroupLock.Lock()
		delete(this.goodGroupClients, remoteAddr)
		this.goodGroupLock.Unlock()

		this.closedConnCh <- remoteAddr

		this.idleConnsWg.Done()
		atomic.AddInt32(&this.activeConnN, -1)

		this.idleConnsLock.Lock()
		delete(this.idleConns, c)
		this.idleConnsLock.Unlock()
	}
}

func (this *subServer) waitExit(exit <-chan struct{}) {
	<-exit

	if this.httpServer != nil {
		// HTTP response will have "Connection: close"
		this.httpServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpServer.Addr)
	}

	if this.httpsServer != nil {
		// HTTP response will have "Connection: close"
		this.httpsServer.SetKeepAlivesEnabled(false)

		// avoid new connections
		if err := this.httpsListener.Close(); err != nil {
			log.Error(err.Error())
		}

		log.Trace("%s on %s listener closed", this.name, this.httpsServer.Addr)
	}

	this.idleConnsLock.Lock()
	t := time.Now().Add(time.Millisecond * 100)
	for c := range this.idleConns {
		c.SetReadDeadline(t)
	}
	this.idleConnsLock.Unlock()

	if this.idleConnsWg.WaitTimeout(Options.SubTimeout) {
		log.Warn("%s waiting for all connected client close timeout: %s",
			this.name, Options.SubTimeout)
	}

	this.subMetrics.Flush()
	this.timer.Stop()

	this.gw.wg.Done()
	close(this.closed)
}

func (this *subServer) ackCommitter() {
	ticker := time.NewTicker(time.Second * 30)
	defer func() {
		ticker.Stop()
		log.Debug("ack committer done")
		this.gw.wg.Done()
	}()

	var n int64
	const flushPerN int64 = 4000
	for {
		select {
		case <-this.gw.shutdownCh:
			this.shutdownOnce.Do(func() {
				atomic.AddInt32(&this.ackShutdown, -1)

				for {
					// waiting for all ack handlers finish
					if atomic.LoadInt32(&this.ackShutdown) <= -1 {
						break
					}

					time.Sleep(time.Millisecond * 10)
				}
				close(this.ackCh)
			})

		case acks, ok := <-this.ackCh:
			if ok {
				for _, ack := range acks {
					if _, present := this.ackedOffsets[ack.cluster]; !present {
						this.ackedOffsets[ack.cluster] = make(map[string]map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.cluster][ack.topic]; !present {
						this.ackedOffsets[ack.cluster][ack.topic] = make(map[string]map[int]int64)
					}
					if _, present := this.ackedOffsets[ack.topic][ack.group]; !present {
						this.ackedOffsets[ack.cluster][ack.topic][ack.group] = make(map[int]int64)
					}

					// TODO validation
					this.ackedOffsets[ack.cluster][ack.topic][ack.group][ack.Partition] = ack.Offset
				}

				n++
				if n%flushPerN == 0 {
					// TODO reset the tick timer?
					this.commitOffsets()
				}
			} else {
				// channel buffer drained, flush all offsets
				// zk is still alive, safe to commit offsets
				this.commitOffsets()
				return
			}

		case <-ticker.C:
			this.commitOffsets()
		}
	}

}

func (this *subServer) commitOffsets() {
	for cluster, clusterTopic := range this.ackedOffsets {
		zkcluster := meta.Default.ZkCluster(cluster)

		for topic, groupPartition := range clusterTopic {
			for group, partitionOffset := range groupPartition {
				for partition, offset := range partitionOffset {
					if offset == -1 {
						// this slot is empty
						continue
					}

					log.Debug("cluster[%s] group[%s] commit offset {T:%s/%d O:%d}", cluster, group, topic, partition, offset)

					if err := zkcluster.ResetConsumerGroupOffset(topic, group, strconv.Itoa(partition), offset); err != nil {
						log.Error("cluster[%s] group[%s] commit offset {T:%s/%d O:%d} %v", cluster, group, topic, partition, offset, err)

						if err == zk.ErrNoNode {
							// invalid offset commit request, will not retry
							this.ackedOffsets[cluster][topic][group][partition] = -1
						}
					} else {
						// mark this slot empty
						this.ackedOffsets[cluster][topic][group][partition] = -1
					}
				}
			}
		}
	}

}
<|endoftext|>


<commit_before>package i2cm

import (
	"errors"
)

type Transactor8x8 interface {
	Transact8x8(addr Addr, regaddr uint8, w []byte, r []byte) (nw, nr int, err error)
}

type transactor8x8 struct {
	m I2CMaster
}

func NewTransact8x8(m I2CMaster) Transactor8x8 {
	if t, ok := m.(Transactor8x8); ok {
		return t
	}
	return transactor8x8{m}
}

func (t transactor8x8) Transact8x8(addr Addr, regaddr uint8, w []byte, r []byte) (int, int, error) {
	return I2CMasterTransact8x8(t.m, addr, regaddr, w, r)
}

func I2CMasterTransact8x8(m I2CMaster, addr Addr, regaddr uint8, w []byte, r []byte) (int, int, error) {
	nr := 0
	nw := 0

	if addr.GetAddrLen() != 7 {
		return nw, nr, errors.New("I2CMasterTransact8x8: only 7 bit addresses are supported")
	}

	if err := m.Start(); err != nil {
		return nw, nr, err
	}

	// inner function handles the whole transaction between
	// but not including the start and the stop bit
	err := func() error {
		// address device
		addrb := uint8(addr.GetBaseAddr() << 1)
		if err := m.WriteByte(addrb); err != nil {
			if err == NACKReceived {
				return NoSuchDevice
			}
			return err
		}

		// write regaddr
		if err := m.WriteByte(regaddr); err != nil {
			return err
		}

		// write w
		for _, b := range w {
			if err := m.WriteByte(b); err != nil {
				return err
			}

			nw++
		}

		// read part of transaction is only performed if desired
		if len(r) > 0 {
			// start again
			if err := m.Start(); err != nil {
				return err
			}

			// write device's read address
			if err := m.WriteByte(addrb | 0x01); err != nil {
				if err == NACKReceived {
					return NoSuchDevice
				}
				return err
			}

			for i := 0; i < len(r); i++ {
				ack := true
				if i == len(r)-1 {
					ack = false
				}
				rb, err := m.ReadByte(ack)
				if err != nil {
					return err
				}

				r[i] = rb

				nr++
			}

		}

		return nil
	}()

	if err != nil {
		// if there already was an error, the error from stop is ignored
		// and the first error is reported
		m.Stop()
	} else {
		err = m.Stop()
	}

	return nw, nr, err
}
<commit_msg>add consolidating Tranasctor<commit_after>package i2cm

import (
	"errors"
)

// Transactor encompasses all implemented I2C bus transaction
// types.
type Transactor interface {
	Transactor8x8
}

type transactor struct {
	Transactor8x8
}

// NewTransactor returns all implemented I2C transactors
// based on the argument I2CMaster. This is a convenience
// function which consolidates the results of the
// NewTransact*x* family of functions.
func NewTransactor(m I2CMaster) Transactor {
	var t transactor
	t.Transactor8x8 = NewTransact8x8(m)

	return &t
}

type Transactor8x8 interface {
	Transact8x8(addr Addr, regaddr uint8, w []byte, r []byte) (nw, nr int, err error)
}

type transactor8x8 struct {
	m I2CMaster
}

func NewTransact8x8(m I2CMaster) Transactor8x8 {
	if t, ok := m.(Transactor8x8); ok {
		return t
	}
	return transactor8x8{m}
}

func (t transactor8x8) Transact8x8(addr Addr, regaddr uint8, w []byte, r []byte) (int, int, error) {
	return I2CMasterTransact8x8(t.m, addr, regaddr, w, r)
}

func I2CMasterTransact8x8(m I2CMaster, addr Addr, regaddr uint8, w []byte, r []byte) (int, int, error) {
	nr := 0
	nw := 0

	if addr.GetAddrLen() != 7 {
		return nw, nr, errors.New("I2CMasterTransact8x8: only 7 bit addresses are supported")
	}

	if err := m.Start(); err != nil {
		return nw, nr, err
	}

	// inner function handles the whole transaction between
	// but not including the start and the stop bit
	err := func() error {
		// address device
		addrb := uint8(addr.GetBaseAddr() << 1)
		if err := m.WriteByte(addrb); err != nil {
			if err == NACKReceived {
				return NoSuchDevice
			}
			return err
		}

		// write regaddr
		if err := m.WriteByte(regaddr); err != nil {
			return err
		}

		// write w
		for _, b := range w {
			if err := m.WriteByte(b); err != nil {
				return err
			}

			nw++
		}

		// read part of transaction is only performed if desired
		if len(r) > 0 {
			// start again
			if err := m.Start(); err != nil {
				return err
			}

			// write device's read address
			if err := m.WriteByte(addrb | 0x01); err != nil {
				if err == NACKReceived {
					return NoSuchDevice
				}
				return err
			}

			for i := 0; i < len(r); i++ {
				ack := true
				if i == len(r)-1 {
					ack = false
				}
				rb, err := m.ReadByte(ack)
				if err != nil {
					return err
				}

				r[i] = rb

				nr++
			}

		}

		return nil
	}()

	if err != nil {
		// if there already was an error, the error from stop is ignored
		// and the first error is reported
		m.Stop()
	} else {
		err = m.Stop()
	}

	return nw, nr, err
}
<|endoftext|>


<commit_before>package system

import (
	"fmt"

	"github.com/influxdb/tivan/plugins"
	"github.com/influxdb/tivan/plugins/system/ps/common"
	"github.com/influxdb/tivan/plugins/system/ps/cpu"
	"github.com/influxdb/tivan/plugins/system/ps/disk"
	"github.com/influxdb/tivan/plugins/system/ps/docker"
	"github.com/influxdb/tivan/plugins/system/ps/load"
	"github.com/influxdb/tivan/plugins/system/ps/mem"
	"github.com/influxdb/tivan/plugins/system/ps/net"
)

type DockerContainerStat struct {
	Name string
	CPU  *cpu.CPUTimesStat
	Mem  *docker.CgroupMemStat
}

type PS interface {
	LoadAvg() (*load.LoadAvgStat, error)
	CPUTimes() ([]cpu.CPUTimesStat, error)
	DiskUsage() ([]*disk.DiskUsageStat, error)
	NetIO() ([]net.NetIOCountersStat, error)
	DiskIO() (map[string]disk.DiskIOCountersStat, error)
	VMStat() (*mem.VirtualMemoryStat, error)
	SwapStat() (*mem.SwapMemoryStat, error)
	DockerStat() ([]*DockerContainerStat, error)
}

type SystemStats struct {
	ps PS
}

func (s *SystemStats) add(acc plugins.Accumulator,
	name string, val float64, tags map[string]string) {
	if val >= 0 {
		acc.Add(name, val, tags)
	}
}

func (s *SystemStats) Gather(acc plugins.Accumulator) error {
	lv, err := s.ps.LoadAvg()
	if err != nil {
		return err
	}

	acc.Add("load1", lv.Load1, nil)
	acc.Add("load5", lv.Load5, nil)
	acc.Add("load15", lv.Load15, nil)

	times, err := s.ps.CPUTimes()
	if err != nil {
		return fmt.Errorf("error getting CPU info: %s", err)
	}

	for _, cts := range times {
		tags := map[string]string{
			"cpu": cts.CPU,
		}

		s.add(acc, "user", cts.User, tags)
		s.add(acc, "system", cts.System, tags)
		s.add(acc, "idle", cts.Idle, tags)
		s.add(acc, "nice", cts.Nice, tags)
		s.add(acc, "iowait", cts.Iowait, tags)
		s.add(acc, "irq", cts.Irq, tags)
		s.add(acc, "softirq", cts.Softirq, tags)
		s.add(acc, "steal", cts.Steal, tags)
		s.add(acc, "guest", cts.Guest, tags)
		s.add(acc, "guestNice", cts.GuestNice, tags)
		s.add(acc, "stolen", cts.Stolen, tags)
	}

	disks, err := s.ps.DiskUsage()
	if err != nil {
		return err
	}

	for _, du := range disks {
		tags := map[string]string{
			"path": du.Path,
		}

		acc.Add("total", du.Total, tags)
		acc.Add("free", du.Free, tags)
		acc.Add("used", du.Total-du.Free, tags)
		acc.Add("inodes_total", du.InodesTotal, tags)
		acc.Add("inodes_free", du.InodesFree, tags)
		acc.Add("inodes_used", du.InodesTotal-du.InodesFree, tags)
	}

	diskio, err := s.ps.DiskIO()
	if err != nil {
		return err
	}

	for _, io := range diskio {
		tags := map[string]string{
			"name":   io.Name,
			"serial": io.SerialNumber,
		}

		acc.Add("reads", io.ReadCount, tags)
		acc.Add("writes", io.WriteCount, tags)
		acc.Add("read_bytes", io.ReadBytes, tags)
		acc.Add("write_bytes", io.WriteBytes, tags)
		acc.Add("read_time", io.ReadTime, tags)
		acc.Add("write_time", io.WriteTime, tags)
		acc.Add("io_time", io.IoTime, tags)
	}

	netio, err := s.ps.NetIO()
	if err != nil {
		return err
	}

	for _, io := range netio {
		tags := map[string]string{
			"interface": io.Name,
		}

		acc.Add("bytes_sent", io.BytesSent, tags)
		acc.Add("bytes_recv", io.BytesRecv, tags)
		acc.Add("packets_sent", io.PacketsSent, tags)
		acc.Add("packets_recv", io.PacketsRecv, tags)
		acc.Add("err_in", io.Errin, tags)
		acc.Add("err_out", io.Errout, tags)
		acc.Add("drop_in", io.Dropin, tags)
		acc.Add("drop_out", io.Dropout, tags)
	}

	vm, err := s.ps.VMStat()
	if err != nil {
		return err
	}

	vmtags := map[string]string{
		"memory": "virtual",
	}

	acc.Add("total", vm.Total, vmtags)
	acc.Add("available", vm.Available, vmtags)
	acc.Add("used", vm.Used, vmtags)
	acc.Add("used_prec", vm.UsedPercent, vmtags)
	acc.Add("free", vm.Free, vmtags)
	acc.Add("active", vm.Active, vmtags)
	acc.Add("inactive", vm.Inactive, vmtags)
	acc.Add("buffers", vm.Buffers, vmtags)
	acc.Add("cached", vm.Cached, vmtags)
	acc.Add("wired", vm.Wired, vmtags)
	acc.Add("shared", vm.Shared, vmtags)

	swap, err := s.ps.SwapStat()
	if err != nil {
		return err
	}

	swaptags := map[string]string{
		"memory": "swap",
	}

	acc.Add("total", swap.Total, swaptags)
	acc.Add("used", swap.Used, swaptags)
	acc.Add("free", swap.Free, swaptags)
	acc.Add("used_perc", swap.UsedPercent, swaptags)
	acc.Add("swap_in", swap.Sin, swaptags)
	acc.Add("swap_out", swap.Sout, swaptags)

	containers, err := s.ps.DockerStat()
	if err != nil {
		return err
	}

	for _, cont := range containers {
		tags := map[string]string{
			"docker": cont.Name,
		}

		cts := cont.CPU

		acc.Add("user", cts.User, tags)
		acc.Add("system", cts.System, tags)
		acc.Add("idle", cts.Idle, tags)
		acc.Add("nice", cts.Nice, tags)
		acc.Add("iowait", cts.Iowait, tags)
		acc.Add("irq", cts.Irq, tags)
		acc.Add("softirq", cts.Softirq, tags)
		acc.Add("steal", cts.Steal, tags)
		acc.Add("guest", cts.Guest, tags)
		acc.Add("guestNice", cts.GuestNice, tags)
		acc.Add("stolen", cts.Stolen, tags)

		acc.Add("cache", cont.Mem.Cache, tags)
		acc.Add("rss", cont.Mem.RSS, tags)
		acc.Add("rss_huge", cont.Mem.RSSHuge, tags)
		acc.Add("mapped_file", cont.Mem.MappedFile, tags)
		acc.Add("swap_in", cont.Mem.Pgpgin, tags)
		acc.Add("swap_out", cont.Mem.Pgpgout, tags)
		acc.Add("page_fault", cont.Mem.Pgfault, tags)
		acc.Add("page_major_fault", cont.Mem.Pgmajfault, tags)
		acc.Add("inactive_anon", cont.Mem.InactiveAnon, tags)
		acc.Add("active_anon", cont.Mem.ActiveAnon, tags)
		acc.Add("inactive_file", cont.Mem.InactiveFile, tags)
		acc.Add("active_file", cont.Mem.ActiveFile, tags)
		acc.Add("unevictable", cont.Mem.Unevictable, tags)
		acc.Add("memory_limit", cont.Mem.HierarchicalMemoryLimit, tags)
		acc.Add("total_cache", cont.Mem.TotalCache, tags)
		acc.Add("total_rss", cont.Mem.TotalRSS, tags)
		acc.Add("total_rss_huge", cont.Mem.TotalRSSHuge, tags)
		acc.Add("total_mapped_file", cont.Mem.TotalMappedFile, tags)
		acc.Add("total_swap_in", cont.Mem.TotalPgpgIn, tags)
		acc.Add("total_swap_out", cont.Mem.TotalPgpgOut, tags)
		acc.Add("total_page_fault", cont.Mem.TotalPgFault, tags)
		acc.Add("total_page_major_fault", cont.Mem.TotalPgMajFault, tags)
		acc.Add("total_inactive_anon", cont.Mem.TotalInactiveAnon, tags)
		acc.Add("total_active_anon", cont.Mem.TotalActiveAnon, tags)
		acc.Add("total_inactive_file", cont.Mem.TotalInactiveFile, tags)
		acc.Add("total_active_file", cont.Mem.TotalActiveFile, tags)
		acc.Add("total_unevictable", cont.Mem.TotalUnevictable, tags)
	}

	return nil
}

type systemPS struct{}

func (s *systemPS) LoadAvg() (*load.LoadAvgStat, error) {
	return load.LoadAvg()
}

func (s *systemPS) CPUTimes() ([]cpu.CPUTimesStat, error) {
	return cpu.CPUTimes(true)
}

func (s *systemPS) DiskUsage() ([]*disk.DiskUsageStat, error) {
	parts, err := disk.DiskPartitions(true)
	if err != nil {
		return nil, err
	}

	var usage []*disk.DiskUsageStat

	for _, p := range parts {
		du, err := disk.DiskUsage(p.Mountpoint)
		if err != nil {
			return nil, err
		}

		usage = append(usage, du)
	}

	return usage, nil
}

func (s *systemPS) NetIO() ([]net.NetIOCountersStat, error) {
	return net.NetIOCounters(true)
}

func (s *systemPS) DiskIO() (map[string]disk.DiskIOCountersStat, error) {
	m, err := disk.DiskIOCounters()
	if err == common.NotImplementedError {
		return nil, nil
	}

	return m, err
}

func (s *systemPS) VMStat() (*mem.VirtualMemoryStat, error) {
	return mem.VirtualMemory()
}

func (s *systemPS) SwapStat() (*mem.SwapMemoryStat, error) {
	return mem.SwapMemory()
}

func (s *systemPS) DockerStat() ([]*DockerContainerStat, error) {
	list, err := docker.GetDockerIDList()
	if err != nil {
		if err == common.NotImplementedError {
			return nil, nil
		}

		return nil, err
	}

	var stats []*DockerContainerStat

	for _, cont := range list {
		ctu, err := docker.CgroupCPUDocker(cont)
		if err != nil {
			return nil, err
		}

		mem, err := docker.CgroupMemDocker(cont)
		if err != nil {
			return nil, err
		}

		stats = append(stats, &DockerContainerStat{cont, ctu, mem})
	}

	return stats, nil
}

func init() {
	plugins.Add("system", func() plugins.Plugin {
		return &SystemStats{ps: &systemPS{}}
	})
}
<commit_msg>Report better errors where system stats can't be gathered<commit_after>package system

import (
	"fmt"

	"github.com/influxdb/tivan/plugins"
	"github.com/influxdb/tivan/plugins/system/ps/common"
	"github.com/influxdb/tivan/plugins/system/ps/cpu"
	"github.com/influxdb/tivan/plugins/system/ps/disk"
	"github.com/influxdb/tivan/plugins/system/ps/docker"
	"github.com/influxdb/tivan/plugins/system/ps/load"
	"github.com/influxdb/tivan/plugins/system/ps/mem"
	"github.com/influxdb/tivan/plugins/system/ps/net"
)

type DockerContainerStat struct {
	Name string
	CPU  *cpu.CPUTimesStat
	Mem  *docker.CgroupMemStat
}

type PS interface {
	LoadAvg() (*load.LoadAvgStat, error)
	CPUTimes() ([]cpu.CPUTimesStat, error)
	DiskUsage() ([]*disk.DiskUsageStat, error)
	NetIO() ([]net.NetIOCountersStat, error)
	DiskIO() (map[string]disk.DiskIOCountersStat, error)
	VMStat() (*mem.VirtualMemoryStat, error)
	SwapStat() (*mem.SwapMemoryStat, error)
	DockerStat() ([]*DockerContainerStat, error)
}

type SystemStats struct {
	ps PS
}

func (s *SystemStats) add(acc plugins.Accumulator,
	name string, val float64, tags map[string]string) {
	if val >= 0 {
		acc.Add(name, val, tags)
	}
}

func (s *SystemStats) Gather(acc plugins.Accumulator) error {
	lv, err := s.ps.LoadAvg()
	if err != nil {
		return err
	}

	acc.Add("load1", lv.Load1, nil)
	acc.Add("load5", lv.Load5, nil)
	acc.Add("load15", lv.Load15, nil)

	times, err := s.ps.CPUTimes()
	if err != nil {
		return fmt.Errorf("error getting CPU info: %s", err)
	}

	for _, cts := range times {
		tags := map[string]string{
			"cpu": cts.CPU,
		}

		s.add(acc, "user", cts.User, tags)
		s.add(acc, "system", cts.System, tags)
		s.add(acc, "idle", cts.Idle, tags)
		s.add(acc, "nice", cts.Nice, tags)
		s.add(acc, "iowait", cts.Iowait, tags)
		s.add(acc, "irq", cts.Irq, tags)
		s.add(acc, "softirq", cts.Softirq, tags)
		s.add(acc, "steal", cts.Steal, tags)
		s.add(acc, "guest", cts.Guest, tags)
		s.add(acc, "guestNice", cts.GuestNice, tags)
		s.add(acc, "stolen", cts.Stolen, tags)
	}

	disks, err := s.ps.DiskUsage()
	if err != nil {
		return fmt.Errorf("error getting disk usage info: %s", err)
	}

	for _, du := range disks {
		tags := map[string]string{
			"path": du.Path,
		}

		acc.Add("total", du.Total, tags)
		acc.Add("free", du.Free, tags)
		acc.Add("used", du.Total-du.Free, tags)
		acc.Add("inodes_total", du.InodesTotal, tags)
		acc.Add("inodes_free", du.InodesFree, tags)
		acc.Add("inodes_used", du.InodesTotal-du.InodesFree, tags)
	}

	diskio, err := s.ps.DiskIO()
	if err != nil {
		return fmt.Errorf("error getting disk io info: %s", err)
	}

	for _, io := range diskio {
		tags := map[string]string{
			"name":   io.Name,
			"serial": io.SerialNumber,
		}

		acc.Add("reads", io.ReadCount, tags)
		acc.Add("writes", io.WriteCount, tags)
		acc.Add("read_bytes", io.ReadBytes, tags)
		acc.Add("write_bytes", io.WriteBytes, tags)
		acc.Add("read_time", io.ReadTime, tags)
		acc.Add("write_time", io.WriteTime, tags)
		acc.Add("io_time", io.IoTime, tags)
	}

	netio, err := s.ps.NetIO()
	if err != nil {
		return fmt.Errorf("error getting net io info: %s", err)
	}

	for _, io := range netio {
		tags := map[string]string{
			"interface": io.Name,
		}

		acc.Add("bytes_sent", io.BytesSent, tags)
		acc.Add("bytes_recv", io.BytesRecv, tags)
		acc.Add("packets_sent", io.PacketsSent, tags)
		acc.Add("packets_recv", io.PacketsRecv, tags)
		acc.Add("err_in", io.Errin, tags)
		acc.Add("err_out", io.Errout, tags)
		acc.Add("drop_in", io.Dropin, tags)
		acc.Add("drop_out", io.Dropout, tags)
	}

	vm, err := s.ps.VMStat()
	if err != nil {
		return fmt.Errorf("error getting virtual memory info: %s", err)
	}

	vmtags := map[string]string{
		"memory": "virtual",
	}

	acc.Add("total", vm.Total, vmtags)
	acc.Add("available", vm.Available, vmtags)
	acc.Add("used", vm.Used, vmtags)
	acc.Add("used_prec", vm.UsedPercent, vmtags)
	acc.Add("free", vm.Free, vmtags)
	acc.Add("active", vm.Active, vmtags)
	acc.Add("inactive", vm.Inactive, vmtags)
	acc.Add("buffers", vm.Buffers, vmtags)
	acc.Add("cached", vm.Cached, vmtags)
	acc.Add("wired", vm.Wired, vmtags)
	acc.Add("shared", vm.Shared, vmtags)

	swap, err := s.ps.SwapStat()
	if err != nil {
		return fmt.Errorf("error getting swap memory info: %s", err)
	}

	swaptags := map[string]string{
		"memory": "swap",
	}

	acc.Add("total", swap.Total, swaptags)
	acc.Add("used", swap.Used, swaptags)
	acc.Add("free", swap.Free, swaptags)
	acc.Add("used_perc", swap.UsedPercent, swaptags)
	acc.Add("swap_in", swap.Sin, swaptags)
	acc.Add("swap_out", swap.Sout, swaptags)

	containers, err := s.ps.DockerStat()
	if err != nil {
		return err
	}

	for _, cont := range containers {
		tags := map[string]string{
			"docker": cont.Name,
		}

		cts := cont.CPU

		acc.Add("user", cts.User, tags)
		acc.Add("system", cts.System, tags)
		acc.Add("idle", cts.Idle, tags)
		acc.Add("nice", cts.Nice, tags)
		acc.Add("iowait", cts.Iowait, tags)
		acc.Add("irq", cts.Irq, tags)
		acc.Add("softirq", cts.Softirq, tags)
		acc.Add("steal", cts.Steal, tags)
		acc.Add("guest", cts.Guest, tags)
		acc.Add("guestNice", cts.GuestNice, tags)
		acc.Add("stolen", cts.Stolen, tags)

		acc.Add("cache", cont.Mem.Cache, tags)
		acc.Add("rss", cont.Mem.RSS, tags)
		acc.Add("rss_huge", cont.Mem.RSSHuge, tags)
		acc.Add("mapped_file", cont.Mem.MappedFile, tags)
		acc.Add("swap_in", cont.Mem.Pgpgin, tags)
		acc.Add("swap_out", cont.Mem.Pgpgout, tags)
		acc.Add("page_fault", cont.Mem.Pgfault, tags)
		acc.Add("page_major_fault", cont.Mem.Pgmajfault, tags)
		acc.Add("inactive_anon", cont.Mem.InactiveAnon, tags)
		acc.Add("active_anon", cont.Mem.ActiveAnon, tags)
		acc.Add("inactive_file", cont.Mem.InactiveFile, tags)
		acc.Add("active_file", cont.Mem.ActiveFile, tags)
		acc.Add("unevictable", cont.Mem.Unevictable, tags)
		acc.Add("memory_limit", cont.Mem.HierarchicalMemoryLimit, tags)
		acc.Add("total_cache", cont.Mem.TotalCache, tags)
		acc.Add("total_rss", cont.Mem.TotalRSS, tags)
		acc.Add("total_rss_huge", cont.Mem.TotalRSSHuge, tags)
		acc.Add("total_mapped_file", cont.Mem.TotalMappedFile, tags)
		acc.Add("total_swap_in", cont.Mem.TotalPgpgIn, tags)
		acc.Add("total_swap_out", cont.Mem.TotalPgpgOut, tags)
		acc.Add("total_page_fault", cont.Mem.TotalPgFault, tags)
		acc.Add("total_page_major_fault", cont.Mem.TotalPgMajFault, tags)
		acc.Add("total_inactive_anon", cont.Mem.TotalInactiveAnon, tags)
		acc.Add("total_active_anon", cont.Mem.TotalActiveAnon, tags)
		acc.Add("total_inactive_file", cont.Mem.TotalInactiveFile, tags)
		acc.Add("total_active_file", cont.Mem.TotalActiveFile, tags)
		acc.Add("total_unevictable", cont.Mem.TotalUnevictable, tags)
	}

	return nil
}

type systemPS struct{}

func (s *systemPS) LoadAvg() (*load.LoadAvgStat, error) {
	return load.LoadAvg()
}

func (s *systemPS) CPUTimes() ([]cpu.CPUTimesStat, error) {
	return cpu.CPUTimes(true)
}

func (s *systemPS) DiskUsage() ([]*disk.DiskUsageStat, error) {
	parts, err := disk.DiskPartitions(true)
	if err != nil {
		return nil, err
	}

	var usage []*disk.DiskUsageStat

	for _, p := range parts {
		du, err := disk.DiskUsage(p.Mountpoint)
		if err != nil {
			return nil, err
		}

		usage = append(usage, du)
	}

	return usage, nil
}

func (s *systemPS) NetIO() ([]net.NetIOCountersStat, error) {
	return net.NetIOCounters(true)
}

func (s *systemPS) DiskIO() (map[string]disk.DiskIOCountersStat, error) {
	m, err := disk.DiskIOCounters()
	if err == common.NotImplementedError {
		return nil, nil
	}

	return m, err
}

func (s *systemPS) VMStat() (*mem.VirtualMemoryStat, error) {
	return mem.VirtualMemory()
}

func (s *systemPS) SwapStat() (*mem.SwapMemoryStat, error) {
	return mem.SwapMemory()
}

func (s *systemPS) DockerStat() ([]*DockerContainerStat, error) {
	list, err := docker.GetDockerIDList()
	if err != nil {
		if err == common.NotImplementedError {
			return nil, nil
		}

		return nil, err
	}

	var stats []*DockerContainerStat

	for _, cont := range list {
		ctu, err := docker.CgroupCPUDocker(cont)
		if err != nil {
			return nil, err
		}

		mem, err := docker.CgroupMemDocker(cont)
		if err != nil {
			return nil, err
		}

		stats = append(stats, &DockerContainerStat{cont, ctu, mem})
	}

	return stats, nil
}

func init() {
	plugins.Add("system", func() plugins.Plugin {
		return &SystemStats{ps: &systemPS{}}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
)

type Client struct {
	conn     net.Conn
	incoming chan string // Channel for incoming data from client
}

type Server struct {
	clients []*Client
	joins   chan net.Conn // Channel for new connections
}

func (client *Client) listen() {
	reader := bufio.NewReader(client.conn)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Print("Closing connection")
			// TODO: Remove client from list
			client.conn.Close()
			return
		}
		log.Print(line)
	}
}

// Add new connection to server
// TODO: Write authentication by key
func (server *Server) addConnection(conn net.Conn) {
	log.Print("Adding new connection")
	client := &Client{
		conn: conn,
	}
	go client.listen()
	server.clients = append(server.clients, client)
}

// Listen channels
func (server *Server) listen() {
	for {
		select {
		case conn := <-server.joins:
			server.addConnection(conn)
		}
	}
}

func NewServer() *Server {
	server := &Server{
		joins: make(chan net.Conn),
	}
	go server.listen()
	return server
}

func main() {
	fmt.Print("Initializing server")
	listener, err := net.Listen("tcp", ":9999")
	if err != nil {
		log.Fatal("Error starting TCP server.")
	}

	server := NewServer()

	for {
		conn, _ := listener.Accept()
		server.joins <- conn
	}
}
<commit_msg>Added todo<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
)

type Client struct {
	conn     net.Conn
	incoming chan string // Channel for incoming data from client
}

type Server struct {
	clients []*Client
	joins   chan net.Conn // Channel for new connections
}

func (client *Client) listen() {
	reader := bufio.NewReader(client.conn)
	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			log.Print("Closing connection")
			// TODO: Remove client from list
			client.conn.Close()
			return
		}
		// Send data to elastic
		log.Print(line)
	}
}

// Add new connection to server
// TODO: Write authentication by key
func (server *Server) addConnection(conn net.Conn) {
	log.Print("Adding new connection")
	client := &Client{
		conn: conn,
	}
	go client.listen()
	server.clients = append(server.clients, client)
}

// Listen channels
func (server *Server) listen() {
	for {
		select {
		case conn := <-server.joins:
			server.addConnection(conn)
		}
	}
}

func NewServer() *Server {
	server := &Server{
		joins: make(chan net.Conn),
	}
	go server.listen()
	return server
}

func main() {
	fmt.Print("Initializing server")
	listener, err := net.Listen("tcp", ":9999")
	if err != nil {
		log.Fatal("Error starting TCP server.")
	}

	server := NewServer()

	for {
		conn, _ := listener.Accept()
		server.joins <- conn
	}
}
<|endoftext|>


<commit_before>// This example demonstrates how to authenticate with Spotify.
// In order to run this example yourself, you'll need to:
//
//  1. Register an application at: https://developer.spotify.com/my-applications/
//       - Use "http://localhost:8080/callback" as the redirect URI
//  2. Set the SPOTIFY_ID environment variable to the client ID you got in step 1.
//  3. Set the SPOTIFY_SECRET environment variable to the client secret from step 1.
package main

import (
	"container/heap"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/pressly/chi"
	"github.com/zmb3/spotify"
	"github.com/coolbry95/partydj/backend/pool"
	"strconv"
)

// redirectURI is the OAuth redirect URI for the application.
// You must register an application at Spotify's developer portal
// and enter this value.
const redirectURI = "https://linode.shellcode.in/callback"

var (
	auth  = spotify.NewAuthenticator(redirectURI, spotify.ScopeUserLibraryModify, spotify.ScopePlaylistModifyPrivate,
	spotify.ScopePlaylistModifyPublic)
	ch    = make(chan spotify.Client)
	state = "stateless"
)

type DI struct {
	client spotify.Client
	pool   pool.Pool
}

var PoolShortIDToLongID map[int]string
var UserIDToPoolID map[string]string

func main() {
	// Initialize the pool (single for demonstration)
	var d *DI
	d = new(DI)

	// setup the router and paths
	r := chi.NewRouter()
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		log.Println(r.URL.String())
	})
	r.Get("/callback", completeAuth)
	r.Get("/getpool", d.getPool)
	r.Post("/createpool", d.createPool)
	r.Post("/add_song/:poolID/:songID", d.addSong)
	r.Post("/upvote/:poolID/:songID", d.upVote)
	r.Post("/downvote/:poolID/:songID", d.downVote)
	r.Post("/join_pool", d.joinPool)

	// Authenticate the users spotify account
	url := auth.AuthURL(state)
	fmt.Println("Please log in to Spotify by visiting the following page in your browser:", url)

	go http.ListenAndServe(":6060", r)

	// wait for auth to complete
	d.client = <-ch
	d.pool.PlaylistID = spotify.ID("0nXlYUH7zBAzubO9Yub4rR") // change to spartahack playlist

	// Initialize the short id to long id map
	PoolShortIDToLongID = make(map[int]string)
	// hard code the one join digit for this sample pool to 121
	PoolShortIDToLongID[121] = "0nXlYUH7zBAzubO9Yub4rR"

	// initialize the user map
	UserIDToPoolID = make(map[string]string)

	userID, err := d.client.CurrentUser()
	if err != nil {
		log.Println(err)
	}

	d.pool.UserID = userID.ID

	block := make(chan struct{})
	<-block
}

func (d *DI) joinPool(w http.ResponseWriter, r *http.Request) {
	userIDNumber := r.PostFormValue("userId")
	poolIDNumberStr := r.PostFormValue("poolShortId")

	// Both must be present
	if len(userIDNumber) == 0 || len(poolIDNumberStr) == 0 {
		w.WriteHeader(http.StatusPartialContent)
		fmt.Println("One of the parameters were empty: ", "userId:",
			userIDNumber, "poolShortId:", poolIDNumberStr)
		return
	}

	poolIDNumber, convErr := strconv.Atoi(poolIDNumberStr)

	// Make sure the poolIDNumber is a valid integer, otherwise
	// respond with Status partial content
	if convErr != nil {
		w.WriteHeader(http.StatusPartialContent)
		fmt.Println("Conversion error: ", convErr)
		return
	}

	// Check if pool exists, if not respond with status not found
	if _, ok := PoolShortIDToLongID[poolIDNumber]; ok != true {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	UserIDToPoolID[userIDNumber] = poolIDNumberStr
	// Let the user know the request was accepted
	w.WriteHeader(http.StatusAccepted)
}

func (d *DI) addSong(w http.ResponseWriter, r *http.Request) {
	songID := chi.URLParam(r, "songID")
	//poolID := chi.URLParam(r, "poolID")

	s := &pool.Song{ID: spotify.ID(songID)}
	heap.Push(&d.pool, s)
}

func (d *DI) upVote(w http.ResponseWriter, r *http.Request) {
	// TODO check for user ID to see if they already voted
	songID := chi.URLParam(r, "songID")

	d.pool.UpVote(spotify.ID(songID))
	d.pool.UpdateSpotifyPlaylist(&d.client, d.pool.PlaylistID)
}

func (d *DI) downVote(w http.ResponseWriter, r *http.Request) {
	// TODO check for user ID to see if they already voted
	songID := chi.URLParam(r, "songID")
	//poolID := chi.URLParam(r, "poolID")

	d.pool.DownVote(spotify.ID(songID))
	d.pool.UpdateSpotifyPlaylist(&d.client, d.pool.PlaylistID)
}

func (d *DI) createPool(w http.ResponseWriter, r *http.Request) {
}

func (d *DI) getPool(w http.ResponseWriter, r *http.Request) {
	userid, err := d.client.CurrentUser()
	if err != nil {
		log.Println(err)
	}

	// TODO: instead of using existing playlist we require a new playlist to be created
	// playlist, err := d.client.CreatePlaylistForUser(userid.ID, playlistName, true)
	playlist, err := d.client.GetPlaylistTracks(userid.ID, "0nXlYUH7zBAzubO9Yub4rR")
	//spartahackPlaylistName := "Sparthack Sample Playlist!"
	//newPlaylist, err := d.client.CreatePlaylistForUser(userid.ID, spartahackPlaylistName, true)
	if err != nil {
		log.Println(err)
	}

	d.pool.SongHeap = make([]*pool.Song, 0, 2)
	for i, track := range playlist.Tracks {
		s := &pool.Song{
			ID:       track.Track.ID,
			Name:     track.Track.Name,
			Duration: track.Track.Duration,
			Album:    track.Track.Album.Name,
			Images:   track.Track.Album.Images,
			Artists:  track.Track.Artists,
			Priority: i,
		}
		d.pool.SongHeap = append(d.pool.SongHeap, s)
	}

	json.NewEncoder(w).Encode(d.pool)
	return
}

func completeAuth(w http.ResponseWriter, r *http.Request) {
	tok, err := auth.Token(state, r)
	if err != nil {
		http.Error(w, "Couldn't get token", http.StatusForbidden)
		log.Fatal(err)
	}

	if st := r.FormValue("state"); st != state {
		http.NotFound(w, r)
		log.Fatalf("State mismatch: %s != %s\n", st, state)
	}

	// use the token to get an authenticated client
	client := auth.NewClient(tok)
	ch <- client
	http.Redirect(w, r, "/getpool", 301)
}
<commit_msg>Joining a pool is now possible using the unique short id of the pool.<commit_after>// This example demonstrates how to authenticate with Spotify.
// In order to run this example yourself, you'll need to:
//
//  1. Register an application at: https://developer.spotify.com/my-applications/
//       - Use "http://localhost:8080/callback" as the redirect URI
//  2. Set the SPOTIFY_ID environment variable to the client ID you got in step 1.
//  3. Set the SPOTIFY_SECRET environment variable to the client secret from step 1.
package main

import (
	"container/heap"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/pressly/chi"
	"github.com/zmb3/spotify"
	"github.com/coolbry95/partydj/backend/pool"
	"strconv"
)

// redirectURI is the OAuth redirect URI for the application.
// You must register an application at Spotify's developer portal
// and enter this value.
const redirectURI = "https://linode.shellcode.in/callback"

var (
	auth = spotify.NewAuthenticator(redirectURI, spotify.ScopeUserLibraryModify, spotify.ScopePlaylistModifyPrivate,
		spotify.ScopePlaylistModifyPublic)
	ch    = make(chan spotify.Client)
	state = "stateless"
)

type DI struct {
	client spotify.Client
	pool   pool.Pool
}

var PoolShortIDToLongID map[int]string
var UserIDToPoolID map[string]string

func main() {
	// Initialize the pool (single for demonstration)
	var d *DI
	d = new(DI)

	// setup the router and paths
	r := chi.NewRouter()
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		log.Println(r.URL.String())
	})
	r.Get("/callback", completeAuth)
	r.Get("/getpool", d.getPool)
	r.Post("/createpool", d.createPool)
	r.Post("/add_song/:poolID/:songID", d.addSong)
	r.Post("/upvote/:poolID/:songID", d.upVote)
	r.Post("/downvote/:poolID/:songID", d.downVote)
	r.Post("/join_pool", d.joinPool)

	// Authenticate the users spotify account
	url := auth.AuthURL(state)
	fmt.Println("Please log in to Spotify by visiting the following page in your browser:", url)

	go http.ListenAndServe(":6060", r)

	// wait for auth to complete
	d.client = <-ch
	d.pool.PlaylistID = spotify.ID("0nXlYUH7zBAzubO9Yub4rR") // change to spartahack playlist

	// Initialize the short id to long id map
	PoolShortIDToLongID = make(map[int]string)
	// hard code the one join digit for this sample pool to 121
	PoolShortIDToLongID[121] = "0nXlYUH7zBAzubO9Yub4rR"

	// initialize the user map
	UserIDToPoolID = make(map[string]string)

	userID, err := d.client.CurrentUser()
	if err != nil {
		log.Println(err)
	}

	d.pool.UserID = userID.ID

	block := make(chan struct{})
	<-block
}

func (d *DI) joinPool(w http.ResponseWriter, r *http.Request) {
	userIDNumber := r.PostFormValue("userId")
	poolIDNumberStr := r.PostFormValue("poolShortId")

	// Both must be present
	if len(userIDNumber) == 0 || len(poolIDNumberStr) == 0 {
		w.WriteHeader(http.StatusPartialContent)
		fmt.Println("One of the parameters were empty: ", "userId:",
			userIDNumber, "poolShortId:", poolIDNumberStr)
		return
	}

	poolIDNumber, convErr := strconv.Atoi(poolIDNumberStr)

	// Make sure the poolIDNumber is a valid integer, otherwise
	// respond with Status partial content
	if convErr != nil {
		w.WriteHeader(http.StatusPartialContent)
		fmt.Println("Conversion error: ", convErr)
		return
	}

	// Check if pool exists, if not respond with status not found
	if _, ok := PoolShortIDToLongID[poolIDNumber]; ok != true {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	UserIDToPoolID[userIDNumber] = poolIDNumberStr
	// Let the user know the request was accepted
	w.WriteHeader(http.StatusAccepted)
}

func (d *DI) addSong(w http.ResponseWriter, r *http.Request) {
	songID := chi.URLParam(r, "songID")
	//poolID := chi.URLParam(r, "poolID")

	s := &pool.Song{ID: spotify.ID(songID)}
	heap.Push(&d.pool, s)
}

func (d *DI) upVote(w http.ResponseWriter, r *http.Request) {
	// TODO check for user ID to see if they already voted
	songID := chi.URLParam(r, "songID")

	d.pool.UpVote(spotify.ID(songID))
	d.pool.UpdateSpotifyPlaylist(&d.client, d.pool.PlaylistID)
}

func (d *DI) downVote(w http.ResponseWriter, r *http.Request) {
	// TODO check for user ID to see if they already voted
	songID := chi.URLParam(r, "songID")
	//poolID := chi.URLParam(r, "poolID")

	d.pool.DownVote(spotify.ID(songID))
	d.pool.UpdateSpotifyPlaylist(&d.client, d.pool.PlaylistID)
}

func (d *DI) createPool(w http.ResponseWriter, r *http.Request) {
}

func (d *DI) getPool(w http.ResponseWriter, r *http.Request) {
	userid, err := d.client.CurrentUser()
	if err != nil {
		log.Println(err)
	}

	// TODO: instead of using existing playlist we require a new playlist to be created
	// playlist, err := d.client.CreatePlaylistForUser(userid.ID, playlistName, true)
	playlist, err := d.client.GetPlaylistTracks(userid.ID, "0nXlYUH7zBAzubO9Yub4rR")
	//spartahackPlaylistName := "Sparthack Sample Playlist!"
	//newPlaylist, err := d.client.CreatePlaylistForUser(userid.ID, spartahackPlaylistName, true)
	if err != nil {
		log.Println(err)
	}

	d.pool.SongHeap = make([]*pool.Song, 0, 10)
	for i, track := range playlist.Tracks {
		d.pool.SongHeap = append(d.pool.SongHeap, pool.TrackToSong(&track.Track.SimpleTrack, i))
	}

	//TODO: only call this function only after the the current song finishes
	//d.pool.AddNextSong(&d.client)

	json.NewEncoder(w).Encode(d.pool)
	return
}

func completeAuth(w http.ResponseWriter, r *http.Request) {
	tok, err := auth.Token(state, r)
	if err != nil {
		http.Error(w, "Couldn't get token", http.StatusForbidden)
		log.Fatal(err)
	}

	if st := r.FormValue("state"); st != state {
		http.NotFound(w, r)
		log.Fatalf("State mismatch: %s != %s\n", st, state)
	}

	// use the token to get an authenticated client
	client := auth.NewClient(tok)
	ch <- client
	http.Redirect(w, r, "/getpool", 301)
}
<|endoftext|>


<commit_before><commit_msg>Delete debug msg<commit_after><|endoftext|>


<commit_before>package bflag

import (
	"errors"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

type bflag struct {
	bools       map[string]*BoolOption
	strings     map[string]*StringOption
	ints        map[string]*IntOption
	floats      map[string]*FloatOption
	args        []string
	invalidArgs []string
}

var bf *bflag

func init() {
	bf = &bflag{
		bools:       make(map[string]*BoolOption),
		strings:     make(map[string]*StringOption),
		ints:        make(map[string]*IntOption),
		floats:      make(map[string]*FloatOption),
		args:        make([]string, 0),
		invalidArgs: make([]string, 0),
	}
}

// BoolOption type
type BoolOption struct {
	Name      string
	ShortName byte
	Value     bool
	HelpText  string
}

// StringOption type
type StringOption struct {
	Name      string
	ShortName byte
	Value     string
	HelpText  string
}

// IntOption type
type IntOption struct {
	Name      string
	ShortName byte
	Value     int
	HelpText  string
}

// FloatOption type
type FloatOption struct {
	Name      string
	ShortName byte
	Value     float64
	HelpText  string
}

// DefineBool func
func DefineBool(name string, shortName byte, defaultVal bool) *BoolOption {
	data := &BoolOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.bools[name] = data
	bf.bools[string(shortName)] = data
	return data
}

// DefineString func
func DefineString(name string, shortName byte, defaultVal string) *StringOption {
	data := &StringOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.strings[name] = data
	bf.strings[string(shortName)] = data
	return data
}

// DefineInt func
func DefineInt(name string, shortName byte, defaultVal int) *IntOption {
	data := &IntOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.ints[name] = data
	bf.ints[string(shortName)] = data
	return data
}

// DefineFloat func
func DefineFloat(name string, shortName byte, defaultVal float64) *FloatOption {
	data := &FloatOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.floats[name] = data
	bf.floats[string(shortName)] = data
	return data
}

// Args returns non-option arguments
func Args() []string {
	return bf.args
}

// Invalid returns invalid options
func Invalid() []string {
	return bf.invalidArgs
}

// GetBool returns the specified *BoolOption if defined
func GetBool(name string) *BoolOption {
	if b, ok := bf.bools[name]; ok {
		return b
	}
	return nil
}

// SetBool func
func SetBool(name string, value bool) error {
	if obj, ok := bf.bools[name]; ok {
		obj.Value = value
	} else {
		return errors.New("cannot set undefined bool: " + name)
	}
	return nil
}

// GetString returns the specified *StringOption if defined
func GetString(name string) *StringOption {
	if s, ok := bf.strings[name]; ok {
		return s
	}
	return nil
}

// SetString func
func SetString(name string, value string) error {
	if s, ok := bf.strings[name]; ok {
		s.Value = value
	} else {
		return errors.New("cannot set undefined string: " + name)
	}
	return nil
}

// GetInt returns the specified *IntOption if defined
func GetInt(name string) *IntOption {
	if i, ok := bf.ints[name]; ok {
		return i
	}
	return nil
}

// SetInt func
func SetInt(name string, value int) error {
	if i, ok := bf.ints[name]; ok {
		i.Value = value
	} else {
		return errors.New("cannot set undefined int: " + name)
	}
	return nil
}

// GetFloat returns the specified *FloatOption if defined
func GetFloat(name string) *FloatOption {
	if f, ok := bf.floats[name]; ok {
		return f
	}
	return nil
}

// SetFloat func
func SetFloat(name string, value float64) error {
	if f, ok := bf.floats[name]; ok {
		f.Value = value
	} else {
		return errors.New("cannot set undefined float: " + name)
	}
	return nil
}

// Parse will merge the command-line arguments into user-defined options
func Parse() {
	osArgs := os.Args[1:]
	removeSet := map[int]bool{}
	invalidArgs := []string{}

	// Expand short bool options
	args := make([]string, 0, len(osArgs))
	for _, arg := range osArgs {
		name, _ := parseOption(arg)
		if isValidCombinedFlag(arg) {
			nameArr := strings.Split(name, "")
			for _, v := range nameArr {
				args = append(args, "-"+v)
			}
		} else {
			args = append(args, arg)
		}
	}

	for i, arg := range args {
		// Check to see if there is a valid option set
		if isValidOption(arg) {
			name, val := parseOption(arg)

			// Parse boolean option
			if boolOpt := GetBool(name); boolOpt != nil {
				if isValidFlag(arg) {
					if val == "true" {
						SetBool(name, true)
					} else if val == "false" {
						SetBool(name, false)
					} else {
						SetBool(name, true)
					}
					removeSet[i] = true
					continue
				} else {
					invalidArgs = append(invalidArgs, "invalid option '"+arg+"'")
					removeSet[i] = true
					continue
				}
			} else {
				if isValidFlag(arg) {
					// bool flag passed but does not match any defined bool types
					invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")
					removeSet[i] = true
					continue
				}
			}

			// Check `option=value` vs. `option value` for non-boolean flags
			if strings.Index(arg, "=") == -1 {
				// must get the next command-line argument for value
				if i == len(args)-1 {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is missing a value")
					removeSet[i] = true
					continue
				}
				val = args[i+1]
				removeSet[i] = true
				removeSet[i+1] = true
			} else {
				removeSet[i] = true
			}

			// Parse string option
			if strOpt := GetString(name); strOpt != nil {
				SetString(name, val)
				continue
			}

			// Parse int option
			if intOpt := GetInt(name); intOpt != nil {
				if intVal, err := strconv.Atoi(val); err != nil {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is an invalid int ("+val+")")
				} else {
					SetInt(name, intVal)
				}
				continue
			}

			// Parse float option
			if floatOpt := GetFloat(name); floatOpt != nil {
				if floatVal, err := strconv.ParseFloat(val, 64); err != nil {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is an invalid float ("+val+")")
				} else {
					SetFloat(name, floatVal)
				}
				continue
			}

			// Doesn't match anything else; invalid option
			invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")

		} else {
			// Invalid option found, check to see if it begins with - or --
			isOption, _ := regexp.MatchString("\\A-{1,2}", arg)
			if isOption {
				invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")
				removeSet[i] = true
			}
		}
	}

	// Construct list of sorted indexes to remove based on removeSet
	var removeIndexes []int
	for i := range removeSet {
		removeIndexes = append(removeIndexes, i)
	}
	sort.Ints(removeIndexes)

	// Populate Args with all other non-option command-line arguments
	for i := len(removeIndexes) - 1; i >= 0; i-- {
		j := removeIndexes[i]
		if j == len(args)-1 {
			args = args[:j]
		} else {
			args = append(args[:j], args[j+1:]...)
		}
	}
	bf.args = args
	bf.invalidArgs = invalidArgs
}
<commit_msg>Fixed bug where a bool flag at the end could expect a value for the next arg<commit_after>package bflag

import (
	"errors"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

type bflag struct {
	bools       map[string]*BoolOption
	strings     map[string]*StringOption
	ints        map[string]*IntOption
	floats      map[string]*FloatOption
	args        []string
	invalidArgs []string
}

var bf *bflag

func init() {
	bf = &bflag{
		bools:       make(map[string]*BoolOption),
		strings:     make(map[string]*StringOption),
		ints:        make(map[string]*IntOption),
		floats:      make(map[string]*FloatOption),
		args:        make([]string, 0),
		invalidArgs: make([]string, 0),
	}
}

// BoolOption type
type BoolOption struct {
	Name      string
	ShortName byte
	Value     bool
	HelpText  string
}

// StringOption type
type StringOption struct {
	Name      string
	ShortName byte
	Value     string
	HelpText  string
}

// IntOption type
type IntOption struct {
	Name      string
	ShortName byte
	Value     int
	HelpText  string
}

// FloatOption type
type FloatOption struct {
	Name      string
	ShortName byte
	Value     float64
	HelpText  string
}

// DefineBool func
func DefineBool(name string, shortName byte, defaultVal bool) *BoolOption {
	data := &BoolOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.bools[name] = data
	bf.bools[string(shortName)] = data
	return data
}

// DefineString func
func DefineString(name string, shortName byte, defaultVal string) *StringOption {
	data := &StringOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.strings[name] = data
	bf.strings[string(shortName)] = data
	return data
}

// DefineInt func
func DefineInt(name string, shortName byte, defaultVal int) *IntOption {
	data := &IntOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.ints[name] = data
	bf.ints[string(shortName)] = data
	return data
}

// DefineFloat func
func DefineFloat(name string, shortName byte, defaultVal float64) *FloatOption {
	data := &FloatOption{
		Name:      name,
		ShortName: shortName,
		Value:     defaultVal,
		HelpText:  "",
	}
	bf.floats[name] = data
	bf.floats[string(shortName)] = data
	return data
}

// Args returns non-option arguments
func Args() []string {
	return bf.args
}

// Invalid returns invalid options
func Invalid() []string {
	return bf.invalidArgs
}

// GetBool returns the specified *BoolOption if defined
func GetBool(name string) *BoolOption {
	if b, ok := bf.bools[name]; ok {
		return b
	}
	return nil
}

// SetBool func
func SetBool(name string, value bool) error {
	if obj, ok := bf.bools[name]; ok {
		obj.Value = value
	} else {
		return errors.New("cannot set undefined bool: " + name)
	}
	return nil
}

// GetString returns the specified *StringOption if defined
func GetString(name string) *StringOption {
	if s, ok := bf.strings[name]; ok {
		return s
	}
	return nil
}

// SetString func
func SetString(name string, value string) error {
	if s, ok := bf.strings[name]; ok {
		s.Value = value
	} else {
		return errors.New("cannot set undefined string: " + name)
	}
	return nil
}

// GetInt returns the specified *IntOption if defined
func GetInt(name string) *IntOption {
	if i, ok := bf.ints[name]; ok {
		return i
	}
	return nil
}

// SetInt func
func SetInt(name string, value int) error {
	if i, ok := bf.ints[name]; ok {
		i.Value = value
	} else {
		return errors.New("cannot set undefined int: " + name)
	}
	return nil
}

// GetFloat returns the specified *FloatOption if defined
func GetFloat(name string) *FloatOption {
	if f, ok := bf.floats[name]; ok {
		return f
	}
	return nil
}

// SetFloat func
func SetFloat(name string, value float64) error {
	if f, ok := bf.floats[name]; ok {
		f.Value = value
	} else {
		return errors.New("cannot set undefined float: " + name)
	}
	return nil
}

// Parse will merge the command-line arguments into user-defined options
func Parse() {
	osArgs := os.Args[1:]
	removeSet := map[int]bool{}
	invalidArgs := []string{}

	// Expand short bool options
	args := make([]string, 0, len(osArgs))
	for _, arg := range osArgs {
		name, _ := parseOption(arg)
		if isValidCombinedFlag(arg) {
			nameArr := strings.Split(name, "")
			for _, v := range nameArr {
				args = append(args, "-"+v)
			}
		} else {
			args = append(args, arg)
		}
	}

	for i, arg := range args {
		// Check to see if there is a valid option set
		if isValidOption(arg) {
			name, val := parseOption(arg)

			// Parse boolean option
			if boolOpt := GetBool(name); boolOpt != nil {
				if isValidFlag(arg) {
					if val == "true" {
						SetBool(name, true)
					} else if val == "false" {
						SetBool(name, false)
					} else {
						SetBool(name, true)
					}
					removeSet[i] = true
					continue
				} else {
					invalidArgs = append(invalidArgs, "invalid option '"+arg+"'")
					removeSet[i] = true
					continue
				}
			} else {
				if isValidFlag(arg) {
					isStr := GetString(name)
					isInt := GetInt(name)
					isFloat := GetFloat(name)

					// The argument is a valid single-character option (e.g. -a)
					// but does not match any other types; invalid option
					if isStr == nil && isInt == nil && isFloat == nil {
						// bool flag passed but does not match any defined bool types
						invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")
						removeSet[i] = true
						continue
					}
				}
			}

			// Check `option=value` vs. `option value` for non-boolean flags
			if strings.Index(arg, "=") == -1 {
				// must get the next command-line argument for value
				if i == len(args)-1 {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is missing a value")
					removeSet[i] = true
					continue
				}
				val = args[i+1]
				removeSet[i] = true
				removeSet[i+1] = true
			} else {
				removeSet[i] = true
			}

			// Parse string option
			if strOpt := GetString(name); strOpt != nil {
				SetString(name, val)
				continue
			}

			// Parse int option
			if intOpt := GetInt(name); intOpt != nil {
				if intVal, err := strconv.Atoi(val); err != nil {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is an invalid int ("+val+")")
				} else {
					SetInt(name, intVal)
				}
				continue
			}

			// Parse float option
			if floatOpt := GetFloat(name); floatOpt != nil {
				if floatVal, err := strconv.ParseFloat(val, 64); err != nil {
					invalidArgs = append(invalidArgs, "option '"+arg+"' is an invalid float ("+val+")")
				} else {
					SetFloat(name, floatVal)
				}
				continue
			}

			// Doesn't match anything else; invalid option
			invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")

		} else {
			// Invalid option found, check to see if it begins with - or --
			isOption, _ := regexp.MatchString("\\A-{1,2}", arg)
			if isOption {
				invalidArgs = append(invalidArgs, "'"+arg+"' is not a valid option")
				removeSet[i] = true
			}
		}
	}

	// Construct list of sorted indexes to remove based on removeSet
	var removeIndexes []int
	for i := range removeSet {
		removeIndexes = append(removeIndexes, i)
	}
	sort.Ints(removeIndexes)

	// Populate Args with all other non-option command-line arguments
	for i := len(removeIndexes) - 1; i >= 0; i-- {
		j := removeIndexes[i]
		if j == len(args)-1 {
			args = args[:j]
		} else {
			args = append(args[:j], args[j+1:]...)
		}
	}
	bf.args = args
	bf.invalidArgs = invalidArgs
}
<|endoftext|>