<commit_before>package model

// Model holds an abstracted data model representing the translation
// of various types of Kubernetes config to Cilium config.
type Model struct {
	HTTP []HTTPListener `json:"http,omitempty"`
}

// HTTPListener holds configuration for any listener that terminates and proxies HTTP
// including HTTP and HTTPS.
// Each holds the configuration info for one distinct HTTP listener, by
//   - Hostname
//   - TLS
//   - Address
//   - Port
type HTTPListener struct {
	// Name of the HTTPListener
	Name string `json:"name,omitempty"`
	// Sources is a slice of fully qualified resources this HTTPListener is sourced
	// from.
	Sources []FullyQualifiedResource `json:"sources,omitempty"`
	// IPAddress that the listener should listen on.
	// The string must be parseable as an IP address.
	Address string `json:"address,omitempty"`
	// Port on which the service can be expected to be accessed by clients.
	Port uint32 `json:"port,omitempty"`
	// Hostname that the listener should match.
	// Wildcards are supported in prefix or suffix forms, or the special wildcard `*`.
	// An empty list means that the Listener should match all hostnames.
	Hostname string `json:"hostname,omitempty"`
	// TLS Certificate information. If omitted, then the listener is a cleartext HTTP listener.
	TLS *TLSSecret `json:"tls,omitempty"`
	// Routes associated with HTTP traffic to the service.
	// An empty list means that traffic will not be routed.
	Routes []HTTPRoute `json:"routes,omitempty"`
}

// FullyQualifiedResource stores the full details of a Kubernetes resource, including
// the Group, Version, and Kind.
// Namespace must be set to the empty string for cluster-scoped resources.
type FullyQualifiedResource struct {
	Name      string `json:"name,omitempty"`
	Namespace string `json:"namespace,omitempty"`
	Group     string `json:"group,omitempty"`
	Version   string `json:"version,omitempty"`
	Kind      string `json:"kind,omitempty"`
}

// TLSSecret holds a reference to a secret containing a TLS keypair.
type TLSSecret struct {
	Name      string `json:"name,omitempty"`
	Namespace string `json:"namespace,omitempty"`
}

// HTTPRoute holds all the details needed to route HTTP traffic to a backend.
type HTTPRoute struct {
	Name string `json:"name,omitempty"`
	// PathMatch specifies that the HTTPRoute should match a path.
	PathMatch StringMatch `json:"path_match"`
	Backends  []Backend   `json:"backends,omitempty"`
}

// StringMatch describes various types of string matching.
// Only one field may be set.
// If no fields are set, all paths should match (no path match criteria should
// be generated for Envoy.)
type StringMatch struct {
	Prefix string `json:"prefix,omitempty"`
	Exact  string `json:"exact,omitempty"`
	Regex  string `json:"regex,omitempty"`
}

// Backend holds a Kubernetes Service that points to a backend for traffic.
type Backend struct {
	// Name of the Service.
	Name string `json:"name,omitempty"`
	// Namespace of the Service.
	Namespace string `json:"namespace,omitempty"`
	// Port contains the details of the port on the Service to connect to
	// If unset, the same port as the top-level Listener will be used.
	Port *BackendPort `json:"port,omitempty"`
}

// BackendPort holds the details of what port on the Service to connect to.
// Only one of Port or Name can be set.
type BackendPort struct {
	// Port holds the numeric port to connect to.
	Port uint32 `json:"port,omitempty"`
	// Name holds a string which will be used to connect to the port with a
	// matching spec.ports[].name in the target Service.
	Name string `json:"name,omitempty"`
}
<commit_msg>ingress/model: Add GetPort function for BackendPort<commit_after>package model

import "strconv"

// Model holds an abstracted data model representing the translation
// of various types of Kubernetes config to Cilium config.
type Model struct {
	HTTP []HTTPListener `json:"http,omitempty"`
}

// HTTPListener holds configuration for any listener that terminates and proxies HTTP
// including HTTP and HTTPS.
// Each holds the configuration info for one distinct HTTP listener, by
//   - Hostname
//   - TLS
//   - Address
//   - Port
type HTTPListener struct {
	// Name of the HTTPListener
	Name string `json:"name,omitempty"`
	// Sources is a slice of fully qualified resources this HTTPListener is sourced
	// from.
	Sources []FullyQualifiedResource `json:"sources,omitempty"`
	// IPAddress that the listener should listen on.
	// The string must be parseable as an IP address.
	Address string `json:"address,omitempty"`
	// Port on which the service can be expected to be accessed by clients.
	Port uint32 `json:"port,omitempty"`
	// Hostname that the listener should match.
	// Wildcards are supported in prefix or suffix forms, or the special wildcard `*`.
	// An empty list means that the Listener should match all hostnames.
	Hostname string `json:"hostname,omitempty"`
	// TLS Certificate information. If omitted, then the listener is a cleartext HTTP listener.
	TLS *TLSSecret `json:"tls,omitempty"`
	// Routes associated with HTTP traffic to the service.
	// An empty list means that traffic will not be routed.
	Routes []HTTPRoute `json:"routes,omitempty"`
}

// FullyQualifiedResource stores the full details of a Kubernetes resource, including
// the Group, Version, and Kind.
// Namespace must be set to the empty string for cluster-scoped resources.
type FullyQualifiedResource struct {
	Name      string `json:"name,omitempty"`
	Namespace string `json:"namespace,omitempty"`
	Group     string `json:"group,omitempty"`
	Version   string `json:"version,omitempty"`
	Kind      string `json:"kind,omitempty"`
}

// TLSSecret holds a reference to a secret containing a TLS keypair.
type TLSSecret struct {
	Name      string `json:"name,omitempty"`
	Namespace string `json:"namespace,omitempty"`
}

// HTTPRoute holds all the details needed to route HTTP traffic to a backend.
type HTTPRoute struct {
	Name string `json:"name,omitempty"`
	// PathMatch specifies that the HTTPRoute should match a path.
	PathMatch StringMatch `json:"path_match"`
	Backends  []Backend   `json:"backends,omitempty"`
}

// StringMatch describes various types of string matching.
// Only one field may be set.
// If no fields are set, all paths should match (no path match criteria should
// be generated for Envoy.)
type StringMatch struct {
	Prefix string `json:"prefix,omitempty"`
	Exact  string `json:"exact,omitempty"`
	Regex  string `json:"regex,omitempty"`
}

// Backend holds a Kubernetes Service that points to a backend for traffic.
type Backend struct {
	// Name of the Service.
	Name string `json:"name,omitempty"`
	// Namespace of the Service.
	Namespace string `json:"namespace,omitempty"`
	// Port contains the details of the port on the Service to connect to
	// If unset, the same port as the top-level Listener will be used.
	Port *BackendPort `json:"port,omitempty"`
}

// BackendPort holds the details of what port on the Service to connect to.
// Only one of Port or Name can be set.
type BackendPort struct {
	// Port holds the numeric port to connect to.
	Port uint32 `json:"port,omitempty"`
	// Name holds a string which will be used to connect to the port with a
	// matching spec.ports[].name in the target Service.
	Name string `json:"name,omitempty"`
}

// GetPort return the string representation of the port (either the port number or the port name)
func (be *BackendPort) GetPort() string {
	if be.Port != 0 {
		return strconv.Itoa(int(be.Port))
	}
	return be.Name
}
<|endoftext|>


<commit_before>// robustsession represents a RobustIRC session and handles all communication
// to the RobustIRC network.
package robustsession

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/robustirc/robustirc/types"

	"github.com/sorcix/irc"
)

const (
	pathCreateSession = "/robustirc/v1/session"
	pathDeleteSession = "/robustirc/v1/%s"
	pathPostMessage   = "/robustirc/v1/%s/message"
	pathGetMessages   = "/robustirc/v1/%s/messages?lastseen=%s"
)

var (
	NoSuchSession = errors.New("No such RobustIRC session (killed by the network?)")

	networks   = make(map[string]*network)
	networksMu sync.Mutex
)

type backoffState struct {
	exp  float64
	next time.Time
}

type network struct {
	servers []string
	mu      sync.RWMutex
	backoff map[string]backoffState
}

func newNetwork(networkname string) (*network, error) {
	var servers []string

	parts := strings.Split(networkname, ",")
	if len(parts) > 1 {
		log.Printf("Interpreting %q as list of servers instead of network name\n", networkname)
		servers = parts
	} else {
		// Try to resolve the DNS name up to 5 times. This is to be nice to
		// people in environments with flaky network connections at boot, who,
		// for some reason, don’t run this program under systemd with
		// Restart=on-failure.
		try := 0
		for {
			_, addrs, err := net.LookupSRV("robustirc", "tcp", networkname)
			if err != nil {
				log.Println(err)
				if try < 4 {
					time.Sleep(time.Duration(int64(math.Pow(2, float64(try)))) * time.Second)
				} else {
					return nil, fmt.Errorf("DNS lookup of %q failed 5 times", networkname)
				}
				try++
				continue
			}
			// Randomly shuffle the addresses.
			for i := range addrs {
				j := rand.Intn(i + 1)
				addrs[i], addrs[j] = addrs[j], addrs[i]
			}
			for _, addr := range addrs {
				target := addr.Target
				if target[len(target)-1] == '.' {
					target = target[:len(target)-1]
				}
				servers = append(servers, fmt.Sprintf("%s:%d", target, addr.Port))
			}
			break
		}
	}

	return &network{
		servers: servers,
		backoff: make(map[string]backoffState),
	}, nil
}

// server (eventually) returns the host:port to which we should connect to. In
// case back-off prevents us from connecting anywhere right now, the function
// blocks until back-off is over.
func (n *network) server() string {
	n.mu.RLock()
	defer n.mu.RUnlock()

	for {
		soonest := time.Duration(math.MaxInt64)
		for _, server := range n.servers {
			wait := n.backoff[server].next.Sub(time.Now())
			if wait <= 0 {
				return server
			}
			if wait < soonest {
				soonest = wait
			}
		}

		time.Sleep(soonest)
	}
}

func (n *network) setServers(servers []string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	// TODO(secure): we should clean up n.backoff from servers which no longer exist
	n.servers = servers
}

// prefer adds the specified server to the front of the servers list, thereby
// trying to prefer it over other servers for the next request. Note that
// exponential backoff overrides this, so this is only a hint, not a guarantee.
func (n *network) prefer(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	n.servers = append([]string{server}, n.servers...)
}

func (n *network) failed(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	b := n.backoff[server]
	b.exp++
	b.next = time.Now().Add(time.Duration(math.Pow(2, b.exp)) * time.Second)
	n.backoff[server] = b
}

func (n *network) succeeded(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	delete(n.backoff, server)
}

func discardResponse(resp *http.Response) {
	// We need to read the entire body, otherwise net/http will not
	// re-use this connection.
	ioutil.ReadAll(resp.Body)
	resp.Body.Close()
}

type RobustSession struct {
	IrcPrefix *irc.Prefix
	Messages  chan string
	Errors    chan error

	sessionId   string
	sessionAuth string
	deleted     bool
	done        chan bool
	network     *network
	client      *http.Client
}

func (s *RobustSession) sendRequest(method, path string, data []byte) (string, *http.Response, error) {
	for !s.deleted {
		target := s.network.server()
		req, err := http.NewRequest(method, fmt.Sprintf("https://%s%s", target, path), bytes.NewBuffer(data))
		if err != nil {
			return "", nil, err
		}
		req.Header.Set("X-Session-Auth", s.sessionAuth)
		req.Header.Set("Content-Type", "application/json")
		resp, err := s.client.Do(req)
		if err != nil {
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed: %v\n", path, err)
			continue
		}
		if resp.StatusCode == http.StatusNotFound {
			s.network.failed(target)
			return "", nil, NoSuchSession
		}
		// Server errors, temporary.
		if resp.StatusCode >= 500 && resp.StatusCode < 600 {
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed with %v (retrying)\n", path, resp.Status)
			continue
		}
		// Client errors and anything unexpected.
		if resp.StatusCode != http.StatusOK {
			message, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed with %v: %q\n", path, resp.Status, message)
			continue
		}
		return target, resp, nil
	}

	return "", nil, NoSuchSession
}

// Create creates a new RobustIRC session. It resolves the given network name
// (e.g. "robustirc.net") to a set of servers by querying the
// _robustirc._tcp.<network> SRV record and sends the CreateSession request.
//
// When err == nil, the caller MUST read the RobustSession.Messages and
// RobustSession.Errors channels.
func Create(network string, tlsCAFile string) (*RobustSession, error) {
	networksMu.Lock()
	n, ok := networks[network]
	if !ok {
		var err error
		n, err = newNetwork(network)
		if err != nil {
			return nil, err
		}
		networks[network] = n
	}
	networksMu.Unlock()

	var client *http.Client

	if tlsCAFile != "" {
		roots := x509.NewCertPool()
		contents, err := ioutil.ReadFile(tlsCAFile)
		if err != nil {
			log.Fatalf("Could not read cert.pem: %v", err)
		}
		if !roots.AppendCertsFromPEM(contents) {
			log.Fatalf("Could not parse %q", tlsCAFile)
		}

		client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig:     &tls.Config{RootCAs: roots},
				MaxIdleConnsPerHost: 1,
			},
		}
	} else {
		client = &http.Client{
			Transport: &http.Transport{
				MaxIdleConnsPerHost: 1,
			},
		}
	}

	s := &RobustSession{
		Messages: make(chan string),
		Errors:   make(chan error),
		done:     make(chan bool, 1),
		network:  n,
		client:   client,
	}

	_, resp, err := s.sendRequest("POST", pathCreateSession, nil)
	if err != nil {
		return nil, err
	}
	defer discardResponse(resp)

	var createSessionReply struct {
		Sessionid   string
		Sessionauth string
		Prefix      string
	}

	if err := json.NewDecoder(resp.Body).Decode(&createSessionReply); err != nil {
		return nil, err
	}

	if cl := resp.Header.Get("Content-Location"); cl != "" {
		if location, err := url.Parse(cl); err == nil {
			log.Printf("Preferring %q (current leader)\n", location.Host)
			s.network.prefer(location.Host)
		}
	}

	s.sessionId = createSessionReply.Sessionid
	s.sessionAuth = createSessionReply.Sessionauth
	s.IrcPrefix = &irc.Prefix{Name: createSessionReply.Prefix}

	go s.getMessages()

	return s, nil
}

func (s *RobustSession) getMessages() {
	var lastseen types.RobustId

	for !s.deleted {
		target, resp, err := s.sendRequest("GET", fmt.Sprintf(pathGetMessages, s.sessionId, lastseen.String()), nil)
		if err != nil {
			s.Errors <- err
			return
		}

		msgchan := make(chan types.RobustMessage, 1)
		errchan := make(chan error)
		go func() {
			dec := json.NewDecoder(resp.Body)
			for {
				var msg types.RobustMessage
				if err := dec.Decode(&msg); err != nil {
					errchan <- err
					return
				}
				msgchan <- msg
			}
		}()

	ReadLoop:
		for !s.deleted {
			select {
			case msg := <-msgchan:
				if msg.Type == types.RobustPing {
					s.network.setServers(msg.Servers)
				} else if msg.Type == types.RobustIRCToClient {
					// TODO: remove/debug
					log.Printf("<-robustirc: %q\n", msg.Data)
					s.Messages <- msg.Data
					lastseen = msg.Id
				}

			case err := <-errchan:
				log.Printf("Protocol error on %q: Could not decode response chunk as JSON: %v\n", target, err)
				s.network.failed(target)
				break ReadLoop

			case <-time.After(1 * time.Minute):
				log.Printf("Timeout (60s) on GetMessages, reconnecting…\n")
				s.network.failed(target)
				break ReadLoop

			case <-s.done:
				break ReadLoop
			}
		}
		resp.Body.Close()

		// Delay reconnecting for somewhere in between [250, 500) ms to avoid
		// overloading the remaining servers from many clients at once when one
		// server fails.
		time.Sleep(time.Duration(250+rand.Int63n(250)) * time.Millisecond)
	}
}

// PostMessage posts the given IRC message.
func (s *RobustSession) PostMessage(message string) error {
	type postMessageRequest struct {
		Data            string
		ClientMessageId uint64
	}

	h := fnv.New32()
	h.Write([]byte(message))
	// The message id should be unique across separate instances of the bridge,
	// even if they were attached to the same session. A collision in this case
	// means one bridge instance (with the same session) is unable to send a
	// message because the message id is equal to the one the other bridge
	// instance just sent. With the hash of the message itself, such a
	// collision can only occur when both instances try to send exactly the
	// same message _and_ the random value is the same for both instances.
	msgid := (uint64(h.Sum32()) << 32) | uint64(rand.Int31n(math.MaxInt32))

	b, err := json.Marshal(postMessageRequest{
		Data:            message,
		ClientMessageId: msgid,
	})
	if err != nil {
		return fmt.Errorf("Message could not be encoded as JSON: %v\n", err)
	}

	target, resp, err := s.sendRequest("POST", fmt.Sprintf(pathPostMessage, s.sessionId), b)
	if err != nil {
		return err
	}
	discardResponse(resp)
	s.network.succeeded(target)
	return nil
}

// Delete sends a delete request for this session on the server.
//
// This session MUST not be used after this method returns. Even if the delete
// request did not succeed, the session is deleted from the client’s point of
// view.
func (s *RobustSession) Delete(quitmessage string) error {
	defer func() {
		s.deleted = true
		s.done <- true
	}()

	b, err := json.Marshal(struct{ Quitmessage string }{quitmessage})
	if err != nil {
		return err
	}
	_, resp, err := s.sendRequest("DELETE", fmt.Sprintf(pathDeleteSession, s.sessionId), b)
	if err != nil {
		return err
	}
	discardResponse(resp)
	return nil
}
<commit_msg>close and drain channels to avoid hanging goroutines<commit_after>// robustsession represents a RobustIRC session and handles all communication
// to the RobustIRC network.
package robustsession

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"hash/fnv"
	"io/ioutil"
	"log"
	"math"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/robustirc/robustirc/types"

	"github.com/sorcix/irc"
)

const (
	pathCreateSession = "/robustirc/v1/session"
	pathDeleteSession = "/robustirc/v1/%s"
	pathPostMessage   = "/robustirc/v1/%s/message"
	pathGetMessages   = "/robustirc/v1/%s/messages?lastseen=%s"
)

var (
	NoSuchSession = errors.New("No such RobustIRC session (killed by the network?)")

	networks   = make(map[string]*network)
	networksMu sync.Mutex
)

type backoffState struct {
	exp  float64
	next time.Time
}

type network struct {
	servers []string
	mu      sync.RWMutex
	backoff map[string]backoffState
}

func newNetwork(networkname string) (*network, error) {
	var servers []string

	parts := strings.Split(networkname, ",")
	if len(parts) > 1 {
		log.Printf("Interpreting %q as list of servers instead of network name\n", networkname)
		servers = parts
	} else {
		// Try to resolve the DNS name up to 5 times. This is to be nice to
		// people in environments with flaky network connections at boot, who,
		// for some reason, don’t run this program under systemd with
		// Restart=on-failure.
		try := 0
		for {
			_, addrs, err := net.LookupSRV("robustirc", "tcp", networkname)
			if err != nil {
				log.Println(err)
				if try < 4 {
					time.Sleep(time.Duration(int64(math.Pow(2, float64(try)))) * time.Second)
				} else {
					return nil, fmt.Errorf("DNS lookup of %q failed 5 times", networkname)
				}
				try++
				continue
			}
			// Randomly shuffle the addresses.
			for i := range addrs {
				j := rand.Intn(i + 1)
				addrs[i], addrs[j] = addrs[j], addrs[i]
			}
			for _, addr := range addrs {
				target := addr.Target
				if target[len(target)-1] == '.' {
					target = target[:len(target)-1]
				}
				servers = append(servers, fmt.Sprintf("%s:%d", target, addr.Port))
			}
			break
		}
	}

	return &network{
		servers: servers,
		backoff: make(map[string]backoffState),
	}, nil
}

// server (eventually) returns the host:port to which we should connect to. In
// case back-off prevents us from connecting anywhere right now, the function
// blocks until back-off is over.
func (n *network) server() string {
	n.mu.RLock()
	defer n.mu.RUnlock()

	for {
		soonest := time.Duration(math.MaxInt64)
		for _, server := range n.servers {
			wait := n.backoff[server].next.Sub(time.Now())
			if wait <= 0 {
				return server
			}
			if wait < soonest {
				soonest = wait
			}
		}

		time.Sleep(soonest)
	}
}

func (n *network) setServers(servers []string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	// TODO(secure): we should clean up n.backoff from servers which no longer exist
	n.servers = servers
}

// prefer adds the specified server to the front of the servers list, thereby
// trying to prefer it over other servers for the next request. Note that
// exponential backoff overrides this, so this is only a hint, not a guarantee.
func (n *network) prefer(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	n.servers = append([]string{server}, n.servers...)
}

func (n *network) failed(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	b := n.backoff[server]
	b.exp++
	b.next = time.Now().Add(time.Duration(math.Pow(2, b.exp)) * time.Second)
	n.backoff[server] = b
}

func (n *network) succeeded(server string) {
	n.mu.Lock()
	defer n.mu.Unlock()

	delete(n.backoff, server)
}

func discardResponse(resp *http.Response) {
	// We need to read the entire body, otherwise net/http will not
	// re-use this connection.
	ioutil.ReadAll(resp.Body)
	resp.Body.Close()
}

type RobustSession struct {
	IrcPrefix *irc.Prefix
	Messages  chan string
	Errors    chan error

	sessionId   string
	sessionAuth string
	deleted     bool
	done        chan bool
	network     *network
	client      *http.Client
}

func (s *RobustSession) sendRequest(method, path string, data []byte) (string, *http.Response, error) {
	for !s.deleted {
		target := s.network.server()
		req, err := http.NewRequest(method, fmt.Sprintf("https://%s%s", target, path), bytes.NewBuffer(data))
		if err != nil {
			return "", nil, err
		}
		req.Header.Set("X-Session-Auth", s.sessionAuth)
		req.Header.Set("Content-Type", "application/json")
		resp, err := s.client.Do(req)
		if err != nil {
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed: %v\n", path, err)
			continue
		}
		if resp.StatusCode == http.StatusNotFound {
			s.network.failed(target)
			return "", nil, NoSuchSession
		}
		// Server errors, temporary.
		if resp.StatusCode >= 500 && resp.StatusCode < 600 {
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed with %v (retrying)\n", path, resp.Status)
			continue
		}
		// Client errors and anything unexpected.
		if resp.StatusCode != http.StatusOK {
			message, _ := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			s.network.failed(target)
			log.Printf("sendRequest(%q) failed with %v: %q\n", path, resp.Status, message)
			continue
		}
		return target, resp, nil
	}

	return "", nil, NoSuchSession
}

// Create creates a new RobustIRC session. It resolves the given network name
// (e.g. "robustirc.net") to a set of servers by querying the
// _robustirc._tcp.<network> SRV record and sends the CreateSession request.
//
// When err == nil, the caller MUST read the RobustSession.Messages and
// RobustSession.Errors channels.
func Create(network string, tlsCAFile string) (*RobustSession, error) {
	networksMu.Lock()
	n, ok := networks[network]
	if !ok {
		var err error
		n, err = newNetwork(network)
		if err != nil {
			return nil, err
		}
		networks[network] = n
	}
	networksMu.Unlock()

	var client *http.Client

	if tlsCAFile != "" {
		roots := x509.NewCertPool()
		contents, err := ioutil.ReadFile(tlsCAFile)
		if err != nil {
			log.Fatalf("Could not read cert.pem: %v", err)
		}
		if !roots.AppendCertsFromPEM(contents) {
			log.Fatalf("Could not parse %q", tlsCAFile)
		}

		client = &http.Client{
			Transport: &http.Transport{
				TLSClientConfig:     &tls.Config{RootCAs: roots},
				MaxIdleConnsPerHost: 1,
			},
		}
	} else {
		client = &http.Client{
			Transport: &http.Transport{
				MaxIdleConnsPerHost: 1,
			},
		}
	}

	s := &RobustSession{
		Messages: make(chan string),
		Errors:   make(chan error),
		done:     make(chan bool),
		network:  n,
		client:   client,
	}

	_, resp, err := s.sendRequest("POST", pathCreateSession, nil)
	if err != nil {
		return nil, err
	}
	defer discardResponse(resp)

	var createSessionReply struct {
		Sessionid   string
		Sessionauth string
		Prefix      string
	}

	if err := json.NewDecoder(resp.Body).Decode(&createSessionReply); err != nil {
		return nil, err
	}

	if cl := resp.Header.Get("Content-Location"); cl != "" {
		if location, err := url.Parse(cl); err == nil {
			log.Printf("Preferring %q (current leader)\n", location.Host)
			s.network.prefer(location.Host)
		}
	}

	s.sessionId = createSessionReply.Sessionid
	s.sessionAuth = createSessionReply.Sessionauth
	s.IrcPrefix = &irc.Prefix{Name: createSessionReply.Prefix}

	go s.getMessages()

	return s, nil
}

func (s *RobustSession) getMessages() {
	var lastseen types.RobustId

	defer func() {
		for _ = range s.done {
		}
	}()

	for !s.deleted {
		target, resp, err := s.sendRequest("GET", fmt.Sprintf(pathGetMessages, s.sessionId, lastseen.String()), nil)
		if err != nil {
			s.Errors <- err
			return
		}

		msgchan := make(chan types.RobustMessage)
		errchan := make(chan error)
		done := false
		go func() {
			defer close(msgchan)
			defer close(errchan)
			dec := json.NewDecoder(resp.Body)
			for !done {
				var msg types.RobustMessage
				if err := dec.Decode(&msg); err != nil {
					errchan <- err
					return
				}
				msgchan <- msg
			}
		}()

	ReadLoop:
		for !s.deleted {
			select {
			case msg := <-msgchan:
				if msg.Type == types.RobustPing {
					s.network.setServers(msg.Servers)
				} else if msg.Type == types.RobustIRCToClient {
					// TODO: remove/debug
					log.Printf("<-robustirc: %q\n", msg.Data)
					s.Messages <- msg.Data
					lastseen = msg.Id
				}

			case err := <-errchan:
				log.Printf("Protocol error on %q: Could not decode response chunk as JSON: %v\n", target, err)
				s.network.failed(target)
				break ReadLoop

			case <-time.After(1 * time.Minute):
				log.Printf("Timeout (60s) on GetMessages, reconnecting…\n")
				s.network.failed(target)
				break ReadLoop

			case <-s.done:
				break ReadLoop
			}
		}
		done = true
		go func() {
			for _ = range msgchan {
			}
		}()
		go func() {
			for _ = range errchan {
			}
		}()
		// Cannot use discardResponse() because the response never completes.
		resp.Body.Close()

		// Delay reconnecting for somewhere in between [250, 500) ms to avoid
		// overloading the remaining servers from many clients at once when one
		// server fails.
		time.Sleep(time.Duration(250+rand.Int63n(250)) * time.Millisecond)
	}
}

// PostMessage posts the given IRC message.
func (s *RobustSession) PostMessage(message string) error {
	type postMessageRequest struct {
		Data            string
		ClientMessageId uint64
	}

	h := fnv.New32()
	h.Write([]byte(message))
	// The message id should be unique across separate instances of the bridge,
	// even if they were attached to the same session. A collision in this case
	// means one bridge instance (with the same session) is unable to send a
	// message because the message id is equal to the one the other bridge
	// instance just sent. With the hash of the message itself, such a
	// collision can only occur when both instances try to send exactly the
	// same message _and_ the random value is the same for both instances.
	msgid := (uint64(h.Sum32()) << 32) | uint64(rand.Int31n(math.MaxInt32))

	b, err := json.Marshal(postMessageRequest{
		Data:            message,
		ClientMessageId: msgid,
	})
	if err != nil {
		return fmt.Errorf("Message could not be encoded as JSON: %v\n", err)
	}

	target, resp, err := s.sendRequest("POST", fmt.Sprintf(pathPostMessage, s.sessionId), b)
	if err != nil {
		return err
	}
	discardResponse(resp)
	s.network.succeeded(target)
	return nil
}

// Delete sends a delete request for this session on the server.
//
// This session MUST not be used after this method returns. Even if the delete
// request did not succeed, the session is deleted from the client’s point of
// view.
func (s *RobustSession) Delete(quitmessage string) error {
	defer func() {
		s.deleted = true

		// Make sure nobody is blocked on sending to the channel by closing them
		// and reading all remaining values.
		go func() {
			for _ = range s.Messages {
			}
		}()
		go func() {
			for _ = range s.Errors {
			}
		}()

		// This will be read by getMessages(), which will not send on the
		// channels after reading it, so we can safely close them.
		s.done <- true
		close(s.done)

		close(s.Messages)
		close(s.Errors)
	}()

	b, err := json.Marshal(struct{ Quitmessage string }{quitmessage})
	if err != nil {
		return err
	}
	_, resp, err := s.sendRequest("DELETE", fmt.Sprintf(pathDeleteSession, s.sessionId), b)
	if err != nil {
		return err
	}
	discardResponse(resp)
	return nil
}
<|endoftext|>


<commit_before>package greedo

import (
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	"github.com/tmc/scp"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

func getAgent() (agent.Agent, error) {
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	return agent.NewClient(agentConn), err
}

const (
	greedoAddr = `greedo.xeserv.us:22`
	greedoUser = `xena`
)

// Dial opens a SSH client to greedo.
func Dial() (*ssh.Client, error) {
	agent, err := getAgent()
	if err != nil {
		return nil, err
	}

	client, err := ssh.Dial("tcp", greedoAddr, &ssh.ClientConfig{
		User: greedoUser,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeysCallback(agent.Signers),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
	if err != nil {
		return nil, err
	}

	return client, nil
}

// Copy copies a local file reader to the remote destination path. This copies the enitre contents of contents into ram. Don't use this function if doing so is a bad idea. Only works on files less than 2 GB.
func Copy(mode os.FileMode, fileName string, contents io.Reader, destinationPath string) error {
	data, err := ioutil.ReadAll(contents)
	if err != nil {
		return err
	}

	client, err := Dial()
	if err != nil {
		return err
	}

	session, err := client.NewSession()
	if err != nil {
		return err
	}

	err = scp.Copy(int64(len(data)), mode, fileName, bytes.NewBuffer(data), destinationPath, session)
	if err != nil {
		return err
	}

	return nil
}

// CopyFile copies a file to Greedo's public files folder and returns its public-facing URL.
func CopyFile(fileName string, contents io.Reader) (string, error) {
	err := Copy(0644, fileName, contents, filepath.Join("public_html", "files", fileName))
	if err != nil {
		return "", err
	}

	return WebURL(fileName), nil
}

// WebURL constructs a public-facing URL for a given resource by fragment.
func WebURL(fragment string) string {
	return "https://xena.greedo.xeserv.us/files/" + fragment
}
<commit_msg>internal/greedo: add slugs prefix to copied files<commit_after>package greedo

import (
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"path/filepath"

	"github.com/tmc/scp"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
)

func getAgent() (agent.Agent, error) {
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	return agent.NewClient(agentConn), err
}

const (
	greedoAddr = `greedo.xeserv.us:22`
	greedoUser = `xena`
)

// Dial opens a SSH client to greedo.
func Dial() (*ssh.Client, error) {
	agent, err := getAgent()
	if err != nil {
		return nil, err
	}

	client, err := ssh.Dial("tcp", greedoAddr, &ssh.ClientConfig{
		User: greedoUser,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeysCallback(agent.Signers),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	})
	if err != nil {
		return nil, err
	}

	return client, nil
}

// Copy copies a local file reader to the remote destination path. This copies the enitre contents of contents into ram. Don't use this function if doing so is a bad idea. Only works on files less than 2 GB.
func Copy(mode os.FileMode, fileName string, contents io.Reader, destinationPath string) error {
	data, err := ioutil.ReadAll(contents)
	if err != nil {
		return err
	}

	client, err := Dial()
	if err != nil {
		return err
	}

	session, err := client.NewSession()
	if err != nil {
		return err
	}

	err = scp.Copy(int64(len(data)), mode, fileName, bytes.NewBuffer(data), destinationPath, session)
	if err != nil {
		return err
	}

	return nil
}

// CopyFile copies a file to Greedo's public files folder and returns its public-facing URL.
func CopyFile(fileName string, contents io.Reader) (string, error) {
	err := Copy(0644, fileName, contents, filepath.Join("public_html", "files", "slugs", fileName))
	if err != nil {
		return "", err
	}

	return WebURL(fileName), nil
}

// WebURL constructs a public-facing URL for a given resource by fragment.
func WebURL(fragment string) string {
	return "https://xena.greedo.xeserv.us/files/" + fragment
}
<|endoftext|>


<commit_before>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package diff supports a pluggable diff algorithm.
package diff

import (
	"bytes"
	"sort"

	"golang.org/x/tools/internal/span"
)

// TextEdit represents a change to a section of a document.
// The text within the specified span should be replaced by the supplied new text.
type TextEdit struct {
	Span    span.Span
	NewText string
}

// ComputeEdits is the type for a function that produces a set of edits that
// convert from the before content to the after content.
type ComputeEdits func(uri span.URI, before, after string) []TextEdit

// SortTextEdits attempts to order all edits by their starting points.
// The sort is stable so that edits with the same starting point will not
// be reordered.
func SortTextEdits(d []TextEdit) {
	// Use a stable sort to maintain the order of edits inserted at the same position.
	sort.SliceStable(d, func(i int, j int) bool {
		return span.Compare(d[i].Span, d[j].Span) < 0
	})
}

// ApplyEdits applies the set of edits to the before and returns the resulting
// content.
// It may panic or produce garbage if the edits are not valid for the provided
// before content.
func ApplyEdits(before string, edits []TextEdit) string {
	// Preconditions:
	//   - all of the edits apply to before
	//   - and all the spans for each TextEdit have the same URI

	// copy edits so we don't make a mess of the caller's slice
	s := make([]TextEdit, len(edits))
	copy(s, edits)
	edits = s

	// TODO(matloob): Initialize the Converter Once?
	var conv span.Converter = span.NewContentConverter("", []byte(before))
	offset := func(point span.Point) int {
		if point.HasOffset() {
			return point.Offset()
		}
		offset, err := conv.ToOffset(point.Line(), point.Column())
		if err != nil {
			panic(err)
		}
		return offset
	}

	// sort the copy
	sort.Slice(edits, func(i, j int) bool { return offset(edits[i].Span.Start()) < offset(edits[j].Span.Start()) })

	var after bytes.Buffer
	beforeOffset := 0
	for _, edit := range edits {
		if offset(edit.Span.Start()) < beforeOffset {
			panic("overlapping edits") // TODO(matloob): ApplyEdits doesn't return an error. What do we do?
		} else if offset(edit.Span.Start()) > beforeOffset {
			after.WriteString(before[beforeOffset:offset(edit.Span.Start())])
			beforeOffset = offset(edit.Span.Start())
		}
		// offset(edit.Span.Start) is now equal to beforeOffset
		after.WriteString(edit.NewText)
		beforeOffset += offset(edit.Span.End()) - offset(edit.Span.Start())
	}
	if beforeOffset < len(before) {
		after.WriteString(before[beforeOffset:])
		beforeOffset = len(before[beforeOffset:]) // just to preserve invariants
	}
	return after.String()
}
<commit_msg>internal/lsp: clean up ApplyEdits<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package diff supports a pluggable diff algorithm.
package diff

import (
	"sort"
	"strings"

	"golang.org/x/tools/internal/span"
)

// TextEdit represents a change to a section of a document.
// The text within the specified span should be replaced by the supplied new text.
type TextEdit struct {
	Span    span.Span
	NewText string
}

// ComputeEdits is the type for a function that produces a set of edits that
// convert from the before content to the after content.
type ComputeEdits func(uri span.URI, before, after string) []TextEdit

// SortTextEdits attempts to order all edits by their starting points.
// The sort is stable so that edits with the same starting point will not
// be reordered.
func SortTextEdits(d []TextEdit) {
	// Use a stable sort to maintain the order of edits inserted at the same position.
	sort.SliceStable(d, func(i int, j int) bool {
		return span.Compare(d[i].Span, d[j].Span) < 0
	})
}

// ApplyEdits applies the set of edits to the before and returns the resulting
// content.
// It may panic or produce garbage if the edits are not valid for the provided
// before content.
func ApplyEdits(before string, edits []TextEdit) string {
	// Preconditions:
	//   - all of the edits apply to before
	//   - and all the spans for each TextEdit have the same URI
	if len(edits) == 0 {
		return before
	}
	edits = prepareEdits(edits)
	c := span.NewContentConverter("", []byte(before))
	after := strings.Builder{}
	last := 0
	for _, edit := range edits {
		spn, _ := edit.Span.WithAll(c)
		start := spn.Start().Offset()
		if start > last {
			after.WriteString(before[last:start])
			last = start
		}
		after.WriteString(edit.NewText)
		last = spn.End().Offset()
	}
	if last < len(before) {
		after.WriteString(before[last:])
	}
	return after.String()
}

// prepareEdits returns a sorted copy of the edits
func prepareEdits(edits []TextEdit) []TextEdit {
	copied := make([]TextEdit, len(edits))
	copy(copied, edits)
	SortTextEdits(copied)
	return copied
}
<|endoftext|>


<commit_before>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package server

import (
	"expvar"
	"fmt"
	"net/http"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/maier/go-appstats"
	"github.com/spf13/viper"
)

func (s *Server) router(w http.ResponseWriter, r *http.Request) {
	_ = appstats.IncrementInt("requests_total")

	s.logger.Debug().Str("method", r.Method).Str("url", r.URL.String()).Msg("request")

	switch r.Method {
	case "GET":
		switch {
		case r.URL.Path == "/health", r.URL.Path == "/health/":
			w.WriteHeader(http.StatusOK)
			_, _ = fmt.Fprintln(w, "Alive")
		case pluginPathRx.MatchString(r.URL.Path): // run plugin(s)
			if viper.GetBool(config.KeyMultiAgent) {
				http.Error(w, "not allowed when multi-agent enabled", http.StatusForbidden)
				return
			}
			s.run(w, r)
		case inventoryPathRx.MatchString(r.URL.Path): // plugin inventory
			s.inventory(w)
		case statsPathRx.MatchString(r.URL.Path): // app stats
			expvar.Handler().ServeHTTP(w, r)
		case promPathRx.MatchString(r.URL.Path): // output prom format...
			s.promOutput(w)
		default:
			_ = appstats.IncrementInt("requests_bad")
			s.logger.Warn().Str("method", r.Method).Str("url", r.URL.String()).Msg("not found")
			http.NotFound(w, r)
		}
	case "POST":
		fallthrough
	case "PUT":
		switch {
		case writePathRx.MatchString(r.URL.Path):
			s.write(w, r)
		case promPathRx.MatchString(r.URL.Path):
			s.promReceiver(w, r)
		default:
			_ = appstats.IncrementInt("requests_bad")
			s.logger.Warn().Str("method", r.Method).Str("url", r.URL.String()).Msg("not found")
			http.NotFound(w, r)
		}
	default:
		_ = appstats.IncrementInt("requests_bad")
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}
<commit_msg>add: last run req stat<commit_after>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package server

import (
	"expvar"
	"fmt"
	"net/http"
	"time"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/maier/go-appstats"
	"github.com/spf13/viper"
)

func (s *Server) router(w http.ResponseWriter, r *http.Request) {
	_ = appstats.IncrementInt("server.requests_total")

	s.logger.Debug().Str("method", r.Method).Str("url", r.URL.String()).Msg("request")

	switch r.Method {
	case "GET":
		switch {
		case r.URL.Path == "/health", r.URL.Path == "/health/":
			w.WriteHeader(http.StatusOK)
			_, _ = fmt.Fprintln(w, "Alive")
		case pluginPathRx.MatchString(r.URL.Path): // run plugin(s)
			if viper.GetBool(config.KeyMultiAgent) {
				http.Error(w, "not allowed when multi-agent enabled", http.StatusForbidden)
				return
			}
			if err := appstats.SetString("server.last_run_request", time.Now().String()); err != nil {
				s.logger.Warn().Err(err).Msg("setting app stat - last_run_request")
			}
			s.run(w, r)
		case inventoryPathRx.MatchString(r.URL.Path): // plugin inventory
			s.inventory(w)
		case statsPathRx.MatchString(r.URL.Path): // app stats
			expvar.Handler().ServeHTTP(w, r)
		case promPathRx.MatchString(r.URL.Path): // output prom format...
			s.promOutput(w)
		default:
			_ = appstats.IncrementInt("server.requests_bad")
			s.logger.Warn().Str("method", r.Method).Str("url", r.URL.String()).Msg("not found")
			http.NotFound(w, r)
		}
	case "POST":
		fallthrough
	case "PUT":
		switch {
		case writePathRx.MatchString(r.URL.Path):
			s.write(w, r)
		case promPathRx.MatchString(r.URL.Path):
			s.promReceiver(w, r)
		default:
			_ = appstats.IncrementInt("server.requests_bad")
			s.logger.Warn().Str("method", r.Method).Str("url", r.URL.String()).Msg("not found")
			http.NotFound(w, r)
		}
	default:
		_ = appstats.IncrementInt("server.requests_bad")
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssa

// flagalloc allocates the flag register among all the flag-generating
// instructions. Flag values are recomputed if they need to be
// spilled/restored.
func flagalloc(f *Func) {
	// Compute the in-register flag value we want at the end of
	// each block. This is basically a best-effort live variable
	// analysis, so it can be much simpler than a full analysis.
	// TODO: do we really need to keep flag values live across blocks?
	// Could we force the flags register to be unused at basic block
	// boundaries?  Then we wouldn't need this computation.
	end := make([]*Value, f.NumBlocks())
	for n := 0; n < 2; n++ {
		// Walk blocks backwards. Poor-man's postorder traversal.
		for i := len(f.Blocks) - 1; i >= 0; i-- {
			b := f.Blocks[i]
			// Walk values backwards to figure out what flag
			// value we want in the flag register at the start
			// of the block.
			flag := end[b.ID]
			if b.Control != nil && b.Control.Type.IsFlags() {
				flag = b.Control
			}
			for j := len(b.Values) - 1; j >= 0; j-- {
				v := b.Values[j]
				if v == flag {
					flag = nil
				}
				if opcodeTable[v.Op].clobberFlags {
					flag = nil
				}
				for _, a := range v.Args {
					if a.Type.IsFlags() {
						flag = a
					}
				}
			}
			if flag != nil {
				for _, e := range b.Preds {
					p := e.b
					end[p.ID] = flag
				}
			}
		}
	}

	// For blocks which have a flags control value, that's the only value
	// we can leave in the flags register at the end of the block. (There
	// is no place to put a flag regeneration instruction.)
	for _, b := range f.Blocks {
		v := b.Control
		if v != nil && v.Type.IsFlags() && end[b.ID] != v {
			end[b.ID] = nil
		}
		if b.Kind == BlockDefer {
			// Defer blocks internally use/clobber the flags value.
			end[b.ID] = nil
		}
	}

	// Add flag recomputations where they are needed.
	// TODO: Remove original instructions if they are never used.
	var oldSched []*Value
	for _, b := range f.Blocks {
		oldSched = append(oldSched[:0], b.Values...)
		b.Values = b.Values[:0]
		// The current live flag value the pre-flagalloc copy).
		var flag *Value
		if len(b.Preds) > 0 {
			flag = end[b.Preds[0].b.ID]
			// Note: the following condition depends on the lack of critical edges.
			for _, e := range b.Preds[1:] {
				p := e.b
				if end[p.ID] != flag {
					f.Fatalf("live flag in %s's predecessors not consistent", b)
				}
			}
		}
		for _, v := range oldSched {
			if v.Op == OpPhi && v.Type.IsFlags() {
				f.Fatalf("phi of flags not supported: %s", v.LongString())
			}
			// Make sure any flag arg of v is in the flags register.
			// If not, recompute it.
			for i, a := range v.Args {
				if !a.Type.IsFlags() {
					continue
				}
				if a == flag {
					continue
				}
				// Recalculate a
				c := copyFlags(a, b)
				// Update v.
				v.SetArg(i, c)
				// Remember the most-recently computed flag value.
				flag = a
			}
			// Issue v.
			b.Values = append(b.Values, v)
			if opcodeTable[v.Op].clobberFlags {
				flag = nil
			}
			if v.Type.IsFlags() {
				flag = v
			}
		}
		if v := b.Control; v != nil && v != flag && v.Type.IsFlags() {
			// Recalculate control value.
			c := v.copyInto(b)
			b.SetControl(c)
			flag = v
		}
		if v := end[b.ID]; v != nil && v != flag {
			// Need to reissue flag generator for use by
			// subsequent blocks.
			copyFlags(v, b)
			// Note: this flag generator is not properly linked up
			// with the flag users. This breaks the SSA representation.
			// We could fix up the users with another pass, but for now
			// we'll just leave it.  (Regalloc has the same issue for
			// standard regs, and it runs next.)
		}
	}

	// Save live flag state for later.
	for _, b := range f.Blocks {
		b.FlagsLiveAtEnd = end[b.ID] != nil
	}
}

// copyFlags copies v (flag generator) into b, returns the copy.
// If v's arg is also flags, copy recursively.
func copyFlags(v *Value, b *Block) *Value {
	flagsArgs := make(map[int]*Value)
	for i, a := range v.Args {
		if a.Type.IsFlags() || a.Type.IsTuple() {
			flagsArgs[i] = copyFlags(a, b)
		}
	}
	c := v.copyInto(b)
	for i, a := range flagsArgs {
		c.SetArg(i, a)
	}
	return c
}
<commit_msg>cmd/compile: fix tuple-generating flag ops as clobbering flags<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ssa

// flagalloc allocates the flag register among all the flag-generating
// instructions. Flag values are recomputed if they need to be
// spilled/restored.
func flagalloc(f *Func) {
	// Compute the in-register flag value we want at the end of
	// each block. This is basically a best-effort live variable
	// analysis, so it can be much simpler than a full analysis.
	// TODO: do we really need to keep flag values live across blocks?
	// Could we force the flags register to be unused at basic block
	// boundaries?  Then we wouldn't need this computation.
	end := make([]*Value, f.NumBlocks())
	for n := 0; n < 2; n++ {
		// Walk blocks backwards. Poor-man's postorder traversal.
		for i := len(f.Blocks) - 1; i >= 0; i-- {
			b := f.Blocks[i]
			// Walk values backwards to figure out what flag
			// value we want in the flag register at the start
			// of the block.
			flag := end[b.ID]
			if b.Control != nil && b.Control.Type.IsFlags() {
				flag = b.Control
			}
			for j := len(b.Values) - 1; j >= 0; j-- {
				v := b.Values[j]
				if v == flag {
					flag = nil
				}
				if v.clobbersFlags() {
					flag = nil
				}
				for _, a := range v.Args {
					if a.Type.IsFlags() {
						flag = a
					}
				}
			}
			if flag != nil {
				for _, e := range b.Preds {
					p := e.b
					end[p.ID] = flag
				}
			}
		}
	}

	// For blocks which have a flags control value, that's the only value
	// we can leave in the flags register at the end of the block. (There
	// is no place to put a flag regeneration instruction.)
	for _, b := range f.Blocks {
		v := b.Control
		if v != nil && v.Type.IsFlags() && end[b.ID] != v {
			end[b.ID] = nil
		}
		if b.Kind == BlockDefer {
			// Defer blocks internally use/clobber the flags value.
			end[b.ID] = nil
		}
	}

	// Add flag recomputations where they are needed.
	// TODO: Remove original instructions if they are never used.
	var oldSched []*Value
	for _, b := range f.Blocks {
		oldSched = append(oldSched[:0], b.Values...)
		b.Values = b.Values[:0]
		// The current live flag value the pre-flagalloc copy).
		var flag *Value
		if len(b.Preds) > 0 {
			flag = end[b.Preds[0].b.ID]
			// Note: the following condition depends on the lack of critical edges.
			for _, e := range b.Preds[1:] {
				p := e.b
				if end[p.ID] != flag {
					f.Fatalf("live flag in %s's predecessors not consistent", b)
				}
			}
		}
		for _, v := range oldSched {
			if v.Op == OpPhi && v.Type.IsFlags() {
				f.Fatalf("phi of flags not supported: %s", v.LongString())
			}
			// Make sure any flag arg of v is in the flags register.
			// If not, recompute it.
			for i, a := range v.Args {
				if !a.Type.IsFlags() {
					continue
				}
				if a == flag {
					continue
				}
				// Recalculate a
				c := copyFlags(a, b)
				// Update v.
				v.SetArg(i, c)
				// Remember the most-recently computed flag value.
				flag = a
			}
			// Issue v.
			b.Values = append(b.Values, v)
			if v.clobbersFlags() {
				flag = nil
			}
			if v.Type.IsFlags() {
				flag = v
			}
		}
		if v := b.Control; v != nil && v != flag && v.Type.IsFlags() {
			// Recalculate control value.
			c := v.copyInto(b)
			b.SetControl(c)
			flag = v
		}
		if v := end[b.ID]; v != nil && v != flag {
			// Need to reissue flag generator for use by
			// subsequent blocks.
			copyFlags(v, b)
			// Note: this flag generator is not properly linked up
			// with the flag users. This breaks the SSA representation.
			// We could fix up the users with another pass, but for now
			// we'll just leave it.  (Regalloc has the same issue for
			// standard regs, and it runs next.)
		}
	}

	// Save live flag state for later.
	for _, b := range f.Blocks {
		b.FlagsLiveAtEnd = end[b.ID] != nil
	}
}

func (v *Value) clobbersFlags() bool {
	if opcodeTable[v.Op].clobberFlags {
		return true
	}
	if v.Type.IsTuple() && (v.Type.FieldType(0).IsFlags() || v.Type.FieldType(1).IsFlags()) {
		// This case handles the possibility where a flag value is generated but never used.
		// In that case, there's no corresponding Select to overwrite the flags value,
		// so we must consider flags clobbered by the tuple-generating instruction.
		return true
	}
	return false
}

// copyFlags copies v (flag generator) into b, returns the copy.
// If v's arg is also flags, copy recursively.
func copyFlags(v *Value, b *Block) *Value {
	flagsArgs := make(map[int]*Value)
	for i, a := range v.Args {
		if a.Type.IsFlags() || a.Type.IsTuple() {
			flagsArgs[i] = copyFlags(a, b)
		}
	}
	c := v.copyInto(b)
	for i, a := range flagsArgs {
		c.SetArg(i, a)
	}
	return c
}
<|endoftext|>


<commit_before>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"time"

	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

// NATStatus is the state of NAT as detected by the ambient service.
type NATStatus int

const (
	// NAT status is unknown; this means that the ambient service has not been
	// able to decide the presence of NAT in the most recent attempt to test
	// dial through known autonat peers.  initial state.
	NATStatusUnknown NATStatus = iota
	// NAT status is publicly dialable
	NATStatusPublic
	// NAT status is private network
	NATStatusPrivate
)

var (
	AutoNATBootDelay       = 15 * time.Second
	AutoNATRetryInterval   = 90 * time.Second
	AutoNATRefreshInterval = 15 * time.Minute
	AutoNATRequestTimeout  = 30 * time.Second
)

// AutoNAT is the interface for ambient NAT autodiscovery
type AutoNAT interface {
	// Status returns the current NAT status
	Status() NATStatus
	// PublicAddr returns the public dial address when NAT status is public and an
	// error otherwise
	PublicAddr() (ma.Multiaddr, error)
}

// AmbientAutoNAT is the implementation of ambient NAT autodiscovery
type AmbientAutoNAT struct {
	ctx  context.Context
	host host.Host

	getAddrs GetAddrs

	mx     sync.Mutex
	peers  map[peer.ID][]ma.Multiaddr
	status NATStatus
	addr   ma.Multiaddr
	// Reflects the confidence on of the NATStatus being private, as a single
	// dialback may fail for reasons unrelated to NAT.
	// If it is <3, then multiple autoNAT peers may be contacted for dialback
	// If only a single autoNAT peer is known, then the confidence increases
	// for each failure until it reaches 3.
	confidence int
}

// NewAutoNAT creates a new ambient NAT autodiscovery instance attached to a host
// If getAddrs is nil, h.Addrs will be used
func NewAutoNAT(ctx context.Context, h host.Host, getAddrs GetAddrs) AutoNAT {
	if getAddrs == nil {
		getAddrs = h.Addrs
	}

	as := &AmbientAutoNAT{
		ctx:      ctx,
		host:     h,
		getAddrs: getAddrs,
		peers:    make(map[peer.ID][]ma.Multiaddr),
		status:   NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AmbientAutoNAT) Status() NATStatus {
	as.mx.Lock()
	defer as.mx.Unlock()
	return as.status
}

func (as *AmbientAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AmbientAutoNAT) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	select {
	case <-time.After(AutoNATBootDelay):
	case <-as.ctx.Done():
		return
	}

	for {
		as.autodetect()

		delay := AutoNATRefreshInterval
		if as.status == NATStatusUnknown {
			delay = AutoNATRetryInterval
		}

		select {
		case <-time.After(delay):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AmbientAutoNAT) autodetect() {
	peers := as.getPeers()

	if len(peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	cli := NewAutoNATClient(as.host, as.getAddrs)
	ctx, cancel := context.WithTimeout(as.ctx, AutoNATRequestTimeout)
	defer cancel()

	var result struct {
		sync.Mutex
		private int
		public  int
		pubaddr ma.Multiaddr
	}

	probe := 3 - as.confidence
	if probe == 0 {
		probe = 1
	}
	if probe > len(peers) {
		probe = len(peers)
	}

	var wg sync.WaitGroup

	for _, pi := range peers[:probe] {
		wg.Add(1)
		go func(pi pstore.PeerInfo) {
			as.host.Peerstore().AddAddrs(pi.ID, pi.Addrs, pstore.TempAddrTTL)
			a, err := cli.DialBack(ctx, pi.ID)

			switch {
			case err == nil:
				log.Debugf("Dialback through %s successful; public address is %s", pi.ID.Pretty(), a.String())
				result.Lock()
				result.public++
				result.pubaddr = a
				result.Unlock()

			case IsDialError(err):
				log.Debugf("Dialback through %s failed", pi.ID.Pretty())
				result.Lock()
				result.private++
				result.Unlock()

			default:
				log.Debugf("Dialback error through %s: %s", pi.ID.Pretty(), err)
			}

			wg.Done()
		}(pi)
	}

	wg.Wait()

	as.mx.Lock()
	if result.public > 0 {
		log.Debugf("NAT status is public")
		if as.status == NATStatusPrivate {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
		} else if as.confidence < 3 {
			as.confidence++
		}
		as.status = NATStatusPublic
		as.addr = result.pubaddr
	} else if result.private > 0 {
		log.Debugf("NAT status is private")
		if as.status == NATStatusPublic {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
		} else if as.confidence < 3 {
			as.confidence++
		}
		as.status = NATStatusPrivate
		as.addr = nil
	} else if as.confidence > 0 {
		// don't just flip to unknown, reduce confidence first
		as.confidence--
	} else {
		log.Debugf("NAT status is unknown")
		as.status = NATStatusUnknown
		as.addr = nil
	}
	as.mx.Unlock()
}

func (as *AmbientAutoNAT) getPeers() []pstore.PeerInfo {
	as.mx.Lock()
	defer as.mx.Unlock()

	if len(as.peers) == 0 {
		return nil
	}

	var connected, others []pstore.PeerInfo

	for p, addrs := range as.peers {
		if as.host.Network().Connectedness(p) == inet.Connected {
			connected = append(connected, pstore.PeerInfo{ID: p, Addrs: addrs})
		} else {
			others = append(others, pstore.PeerInfo{ID: p, Addrs: addrs})
		}
	}

	shufflePeers(connected)

	if len(connected) < 3 {
		shufflePeers(others)
		return append(connected, others...)
	} else {
		return connected
	}
}

func shufflePeers(peers []pstore.PeerInfo) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<commit_msg>defer wg.Done<commit_after>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"time"

	host "github.com/libp2p/go-libp2p-host"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	ma "github.com/multiformats/go-multiaddr"
)

// NATStatus is the state of NAT as detected by the ambient service.
type NATStatus int

const (
	// NAT status is unknown; this means that the ambient service has not been
	// able to decide the presence of NAT in the most recent attempt to test
	// dial through known autonat peers.  initial state.
	NATStatusUnknown NATStatus = iota
	// NAT status is publicly dialable
	NATStatusPublic
	// NAT status is private network
	NATStatusPrivate
)

var (
	AutoNATBootDelay       = 15 * time.Second
	AutoNATRetryInterval   = 90 * time.Second
	AutoNATRefreshInterval = 15 * time.Minute
	AutoNATRequestTimeout  = 30 * time.Second
)

// AutoNAT is the interface for ambient NAT autodiscovery
type AutoNAT interface {
	// Status returns the current NAT status
	Status() NATStatus
	// PublicAddr returns the public dial address when NAT status is public and an
	// error otherwise
	PublicAddr() (ma.Multiaddr, error)
}

// AmbientAutoNAT is the implementation of ambient NAT autodiscovery
type AmbientAutoNAT struct {
	ctx  context.Context
	host host.Host

	getAddrs GetAddrs

	mx     sync.Mutex
	peers  map[peer.ID][]ma.Multiaddr
	status NATStatus
	addr   ma.Multiaddr
	// Reflects the confidence on of the NATStatus being private, as a single
	// dialback may fail for reasons unrelated to NAT.
	// If it is <3, then multiple autoNAT peers may be contacted for dialback
	// If only a single autoNAT peer is known, then the confidence increases
	// for each failure until it reaches 3.
	confidence int
}

// NewAutoNAT creates a new ambient NAT autodiscovery instance attached to a host
// If getAddrs is nil, h.Addrs will be used
func NewAutoNAT(ctx context.Context, h host.Host, getAddrs GetAddrs) AutoNAT {
	if getAddrs == nil {
		getAddrs = h.Addrs
	}

	as := &AmbientAutoNAT{
		ctx:      ctx,
		host:     h,
		getAddrs: getAddrs,
		peers:    make(map[peer.ID][]ma.Multiaddr),
		status:   NATStatusUnknown,
	}

	h.Network().Notify(as)
	go as.background()

	return as
}

func (as *AmbientAutoNAT) Status() NATStatus {
	as.mx.Lock()
	defer as.mx.Unlock()
	return as.status
}

func (as *AmbientAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	as.mx.Lock()
	defer as.mx.Unlock()

	if as.status != NATStatusPublic {
		return nil, errors.New("NAT Status is not public")
	}

	return as.addr, nil
}

func (as *AmbientAutoNAT) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	select {
	case <-time.After(AutoNATBootDelay):
	case <-as.ctx.Done():
		return
	}

	for {
		as.autodetect()

		delay := AutoNATRefreshInterval
		if as.status == NATStatusUnknown {
			delay = AutoNATRetryInterval
		}

		select {
		case <-time.After(delay):
		case <-as.ctx.Done():
			return
		}
	}
}

func (as *AmbientAutoNAT) autodetect() {
	peers := as.getPeers()

	if len(peers) == 0 {
		log.Debugf("skipping NAT auto detection; no autonat peers")
		return
	}

	cli := NewAutoNATClient(as.host, as.getAddrs)
	ctx, cancel := context.WithTimeout(as.ctx, AutoNATRequestTimeout)
	defer cancel()

	var result struct {
		sync.Mutex
		private int
		public  int
		pubaddr ma.Multiaddr
	}

	probe := 3 - as.confidence
	if probe == 0 {
		probe = 1
	}
	if probe > len(peers) {
		probe = len(peers)
	}

	var wg sync.WaitGroup

	for _, pi := range peers[:probe] {
		wg.Add(1)
		go func(pi pstore.PeerInfo) {
			defer wg.Done()

			as.host.Peerstore().AddAddrs(pi.ID, pi.Addrs, pstore.TempAddrTTL)
			a, err := cli.DialBack(ctx, pi.ID)

			switch {
			case err == nil:
				log.Debugf("Dialback through %s successful; public address is %s", pi.ID.Pretty(), a.String())
				result.Lock()
				result.public++
				result.pubaddr = a
				result.Unlock()

			case IsDialError(err):
				log.Debugf("Dialback through %s failed", pi.ID.Pretty())
				result.Lock()
				result.private++
				result.Unlock()

			default:
				log.Debugf("Dialback error through %s: %s", pi.ID.Pretty(), err)
			}
		}(pi)
	}

	wg.Wait()

	as.mx.Lock()
	if result.public > 0 {
		log.Debugf("NAT status is public")
		if as.status == NATStatusPrivate {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
		} else if as.confidence < 3 {
			as.confidence++
		}
		as.status = NATStatusPublic
		as.addr = result.pubaddr
	} else if result.private > 0 {
		log.Debugf("NAT status is private")
		if as.status == NATStatusPublic {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
		} else if as.confidence < 3 {
			as.confidence++
		}
		as.status = NATStatusPrivate
		as.addr = nil
	} else if as.confidence > 0 {
		// don't just flip to unknown, reduce confidence first
		as.confidence--
	} else {
		log.Debugf("NAT status is unknown")
		as.status = NATStatusUnknown
		as.addr = nil
	}
	as.mx.Unlock()
}

func (as *AmbientAutoNAT) getPeers() []pstore.PeerInfo {
	as.mx.Lock()
	defer as.mx.Unlock()

	if len(as.peers) == 0 {
		return nil
	}

	var connected, others []pstore.PeerInfo

	for p, addrs := range as.peers {
		if as.host.Network().Connectedness(p) == inet.Connected {
			connected = append(connected, pstore.PeerInfo{ID: p, Addrs: addrs})
		} else {
			others = append(others, pstore.PeerInfo{ID: p, Addrs: addrs})
		}
	}

	shufflePeers(connected)

	if len(connected) < 3 {
		shufflePeers(others)
		return append(connected, others...)
	} else {
		return connected
	}
}

func shufflePeers(peers []pstore.PeerInfo) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}
<|endoftext|>


<commit_before>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/libp2p/go-eventbus"
	"github.com/libp2p/go-libp2p-core/event"
	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"

	logging "github.com/ipfs/go-log"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var log = logging.Logger("autonat")

// AmbientAutoNAT is the implementation of ambient NAT autodiscovery
type AmbientAutoNAT struct {
	ctx  context.Context
	host host.Host

	*config

	inboundConn  chan network.Conn
	observations chan autoNATResult
	// status is an autoNATResult reflecting current status.
	status atomic.Value
	// Reflects the confidence on of the NATStatus being private, as a single
	// dialback may fail for reasons unrelated to NAT.
	// If it is <3, then multiple autoNAT peers may be contacted for dialback
	// If only a single autoNAT peer is known, then the confidence increases
	// for each failure until it reaches 3.
	confidence   int
	lastInbound  time.Time
	lastProbeTry time.Time
	lastProbe    time.Time
	recentProbes map[peer.ID]time.Time

	service *autoNATService

	emitReachabilityChanged event.Emitter
	subscriber              event.Subscription
}

// StaticAutoNAT is a simple AutoNAT implementation when a single NAT status is desired.
type StaticAutoNAT struct {
	ctx          context.Context
	host         host.Host
	reachability network.Reachability
	service      *autoNATService
}

type autoNATResult struct {
	network.Reachability
	address ma.Multiaddr
}

// New creates a new NAT autodiscovery system attached to a host
func New(ctx context.Context, h host.Host, options ...Option) (AutoNAT, error) {
	var err error
	conf := new(config)
	conf.host = h
	conf.dialPolicy.host = h

	if err = defaults(conf); err != nil {
		return nil, err
	}
	if conf.addressFunc == nil {
		conf.addressFunc = h.Addrs
	}

	for _, o := range options {
		if err = o(conf); err != nil {
			return nil, err
		}
	}
	emitReachabilityChanged, _ := h.EventBus().Emitter(new(event.EvtLocalReachabilityChanged), eventbus.Stateful)

	var service *autoNATService
	if (!conf.forceReachability || conf.reachability == network.ReachabilityPublic) && conf.dialer != nil {
		service, err = newAutoNATService(ctx, conf)
		if err != nil {
			return nil, err
		}
		service.Enable()
	}

	if conf.forceReachability {
		emitReachabilityChanged.Emit(event.EvtLocalReachabilityChanged{Reachability: conf.reachability})

		return &StaticAutoNAT{
			ctx:          ctx,
			host:         h,
			reachability: conf.reachability,
			service:      service,
		}, nil
	}

	as := &AmbientAutoNAT{
		ctx:          ctx,
		host:         h,
		config:       conf,
		inboundConn:  make(chan network.Conn, 5),
		observations: make(chan autoNATResult, 1),

		emitReachabilityChanged: emitReachabilityChanged,
		service:                 service,
		recentProbes:            make(map[peer.ID]time.Time),
	}
	as.status.Store(autoNATResult{network.ReachabilityUnknown, nil})

	subscriber, err := as.host.EventBus().Subscribe([]interface{}{new(event.EvtLocalAddressesUpdated), new(event.EvtPeerIdentificationCompleted)})
	if err != nil {
		return nil, err
	}
	as.subscriber = subscriber

	h.Network().Notify(as)
	go as.background()

	return as, nil
}

// Status returns the AutoNAT observed reachability status.
func (as *AmbientAutoNAT) Status() network.Reachability {
	s := as.status.Load().(autoNATResult)
	return s.Reachability
}

func (as *AmbientAutoNAT) emitStatus() {
	status := as.status.Load().(autoNATResult)
	as.emitReachabilityChanged.Emit(event.EvtLocalReachabilityChanged{Reachability: status.Reachability})
}

// PublicAddr returns the publicly connectable Multiaddr of this node if one is known.
func (as *AmbientAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	s := as.status.Load().(autoNATResult)
	if s.Reachability != network.ReachabilityPublic {
		return nil, errors.New("NAT status is not public")
	}

	return s.address, nil
}

func ipInList(candidate ma.Multiaddr, list []ma.Multiaddr) bool {
	candidateIP, _ := manet.ToIP(candidate)
	for _, i := range list {
		if ip, err := manet.ToIP(i); err == nil && ip.Equal(candidateIP) {
			return true
		}
	}
	return false
}

func (as *AmbientAutoNAT) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	delay := as.config.bootDelay

	var lastAddrUpdated time.Time
	subChan := as.subscriber.Out()
	defer as.subscriber.Close()
	defer as.emitReachabilityChanged.Close()

	timer := time.NewTimer(delay)
	defer timer.Stop()
	timerRunning := true
	for {
		select {
		// new inbound connection.
		case conn := <-as.inboundConn:
			localAddrs := as.host.Addrs()
			ca := as.status.Load().(autoNATResult)
			if ca.address != nil {
				localAddrs = append(localAddrs, ca.address)
			}
			if manet.IsPublicAddr(conn.RemoteMultiaddr()) &&
				!ipInList(conn.RemoteMultiaddr(), localAddrs) {
				as.lastInbound = time.Now()
			}

		case e := <-subChan:
			switch e := e.(type) {
			case event.EvtLocalAddressesUpdated:
				if !lastAddrUpdated.Add(time.Second).After(time.Now()) {
					lastAddrUpdated = time.Now()
					if as.confidence > 1 {
						as.confidence--
					}
				}
			case event.EvtPeerIdentificationCompleted:
				if s, err := as.host.Peerstore().SupportsProtocols(e.Peer, AutoNATProto); err == nil && len(s) > 0 {
					currentStatus := as.status.Load().(autoNATResult)
					if currentStatus.Reachability == network.ReachabilityUnknown {
						as.tryProbe(e.Peer)
					}
				}
			default:
				log.Errorf("unknown event type: %T", e)
			}

		// probe finished.
		case result, ok := <-as.observations:
			if !ok {
				return
			}
			as.recordObservation(result)
		case <-timer.C:
			peer := as.getPeerToProbe()
			as.tryProbe(peer)
			timerRunning = false
		case <-as.ctx.Done():
			return
		}

		// Drain the timer channel if it hasn't fired in preparation for Resetting it.
		if timerRunning && !timer.Stop() {
			<-timer.C
		}
		timer.Reset(as.scheduleProbe())
		timerRunning = true
	}
}

func (as *AmbientAutoNAT) cleanupRecentProbes() {
	fixedNow := time.Now()
	for k, v := range as.recentProbes {
		if fixedNow.Sub(v) > as.throttlePeerPeriod {
			delete(as.recentProbes, k)
		}
	}
}

// scheduleProbe calculates when the next probe should be scheduled for.
func (as *AmbientAutoNAT) scheduleProbe() time.Duration {
	// Our baseline is a probe every 'AutoNATRefreshInterval'
	// This is modulated by:
	// * if we are in an unknown state, or have low confidence, that should drop to 'AutoNATRetryInterval'
	// * recent inbound connections (implying continued connectivity) should decrease the retry when public
	// * recent inbound connections when not public mean we should try more actively to see if we're public.
	fixedNow := time.Now()
	currentStatus := as.status.Load().(autoNATResult)

	nextProbe := fixedNow
	// Don't look for peers in the peer store more than once per second.
	if !as.lastProbeTry.IsZero() {
		backoff := as.lastProbeTry.Add(time.Second)
		if backoff.After(nextProbe) {
			nextProbe = backoff
		}
	}
	if !as.lastProbe.IsZero() {
		untilNext := as.config.refreshInterval
		if currentStatus.Reachability == network.ReachabilityUnknown {
			untilNext = as.config.retryInterval
		} else if as.confidence < 3 {
			untilNext = as.config.retryInterval
		} else if currentStatus.Reachability == network.ReachabilityPublic && as.lastInbound.After(as.lastProbe) {
			untilNext *= 2
		} else if currentStatus.Reachability != network.ReachabilityPublic && as.lastInbound.After(as.lastProbe) {
			untilNext /= 5
		}

		if as.lastProbe.Add(untilNext).After(nextProbe) {
			nextProbe = as.lastProbe.Add(untilNext)
		}
	}

	return nextProbe.Sub(fixedNow)
}

// Update the current status based on an observed result.
func (as *AmbientAutoNAT) recordObservation(observation autoNATResult) {
	currentStatus := as.status.Load().(autoNATResult)
	if observation.Reachability == network.ReachabilityPublic {
		log.Debugf("NAT status is public")
		changed := false
		if currentStatus.Reachability != network.ReachabilityPublic {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
			if as.service != nil {
				as.service.Enable()
			}
			changed = true
		} else if as.confidence < 3 {
			as.confidence++
		}
		if observation.address != nil {
			if !changed && currentStatus.address != nil && !observation.address.Equal(currentStatus.address) {
				as.confidence--
			}
			if currentStatus.address == nil || !observation.address.Equal(currentStatus.address) {
				changed = true
			}
			as.status.Store(observation)
		}
		if observation.address != nil && changed {
			as.emitStatus()
		}
	} else if observation.Reachability == network.ReachabilityPrivate {
		log.Debugf("NAT status is private")
		if currentStatus.Reachability == network.ReachabilityPublic {
			if as.confidence > 0 {
				as.confidence--
			} else {
				// we are flipping our NATStatus, so confidence drops to 0
				as.confidence = 0
				as.status.Store(observation)
				if as.service != nil {
					as.service.Disable()
				}
				as.emitStatus()
			}
		} else if as.confidence < 3 {
			as.confidence++
			as.status.Store(observation)
			if currentStatus.Reachability != network.ReachabilityPrivate {
				as.emitStatus()
			}
		}
	} else if as.confidence > 0 {
		// don't just flip to unknown, reduce confidence first
		as.confidence--
	} else {
		log.Debugf("NAT status is unknown")
		as.status.Store(autoNATResult{network.ReachabilityUnknown, nil})
		if currentStatus.Reachability != network.ReachabilityUnknown {
			if as.service != nil {
				as.service.Enable()
			}
			as.emitStatus()
		}
	}
}

func (as *AmbientAutoNAT) tryProbe(p peer.ID) bool {
	as.lastProbeTry = time.Now()
	if p.Validate() != nil {
		return false
	}

	if lastTime, ok := as.recentProbes[p]; ok {
		if time.Since(lastTime) < as.throttlePeerPeriod {
			return false
		}
	}
	as.cleanupRecentProbes()

	info := as.host.Peerstore().PeerInfo(p)

	if !as.config.dialPolicy.skipPeer(info.Addrs) {
		as.recentProbes[p] = time.Now()
		as.lastProbe = time.Now()
		go as.probe(&info)
		return true
	}
	return false
}

func (as *AmbientAutoNAT) probe(pi *peer.AddrInfo) {
	cli := NewAutoNATClient(as.host, as.config.addressFunc)
	ctx, cancel := context.WithTimeout(as.ctx, as.config.requestTimeout)
	defer cancel()

	a, err := cli.DialBack(ctx, pi.ID)

	var result autoNATResult
	switch {
	case err == nil:
		log.Debugf("Dialback through %s successful; public address is %s", pi.ID.Pretty(), a.String())
		result.Reachability = network.ReachabilityPublic
		result.address = a
	case IsDialError(err):
		log.Debugf("Dialback through %s failed", pi.ID.Pretty())
		result.Reachability = network.ReachabilityPrivate
	default:
		result.Reachability = network.ReachabilityUnknown
	}

	select {
	case as.observations <- result:
	case <-as.ctx.Done():
		return
	}
}

func (as *AmbientAutoNAT) getPeerToProbe() peer.ID {
	peers := as.host.Network().Peers()
	if len(peers) == 0 {
		return ""
	}

	candidates := make([]peer.ID, 0, len(peers))

	for _, p := range peers {
		info := as.host.Peerstore().PeerInfo(p)
		// Exclude peers which don't support the autonat protocol.
		if proto, err := as.host.Peerstore().SupportsProtocols(p, AutoNATProto); len(proto) == 0 || err != nil {
			continue
		}

		// Exclude peers in backoff.
		if lastTime, ok := as.recentProbes[p]; ok {
			if time.Since(lastTime) < as.throttlePeerPeriod {
				continue
			}
		}

		if as.config.dialPolicy.skipPeer(info.Addrs) {
			continue
		}
		candidates = append(candidates, p)
	}

	if len(candidates) == 0 {
		return ""
	}

	shufflePeers(candidates)
	return candidates[0]
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}

// Status returns the AutoNAT observed reachability status.
func (s *StaticAutoNAT) Status() network.Reachability {
	return s.reachability
}

// PublicAddr returns the publicly connectable Multiaddr of this node if one is known.
func (s *StaticAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	if s.reachability != network.ReachabilityPublic {
		return nil, errors.New("NAT status is not public")
	}
	addrs := s.host.Addrs()
	if len(addrs) > 0 {
		return s.host.Addrs()[0], nil
	}
	return nil, errors.New("No available address")
}
<commit_msg>static nat shouldn't call host.Addrs()<commit_after>package autonat

import (
	"context"
	"errors"
	"math/rand"
	"sync/atomic"
	"time"

	"github.com/libp2p/go-eventbus"
	"github.com/libp2p/go-libp2p-core/event"
	"github.com/libp2p/go-libp2p-core/host"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"

	logging "github.com/ipfs/go-log"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

var log = logging.Logger("autonat")

// AmbientAutoNAT is the implementation of ambient NAT autodiscovery
type AmbientAutoNAT struct {
	ctx  context.Context
	host host.Host

	*config

	inboundConn  chan network.Conn
	observations chan autoNATResult
	// status is an autoNATResult reflecting current status.
	status atomic.Value
	// Reflects the confidence on of the NATStatus being private, as a single
	// dialback may fail for reasons unrelated to NAT.
	// If it is <3, then multiple autoNAT peers may be contacted for dialback
	// If only a single autoNAT peer is known, then the confidence increases
	// for each failure until it reaches 3.
	confidence   int
	lastInbound  time.Time
	lastProbeTry time.Time
	lastProbe    time.Time
	recentProbes map[peer.ID]time.Time

	service *autoNATService

	emitReachabilityChanged event.Emitter
	subscriber              event.Subscription
}

// StaticAutoNAT is a simple AutoNAT implementation when a single NAT status is desired.
type StaticAutoNAT struct {
	ctx          context.Context
	host         host.Host
	reachability network.Reachability
	service      *autoNATService
}

type autoNATResult struct {
	network.Reachability
	address ma.Multiaddr
}

// New creates a new NAT autodiscovery system attached to a host
func New(ctx context.Context, h host.Host, options ...Option) (AutoNAT, error) {
	var err error
	conf := new(config)
	conf.host = h
	conf.dialPolicy.host = h

	if err = defaults(conf); err != nil {
		return nil, err
	}
	if conf.addressFunc == nil {
		conf.addressFunc = h.Addrs
	}

	for _, o := range options {
		if err = o(conf); err != nil {
			return nil, err
		}
	}
	emitReachabilityChanged, _ := h.EventBus().Emitter(new(event.EvtLocalReachabilityChanged), eventbus.Stateful)

	var service *autoNATService
	if (!conf.forceReachability || conf.reachability == network.ReachabilityPublic) && conf.dialer != nil {
		service, err = newAutoNATService(ctx, conf)
		if err != nil {
			return nil, err
		}
		service.Enable()
	}

	if conf.forceReachability {
		emitReachabilityChanged.Emit(event.EvtLocalReachabilityChanged{Reachability: conf.reachability})

		return &StaticAutoNAT{
			ctx:          ctx,
			host:         h,
			reachability: conf.reachability,
			service:      service,
		}, nil
	}

	as := &AmbientAutoNAT{
		ctx:          ctx,
		host:         h,
		config:       conf,
		inboundConn:  make(chan network.Conn, 5),
		observations: make(chan autoNATResult, 1),

		emitReachabilityChanged: emitReachabilityChanged,
		service:                 service,
		recentProbes:            make(map[peer.ID]time.Time),
	}
	as.status.Store(autoNATResult{network.ReachabilityUnknown, nil})

	subscriber, err := as.host.EventBus().Subscribe([]interface{}{new(event.EvtLocalAddressesUpdated), new(event.EvtPeerIdentificationCompleted)})
	if err != nil {
		return nil, err
	}
	as.subscriber = subscriber

	h.Network().Notify(as)
	go as.background()

	return as, nil
}

// Status returns the AutoNAT observed reachability status.
func (as *AmbientAutoNAT) Status() network.Reachability {
	s := as.status.Load().(autoNATResult)
	return s.Reachability
}

func (as *AmbientAutoNAT) emitStatus() {
	status := as.status.Load().(autoNATResult)
	as.emitReachabilityChanged.Emit(event.EvtLocalReachabilityChanged{Reachability: status.Reachability})
}

// PublicAddr returns the publicly connectable Multiaddr of this node if one is known.
func (as *AmbientAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	s := as.status.Load().(autoNATResult)
	if s.Reachability != network.ReachabilityPublic {
		return nil, errors.New("NAT status is not public")
	}

	return s.address, nil
}

func ipInList(candidate ma.Multiaddr, list []ma.Multiaddr) bool {
	candidateIP, _ := manet.ToIP(candidate)
	for _, i := range list {
		if ip, err := manet.ToIP(i); err == nil && ip.Equal(candidateIP) {
			return true
		}
	}
	return false
}

func (as *AmbientAutoNAT) background() {
	// wait a bit for the node to come online and establish some connections
	// before starting autodetection
	delay := as.config.bootDelay

	var lastAddrUpdated time.Time
	subChan := as.subscriber.Out()
	defer as.subscriber.Close()
	defer as.emitReachabilityChanged.Close()

	timer := time.NewTimer(delay)
	defer timer.Stop()
	timerRunning := true
	for {
		select {
		// new inbound connection.
		case conn := <-as.inboundConn:
			localAddrs := as.host.Addrs()
			ca := as.status.Load().(autoNATResult)
			if ca.address != nil {
				localAddrs = append(localAddrs, ca.address)
			}
			if manet.IsPublicAddr(conn.RemoteMultiaddr()) &&
				!ipInList(conn.RemoteMultiaddr(), localAddrs) {
				as.lastInbound = time.Now()
			}

		case e := <-subChan:
			switch e := e.(type) {
			case event.EvtLocalAddressesUpdated:
				if !lastAddrUpdated.Add(time.Second).After(time.Now()) {
					lastAddrUpdated = time.Now()
					if as.confidence > 1 {
						as.confidence--
					}
				}
			case event.EvtPeerIdentificationCompleted:
				if s, err := as.host.Peerstore().SupportsProtocols(e.Peer, AutoNATProto); err == nil && len(s) > 0 {
					currentStatus := as.status.Load().(autoNATResult)
					if currentStatus.Reachability == network.ReachabilityUnknown {
						as.tryProbe(e.Peer)
					}
				}
			default:
				log.Errorf("unknown event type: %T", e)
			}

		// probe finished.
		case result, ok := <-as.observations:
			if !ok {
				return
			}
			as.recordObservation(result)
		case <-timer.C:
			peer := as.getPeerToProbe()
			as.tryProbe(peer)
			timerRunning = false
		case <-as.ctx.Done():
			return
		}

		// Drain the timer channel if it hasn't fired in preparation for Resetting it.
		if timerRunning && !timer.Stop() {
			<-timer.C
		}
		timer.Reset(as.scheduleProbe())
		timerRunning = true
	}
}

func (as *AmbientAutoNAT) cleanupRecentProbes() {
	fixedNow := time.Now()
	for k, v := range as.recentProbes {
		if fixedNow.Sub(v) > as.throttlePeerPeriod {
			delete(as.recentProbes, k)
		}
	}
}

// scheduleProbe calculates when the next probe should be scheduled for.
func (as *AmbientAutoNAT) scheduleProbe() time.Duration {
	// Our baseline is a probe every 'AutoNATRefreshInterval'
	// This is modulated by:
	// * if we are in an unknown state, or have low confidence, that should drop to 'AutoNATRetryInterval'
	// * recent inbound connections (implying continued connectivity) should decrease the retry when public
	// * recent inbound connections when not public mean we should try more actively to see if we're public.
	fixedNow := time.Now()
	currentStatus := as.status.Load().(autoNATResult)

	nextProbe := fixedNow
	// Don't look for peers in the peer store more than once per second.
	if !as.lastProbeTry.IsZero() {
		backoff := as.lastProbeTry.Add(time.Second)
		if backoff.After(nextProbe) {
			nextProbe = backoff
		}
	}
	if !as.lastProbe.IsZero() {
		untilNext := as.config.refreshInterval
		if currentStatus.Reachability == network.ReachabilityUnknown {
			untilNext = as.config.retryInterval
		} else if as.confidence < 3 {
			untilNext = as.config.retryInterval
		} else if currentStatus.Reachability == network.ReachabilityPublic && as.lastInbound.After(as.lastProbe) {
			untilNext *= 2
		} else if currentStatus.Reachability != network.ReachabilityPublic && as.lastInbound.After(as.lastProbe) {
			untilNext /= 5
		}

		if as.lastProbe.Add(untilNext).After(nextProbe) {
			nextProbe = as.lastProbe.Add(untilNext)
		}
	}

	return nextProbe.Sub(fixedNow)
}

// Update the current status based on an observed result.
func (as *AmbientAutoNAT) recordObservation(observation autoNATResult) {
	currentStatus := as.status.Load().(autoNATResult)
	if observation.Reachability == network.ReachabilityPublic {
		log.Debugf("NAT status is public")
		changed := false
		if currentStatus.Reachability != network.ReachabilityPublic {
			// we are flipping our NATStatus, so confidence drops to 0
			as.confidence = 0
			if as.service != nil {
				as.service.Enable()
			}
			changed = true
		} else if as.confidence < 3 {
			as.confidence++
		}
		if observation.address != nil {
			if !changed && currentStatus.address != nil && !observation.address.Equal(currentStatus.address) {
				as.confidence--
			}
			if currentStatus.address == nil || !observation.address.Equal(currentStatus.address) {
				changed = true
			}
			as.status.Store(observation)
		}
		if observation.address != nil && changed {
			as.emitStatus()
		}
	} else if observation.Reachability == network.ReachabilityPrivate {
		log.Debugf("NAT status is private")
		if currentStatus.Reachability == network.ReachabilityPublic {
			if as.confidence > 0 {
				as.confidence--
			} else {
				// we are flipping our NATStatus, so confidence drops to 0
				as.confidence = 0
				as.status.Store(observation)
				if as.service != nil {
					as.service.Disable()
				}
				as.emitStatus()
			}
		} else if as.confidence < 3 {
			as.confidence++
			as.status.Store(observation)
			if currentStatus.Reachability != network.ReachabilityPrivate {
				as.emitStatus()
			}
		}
	} else if as.confidence > 0 {
		// don't just flip to unknown, reduce confidence first
		as.confidence--
	} else {
		log.Debugf("NAT status is unknown")
		as.status.Store(autoNATResult{network.ReachabilityUnknown, nil})
		if currentStatus.Reachability != network.ReachabilityUnknown {
			if as.service != nil {
				as.service.Enable()
			}
			as.emitStatus()
		}
	}
}

func (as *AmbientAutoNAT) tryProbe(p peer.ID) bool {
	as.lastProbeTry = time.Now()
	if p.Validate() != nil {
		return false
	}

	if lastTime, ok := as.recentProbes[p]; ok {
		if time.Since(lastTime) < as.throttlePeerPeriod {
			return false
		}
	}
	as.cleanupRecentProbes()

	info := as.host.Peerstore().PeerInfo(p)

	if !as.config.dialPolicy.skipPeer(info.Addrs) {
		as.recentProbes[p] = time.Now()
		as.lastProbe = time.Now()
		go as.probe(&info)
		return true
	}
	return false
}

func (as *AmbientAutoNAT) probe(pi *peer.AddrInfo) {
	cli := NewAutoNATClient(as.host, as.config.addressFunc)
	ctx, cancel := context.WithTimeout(as.ctx, as.config.requestTimeout)
	defer cancel()

	a, err := cli.DialBack(ctx, pi.ID)

	var result autoNATResult
	switch {
	case err == nil:
		log.Debugf("Dialback through %s successful; public address is %s", pi.ID.Pretty(), a.String())
		result.Reachability = network.ReachabilityPublic
		result.address = a
	case IsDialError(err):
		log.Debugf("Dialback through %s failed", pi.ID.Pretty())
		result.Reachability = network.ReachabilityPrivate
	default:
		result.Reachability = network.ReachabilityUnknown
	}

	select {
	case as.observations <- result:
	case <-as.ctx.Done():
		return
	}
}

func (as *AmbientAutoNAT) getPeerToProbe() peer.ID {
	peers := as.host.Network().Peers()
	if len(peers) == 0 {
		return ""
	}

	candidates := make([]peer.ID, 0, len(peers))

	for _, p := range peers {
		info := as.host.Peerstore().PeerInfo(p)
		// Exclude peers which don't support the autonat protocol.
		if proto, err := as.host.Peerstore().SupportsProtocols(p, AutoNATProto); len(proto) == 0 || err != nil {
			continue
		}

		// Exclude peers in backoff.
		if lastTime, ok := as.recentProbes[p]; ok {
			if time.Since(lastTime) < as.throttlePeerPeriod {
				continue
			}
		}

		if as.config.dialPolicy.skipPeer(info.Addrs) {
			continue
		}
		candidates = append(candidates, p)
	}

	if len(candidates) == 0 {
		return ""
	}

	shufflePeers(candidates)
	return candidates[0]
}

func shufflePeers(peers []peer.ID) {
	for i := range peers {
		j := rand.Intn(i + 1)
		peers[i], peers[j] = peers[j], peers[i]
	}
}

// Status returns the AutoNAT observed reachability status.
func (s *StaticAutoNAT) Status() network.Reachability {
	return s.reachability
}

// PublicAddr returns the publicly connectable Multiaddr of this node if one is known.
func (s *StaticAutoNAT) PublicAddr() (ma.Multiaddr, error) {
	if s.reachability != network.ReachabilityPublic {
		return nil, errors.New("NAT status is not public")
	}
	return nil, errors.New("No available address")
}
<|endoftext|>


<commit_before>package libp2pquic

import (
	"net"
	"sync"
	"sync/atomic"

	"github.com/vishvananda/netlink"
)

type reuseConn struct {
	net.PacketConn
	refCount int32 // to be used as an atomic
}

func newReuseConn(conn net.PacketConn) *reuseConn {
	return &reuseConn{PacketConn: conn}
}

func (c *reuseConn) IncreaseCount() { atomic.AddInt32(&c.refCount, 1) }
func (c *reuseConn) DecreaseCount() { atomic.AddInt32(&c.refCount, -1) }
func (c *reuseConn) GetCount() int  { return int(atomic.LoadInt32(&c.refCount)) }

type reuse struct {
	mutex sync.Mutex

	handle *netlink.Handle // Only set on Linux. nil on other systems.

	unicast map[string] /* IP.String() */ map[int] /* port */ *reuseConn
	// global contains connections that are listening on 0.0.0.0 / ::
	global map[int]*reuseConn
}

func newReuse() (*reuse, error) {
	// On non-Linux systems, this will return ErrNotImplemented.
	handle, err := netlink.NewHandle()
	if err == netlink.ErrNotImplemented {
		handle = nil
	} else if err != nil {
		return nil, err
	}
	return &reuse{
		unicast: make(map[string]map[int]*reuseConn),
		global:  make(map[int]*reuseConn),
		handle:  handle,
	}, nil
}

// Get the source IP that the kernel would use for dialing.
// This only works on Linux.
// On other systems, this returns an empty slice of IP addresses.
func (r *reuse) getSourceIPs(network string, raddr *net.UDPAddr) ([]net.IP, error) {
	if r.handle == nil {
		return nil, nil
	}

	routes, err := r.handle.RouteGet(raddr.IP)
	if err != nil {
		return nil, err
	}

	ips := make([]net.IP, 0, len(routes))
	for _, route := range routes {
		ips = append(ips, route.Src)
	}
	return ips, nil
}

func (r *reuse) Dial(network string, raddr *net.UDPAddr) (*reuseConn, error) {
	ips, err := r.getSourceIPs(network, raddr)
	if err != nil {
		return nil, err
	}

	r.mutex.Lock()
	defer r.mutex.Unlock()

	conn, err := r.dialLocked(network, raddr, ips)
	if err != nil {
		return nil, err
	}
	conn.IncreaseCount()
	return conn, nil
}

func (r *reuse) dialLocked(network string, raddr *net.UDPAddr, ips []net.IP) (*reuseConn, error) {
	for _, ip := range ips {
		// We already have at least one suitable connection...
		if conns, ok := r.unicast[ip.String()]; ok {
			// ... we don't care which port we're dialing from. Just use the first.
			for _, c := range conns {
				return c, nil
			}
		}
	}

	// Use a connection listening on 0.0.0.0 (or ::).
	// Again, we don't care about the port number.
	for _, conn := range r.global {
		return conn, nil
	}

	// We don't have a connection that we can use for dialing.
	// Dial a new connection from a random port.
	var addr *net.UDPAddr
	switch network {
	case "udp4":
		addr = &net.UDPAddr{IP: net.IPv4zero, Port: 0}
	case "udp6":
		addr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
	}
	conn, err := net.ListenUDP(network, addr)
	if err != nil {
		return nil, err
	}
	rconn := newReuseConn(conn)
	r.global[conn.LocalAddr().(*net.UDPAddr).Port] = rconn
	return rconn, nil
}

func (r *reuse) Listen(network string, laddr *net.UDPAddr) (*reuseConn, error) {
	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}
	localAddr := conn.LocalAddr().(*net.UDPAddr)

	rconn := newReuseConn(conn)
	rconn.IncreaseCount()

	r.mutex.Lock()
	defer r.mutex.Unlock()

	// Deal with listen on a global address
	if laddr.IP.IsUnspecified() {
		// The kernel already checked that the laddr is not already listen
		// so we need not check here (when we create ListenUDP).
		r.global[laddr.Port] = rconn
		return rconn, err
	}

	// Deal with listen on a unicast address
	if _, ok := r.unicast[localAddr.IP.String()]; !ok {
		r.unicast[laddr.IP.String()] = make(map[int]*reuseConn)
	}

	// The kernel already checked that the laddr is not already listen
	// so we need not check here (when we create ListenUDP).
	r.unicast[laddr.IP.String()][localAddr.Port] = rconn
	return rconn, err
}
<commit_msg>fix saving of listening connections when listening on IP:0<commit_after>package libp2pquic

import (
	"net"
	"sync"
	"sync/atomic"

	"github.com/vishvananda/netlink"
)

type reuseConn struct {
	net.PacketConn
	refCount int32 // to be used as an atomic
}

func newReuseConn(conn net.PacketConn) *reuseConn {
	return &reuseConn{PacketConn: conn}
}

func (c *reuseConn) IncreaseCount() { atomic.AddInt32(&c.refCount, 1) }
func (c *reuseConn) DecreaseCount() { atomic.AddInt32(&c.refCount, -1) }
func (c *reuseConn) GetCount() int  { return int(atomic.LoadInt32(&c.refCount)) }

type reuse struct {
	mutex sync.Mutex

	handle *netlink.Handle // Only set on Linux. nil on other systems.

	unicast map[string] /* IP.String() */ map[int] /* port */ *reuseConn
	// global contains connections that are listening on 0.0.0.0 / ::
	global map[int]*reuseConn
}

func newReuse() (*reuse, error) {
	// On non-Linux systems, this will return ErrNotImplemented.
	handle, err := netlink.NewHandle()
	if err == netlink.ErrNotImplemented {
		handle = nil
	} else if err != nil {
		return nil, err
	}
	return &reuse{
		unicast: make(map[string]map[int]*reuseConn),
		global:  make(map[int]*reuseConn),
		handle:  handle,
	}, nil
}

// Get the source IP that the kernel would use for dialing.
// This only works on Linux.
// On other systems, this returns an empty slice of IP addresses.
func (r *reuse) getSourceIPs(network string, raddr *net.UDPAddr) ([]net.IP, error) {
	if r.handle == nil {
		return nil, nil
	}

	routes, err := r.handle.RouteGet(raddr.IP)
	if err != nil {
		return nil, err
	}

	ips := make([]net.IP, 0, len(routes))
	for _, route := range routes {
		ips = append(ips, route.Src)
	}
	return ips, nil
}

func (r *reuse) Dial(network string, raddr *net.UDPAddr) (*reuseConn, error) {
	ips, err := r.getSourceIPs(network, raddr)
	if err != nil {
		return nil, err
	}

	r.mutex.Lock()
	defer r.mutex.Unlock()

	conn, err := r.dialLocked(network, raddr, ips)
	if err != nil {
		return nil, err
	}
	conn.IncreaseCount()
	return conn, nil
}

func (r *reuse) dialLocked(network string, raddr *net.UDPAddr, ips []net.IP) (*reuseConn, error) {
	for _, ip := range ips {
		// We already have at least one suitable connection...
		if conns, ok := r.unicast[ip.String()]; ok {
			// ... we don't care which port we're dialing from. Just use the first.
			for _, c := range conns {
				return c, nil
			}
		}
	}

	// Use a connection listening on 0.0.0.0 (or ::).
	// Again, we don't care about the port number.
	for _, conn := range r.global {
		return conn, nil
	}

	// We don't have a connection that we can use for dialing.
	// Dial a new connection from a random port.
	var addr *net.UDPAddr
	switch network {
	case "udp4":
		addr = &net.UDPAddr{IP: net.IPv4zero, Port: 0}
	case "udp6":
		addr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
	}
	conn, err := net.ListenUDP(network, addr)
	if err != nil {
		return nil, err
	}
	rconn := newReuseConn(conn)
	r.global[conn.LocalAddr().(*net.UDPAddr).Port] = rconn
	return rconn, nil
}

func (r *reuse) Listen(network string, laddr *net.UDPAddr) (*reuseConn, error) {
	conn, err := net.ListenUDP(network, laddr)
	if err != nil {
		return nil, err
	}
	localAddr := conn.LocalAddr().(*net.UDPAddr)

	rconn := newReuseConn(conn)
	rconn.IncreaseCount()

	r.mutex.Lock()
	defer r.mutex.Unlock()

	// Deal with listen on a global address
	if localAddr.IP.IsUnspecified() {
		// The kernel already checked that the laddr is not already listen
		// so we need not check here (when we create ListenUDP).
		r.global[localAddr.Port] = rconn
		return rconn, err
	}

	// Deal with listen on a unicast address
	if _, ok := r.unicast[localAddr.IP.String()]; !ok {
		r.unicast[localAddr.IP.String()] = make(map[int]*reuseConn)
	}

	// The kernel already checked that the laddr is not already listen
	// so we need not check here (when we create ListenUDP).
	r.unicast[localAddr.IP.String()][localAddr.Port] = rconn
	return rconn, err
}
<|endoftext|>