<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/iana"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "eth0", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 30 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	debug := func(string, ...interface{}) {}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}
	// DHCPv6
	if *useV6 {
		log.Printf("Trying to obtain a DHCPv6 lease on %s", *ifname)
		log.Printf("Waiting for network interface %s to come up", *ifname)
		start := time.Now()
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv6: IfUp failed: %v", err)
		}
		debug("Interface %s is up after %v", *ifname, time.Since(start))
		var (
			netconf  *netboot.NetConf
			bootfile string
		)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			// send a netboot request via DHCP
			modifiers := []dhcpv6.Modifier{
				dhcpv6.WithArchType(iana.EFI_X86_64),
			}
			if *userClass != "" {
				modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
			}
			conversation, err := netboot.RequestNetbootv6(*ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv6: netboot request for interface %s failed: %v", *ifname, err)
			}
			// get network configuration and boot file
			netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
			if err != nil {
				log.Fatalf("DHCPv6: failed to extract network configuration for %s: %v", *ifname, err)
			}
			debug("DHCPv6: network configuration: %+v", netconf)
			if !*dryRun {
				// Set up IP addresses
				log.Printf("DHCPv6: configuring network interface %s", *ifname)
				if err = netboot.ConfigureInterface(*ifname, netconf); err != nil {
					log.Fatalf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", *ifname, err)
				}
				// Set up DNS
			}
			if *overrideNetbootURL != "" {
				bootfile = *overrideNetbootURL
			}
			log.Printf("DHCPv6: boot file for interface %s is %s", *ifname, bootfile)
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		debug("DHCPv6: boot file URL is %s", bootfile)
		// check for supported schemes
		if !strings.HasPrefix(bootfile, "http://") {
			log.Fatal("DHCPv6: can only handle http scheme")
		}

		log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
		resp, err := http.Get(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
		}
		// FIXME this will not be called if something fails after this point
		defer resp.Body.Close()
		if resp.StatusCode != 200 {
			log.Fatalf("Status code is not 200 OK: %d", resp.StatusCode)
		}
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			log.Fatalf("DHCPv6: cannot read boot file from the network: %v", err)
		}
		u, err := url.Parse(bootfile)
		if err != nil {
			log.Fatalf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
		}
		// extract file name component
		if strings.HasSuffix(u.Path, "/") {
			log.Fatalf("Invalid file path, cannot end with '/': %s", u.Path)
		}
		filename := filepath.Base(u.Path)
		if filename == "." || filename == "" {
			log.Fatalf("Invalid empty file name extracted from file path %s", u.Path)
		}
		if err = ioutil.WriteFile(filename, body, 0400); err != nil {
			log.Fatalf("DHCPv6: cannot write to file %s: %v", filename, err)
		}
		debug("DHCPv6: saved boot file to %s", filename)
		if !*dryRun {
			log.Printf("DHCPv6: kexec'ing into %s", filename)
			kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
			if err != nil {
				log.Fatalf("DHCPv6: cannot open file %s: %v", filename, err)
			}
			if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
				log.Fatalf("DHCPv6: kexec.FileLoad failed: %v", err)
			}
			if err = kexec.Reboot(); err != nil {
				log.Fatalf("DHCPv6: kexec.Reboot failed: %v", err)
			}
		}
	}
	// DHCPv4
	if *useV4 {
		log.Printf("Trying to obtain a DHCPv4 lease on %s", *ifname)
		_, err := netboot.IfUp(*ifname, interfaceUpTimeout)
		if err != nil {
			log.Fatalf("DHCPv4: IfUp failed: %v", err)
		}
		debug("DHCPv4: interface %s is up", *ifname)
		if *skipDHCP {
			log.Print("Skipping DHCP")
		} else {
			log.Print("DHCPv4: sending request")
			client := dhcpv4.NewClient()
			// TODO add options to request to netboot
			conversation, err := client.Exchange(*ifname)
			for _, m := range conversation {
				debug(m.Summary())
			}
			if err != nil {
				log.Fatalf("DHCPv4: Exchange failed: %v", err)
			}
			// TODO configure the network and DNS
			// TODO extract the next server and boot file and fetch it
			// TODO kexec into the NBP
		}
	}

}
<commit_msg>RFC: netboot from all interfaces (#73)<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/insomniacslk/dhcp/dhcpv4"
	"github.com/insomniacslk/dhcp/dhcpv6"
	"github.com/insomniacslk/dhcp/iana"
	"github.com/insomniacslk/dhcp/interfaces"
	"github.com/insomniacslk/dhcp/netboot"
	"github.com/u-root/u-root/pkg/kexec"
)

var (
	useV4              = flag.Bool("4", false, "Get a DHCPv4 lease")
	useV6              = flag.Bool("6", true, "Get a DHCPv6 lease")
	ifname             = flag.String("i", "", "Interface to send packets through")
	dryRun             = flag.Bool("dryrun", false, "Do everything except assigning IP addresses, changing DNS, and kexec")
	doDebug            = flag.Bool("d", false, "Print debug output")
	skipDHCP           = flag.Bool("skip-dhcp", false, "Skip DHCP and rely on SLAAC for network configuration. This requires -netboot-url")
	overrideNetbootURL = flag.String("netboot-url", "", "Override the netboot URL normally obtained via DHCP")
	readTimeout        = flag.Int("timeout", 3, "Read timeout in seconds")
	dhcpRetries        = flag.Int("retries", 3, "Number of times a DHCP request is retried")
	userClass          = flag.String("userclass", "", "Override DHCP User Class option")
)

const (
	interfaceUpTimeout = 10 * time.Second
)

var banner = `

 _________________________________
< Net booting is so hot right now >
 ---------------------------------
        \   ^__^
         \  (oo)\_______
            (__)\       )\/\
                ||----w |
                ||     ||

`
var debug = func(string, ...interface{}) {}

func main() {
	flag.Parse()
	if *skipDHCP && *overrideNetbootURL == "" {
		log.Fatal("-skip-dhcp requires -netboot-url")
	}
	if *doDebug {
		debug = log.Printf
	}
	log.Print(banner)

	if !*useV6 && !*useV4 {
		log.Fatal("At least one of DHCPv6 and DHCPv4 is required")
	}

	iflist := []net.Interface{}
	if *ifname != "" {
		var iface *net.Interface
		var err error
		if iface, err = net.InterfaceByName(*ifname); err != nil {
			log.Fatalf("Could not find interface %s: %v", *ifname, err)
		}
		iflist = append(iflist, *iface)
	} else {
		var err error
		if iflist, err = interfaces.GetNonLoopbackInterfaces(); err != nil {
			log.Fatalf("Could not obtain the list of network interfaces: %v", err)
		}
	}

	for _, iface := range iflist {
		if *useV6 {
			if err := boot6(iface.Name); err != nil {
				log.Printf("Could not boot from %s: %v", iface.Name, err)
			}
		}
		if *useV4 {
			if err := boot4(iface.Name); err != nil {
				log.Printf("Could not boot from %s: %v", iface.Name, err)
			}
		}
	}

	log.Fatalln("Could not boot from any interfaces")
}

func boot6(ifname string) error {
	log.Printf("Trying to obtain a DHCPv6 lease on %s", ifname)
	log.Printf("Waiting for network interface %s to come up", ifname)
	start := time.Now()
	_, err := netboot.IfUp(ifname, interfaceUpTimeout)
	if err != nil {
		return fmt.Errorf("DHCPv6: IfUp failed: %v", err)
	}
	debug("Interface %s is up after %v", ifname, time.Since(start))
	var (
		netconf  *netboot.NetConf
		bootfile string
	)
	if *skipDHCP {
		log.Print("Skipping DHCP")
	} else {
		// send a netboot request via DHCP
		modifiers := []dhcpv6.Modifier{
			dhcpv6.WithArchType(iana.EFI_X86_64),
		}
		if *userClass != "" {
			modifiers = append(modifiers, dhcpv6.WithUserClass([]byte(*userClass)))
		}
		conversation, err := netboot.RequestNetbootv6(ifname, time.Duration(*readTimeout)*time.Second, *dhcpRetries, modifiers...)
		for _, m := range conversation {
			debug(m.Summary())
		}
		if err != nil {
			return fmt.Errorf("DHCPv6: netboot request for interface %s failed: %v", ifname, err)
		}
		// get network configuration and boot file
		netconf, bootfile, err = netboot.ConversationToNetconf(conversation)
		if err != nil {
			return fmt.Errorf("DHCPv6: failed to extract network configuration for %s: %v", ifname, err)
		}
		debug("DHCPv6: network configuration: %+v", netconf)
		if !*dryRun {
			// Set up IP addresses
			log.Printf("DHCPv6: configuring network interface %s", ifname)
			if err = netboot.ConfigureInterface(ifname, netconf); err != nil {
				return fmt.Errorf("DHCPv6: cannot configure IPv6 addresses on interface %s: %v", ifname, err)
			}
			// Set up DNS
		}
		if *overrideNetbootURL != "" {
			bootfile = *overrideNetbootURL
		}
		log.Printf("DHCPv6: boot file for interface %s is %s", ifname, bootfile)
	}
	if *overrideNetbootURL != "" {
		bootfile = *overrideNetbootURL
	}
	debug("DHCPv6: boot file URL is %s", bootfile)
	// check for supported schemes
	if !strings.HasPrefix(bootfile, "http://") {
		return fmt.Errorf("DHCPv6: can only handle http scheme")
	}

	log.Printf("DHCPv6: fetching boot file URL: %s", bootfile)
	resp, err := http.Get(bootfile)
	if err != nil {
		return fmt.Errorf("DHCPv6: http.Get of %s failed: %v", bootfile, err)
	}
	// FIXME this will not be called if something fails after this point
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("Status code is not 200 OK: %d", resp.StatusCode)
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("DHCPv6: cannot read boot file from the network: %v", err)
	}
	u, err := url.Parse(bootfile)
	if err != nil {
		return fmt.Errorf("DHCPv6: cannot parse URL %s: %v", bootfile, err)
	}
	// extract file name component
	if strings.HasSuffix(u.Path, "/") {
		return fmt.Errorf("Invalid file path, cannot end with '/': %s", u.Path)
	}
	filename := filepath.Base(u.Path)
	if filename == "." || filename == "" {
		return fmt.Errorf("Invalid empty file name extracted from file path %s", u.Path)
	}
	if err = ioutil.WriteFile(filename, body, 0400); err != nil {
		return fmt.Errorf("DHCPv6: cannot write to file %s: %v", filename, err)
	}
	debug("DHCPv6: saved boot file to %s", filename)
	if !*dryRun {
		log.Printf("DHCPv6: kexec'ing into %s", filename)
		kernel, err := os.OpenFile(filename, os.O_RDONLY, 0)
		if err != nil {
			return fmt.Errorf("DHCPv6: cannot open file %s: %v", filename, err)
		}
		if err = kexec.FileLoad(kernel, nil /* ramfs */, "" /* cmdline */); err != nil {
			return fmt.Errorf("DHCPv6: kexec.FileLoad failed: %v", err)
		}
		if err = kexec.Reboot(); err != nil {
			return fmt.Errorf("DHCPv6: kexec.Reboot failed: %v", err)
		}
	}
	return nil
}

func boot4(ifname string) error {
	log.Printf("Trying to obtain a DHCPv4 lease on %s", ifname)
	_, err := netboot.IfUp(ifname, interfaceUpTimeout)
	if err != nil {
		return fmt.Errorf("DHCPv4: IfUp failed: %v", err)
	}
	debug("DHCPv4: interface %s is up", ifname)
	if *skipDHCP {
		log.Print("Skipping DHCP")
	} else {
		log.Print("DHCPv4: sending request")
		client := dhcpv4.NewClient()
		// TODO add options to request to netboot
		conversation, err := client.Exchange(ifname)
		for _, m := range conversation {
			debug(m.Summary())
		}
		if err != nil {
			return fmt.Errorf("DHCPv4: Exchange failed: %v", err)
		}
		// TODO configure the network and DNS
		// TODO extract the next server and boot file and fetch it
		// TODO kexec into the NBP
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 config

import (
	"crypto/tls"

	"github.com/skydive-project/skydive/common"
)

// GetTLSClientConfig returns TLS config to be used by a server
func GetTLSClientConfig(setupRootCA bool) (*tls.Config, error) {
	certPEM := GetString("agent.X509_cert")
	keyPEM := GetString("agent.X509_key")
	var tlsConfig *tls.Config
	if certPEM != "" && keyPEM != "" {
		var err error
		tlsConfig, err = common.SetupTLSClientConfig(certPEM, keyPEM)
		if err != nil {
			return nil, err
		}
		if setupRootCA {
			analyzerCertPEM := GetString("analyzer.X509_cert")
			tlsConfig.RootCAs, err = common.SetupTLSLoadCertificate(analyzerCertPEM)
			if err != nil {
				return nil, err
			}
		}
	}
	return tlsConfig, nil
}

// GetTLSServerConfig returns TLS config to be used by a server
func GetTLSServerConfig(setupRootCA bool) (*tls.Config, error) {
	certPEM := GetString("analyzer.X509_cert")
	keyPEM := GetString("analyzer.X509_key")
	agentCertPEM := GetString("agent.X509_cert")
	tlsConfig, err := common.SetupTLSServerConfig(certPEM, keyPEM)
	if err != nil {
		return nil, err
	}
	tlsConfig.ClientCAs, err = common.SetupTLSLoadCertificate(agentCertPEM)
	if err != nil {
		return nil, err
	}
	return tlsConfig, nil
}
<commit_msg>Refactor TLS configuration<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 config

import (
	"crypto/tls"

	"github.com/skydive-project/skydive/common"
)

// GetTLSClientConfig returns TLS config to be used by client
func GetTLSClientConfig(setupRootCA bool) (*tls.Config, error) {
	certPEM := GetString("tls.client_cert")
	keyPEM := GetString("tls.client_key")
	var tlsConfig *tls.Config
	if certPEM != "" && keyPEM != "" {
		var err error
		tlsConfig, err = common.SetupTLSClientConfig(certPEM, keyPEM)
		if err != nil {
			return nil, err
		}
		if setupRootCA {
			rootCaPEM := GetString("tls.ca_cert")
			tlsConfig.RootCAs, err = common.SetupTLSLoadCA(rootCaPEM)
			if err != nil {
				return nil, err
			}
		}
	}
	return tlsConfig, nil
}

// GetTLSServerConfig returns TLS config to be used by server
func GetTLSServerConfig(setupRootCA bool) (*tls.Config, error) {
	certPEM := GetString("tls.server_cert")
	keyPEM := GetString("tls.server_key")

	tlsConfig, err := common.SetupTLSServerConfig(certPEM, keyPEM)
	if err != nil {
		return nil, err
	}
	if setupRootCA {
		rootCaPEM := GetString("tls.ca_cert")
		tlsConfig.ClientCAs, err = common.SetupTLSLoadCA(rootCaPEM)
		if err != nil {
			return nil, err
		}
	}
	return tlsConfig, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	ini "gopkg.in/ini.v1"
)

type ConfigClient interface {
	GetAll() (map[string]string, error)
	Unset(system bool, key string) error
	Set(system bool, key, value string) error
	Get(key string) (string, error)
}

type ConfigGetter interface {
	Get(key string) (string, error)
}

type RealConfigClient struct {
	systemConfig string
	userConfig   string
}

func (rcc *RealConfigClient) Set(system bool, key, value string) error {
	var configPath string

	if system {
		configPath = rcc.systemConfig
	} else {
		configPath = rcc.userConfig
	}

	cfg, err := ini.LooseLoad(configPath)
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	section, k := splitKey(key)

	_, err = cfg.Section(section).NewKey(k, value)
	if err != nil {
		return fmt.Errorf("unable to set new key: %v", err)
	}

	err = cfg.SaveToIndent(configPath, "    ")
	if err != nil {
		return fmt.Errorf("unable to save: %v", err)
	}

	return nil
}

func (rcc *RealConfigClient) Unset(system bool, key string) error {
	var configPath string

	if system {
		configPath = rcc.systemConfig
	} else {
		configPath = rcc.userConfig
	}

	cfg, err := ini.LooseLoad(configPath)
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	section, k := splitKey(key)

	cfg.Section(section).DeleteKey(k)
	if len(cfg.Section(section).Keys()) == 0 {
		cfg.DeleteSection(section)
	}

	err = cfg.SaveToIndent(configPath, "    ")
	if err != nil {
		return fmt.Errorf("unable to save: %v", err)
	}

	return nil
}

func (rcc *RealConfigClient) Get(key string) (string, error) {
	cfg, err := ini.LooseLoad(rcc.systemConfig, rcc.userConfig)
	if err != nil {
		return "", fmt.Errorf("failure to load config: %v", err)
	}

	section, k := splitKey(key)
	ck := cfg.Section(section).Key(k)

	return ck.Value(), nil
}

func (rcc *RealConfigClient) GetAll() (map[string]string, error) {
	all := make(map[string]string)

	cfg, err := ini.LooseLoad(rcc.systemConfig, rcc.userConfig)
	if err != nil {
		return all, fmt.Errorf("failure to load config: %v", err)
	}

	for _, section := range cfg.Sections() {
		if len(section.Keys()) > 0 {
			for _, key := range section.Keys() {
				all[fmt.Sprintf("%s.%s", section.Name(), key.Name())] = key.Value()
			}
		}
	}

	return all, nil
}

func splitKey(key string) (string, string) {
	parts := strings.SplitN(key, ".", 2)

	return parts[0], parts[1]
}

func NewDefaultConfigClient(system System) (*RealConfigClient, error) {
	var baseDir string
	if xdgConfigHome := system.Getenv("XDG_CONFIG_HOME"); len(xdgConfigHome) > 0 {
		baseDir = filepath.Join(xdgConfigHome, "holen")
	} else {
		var home string
		if home = system.Getenv("HOME"); len(home) == 0 {
			return nil, fmt.Errorf("$HOME environment variable not found")
		}
		baseDir = filepath.Join(home, ".config", "holen")
		os.MkdirAll(baseDir, 0755)
	}

	systemHome := "/etc"
	if holenSystemEnv := system.Getenv("HOLEN_SYSTEM_CONFIG"); len(holenSystemEnv) > 0 {
		systemHome = holenSystemEnv
	}

	systemConfigPath, err := filepath.EvalSymlinks(filepath.Join(systemHome, "holenconfig"))
	if err != nil {
		return nil, err
	}
	userConfigPath, err := filepath.EvalSymlinks(filepath.Join(baseDir, "config"))
	if err != nil {
		return nil, err
	}

	configClient := RealConfigClient{
		systemConfig: systemConfigPath,
		userConfig:   userConfigPath,
	}

	return &configClient, nil
}
<commit_msg>only desymlink when file exists<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	ini "gopkg.in/ini.v1"
)

type ConfigClient interface {
	GetAll() (map[string]string, error)
	Unset(system bool, key string) error
	Set(system bool, key, value string) error
	Get(key string) (string, error)
}

type ConfigGetter interface {
	Get(key string) (string, error)
}

type RealConfigClient struct {
	systemConfig string
	userConfig   string
}

func (rcc *RealConfigClient) Set(system bool, key, value string) error {
	var configPath string

	if system {
		configPath = rcc.systemConfig
	} else {
		configPath = rcc.userConfig
	}

	cfg, err := ini.LooseLoad(configPath)
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	section, k := splitKey(key)

	_, err = cfg.Section(section).NewKey(k, value)
	if err != nil {
		return fmt.Errorf("unable to set new key: %v", err)
	}

	err = cfg.SaveToIndent(configPath, "    ")
	if err != nil {
		return fmt.Errorf("unable to save: %v", err)
	}

	return nil
}

func (rcc *RealConfigClient) Unset(system bool, key string) error {
	var configPath string

	if system {
		configPath = rcc.systemConfig
	} else {
		configPath = rcc.userConfig
	}

	cfg, err := ini.LooseLoad(configPath)
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	section, k := splitKey(key)

	cfg.Section(section).DeleteKey(k)
	if len(cfg.Section(section).Keys()) == 0 {
		cfg.DeleteSection(section)
	}

	err = cfg.SaveToIndent(configPath, "    ")
	if err != nil {
		return fmt.Errorf("unable to save: %v", err)
	}

	return nil
}

func (rcc *RealConfigClient) Get(key string) (string, error) {
	cfg, err := ini.LooseLoad(rcc.systemConfig, rcc.userConfig)
	if err != nil {
		return "", fmt.Errorf("failure to load config: %v", err)
	}

	section, k := splitKey(key)
	ck := cfg.Section(section).Key(k)

	return ck.Value(), nil
}

func (rcc *RealConfigClient) GetAll() (map[string]string, error) {
	all := make(map[string]string)

	cfg, err := ini.LooseLoad(rcc.systemConfig, rcc.userConfig)
	if err != nil {
		return all, fmt.Errorf("failure to load config: %v", err)
	}

	for _, section := range cfg.Sections() {
		if len(section.Keys()) > 0 {
			for _, key := range section.Keys() {
				all[fmt.Sprintf("%s.%s", section.Name(), key.Name())] = key.Value()
			}
		}
	}

	return all, nil
}

func splitKey(key string) (string, string) {
	parts := strings.SplitN(key, ".", 2)

	return parts[0], parts[1]
}

func NewDefaultConfigClient(system System) (*RealConfigClient, error) {
	var baseDir string
	if xdgConfigHome := system.Getenv("XDG_CONFIG_HOME"); len(xdgConfigHome) > 0 {
		baseDir = filepath.Join(xdgConfigHome, "holen")
	} else {
		var home string
		if home = system.Getenv("HOME"); len(home) == 0 {
			return nil, fmt.Errorf("$HOME environment variable not found")
		}
		baseDir = filepath.Join(home, ".config", "holen")
		os.MkdirAll(baseDir, 0755)
	}

	systemHome := "/etc"
	if holenSystemEnv := system.Getenv("HOLEN_SYSTEM_CONFIG"); len(holenSystemEnv) > 0 {
		systemHome = holenSystemEnv
	}

	var err error
	systemConfigPath := filepath.Join(systemHome, "holenconfig")
	if system.FileExists(systemConfigPath) {
		systemConfigPath, err = filepath.EvalSymlinks(systemConfigPath)
		if err != nil {
			return nil, err
		}
	}

	userConfigPath := filepath.Join(baseDir, "config")
	if system.FileExists(userConfigPath) {
		userConfigPath, err = filepath.EvalSymlinks(userConfigPath)
		if err != nil {
			return nil, err
		}
	}

	configClient := RealConfigClient{
		systemConfig: systemConfigPath,
		userConfig:   userConfigPath,
	}

	return &configClient, nil
}
<|endoftext|>


<commit_before>package gobrake

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

var defaultContext map[string]interface{}

func getDefaultContext() map[string]interface{} {
	if defaultContext != nil {
		return defaultContext
	}

	defaultContext = map[string]interface{}{
		"notifier": map[string]interface{}{
			"name":    "gobrake",
			"version": "2.0.4",
			"url":     "https://github.com/airbrake/gobrake",
		},

		"language":     runtime.Version(),
		"os":           runtime.GOOS,
		"architecture": runtime.GOARCH,
	}
	if s, err := os.Hostname(); err == nil {
		defaultContext["hostname"] = s
	}
	if s, ok := os.LookupEnv("GOPATH"); ok {
		list := filepath.SplitList(s)
		// TODO: multiple root dirs?
		defaultContext["rootDirectory"] = list[0]
	}
	return defaultContext
}

type Error struct {
	Type      string       `json:"type"`
	Message   string       `json:"message"`
	Backtrace []StackFrame `json:"backtrace"`
}

type Notice struct {
	Errors  []Error                `json:"errors"`
	Context map[string]interface{} `json:"context"`
	Env     map[string]interface{} `json:"environment"`
	Session map[string]interface{} `json:"session"`
	Params  map[string]interface{} `json:"params"`
}

func (n *Notice) String() string {
	if len(n.Errors) == 0 {
		return fmt.Sprint(n)
	}
	e := n.Errors[0]
	return fmt.Sprintf("%s: %s", e.Type, e.Message)
}

func NewNotice(e interface{}, req *http.Request, depth int) *Notice {
	notice := &Notice{
		Errors: []Error{{
			Type:      fmt.Sprintf("%T", e),
			Message:   fmt.Sprint(e),
			Backtrace: stack(depth),
		}},
		Context: map[string]interface{}{},
		Env:     map[string]interface{}{},
		Session: map[string]interface{}{},
		Params:  map[string]interface{}{},
	}

	for k, v := range getDefaultContext() {
		notice.Context[k] = v
	}

	if req != nil {
		notice.Context["url"] = req.URL.String()
		if ua := req.Header.Get("User-Agent"); ua != "" {
			notice.Context["userAgent"] = ua
		}

		for k, v := range req.Header {
			if len(v) == 1 {
				notice.Env[k] = v[0]
			} else {
				notice.Env[k] = v
			}
		}

		if err := req.ParseForm(); err == nil {
			for k, v := range req.Form {
				if len(v) == 1 {
					notice.Params[k] = v[0]
				} else {
					notice.Params[k] = v
				}
			}
		}
	}

	return notice
}
<commit_msg>Fix vet warning.<commit_after>package gobrake

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

var defaultContext map[string]interface{}

func getDefaultContext() map[string]interface{} {
	if defaultContext != nil {
		return defaultContext
	}

	defaultContext = map[string]interface{}{
		"notifier": map[string]interface{}{
			"name":    "gobrake",
			"version": "2.0.4",
			"url":     "https://github.com/airbrake/gobrake",
		},

		"language":     runtime.Version(),
		"os":           runtime.GOOS,
		"architecture": runtime.GOARCH,
	}
	if s, err := os.Hostname(); err == nil {
		defaultContext["hostname"] = s
	}
	if s, ok := os.LookupEnv("GOPATH"); ok {
		list := filepath.SplitList(s)
		// TODO: multiple root dirs?
		defaultContext["rootDirectory"] = list[0]
	}
	return defaultContext
}

type Error struct {
	Type      string       `json:"type"`
	Message   string       `json:"message"`
	Backtrace []StackFrame `json:"backtrace"`
}

type Notice struct {
	Errors  []Error                `json:"errors"`
	Context map[string]interface{} `json:"context"`
	Env     map[string]interface{} `json:"environment"`
	Session map[string]interface{} `json:"session"`
	Params  map[string]interface{} `json:"params"`
}

func (n *Notice) String() string {
	if len(n.Errors) == 0 {
		return "Notice<no errors>"
	}
	e := n.Errors[0]
	return fmt.Sprintf("Notice<%s: %s>", e.Type, e.Message)
}

func NewNotice(e interface{}, req *http.Request, depth int) *Notice {
	notice := &Notice{
		Errors: []Error{{
			Type:      fmt.Sprintf("%T", e),
			Message:   fmt.Sprint(e),
			Backtrace: stack(depth),
		}},
		Context: map[string]interface{}{},
		Env:     map[string]interface{}{},
		Session: map[string]interface{}{},
		Params:  map[string]interface{}{},
	}

	for k, v := range getDefaultContext() {
		notice.Context[k] = v
	}

	if req != nil {
		notice.Context["url"] = req.URL.String()
		if ua := req.Header.Get("User-Agent"); ua != "" {
			notice.Context["userAgent"] = ua
		}

		for k, v := range req.Header {
			if len(v) == 1 {
				notice.Env[k] = v[0]
			} else {
				notice.Env[k] = v
			}
		}

		if err := req.ParseForm(); err == nil {
			for k, v := range req.Form {
				if len(v) == 1 {
					notice.Params[k] = v[0]
				} else {
					notice.Params[k] = v
				}
			}
		}
	}

	return notice
}
<|endoftext|>


<commit_before>package network

import (
	"container/list"
	"errors"
	"sync"
	"time"

	pb "github.com/micro/go-micro/network/proto"
)

var (
	// MaxDepth defines max depth of peer topology
	MaxDepth uint = 3
)

var (
	// ErrPeerExists is returned when adding a peer which already exists
	ErrPeerExists = errors.New("peer already exists")
	// ErrPeerNotFound is returned when a peer could not be found in node topology
	ErrPeerNotFound = errors.New("peer not found")
)

// node is network node
type node struct {
	sync.RWMutex
	// id is node id
	id string
	// address is node address
	address string
	// peers are nodes with direct link to this node
	peers map[string]*node
	// network returns the node network
	network Network
	// lastSeen keeps track of node lifetime and updates
	lastSeen time.Time
}

// Id is node ide
func (n *node) Id() string {
	return n.id
}

// Address returns node address
func (n *node) Address() string {
	return n.address
}

// Network returns node network
func (n *node) Network() Network {
	return n.network
}

// walk walks the node graph until some condition is met
func (n *node) walk(until func(peer *node) bool, action func(parent, peer *node)) map[string]*node {
	// track the visited nodes
	visited := make(map[string]*node)
	// queue of the nodes to visit
	queue := list.New()

	// push node to the back of queue
	queue.PushBack(n)
	// mark the node as visited
	visited[n.id] = n

	// keep iterating over the queue until its empty
	for queue.Len() > 0 {
		// pop the node from the front of the queue
		qnode := queue.Front()
		if until(qnode.Value.(*node)) {
			return visited
		}
		// iterate through all of the node peers
		// mark the visited nodes; enqueue the non-visted
		for id, peer := range qnode.Value.(*node).peers {
			if _, ok := visited[id]; !ok {
				visited[id] = peer
				action(qnode.Value.(*node), peer)
				queue.PushBack(peer)
			}
		}
		// remove the node from the queue
		queue.Remove(qnode)
	}

	return visited
}

// AddPeer adds a new peer to node topology
// It returns false if the peer already exists
func (n *node) AddPeer(peer *node) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.peers[peer.id]; !ok {
		n.peers[peer.id] = peer
		return nil
	}

	return ErrPeerExists
}

// DeletePeer deletes a peer from node peers
// It returns true if the peer has been deleted
func (n *node) DeletePeer(id string) bool {
	n.Lock()
	defer n.Unlock()

	delete(n.peers, id)

	return true
}

// UpdatePeer updates a peer if it already exists
// It returns error if the peer does not exist
func (n *node) UpdatePeer(peer *node) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.peers[peer.id]; ok {
		n.peers[peer.id] = peer
		return nil
	}

	return ErrPeerNotFound
}

// RefreshPeer updates node timestamp
// It returns false if the peer has not been found.
func (n *node) RefreshPeer(id string, now time.Time) error {
	n.Lock()
	defer n.Unlock()

	peer, ok := n.peers[id]
	if !ok {
		return ErrPeerNotFound
	}

	if peer.lastSeen.Before(now) {
		peer.lastSeen = now
	}

	return nil
}

// Nodes returns a slice of all nodes in the whole node topology
func (n *node) Nodes() []Node {
	// we need to freeze the network graph here
	// otherwise we might get inconsisten results
	n.RLock()
	defer n.RUnlock()

	// NOTE: this should never be true
	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}
	justWalk := func(parent, node *node) {}

	visited := n.walk(untilNoMorePeers, justWalk)

	var nodes []Node
	// collect all the nodes and return them
	for _, node := range visited {
		nodes = append(nodes, node)
	}

	return nodes
}

// GetPeerNode returns a node from node MaxDepth topology
// It returns nil if the peer was not found
func (n *node) GetPeerNode(id string) *node {
	// get node topology up to MaxDepth
	top := n.Topology(MaxDepth)

	untilFoundPeer := func(n *node) bool {
		return n.id == id
	}
	justWalk := func(paent, node *node) {}

	visited := top.walk(untilFoundPeer, justWalk)

	peerNode, ok := visited[id]
	if !ok {
		return nil
	}

	return peerNode
}

// DeletePeerNode removes peer node from node topology
func (n *node) DeletePeerNode(id string) error {
	n.Lock()
	defer n.Unlock()

	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}

	deleted := make(map[string]*node)
	deletePeer := func(parent, node *node) {
		if node.id != n.id && node.id == id {
			delete(parent.peers, node.id)
			deleted[node.id] = node
		}
	}

	n.walk(untilNoMorePeers, deletePeer)

	if _, ok := deleted[id]; !ok {
		return ErrPeerNotFound
	}

	return nil
}

// PruneStalePeerNodes prune the peers that have not been seen for longer than given time
// It returns a map of the the nodes that got pruned
func (n *node) PruneStalePeerNodes(pruneTime time.Duration) map[string]*node {
	n.Lock()
	defer n.Unlock()

	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}

	pruned := make(map[string]*node)
	pruneStalePeer := func(parent, node *node) {
		if node.id != n.id && time.Since(node.lastSeen) > PruneTime {
			delete(parent.peers, node.id)
			pruned[node.id] = node
		}
	}

	n.walk(untilNoMorePeers, pruneStalePeer)

	return pruned
}

// Topology returns a copy of the node topology down to given depth
// NOTE: the returned node is a node graph - not a single node
func (n *node) Topology(depth uint) *node {
	n.RLock()
	defer n.RUnlock()

	// make a copy of yourself
	node := &node{
		id:       n.id,
		address:  n.address,
		peers:    make(map[string]*node),
		network:  n.network,
		lastSeen: n.lastSeen,
	}

	// return if we reach requested depth or we have no more peers
	if depth == 0 || len(n.peers) == 0 {
		return node
	}

	// decrement the depth
	depth--

	// iterate through our peers and update the node peers
	for _, peer := range n.peers {
		nodePeer := peer.Topology(depth)
		if _, ok := node.peers[nodePeer.id]; !ok {
			node.peers[nodePeer.id] = nodePeer
		}
	}

	return node
}

// Peers returns node peers up to MaxDepth
func (n *node) Peers() []Node {
	n.RLock()
	defer n.RUnlock()

	var peers []Node
	for _, nodePeer := range n.peers {
		peer := nodePeer.Topology(MaxDepth)
		peers = append(peers, peer)
	}

	return peers
}

// UnpackPeerTopology unpacks pb.Peer into node topology of given depth
func UnpackPeerTopology(pbPeer *pb.Peer, lastSeen time.Time, depth uint) *node {
	peerNode := &node{
		id:       pbPeer.Node.Id,
		address:  pbPeer.Node.Address,
		peers:    make(map[string]*node),
		lastSeen: lastSeen,
	}

	// return if have either reached the depth or have no more peers
	if depth == 0 || len(pbPeer.Peers) == 0 {
		return peerNode
	}

	// decrement the depth
	depth--

	peers := make(map[string]*node)
	for _, pbPeer := range pbPeer.Peers {
		peer := UnpackPeerTopology(pbPeer, lastSeen, depth)
		peers[pbPeer.Node.Id] = peer
	}

	peerNode.peers = peers

	return peerNode
}

func peerProtoTopology(peer Node, depth uint) *pb.Peer {
	node := &pb.Node{
		Id:      peer.Id(),
		Address: peer.Address(),
	}

	// set the network name if network is not nil
	if peer.Network() != nil {
		node.Network = peer.Network().Name()
	}

	pbPeers := &pb.Peer{
		Node:  node,
		Peers: make([]*pb.Peer, 0),
	}

	// return if we reached the end of topology or depth
	if depth == 0 || len(peer.Peers()) == 0 {
		return pbPeers
	}

	// decrement the depth
	depth--

	// iterate through peers of peers aka pops
	for _, pop := range peer.Peers() {
		peer := peerProtoTopology(pop, depth)
		pbPeers.Peers = append(pbPeers.Peers, peer)
	}

	return pbPeers
}

// PeersToProto returns node peers graph encoded into protobuf
func PeersToProto(node Node, depth uint) *pb.Peer {
	// network node aka root node
	pbNode := &pb.Node{
		Id:      node.Id(),
		Address: node.Address(),
	}

	// set the network name if network is not nil
	if node.Network() != nil {
		pbNode.Network = node.Network().Name()
	}

	// we will build proto topology into this
	pbPeers := &pb.Peer{
		Node:  pbNode,
		Peers: make([]*pb.Peer, 0),
	}

	for _, peer := range node.Peers() {
		pbPeer := peerProtoTopology(peer, depth)
		pbPeers.Peers = append(pbPeers.Peers, pbPeer)
	}

	return pbPeers
}
<commit_msg>Avoid recursive calls to RLock()<commit_after>package network

import (
	"container/list"
	"errors"
	"sync"
	"time"

	pb "github.com/micro/go-micro/network/proto"
)

var (
	// MaxDepth defines max depth of peer topology
	MaxDepth uint = 3
)

var (
	// ErrPeerExists is returned when adding a peer which already exists
	ErrPeerExists = errors.New("peer already exists")
	// ErrPeerNotFound is returned when a peer could not be found in node topology
	ErrPeerNotFound = errors.New("peer not found")
)

// node is network node
type node struct {
	sync.RWMutex
	// id is node id
	id string
	// address is node address
	address string
	// peers are nodes with direct link to this node
	peers map[string]*node
	// network returns the node network
	network Network
	// lastSeen keeps track of node lifetime and updates
	lastSeen time.Time
}

// Id is node ide
func (n *node) Id() string {
	return n.id
}

// Address returns node address
func (n *node) Address() string {
	return n.address
}

// Network returns node network
func (n *node) Network() Network {
	return n.network
}

// walk walks the node graph until some condition is met
func (n *node) walk(until func(peer *node) bool, action func(parent, peer *node)) map[string]*node {
	// track the visited nodes
	visited := make(map[string]*node)
	// queue of the nodes to visit
	queue := list.New()

	// push node to the back of queue
	queue.PushBack(n)
	// mark the node as visited
	visited[n.id] = n

	// keep iterating over the queue until its empty
	for queue.Len() > 0 {
		// pop the node from the front of the queue
		qnode := queue.Front()
		if until(qnode.Value.(*node)) {
			return visited
		}
		// iterate through all of the node peers
		// mark the visited nodes; enqueue the non-visted
		for id, peer := range qnode.Value.(*node).peers {
			if _, ok := visited[id]; !ok {
				visited[id] = peer
				action(qnode.Value.(*node), peer)
				queue.PushBack(peer)
			}
		}
		// remove the node from the queue
		queue.Remove(qnode)
	}

	return visited
}

// AddPeer adds a new peer to node topology
// It returns false if the peer already exists
func (n *node) AddPeer(peer *node) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.peers[peer.id]; !ok {
		n.peers[peer.id] = peer
		return nil
	}

	return ErrPeerExists
}

// DeletePeer deletes a peer from node peers
// It returns true if the peer has been deleted
func (n *node) DeletePeer(id string) bool {
	n.Lock()
	defer n.Unlock()

	delete(n.peers, id)

	return true
}

// UpdatePeer updates a peer if it already exists
// It returns error if the peer does not exist
func (n *node) UpdatePeer(peer *node) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.peers[peer.id]; ok {
		n.peers[peer.id] = peer
		return nil
	}

	return ErrPeerNotFound
}

// RefreshPeer updates node timestamp
// It returns false if the peer has not been found.
func (n *node) RefreshPeer(id string, now time.Time) error {
	n.Lock()
	defer n.Unlock()

	peer, ok := n.peers[id]
	if !ok {
		return ErrPeerNotFound
	}

	if peer.lastSeen.Before(now) {
		peer.lastSeen = now
	}

	return nil
}

// Nodes returns a slice of all nodes in the whole node topology
func (n *node) Nodes() []Node {
	// we need to freeze the network graph here
	// otherwise we might get inconsisten results
	n.RLock()
	defer n.RUnlock()

	// NOTE: this should never be true
	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}
	justWalk := func(parent, node *node) {}

	visited := n.walk(untilNoMorePeers, justWalk)

	var nodes []Node
	// collect all the nodes and return them
	for _, node := range visited {
		nodes = append(nodes, node)
	}

	return nodes
}

// GetPeerNode returns a node from node MaxDepth topology
// It returns nil if the peer was not found
func (n *node) GetPeerNode(id string) *node {
	// get node topology up to MaxDepth
	top := n.Topology(MaxDepth)

	untilFoundPeer := func(n *node) bool {
		return n.id == id
	}
	justWalk := func(paent, node *node) {}

	visited := top.walk(untilFoundPeer, justWalk)

	peerNode, ok := visited[id]
	if !ok {
		return nil
	}

	return peerNode
}

// DeletePeerNode removes peer node from node topology
func (n *node) DeletePeerNode(id string) error {
	n.Lock()
	defer n.Unlock()

	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}

	deleted := make(map[string]*node)
	deletePeer := func(parent, node *node) {
		if node.id != n.id && node.id == id {
			delete(parent.peers, node.id)
			deleted[node.id] = node
		}
	}

	n.walk(untilNoMorePeers, deletePeer)

	if _, ok := deleted[id]; !ok {
		return ErrPeerNotFound
	}

	return nil
}

// PruneStalePeerNodes prune the peers that have not been seen for longer than given time
// It returns a map of the the nodes that got pruned
func (n *node) PruneStalePeerNodes(pruneTime time.Duration) map[string]*node {
	n.Lock()
	defer n.Unlock()

	untilNoMorePeers := func(node *node) bool {
		return node == nil
	}

	pruned := make(map[string]*node)
	pruneStalePeer := func(parent, node *node) {
		if node.id != n.id && time.Since(node.lastSeen) > PruneTime {
			delete(parent.peers, node.id)
			pruned[node.id] = node
		}
	}

	n.walk(untilNoMorePeers, pruneStalePeer)

	return pruned
}

// getTopology traverses node graph and builds node topology
// NOTE: this function is not thread safe
func (n *node) getTopology(depth uint) *node {
	// make a copy of yourself
	node := &node{
		id:       n.id,
		address:  n.address,
		peers:    make(map[string]*node),
		network:  n.network,
		lastSeen: n.lastSeen,
	}

	// return if we reach requested depth or we have no more peers
	if depth == 0 || len(n.peers) == 0 {
		return node
	}

	// decrement the depth
	depth--

	// iterate through our peers and update the node peers
	for _, peer := range n.peers {
		nodePeer := peer.getTopology(depth)
		if _, ok := node.peers[nodePeer.id]; !ok {
			node.peers[nodePeer.id] = nodePeer
		}
	}

	return node
}

// Topology returns a copy of the node topology down to given depth
// NOTE: the returned node is a node graph - not a single node
func (n *node) Topology(depth uint) *node {
	n.RLock()
	defer n.RUnlock()

	return n.getTopology(depth)
}

// Peers returns node peers up to MaxDepth
func (n *node) Peers() []Node {
	n.RLock()
	defer n.RUnlock()

	var peers []Node
	for _, nodePeer := range n.peers {
		peer := nodePeer.getTopology(MaxDepth)
		peers = append(peers, peer)
	}

	return peers
}

// UnpackPeerTopology unpacks pb.Peer into node topology of given depth
func UnpackPeerTopology(pbPeer *pb.Peer, lastSeen time.Time, depth uint) *node {
	peerNode := &node{
		id:       pbPeer.Node.Id,
		address:  pbPeer.Node.Address,
		peers:    make(map[string]*node),
		lastSeen: lastSeen,
	}

	// return if have either reached the depth or have no more peers
	if depth == 0 || len(pbPeer.Peers) == 0 {
		return peerNode
	}

	// decrement the depth
	depth--

	peers := make(map[string]*node)
	for _, pbPeer := range pbPeer.Peers {
		peer := UnpackPeerTopology(pbPeer, lastSeen, depth)
		peers[pbPeer.Node.Id] = peer
	}

	peerNode.peers = peers

	return peerNode
}

func peerProtoTopology(peer Node, depth uint) *pb.Peer {
	node := &pb.Node{
		Id:      peer.Id(),
		Address: peer.Address(),
	}

	// set the network name if network is not nil
	if peer.Network() != nil {
		node.Network = peer.Network().Name()
	}

	pbPeers := &pb.Peer{
		Node:  node,
		Peers: make([]*pb.Peer, 0),
	}

	// return if we reached the end of topology or depth
	if depth == 0 || len(peer.Peers()) == 0 {
		return pbPeers
	}

	// decrement the depth
	depth--

	// iterate through peers of peers aka pops
	for _, pop := range peer.Peers() {
		peer := peerProtoTopology(pop, depth)
		pbPeers.Peers = append(pbPeers.Peers, peer)
	}

	return pbPeers
}

// PeersToProto returns node peers graph encoded into protobuf
func PeersToProto(node Node, depth uint) *pb.Peer {
	// network node aka root node
	pbNode := &pb.Node{
		Id:      node.Id(),
		Address: node.Address(),
	}

	// set the network name if network is not nil
	if node.Network() != nil {
		pbNode.Network = node.Network().Name()
	}

	// we will build proto topology into this
	pbPeers := &pb.Peer{
		Node:  pbNode,
		Peers: make([]*pb.Peer, 0),
	}

	for _, peer := range node.Peers() {
		pbPeer := peerProtoTopology(peer, depth)
		pbPeers.Peers = append(pbPeers.Peers, pbPeer)
	}

	return pbPeers
}
<|endoftext|>


<commit_before>package apns

import (
	"crypto/tls"
	"encoding/binary"
	"log"
	"time"
)

//ResponseQueueSize indicates how many APNS responses may be buffered.
var ResponseQueueSize = 10000

//SentBufferSize is the maximum number of sent notifications which may be buffered.
var SentBufferSize = 10000

var maxBackoff = 20 * time.Second

//Connection represents a single connection to APNS.
type Connection struct {
	Client
	conn      *tls.Conn
	queue     chan PushNotification
	errors    chan BadPushNotification
	responses chan Response
}

//NewConnection initializes an APNS connection. Use Connection.Start() to actually start sending notifications.
func NewConnection(client *Client) *Connection {
	c := new(Connection)
	c.Client = *client
	queue := make(chan PushNotification)
	errors := make(chan BadPushNotification)
	responses := make(chan Response, ResponseQueueSize)
	c.queue = queue
	c.errors = errors
	c.responses = responses
	return c
}

//Response is a reply from APNS - see apns.ApplePushResponses.
type Response struct {
	Status     uint8
	Identifier uint32
}

func newResponse() Response {
	r := Response{}
	return r
}

//BadPushNotification represents a notification which APNS didn't like.
type BadPushNotification struct {
	PushNotification
	Status uint8
}

//Enqueue adds a push notification to the end of the "sending" queue.
func (conn *Connection) Enqueue(pn *PushNotification) {
	go func(pn *PushNotification) {
		conn.queue <- *pn
	}(pn)
}

//Errors gives you a channel of the push notifications Apple rejected.
func (conn *Connection) Errors() (errors <-chan BadPushNotification) {
	return conn.errors
}

//Start initiates a connection to APNS and asnchronously sends notifications which have been queued.
func (conn *Connection) Start() error {
	//Connect to APNS. The reason this is here as well as in sender is that this probably catches any unavoidable errors in a synchronous fashion, while in sender it can reconnect after temporary errors (which should work most of the time.)
	err := conn.connect()
	if err != nil {
		return err
	}
	//Start sender goroutine
	sent := make(chan PushNotification)
	go conn.sender(conn.queue, sent)
	//Start limbo goroutine
	go conn.limbo(sent, conn.responses, conn.errors, conn.queue)
	return nil
}

//Stop gracefully closes the connection - it waits for the sending queue to clear, and then shuts down.
func (conn *Connection) Stop() {
	//We can't just close the main queue channel, because retries might still need to be sent there.
	//
}

func (conn *Connection) sender(queue <-chan PushNotification, sent chan PushNotification) {
	i := 0
	defer conn.conn.Close()
	var backoff = time.Duration(100)
	for {
		pn, ok := <-conn.queue
		if !ok {
			log.Println("Not okay; queue closed.")
			//That means the Connection is stopped
			//close sent?
			return
		} else {
			//If not connected, connect
			if conn.conn == nil {
				for {
					log.Println("Connection lost; reconnecting.")
					err := conn.connect()
					if err != nil {
						//Exponential backoff up to a limit
						log.Println("APNS: Error connecting to server: ", err)
						backoff = backoff * 2
						if backoff > maxBackoff {
							backoff = maxBackoff
						}
						time.Sleep(backoff)
					} else {
						backoff = 100
						log.Println("Connected...")
						break
					}
				}
			}
			//Then send the push notification
			payload, err := pn.ToBytes()
			if err != nil {
				log.Println(err)
				//Should report this on the bad notifications channel probably
			} else {
				n, err := conn.conn.Write(payload)
				if err != nil {
					log.Println(err)
					//Disconnect?
				} else {
					i++
					log.Println("Sent count:", i)
					log.Println("Wrote bytes:", n, "of notification:", pn.Identifier)
					sent <- pn
				}
			}
		}
	}
}

func (conn *Connection) reader(responses chan<- Response) {
	buffer := make([]byte, 6)
	for {
		n, err := conn.conn.Read(buffer)
		if err != nil && n < 6 {
			log.Println("APNS: Error before reading complete response", n, err)
			conn.conn.Close()
			conn.conn = nil
			return
		}
		command := uint8(buffer[0])
		if command != 8 {
			log.Println("Something went wrong: command should have been 8; it was actually", command)
		}
		resp := newResponse()
		resp.Identifier = binary.BigEndian.Uint32(buffer[2:6])
		resp.Status = uint8(buffer[1])
		responses <- resp
		conn.conn.Close()
		conn.conn = nil
		return
	}
}

func (conn *Connection) limbo(sent <-chan PushNotification, responses chan Response, errors chan BadPushNotification, queue chan PushNotification) {
	limbo := make(chan PushNotification, SentBufferSize)
	ticker := time.NewTicker(1 * time.Second)
	timeNextNotification := true
	for {
		select {
		case pn := <-sent:
			//Drop it into the array
			limbo <- pn
			if timeNextNotification {
				//Is there a cleaner way of doing this?
				go func(pn PushNotification) {
					<-time.After(TimeoutSeconds * time.Second)
					successResp := newResponse()
					successResp.Identifier = pn.Identifier
					responses <- successResp
				}(pn)
				timeNextNotification = false
			}
		case resp, ok := <-responses:
			if !ok {
				//If the responses channel is closed,
				//that means we're shutting down the connection.
			}
			switch {
			case resp.Status == 0:
				log.Println("Got a 0; flushing old notifications up to:", resp.Identifier)
				//Status 0 is a "success" response generated by a timeout in the library.
				if len(limbo) == 0 {
					log.Println("Nothing to flush.")
					break
				}
			FlushLoop:
				for pn := range limbo {
					//Drop all the notifications until we get to the timed-out one.
					//(and leave the others in limbo)
					switch {
					case pn.Identifier == resp.Identifier:
						log.Println("This notification is the last one we know is OK:", pn.Identifier)
						break FlushLoop
					default:
						log.Println("This notification was sent before a successful one, so it's OK:", pn.Identifier)
						//This one was successful; drop silently.
					}
					if len(limbo) == 0 {
						log.Println("Flushed all notifications.")
						break
					}
				}
			default:
				hit := false
				log.Println("Saw an error response; flushing up to", resp.Identifier)
				for pn := range limbo {
					switch {
					case pn.Identifier != resp.Identifier && !hit:
						log.Println("This notification was sent before a failed one, so it's OK:", pn.Identifier)
						//We haven't seen the identified notification yet
						//so these are all successful (drop silently)
					case pn.Identifier == resp.Identifier:
						log.Println("This notification was a bad one:", pn.Identifier)
						hit = true
						if resp.Status != 10 {
							//It was an error, we should report this on the error channel
							bad := BadPushNotification{PushNotification: pn, Status: resp.Status}
							go func(bad BadPushNotification) {
								errors <- bad
							}(bad)
						}
					case pn.Identifier != resp.Identifier && hit:
						log.Println("This notification came after a failed one, so requeue:", pn.Identifier)
						//We've already seen the identified notification,
						//so these should be requeued
						conn.Enqueue(&pn)
					}
					if len(limbo) == 0 {
						log.Println("Flushed all notifications.")
						break
					}
				}
			}
		case <-ticker.C:
			timeNextNotification = true
		}
	}
}

func (conn *Connection) connect() error {
	if conn.conn != nil {
		conn.conn.Close()
	}

	var cert tls.Certificate
	var err error
	if len(conn.CertificateBase64) == 0 && len(conn.KeyBase64) == 0 {
		// The user did not specify raw block contents, so check the filesystem.
		cert, err = tls.LoadX509KeyPair(conn.CertificateFile, conn.KeyFile)
	} else {
		// The user provided the raw block contents, so use that.
		cert, err = tls.X509KeyPair([]byte(conn.CertificateBase64), []byte(conn.KeyBase64))
	}

	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	tlsConn, err := tls.Dial("tcp", conn.Gateway, conf)
	if err != nil {
		return err
	}
	err = tlsConn.Handshake()
	if err != nil {
		_ = tlsConn.Close()
		return err
	}
	conn.conn = tlsConn
	//Start reader goroutine
	go conn.reader(conn.responses)
	return nil
}
<commit_msg>Change the way sent notifications are recorded, from a channel to a slice; this enables correct handling of timeouts.<commit_after>package apns

import (
	"crypto/tls"
	"encoding/binary"
	"log"
	"time"
)

//ResponseQueueSize indicates how many APNS responses may be buffered.
var ResponseQueueSize = 10000

//SentBufferSize is the maximum number of sent notifications which may be buffered.
var SentBufferSize = 10000

var maxBackoff = 20 * time.Second

//Connection represents a single connection to APNS.
type Connection struct {
	Client
	conn      *tls.Conn
	queue     chan PushNotification
	errors    chan BadPushNotification
	responses chan Response
}

//NewConnection initializes an APNS connection. Use Connection.Start() to actually start sending notifications.
func NewConnection(client *Client) *Connection {
	c := new(Connection)
	c.Client = *client
	queue := make(chan PushNotification)
	errors := make(chan BadPushNotification)
	responses := make(chan Response, ResponseQueueSize)
	c.queue = queue
	c.errors = errors
	c.responses = responses
	return c
}

//Response is a reply from APNS - see apns.ApplePushResponses.
type Response struct {
	Status     uint8
	Identifier uint32
}

func newResponse() Response {
	r := Response{}
	return r
}

//BadPushNotification represents a notification which APNS didn't like.
type BadPushNotification struct {
	PushNotification
	Status uint8
}

type timedPushNotification struct {
	PushNotification
	time.Time
}

func (pn PushNotification) timed() timedPushNotification {
	return timedPushNotification{PushNotification: pn, Time: time.Now()}
}

//Enqueue adds a push notification to the end of the "sending" queue.
func (conn *Connection) Enqueue(pn *PushNotification) {
	go func(pn *PushNotification) {
		conn.queue <- *pn
	}(pn)
}

//Errors gives you a channel of the push notifications Apple rejected.
func (conn *Connection) Errors() (errors <-chan BadPushNotification) {
	return conn.errors
}

//Start initiates a connection to APNS and asnchronously sends notifications which have been queued.
func (conn *Connection) Start() error {
	//Connect to APNS. The reason this is here as well as in sender is that this probably catches any unavoidable errors in a synchronous fashion, while in sender it can reconnect after temporary errors (which should work most of the time.)
	err := conn.connect()
	if err != nil {
		return err
	}
	//Start sender goroutine
	sent := make(chan PushNotification)
	go conn.sender(conn.queue, sent)
	//Start limbo goroutine
	go conn.limbo(sent, conn.responses, conn.errors, conn.queue)
	return nil
}

//Stop gracefully closes the connection - it waits for the sending queue to clear, and then shuts down.
func (conn *Connection) Stop() {
	//We can't just close the main queue channel, because retries might still need to be sent there.
	//
}

func (conn *Connection) sender(queue <-chan PushNotification, sent chan PushNotification) {
	i := 0
	defer conn.conn.Close()
	var backoff = time.Duration(100)
	for {
		pn, ok := <-conn.queue
		if !ok {
			log.Println("Not okay; queue closed.")
			//That means the Connection is stopped
			//close sent?
			return
		} else {
			//If not connected, connect
			if conn.conn == nil {
				for {
					log.Println("Connection lost; reconnecting.")
					err := conn.connect()
					if err != nil {
						//Exponential backoff up to a limit
						log.Println("APNS: Error connecting to server: ", err)
						backoff = backoff * 2
						if backoff > maxBackoff {
							backoff = maxBackoff
						}
						time.Sleep(backoff)
					} else {
						backoff = 100
						log.Println("Connected...")
						break
					}
				}
			}
			//Then send the push notification
			payload, err := pn.ToBytes()
			if err != nil {
				log.Println(err)
				//Should report this on the bad notifications channel probably
			} else {
				n, err := conn.conn.Write(payload)
				if err != nil {
					log.Println(err)
					//Disconnect?
				} else {
					i++
					log.Println("Sent count:", i)
					log.Println("Wrote bytes:", n, "of notification:", pn.Identifier)
					sent <- pn
				}
			}
		}
	}
}

func (conn *Connection) reader(responses chan<- Response) {
	buffer := make([]byte, 6)
	for {
		n, err := conn.conn.Read(buffer)
		if err != nil && n < 6 {
			log.Println("APNS: Error before reading complete response", n, err)
			conn.conn.Close()
			conn.conn = nil
			return
		}
		command := uint8(buffer[0])
		if command != 8 {
			log.Println("Something went wrong: command should have been 8; it was actually", command)
		}
		resp := newResponse()
		resp.Identifier = binary.BigEndian.Uint32(buffer[2:6])
		resp.Status = uint8(buffer[1])
		responses <- resp
		conn.conn.Close()
		conn.conn = nil
		return
	}
}

func (conn *Connection) limbo(sent <-chan PushNotification, responses chan Response, errors chan BadPushNotification, queue chan PushNotification) {
	limbo := make([]timedPushNotification, 0, SentBufferSize)
	ticker := time.NewTicker(1 * time.Second)
	for {
		select {
		case pn := <-sent:
			//Drop it into the array
			limbo = append(limbo, pn.timed())
		case resp, ok := <-responses:
			if !ok {
				//If the responses channel is closed,
				//that means we're shutting down the connection.
			}
			log.Println("Saw an error response; flushing up to", resp.Identifier)
			log.Println("PNs in limbo:", len(limbo))
			for i, pn := range limbo {
				if pn.Identifier == resp.Identifier {
					log.Println("Got the bad one!!! At index:", i, "identifier:", pn.Identifier)
					if resp.Status != 10 {
						//It was an error, we should report this on the error channel
						bad := BadPushNotification{PushNotification: pn.PushNotification, Status: resp.Status}
						go func(bad BadPushNotification) {
							errors <- bad
						}(bad)
					}
					if len(limbo) > i {
						log.Printf("Requeueing %d notifications\n", len(limbo)-(i+1))
						conn.requeue(limbo[i+1:])
					}
				}
			}
			limbo = make([]timedPushNotification, 0, SentBufferSize)
		case <-ticker.C:
			flushed := false
			log.Println("Notifications in limbo:", len(limbo))
			for i := range limbo {
				if limbo[i].After(time.Now().Add(-TimeoutSeconds * time.Second)) {
					log.Printf("The first %d notifications timed out ok.\n", i)
					newLimbo := make([]timedPushNotification, len(limbo[i:]), SentBufferSize)
					copy(newLimbo, limbo[i:])
					limbo = newLimbo
					log.Println("Size of limbo is now:", len(limbo))
					flushed = true
					break
				}
			}
			if !flushed {
				log.Println("Looks like every notification has timed out.")
				limbo = make([]timedPushNotification, 0, SentBufferSize)
			}
		}
	}
}

func (conn *Connection) requeue(queue []timedPushNotification) {
	for _, pn := range queue {
		conn.Enqueue(&pn.PushNotification)
	}
}

func (conn *Connection) connect() error {
	if conn.conn != nil {
		conn.conn.Close()
	}

	var cert tls.Certificate
	var err error
	if len(conn.CertificateBase64) == 0 && len(conn.KeyBase64) == 0 {
		// The user did not specify raw block contents, so check the filesystem.
		cert, err = tls.LoadX509KeyPair(conn.CertificateFile, conn.KeyFile)
	} else {
		// The user provided the raw block contents, so use that.
		cert, err = tls.X509KeyPair([]byte(conn.CertificateBase64), []byte(conn.KeyBase64))
	}

	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{cert},
	}

	tlsConn, err := tls.Dial("tcp", conn.Gateway, conf)
	if err != nil {
		return err
	}
	err = tlsConn.Handshake()
	if err != nil {
		_ = tlsConn.Close()
		return err
	}
	conn.conn = tlsConn
	//Start reader goroutine
	go conn.reader(conn.responses)
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Sean Treadway, SoundCloud Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Source code and contact info at http://github.com/streadway/amqp

package amqp

import (
	"bufio"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
)

// Used along side NewConnection to specify the desired tuning parameters used
// during a Connection handshake.  The negotiated tuning will be stored in the
// resultant connection.
type Config struct {
	// The SASL mechanisms to try in the client request, and the successful
	// mechanism used on the Connection object
	SASL []Authentication

	Vhost string // The Vhost the Auth credentials are permitted to open

	MaxChannels       int // Maximum number of channels the client intends to open - defaults to 0
	MaxFrameSize      int // Maximum frame size the client intends to send - default to 0
	HeartbeatInterval int // Frequency in seconds the client wishes the server to send heartbeats - defaults to 0 (no heartbeats)
}

// Manages the serialization and deserialization of frames from IO and dispatches the frames to the appropriate channel.
type Connection struct {
	conn io.ReadWriteCloser

	in            chan message
	shutdownMutex sync.Mutex

	state state // TODO not goroutine-safe; refactor with stateMutex?

	writer      *writer
	writerMutex sync.Mutex

	increment sync.Mutex
	sequence  uint16

	channels map[uint16]*Channel

	Config Config // The negotiated Config after connection.open

	VersionMajor int   // The server's major version
	VersionMinor int   // The server's minor version
	Properties   Table // Server properties
}

// Dial accepts a string in the AMQP URI format, and returns a new Connection
// over TCP using PlainAuth.
func Dial(amqp string) (*Connection, error) {
	uri, err := ParseURI(amqp)
	if err != nil {
		return nil, err
	}

	addr := net.JoinHostPort(uri.Host, strconv.FormatInt(int64(uri.Port), 10))

	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	return NewConnection(conn, Config{
		SASL:  []Authentication{uri.PlainAuth()},
		Vhost: uri.Vhost,
	})
}

func NewConnection(conn io.ReadWriteCloser, config Config) (me *Connection, err error) {
	me = &Connection{
		conn:     conn,
		writer:   &writer{bufio.NewWriter(conn)},
		in:       make(chan message),
		channels: make(map[uint16]*Channel),
	}

	go me.reader()

	return me, me.open(config)
}

func (me *Connection) nextChannelId() uint16 {
	me.increment.Lock()
	defer me.increment.Unlock()
	me.sequence++
	return me.sequence
}

func (me *Connection) Close() (err error) {
	if me.state != closed {
		if err = me.send(
			&methodFrame{
				ChannelId: 0,
				Method: &connectionClose{
					ReplyCode: ReplySuccess,
					ReplyText: "bye",
				},
			},
		); err != nil {
			return
		}

		switch (<-me.in).(type) {
		case *connectionCloseOk:
			me.shutdown()
			return nil
		default:
			return ErrBadProtocol
		}
	}

	return ErrAlreadyClosed
}

func (me *Connection) dispatch() {
	for {

		switch msg := <-me.in; msg.(type) {
		// handle the 4 way shutdown
		case *connectionClose: // request from server
			me.send(&methodFrame{
				ChannelId: 0,
				Method:    &connectionCloseOk{},
			})
			me.Close()
		case *connectionCloseOk: // response to our Close() request
			me.in <- msg // forward to Close() method
			return
		case nil: // closed
			return
		}
	}
}

func (me *Connection) send(f frame) (err error) {
	//fmt.Println("send:", f)

	me.writerMutex.Lock()
	defer me.writerMutex.Unlock()

	if me.state > open {
		return ErrAlreadyClosed
	}

	if err = me.writer.WriteFrame(f); err != nil {
		// TODO handle write failure to cleanly shutdown the connection
		me.shutdown()
	}

	return nil
}

func (me *Connection) shutdown() {
	me.shutdownMutex.Lock()
	defer me.shutdownMutex.Unlock()

	if me.state != closed {
		me.state = closing
		for i, c := range me.channels {
			delete(me.channels, i)
			c.shutdown()
		}
		close(me.in)
		me.conn.Close()
		me.state = closed
	}
}

// All methods sent to the connection channel should be synchronous so we
// can handle them directly without a framing component
func (me *Connection) demux(f frame) {
	if me.state != closed {
		if f.channel() == 0 {
			// TODO send hard error if any content frames/async frames are sent here
			switch mf := f.(type) {
			case *methodFrame:
				me.in <- mf.Method
			default:
				panic("TODO close with hard-error")
			}
		} else {
			channel, ok := me.channels[f.channel()]
			if ok {
				channel.recv(channel, f)
			} else {
				// TODO handle unknown channel for now drop
				println("XXX unknown channel", f.channel())
				panic("XXX unknown channel")
			}
		}
	}
}

// Reads each frame off the IO and hand off to the connection object that
// will demux the streams and dispatch to one of the opened channels or
// handle on channel 0 (the connection channel).
func (me *Connection) reader() {
	buf := bufio.NewReader(me.conn)
	frames := &reader{buf}

	for {
		frame, err := frames.ReadFrame()

		if err != nil {
			me.shutdown()
			return
		}

		me.demux(frame)
	}
}

// Convienence method to inspect the Connection.Properties["capabilities"]
// Table for server identified capabilities like "basic.ack" or
// "confirm.select".
func (me *Connection) IsCapable(featureName string) bool {
	if me.Properties != nil {
		if v, ok := me.Properties["capabilities"]; ok {
			if capabilities, ok := v.(Table); ok {
				if feature, ok := capabilities[featureName]; ok {
					if has, ok := feature.(bool); ok && has {
						return true
					}
				}
			}
		}
	}
	return false
}

// Constructs and opens a unique channel for concurrent operations
func (me *Connection) Channel() (channel *Channel, err error) {
	id := me.nextChannelId()
	channel, err = newChannel(me, id)
	me.channels[id] = channel
	return channel, channel.open()
}

//    Connection          = open-Connection *use-Connection close-Connection
//    open-Connection     = C:protocol-header
//                          S:START C:START-OK
//                          *challenge
//                          S:TUNE C:TUNE-OK
//                          C:OPEN S:OPEN-OK
//    challenge           = S:SECURE C:SECURE-OK
//    use-Connection      = *channel
//    close-Connection    = C:CLOSE S:CLOSE-OK
//                        / S:CLOSE C:CLOSE-OK
func (me *Connection) open(config Config) (err error) {
	me.state = handshaking

	if _, err = me.conn.Write(protocolHeader); err != nil {
		return
	}

	return me.openStart(config)
}

func (me *Connection) openStart(config Config) (err error) {
	switch start := (<-me.in).(type) {
	case *connectionStart:
		me.VersionMajor = int(start.VersionMajor)
		me.VersionMinor = int(start.VersionMinor)
		me.Properties = Table(start.ServerProperties)

		// TODO support challenge/response
		auth, ok := pickSASLMechanism(config.SASL, strings.Split(start.Mechanisms, " "))
		if !ok {
			return ErrUnsupportedMechanism
		}

		// Save this mechanism off as the one we chose
		me.Config.SASL = []Authentication{auth}

		if err = me.send(&methodFrame{
			ChannelId: 0,
			Method: &connectionStartOk{
				Mechanism: auth.Mechanism(),
				Response:  auth.Response(),
			},
		}); err != nil {
			return
		}

		return me.openTune(config)
	case nil:
		return ErrBadProtocol
	}
	return ErrBadProtocol
}

func (me *Connection) openTune(config Config) (err error) {
	switch tune := (<-me.in).(type) {
	// TODO SECURE HANDSHAKE
	case *connectionTune:
		// When this is bounded, share the bound.  We're effectively only bounded
		// by MaxUint16.  If you hit a wrap around bug, throw a small party then
		// make an github issue.
		if int(tune.ChannelMax) > 0 {
			me.Config.MaxChannels = int(tune.ChannelMax)
		}

		// Frame size includes headers and end byte (len(payload)+8), even if
		// this is less than FrameMinSize, use what the server sends because the
		// alternative is to stop the handshake here.
		if int(tune.FrameMax) > 0 {
			if config.MaxFrameSize <= 0 || int(tune.FrameMax) < config.MaxFrameSize {
				me.Config.MaxFrameSize = int(tune.FrameMax)
			} else {
				me.Config.MaxFrameSize = config.MaxFrameSize
			}
		} else {
			if config.MaxFrameSize > 0 {
				me.Config.MaxFrameSize = config.MaxFrameSize
			} else {
				// Client and Server are unlimited.  We'll bound the frame size to
				// 256KB to fit within the bandwidth delay product of a satellite
				// uplink on a scaled TCP window up to 2Mb/s
				me.Config.MaxFrameSize = 256 * 1024
			}
		}

		// This the interval the server wishes us to send at.  config.Heartbeat
		// is the interval what the client wishes the server to send at.
		me.Config.HeartbeatInterval = int(tune.Heartbeat)

		if err = me.send(&methodFrame{
			ChannelId: 0,
			Method: &connectionTuneOk{
				ChannelMax: uint16(me.Config.MaxChannels),
				FrameMax:   uint32(me.Config.MaxFrameSize),
				Heartbeat:  uint16(config.HeartbeatInterval),
			},
		}); err != nil {
			return
		}

		return me.openVhost(config)
	case nil:
		return ErrBadCredentials
	}
	return ErrBadProtocol
}

func (me *Connection) openVhost(config Config) (err error) {
	me.Config.Vhost = config.Vhost

	if err = me.send(&methodFrame{
		ChannelId: 0,
		Method: &connectionOpen{
			VirtualHost: me.Config.Vhost,
		},
	}); err != nil {
		return
	}

	switch (<-me.in).(type) {
	case *connectionOpenOk:
		me.state = open
		go me.dispatch()
		return nil // connection.open handshake finished
	case nil:
		return ErrBadVhost
	}

	return ErrBadProtocol
}
<commit_msg>Implement heartbeats and SetReadTimeout Closes #22 Changes the behavior of `Dial` to wait 30s before giving up on connect and 30s before connection.open is received.<commit_after>// Copyright (c) 2012, Sean Treadway, SoundCloud Ltd.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Source code and contact info at http://github.com/streadway/amqp

package amqp

import (
	"bufio"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Used along side NewConnection to specify the desired tuning parameters used
// during a Connection handshake.  The negotiated tuning will be stored in the
// resultant connection.
type Config struct {
	// The SASL mechanisms to try in the client request, and the successful
	// mechanism used on the Connection object
	SASL []Authentication

	Vhost string // The Vhost the Auth credentials are permitted to open

	MaxChannels       int // Maximum number of channels the client intends to open - defaults to 0
	MaxFrameSize      int // Maximum frame size the client intends to send - default to 0
	HeartbeatInterval int // Frequency in seconds the client wishes the server to send heartbeats - defaults to 0 (no heartbeats)
}

// Manages the serialization and deserialization of frames from IO and dispatches the frames to the appropriate channel.
type Connection struct {
	conn io.ReadWriteCloser

	in            chan message
	shutdownMutex sync.Mutex

	state state // TODO not goroutine-safe; refactor with stateMutex?

	writer      *writer
	writerMutex sync.Mutex

	increment sync.Mutex
	sequence  uint16

	channels map[uint16]*Channel
	sends    chan time.Time // used in place of heartbeats

	Config Config // The negotiated Config after connection.open

	VersionMajor int   // The server's major version
	VersionMinor int   // The server's minor version
	Properties   Table // Server properties
}

type readDeadliner interface {
	SetReadDeadline(time.Time) error
}

// Dial accepts a string in the AMQP URI format, and returns a new Connection
// over TCP using PlainAuth.  Defaults to a server heartbeat interval of 10
// seconds and sets the initial read deadline to 30 seconds.
func Dial(amqp string) (*Connection, error) {
	timeout := time.Duration(30) * time.Second

	uri, err := ParseURI(amqp)
	if err != nil {
		return nil, err
	}

	addr := net.JoinHostPort(uri.Host, strconv.FormatInt(int64(uri.Port), 10))

	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return nil, err
	}

	if err := conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
		return nil, err
	}

	return NewConnection(conn, Config{
		SASL:              []Authentication{uri.PlainAuth()},
		Vhost:             uri.Vhost,
		HeartbeatInterval: 10, // seconds
	})
}

func NewConnection(conn io.ReadWriteCloser, config Config) (me *Connection, err error) {
	me = &Connection{
		conn:     conn,
		writer:   &writer{bufio.NewWriter(conn)},
		in:       make(chan message),
		channels: make(map[uint16]*Channel),
		sends:    make(chan time.Time), // buffered - heartbeat() is reentrant via send()
	}

	go me.reader()

	return me, me.open(config)
}

func (me *Connection) nextChannelId() uint16 {
	me.increment.Lock()
	defer me.increment.Unlock()
	me.sequence++
	return me.sequence
}

func (me *Connection) Close() (err error) {
	if me.state != closed {
		if err = me.send(
			&methodFrame{
				ChannelId: 0,
				Method: &connectionClose{
					ReplyCode: ReplySuccess,
					ReplyText: "bye",
				},
			},
		); err != nil {
			return
		}

		switch (<-me.in).(type) {
		case *connectionCloseOk:
			me.shutdown()
			return nil
		default:
			return ErrBadProtocol
		}
	}

	return ErrAlreadyClosed
}

func (me *Connection) send(f frame) (err error) {
	//fmt.Println("send:", f)

	me.writerMutex.Lock()
	defer me.writerMutex.Unlock()

	if me.state > open {
		return ErrAlreadyClosed
	}

	if err = me.writer.WriteFrame(f); err != nil {
		// TODO handle write failure to cleanly shutdown the connection
		me.shutdown()
	} else {
		me.sent()
	}

	return nil
}

// Broadcast when sent some bytes, reducing heartbeats, only
// if there is something that can receive - like a non-reentrant
// call or if the heartbeater is running
func (me *Connection) sent() {
	select {
	case me.sends <- time.Now():
	default: // drop
	}
}

func (me *Connection) shutdown() {
	me.shutdownMutex.Lock()
	defer me.shutdownMutex.Unlock()

	if me.state != closed {
		me.state = closing
		for i, c := range me.channels {
			delete(me.channels, i)
			c.shutdown()
		}
		close(me.in)
		me.conn.Close()
		me.state = closed
	}
}

// Reset the blocking read deadline on the underlying connection when it
// implements SetReadDeadline to three times the requested heartbeat interval.
// On error, resort to blocking reads.
func (me *Connection) resetDeadline() {
	if beat := time.Duration(me.Config.HeartbeatInterval); beat > 0 {
		if c, ok := me.conn.(readDeadliner); ok {
			c.SetReadDeadline(time.Now().Add(beat * time.Second))
		}
	}
}

// All methods sent to the connection channel should be synchronous so we
// can handle them directly without a framing component
func (me *Connection) demux(f frame) {
	if me.state != closed {
		if f.channel() == 0 {
			// TODO send hard error if any content frames/async frames are sent here
			switch mf := f.(type) {
			case *methodFrame:
				me.in <- mf.Method
			case *heartbeatFrame:
				// kthxbai - all reads reset our deadline.  just drop this one.
			default:
				panic("TODO close with hard-error on unhandled connection frame type")
			}
		} else {
			channel, ok := me.channels[f.channel()]
			if ok {
				channel.recv(channel, f)
			} else {
				// TODO handle unknown channel for now drop
				println("XXX unknown channel", f.channel())
				panic("XXX unknown channel")
			}
		}
	}
}

func (me *Connection) dispatch() {
	for {
		switch msg := <-me.in; msg.(type) {
		// handle the 4 way shutdown
		case *connectionClose: // request from server
			me.send(&methodFrame{
				ChannelId: 0,
				Method:    &connectionCloseOk{},
			})
			me.Close()
		case *connectionCloseOk: // response to our Close() request
			me.in <- msg // forward to Close() method
			return
		case nil: // closed
			return
		}
	}
}

// Reads each frame off the IO and hand off to the connection object that
// will demux the streams and dispatch to one of the opened channels or
// handle on channel 0 (the connection channel).
func (me *Connection) reader() {
	buf := bufio.NewReader(me.conn)
	frames := &reader{buf}

	for {
		frame, err := frames.ReadFrame()

		if err != nil {
			me.shutdown()
			return
		}

		me.demux(frame)

		me.resetDeadline()
	}
}

// Ensures that at least one frame is being sent at the tuned interval with a
// jitter tolerance of 250ms
func (me *Connection) heartbeat(interval time.Duration) {
	last := time.Now()
	tick := time.Tick(interval)

	for {
		select {
		case at := <-tick:
			if at.Sub(last) > interval-250*time.Millisecond {
				if err := me.send(&heartbeatFrame{}); err != nil {
					// send heartbeats even after close/closeOk so we
					// tick until we can't tick no more.
					return
				}
			}
		case at := <-me.sends:
			last = at
		}
	}
}

// Convienence method to inspect the Connection.Properties["capabilities"]
// Table for server identified capabilities like "basic.ack" or
// "confirm.select".
func (me *Connection) IsCapable(featureName string) bool {
	if me.Properties != nil {
		if v, ok := me.Properties["capabilities"]; ok {
			if capabilities, ok := v.(Table); ok {
				if feature, ok := capabilities[featureName]; ok {
					if has, ok := feature.(bool); ok && has {
						return true
					}
				}
			}
		}
	}
	return false
}

// Constructs and opens a unique channel for concurrent operations
func (me *Connection) Channel() (channel *Channel, err error) {
	id := me.nextChannelId()
	channel, err = newChannel(me, id)
	me.channels[id] = channel
	return channel, channel.open()
}

//    Connection          = open-Connection *use-Connection close-Connection
//    open-Connection     = C:protocol-header
//                          S:START C:START-OK
//                          *challenge
//                          S:TUNE C:TUNE-OK
//                          C:OPEN S:OPEN-OK
//    challenge           = S:SECURE C:SECURE-OK
//    use-Connection      = *channel
//    close-Connection    = C:CLOSE S:CLOSE-OK
//                        / S:CLOSE C:CLOSE-OK
func (me *Connection) open(config Config) (err error) {
	me.state = handshaking

	if _, err = me.conn.Write(protocolHeader); err != nil {
		return
	}

	return me.openStart(config)
}

func (me *Connection) openStart(config Config) (err error) {
	switch start := (<-me.in).(type) {
	case *connectionStart:
		me.VersionMajor = int(start.VersionMajor)
		me.VersionMinor = int(start.VersionMinor)
		me.Properties = Table(start.ServerProperties)

		// TODO support challenge/response
		auth, ok := pickSASLMechanism(config.SASL, strings.Split(start.Mechanisms, " "))
		if !ok {
			return ErrUnsupportedMechanism
		}

		// Save this mechanism off as the one we chose
		me.Config.SASL = []Authentication{auth}

		if err = me.send(&methodFrame{
			ChannelId: 0,
			Method: &connectionStartOk{
				Mechanism: auth.Mechanism(),
				Response:  auth.Response(),
			},
		}); err != nil {
			return
		}

		return me.openTune(config)
	case nil:
		return ErrBadProtocol
	}
	return ErrBadProtocol
}

func (me *Connection) openTune(config Config) (err error) {
	switch tune := (<-me.in).(type) {
	// TODO SECURE HANDSHAKE
	case *connectionTune:
		// When this is bounded, share the bound.  We're effectively only bounded
		// by MaxUint16.  If you hit a wrap around bug, throw a small party then
		// make an github issue.
		if int(tune.ChannelMax) > 0 {
			me.Config.MaxChannels = int(tune.ChannelMax)
		}

		// Frame size includes headers and end byte (len(payload)+8), even if
		// this is less than FrameMinSize, use what the server sends because the
		// alternative is to stop the handshake here.
		if int(tune.FrameMax) > 0 {
			if config.MaxFrameSize <= 0 || int(tune.FrameMax) < config.MaxFrameSize {
				me.Config.MaxFrameSize = int(tune.FrameMax)
			} else {
				me.Config.MaxFrameSize = config.MaxFrameSize
			}
		} else {
			if config.MaxFrameSize > 0 {
				me.Config.MaxFrameSize = config.MaxFrameSize
			} else {
				// Client and Server are unlimited.  We'll bound the frame size to
				// 256KB to fit within the bandwidth delay product of a satellite
				// uplink on a scaled TCP window up to 2Mb/s
				me.Config.MaxFrameSize = 256 * 1024
			}
		}

		// Save this off for resetDeadline()
		me.Config.HeartbeatInterval = config.HeartbeatInterval

		// "The client should start sending heartbeats after receiving a
		// Connection.Tune method"
		if tune.Heartbeat > 0 {
			go me.heartbeat(time.Duration(tune.Heartbeat) * time.Second)
		}

		if err = me.send(&methodFrame{
			ChannelId: 0,
			Method: &connectionTuneOk{
				ChannelMax: uint16(me.Config.MaxChannels),
				FrameMax:   uint32(me.Config.MaxFrameSize),
				Heartbeat:  uint16(config.HeartbeatInterval),
			},
		}); err != nil {
			return
		}

		return me.openVhost(config)
	case nil:
		return ErrBadCredentials
	}
	return ErrBadProtocol
}

func (me *Connection) openVhost(config Config) (err error) {
	me.Config.Vhost = config.Vhost

	if err = me.send(&methodFrame{
		ChannelId: 0,
		Method: &connectionOpen{
			VirtualHost: me.Config.Vhost,
		},
	}); err != nil {
		return
	}

	switch (<-me.in).(type) {
	case *connectionOpenOk:
		me.state = open
		go me.dispatch()
		return nil // connection.open handshake finished
	case nil:
		return ErrBadVhost
	}

	return ErrBadProtocol
}
<|endoftext|>


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

package fuse

import (
	"fmt"
	"log"
	"path"
	"runtime"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/fuse/fuseops"
)

// A connection to the fuse kernel process.
type Connection struct {
	logger      *log.Logger
	wrapped     *bazilfuse.Conn
	opsInFlight sync.WaitGroup

	// The context from which all op contexts inherit.
	parentCtx context.Context

	// For logging purposes only.
	nextOpID uint32
}

// Responsibility for closing the wrapped connection is transferred to the
// result. You must call c.close() eventually.
func newConnection(
	parentCtx context.Context,
	logger *log.Logger,
	wrapped *bazilfuse.Conn) (c *Connection, err error) {
	c = &Connection{
		logger:    logger,
		wrapped:   wrapped,
		parentCtx: parentCtx,
	}

	return
}

// Log information for an operation with the given ID. calldepth is the depth
// to use when recovering file:line information with runtime.Caller.
func (c *Connection) log(
	opID uint32,
	calldepth int,
	format string,
	v ...interface{}) {
	// Get file:line info.
	var file string
	var line int
	var ok bool

	_, file, line, ok = runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}

	fileLine := fmt.Sprintf("%v:%v", path.Base(file), line)

	// Format the actual message to be printed.
	msg := fmt.Sprintf(
		"Op 0x%08x %24s] %v",
		opID,
		fileLine,
		fmt.Sprintf(format, v...))

	// Print it.
	c.logger.Println(msg)
}

// Set up state for an op that is about to be returned to the user, given its
// bazilfuse request ID.
//
// Return a context that should be used for the op.
func (c *Connection) beginOp(reqID bazilfuse.RequestID) (ctx context.Context) {
	c.opsInFlight.Add(1)

	// TODO(jacobsa): Use WithCancel and stash a cancellation function.
	ctx = c.parentCtx

	return
}

// Clean up all state associated with an op to which the user has responded.
func (c *Connection) finishOp(reqID bazilfuse.RequestID) {
	c.opsInFlight.Done()
}

// Read the next op from the kernel process. Return io.EOF if the kernel has
// closed the connection.
//
// This function delivers ops in exactly the order they are received from
// /dev/fuse. It must not be called multiple times concurrently.
func (c *Connection) ReadOp() (op fuseops.Op, err error) {
	// Keep going until we find a request we know how to convert.
	for {
		// Read a bazilfuse request.
		var bfReq bazilfuse.Request
		bfReq, err = c.wrapped.ReadRequest()

		if err != nil {
			return
		}

		// Choose an ID for this operation.
		opID := c.nextOpID
		c.nextOpID++

		// Log the receipt of the operation.
		c.log(opID, 1, "<- %v", bfReq)

		// Special case: responding to this is required to make mounting work on OS
		// X. We don't currently expose the capability for the file system to
		// intercept this.
		if statfsReq, ok := bfReq.(*bazilfuse.StatfsRequest); ok {
			c.log(opID, 1, "-> (Statfs) OK")
			statfsReq.Respond(&bazilfuse.StatfsResponse{})
			continue
		}

		// Set up op dependencies.
		var reqID bazilfuse.RequestID = bfReq.Hdr().ID
		opCtx := c.beginOp(reqID)

		logForOp := func(calldepth int, format string, v ...interface{}) {
			c.log(opID, calldepth+1, format, v...)
		}

		finished := func(err error) { c.finishOp(reqID) }

		op = fuseops.Convert(opCtx, bfReq, logForOp, finished)
		return
	}
}

func (c *Connection) waitForReady() (err error) {
	<-c.wrapped.Ready
	err = c.wrapped.MountError
	return
}

// Close the connection and wait for in-flight ops.
func (c *Connection) close() (err error) {
	err = c.wrapped.Close()
	c.opsInFlight.Wait()
	return
}
<commit_msg>Arrange for a cancellation channel.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fuse

import (
	"fmt"
	"log"
	"path"
	"runtime"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/fuse/fuseops"
)

// A connection to the fuse kernel process.
type Connection struct {
	logger      *log.Logger
	wrapped     *bazilfuse.Conn
	opsInFlight sync.WaitGroup

	// The context from which all op contexts inherit.
	parentCtx context.Context

	// For logging purposes only.
	nextOpID uint32

	mu sync.Mutex

	// A map from bazilfuse request ID (*not* the op ID for logging used above)
	// to a function that cancel's its associated context.
	//
	// GUARDED_BY(mu)
	cancelFuncs map[bazilfuse.RequestID]func()
}

// Responsibility for closing the wrapped connection is transferred to the
// result. You must call c.close() eventually.
func newConnection(
	parentCtx context.Context,
	logger *log.Logger,
	wrapped *bazilfuse.Conn) (c *Connection, err error) {
	c = &Connection{
		logger:    logger,
		wrapped:   wrapped,
		parentCtx: parentCtx,
	}

	return
}

// Log information for an operation with the given ID. calldepth is the depth
// to use when recovering file:line information with runtime.Caller.
func (c *Connection) log(
	opID uint32,
	calldepth int,
	format string,
	v ...interface{}) {
	// Get file:line info.
	var file string
	var line int
	var ok bool

	_, file, line, ok = runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}

	fileLine := fmt.Sprintf("%v:%v", path.Base(file), line)

	// Format the actual message to be printed.
	msg := fmt.Sprintf(
		"Op 0x%08x %24s] %v",
		opID,
		fileLine,
		fmt.Sprintf(format, v...))

	// Print it.
	c.logger.Println(msg)
}

// LOCKS_EXCLUDED(c.mu)
func (c *Connection) recordCancelFunc(
	reqID bazilfuse.RequestID,
	f func())

// Set up state for an op that is about to be returned to the user, given its
// bazilfuse request ID.
//
// Return a context that should be used for the op.
//
// LOCKS_EXCLUDED(c.mu)
func (c *Connection) beginOp(reqID bazilfuse.RequestID) (ctx context.Context) {
	// Note that the op is in flight.
	c.opsInFlight.Add(1)

	// Set up a cancellation function.
	ctx, cancel := context.WithCancel(c.parentCtx)
	c.recordCancelFunc(reqID, cancel)

	return
}

// Clean up all state associated with an op to which the user has responded.
//
// LOCKS_EXCLUDED(c.mu)
func (c *Connection) finishOp(reqID bazilfuse.RequestID) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// Even though the op is finished, context.WithCancel requires us to arrange
	// for the cancellation function to be invoked. We also must remove it from
	// our map.
	cancel, ok := c.cancelFuncs[reqID]
	if !ok {
		panic(fmt.Sprintf("Unknown request ID in finishOp: %v", reqID))
	}

	cancel()
	delete(c.cancelFuncs, reqID)

	// Decrement the in-flight counter.
	c.opsInFlight.Done()
}

// Read the next op from the kernel process. Return io.EOF if the kernel has
// closed the connection.
//
// This function delivers ops in exactly the order they are received from
// /dev/fuse. It must not be called multiple times concurrently.
//
// LOCKS_EXCLUDED(c.mu)
func (c *Connection) ReadOp() (op fuseops.Op, err error) {
	// Keep going until we find a request we know how to convert.
	for {
		// Read a bazilfuse request.
		var bfReq bazilfuse.Request
		bfReq, err = c.wrapped.ReadRequest()

		if err != nil {
			return
		}

		// Choose an ID for this operation.
		opID := c.nextOpID
		c.nextOpID++

		// Log the receipt of the operation.
		c.log(opID, 1, "<- %v", bfReq)

		// Special case: responding to this is required to make mounting work on OS
		// X. We don't currently expose the capability for the file system to
		// intercept this.
		if statfsReq, ok := bfReq.(*bazilfuse.StatfsRequest); ok {
			c.log(opID, 1, "-> (Statfs) OK")
			statfsReq.Respond(&bazilfuse.StatfsResponse{})
			continue
		}

		// Set up op dependencies.
		var reqID bazilfuse.RequestID = bfReq.Hdr().ID
		opCtx := c.beginOp(reqID)

		logForOp := func(calldepth int, format string, v ...interface{}) {
			c.log(opID, calldepth+1, format, v...)
		}

		finished := func(err error) { c.finishOp(reqID) }

		op = fuseops.Convert(opCtx, bfReq, logForOp, finished)
		return
	}
}

func (c *Connection) waitForReady() (err error) {
	<-c.wrapped.Ready
	err = c.wrapped.MountError
	return
}

// Close the connection and wait for in-flight ops.
func (c *Connection) close() (err error) {
	err = c.wrapped.Close()
	c.opsInFlight.Wait()
	return
}
<|endoftext|>