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

package hostagent

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"

	"encoding/json"

	uuid "github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

type opflexFault struct {
	FaultUUID   string `json:"fault_uuid"`
	Severity    string `json:"severity"`
	Description string `json:"description"`
	FaultCode   int    `json:"faultCode"`
}

func writeFault(faultfile string, ep *opflexFault) (bool, error) {
	newdata, err := json.MarshalIndent(ep, "", "  ")
	if err != nil {
		return true, err
	}
	existingdata, err := ioutil.ReadFile(faultfile)
	if err == nil && reflect.DeepEqual(existingdata, newdata) {
		return false, nil
	}
	err = ioutil.WriteFile(faultfile, newdata, 0644)
	return true, err
}

func (agent *HostAgent) createFaultOnAgent(description string, faultCode int) {
	if agent.config.OpFlexFaultDir == "" {
		agent.log.Error("OpFlex Fault directory not set")
		return
	}
	Uuid := uuid.New().String()
	faultFilePath := filepath.Join(agent.config.OpFlexFaultDir, description+".fs")
	faultFileExists := fileExists(faultFilePath)
	if faultFileExists {
		agent.log.Debug("Fault file exist at: ", faultFilePath)
		return
	}
	desc := strings.Replace(description, "_", " ", -1)
	fault := &opflexFault{
		FaultUUID:   Uuid,
		Severity:    "critical",
		Description: desc,
		FaultCode:   faultCode,
	}
	wrote, err := writeFault(faultFilePath, fault)
	if err != nil {
		agent.log.Warn("Unable to write fault file: ", err.Error())
	} else if wrote {
		agent.log.Debug("Created fault files at the location: ", faultFilePath)
	}
	return
}

func (agent *HostAgent) discoverHostConfig() (conf *HostAgentNodeConfig) {
	if agent.config.OpflexMode == "overlay" {
		conf = &HostAgentNodeConfig{}
		conf.OpflexPeerIp = "127.0.0.1"
		agent.log.Debug("\n  == Opflex: Running in overlay mode ==\n")
		return
	}

	links, err := netlink.LinkList()
	if err != nil {
		agent.log.Error("Could not enumerate interfaces: ", err)
		description := "Could_not_enumerate_interfaces"
		agent.createFaultOnAgent(description, 3)
		return
	}

	for _, link := range links {
		switch link := link.(type) {
		case *netlink.Vlan:
			// find link with matching vlan
			if link.VlanId != int(agent.config.AciInfraVlan) {
				continue
			}
			// if the interface MTU was not explicitly set by
			// the user, use the link MTU
			if agent.config.InterfaceMtu == 0 {
				agent.config.InterfaceMtu = link.MTU - agent.config.InterfaceMtuHeadroom
			}
			// giving extra headroom of 100 bytes unless specified otherwise
			configMtu := agent.config.InterfaceMtuHeadroom + agent.config.InterfaceMtu
			if link.MTU < configMtu {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
					"vlan": agent.config.AciInfraVlan,
					"mtu":  link.MTU,
				}).Error("OpFlex link MTU must be >= ", configMtu)
				description := "User_configured_MTU_exceeds_opflex_MTU"
				agent.createFaultOnAgent(description, 4)
				return
			}

			// find parent link
			var parent netlink.Link
			for _, plink := range links {
				if plink.Attrs().Index != link.ParentIndex {
					continue
				}

				parent = plink
				if parent.Attrs().MTU < configMtu {
					agent.log.WithFields(logrus.Fields{
						"name": parent.Attrs().Name,
						"vlan": agent.config.AciInfraVlan,
						"mtu":  parent.Attrs().MTU,
					}).Error("Uplink MTU must be >= ", configMtu)
					description := "User_configured_MTU_exceed_uplink_MTU"
					agent.createFaultOnAgent(description, 5)
					return
				}
			}
			if parent == nil {
				agent.log.WithFields(logrus.Fields{
					"index": link.ParentIndex,
					"name":  link.Name,
				}).Error("Could not find parent link for OpFlex interface")
				description := "Could_not_find_parent_link_for_OpFlex_interface"
				agent.createFaultOnAgent(description, 6)
				return
			}

			// Find address of link to compute anycast and peer IPs
			addrs, err := netlink.AddrList(link, 2)
			if err != nil {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
				}).Error("Could not enumerate link addresses: ", err)
				description := "Could_not_enumerate_link_addresses"
				agent.createFaultOnAgent(description, 7)
				return
			}
			var anycast net.IP
			var peerIp net.IP
			for _, addr := range addrs {
				if addr.IP.To4() == nil || addr.IP.IsLoopback() {
					continue
				}
				anycast = addr.IP.Mask(addr.Mask)
				anycast[len(anycast)-1] = 32
				peerIp = addr.IP.Mask(addr.Mask)
				peerIp[len(peerIp)-1] = 30
			}

			if anycast == nil {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
					"vlan": agent.config.AciInfraVlan,
				}).Error("IP address not set for OpFlex link")
				description := "IP_address_not_set_for_OpFlex_link"
				agent.createFaultOnAgent(description, 8)
				return
			}

			conf = &HostAgentNodeConfig{}
			conf.VxlanIface = link.Name
			conf.UplinkIface = parent.Attrs().Name
			conf.VxlanAnycastIp = anycast.String()
			conf.OpflexPeerIp = peerIp.String()
		}
	}

	if conf != nil {
		intf, err := net.InterfaceByName(conf.UplinkIface)
		if err == nil {
			conf.UplinkMacAdress = intf.HardwareAddr.String()
			return
		}
	}

	agent.log.WithFields(logrus.Fields{"vlan": agent.config.AciInfraVlan}).
		Error("Could not find suitable host uplink interface for vlan")
	description := "Could_not_find_suitable_host_uplink_interface_for_vlan"
	agent.createFaultOnAgent(description, 9)
	return
}

var opflexConfigBase = initTempl("opflex-config-base", `{
    "opflex": {
        "name": "{{.NodeName | js}}",
        "domain": "{{print "comp/prov-" .AciVmmDomainType "/ctrlr-[" .AciVmmDomain "]-" .AciVmmController "/sw-InsiemeLSOid" | js}}",
        "peers": [
            {"hostname": "{{.OpflexPeerIp | js}}", "port": "8009"}
        ]
    } ,
    "endpoint-sources": {
        "filesystem": ["{{.OpFlexEndpointDir | js}}"]
    },
    "service-sources": {
        "filesystem": ["{{.OpFlexServiceDir | js}}"]
    },
    "snat-sources": {
        "filesystem": ["{{.OpFlexSnatDir | js}}"]
    },
    "drop-log-config-sources": {
        "filesystem": ["{{.OpFlexDropLogConfigDir | js}}"]
    },
    "packet-event-notif": {
        "socket-name": ["{{.PacketEventNotificationSock | js}}"]
    },
    "host-agent-fault-sources": {
        "filesystem": ["{{.OpFlexFaultDir | js}}"]
    }
}
`)

var opflexConfigVxlan = initTempl("opflex-config-vxlan", `{
    "renderers": {
        "stitched-mode": {
            "int-bridge-name": "{{.IntBridgeName | js}}",
            "access-bridge-name": "{{.AccessBridgeName | js}}",
            "encap": {
                "vxlan" : {
                    "encap-iface": "vxlan0",
                    "uplink-iface": "{{.VxlanIface | js}}",
                    "uplink-vlan": "{{.AciInfraVlan}}",
                    "remote-ip": "{{.VxlanAnycastIp | js}}",
                    "remote-port": 8472
                }
            },
            "flowid-cache-dir": "{{.OpFlexFlowIdCacheDir | js}}",
            "mcast-group-file": "{{.OpFlexMcastFile | js}}",
            "drop-log": {
		"geneve" : {
		    "int-br-iface": "{{.DropLogIntInterface | js}}",
		    "access-br-iface": "{{.DropLogAccessInterface | js}}",
		    "remote-ip": "{{.OpFlexDropLogRemoteIp | js}}"
		}
	    }
        }
    }
}
`)

var opflexConfigVlan = initTempl("opflex-config-vlan", `{
    "renderers": {
        "stitched-mode": {
            "int-bridge-name": "{{.IntBridgeName | js}}",
            "access-bridge-name": "{{.AccessBridgeName | js}}",
            "encap": {
                "vlan" : {
                    "encap-iface": "{{.UplinkIface | js}}"
                }
            },
            "flowid-cache-dir": "{{.OpFlexFlowIdCacheDir | js}}",
            "mcast-group-file": "{{.OpFlexMcastFile | js}}",
            "drop-log": {
		"geneve" : {
		    "int-br-iface": "{{.DropLogIntInterface | js}}",
		    "access-br-iface": "{{.DropLogAccessInterface | js}}",
		    "remote-ip": "{{.OpFlexDropLogRemoteIp | js}}"
		}
	    }
        }
    }
}
`)

func initTempl(name string, templ string) *template.Template {
	return template.Must(template.New(name).Parse(templ))
}

func (agent *HostAgent) writeConfigFile(name string,
	templ *template.Template) error {

	var buffer bytes.Buffer
	templ.Execute(&buffer, agent.config)

	path := filepath.Join(agent.config.OpFlexConfigPath, name)

	existing, err := ioutil.ReadFile(path)
	if err != nil {
		if bytes.Equal(existing, buffer.Bytes()) {
			agent.log.Info("OpFlex agent configuration file ",
				path, " unchanged")
			return nil
		}
	}

	f, err := os.Create(path)
	if err != nil {
		return err
	}
	// in case there's an error in the write
	defer f.Close()
	_, err = f.Write(buffer.Bytes())
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	agent.log.Info("Wrote OpFlex agent configuration file ", path)

	return nil
}

func (agent *HostAgent) updateOpflexConfig() {
	if agent.config.OpFlexConfigPath == "" {
		agent.log.Debug("OpFlex agent configuration path not set")
		return
	}
	if agent.config.OpFlexFaultDir == "" {
		agent.log.Warn("OpFlex Fault directory not set")
	} else {
		err := agent.removeAllFiles(agent.config.OpFlexFaultDir)
		if err != nil {
			agent.log.Error("Not able to clear Fault files on agent: ", err.Error())
		}
	}

	newNodeConfig := agent.discoverHostConfig()
	if newNodeConfig == nil {
		panic(errors.New("Node configuration autodiscovery failed"))
	}
	var update bool

	agent.indexMutex.Lock()
	if !reflect.DeepEqual(*newNodeConfig, agent.config.HostAgentNodeConfig) ||
		!agent.opflexConfigWritten {

		// reset opflexConfigWritten flag when node-config differs
		agent.opflexConfigWritten = false

		agent.config.HostAgentNodeConfig = *newNodeConfig
		agent.log.WithFields(logrus.Fields{
			"uplink-iface":     newNodeConfig.UplinkIface,
			"vxlan-iface":      newNodeConfig.VxlanIface,
			"vxlan-anycast-ip": newNodeConfig.VxlanAnycastIp,
			"opflex-peer-ip":   newNodeConfig.OpflexPeerIp,
			"opflex-mode":      agent.config.OpflexMode,
		}).Info("Discovered node configuration")
		if err := agent.writeOpflexConfig(); err == nil {
			agent.opflexConfigWritten = true
		} else {
			agent.log.Error("Failed to write OpFlex agent config: ", err)
		}
	}
	agent.indexMutex.Unlock()

	if update {
		agent.updateAllServices()
	}
}

func (agent *HostAgent) writeOpflexConfig() error {
	err := agent.writeConfigFile("01-base.conf", opflexConfigBase)
	if err != nil {
		return err
	}

	var rtempl *template.Template
	if agent.config.EncapType == "vlan" {
		rtempl = opflexConfigVlan
	} else if agent.config.EncapType == "vxlan" {
		rtempl = opflexConfigVxlan
	} else {
		panic("Unsupported encap type: " + agent.config.EncapType)
	}

	err = agent.writeConfigFile("10-renderer.conf", rtempl)
	if err != nil {
		return err
	}
	return nil
}

func (agent *HostAgent) removeAllFiles(dir string) error {
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			agent.log.Error("Not able to clear the Fault Files  ", err)
			return err
		}
	}
	return nil
}
<commit_msg>Disable service / service flow stats<commit_after>// Copyright 2017 Cisco Systems, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package hostagent

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"

	"encoding/json"

	uuid "github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

type opflexFault struct {
	FaultUUID   string `json:"fault_uuid"`
	Severity    string `json:"severity"`
	Description string `json:"description"`
	FaultCode   int    `json:"faultCode"`
}

func writeFault(faultfile string, ep *opflexFault) (bool, error) {
	newdata, err := json.MarshalIndent(ep, "", "  ")
	if err != nil {
		return true, err
	}
	existingdata, err := ioutil.ReadFile(faultfile)
	if err == nil && reflect.DeepEqual(existingdata, newdata) {
		return false, nil
	}
	err = ioutil.WriteFile(faultfile, newdata, 0644)
	return true, err
}

func (agent *HostAgent) createFaultOnAgent(description string, faultCode int) {
	if agent.config.OpFlexFaultDir == "" {
		agent.log.Error("OpFlex Fault directory not set")
		return
	}
	Uuid := uuid.New().String()
	faultFilePath := filepath.Join(agent.config.OpFlexFaultDir, description+".fs")
	faultFileExists := fileExists(faultFilePath)
	if faultFileExists {
		agent.log.Debug("Fault file exist at: ", faultFilePath)
		return
	}
	desc := strings.Replace(description, "_", " ", -1)
	fault := &opflexFault{
		FaultUUID:   Uuid,
		Severity:    "critical",
		Description: desc,
		FaultCode:   faultCode,
	}
	wrote, err := writeFault(faultFilePath, fault)
	if err != nil {
		agent.log.Warn("Unable to write fault file: ", err.Error())
	} else if wrote {
		agent.log.Debug("Created fault files at the location: ", faultFilePath)
	}
	return
}

func (agent *HostAgent) discoverHostConfig() (conf *HostAgentNodeConfig) {
	if agent.config.OpflexMode == "overlay" {
		conf = &HostAgentNodeConfig{}
		conf.OpflexPeerIp = "127.0.0.1"
		agent.log.Debug("\n  == Opflex: Running in overlay mode ==\n")
		return
	}

	links, err := netlink.LinkList()
	if err != nil {
		agent.log.Error("Could not enumerate interfaces: ", err)
		description := "Could_not_enumerate_interfaces"
		agent.createFaultOnAgent(description, 3)
		return
	}

	for _, link := range links {
		switch link := link.(type) {
		case *netlink.Vlan:
			// find link with matching vlan
			if link.VlanId != int(agent.config.AciInfraVlan) {
				continue
			}
			// if the interface MTU was not explicitly set by
			// the user, use the link MTU
			if agent.config.InterfaceMtu == 0 {
				agent.config.InterfaceMtu = link.MTU - agent.config.InterfaceMtuHeadroom
			}
			// giving extra headroom of 100 bytes unless specified otherwise
			configMtu := agent.config.InterfaceMtuHeadroom + agent.config.InterfaceMtu
			if link.MTU < configMtu {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
					"vlan": agent.config.AciInfraVlan,
					"mtu":  link.MTU,
				}).Error("OpFlex link MTU must be >= ", configMtu)
				description := "User_configured_MTU_exceeds_opflex_MTU"
				agent.createFaultOnAgent(description, 4)
				return
			}

			// find parent link
			var parent netlink.Link
			for _, plink := range links {
				if plink.Attrs().Index != link.ParentIndex {
					continue
				}

				parent = plink
				if parent.Attrs().MTU < configMtu {
					agent.log.WithFields(logrus.Fields{
						"name": parent.Attrs().Name,
						"vlan": agent.config.AciInfraVlan,
						"mtu":  parent.Attrs().MTU,
					}).Error("Uplink MTU must be >= ", configMtu)
					description := "User_configured_MTU_exceed_uplink_MTU"
					agent.createFaultOnAgent(description, 5)
					return
				}
			}
			if parent == nil {
				agent.log.WithFields(logrus.Fields{
					"index": link.ParentIndex,
					"name":  link.Name,
				}).Error("Could not find parent link for OpFlex interface")
				description := "Could_not_find_parent_link_for_OpFlex_interface"
				agent.createFaultOnAgent(description, 6)
				return
			}

			// Find address of link to compute anycast and peer IPs
			addrs, err := netlink.AddrList(link, 2)
			if err != nil {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
				}).Error("Could not enumerate link addresses: ", err)
				description := "Could_not_enumerate_link_addresses"
				agent.createFaultOnAgent(description, 7)
				return
			}
			var anycast net.IP
			var peerIp net.IP
			for _, addr := range addrs {
				if addr.IP.To4() == nil || addr.IP.IsLoopback() {
					continue
				}
				anycast = addr.IP.Mask(addr.Mask)
				anycast[len(anycast)-1] = 32
				peerIp = addr.IP.Mask(addr.Mask)
				peerIp[len(peerIp)-1] = 30
			}

			if anycast == nil {
				agent.log.WithFields(logrus.Fields{
					"name": link.Name,
					"vlan": agent.config.AciInfraVlan,
				}).Error("IP address not set for OpFlex link")
				description := "IP_address_not_set_for_OpFlex_link"
				agent.createFaultOnAgent(description, 8)
				return
			}

			conf = &HostAgentNodeConfig{}
			conf.VxlanIface = link.Name
			conf.UplinkIface = parent.Attrs().Name
			conf.VxlanAnycastIp = anycast.String()
			conf.OpflexPeerIp = peerIp.String()
		}
	}

	if conf != nil {
		intf, err := net.InterfaceByName(conf.UplinkIface)
		if err == nil {
			conf.UplinkMacAdress = intf.HardwareAddr.String()
			return
		}
	}

	agent.log.WithFields(logrus.Fields{"vlan": agent.config.AciInfraVlan}).
		Error("Could not find suitable host uplink interface for vlan")
	description := "Could_not_find_suitable_host_uplink_interface_for_vlan"
	agent.createFaultOnAgent(description, 9)
	return
}

var opflexConfigBase = initTempl("opflex-config-base", `{
    "opflex": {
        "name": "{{.NodeName | js}}",
        "domain": "{{print "comp/prov-" .AciVmmDomainType "/ctrlr-[" .AciVmmDomain "]-" .AciVmmController "/sw-InsiemeLSOid" | js}}",
        "peers": [
            {"hostname": "{{.OpflexPeerIp | js}}", "port": "8009"}
        ]
    } ,
    "endpoint-sources": {
        "filesystem": ["{{.OpFlexEndpointDir | js}}"]
    },
    "service-sources": {
        "filesystem": ["{{.OpFlexServiceDir | js}}"]
    },
    "snat-sources": {
        "filesystem": ["{{.OpFlexSnatDir | js}}"]
    },
    "drop-log-config-sources": {
        "filesystem": ["{{.OpFlexDropLogConfigDir | js}}"]
    },
    "packet-event-notif": {
        "socket-name": ["{{.PacketEventNotificationSock | js}}"]
    },
    "host-agent-fault-sources": {
        "filesystem": ["{{.OpFlexFaultDir | js}}"]
    }
}
`)

var opflexConfigVxlan = initTempl("opflex-config-vxlan", `{
    "renderers": {
        "stitched-mode": {
            "int-bridge-name": "{{.IntBridgeName | js}}",
            "access-bridge-name": "{{.AccessBridgeName | js}}",
            "encap": {
                "vxlan" : {
                    "encap-iface": "vxlan0",
                    "uplink-iface": "{{.VxlanIface | js}}",
                    "uplink-vlan": "{{.AciInfraVlan}}",
                    "remote-ip": "{{.VxlanAnycastIp | js}}",
                    "remote-port": 8472
                }
            },
            "flowid-cache-dir": "{{.OpFlexFlowIdCacheDir | js}}",
            "mcast-group-file": "{{.OpFlexMcastFile | js}}",
            "drop-log": {
                "geneve" : {
                    "int-br-iface": "{{.DropLogIntInterface | js}}",
                    "access-br-iface": "{{.DropLogAccessInterface | js}}",
                    "remote-ip": "{{.OpFlexDropLogRemoteIp | js}}"
                }
            },
            "statistics": {
                "service": {
                    "flow-disabled": "true",
                    "enabled": "false"
                }
            }
        }
    }
}
`)

var opflexConfigVlan = initTempl("opflex-config-vlan", `{
    "renderers": {
        "stitched-mode": {
            "int-bridge-name": "{{.IntBridgeName | js}}",
            "access-bridge-name": "{{.AccessBridgeName | js}}",
            "encap": {
                "vlan" : {
                    "encap-iface": "{{.UplinkIface | js}}"
                }
            },
            "flowid-cache-dir": "{{.OpFlexFlowIdCacheDir | js}}",
            "mcast-group-file": "{{.OpFlexMcastFile | js}}",
            "drop-log": {
		"geneve" : {
		    "int-br-iface": "{{.DropLogIntInterface | js}}",
		    "access-br-iface": "{{.DropLogAccessInterface | js}}",
		    "remote-ip": "{{.OpFlexDropLogRemoteIp | js}}"
		}
	    }
        }
    }
}
`)

func initTempl(name string, templ string) *template.Template {
	return template.Must(template.New(name).Parse(templ))
}

func (agent *HostAgent) writeConfigFile(name string,
	templ *template.Template) error {

	var buffer bytes.Buffer
	templ.Execute(&buffer, agent.config)

	path := filepath.Join(agent.config.OpFlexConfigPath, name)

	existing, err := ioutil.ReadFile(path)
	if err != nil {
		if bytes.Equal(existing, buffer.Bytes()) {
			agent.log.Info("OpFlex agent configuration file ",
				path, " unchanged")
			return nil
		}
	}

	f, err := os.Create(path)
	if err != nil {
		return err
	}
	// in case there's an error in the write
	defer f.Close()
	_, err = f.Write(buffer.Bytes())
	if err != nil {
		return err
	}
	err = f.Close()
	if err != nil {
		return err
	}

	agent.log.Info("Wrote OpFlex agent configuration file ", path)

	return nil
}

func (agent *HostAgent) updateOpflexConfig() {
	if agent.config.OpFlexConfigPath == "" {
		agent.log.Debug("OpFlex agent configuration path not set")
		return
	}
	if agent.config.OpFlexFaultDir == "" {
		agent.log.Warn("OpFlex Fault directory not set")
	} else {
		err := agent.removeAllFiles(agent.config.OpFlexFaultDir)
		if err != nil {
			agent.log.Error("Not able to clear Fault files on agent: ", err.Error())
		}
	}

	newNodeConfig := agent.discoverHostConfig()
	if newNodeConfig == nil {
		panic(errors.New("Node configuration autodiscovery failed"))
	}
	var update bool

	agent.indexMutex.Lock()
	if !reflect.DeepEqual(*newNodeConfig, agent.config.HostAgentNodeConfig) ||
		!agent.opflexConfigWritten {

		// reset opflexConfigWritten flag when node-config differs
		agent.opflexConfigWritten = false

		agent.config.HostAgentNodeConfig = *newNodeConfig
		agent.log.WithFields(logrus.Fields{
			"uplink-iface":     newNodeConfig.UplinkIface,
			"vxlan-iface":      newNodeConfig.VxlanIface,
			"vxlan-anycast-ip": newNodeConfig.VxlanAnycastIp,
			"opflex-peer-ip":   newNodeConfig.OpflexPeerIp,
			"opflex-mode":      agent.config.OpflexMode,
		}).Info("Discovered node configuration")
		if err := agent.writeOpflexConfig(); err == nil {
			agent.opflexConfigWritten = true
		} else {
			agent.log.Error("Failed to write OpFlex agent config: ", err)
		}
	}
	agent.indexMutex.Unlock()

	if update {
		agent.updateAllServices()
	}
}

func (agent *HostAgent) writeOpflexConfig() error {
	err := agent.writeConfigFile("01-base.conf", opflexConfigBase)
	if err != nil {
		return err
	}

	var rtempl *template.Template
	if agent.config.EncapType == "vlan" {
		rtempl = opflexConfigVlan
	} else if agent.config.EncapType == "vxlan" {
		rtempl = opflexConfigVxlan
	} else {
		panic("Unsupported encap type: " + agent.config.EncapType)
	}

	err = agent.writeConfigFile("10-renderer.conf", rtempl)
	if err != nil {
		return err
	}
	return nil
}

func (agent *HostAgent) removeAllFiles(dir string) error {
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			agent.log.Error("Not able to clear the Fault Files  ", err)
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package hub

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"net/http"

	"github.com/traefik/traefik/v2/pkg/config/dynamic"
	"github.com/traefik/traefik/v2/pkg/log"
	"github.com/traefik/traefik/v2/pkg/provider"
	"github.com/traefik/traefik/v2/pkg/safe"
	ttls "github.com/traefik/traefik/v2/pkg/tls"
)

var _ provider.Provider = (*Provider)(nil)

// Entrypoints created for Hub.
const (
	APIEntrypoint    = "traefikhub-api"
	TunnelEntrypoint = "traefikhub-tunl"
)

// Provider holds configurations of the provider.
type Provider struct {
	TLS *TLS `description:"TLS configuration for mTLS communication between Traefik and Hub Agent." json:"tls,omitempty" toml:"tls,omitempty" yaml:"tls,omitempty" export:"true"`

	server *http.Server
}

// TLS configures the mTLS connection between Traefik Proxy and the Traefik Hub Agent.
type TLS struct {
	Insecure bool               `description:"Enables an insecure TLS connection that uses default credentials, and which has no peer authentication between Traefik Proxy and the Traefik Hub Agent." json:"insecure,omitempty" toml:"insecure,omitempty" yaml:"insecure,omitempty" export:"true"`
	CA       ttls.FileOrContent `description:"The certificate authority authenticates the Traefik Hub Agent certificate." json:"ca,omitempty" toml:"ca,omitempty" yaml:"ca,omitempty" loggable:"false"`
	Cert     ttls.FileOrContent `description:"The TLS certificate for Traefik Proxy as a TLS client." json:"cert,omitempty" toml:"cert,omitempty" yaml:"cert,omitempty" loggable:"false"`
	Key      ttls.FileOrContent `description:"The TLS key for Traefik Proxy as a TLS client." json:"key,omitempty" toml:"key,omitempty" yaml:"key,omitempty" loggable:"false"`
}

// Init the provider.
func (p *Provider) Init() error {
	return nil
}

// Provide allows the hub provider to provide configurations to traefik using the given configuration channel.
func (p *Provider) Provide(configurationChan chan<- dynamic.Message, _ *safe.Pool) error {
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("listener: %w", err)
	}
	port := listener.Addr().(*net.TCPAddr).Port

	client, err := createAgentClient(p.TLS)
	if err != nil {
		return fmt.Errorf("creating Hub Agent HTTP client: %w", err)
	}

	p.server = &http.Server{Handler: newHandler(APIEntrypoint, port, configurationChan, p.TLS, client)}

	// TODO: this is going to be leaky (because no context to make it terminate)
	// if/when Provide lifecycle differs with Traefik lifecycle.
	go func() {
		if err = p.server.Serve(listener); err != nil {
			log.WithoutContext().WithField(log.ProviderName, "hub").Errorf("Unexpected error while running server: %v", err)
			return
		}
	}()

	exposeAPIAndMetrics(configurationChan, APIEntrypoint, port, p.TLS)

	return nil
}

func exposeAPIAndMetrics(cfgChan chan<- dynamic.Message, ep string, port int, tlsCfg *TLS) {
	cfg := emptyDynamicConfiguration()

	patchDynamicConfiguration(cfg, ep, port, tlsCfg)

	cfgChan <- dynamic.Message{ProviderName: "hub", Configuration: cfg}
}

func patchDynamicConfiguration(cfg *dynamic.Configuration, ep string, port int, tlsCfg *TLS) {
	cfg.HTTP.Routers["traefik-hub-agent-api"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "api@internal",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/api`)",
	}
	cfg.HTTP.Routers["traefik-hub-agent-metrics"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "prometheus@internal",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/metrics`)",
	}

	cfg.HTTP.Routers["traefik-hub-agent-service"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "traefik-hub-agent-service",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/config`, `/discover-ip`, `/state`)",
	}

	cfg.HTTP.Services["traefik-hub-agent-service"] = &dynamic.Service{
		LoadBalancer: &dynamic.ServersLoadBalancer{
			Servers: []dynamic.Server{
				{
					URL: fmt.Sprintf("http://127.0.0.1:%d", port),
				},
			},
		},
	}

	if tlsCfg == nil {
		return
	}

	if tlsCfg.Insecure {
		cfg.TLS.Options["traefik-hub"] = ttls.Options{
			MinVersion: "VersionTLS13",
		}

		return
	}

	cfg.TLS.Options["traefik-hub"] = ttls.Options{
		ClientAuth: ttls.ClientAuth{
			CAFiles:        []ttls.FileOrContent{tlsCfg.CA},
			ClientAuthType: "RequireAndVerifyClientCert",
		},
		SniStrict:  true,
		MinVersion: "VersionTLS13",
	}

	cfg.TLS.Certificates = append(cfg.TLS.Certificates, &ttls.CertAndStores{
		Certificate: ttls.Certificate{
			CertFile: tlsCfg.Cert,
			KeyFile:  tlsCfg.Key,
		},
	})
}

func emptyDynamicConfiguration() *dynamic.Configuration {
	return &dynamic.Configuration{
		HTTP: &dynamic.HTTPConfiguration{
			Routers:           make(map[string]*dynamic.Router),
			Middlewares:       make(map[string]*dynamic.Middleware),
			Services:          make(map[string]*dynamic.Service),
			ServersTransports: make(map[string]*dynamic.ServersTransport),
		},
		TCP: &dynamic.TCPConfiguration{
			Routers:  make(map[string]*dynamic.TCPRouter),
			Services: make(map[string]*dynamic.TCPService),
		},
		TLS: &dynamic.TLSConfiguration{
			Stores:  make(map[string]ttls.Store),
			Options: make(map[string]ttls.Options),
		},
		UDP: &dynamic.UDPConfiguration{
			Routers:  make(map[string]*dynamic.UDPRouter),
			Services: make(map[string]*dynamic.UDPService),
		},
	}
}

func createAgentClient(tlsCfg *TLS) (http.Client, error) {
	var client http.Client
	if tlsCfg.Insecure {
		client.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
				MinVersion:         tls.VersionTLS13,
			},
		}

		return client, nil
	}

	caContent, err := tlsCfg.CA.Read()
	if err != nil {
		return client, fmt.Errorf("reading CA: %w", err)
	}

	roots := x509.NewCertPool()
	if ok := roots.AppendCertsFromPEM(caContent); !ok {
		return client, errors.New("appending CA error")
	}

	certContent, err := tlsCfg.Cert.Read()
	if err != nil {
		return client, fmt.Errorf("reading Cert: %w", err)
	}
	keyContent, err := tlsCfg.Key.Read()
	if err != nil {
		return client, fmt.Errorf("reading Key: %w", err)
	}

	certificate, err := tls.X509KeyPair(certContent, keyContent)
	if err != nil {
		return client, fmt.Errorf("creating key pair: %w", err)
	}

	// mTLS
	client.Transport = &http.Transport{
		TLSClientConfig: &tls.Config{
			RootCAs:      roots,
			Certificates: []tls.Certificate{certificate},
			ServerName:   "agent.traefik",
			ClientAuth:   tls.RequireAndVerifyClientCert,
			MinVersion:   tls.VersionTLS13,
		},
	}

	return client, nil
}
<commit_msg>fix: skip Provide when TLS is nil<commit_after>package hub

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"net"
	"net/http"

	"github.com/traefik/traefik/v2/pkg/config/dynamic"
	"github.com/traefik/traefik/v2/pkg/log"
	"github.com/traefik/traefik/v2/pkg/provider"
	"github.com/traefik/traefik/v2/pkg/safe"
	ttls "github.com/traefik/traefik/v2/pkg/tls"
)

var _ provider.Provider = (*Provider)(nil)

// Entrypoints created for Hub.
const (
	APIEntrypoint    = "traefikhub-api"
	TunnelEntrypoint = "traefikhub-tunl"
)

// Provider holds configurations of the provider.
type Provider struct {
	TLS *TLS `description:"TLS configuration for mTLS communication between Traefik and Hub Agent." json:"tls,omitempty" toml:"tls,omitempty" yaml:"tls,omitempty" export:"true"`

	server *http.Server
}

// TLS configures the mTLS connection between Traefik Proxy and the Traefik Hub Agent.
type TLS struct {
	Insecure bool               `description:"Enables an insecure TLS connection that uses default credentials, and which has no peer authentication between Traefik Proxy and the Traefik Hub Agent." json:"insecure,omitempty" toml:"insecure,omitempty" yaml:"insecure,omitempty" export:"true"`
	CA       ttls.FileOrContent `description:"The certificate authority authenticates the Traefik Hub Agent certificate." json:"ca,omitempty" toml:"ca,omitempty" yaml:"ca,omitempty" loggable:"false"`
	Cert     ttls.FileOrContent `description:"The TLS certificate for Traefik Proxy as a TLS client." json:"cert,omitempty" toml:"cert,omitempty" yaml:"cert,omitempty" loggable:"false"`
	Key      ttls.FileOrContent `description:"The TLS key for Traefik Proxy as a TLS client." json:"key,omitempty" toml:"key,omitempty" yaml:"key,omitempty" loggable:"false"`
}

// Init the provider.
func (p *Provider) Init() error {
	return nil
}

// Provide allows the hub provider to provide configurations to traefik using the given configuration channel.
func (p *Provider) Provide(configurationChan chan<- dynamic.Message, _ *safe.Pool) error {
	if p.TLS == nil {
		return nil
	}

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("listener: %w", err)
	}

	port := listener.Addr().(*net.TCPAddr).Port

	client, err := createAgentClient(p.TLS)
	if err != nil {
		return fmt.Errorf("creating Hub Agent HTTP client: %w", err)
	}

	p.server = &http.Server{Handler: newHandler(APIEntrypoint, port, configurationChan, p.TLS, client)}

	// TODO: this is going to be leaky (because no context to make it terminate)
	// if/when Provide lifecycle differs with Traefik lifecycle.
	go func() {
		if err = p.server.Serve(listener); err != nil {
			log.WithoutContext().WithField(log.ProviderName, "hub").Errorf("Unexpected error while running server: %v", err)
			return
		}
	}()

	exposeAPIAndMetrics(configurationChan, APIEntrypoint, port, p.TLS)

	return nil
}

func exposeAPIAndMetrics(cfgChan chan<- dynamic.Message, ep string, port int, tlsCfg *TLS) {
	cfg := emptyDynamicConfiguration()

	patchDynamicConfiguration(cfg, ep, port, tlsCfg)

	cfgChan <- dynamic.Message{ProviderName: "hub", Configuration: cfg}
}

func patchDynamicConfiguration(cfg *dynamic.Configuration, ep string, port int, tlsCfg *TLS) {
	cfg.HTTP.Routers["traefik-hub-agent-api"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "api@internal",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/api`)",
	}
	cfg.HTTP.Routers["traefik-hub-agent-metrics"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "prometheus@internal",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/metrics`)",
	}

	cfg.HTTP.Routers["traefik-hub-agent-service"] = &dynamic.Router{
		EntryPoints: []string{ep},
		Service:     "traefik-hub-agent-service",
		Rule:        "Host(`proxy.traefik`) && PathPrefix(`/config`, `/discover-ip`, `/state`)",
	}

	cfg.HTTP.Services["traefik-hub-agent-service"] = &dynamic.Service{
		LoadBalancer: &dynamic.ServersLoadBalancer{
			Servers: []dynamic.Server{
				{
					URL: fmt.Sprintf("http://127.0.0.1:%d", port),
				},
			},
		},
	}

	if tlsCfg == nil {
		return
	}

	if tlsCfg.Insecure {
		cfg.TLS.Options["traefik-hub"] = ttls.Options{
			MinVersion: "VersionTLS13",
		}

		return
	}

	cfg.TLS.Options["traefik-hub"] = ttls.Options{
		ClientAuth: ttls.ClientAuth{
			CAFiles:        []ttls.FileOrContent{tlsCfg.CA},
			ClientAuthType: "RequireAndVerifyClientCert",
		},
		SniStrict:  true,
		MinVersion: "VersionTLS13",
	}

	cfg.TLS.Certificates = append(cfg.TLS.Certificates, &ttls.CertAndStores{
		Certificate: ttls.Certificate{
			CertFile: tlsCfg.Cert,
			KeyFile:  tlsCfg.Key,
		},
	})
}

func emptyDynamicConfiguration() *dynamic.Configuration {
	return &dynamic.Configuration{
		HTTP: &dynamic.HTTPConfiguration{
			Routers:           make(map[string]*dynamic.Router),
			Middlewares:       make(map[string]*dynamic.Middleware),
			Services:          make(map[string]*dynamic.Service),
			ServersTransports: make(map[string]*dynamic.ServersTransport),
		},
		TCP: &dynamic.TCPConfiguration{
			Routers:  make(map[string]*dynamic.TCPRouter),
			Services: make(map[string]*dynamic.TCPService),
		},
		TLS: &dynamic.TLSConfiguration{
			Stores:  make(map[string]ttls.Store),
			Options: make(map[string]ttls.Options),
		},
		UDP: &dynamic.UDPConfiguration{
			Routers:  make(map[string]*dynamic.UDPRouter),
			Services: make(map[string]*dynamic.UDPService),
		},
	}
}

func createAgentClient(tlsCfg *TLS) (http.Client, error) {
	var client http.Client
	if tlsCfg.Insecure {
		client.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
				MinVersion:         tls.VersionTLS13,
			},
		}

		return client, nil
	}

	caContent, err := tlsCfg.CA.Read()
	if err != nil {
		return client, fmt.Errorf("reading CA: %w", err)
	}

	roots := x509.NewCertPool()
	if ok := roots.AppendCertsFromPEM(caContent); !ok {
		return client, errors.New("appending CA error")
	}

	certContent, err := tlsCfg.Cert.Read()
	if err != nil {
		return client, fmt.Errorf("reading Cert: %w", err)
	}
	keyContent, err := tlsCfg.Key.Read()
	if err != nil {
		return client, fmt.Errorf("reading Key: %w", err)
	}

	certificate, err := tls.X509KeyPair(certContent, keyContent)
	if err != nil {
		return client, fmt.Errorf("creating key pair: %w", err)
	}

	// mTLS
	client.Transport = &http.Transport{
		TLSClientConfig: &tls.Config{
			RootCAs:      roots,
			Certificates: []tls.Certificate{certificate},
			ServerName:   "agent.traefik",
			ClientAuth:   tls.RequireAndVerifyClientCert,
			MinVersion:   tls.VersionTLS13,
		},
	}

	return client, nil
}
<|endoftext|>


<commit_before>package scm_test

import (
	"capsulecd/pkg/config/mock"
	"capsulecd/pkg/pipeline"
	"capsulecd/pkg/scm"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"testing"
)

// Define the suite, and absorb the built-in basic suite
// functionality from testify - including a T() method which
// returns the current testing context
type ScmTestSuite struct {
	suite.Suite
	MockCtrl     *gomock.Controller
	Config       *mock_config.MockInterface
	PipelineData *pipeline.Data
}

// Make sure that VariableThatShouldStartAtFive is set to five
// before each test
func (suite *ScmTestSuite) SetupTest() {
	suite.MockCtrl = gomock.NewController(suite.T())

	suite.PipelineData = new(pipeline.Data)

	suite.Config = mock_config.NewMockInterface(suite.MockCtrl)

}

func (suite *ScmTestSuite) TearDownTest() {
	suite.MockCtrl.Finish()
}

func (suite *ScmTestSuite) TestCreate_Invalid() {
	//test
	testEngine, cerr := scm.Create("invalidtype", suite.PipelineData, suite.Config, nil)

	//assert
	require.Error(suite.T(), cerr, "should return an erro")
	require.Nil(suite.T(), testEngine, "engine should be nil")
}

func (suite *ScmTestSuite) TestCreate_Github() {
	//setup
	suite.Config.EXPECT().GetString("scm_github_access_token").Return("placeholder")
	suite.Config.EXPECT().IsSet("scm_github_api_endpoint").Return(false)
	suite.Config.EXPECT().IsSet("scm_github_access_token").Return(true)
	suite.Config.EXPECT().IsSet("scm_git_parent_path").Return(false)

	//test
	testScm, cerr := scm.Create("github", suite.PipelineData, suite.Config, nil)

	//assert
	require.NoError(suite.T(), cerr)
	require.NotNil(suite.T(), testScm)
}

func (suite *ScmTestSuite) TestCreate_Bitbucket() {
	//setup
	//suite.Config.EXPECT().SetDefault(gomock.Any(),gomock.Any()).MinTimes(1)

	//test
	testScm, cerr := scm.Create("bitbucket", suite.PipelineData, suite.Config, nil)

	//assert
	require.NoError(suite.T(), cerr)
	require.NotNil(suite.T(), testScm)
}

// In order for 'go test' to run this suite, we need to create
// a normal test function and pass our suite to suite.Run
func TestFactoryTestSuite(t *testing.T) {
	suite.Run(t, new(ScmTestSuite))
}
<commit_msg>adding placeholder for bitbucket create in factory.<commit_after>package scm_test

import (
	"capsulecd/pkg/config/mock"
	"capsulecd/pkg/pipeline"
	"capsulecd/pkg/scm"
	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"testing"
)

// Define the suite, and absorb the built-in basic suite
// functionality from testify - including a T() method which
// returns the current testing context
type ScmTestSuite struct {
	suite.Suite
	MockCtrl     *gomock.Controller
	Config       *mock_config.MockInterface
	PipelineData *pipeline.Data
}

// Make sure that VariableThatShouldStartAtFive is set to five
// before each test
func (suite *ScmTestSuite) SetupTest() {
	suite.MockCtrl = gomock.NewController(suite.T())

	suite.PipelineData = new(pipeline.Data)

	suite.Config = mock_config.NewMockInterface(suite.MockCtrl)

}

func (suite *ScmTestSuite) TearDownTest() {
	suite.MockCtrl.Finish()
}

func (suite *ScmTestSuite) TestCreate_Invalid() {
	//test
	testEngine, cerr := scm.Create("invalidtype", suite.PipelineData, suite.Config, nil)

	//assert
	require.Error(suite.T(), cerr, "should return an erro")
	require.Nil(suite.T(), testEngine, "engine should be nil")
}

func (suite *ScmTestSuite) TestCreate_Github() {
	//setup
	suite.Config.EXPECT().GetString("scm_github_access_token").Return("placeholder")
	suite.Config.EXPECT().IsSet("scm_github_api_endpoint").Return(false)
	suite.Config.EXPECT().IsSet("scm_github_access_token").Return(true)
	suite.Config.EXPECT().IsSet("scm_git_parent_path").Return(false)

	//test
	testScm, cerr := scm.Create("github", suite.PipelineData, suite.Config, nil)

	//assert
	require.NoError(suite.T(), cerr)
	require.NotNil(suite.T(), testScm)
}

func (suite *ScmTestSuite) TestCreate_Bitbucket() {
	//setup
	suite.Config.EXPECT().IsSet("scm_bitbucket_username").Return(true)
	suite.Config.EXPECT().IsSet("scm_bitbucket_password").Return(true)
	suite.Config.EXPECT().GetString("scm_bitbucket_username").Return("placeholder")
	suite.Config.EXPECT().GetString("scm_bitbucket_password").Return("placeholder")

	//test
	testScm, cerr := scm.Create("bitbucket", suite.PipelineData, suite.Config, nil)

	//assert
	require.NoError(suite.T(), cerr)
	require.NotNil(suite.T(), testScm)
}

// In order for 'go test' to run this suite, we need to create
// a normal test function and pass our suite to suite.Run
func TestFactoryTestSuite(t *testing.T) {
	suite.Run(t, new(ScmTestSuite))
}
<|endoftext|>


<commit_before>package mysql

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"sync"

	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/go-xorm/core"
	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
)

type MysqlExecutor struct {
	datasource *models.DataSource
	engine     *xorm.Engine
	log        log.Logger
}

type engineCacheType struct {
	cache    map[int64]*xorm.Engine
	versions map[int64]int
	sync.Mutex
}

var engineCache = engineCacheType{
	cache:    make(map[int64]*xorm.Engine),
	versions: make(map[int64]int),
}

func init() {
	tsdb.RegisterExecutor("mysql", NewMysqlExecutor)
}

func NewMysqlExecutor(datasource *models.DataSource) (tsdb.Executor, error) {
	executor := &MysqlExecutor{
		datasource: datasource,
		log:        log.New("tsdb.mysql"),
	}

	err := executor.initEngine()
	if err != nil {
		return nil, err
	}

	return executor, nil
}

func (e *MysqlExecutor) initEngine() error {
	engineCache.Lock()
	defer engineCache.Unlock()

	if engine, present := engineCache.cache[e.datasource.Id]; present {
		if version, _ := engineCache.versions[e.datasource.Id]; version == e.datasource.Version {
			e.engine = engine
			return nil
		}
	}

	cnnstr := fmt.Sprintf("%s:%s@%s(%s)/%s?charset=utf8mb4&parseTime=true&loc=UTC", e.datasource.User, e.datasource.Password, "tcp", e.datasource.Url, e.datasource.Database)
	e.log.Debug("getEngine", "connection", cnnstr)

	engine, err := xorm.NewEngine("mysql", cnnstr)
	engine.SetMaxOpenConns(10)
	engine.SetMaxIdleConns(10)
	if err != nil {
		return err
	}

	engineCache.cache[e.datasource.Id] = engine
	e.engine = engine
	return nil
}

func (e *MysqlExecutor) Execute(ctx context.Context, queries tsdb.QuerySlice, context *tsdb.QueryContext) *tsdb.BatchResult {
	result := &tsdb.BatchResult{
		QueryResults: make(map[string]*tsdb.QueryResult),
	}

	macroEngine := NewMysqlMacroEngine(context.TimeRange)
	session := e.engine.NewSession()
	defer session.Close()
	db := session.DB()

	for _, query := range queries {
		rawSql := query.Model.Get("rawSql").MustString()
		if rawSql == "" {
			continue
		}

		queryResult := &tsdb.QueryResult{Meta: simplejson.New(), RefId: query.RefId}
		result.QueryResults[query.RefId] = queryResult

		rawSql, err := macroEngine.Interpolate(rawSql)
		if err != nil {
			queryResult.Error = err
			continue
		}

		queryResult.Meta.Set("sql", rawSql)

		rows, err := db.Query(rawSql)
		if err != nil {
			queryResult.Error = err
			continue
		}

		defer rows.Close()

		format := query.Model.Get("format").MustString("time_series")

		switch format {
		case "time_series":
			err := e.TransformToTimeSeries(query, rows, queryResult)
			if err != nil {
				queryResult.Error = err
				continue
			}
		case "table":
			err := e.TransformToTable(query, rows, queryResult)
			if err != nil {
				queryResult.Error = err
				continue
			}
		}
	}

	return result
}

func (e MysqlExecutor) TransformToTable(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	columnNames, err := rows.Columns()
	columnCount := len(columnNames)

	if err != nil {
		return err
	}

	table := &tsdb.Table{
		Columns: make([]tsdb.TableColumn, columnCount),
		Rows:    make([]tsdb.RowValues, 0),
	}

	for i, name := range columnNames {
		table.Columns[i].Text = name
	}

	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return err
	}

	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount += 1 {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		values, err := e.getTypedRowData(columnTypes, rows)
		if err != nil {
			return err
		}

		table.Rows = append(table.Rows, values)
	}

	result.Tables = append(result.Tables, table)
	result.Meta.Set("rowCount", rowCount)
	return nil
}

func (e MysqlExecutor) getTypedRowData(types []*sql.ColumnType, rows *core.Rows) (tsdb.RowValues, error) {
	values := make([]interface{}, len(types))

	for i, stype := range types {
		switch stype.DatabaseTypeName() {
		case mysql.FieldTypeNameTiny:
			values[i] = new(int8)
		case mysql.FieldTypeNameInt24:
			values[i] = new(int32)
		case mysql.FieldTypeNameShort:
			values[i] = new(int16)
		case mysql.FieldTypeNameVarString:
			values[i] = new(string)
		case mysql.FieldTypeNameVarChar:
			values[i] = new(string)
		case mysql.FieldTypeNameLongLong:
			values[i] = new(int64)
		case mysql.FieldTypeNameDouble:
			values[i] = new(float64)
		case mysql.FieldTypeNameDecimal:
			values[i] = new(float32)
		case mysql.FieldTypeNameNewDecimal:
			values[i] = new(float64)
		case mysql.FieldTypeNameTimestamp:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameDateTime:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameTime:
			values[i] = new(time.Duration)
		case mysql.FieldTypeNameYear:
			values[i] = new(int16)
		case mysql.FieldTypeNameNULL:
			values[i] = nil
		default:
			return nil, fmt.Errorf("Database type %s not supported", stype.DatabaseTypeName())
		}
	}

	if err := rows.Scan(values...); err != nil {
		return nil, err
	}

	return values, nil
}

func (e MysqlExecutor) TransformToTimeSeries(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	pointsBySeries := make(map[string]*tsdb.TimeSeries)
	columnNames, err := rows.Columns()

	if err != nil {
		return err
	}

	rowData := NewStringStringScan(columnNames)
	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount += 1 {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		err := rowData.Update(rows.Rows)
		if err != nil {
			e.log.Error("MySQL response parsing", "error", err)
			return fmt.Errorf("MySQL response parsing error %v", err)
		}

		if rowData.metric == "" {
			rowData.metric = "Unknown"
		}

		//e.log.Debug("Rows", "metric", rowData.metric, "time", rowData.time, "value", rowData.value)

		if !rowData.time.Valid {
			return fmt.Errorf("Found row with no time value")
		}

		if series, exist := pointsBySeries[rowData.metric]; exist {
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
		} else {
			series := &tsdb.TimeSeries{Name: rowData.metric}
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
			pointsBySeries[rowData.metric] = series
		}
	}

	for _, value := range pointsBySeries {
		result.Series = append(result.Series, value)
	}

	result.Meta.Set("rowCount", rowCount)
	return nil
}

type stringStringScan struct {
	rowPtrs     []interface{}
	rowValues   []string
	columnNames []string
	columnCount int

	time   null.Float
	value  null.Float
	metric string
}

func NewStringStringScan(columnNames []string) *stringStringScan {
	s := &stringStringScan{
		columnCount: len(columnNames),
		columnNames: columnNames,
		rowPtrs:     make([]interface{}, len(columnNames)),
		rowValues:   make([]string, len(columnNames)),
	}

	for i := 0; i < s.columnCount; i++ {
		s.rowPtrs[i] = new(sql.RawBytes)
	}

	return s
}

func (s *stringStringScan) Update(rows *sql.Rows) error {
	if err := rows.Scan(s.rowPtrs...); err != nil {
		return err
	}

	for i := 0; i < s.columnCount; i++ {
		if rb, ok := s.rowPtrs[i].(*sql.RawBytes); ok {
			s.rowValues[i] = string(*rb)

			switch s.columnNames[i] {
			case "time_sec":
				if sec, err := strconv.ParseInt(s.rowValues[i], 10, 64); err == nil {
					s.time = null.FloatFrom(float64(sec * 1000))
				}
			case "value":
				if value, err := strconv.ParseFloat(s.rowValues[i], 64); err == nil {
					s.value = null.FloatFrom(value)
				}
			case "metric":
				s.metric = s.rowValues[i]
			}

			*rb = nil // reset pointer to discard current value to avoid a bug
		} else {
			return fmt.Errorf("Cannot convert index %d column %s to type *sql.RawBytes", i, s.columnNames[i])
		}
	}
	return nil
}
<commit_msg>mysql: adds mapping for int/long<commit_after>package mysql

import (
	"context"
	"database/sql"
	"fmt"
	"strconv"
	"sync"

	"time"

	"github.com/go-sql-driver/mysql"
	"github.com/go-xorm/core"
	"github.com/go-xorm/xorm"
	"github.com/grafana/grafana/pkg/components/null"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/tsdb"
)

type MysqlExecutor struct {
	datasource *models.DataSource
	engine     *xorm.Engine
	log        log.Logger
}

type engineCacheType struct {
	cache    map[int64]*xorm.Engine
	versions map[int64]int
	sync.Mutex
}

var engineCache = engineCacheType{
	cache:    make(map[int64]*xorm.Engine),
	versions: make(map[int64]int),
}

func init() {
	tsdb.RegisterExecutor("mysql", NewMysqlExecutor)
}

func NewMysqlExecutor(datasource *models.DataSource) (tsdb.Executor, error) {
	executor := &MysqlExecutor{
		datasource: datasource,
		log:        log.New("tsdb.mysql"),
	}

	err := executor.initEngine()
	if err != nil {
		return nil, err
	}

	return executor, nil
}

func (e *MysqlExecutor) initEngine() error {
	engineCache.Lock()
	defer engineCache.Unlock()

	if engine, present := engineCache.cache[e.datasource.Id]; present {
		if version, _ := engineCache.versions[e.datasource.Id]; version == e.datasource.Version {
			e.engine = engine
			return nil
		}
	}

	cnnstr := fmt.Sprintf("%s:%s@%s(%s)/%s?charset=utf8mb4&parseTime=true&loc=UTC", e.datasource.User, e.datasource.Password, "tcp", e.datasource.Url, e.datasource.Database)
	e.log.Debug("getEngine", "connection", cnnstr)

	engine, err := xorm.NewEngine("mysql", cnnstr)
	engine.SetMaxOpenConns(10)
	engine.SetMaxIdleConns(10)
	if err != nil {
		return err
	}

	engineCache.cache[e.datasource.Id] = engine
	e.engine = engine
	return nil
}

func (e *MysqlExecutor) Execute(ctx context.Context, queries tsdb.QuerySlice, context *tsdb.QueryContext) *tsdb.BatchResult {
	result := &tsdb.BatchResult{
		QueryResults: make(map[string]*tsdb.QueryResult),
	}

	macroEngine := NewMysqlMacroEngine(context.TimeRange)
	session := e.engine.NewSession()
	defer session.Close()
	db := session.DB()

	for _, query := range queries {
		rawSql := query.Model.Get("rawSql").MustString()
		if rawSql == "" {
			continue
		}

		queryResult := &tsdb.QueryResult{Meta: simplejson.New(), RefId: query.RefId}
		result.QueryResults[query.RefId] = queryResult

		rawSql, err := macroEngine.Interpolate(rawSql)
		if err != nil {
			queryResult.Error = err
			continue
		}

		queryResult.Meta.Set("sql", rawSql)

		rows, err := db.Query(rawSql)
		if err != nil {
			queryResult.Error = err
			continue
		}

		defer rows.Close()

		format := query.Model.Get("format").MustString("time_series")

		switch format {
		case "time_series":
			err := e.TransformToTimeSeries(query, rows, queryResult)
			if err != nil {
				queryResult.Error = err
				continue
			}
		case "table":
			err := e.TransformToTable(query, rows, queryResult)
			if err != nil {
				queryResult.Error = err
				continue
			}
		}
	}

	return result
}

func (e MysqlExecutor) TransformToTable(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	columnNames, err := rows.Columns()
	columnCount := len(columnNames)

	if err != nil {
		return err
	}

	table := &tsdb.Table{
		Columns: make([]tsdb.TableColumn, columnCount),
		Rows:    make([]tsdb.RowValues, 0),
	}

	for i, name := range columnNames {
		table.Columns[i].Text = name
	}

	columnTypes, err := rows.ColumnTypes()
	if err != nil {
		return err
	}

	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount += 1 {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		values, err := e.getTypedRowData(columnTypes, rows)
		if err != nil {
			return err
		}

		table.Rows = append(table.Rows, values)
	}

	result.Tables = append(result.Tables, table)
	result.Meta.Set("rowCount", rowCount)
	return nil
}

func (e MysqlExecutor) getTypedRowData(types []*sql.ColumnType, rows *core.Rows) (tsdb.RowValues, error) {
	values := make([]interface{}, len(types))

	for i, stype := range types {
		switch stype.DatabaseTypeName() {
		case mysql.FieldTypeNameTiny:
			values[i] = new(int8)
		case mysql.FieldTypeNameInt24:
			values[i] = new(int32)
		case mysql.FieldTypeNameShort:
			values[i] = new(int16)
		case mysql.FieldTypeNameVarString:
			values[i] = new(string)
		case mysql.FieldTypeNameVarChar:
			values[i] = new(string)
		case mysql.FieldTypeNameLong:
			values[i] = new(int)
		case mysql.FieldTypeNameLongLong:
			values[i] = new(int64)
		case mysql.FieldTypeNameDouble:
			values[i] = new(float64)
		case mysql.FieldTypeNameDecimal:
			values[i] = new(float32)
		case mysql.FieldTypeNameNewDecimal:
			values[i] = new(float64)
		case mysql.FieldTypeNameTimestamp:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameDateTime:
			values[i] = new(time.Time)
		case mysql.FieldTypeNameTime:
			values[i] = new(time.Duration)
		case mysql.FieldTypeNameYear:
			values[i] = new(int16)
		case mysql.FieldTypeNameNULL:
			values[i] = nil
		default:
			return nil, fmt.Errorf("Database type %s not supported", stype.DatabaseTypeName())
		}
	}

	if err := rows.Scan(values...); err != nil {
		return nil, err
	}

	return values, nil
}

func (e MysqlExecutor) TransformToTimeSeries(query *tsdb.Query, rows *core.Rows, result *tsdb.QueryResult) error {
	pointsBySeries := make(map[string]*tsdb.TimeSeries)
	columnNames, err := rows.Columns()

	if err != nil {
		return err
	}

	rowData := NewStringStringScan(columnNames)
	rowLimit := 1000000
	rowCount := 0

	for ; rows.Next(); rowCount += 1 {
		if rowCount > rowLimit {
			return fmt.Errorf("MySQL query row limit exceeded, limit %d", rowLimit)
		}

		err := rowData.Update(rows.Rows)
		if err != nil {
			e.log.Error("MySQL response parsing", "error", err)
			return fmt.Errorf("MySQL response parsing error %v", err)
		}

		if rowData.metric == "" {
			rowData.metric = "Unknown"
		}

		//e.log.Debug("Rows", "metric", rowData.metric, "time", rowData.time, "value", rowData.value)

		if !rowData.time.Valid {
			return fmt.Errorf("Found row with no time value")
		}

		if series, exist := pointsBySeries[rowData.metric]; exist {
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
		} else {
			series := &tsdb.TimeSeries{Name: rowData.metric}
			series.Points = append(series.Points, tsdb.TimePoint{rowData.value, rowData.time})
			pointsBySeries[rowData.metric] = series
		}
	}

	for _, value := range pointsBySeries {
		result.Series = append(result.Series, value)
	}

	result.Meta.Set("rowCount", rowCount)
	return nil
}

type stringStringScan struct {
	rowPtrs     []interface{}
	rowValues   []string
	columnNames []string
	columnCount int

	time   null.Float
	value  null.Float
	metric string
}

func NewStringStringScan(columnNames []string) *stringStringScan {
	s := &stringStringScan{
		columnCount: len(columnNames),
		columnNames: columnNames,
		rowPtrs:     make([]interface{}, len(columnNames)),
		rowValues:   make([]string, len(columnNames)),
	}

	for i := 0; i < s.columnCount; i++ {
		s.rowPtrs[i] = new(sql.RawBytes)
	}

	return s
}

func (s *stringStringScan) Update(rows *sql.Rows) error {
	if err := rows.Scan(s.rowPtrs...); err != nil {
		return err
	}

	for i := 0; i < s.columnCount; i++ {
		if rb, ok := s.rowPtrs[i].(*sql.RawBytes); ok {
			s.rowValues[i] = string(*rb)

			switch s.columnNames[i] {
			case "time_sec":
				if sec, err := strconv.ParseInt(s.rowValues[i], 10, 64); err == nil {
					s.time = null.FloatFrom(float64(sec * 1000))
				}
			case "value":
				if value, err := strconv.ParseFloat(s.rowValues[i], 64); err == nil {
					s.value = null.FloatFrom(value)
				}
			case "metric":
				s.metric = s.rowValues[i]
			}

			*rb = nil // reset pointer to discard current value to avoid a bug
		} else {
			return fmt.Errorf("Cannot convert index %d column %s to type *sql.RawBytes", i, s.columnNames[i])
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package binlog

import (
	"bytes"
	"encoding/binary"
	"io"
	"sync"

	mysql_proto "github.com/dropbox/godropbox/proto/mysql"
)

type MockLogFile struct {
	logBuffer []byte

	mu sync.Mutex
}

type MockLogFileReader struct {
	file   *MockLogFile
	cursor int
}

// Statically verify that MockLogFile implements io.Reader.

var _ io.Reader = &MockLogFileReader{}

func NewMockLogFile() *MockLogFile {
	return &MockLogFile{
		logBuffer: make([]byte, 0),
	}
}

func (mlf *MockLogFile) GetReader() *MockLogFileReader {
	return newMockLogFileReader(mlf)
}

func (mlf *MockLogFile) Write(contents []byte) {
	mlf.logBuffer = append(mlf.logBuffer, contents...)
}

func (mlf *MockLogFile) WriteLogFileMagic() {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	mlf.Write(logFileMagic)
}

func (mlf *MockLogFile) WriteFDE() {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := []byte{
		// binlog version
		4, 0,
		// server version
		53, 46, 53, 46, 51, 52, 45, 51, 50, 46,
		48, 45, 108, 111, 103, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		// created timestamp
		0, 0, 0, 0,
		// total header size
		19,
		// fixed length data size per event type
		56, 13, 0, 8, 0, 18, 0, 4, 4, 4, 4, 18, 0, 0, 84, 0, 4,
		26, 8, 0, 0, 0, 8, 8, 8, 2, 0}

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + len(data)

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_FORMAT_DESCRIPTION_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data)
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteXid(id uint64) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	binary.Write(data, binary.LittleEndian, id)

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_XID_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteRotate(prefix string, num int) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	pos := uint64(4)

	data := &bytes.Buffer{}
	binary.Write(data, binary.LittleEndian, pos)
	data.WriteString(logName(prefix, num))

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_ROTATE_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteQueryWithParam(query string, dbName string) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	data.Write([]byte{
		0, 0, 0, 0, // thread id
		0, 0, 0, 0, // execute duration
		byte(len(dbName)), // db name length
		0, 0,              // error code
		0, 0, // status block size
	})
	data.WriteString(dbName)
	data.WriteByte(0)
	data.Write([]byte(query))

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_QUERY_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteQuery(query string) {
	mlf.WriteQueryWithParam(query, "db")
}

func (mlf *MockLogFile) WriteBegin() {
	mlf.WriteQuery("BEGIN")
}

func (mlf *MockLogFile) WriteRowsQuery(query string) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	data.WriteByte(byte(len(query)))
	data.WriteString(query) // Note: this mimics bug in mysql 5.6

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_ROWS_QUERY_LOG_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteTableMapWithParams(
	tableId int8,
	dbName string,
	tableName string) {

	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	buf := &bytes.Buffer{}
	buf.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		1, 0,
	})

	buf.WriteByte(byte(len(dbName)))
	buf.Write([]byte(dbName))
	buf.WriteByte(0)

	buf.WriteByte(byte(len(tableName)))
	buf.Write([]byte(tableName))
	buf.WriteByte(0)

	buf.Write([]byte{
		// number of columns
		1,
		// a single long fields
		3,
		// metadata size
		0,
		// null bits
		2,
	})

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + buf.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_TABLE_MAP_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		buf.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteTableMap() {
	mlf.WriteTableMapWithParams(
		0,
		"abc",
		"foo")
}

func (mlf *MockLogFile) WriteInsertWithParam(value int, tableId int8) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap
		1,
		// row data's "is null" bit map
		0,
	})
	binary.Write(data, binary.LittleEndian, uint32(value))

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_WRITE_ROWS_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteInsert(value int) {
	mlf.WriteInsertWithParam(value, 0)
}

func (mlf *MockLogFile) WriteDeleteWithParam(value int, tableId int8) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap
		1,
		// row data's "is null" bit map
		0,
	})
	binary.Write(data, binary.LittleEndian, uint32(value))

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_DELETE_ROWS_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteDelete(value int) {
	mlf.WriteDeleteWithParam(value, 0)
}

func (mlf *MockLogFile) WriteUpdateWithParam(
	before int,
	after int,
	tableId int8) {

	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap (before image)
		1,
		// columns used bitmap (after image)
		1,
	})
	// row data's "is null" bit map
	data.WriteByte(0)
	binary.Write(data, binary.LittleEndian, uint32(before))
	// row data's "is null" bit map
	data.WriteByte(0)
	binary.Write(data, binary.LittleEndian, uint32(after))

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + data.Len()

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(mysql_proto.LogEventType_UPDATE_ROWS_EVENT),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		data.Bytes())
	mlf.Write(e)
}

func (mlf *MockLogFile) WriteUpdate(before int, after int) {
	mlf.WriteUpdateWithParam(before, after, 0)
}

func (mlf *MockLogFile) Copy() *MockLogFile {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	bufferCopy := make([]byte, len(mlf.logBuffer))
	copied := copy(bufferCopy, mlf.logBuffer)
	if copied != len(bufferCopy) {
		panic("MockLogFile buffer copy failed")
	}

	return &MockLogFile{
		logBuffer: bufferCopy,
	}
}

func newMockLogFileReader(logFile *MockLogFile) *MockLogFileReader {
	return &MockLogFileReader{
		file:   logFile,
		cursor: 0,
	}
}

func (reader *MockLogFileReader) Read(p []byte) (n int, err error) {
	reader.file.mu.Lock()
	defer reader.file.mu.Unlock()

	if reader.cursor+len(p) > len(reader.file.logBuffer) {
		// We can't read anything.
		return 0, io.EOF
	}

	// We can read something.
	copied := copy(p, reader.file.logBuffer[reader.cursor:])
	if copied != len(p) {
		panic("MockLogFileReader read failed")
	}
	reader.cursor += len(p)
	return len(p), nil
}
<commit_msg>Simplified MockLogFile implementation.<commit_after>package binlog

import (
	"bytes"
	"encoding/binary"
	"io"
	"sync"

	mysql_proto "github.com/dropbox/godropbox/proto/mysql"
)

type MockLogFile struct {
	logBuffer []byte

	mu sync.Mutex
}

type MockLogFileReader struct {
	file   *MockLogFile
	cursor int
}

// Statically verify that MockLogFile implements io.Reader.

var _ io.Reader = &MockLogFileReader{}

func NewMockLogFile() *MockLogFile {
	return &MockLogFile{
		logBuffer: make([]byte, 0),
	}
}

func (mlf *MockLogFile) GetReader() *MockLogFileReader {
	return newMockLogFileReader(mlf)
}

// Every function for writing into the MockLogFile should acquire the lock via either Write() or writeWithHeader().
func (mlf *MockLogFile) Write(contents []byte) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	mlf.logBuffer = append(mlf.logBuffer, contents...)
}

func (mlf *MockLogFile) writeWithHeader(contents []byte, logEventType mysql_proto.LogEventType_Type) {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	nextPosition := len(mlf.logBuffer) + sizeOfBasicV4EventHeader + len(contents)

	e, _ := CreateEventBytes(
		uint32(0),
		uint8(logEventType),
		uint32(1),
		uint32(nextPosition),
		uint16(1),
		contents)
	mlf.logBuffer = append(mlf.logBuffer, e...)
}

func (mlf *MockLogFile) WriteLogFileMagic() {
	mlf.Write(logFileMagic)
}

func (mlf *MockLogFile) WriteFDE() {
	data := []byte{
		// binlog version
		4, 0,
		// server version
		53, 46, 53, 46, 51, 52, 45, 51, 50, 46,
		48, 45, 108, 111, 103, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		// created timestamp
		0, 0, 0, 0,
		// total header size
		19,
		// fixed length data size per event type
		56, 13, 0, 8, 0, 18, 0, 4, 4, 4, 4, 18, 0, 0, 84, 0, 4,
		26, 8, 0, 0, 0, 8, 8, 8, 2, 0}

	mlf.writeWithHeader(data, mysql_proto.LogEventType_FORMAT_DESCRIPTION_EVENT)
}

func (mlf *MockLogFile) WriteXid(id uint64) {
	data := &bytes.Buffer{}
	binary.Write(data, binary.LittleEndian, id)

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_XID_EVENT)
}

func (mlf *MockLogFile) WriteRotate(prefix string, num int) {
	pos := uint64(4)

	data := &bytes.Buffer{}
	binary.Write(data, binary.LittleEndian, pos)
	data.WriteString(logName(prefix, num))

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_ROTATE_EVENT)
}

func (mlf *MockLogFile) WriteQueryWithParam(query string, dbName string) {
	data := &bytes.Buffer{}
	data.Write([]byte{
		0, 0, 0, 0, // thread id
		0, 0, 0, 0, // execute duration
		byte(len(dbName)), // db name length
		0, 0,              // error code
		0, 0, // status block size
	})
	data.WriteString(dbName)
	data.WriteByte(0)
	data.Write([]byte(query))

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_QUERY_EVENT)
}

func (mlf *MockLogFile) WriteQuery(query string) {
	mlf.WriteQueryWithParam(query, "db")
}

func (mlf *MockLogFile) WriteBegin() {
	mlf.WriteQuery("BEGIN")
}

func (mlf *MockLogFile) WriteRowsQuery(query string) {
	data := &bytes.Buffer{}
	data.WriteByte(byte(len(query)))
	data.WriteString(query) // Note: this mimics bug in mysql 5.6

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_ROWS_QUERY_LOG_EVENT)
}

func (mlf *MockLogFile) WriteTableMapWithParams(
	tableId int8,
	dbName string,
	tableName string) {

	buf := &bytes.Buffer{}
	buf.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		1, 0,
	})

	buf.WriteByte(byte(len(dbName)))
	buf.Write([]byte(dbName))
	buf.WriteByte(0)

	buf.WriteByte(byte(len(tableName)))
	buf.Write([]byte(tableName))
	buf.WriteByte(0)

	buf.Write([]byte{
		// number of columns
		1,
		// a single long fields
		3,
		// metadata size
		0,
		// null bits
		2,
	})

	mlf.writeWithHeader(buf.Bytes(), mysql_proto.LogEventType_TABLE_MAP_EVENT)
}

func (mlf *MockLogFile) WriteTableMap() {
	mlf.WriteTableMapWithParams(
		0,
		"abc",
		"foo")
}

func (mlf *MockLogFile) WriteInsertWithParam(value int, tableId int8) {
	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap
		1,
		// row data's "is null" bit map
		0,
	})
	binary.Write(data, binary.LittleEndian, uint32(value))

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_WRITE_ROWS_EVENT)
}

func (mlf *MockLogFile) WriteInsert(value int) {
	mlf.WriteInsertWithParam(value, 0)
}

func (mlf *MockLogFile) WriteDeleteWithParam(value int, tableId int8) {
	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap
		1,
		// row data's "is null" bit map
		0,
	})
	binary.Write(data, binary.LittleEndian, uint32(value))

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_DELETE_ROWS_EVENT)
}

func (mlf *MockLogFile) WriteDelete(value int) {
	mlf.WriteDeleteWithParam(value, 0)
}

func (mlf *MockLogFile) WriteUpdateWithParam(
	before int,
	after int,
	tableId int8) {

	data := &bytes.Buffer{}
	data.Write([]byte{
		// table id
		byte(tableId), 0, 0, 0, 0, 0,
		// flags
		0, 0,
		// empty variable size header
		2, 0,
		// number of columns
		1,
		// columns used bitmap (before image)
		1,
		// columns used bitmap (after image)
		1,
	})
	// row data's "is null" bit map
	data.WriteByte(0)
	binary.Write(data, binary.LittleEndian, uint32(before))
	// row data's "is null" bit map
	data.WriteByte(0)
	binary.Write(data, binary.LittleEndian, uint32(after))

	mlf.writeWithHeader(data.Bytes(), mysql_proto.LogEventType_UPDATE_ROWS_EVENT)
}

func (mlf *MockLogFile) WriteUpdate(before int, after int) {
	mlf.WriteUpdateWithParam(before, after, 0)
}

func (mlf *MockLogFile) Copy() *MockLogFile {
	mlf.mu.Lock()
	defer mlf.mu.Unlock()

	bufferCopy := make([]byte, len(mlf.logBuffer))
	copied := copy(bufferCopy, mlf.logBuffer)
	if copied != len(bufferCopy) {
		panic("MockLogFile buffer copy failed")
	}

	return &MockLogFile{
		logBuffer: bufferCopy,
	}
}

func newMockLogFileReader(logFile *MockLogFile) *MockLogFileReader {
	return &MockLogFileReader{
		file:   logFile,
		cursor: 0,
	}
}

func (reader *MockLogFileReader) Read(p []byte) (n int, err error) {
	reader.file.mu.Lock()
	defer reader.file.mu.Unlock()

	if reader.cursor+len(p) > len(reader.file.logBuffer) {
		// We can't read anything.
		return 0, io.EOF
	}

	// We can read something.
	copied := copy(p, reader.file.logBuffer[reader.cursor:])
	if copied != len(p) {
		panic("MockLogFileReader read failed")
	}
	reader.cursor += len(p)
	return len(p), nil
}
<|endoftext|>


<commit_before>package db

import (
	"time"

	"github.com/Lunchr/luncher-api/db/model"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type RegistrationAccessTokens interface {
	Insert(*model.RegistrationAccessToken) (*model.RegistrationAccessToken, error)
	Exists(model.Token) (bool, error)
}

type registrationAccessTokensCollection struct {
	*mgo.Collection
}

func NewRegistrationAccessTokens(c *Client) (RegistrationAccessTokens, error) {
	collection := c.database.C(model.RegistrationAccessTokenCollectionName)
	tokens := &registrationAccessTokensCollection{collection}
	if err := tokens.ensureTTLIndex(); err != nil {
		return nil, err
	}
	return tokens, nil
}

func (c registrationAccessTokensCollection) Insert(t *model.RegistrationAccessToken) (*model.RegistrationAccessToken, error) {
	if t.ID == "" {
		// TODO copy before changing perhaps. Mutating incoming pointed objects can be bad
		t.ID = bson.NewObjectId()
	}
	return t, c.Collection.Insert(t)
}

func (c registrationAccessTokensCollection) Exists(token model.Token) (bool, error) {
	count, err := c.Find(bson.M{
		"token": token,
	}).Count()
	return count > 0, err
}

func (c registrationAccessTokensCollection) ensureTTLIndex() error {
	return c.EnsureIndex(mgo.Index{
		Key:         []string{"created_at"},
		ExpireAfter: time.Hour * 24 * 7, // 7 days
	})
}
<commit_msg>Extend access token TTL to 30 days<commit_after>package db

import (
	"time"

	"github.com/Lunchr/luncher-api/db/model"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type RegistrationAccessTokens interface {
	Insert(*model.RegistrationAccessToken) (*model.RegistrationAccessToken, error)
	Exists(model.Token) (bool, error)
}

type registrationAccessTokensCollection struct {
	*mgo.Collection
}

func NewRegistrationAccessTokens(c *Client) (RegistrationAccessTokens, error) {
	collection := c.database.C(model.RegistrationAccessTokenCollectionName)
	tokens := &registrationAccessTokensCollection{collection}
	if err := tokens.ensureTTLIndex(); err != nil {
		return nil, err
	}
	return tokens, nil
}

func (c registrationAccessTokensCollection) Insert(t *model.RegistrationAccessToken) (*model.RegistrationAccessToken, error) {
	if t.ID == "" {
		// TODO copy before changing perhaps. Mutating incoming pointed objects can be bad
		t.ID = bson.NewObjectId()
	}
	return t, c.Collection.Insert(t)
}

func (c registrationAccessTokensCollection) Exists(token model.Token) (bool, error) {
	count, err := c.Find(bson.M{
		"token": token,
	}).Count()
	return count > 0, err
}

func (c registrationAccessTokensCollection) ensureTTLIndex() error {
	return c.EnsureIndex(mgo.Index{
		Key:         []string{"created_at"},
		ExpireAfter: time.Hour * 24 * 30, // 30 days
	})
}
<|endoftext|>


<commit_before>package logger

import (
	"io"
	stdlog "log"
	"os"
	"path/filepath"

	"github.com/prasannavl/go-grab/log"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
)

type Options struct {
	VerbosityLevel   int
	LogFile          string
	FallbackFileName string
	FallbackDir      string

	Rolling         bool
	MaxSize         int // megabytes
	MaxBackups      int
	MaxAge          int // days
	CompressBackups bool
	NoColor         bool
	Humanize        bool
}

func DefaultOptions() Options {
	return Options{
		VerbosityLevel:   0,
		LogFile:          "",
		FallbackFileName: "run.log",
		FallbackDir:      "logs",
		Rolling:          true,
		MaxSize:          100,
		MaxBackups:       2,
		MaxAge:           28,
		CompressBackups:  true,
		NoColor:          false,
		Humanize:         true,
	}
}

const (
	StdOut   = ":stdout"
	StdErr   = ":stderr"
	Disabled = ":null"
)

func Init(opts *Options) {
	logFile := opts.LogFile
	if logFile == Disabled {
		return
	}
	level := logLevelFromVerbosityLevel(opts.VerbosityLevel)
	if level == 0 {
		return
	}
	s := createWriteStream(opts)
	var formatter func(r *log.Record) string

	if opts.Humanize {
		if opts.NoColor {
			formatter = log.DefaultTextFormatterForHuman
		} else {
			formatter = log.DefaultColorTextFormatterForHuman
		}
	} else {
		formatter = log.DefaultTextFormatter
	}

	target := log.StreamRecorder{
		Formatter: formatter,
		Stream:    s,
	}

	rec := log.LeveledRecorder{
		MaxLevel: log.InfoLevel,
		Target:   &target,
	}

	log.SetGlobal(&rec)
}

func logLevelFromVerbosityLevel(vLevel int) log.Level {
	switch vLevel {
	case -1:
		return log.ErrorLevel
	case 0:
		return log.WarnLevel
	case 1:
		return log.InfoLevel
	case 2:
		return log.DebugLevel
	case 3:
		return log.TraceLevel
	}
	return log.TraceLevel
}

func createWriteStream(opts *Options) io.Writer {
	var err error
	logFile := opts.LogFile
	const loggerErrFormat = "error: logger => %s"
	if logFile == "" {
		if logFile, err = touchFile(opts.FallbackDir, opts.FallbackFileName); err != nil {
			stdlog.Fatalf(loggerErrFormat, err.Error())
		}
	}
	switch logFile {
	case StdOut:
		return os.Stdout
	case StdErr:
		return os.Stderr
	default:
		if err := touchFilePath(logFile); err != nil {
			stdlog.Fatalf(loggerErrFormat, err.Error())
		}
		if !opts.Rolling {
			fd, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE, os.FileMode(0644))
			if err != nil {
				stdlog.Fatalf(loggerErrFormat, err.Error())
			}
			return fd
		}
		return &lumberjack.Logger{
			Filename:   logFile,
			MaxSize:    opts.MaxSize,
			MaxBackups: opts.MaxBackups,
			MaxAge:     opts.MaxAge,
			Compress:   opts.CompressBackups,
		}
	}
}

func touchFilePath(path string) error {
	var err error
	a, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	d := filepath.Dir(a)
	err = os.MkdirAll(d, os.FileMode(0777))
	if err != nil {
		return err
	}
	return nil
}

func touchFile(dir string, filename string) (string, error) {
	var err error
	d := filepath.Clean(dir)
	f := d + "/" + filepath.Clean(filename)
	err = os.MkdirAll(d, os.FileMode(0777))
	if err != nil {
		return "", err
	}
	fd, err := os.OpenFile(f, os.O_CREATE, os.FileMode(0644))
	if err != nil {
		return "", err
	}
	if err = fd.Close(); err != nil {
		return "", err
	}
	return f, nil
}
<commit_msg>change: better log config defaults<commit_after>package logger

import (
	"io"
	stdlog "log"
	"os"
	"path/filepath"

	"github.com/prasannavl/go-grab/log"
	lumberjack "gopkg.in/natefinch/lumberjack.v2"
)

type Options struct {
	VerbosityLevel   int
	LogFile          string
	FallbackFileName string
	FallbackDir      string

	Rolling         bool
	MaxSize         int // megabytes
	MaxBackups      int
	MaxAge          int // days
	CompressBackups bool
	NoColor         bool
	Humanize        bool
}

func DefaultOptions() Options {
	return Options{
		VerbosityLevel:   0,
		LogFile:          TargetStdOut,
		FallbackFileName: "run.log",
		FallbackDir:      "logs",
		Rolling:          true,
		MaxSize:          100,
		MaxBackups:       2,
		MaxAge:           28,
		CompressBackups:  true,
		NoColor:          false,
		Humanize:         true,
	}
}

const (
	TargetStdOut = ":stdout"
	TargetStdErr = ":stderr"
	TargetNull   = ":null"
)

func Init(opts *Options) {
	logFile := opts.LogFile
	if logFile == TargetNull {
		return
	}
	level := logLevelFromVerbosityLevel(opts.VerbosityLevel)
	if level == 0 {
		return
	}
	s := createWriteStream(opts)
	var formatter func(r *log.Record) string

	if opts.Humanize {
		if opts.NoColor {
			formatter = log.DefaultTextFormatterForHuman
		} else {
			formatter = log.DefaultColorTextFormatterForHuman
		}
	} else {
		formatter = log.DefaultTextFormatter
	}

	target := log.StreamRecorder{
		Formatter: formatter,
		Stream:    s,
	}

	rec := log.LeveledRecorder{
		MaxLevel: log.InfoLevel,
		Target:   &target,
	}

	log.SetGlobal(&rec)
}

func logLevelFromVerbosityLevel(vLevel int) log.Level {
	switch vLevel {
	case -1:
		return log.ErrorLevel
	case 0:
		return log.WarnLevel
	case 1:
		return log.InfoLevel
	case 2:
		return log.DebugLevel
	case 3:
		return log.TraceLevel
	}
	return log.TraceLevel
}

func createWriteStream(opts *Options) io.Writer {
	var err error
	logFile := opts.LogFile
	const loggerErrFormat = "error: logger => %s"
	if logFile == "" {
		if logFile, err = touchFile(opts.FallbackDir, opts.FallbackFileName); err != nil {
			stdlog.Fatalf(loggerErrFormat, err.Error())
		}
	}
	switch logFile {
	case TargetStdOut:
		return os.Stdout
	case TargetStdErr:
		return os.Stderr
	default:
		if err := touchFilePath(logFile); err != nil {
			stdlog.Fatalf(loggerErrFormat, err.Error())
		}
		if !opts.Rolling {
			fd, err := os.OpenFile(logFile, os.O_APPEND|os.O_CREATE, os.FileMode(0644))
			if err != nil {
				stdlog.Fatalf(loggerErrFormat, err.Error())
			}
			return fd
		}
		return &lumberjack.Logger{
			Filename:   logFile,
			MaxSize:    opts.MaxSize,
			MaxBackups: opts.MaxBackups,
			MaxAge:     opts.MaxAge,
			Compress:   opts.CompressBackups,
		}
	}
}

func touchFilePath(path string) error {
	var err error
	a, err := filepath.Abs(path)
	if err != nil {
		return err
	}
	d := filepath.Dir(a)
	err = os.MkdirAll(d, os.FileMode(0777))
	if err != nil {
		return err
	}
	return nil
}

func touchFile(dir string, filename string) (string, error) {
	var err error
	d := filepath.Clean(dir)
	f := d + "/" + filepath.Clean(filename)
	err = os.MkdirAll(d, os.FileMode(0777))
	if err != nil {
		return "", err
	}
	fd, err := os.OpenFile(f, os.O_CREATE, os.FileMode(0644))
	if err != nil {
		return "", err
	}
	if err = fd.Close(); err != nil {
		return "", err
	}
	return f, nil
}
<|endoftext|>


<commit_before>package fluent

import (
	"github.com/Sirupsen/logrus"
	"github.com/evalphobia/logrus_fluent"
)

// default logging hook level
var hookLevel = []logrus.Level{
	logrus.PanicLevel,
	logrus.FatalLevel,
	logrus.ErrorLevel,
	logrus.WarnLevel,
	logrus.InfoLevel,
}

func SetLevels(levels []logrus.Level) {
	hookLevel = levels
}

func AddLevel(level logrus.Level) {
	hookLevel = append(hookLevel, level)
}

func Set(host string, port int) {
	hook := logrus_fluent.NewHook(host, port)
	hook.SetLevels(hookLevel)
	logrus.AddHook(hook)
}
<commit_msg>Update fluent hook to use filter<commit_after>package fluent

import (
	"net/http"

	"github.com/Sirupsen/logrus"
	"github.com/evalphobia/logrus_fluent"
)

// default logging hook level
var hookLevel = []logrus.Level{
	logrus.PanicLevel,
	logrus.FatalLevel,
	logrus.ErrorLevel,
	logrus.WarnLevel,
	logrus.InfoLevel,
}

func SetLevels(levels []logrus.Level) {
	hookLevel = levels
}

func AddLevel(level logrus.Level) {
	hookLevel = append(hookLevel, level)
}

func Set(host string, port int) error {
	hook, err := logrus_fluent.New(host, port)
	if err != nil {
		return err
	}

	hook.SetLevels(hookLevel)
	hook.AddIgnore("context")
	hook.AddFilter("http_request", filterRequest)
	hook.AddFilter("error", logrus_fluent.FilterError)

	logrus.AddHook(hook)
	return nil
}

func filterRequest(v interface{}) interface{} {
	req, ok := v.(*http.Request)
	if !ok {
		return v
	}

	return request{
		Method:     req.Method,
		Host:       req.Host,
		RequestURI: req.RequestURI,
		RemoteAddr: req.RemoteAddr,
	}
}

type request struct {
	Method     string
	Host       string
	RequestURI string
	RemoteAddr string
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"io"
	"log"
	"strings"
	"sync"

	"github.com/fsouza/go-dockerclient"
)

type AttachManager struct {
	sync.Mutex
	attached map[string]*LogPump
	channels map[chan *AttachEvent]struct{}
	client   *docker.Client
}

func NewAttachManager(client *docker.Client) *AttachManager {
	m := &AttachManager{
		attached: make(map[string]*LogPump),
		channels: make(map[chan *AttachEvent]struct{}),
		client:   client,
	}
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	assert(err, "attacher")
	for _, listing := range containers {
		m.attach(listing.ID[:12])
	}
	go func() {
		events := make(chan *docker.APIEvents)
		assert(client.AddEventListener(events), "attacher")
		for msg := range events {
			debug("event:", msg.ID[:12], msg.Status)
			if msg.Status == "start" {
				go m.attach(msg.ID[:12])
			}
		}
		log.Fatal("ruh roh") // todo: loop?
	}()
	return m
}

func (m *AttachManager) attach(id string) {
	container, err := m.client.InspectContainer(id)
	assert(err, "attacher")
	name := container.Name[1:]
	success := make(chan struct{})
	failure := make(chan error)
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	go func() {
		err := m.client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    id,
			OutputStream: outwr,
			ErrorStream:  errwr,
			Stdin:        false,
			Stdout:       true,
			Stderr:       true,
			Stream:       true,
			Success:      success,
		})
		outwr.Close()
		errwr.Close()
		debug("attach:", id, "finished")
		if err != nil {
			close(success)
			failure <- err
		}
		m.send(&AttachEvent{Type: "detach", ID: id, Name: name})
		m.Lock()
		delete(m.attached, id)
		m.Unlock()
	}()
	_, ok := <-success
	if ok {
		m.Lock()
		m.attached[id] = NewLogPump(outrd, errrd, id, name)
		m.Unlock()
		success <- struct{}{}
		m.send(&AttachEvent{ID: id, Name: name, Type: "attach"})
		debug("attach:", id, name, "success")
		return
	}
	debug("attach:", id, "failure:", <-failure)
}

func (m *AttachManager) send(event *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	for ch, _ := range m.channels {
		// TODO: log err after timeout and continue
		ch <- event
	}
}

func (m *AttachManager) addListener(ch chan *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	m.channels[ch] = struct{}{}
	go func() {
		for id, pump := range m.attached {
			ch <- &AttachEvent{ID: id, Name: pump.Name, Type: "attach"}
		}
	}()
}

func (m *AttachManager) removeListener(ch chan *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	delete(m.channels, ch)
}

func (m *AttachManager) Get(id string) *LogPump {
	m.Lock()
	defer m.Unlock()
	return m.attached[id]
}

func (m *AttachManager) Listen(source *Source, logstream chan *Log, closer <-chan bool) {
	if source == nil {
		source = new(Source)
	}
	events := make(chan *AttachEvent)
	m.addListener(events)
	defer m.removeListener(events)
	for {
		select {
		case event := <-events:
			if event.Type == "attach" && (source.All() ||
				(source.ID != "" && strings.HasPrefix(event.ID, source.ID)) ||
				(source.Name != "" && event.Name == source.Name) ||
				(source.Filter != "" && strings.Contains(event.Name, source.Filter))) {
				pump := m.Get(event.ID)
				pump.AddListener(logstream)
				defer func() {
					if pump != nil {
						pump.RemoveListener(logstream)
					}
				}()
			} else if source.ID != "" && event.Type == "detach" &&
				strings.HasPrefix(event.ID, source.ID) {
				return
			}
		case <-closer:
			return
		}
	}
}

type LogPump struct {
	sync.Mutex
	ID       string
	Name     string
	channels map[chan *Log]struct{}
}

func NewLogPump(stdout, stderr io.Reader, id, name string) *LogPump {
	obj := &LogPump{
		ID:       id,
		Name:     name,
		channels: make(map[chan *Log]struct{}),
	}
	pump := func(typ string, source io.Reader) {
		buf := bufio.NewReader(source)
		for {
			data, err := buf.ReadBytes('\n')
			if err != nil {
				if err != io.EOF {
					debug("pump:", id, typ+":", err)
				}
				return
			}
			obj.send(&Log{
				Data: strings.TrimSuffix(string(data), "\n"),
				ID:   id,
				Name: name,
				Type: typ,
			})
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return obj
}

func (o *LogPump) send(log *Log) {
	o.Lock()
	defer o.Unlock()
	for ch, _ := range o.channels {
		// TODO: log err after timeout and continue
		ch <- log
	}
}

func (o *LogPump) AddListener(ch chan *Log) {
	o.Lock()
	defer o.Unlock()
	o.channels[ch] = struct{}{}
}

func (o *LogPump) RemoveListener(ch chan *Log) {
	o.Lock()
	defer o.Unlock()
	delete(o.channels, ch)
}
<commit_msg>fix(logspout): fix deadlock when attaching<commit_after>package main

import (
	"bufio"
	"io"
	"log"
	"strings"
	"sync"

	"github.com/fsouza/go-dockerclient"
)

type AttachManager struct {
	sync.Mutex
	attached map[string]*LogPump
	channels map[chan *AttachEvent]struct{}
	client   *docker.Client
}

func NewAttachManager(client *docker.Client) *AttachManager {
	m := &AttachManager{
		attached: make(map[string]*LogPump),
		channels: make(map[chan *AttachEvent]struct{}),
		client:   client,
	}
	containers, err := client.ListContainers(docker.ListContainersOptions{})
	assert(err, "attacher")
	for _, listing := range containers {
		m.attach(listing.ID[:12])
	}
	go func() {
		events := make(chan *docker.APIEvents)
		assert(client.AddEventListener(events), "attacher")
		for msg := range events {
			debug("event:", msg.ID[:12], msg.Status)
			if msg.Status == "start" {
				go m.attach(msg.ID[:12])
			}
		}
		log.Fatal("ruh roh") // todo: loop?
	}()
	return m
}

func (m *AttachManager) attach(id string) {
	container, err := m.client.InspectContainer(id)
	assert(err, "attacher")
	name := container.Name[1:]
	success := make(chan struct{})
	failure := make(chan error)
	outrd, outwr := io.Pipe()
	errrd, errwr := io.Pipe()
	go func() {
		err := m.client.AttachToContainer(docker.AttachToContainerOptions{
			Container:    id,
			OutputStream: outwr,
			ErrorStream:  errwr,
			Stdin:        false,
			Stdout:       true,
			Stderr:       true,
			Stream:       true,
			Success:      success,
		})
		outwr.Close()
		errwr.Close()
		debug("attach:", id, "finished")
		if err != nil {
			close(success)
			failure <- err
		}
		m.send(&AttachEvent{Type: "detach", ID: id, Name: name})
		m.Lock()
		delete(m.attached, id)
		m.Unlock()
	}()
	_, ok := <-success
	if ok {
		m.Lock()
		m.attached[id] = NewLogPump(outrd, errrd, id, name)
		m.Unlock()
		success <- struct{}{}
		m.send(&AttachEvent{ID: id, Name: name, Type: "attach"})
		debug("attach:", id, name, "success")
		return
	}
	debug("attach:", id, "failure:", <-failure)
}

func (m *AttachManager) send(event *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	for ch, _ := range m.channels {
		// TODO: log err after timeout and continue
		ch <- event
	}
}

func (m *AttachManager) addListener(ch chan *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	m.channels[ch] = struct{}{}
	go func() {
		for id, pump := range m.attached {
			ch <- &AttachEvent{ID: id, Name: pump.Name, Type: "attach"}
		}
	}()
}

func (m *AttachManager) removeListener(ch chan *AttachEvent) {
	m.Lock()
	defer m.Unlock()
	delete(m.channels, ch)
}

func (m *AttachManager) Get(id string) *LogPump {
	m.Lock()
	defer m.Unlock()
	return m.attached[id]
}

func (m *AttachManager) Listen(source *Source, logstream chan *Log, closer <-chan bool) {
	if source == nil {
		source = new(Source)
	}
	depth := len(m.attached)
	if depth == 0 {
		depth = 1
	}
	events := make(chan *AttachEvent, depth)
	m.addListener(events)
	defer m.removeListener(events)
	for {
		select {
		case event := <-events:
			if event.Type == "attach" && (source.All() ||
				(source.ID != "" && strings.HasPrefix(event.ID, source.ID)) ||
				(source.Name != "" && event.Name == source.Name) ||
				(source.Filter != "" && strings.Contains(event.Name, source.Filter))) {
				pump := m.Get(event.ID)
				pump.AddListener(logstream)
				defer func() {
					if pump != nil {
						pump.RemoveListener(logstream)
					}
				}()
			} else if source.ID != "" && event.Type == "detach" &&
				strings.HasPrefix(event.ID, source.ID) {
				return
			}
		case <-closer:
			return
		}
	}
}

type LogPump struct {
	sync.Mutex
	ID       string
	Name     string
	channels map[chan *Log]struct{}
}

func NewLogPump(stdout, stderr io.Reader, id, name string) *LogPump {
	obj := &LogPump{
		ID:       id,
		Name:     name,
		channels: make(map[chan *Log]struct{}),
	}
	pump := func(typ string, source io.Reader) {
		buf := bufio.NewReader(source)
		for {
			data, err := buf.ReadBytes('\n')
			if err != nil {
				if err != io.EOF {
					debug("pump:", id, typ+":", err)
				}
				return
			}
			obj.send(&Log{
				Data: strings.TrimSuffix(string(data), "\n"),
				ID:   id,
				Name: name,
				Type: typ,
			})
		}
	}
	go pump("stdout", stdout)
	go pump("stderr", stderr)
	return obj
}

func (o *LogPump) send(log *Log) {
	o.Lock()
	defer o.Unlock()
	for ch, _ := range o.channels {
		// TODO: log err after timeout and continue
		ch <- log
	}
}

func (o *LogPump) AddListener(ch chan *Log) {
	o.Lock()
	defer o.Unlock()
	o.channels[ch] = struct{}{}
}

func (o *LogPump) RemoveListener(ch chan *Log) {
	o.Lock()
	defer o.Unlock()
	delete(o.channels, ch)
}
<|endoftext|>